daggerml.contrib.supervisor

  1from __future__ import annotations
  2
  3import argparse
  4import json
  5import logging
  6import os
  7import re
  8import signal
  9import subprocess
 10import sys
 11import tempfile
 12import threading
 13import time
 14from pathlib import Path
 15from typing import Any
 16
 17from daggerml import Dml, Ref
 18from daggerml.api import DmlRepoError
 19
 20logger = logging.getLogger(__name__)
 21
 22_CLOUDWATCH_LOG_GROUP = "dml"
 23_CLOUDWATCH_MAX_BATCH_BYTES = 1_048_576
 24_CLOUDWATCH_MAX_MESSAGE_BYTES = 1_048_576
 25_CLOUDWATCH_EVENT_OVERHEAD_BYTES = 26
 26_CLOUDWATCH_MAX_BATCH_COUNT = 10_000
 27
 28
 29def _create_logs_client() -> Any:
 30    from daggerml.util import get_client
 31
 32    return get_client("logs")
 33
 34
 35def _resource_already_exists(exc: Exception) -> bool:
 36    return getattr(exc, "response", {}).get("Error", {}).get("Code") == "ResourceAlreadyExistsException"
 37
 38
 39class _CloudWatchStream:
 40    def __init__(self, *, cache_key: str, execution_id: str, stream_kind: str):
 41        self.cache_key = cache_key
 42        self.execution_id = execution_id
 43        self.stream_kind = stream_kind
 44        self.stream_name = f"/run/{cache_key}/{stream_kind}"
 45        self._client: Any | None = None
 46        self._enabled = True
 47        self._sequence_token: str | None = None
 48        self._pending_events: list[dict[str, Any]] = []
 49        self._pending_bytes = 0
 50        self._lock = threading.Lock()
 51        self._init_client()
 52        self.emit_lifecycle(event="start")
 53
 54    @staticmethod
 55    def _event_bytes(message: str) -> int:
 56        return len(message.encode("utf-8")) + _CLOUDWATCH_EVENT_OVERHEAD_BYTES
 57
 58    @staticmethod
 59    def _split_message(message: str) -> list[str]:
 60        encoded = message.encode("utf-8")
 61        if len(encoded) <= _CLOUDWATCH_MAX_MESSAGE_BYTES:
 62            return [message]
 63
 64        chunks: list[str] = []
 65        start = 0
 66        while start < len(encoded):
 67            end = min(start + _CLOUDWATCH_MAX_MESSAGE_BYTES, len(encoded))
 68            while end > start:
 69                try:
 70                    chunks.append(encoded[start:end].decode("utf-8"))
 71                    start = end
 72                    break
 73                except UnicodeDecodeError:
 74                    end -= 1
 75            else:
 76                raise AssertionError("failed to split UTF-8 message into valid chunks")
 77        return chunks
 78
 79    def _flush_locked(self) -> None:
 80        if not self._enabled or self._client is None or not self._pending_events:
 81            return
 82        params: dict[str, Any] = {
 83            "logGroupName": _CLOUDWATCH_LOG_GROUP,
 84            "logStreamName": self.stream_name,
 85            "logEvents": list(self._pending_events),
 86        }
 87        if self._sequence_token is not None:
 88            params["sequenceToken"] = self._sequence_token
 89        try:
 90            response = self._client.put_log_events(**params)
 91            self._sequence_token = response.get("nextSequenceToken")
 92            self._pending_events.clear()
 93            self._pending_bytes = 0
 94        except Exception as exc:
 95            self._pending_events.clear()
 96            self._pending_bytes = 0
 97            self._disable(f"event delivery failed: {exc}")
 98
 99    def _init_client(self) -> None:
100        try:
101            client = _create_logs_client()
102            try:
103                client.create_log_group(logGroupName=_CLOUDWATCH_LOG_GROUP)
104            except Exception as exc:
105                if not _resource_already_exists(exc):
106                    raise
107            try:
108                client.create_log_stream(logGroupName=_CLOUDWATCH_LOG_GROUP, logStreamName=self.stream_name)
109            except Exception as exc:
110                if not _resource_already_exists(exc):
111                    raise
112            self._client = client
113        except Exception as exc:
114            self._disable(f"initialization failed: {exc}")
115
116    def _disable(self, reason: str) -> None:
117        if not self._enabled:
118            return
119        self._enabled = False
120        self._client = None
121        logger.warning("CloudWatch logging disabled for %s: %s", self.stream_name, reason)
122
123    def emit_lifecycle(self, *, event: str, terminal_status: str | None = None) -> None:
124        payload = {
125            "event": f"stream_{event}",
126            "execution_id": self.execution_id,
127            "cache_key": self.cache_key,
128            "stream": self.stream_kind,
129        }
130        if terminal_status is not None:
131            payload["terminal_status"] = terminal_status
132        self.emit(json.dumps(payload, sort_keys=True))
133
134    def emit(self, message: str) -> None:
135        if not self._enabled or self._client is None:
136            return
137        messages = self._split_message(message)
138        with self._lock:
139            for chunk in messages:
140                event_bytes = self._event_bytes(chunk)
141                if self._pending_events and (
142                    len(self._pending_events) >= _CLOUDWATCH_MAX_BATCH_COUNT
143                    or self._pending_bytes + event_bytes > _CLOUDWATCH_MAX_BATCH_BYTES
144                ):
145                    self._flush_locked()
146                    if not self._enabled:
147                        return
148                event = {"timestamp": round(time.time() * 1000), "message": chunk}
149                self._pending_events.append(event)
150                self._pending_bytes += event_bytes
151
152    def close(self, *, terminal_status: str) -> None:
153        self.emit_lifecycle(event="end", terminal_status=terminal_status)
154        with self._lock:
155            self._flush_locked()
156
157
158def _drain_pipe(pipe: Any, *, local_path: Path, sink: _CloudWatchStream) -> None:
159    with local_path.open("w") as local_file:
160        for line in pipe:
161            local_file.write(line)
162            local_file.flush()
163            sink.emit(line)
164    pipe.close()
165
166
167def _parse_cmd_payload(
168    payload: dict[str, Any],
169) -> tuple[str, str, list[str], dict[str, str], dict[str, str]]:
170    allowed = {"version", "cache_key", "execution_id", "cmd", "remote", "env"}
171    unknown = sorted(set(payload) - allowed)
172    if unknown:
173        raise DmlRepoError(f"Supervisor payload has unknown fields: {', '.join(unknown)}")
174
175    version = payload.get("version")
176    if version != 0:
177        raise DmlRepoError("Supervisor payload version must be 0")
178
179    cache_key = payload.get("cache_key")
180    if not isinstance(cache_key, str) or not cache_key:
181        raise DmlRepoError("Supervisor payload cache_key must be a non-empty string")
182
183    execution_id = payload.get("execution_id")
184    if not isinstance(execution_id, str) or not execution_id:
185        raise DmlRepoError("Supervisor payload execution_id must be a non-empty string")
186
187    cmd = payload.get("cmd")
188    if not isinstance(cmd, list) or not cmd or not all(isinstance(x, str) and x for x in cmd):
189        raise DmlRepoError("Supervisor payload cmd must be a non-empty list[str]")
190
191    remote = payload.get("remote")
192    if not isinstance(remote, dict):
193        raise DmlRepoError("Supervisor payload remote must be a dict")
194    unknown_remote = sorted(set(remote) - {"root"})
195    if unknown_remote:
196        raise DmlRepoError(f"Supervisor payload remote has unknown fields: {', '.join(unknown_remote)}")
197    if not isinstance(remote.get("root"), str):
198        raise DmlRepoError("Supervisor payload remote requires string root")
199
200    env = payload.get("env") or {}
201    if not isinstance(env, dict) or not all(isinstance(k, str) and isinstance(v, str) for k, v in env.items()):
202        raise DmlRepoError("Supervisor payload env must be a dict[str,str]")
203
204    merged_env = os.environ.copy()
205    merged_env.update(env)
206    return cache_key, execution_id, cmd, merged_env, {"root": remote["root"]}
207
208
209def _validate_output(result: Any) -> dict[str, Any]:
210    if not isinstance(result, dict):
211        raise DmlRepoError(f"Supervisor result must be a dict; received: {result!r}")
212    status = result.get("status")
213    if status not in {"succeeded", "failed"}:
214        raise DmlRepoError(
215            "Supervisor result status must be one of succeeded|failed after worker exit; " f"received: {result!r}"
216        )
217    if status == "failed":
218        expected = {"status", "error"}
219        if set(result.keys()) != expected:
220            raise DmlRepoError(f"Supervisor failed result keys must be exactly: status, error; received: {result!r}")
221        error = result.get("error")
222        if error is None:
223            raise DmlRepoError(f"Supervisor result failed requires error; received: {result!r}")
224        return result
225
226    expected = {"status", "error", "dag_id"}
227    if set(result.keys()) != expected:
228        raise DmlRepoError(
229            f"Supervisor succeeded result keys must be exactly: status, error, dag_id; received: {result!r}"
230        )
231    error = result.get("error")
232    if error is not None:
233        raise DmlRepoError(f"Supervisor result succeeded requires error=None; received: {result!r}")
234    dag_id = result.get("dag_id")
235    if not isinstance(dag_id, str) or not re.fullmatch(r"[0-9a-f]{64}", dag_id):
236        raise DmlRepoError(f"Supervisor result succeeded requires real dag_id; received: {result!r}")
237    return result
238
239
240def run(payload: dict[str, Any]) -> dict[str, Any]:
241    """Launch a worker subprocess, wait for it to exit, and return the terminal result."""
242    cache_key, execution_id, cmd, env, remote = _parse_cmd_payload(payload)
243    workdir = tempfile.mkdtemp(prefix=f"dml-supervisor-{execution_id[:8]}-")
244    repo_dir = Path(workdir) / "repo"
245    repo_dir.mkdir(parents=True, exist_ok=True)
246    dml = Dml.init(str(repo_dir), remote_root=remote["root"], user="worker")
247    env = dict(env)
248    env["DML_PROJECT_HOME"] = str(repo_dir)
249    result_path = Path(workdir) / "result.json"
250    stdout_path = Path(workdir) / "stdout.log"
251    stderr_path = Path(workdir) / "stderr.log"
252    stdout_sink = _CloudWatchStream(cache_key=cache_key, execution_id=execution_id, stream_kind="stdout")
253    stderr_sink = _CloudWatchStream(cache_key=cache_key, execution_id=execution_id, stream_kind="stderr")
254    proc = subprocess.Popen(
255        cmd,
256        cwd=workdir,
257        env=env,
258        stdout=subprocess.PIPE,
259        stderr=subprocess.PIPE,
260        text=True,
261        encoding="utf-8",
262        errors="replace",
263        bufsize=1,
264        start_new_session=False,
265        close_fds=True,
266    )
267    assert proc.stdout is not None
268    assert proc.stderr is not None
269    stdout_thread = threading.Thread(
270        target=_drain_pipe,
271        kwargs={"pipe": proc.stdout, "local_path": stdout_path, "sink": stdout_sink},
272        name=f"dml-supervisor-{execution_id[:8]}-stdout",
273    )
274    stderr_thread = threading.Thread(
275        target=_drain_pipe,
276        kwargs={"pipe": proc.stderr, "local_path": stderr_path, "sink": stderr_sink},
277        name=f"dml-supervisor-{execution_id[:8]}-stderr",
278    )
279    stdout_thread.start()
280    stderr_thread.start()
281    execution = Ref(f"index:{execution_id}")
282    while proc.poll() is None:
283        if dml.runtime.read_execution_record(execution)["state"]["lifecycle"].startswith("cancel"):
284            # log to log streams that we're cancelling
285            stdout_sink.emit_lifecycle(event="cancel")
286            stderr_sink.emit_lifecycle(event="cancel")
287            proc.terminate()
288            break
289        time.sleep(0.1)
290    proc.wait()
291    stdout_thread.join()
292    stderr_thread.join()
293
294    result: dict[str, Any]
295    if result_path.exists():
296        try:
297            parsed = json.loads(result_path.read_text())
298            result = _validate_output(parsed)
299        except Exception as e:
300            result = {"status": "failed", "error": f"Supervisor could not read worker result: {e}"}
301    elif proc.returncode is not None and proc.returncode < 0:
302        sig = -proc.returncode
303        try:
304            sig_name = signal.Signals(sig).name
305        except ValueError:
306            sig_name = str(sig)
307        result = {"status": "failed", "error": f"Worker killed by signal {sig_name}"}
308    else:
309        code = proc.returncode if proc.returncode is not None else -1
310        result = {"status": "failed", "error": f"Worker exited without result (code={code})"}
311
312    terminal_status = str(result.get("status", "failed"))
313    stdout_sink.close(terminal_status=terminal_status)
314    stderr_sink.close(terminal_status=terminal_status)
315    return result
316
317
318def _read(path: str) -> str:
319    if path == "-":
320        return sys.stdin.read()
321    return Path(path).read_text()
322
323
324def _write(path: str, data: str) -> None:
325    if path == "-":
326        sys.stdout.write(data)
327        if not data.endswith("\n"):
328            sys.stdout.write("\n")
329        sys.stdout.flush()
330        return
331    Path(path).write_text(data)
332
333
334def main(argv: list[str] | None = None) -> int:
335    parser = argparse.ArgumentParser(description="daggerml contrib supervisor")
336    parser.add_argument("-i", "--input", default="-")
337    parser.add_argument("-o", "--output", default="-")
338    args = parser.parse_args(argv or sys.argv[1:])
339    payload = json.loads(_read(args.input))
340    result = run(payload)
341    _write(args.output, json.dumps(result, separators=(",", ":"), sort_keys=True))
342    return 0
343
344
345if __name__ == "__main__":
346    raise SystemExit(main())