fix(jobs): detect a dead local server and record a terminal server_died error (BE-4751)

comfy_cli/command/job_watcher.py

@@ -40,6 +40,27 @@
 # An unknown status unchanged for this long → terminal error instead of
 # letting the watcher idle for the full 6h ceiling on a status we can't map.
 _UNKNOWN_STALL_S = 300.0
+# Consecutive *unreachable* liveness probes before we declare a local server
+# dead: ~3 polls x 2s ≈ 6s of a refused port, which tolerates a transient blip
+# / a quick server restart while still catching a real crash fast. Only a
+# refused connection counts (see _probe_local_server), so a server that is
+# merely slow can never accumulate this streak.
+_SERVER_DOWN_CONSECUTIVE_LIMIT = 3
+# How long to wait on a liveness probe. Bounds a TCP-reachable but wedged
+# server (stuck in a CUDA kernel) without hanging the poll cycle.
+_PROBE_TIMEOUT_S = 5.0
+# Liveness verdicts. Only _PROBE_UNREACHABLE — nothing listening on the port —
+# counts toward a death. A timeout or a non-200 means something *is* answering
+# and is merely busy (loading a large model into VRAM pins the server for tens
+# of seconds), which must never be mistaken for a crash.
+_PROBE_ALIVE = "alive"
+_PROBE_UNREACHABLE = "unreachable"
+_PROBE_UNRESPONSIVE = "unresponsive"
+# After a confirmed outage the server may come back as a *fresh* process whose
+# queue and history no longer hold this prompt — the job died with the old one.
+# A record that stays missing this long after an outage is a death too, rather
+# than a zombie poll until the 6h ceiling.
+_LOST_AFTER_RESTART_S = 60.0
 
 
 @app.command("_watch-job")
@@ -81,6 +102,15 @@ def watch_job(
     # currently watching, and when we first saw it.
     unknown_status: str | None = None
     unknown_since = 0.0
+    # Local-server death detection: how many consecutive probes found the port
+    # refused. Reset by any probe that reaches the server at all (see
+    # _SERVER_DOWN_CONSECUTIVE_LIMIT). ``saw_outage`` latches once we have seen
+    # a real outage, which is what makes a subsequently-missing record mean
+    # "the server restarted without this job" rather than "not queued yet";
+    # ``missing_since`` times that gap.
+    consecutive_down = 0
+    saw_outage = False
+    missing_since: float | None = None
     while True:
         if time.time() - start > _MAX_RUNTIME_S:
             prior_status = state.status
@@ -96,7 +126,76 @@ def watch_job(
         if where == "cloud":
             terminal = _poll_cloud_once(state, client=cloud_client)
         else:
-            terminal = _poll_local_once(state, host=host, port=port)
+            # A local server that is OOM-killed mid-job takes the queue and the
+            # history record with it, so `_snapshot` just reports "no record
+            # yet" forever and the watcher would poll a corpse until the 6h
+            # ceiling. Probe liveness explicitly so the death gets recorded.
+            h, p = _resolve_watch_target(state, host, port)
+            if _probe_local_server(h, p) == _PROBE_UNREACHABLE:
+                consecutive_down += 1
+                saw_outage = True
+                if consecutive_down < _SERVER_DOWN_CONSECUTIVE_LIMIT:
+                    # A blip or a restart, not a death yet. Skip the poll (it
+                    # can only fail against a port that just refused us) and
+                    # re-probe on the next cycle.
+                    terminal = False
+                else:
+                    # One last poll before the verdict: the probe hits
+                    # `/history`, the poll hits `/queue` + `/history/<id>`, so
+                    # a job that finished right as the probe started failing
+                    # can still be recovered if anything is still answering.
+                    terminal, _ = _poll_local_once(state, host=host, port=port)
+                    if not terminal:
+                        prior = state.status
+                        state = _finalize_server_died(
+                            state,
+                            message=(
+                                f"ComfyUI server {h}:{p} became unreachable while job "
+                                f"{state.prompt_id} was '{prior}'. The server likely crashed or was "
+                                "killed while executing it (e.g. an out-of-memory allocation)."
+                            ),
+                            details={
+                                "host": h,
+                                "port": p,
+                                "last_status": prior,
+                                "consecutive_failed_probes": consecutive_down,
+                            },
+                        )
+                    else:
+                        jobs_state.write(state)
+                    # Server confirmed gone — there is nothing left to watch.
+                    break
+            else:
+                # Something answered the port. Even a timeout or an error
+                # status means the server exists and is merely busy, so the
+                # death streak resets and we poll as normal.
+                consecutive_down = 0
+                terminal, record_seen = _poll_local_once(state, host=host, port=port)
+                if record_seen or not saw_outage:
+                    missing_since = None
+                elif missing_since is None:
+                    # Back up after an outage, but this prompt is gone from the
+                    # server: either it is still settling or the process that
+                    # held the job was replaced. Give it a grace window.
+                    missing_since = time.time()
+                elif (missing_for := time.time() - missing_since) >= _LOST_AFTER_RESTART_S:
+                    prior = state.status
+                    state = _finalize_server_died(
+                        state,
+                        message=(
+                            f"ComfyUI server {h}:{p} restarted while job {state.prompt_id} was "
+                            f"'{prior}' and the new process has no record of it. The job died with "
+                            "the server (it likely crashed or was killed, e.g. by the OOM killer)."
+                        ),
+                        details={
+                            "host": h,
+                            "port": p,
+                            "last_status": prior,
+                            "restarted": True,
+                            "missing_for_s": round(missing_for, 1),
+                        },
+                    )
+                    break
 
         jobs_state.write(state)
         if terminal:
@@ -137,9 +236,12 @@ def watch_job(
 # ---------------------------------------------------------------------------
 
 
-def _poll_local_once(state: jobs_state.JobState, *, host: str | None, port: int | None) -> bool:
-    """Update ``state`` in-place from a local ComfyUI server. Return True if terminal."""
-    from comfy_cli.command import jobs as jobs_module
+def _resolve_watch_target(state: jobs_state.JobState, host: str | None, port: int | None) -> tuple[str, int]:
+    """The local ``(host, port)`` this watcher is bound to.
+
+    Shared by the liveness probe and the poll so the two can never disagree
+    about which server they are talking about.
+    """
     from comfy_cli.local_address import resolve_local_host_port
 
     # Per-job recorded state (state.host/port, captured when the job was
@@ -150,38 +252,114 @@ def _poll_local_once(state: jobs_state.JobState, *, host: str | None, port: int
     # an already-bracketed host, like the `jobs` resolver produces).
     if ":" in h and not h.startswith("["):
         h = f"[{h}]"
+    return h, p
+
+
+def _probe_local_server(host: str, port: int) -> str:
+    """Classify the local ComfyUI server as alive / unreachable / unresponsive.
+
+    ``comfy jobs`` asks ``check_comfy_server_running`` a yes/no question, which
+    collapses "the port is refused" and "the server took too long to answer"
+    into the same ``False``. The watcher cannot: only the first is a crash, and
+    treating a slow server as dead would file a terminal ``server_died`` error
+    against a job that is still running. So it runs its own probe and keeps the
+    distinction, with the same liveness definition (HTTP 200 → alive).
+
+    ``max_items=1`` keeps the response small — the bare ``/history`` body grows
+    without bound on a long-lived server, and a multi-megabyte read every
+    ``_POLL_INTERVAL_S`` is both wasteful and a way to time the probe out.
+    Servers too old to know the parameter ignore it and still answer 200.
+
+    Never raises: an unexpected failure is reported as *unresponsive* (the
+    non-committal verdict), because a probe must not crash the watcher and must
+    not be able to invent a death on its own.
+    """
+    import requests
+
+    try:
+        # ``host`` arrives already bracketed from _resolve_watch_target.
+        resp = requests.get(f"http://{host}:{port}/history?max_items=1", timeout=_PROBE_TIMEOUT_S)
+    except requests.exceptions.Timeout:
+        # Listed before ConnectionError: ConnectTimeout subclasses both, and a
+        # timeout is never evidence that nothing is listening.
+        return _PROBE_UNRESPONSIVE
+    except requests.exceptions.ConnectionError:
+        return _PROBE_UNREACHABLE
+    except Exception:  # noqa: BLE001
+        return _PROBE_UNRESPONSIVE
+    return _PROBE_ALIVE if resp.status_code == 200 else _PROBE_UNRESPONSIVE
+
+
+def _finalize_server_died(state: jobs_state.JobState, *, message: str, details: dict[str, Any]) -> jobs_state.JobState:
+    """Persist the terminal ``server_died`` verdict; return the state to notify on.
+
+    Re-reads the file rather than trusting the watcher's in-memory copy, and
+    does the read and the write as one locked transaction: a concurrent
+    ``comfy jobs cancel`` may have recorded a verdict of its own while we were
+    failing probes, and a dead server must never invalidate a verdict the job
+    already reached. ``jobs_state.write`` takes the same per-file lock, which
+    is reentrant within a thread, so the nested acquire is a no-op.
+
+    Any terminal state — on disk or in memory — wins and is returned unwritten.
+    """
+    from comfy_cli import locking
+
+    lock_path = jobs_state.state_path(state.prompt_id).with_suffix(".lock")
+    with locking.file_lock(lock_path):
+        on_disk = jobs_state.read(state.prompt_id)
+        if on_disk is not None and on_disk.is_terminal:
+            return on_disk
+        if state.is_terminal:
+            return state
+        state.status = "error"
+        state.error = {"code": "server_died", "message": message, "details": details}
+        jobs_state.write(state)
+    return state
+
+
+def _poll_local_once(state: jobs_state.JobState, *, host: str | None, port: int | None) -> tuple[bool, bool]:
+    """Update ``state`` in-place from a local ComfyUI server.
+
+    Returns ``(terminal, record_seen)``. ``record_seen`` distinguishes "the
+    server has no record of this prompt" from "the record says it is still
+    running", which the restart guard in ``watch_job`` needs and a bare
+    terminal flag cannot express.
+    """
+    from comfy_cli.command import jobs as jobs_module
+
+    h, p = _resolve_watch_target(state, host, port)
     try:
         snap = jobs_module._snapshot(h, p, state.prompt_id)
     except Exception as e:  # noqa: BLE001 — never crash the watcher on transient errors
         state.error = {"code": "watcher_poll_error", "message": str(e), "details": {}}
-        return False
+        return False, False
 
     if snap is None:
         # No record yet — keep polling.
-        return False
+        return False, False
 
     # Clear any transient poll error from a previous cycle.
     state.error = None
     snap_status = str(snap.get("status") or "queued")
     state.status = snap_status
     if snap_status == "completed":
         state.outputs = list(snap.get("outputs") or [])
-        return True
+        return True, True
     if snap_status == "error":
         state.error = {
             "code": "execution_error",
             "message": "ComfyUI reported an execution error.",
             "details": snap,
         }
-        return True
+        return True, True
     if snap_status == "cancelled":
         state.error = {
             "code": "cancelled",
             "message": "Job was interrupted/cancelled.",
             "details": {},
         }
-        return True
-    return False
+        return True, True
+    return False, True
 
 
 _CLOUD_STATUS_MAP = {

comfy_cli/error_codes.py

@@ -334,6 +334,12 @@ class ErrorCode:
         "Background watcher encountered a transient error polling the server.",
         "transient — the job is likely still running; re-run `comfy jobs watch <id>`",
     ),
+    ErrorCode(
+        "server_died",
+        "The local ComfyUI server became unreachable (or restarted without the job) while it "
+        "was in flight — the server likely crashed or was killed (e.g. an out-of-memory allocation).",
+        "check the ComfyUI server log, then `comfy launch` and re-submit the workflow",
+    ),
     ErrorCode(
         "unknown_status_stall",
         "Cloud reported a status the CLI does not recognize and it did not change within the stall window.",