fix: keep agent alive across worker crashes

main.py

@@ -6,6 +6,7 @@ import signal
 import subprocess
 import sys
 import time
+from datetime import datetime, timezone
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from pathlib import Path
 
@@ -18,10 +19,18 @@ PORT = int(os.getenv("PORT", "8080"))
 ROOT = Path(__file__).resolve().parent
 WORKER_SCRIPT = ROOT / "modelhub_submmit_api" / "poll_runner.py"
 READINESS_PATH = ROOT / ".modelhub_state" / "readiness.json"
+WORKER_CRASH_PATH = ROOT / ".modelhub_state" / "worker_crashes.jsonl"
+WORKER_STABLE_SECONDS = max(60, int(os.getenv("MODELHUB_AGENT_WORKER_STABLE_SECONDS", "600")))
+WORKER_RESTART_MAX_SECONDS = max(30, int(os.getenv("MODELHUB_AGENT_WORKER_RESTART_MAX_SECONDS", "300")))
 
 shutdown_requested = False
 worker: subprocess.Popen | None = None
 config_error: str | None = None
+worker_start_enabled = False
+worker_started_at: float | None = None
+worker_next_restart_at = 0.0
+worker_restart_count = 0
+worker_last_exit_code: int | None = None
 
 
 def _has_modelhub_auth() -> bool:
@@ -84,7 +93,14 @@ def _config() -> dict[str, object]:
         "modelhub_auth_present": _has_modelhub_auth(),
         "config_error": config_error,
         "worker_running": worker is not None and worker.poll() is None,
-        "worker_return_code": None if worker is None else worker.poll(),
+        "worker_return_code": worker_last_exit_code if worker is None else worker.poll(),
+        "worker_restart_count": worker_restart_count,
+        "worker_restart_pending": bool(
+            worker_start_enabled
+            and worker is None
+            and not shutdown_requested
+            and config_error is None
+        ),
     }
 
 
@@ -99,13 +115,12 @@ def _readiness() -> dict[str, object]:
 class Handler(BaseHTTPRequestHandler):
     def do_GET(self) -> None:
         if self.path == "/health":
+            # This endpoint is a liveness probe for the supervisor. Worker
+            # readiness is reported separately by /ready.
+            status = "ok" if worker is not None and worker.poll() is None else "degraded"
             if config_error:
-                self._send_json({"status": "config_error", "config": _config()}, status=500)
-                return
-            if worker is not None and worker.poll() is not None:
-                self._send_json({"status": "worker_exited", "config": _config()}, status=500)
-                return
-            self._send_json({"status": "ok", "config": _config()})
+                status = "config_error"
+            self._send_json({"status": status, "config": _config()})
             return
 
         if self.path == "/ready":
@@ -152,6 +167,96 @@ def _stop_worker() -> None:
         worker.wait(timeout=5)
 
 
+def _restart_delay(restart_count: int) -> int:
+    return min(WORKER_RESTART_MAX_SECONDS, 5 * (2 ** min(max(0, restart_count - 1), 6)))
+
+
+def _write_supervisor_readiness(reason: str, *, exit_code: int | None = None) -> None:
+    READINESS_PATH.parent.mkdir(parents=True, exist_ok=True)
+    payload = {
+        "ready": False,
+        "reason": reason,
+        "updatedAt": datetime.now(timezone.utc).isoformat(),
+        "supervisor": {
+            "workerRestartCount": worker_restart_count,
+            "lastExitCode": exit_code,
+        },
+    }
+    temporary = READINESS_PATH.with_name(f".{READINESS_PATH.name}.supervisor-{os.getpid()}")
+    temporary.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8")
+    os.replace(temporary, READINESS_PATH)
+
+
+def _record_worker_crash(*, exit_code: int | None, uptime_seconds: float, reason: str) -> None:
+    WORKER_CRASH_PATH.parent.mkdir(parents=True, exist_ok=True)
+    event = {
+        "at": datetime.now(timezone.utc).isoformat(),
+        "exitCode": exit_code,
+        "uptimeSeconds": round(max(0.0, uptime_seconds), 3),
+        "restartCount": worker_restart_count,
+        "reason": reason[:500],
+    }
+    with WORKER_CRASH_PATH.open("a", encoding="utf-8") as handle:
+        handle.write(json.dumps(event, ensure_ascii=False, separators=(",", ":")) + "\n")
+
+
+def _launch_worker(*, now: float | None = None) -> bool:
+    global worker, worker_started_at, worker_last_exit_code
+    now = time.monotonic() if now is None else now
+    cmd = _worker_command()
+    print("starting submission worker: " + " ".join(cmd), flush=True)
+    try:
+        worker = subprocess.Popen(cmd, cwd=str(ROOT))
+    except OSError as exc:
+        worker = None
+        worker_started_at = None
+        worker_last_exit_code = None
+        print(f"submission worker spawn failed: {type(exc).__name__}: {exc}", flush=True)
+        return False
+    worker_started_at = now
+    worker_last_exit_code = None
+    return True
+
+
+def _supervise_worker(*, now: float | None = None) -> None:
+    global worker, worker_started_at, worker_next_restart_at
+    global worker_restart_count, worker_last_exit_code
+    if not worker_start_enabled or shutdown_requested or config_error:
+        return
+    now = time.monotonic() if now is None else now
+    if worker is not None:
+        return_code = worker.poll()
+        if return_code is None:
+            if worker_started_at is not None and now - worker_started_at >= WORKER_STABLE_SECONDS:
+                worker_restart_count = 0
+            return
+        uptime = max(0.0, now - worker_started_at) if worker_started_at is not None else 0.0
+        if uptime >= WORKER_STABLE_SECONDS:
+            worker_restart_count = 0
+        worker_restart_count += 1
+        worker_last_exit_code = int(return_code)
+        delay = _restart_delay(worker_restart_count)
+        worker_next_restart_at = now + delay
+        print(
+            f"submission worker exited code={return_code} uptime={uptime:.1f}s "
+            f"restart_in={delay}s restart_count={worker_restart_count}",
+            flush=True,
+        )
+        _record_worker_crash(exit_code=return_code, uptime_seconds=uptime, reason="worker_exited")
+        _write_supervisor_readiness("worker_restarting", exit_code=return_code)
+        worker = None
+        worker_started_at = None
+        return
+    if now < worker_next_restart_at:
+        return
+    if not _launch_worker(now=now):
+        worker_restart_count += 1
+        delay = _restart_delay(worker_restart_count)
+        worker_next_restart_at = now + delay
+        _record_worker_crash(exit_code=None, uptime_seconds=0.0, reason="worker_spawn_failed")
+        _write_supervisor_readiness("worker_spawn_failed")
+
+
 def _handle_signal(signum: int, _frame: object) -> None:
     global shutdown_requested
     shutdown_requested = True
@@ -160,19 +265,17 @@ def _handle_signal(signum: int, _frame: object) -> None:
 
 
 def main() -> int:
-    global config_error, worker
+    global config_error, worker_start_enabled
     print(f"modelhub-submmit-agent version={AGENT_VERSION}", flush=True)
     signal.signal(signal.SIGTERM, _handle_signal)
     signal.signal(signal.SIGINT, _handle_signal)
 
-    start_worker = os.getenv("MODELHUB_AGENT_START_WORKER", "1").strip().lower() not in {"0", "false", "no"}
-    if start_worker and not _has_modelhub_auth():
+    worker_start_enabled = os.getenv("MODELHUB_AGENT_START_WORKER", "1").strip().lower() not in {"0", "false", "no"}
+    if worker_start_enabled and not _has_modelhub_auth():
         config_error = "missing ModelHub auth: set MODELHUB_XC_TOKEN/XC_TOKEN or MODELHUB_JWT_TOKEN/JWT_TOKEN"
         print(config_error, flush=True)
-    elif start_worker:
-        cmd = _worker_command()
-        print("starting submission worker: " + " ".join(cmd), flush=True)
-        worker = subprocess.Popen(cmd, cwd=str(ROOT))
+    elif worker_start_enabled:
+        _launch_worker()
     else:
         print("submission worker disabled by MODELHUB_AGENT_START_WORKER", flush=True)
 
@@ -183,9 +286,7 @@ def main() -> int:
     try:
         while not shutdown_requested:
             server.handle_request()
-            if worker is not None and worker.poll() is not None:
-                print(f"submission worker exited with code {worker.returncode}", flush=True)
-                return worker.returncode or 1
+            _supervise_worker()
             time.sleep(0.1)
     finally:
         server.server_close()