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            record = Dml(remote_root=payload["remote"]["root"]).runtime.read_execution_record(
 82                Ref(f"index:{payload['execution_id']}")
 83            )
 84            result_ref = record["state"]["result_ref"]
 85            if result_ref is None:
 86                raise DmlRepoError("Successful nested invoke did not publish a result")
 87            cleanup_payload = {**payload, "operation": "cleanup", "result_ref": result_ref}
 88            try:
 89                result = cls.send(**cleanup_payload)
 90                while result.get("status") == "retry":
 91                    state = result.get("adapter_state")
 92                    if not isinstance(state, dict):
 93                        raise DmlRepoError("Retry adapter response requires object adapter_state")
 94                    cleanup_payload["adapter_state"] = state
 95                    time.sleep(0.1)
 96                    result = cls.send(**cleanup_payload)
 97            except Exception as exc:
 98                result = {"status": "failure", "error": f"Nested cleanup failed: {exc}"}
 99        cls._write_output(args.output, json.dumps(result))
100        return 0
101
102
103class LocalAdapter(AdapterBase):
104    name = "local"
105    executable = "dml-local-adapter"
106
107    @classmethod
108    def send(cls, **kw) -> AdapterInvokeResponse | AdapterCleanupResponse | AdapterCancelResponse:
109        from daggerml.contrib.executors._base import get_executor
110
111        return get_executor("local", kw["runnable"]["target"]["uri"]).handle(**kw)
112
113
114class LambdaAdapter(AdapterBase):
115    name = "lambda"
116    executable = "dml-lambda-adapter"
117
118    @classmethod
119    def send(cls, **kw) -> AdapterInvokeResponse | AdapterCleanupResponse | AdapterCancelResponse:
120        client = get_client("lambda")
121        try:
122            response = client.invoke(
123                FunctionName=kw["runnable"]["target"]["uri"],
124                InvocationType="RequestResponse",
125                Payload=json.dumps(kw, separators=(",", ":"), sort_keys=True).encode("utf-8"),
126            )
127        except Exception as exc:
128            code = getattr(exc, "response", {}).get("Error", {}).get("Code")
129            if code not in {"TooManyRequestsException", "ThrottlingException"}:
130                raise
131            headers = getattr(exc, "response", {}).get("ResponseMetadata", {}).get("HTTPHeaders", {})
132            try:
133                retry_after_ms = max(0, int(float(headers.get("retry-after")) * 1000))
134            except (TypeError, ValueError):
135                retry_after_ms = None
136            result: AdapterInvokeResponse = {
137                "status": "retry",
138                "error": None,
139                "adapter_state": kw.get("adapter_state") if isinstance(kw.get("adapter_state"), dict) else {},
140            }
141            if retry_after_ms is not None:
142                result["retry_after_ms"] = retry_after_ms
143            return result
144        stream = response.get("Payload")
145        if stream is None:
146            raise DmlRepoError("Lambda adapter invoke response missing Payload")
147        return json.loads(stream.read().decode("utf-8"))
148
149
150################################################################################
151############################### Adapter registry ###############################
152################################################################################
153ADAPTER_ENTRYPOINT_GROUP = "daggerml.contrib.adapters"
154
155_LOCK = Lock()
156_ADAPTER_SPECS: dict[str, str] = {}
157_PLUGINS_LOADED = False
158
159
160def load_adapter_plugins() -> None:
161    global _PLUGINS_LOADED
162    if _PLUGINS_LOADED:
163        return
164    with _LOCK:
165        if _PLUGINS_LOADED:
166            return
167        for ep in _entry_points(ADAPTER_ENTRYPOINT_GROUP):
168            try:
169                loaded = ep.load()
170                if loaded.name in _ADAPTER_SPECS:
171                    # warn about duplicate adapter registration but allow the last one to win
172                    warn(
173                        f"Adapter: '{loaded.name}' is overwriting existing '{ep.name} ({ep.value})'",
174                        stacklevel=2,
175                    )
176                _ADAPTER_SPECS[loaded.name] = loaded
177            except Exception as e:
178                raise DmlRepoError(f"Adapter plugin '{ep.name} ({ep.value})' failed: {e}") from e
179        _PLUGINS_LOADED = True
180
181
182def get_adapter(name: str) -> Any:
183    load_adapter_plugins()
184    spec = _ADAPTER_SPECS.get(name)
185    if spec is None:
186        raise DmlRepoError(f"Adapter '{name}' is not registered")
187    return spec
188
189
190def list_adapters() -> list[str]:
191    load_adapter_plugins()
192    return sorted(_ADAPTER_SPECS.keys())