daggerml.contrib.adapters

  1from __future__ import annotations
  2
  3import argparse
  4import json
  5import sys
  6import time
  7from pathlib import Path
  8from threading import Lock
  9from typing import Any
 10from urllib.parse import urlparse
 11from warnings import warn
 12
 13from daggerml import Dml, Ref, Runnable
 14from daggerml._core import AdapterCancelResponse, AdapterCleanupResponse, AdapterInvokeResponse
 15from daggerml.api import DmlRepoError, _entry_points
 16from daggerml.contrib.s3 import S3Store, is_s3_uri
 17from daggerml.util import get_client
 18
 19
 20class AdapterBase:
 21    name = ""
 22
 23    @classmethod
 24    def resolve_runnable(cls, uri, kwargs, sub) -> Runnable:
 25        from daggerml.contrib.executors._base import get_executor
 26
 27        return get_executor(cls.name, uri).resolve_runnable(uri, kwargs, sub)
 28
 29    @classmethod
 30    def send(cls, **kw) -> AdapterInvokeResponse | AdapterCleanupResponse | AdapterCancelResponse:
 31        raise NotImplementedError("Adapter send method is not implemented")
 32
 33    @classmethod
 34    def _read_input(cls, input_path: str) -> str:
 35        if input_path == "-":
 36            return sys.stdin.read()
 37        if is_s3_uri(input_path):
 38            return S3Store().get(input_path).decode("utf-8")
 39        return Path(input_path).read_text()
 40
 41    @classmethod
 42    def _write_output(cls, output_path: str, data: str) -> None:
 43        if output_path == "-":
 44            sys.stdout.write(data)
 45            if not data.endswith("\n"):
 46                sys.stdout.write("\n")
 47            sys.stdout.flush()
 48            return
 49        if is_s3_uri(output_path):
 50            parsed = urlparse(output_path)
 51            bucket = parsed.netloc
 52            key = parsed.path.lstrip("/")
 53            get_client("s3").put_object(
 54                Bucket=bucket,
 55                Key=key,
 56                Body=data.encode("utf-8"),
 57                ContentType="application/json",
 58            )
 59            return
 60        Path(output_path).write_text(data)
 61
 62    @classmethod
 63    def cli(cls, argv: list[str] | None = None) -> int:
 64        parser = argparse.ArgumentParser(description=f"{cls.__name__} CLI")
 65        parser.add_argument("-i", "--input", default="-")
 66        parser.add_argument("-o", "--output", default="-")
 67        parser.add_argument("--poll", action="store_true")
 68        # FIXME: `--poll` make `cancel` difficult. We should allow for coordination between caller and this loop.
 69        args = parser.parse_args(argv)
 70        raw = cls._read_input(args.input)
 71        payload = json.loads(raw)
 72        result = cls.send(**payload)
 73        while args.poll and payload.get("operation") == "invoke" and result.get("status") == "retry":
 74            state = result.get("adapter_state")
 75            if not isinstance(state, dict):
 76                raise DmlRepoError("Retry adapter response requires object adapter_state")
 77            payload["adapter_state"] = state
 78            time.sleep(0.1)
 79            result = cls.send(**payload)
 80        if args.poll and payload.get("operation") == "invoke" and result.get("status") == "success":
 81            payload["adapter_state"] = result.get("adapter_state", payload["adapter_state"])
 82            record = Dml(remote_root=payload["remote"]["root"]).runtime.read_execution_record(
 83                Ref(f"index:{payload['execution_id']}")
 84            )
 85            result_ref = record["state"]["result_ref"]
 86            if result_ref is None:
 87                raise DmlRepoError("Successful nested invoke did not publish a result")
 88            cleanup_payload = {**payload, "operation": "cleanup", "result_ref": result_ref}
 89            try:
 90                cleanup_result = cls.send(**cleanup_payload)
 91                while cleanup_result.get("status") == "retry":
 92                    state = cleanup_result.get("adapter_state")
 93                    if not isinstance(state, dict):
 94                        raise DmlRepoError("Retry adapter response requires object adapter_state")
 95                    cleanup_payload["adapter_state"] = state
 96                    time.sleep(0.1)
 97                    cleanup_result = cls.send(**cleanup_payload)
 98                if cleanup_result.get("status") != "success":
 99                    print(
100                        f"Nested cleanup failed for {payload['execution_id']}: {cleanup_result.get('error')}",
101                        file=sys.stderr,
102                    )
103            except Exception as exc:
104                print(f"Nested cleanup failed for {payload['execution_id']}: {exc}", file=sys.stderr)
105        if args.poll and payload.get("operation") == "invoke":
106            result = {key: value for key, value in result.items() if key != "adapter_state"}
107        cls._write_output(args.output, json.dumps(result))
108        return 0
109
110
111class LocalAdapter(AdapterBase):
112    name = "local"
113    executable = "dml-local-adapter"
114
115    @classmethod
116    def send(cls, **kw) -> AdapterInvokeResponse | AdapterCleanupResponse | AdapterCancelResponse:
117        from daggerml.contrib.executors._base import get_executor
118
119        return get_executor("local", kw["runnable"]["target"]["uri"]).handle(**kw)
120
121
122class LambdaAdapter(AdapterBase):
123    name = "lambda"
124    executable = "dml-lambda-adapter"
125
126    @classmethod
127    def send(cls, **kw) -> AdapterInvokeResponse | AdapterCleanupResponse | AdapterCancelResponse:
128        client = get_client("lambda")
129        try:
130            response = client.invoke(
131                FunctionName=kw["runnable"]["target"]["uri"],
132                InvocationType="RequestResponse",
133                Payload=json.dumps(kw, separators=(",", ":"), sort_keys=True).encode("utf-8"),
134            )
135        except Exception as exc:
136            code = getattr(exc, "response", {}).get("Error", {}).get("Code")
137            if code not in {"TooManyRequestsException", "ThrottlingException"}:
138                raise
139            headers = getattr(exc, "response", {}).get("ResponseMetadata", {}).get("HTTPHeaders", {})
140            try:
141                retry_after_ms = max(0, int(float(headers.get("retry-after")) * 1000))
142            except (TypeError, ValueError):
143                retry_after_ms = None
144            result: AdapterInvokeResponse = {
145                "status": "retry",
146                "error": None,
147                "adapter_state": kw.get("adapter_state") if isinstance(kw.get("adapter_state"), dict) else {},
148            }
149            if retry_after_ms is not None:
150                result["retry_after_ms"] = retry_after_ms
151            return result
152        stream = response.get("Payload")
153        if stream is None:
154            raise DmlRepoError("Lambda adapter invoke response missing Payload")
155        return json.loads(stream.read().decode("utf-8"))
156
157
158################################################################################
159############################### Adapter registry ###############################
160################################################################################
161ADAPTER_ENTRYPOINT_GROUP = "daggerml.contrib.adapters"
162
163_LOCK = Lock()
164_ADAPTER_SPECS: dict[str, str] = {}
165_PLUGINS_LOADED = False
166
167
168def load_adapter_plugins() -> None:
169    global _PLUGINS_LOADED
170    if _PLUGINS_LOADED:
171        return
172    with _LOCK:
173        if _PLUGINS_LOADED:
174            return
175        for ep in _entry_points(ADAPTER_ENTRYPOINT_GROUP):
176            try:
177                loaded = ep.load()
178                if loaded.name in _ADAPTER_SPECS:
179                    # warn about duplicate adapter registration but allow the last one to win
180                    warn(
181                        f"Adapter: '{loaded.name}' is overwriting existing '{ep.name} ({ep.value})'",
182                        stacklevel=2,
183                    )
184                _ADAPTER_SPECS[loaded.name] = loaded
185            except Exception as e:
186                raise DmlRepoError(f"Adapter plugin '{ep.name} ({ep.value})' failed: {e}") from e
187        _PLUGINS_LOADED = True
188
189
190def get_adapter(name: str) -> Any:
191    load_adapter_plugins()
192    spec = _ADAPTER_SPECS.get(name)
193    if spec is None:
194        raise DmlRepoError(f"Adapter '{name}' is not registered")
195    return spec
196
197
198def list_adapters() -> list[str]:
199    load_adapter_plugins()
200    return sorted(_ADAPTER_SPECS.keys())