Supervise checkpointed Tower training

scripts/supervise_container_yolo_training.py

@@ -0,0 +1,94 @@
+#!/usr/bin/env python3
+"""Resume one checkpointed YOLO run after container recreation, fail closed."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import subprocess
+import time
+from datetime import UTC, datetime
+from pathlib import Path
+
+
+def container_running(container: str) -> bool:
+    result = subprocess.run(
+        ["docker", "inspect", "--format", "{{.State.Running}}", container],
+        capture_output=True, text=True, check=False,
+    )
+    return result.returncode == 0 and result.stdout.strip() == "true"
+
+
+def training_active(container: str, run_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(
+        "train" in line and run_marker in line
+        for line in result.stdout.splitlines()
+    )
+
+
+def write_state(path: Path, payload: dict) -> None:
+    payload["updated_at"] = datetime.now(UTC).isoformat()
+    temporary = path.with_suffix(path.suffix + ".tmp")
+    temporary.write_text(json.dumps(payload, indent=2), encoding="utf-8")
+    temporary.replace(path)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--container", required=True)
+    parser.add_argument("--host-run-dir", type=Path, required=True)
+    parser.add_argument("--container-checkpoint", required=True)
+    parser.add_argument("--run-marker", required=True)
+    parser.add_argument("--yolo", default="/opt/geointel/venv/bin/yolo")
+    parser.add_argument("--poll-seconds", type=int, default=30)
+    parser.add_argument("--max-resumes", type=int, default=20)
+    parser.add_argument("--once", action="store_true")
+    args = parser.parse_args()
+    if args.poll_seconds < 1 or args.max_resumes < 1:
+        raise SystemExit("poll-seconds and max-resumes must be positive")
+
+    state_path = args.host_run_dir / "supervisor-state.json"
+    state = {"schema_version": 1, "status": "monitoring", "resume_count": 0}
+    if state_path.is_file():
+        state.update(json.loads(state_path.read_text(encoding="utf-8")))
+
+    while True:
+        if (args.host_run_dir / "results.png").is_file():
+            state["status"] = "training_finished"
+            write_state(state_path, state)
+            return 0
+        checkpoint = args.host_run_dir / "weights" / "last.pt"
+        if not checkpoint.is_file() or checkpoint.stat().st_size < 1024 * 1024:
+            state["status"] = "checkpoint_missing_or_incomplete"
+            write_state(state_path, state)
+            return 2
+        if container_running(args.container) and not training_active(args.container, args.run_marker):
+            if int(state["resume_count"]) >= args.max_resumes:
+                state["status"] = "resume_budget_exhausted"
+                write_state(state_path, state)
+                return 3
+            result = subprocess.run(
+                ["docker", "exec", "-d", args.container, args.yolo, "train",
+                 f"resume={args.container_checkpoint}", "device=0"],
+                check=False,
+            )
+            if result.returncode == 0:
+                state["resume_count"] = int(state["resume_count"]) + 1
+                state["status"] = "resumed"
+            else:
+                state["status"] = "resume_command_failed"
+            write_state(state_path, state)
+        else:
+            state["status"] = "monitoring"
+            write_state(state_path, state)
+        if args.once:
+            return 0
+        time.sleep(args.poll_seconds)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())