diff --git a/v4/distributed/client.py b/v4/distributed/client.py index bf0adc9..ff63e4a 100644 --- a/v4/distributed/client.py +++ b/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")) diff --git a/v4/distributed/jobs.db b/v4/distributed/jobs.db index 6b3caa4..6eed7aa 100644 Binary files a/v4/distributed/jobs.db and b/v4/distributed/jobs.db differ diff --git a/v4/distributed/protocol.py b/v4/distributed/protocol.py index c69e775..ea7ca5c 100644 --- a/v4/distributed/protocol.py +++ b/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): diff --git a/v4/distributed/server.py b/v4/distributed/server.py index 7859f8b..955e316 100644 --- a/v4/distributed/server.py +++ b/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(