Implement GPU precheck and client fail streak management

v4/distributed/client.py

@@ -101,6 +101,32 @@ class _Server:
 # GPU info
 # ──────────────────────────────────────────────────────────────
 
+def _gpu_precheck() -> tuple[bool, str]:
+    """Quick allocation + op test on every visible CUDA/HIP device.
+
+    Catches the case where torch.cuda.is_available() returns True but the
+    device is busy (e.g. another process holding it — Ollama, a forgotten
+    notebook, an OS-level driver issue). Returns (ok, message); when ok is
+    False, the caller should defer polling rather than accept a job that will
+    crash on the first .to(device) call.
+    """
+    try:
+        import torch
+    except Exception as e:
+        return False, f"torch import failed: {e}"
+    if not torch.cuda.is_available():
+        return True, "no-cuda (cpu-only client)"
+    try:
+        for i in range(torch.cuda.device_count()):
+            x = torch.zeros(1024, device=f"cuda:{i}")
+            _ = (x + 1).sum().item()  # forces actual kernel launch
+            del x
+        torch.cuda.empty_cache()
+        return True, "ok"
+    except Exception as e:
+        return False, f"{type(e).__name__}: {e}"
+
+
 def _gpu_info() -> str:
     # NVIDIA
     try:
@@ -405,8 +431,22 @@ def main():
         ) else 1)
 
     print(f"[client] polling every {args.poll_interval}s...", flush=True)
+    last_precheck_msg = ""
     while True:
         try:
+            ok, msg = _gpu_precheck()
+            if not ok:
+                if msg != last_precheck_msg:
+                    print(f"[client] gpu precheck failed ({msg}) — deferring polls",
+                          flush=True)
+                    last_precheck_msg = msg
+                server.push_status(StatusPush(state="gpu_busy"))
+                time.sleep(args.poll_interval)
+                continue
+            if last_precheck_msg:
+                print(f"[client] gpu precheck recovered — resuming polls", flush=True)
+                last_precheck_msg = ""
+
             job = server.poll()
             if job is None:
                 server.push_status(StatusPush(state="idle"))

v4/distributed/protocol.py

@@ -32,6 +32,7 @@ class ClientInfo(BaseModel):
     gpu_info: str
     status: StatusPush
     last_seen: str
+    fail_streak: int = 0  # consecutive job failures; resets on success or register
 
 
 class JobSpec(BaseModel):

v4/distributed/server.py

@@ -47,7 +47,11 @@ _TOKEN: str = ""
 _DB_PATH: Path = Path("v4/distributed/jobs.db")
 _REPO_ROOT: Path = Path.cwd()
 _CLIENT_TTL: int = 120  # seconds before a client is considered gone
-_MAX_ATTEMPTS: int = 3  # max times a job is retried before being left as failed
+_MAX_ATTEMPTS: int = 4  # max times a job is retried (cumulatively, across clients)
+                        # before being marked permanently failed
+_CLIENT_FAIL_LIMIT: int = 3  # consecutive job failures from a single client before
+                             # that client is quarantined (no more poll dispatches
+                             # until it re-registers). A successful job resets it.
 _SERVER_START_TS: float = 0.0  # set in main(); used as a reaper grace window
 
 _clients: dict[str, ClientInfo] = {}
@@ -222,6 +226,19 @@ def poll(client_id: str):
     with _clients_lock:
         needs_reregister = _ensure_client(client_id)
         _clients[client_id].last_seen = _now()
+        # Quarantine: a client that has failed _CLIENT_FAIL_LIMIT jobs in a row
+        # is cut off from new dispatches until it re-registers. The client's
+        # poll() helper auto-calls _reregister() when please_reregister=True
+        # arrives, which clears the streak. Successful completions also reset
+        # the streak, so a healthy client never trips this check.
+        if _clients[client_id].fail_streak >= _CLIENT_FAIL_LIMIT:
+            print(
+                f"[server] {client_id} quarantined "
+                f"(fail_streak={_clients[client_id].fail_streak}); "
+                f"requesting re-register before next dispatch",
+                flush=True,
+            )
+            return PollResponse(job=None, please_reregister=True)
 
     with _db() as conn:
         row = conn.execute(
@@ -296,6 +313,28 @@ def get_client(client_id: str):
 
 @app.post("/complete", dependencies=[Depends(_check_token)])
 def complete(result: JobResult):
+    # Pull the worker that ran this job before touching the row, so we can
+    # update its streak regardless of which branch we take below.
+    with _db() as conn:
+        assigned_row = conn.execute(
+            "SELECT assigned_to FROM jobs WHERE job_id=?", (result.job_id,)
+        ).fetchone()
+    assigned_to = assigned_row["assigned_to"] if assigned_row else None
+
+    with _clients_lock:
+        if assigned_to and assigned_to in _clients:
+            if result.success:
+                _clients[assigned_to].fail_streak = 0
+            else:
+                _clients[assigned_to].fail_streak += 1
+                if _clients[assigned_to].fail_streak >= _CLIENT_FAIL_LIMIT:
+                    print(
+                        f"[server] {assigned_to} hit fail_streak="
+                        f"{_clients[assigned_to].fail_streak}; will quarantine "
+                        f"on next poll",
+                        flush=True,
+                    )
+
     with _db() as conn:
         if result.success:
             conn.execute(