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