Recover active training loop checkpoints

scripts/supervise_container_training_loop.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+"""Keep a checkpoint-aware training-loop parent alive across container recreation."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import subprocess
+import time
+from datetime import UTC, datetime
+from pathlib import Path
+
+TERMINAL_STATUSES = {"training_complete"}
+
+
+def read_loop_status(path: Path) -> str | None:
+    if not path.is_file():
+        return None
+    try:
+        value = json.loads(path.read_text(encoding="utf-8"))
+    except (OSError, json.JSONDecodeError):
+        return "invalid"
+    return value.get("status") if isinstance(value, dict) else "invalid"
+
+
+def process_active(container: str, marker: str) -> bool:
+    result = subprocess.run(
+        ["docker", "top", container, "-eo", "pid,args"],
+        capture_output=True, text=True, check=False,
+    )
+    return result.returncode == 0 and any(marker in line for line in result.stdout.splitlines())
+
+
+def write_state(path: Path, state: dict) -> None:
+    state["updated_at"] = datetime.now(UTC).isoformat()
+    temporary = path.with_suffix(path.suffix + ".tmp")
+    temporary.write_text(json.dumps(state, indent=2), encoding="utf-8")
+    temporary.replace(path)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--container", required=True)
+    parser.add_argument("--process-marker", required=True)
+    parser.add_argument("--loop-state", type=Path, required=True)
+    parser.add_argument("--command-json", type=Path, required=True)
+    parser.add_argument("--host-script", type=Path, required=True)
+    parser.add_argument("--container-script", required=True)
+    parser.add_argument("--state", type=Path, required=True)
+    parser.add_argument("--poll-seconds", type=int, default=30)
+    parser.add_argument("--max-launches", type=int, default=20)
+    parser.add_argument("--once", action="store_true")
+    args = parser.parse_args()
+    state = {"schema_version": 1, "status": "monitoring", "launch_count": 0}
+    if args.state.is_file():
+        state.update(json.loads(args.state.read_text(encoding="utf-8")))
+    while True:
+        loop_status = read_loop_status(args.loop_state)
+        state["loop_status"] = loop_status
+        if loop_status in TERMINAL_STATUSES:
+            state["status"] = "loop_finished"
+            write_state(args.state, state)
+            return 0
+        if loop_status == "invalid":
+            state["status"] = "invalid_loop_state"
+            write_state(args.state, state)
+            return 2
+        if not process_active(args.container, args.process_marker):
+            if int(state["launch_count"]) >= args.max_launches:
+                state["status"] = "launch_budget_exhausted"
+                write_state(args.state, state)
+                return 3
+            command = json.loads(args.command_json.read_text(encoding="utf-8"))
+            if not isinstance(command, list) or not command or not all(isinstance(x, str) and x for x in command):
+                state["status"] = "invalid_command"
+                write_state(args.state, state)
+                return 4
+            copied = subprocess.run(
+                ["docker", "cp", str(args.host_script), f"{args.container}:{args.container_script}"],
+                check=False,
+            )
+            launched = subprocess.run(command, check=False) if copied.returncode == 0 else copied
+            if launched.returncode != 0:
+                state["status"] = "launch_failed"
+                write_state(args.state, state)
+                return 5
+            state["launch_count"] = int(state["launch_count"]) + 1
+            state["status"] = "loop_relaunched"
+        else:
+            state["status"] = "monitoring"
+        write_state(args.state, state)
+        if args.once:
+            return 0
+        time.sleep(args.poll_seconds)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())