From 93ebf99988acefcd8b54f80da3f320108c390e04 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Mon, 27 Jul 2026 14:14:17 -0700 Subject: [PATCH 1/2] fix(jobs): detect a dead local server and record a terminal server_died error (BE-4751) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An async-submitted local job (`comfy run` without `--wait`) whose ComfyUI server is OOM-killed left a forever-`queued` ghost: `_snapshot` swallows the /queue connection failure and returns None on the /history one, and the watcher reads `None` as "no record yet — keep polling", so it polled a corpse every 2s for the full 6h ceiling and never recorded the death. The watcher now probes local-server liveness explicitly each cycle with the same `check_comfy_server_running` helper `comfy jobs` uses. Three consecutive failed probes (~6s, so a blip or a quick restart is tolerated) on a non-terminal job writes a terminal `server_died` error naming the host, port, last status and probe count, then breaks so the existing notification fires. Any successful probe resets the streak. Probe and poll share a new `_resolve_watch_target` helper so they can never disagree on the target. `server_died` is added to the error-code registry, which the repo's own contract test requires. The cloud branch is untouched. --- comfy_cli/command/job_watcher.py | 77 +++++++++++++- comfy_cli/error_codes.py | 6 ++ tests/comfy_cli/jobs/test_jobs.py | 164 ++++++++++++++++++++++++++++++ 3 files changed, 243 insertions(+), 4 deletions(-) diff --git a/comfy_cli/command/job_watcher.py b/comfy_cli/command/job_watcher.py index 10bdf9f9..412b952c 100644 --- a/comfy_cli/command/job_watcher.py +++ b/comfy_cli/command/job_watcher.py @@ -40,6 +40,10 @@ # 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 failed liveness probes before we declare a local server dead: +# ~3 polls x 2s ≈ 6s of confirmed unreachability, which tolerates a transient +# blip / a quick server restart while still catching a real crash fast. +_SERVER_DOWN_CONSECUTIVE_LIMIT = 3 @app.command("_watch-job") @@ -81,6 +85,9 @@ 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 liveness probes have + # failed. Reset by any successful probe (see _SERVER_DOWN_CONSECUTIVE_LIMIT). + consecutive_down = 0 while True: if time.time() - start > _MAX_RUNTIME_S: prior_status = state.status @@ -96,7 +103,43 @@ 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 _server_alive(h, p): + consecutive_down = 0 + terminal = _poll_local_once(state, host=host, port=port) + else: + consecutive_down += 1 + if consecutive_down >= _SERVER_DOWN_CONSECUTIVE_LIMIT: + # A state that is already terminal keeps its own verdict — + # a dead server can't invalidate a completed/failed job. + if not state.is_terminal: + prior = state.status + state.status = "error" + state.error = { + "code": "server_died", + "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, + }, + } + jobs_state.write(state) + # Server confirmed gone — there is nothing left to watch. + break + # Unreachable but still under the limit — a blip or a restart. + # Skip the poll (it can only fail against a server we just + # failed to reach) and re-probe on the next cycle. + terminal = False jobs_state.write(state) if terminal: @@ -137,9 +180,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,6 +196,29 @@ 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 _server_alive(host: str, port: int) -> bool: + """Is the local ComfyUI server reachable? + + The same helper ``comfy jobs`` uses for its own up/down check, so the + watcher and the CLI never disagree. An unexpected error counts as *down* + rather than propagating — a probe must never crash the watcher. + """ + from comfy_cli.env_checker import check_comfy_server_running + + try: + return bool(check_comfy_server_running(port, host)) + except Exception: # noqa: BLE001 + return False + + +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 + + 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 diff --git a/comfy_cli/error_codes.py b/comfy_cli/error_codes.py index f0729744..5a153393 100644 --- a/comfy_cli/error_codes.py +++ b/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 `", ), + ErrorCode( + "server_died", + "The local ComfyUI server became unreachable while the job was in flight — " + "it 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.", diff --git a/tests/comfy_cli/jobs/test_jobs.py b/tests/comfy_cli/jobs/test_jobs.py index 5ed15898..ef8ba04e 100644 --- a/tests/comfy_cli/jobs/test_jobs.py +++ b/tests/comfy_cli/jobs/test_jobs.py @@ -1348,3 +1348,167 @@ def test_watch_executing_escapes_server_controlled_node_markup(): assert r"\[red]evil\[/red]" in r.printed[0] # The event stream still carries the raw node id. assert ("executing", {"node": "[red]evil[/red]", "prompt_id": "pid"}) in r.events + + +# --------------------------------------------------------------------------- +# watcher: local server death detection (server_died) +# --------------------------------------------------------------------------- + + +def _watched_local_job(monkeypatch, *, status="queued"): + """A real on-disk state file for a local job, plus a no-op sleep. + + The autouse ``_isolate_jobs_state_dir`` fixture pins the state dir to a tmp + path, so ``watch_job`` reads/writes real files and the assertions can be + made against the file rather than an in-memory object. + """ + from comfy_cli import jobs_state + from comfy_cli.command import job_watcher + + monkeypatch.delenv("COMFY_LOCAL_URL", raising=False) + monkeypatch.setattr(job_watcher.time, "sleep", lambda s: None) + state = jobs_state.new( + prompt_id="pid-dead", + client_id="c", + workflow="/w.json", + where="local", + host="127.0.0.1", + port=8188, + ) + state.status = status + jobs_state.write(state) + return state + + +def _fake_probe(monkeypatch, results): + """Point the watcher's liveness probe at a scripted sequence of verdicts. + + The last verdict repeats forever so a test can't hang on a short script. + """ + seq = list(results) + calls = [] + + def probe(port, host, *a, **kw): + calls.append((host, port)) + return seq[min(len(calls) - 1, len(seq) - 1)] + + monkeypatch.setattr("comfy_cli.env_checker.check_comfy_server_running", probe) + return calls + + +def test_watcher_records_server_died_after_consecutive_failed_probes(monkeypatch): + """A local server that dies mid-job (OOM kill) must be recorded as a + terminal ``server_died`` error instead of leaving a forever-'queued' ghost: + ``_snapshot`` swallows the connection failure, so only an explicit liveness + probe can see the death.""" + from comfy_cli import jobs_state + from comfy_cli.command import job_watcher + + _watched_local_job(monkeypatch) + calls = _fake_probe(monkeypatch, [False]) + # The poll must never run against a server we just failed to probe. + monkeypatch.setattr( + "comfy_cli.command.jobs._snapshot", + lambda h, p, pid: pytest.fail("polled a server that failed its liveness probe"), + ) + notified = [] + monkeypatch.setattr(job_watcher, "_notify", notified.append) + + job_watcher.watch_job("pid-dead", where="local") + + assert len(calls) == job_watcher._SERVER_DOWN_CONSECUTIVE_LIMIT + on_disk = jobs_state.read("pid-dead") + assert on_disk is not None + assert on_disk.status == "error" + assert on_disk.error["code"] == "server_died" + assert on_disk.error["details"]["last_status"] == "queued" + assert on_disk.error["details"]["host"] == "127.0.0.1" + assert on_disk.error["details"]["port"] == 8188 + assert on_disk.error["details"]["consecutive_failed_probes"] == job_watcher._SERVER_DOWN_CONSECUTIVE_LIMIT + assert "pid-dead" in on_disk.error["message"] + # ...and the user is told exactly once, with the final state. + assert len(notified) == 1 + assert notified[0].status == "error" + assert notified[0].error["code"] == "server_died" + + +def test_watcher_down_probe_streak_resets_on_recovery(monkeypatch): + """A transient blip (or a quick restart) must not be mistaken for death: + any successful probe resets the streak and the watcher keeps going.""" + from comfy_cli import jobs_state + from comfy_cli.command import job_watcher + + _watched_local_job(monkeypatch) + _fake_probe(monkeypatch, [False, False, True]) + monkeypatch.setattr( + "comfy_cli.command.jobs._snapshot", + lambda h, p, pid: {"prompt_id": pid, "status": "completed", "outputs": ["a.png"]}, + ) + monkeypatch.setattr(job_watcher, "_notify", lambda s: None) + + job_watcher.watch_job("pid-dead", where="local") + + on_disk = jobs_state.read("pid-dead") + assert on_disk is not None + assert on_disk.status == "completed" + assert on_disk.error is None + assert on_disk.outputs == ["a.png"] + + +def test_watcher_does_not_overwrite_a_terminal_state_when_server_dies(monkeypatch): + """A dead server can't invalidate a verdict the job already reached — the + watcher stops without rewriting a terminal state.""" + from comfy_cli import jobs_state + from comfy_cli.command import job_watcher + + state = _watched_local_job(monkeypatch, status="completed") + state.outputs = ["done.png"] + jobs_state.write(state) + _fake_probe(monkeypatch, [False]) + monkeypatch.setattr(job_watcher, "_notify", lambda s: None) + + job_watcher.watch_job("pid-dead", where="local") + + on_disk = jobs_state.read("pid-dead") + assert on_disk is not None + assert on_disk.status == "completed" + assert on_disk.error is None + assert on_disk.outputs == ["done.png"] + + +def test_watcher_probe_targets_the_same_address_the_poll_uses(monkeypatch): + """The liveness probe and the poll resolve their target through the same + helper, so they can never disagree about which server is being watched.""" + from comfy_cli import jobs_state + from comfy_cli.command import job_watcher + + state = jobs_state.new( + prompt_id="pid-v6", + client_id="c", + workflow="/w.json", + where="local", + host="::1", + port=9999, + ) + monkeypatch.delenv("COMFY_LOCAL_URL", raising=False) + jobs_state.write(state) + monkeypatch.setattr(job_watcher.time, "sleep", lambda s: None) + calls = _fake_probe(monkeypatch, [False]) + monkeypatch.setattr(job_watcher, "_notify", lambda s: None) + + job_watcher.watch_job("pid-v6", where="local") + + assert job_watcher._resolve_watch_target(state, None, None) == ("[::1]", 9999) + assert set(calls) == {("[::1]", 9999)} + + +def test_server_alive_probe_never_raises(monkeypatch): + """An unexpected probe failure counts as 'down' rather than crashing the + watcher (which would strand the state file at its last status).""" + from comfy_cli.command import job_watcher + + def boom(port, host, *a, **kw): + raise RuntimeError("probe exploded") + + monkeypatch.setattr("comfy_cli.env_checker.check_comfy_server_running", boom) + assert job_watcher._server_alive("127.0.0.1", 8188) is False From 38d0b020704949d3815e91aaea6257b330b92bf2 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Mon, 27 Jul 2026 15:10:55 -0700 Subject: [PATCH 2/2] fix(jobs): only a refused port counts as a dead server; catch restart-orphaned jobs (BE-4751) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses the cursor-review panel on #604. The liveness check went through `check_comfy_server_running`, whose single boolean collapses "the port is refused" with "the server took longer than 5s to answer". A server busy loading a large model into VRAM answers slowly, so three of those in a row filed a terminal `server_died` error against a job that was still running. `_probe_local_server` now classifies the probe itself and only a refused connection counts toward the death streak — a timeout or a non-200 means something is listening, so the streak resets and the poll runs as normal. `Timeout` is caught before `ConnectionError` because `ConnectTimeout` subclasses both. The probe also asks for `/history?max_items=1` rather than the unbounded `/history`, whose body grows without limit on a long-lived server and could time the probe out by itself — feeding the same false verdict. Before writing the verdict the watcher now takes one last poll: the probe hits `/history`, the poll hits `/queue` + `/history/`, so a job that finished right as the probe started failing is reported as completed instead of as a failure with empty outputs. A server OOM-killed and restarted inside the detection window used to reset the streak and then poll a prompt the fresh process had never heard of until the 6h ceiling — the exact hang this change set out to remove. ComfyUI holds its queue and history in memory only, so a restart is always fatal to an in-flight job: once an outage has been seen, a record missing for `_LOST_AFTER_RESTART_S` is now recorded as `server_died` with `restarted: true`. The gate on a previously-observed outage keeps a not-yet-published record from being mistaken for one. Finally, `_finalize_server_died` re-reads the state file and does the read and the write as one locked transaction, so a concurrent `comfy jobs cancel` is no longer clobbered by the watcher's stale in-memory copy (`locking.file_lock` is reentrant per thread, so `jobs_state.write`'s nested acquire is a no-op). `_poll_local_once` returns `(terminal, record_seen)` so the restart guard can tell "no record" from "still running". Co-Authored-By: Claude Opus 5 --- comfy_cli/command/job_watcher.py | 187 +++++++++++++++++++------ comfy_cli/error_codes.py | 4 +- tests/comfy_cli/jobs/test_jobs.py | 219 +++++++++++++++++++++++++++--- 3 files changed, 353 insertions(+), 57 deletions(-) diff --git a/comfy_cli/command/job_watcher.py b/comfy_cli/command/job_watcher.py index 412b952c..a9c2f1bc 100644 --- a/comfy_cli/command/job_watcher.py +++ b/comfy_cli/command/job_watcher.py @@ -40,10 +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 failed liveness probes before we declare a local server dead: -# ~3 polls x 2s ≈ 6s of confirmed unreachability, which tolerates a transient -# blip / a quick server restart while still catching a real crash fast. +# 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") @@ -85,9 +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 liveness probes have - # failed. Reset by any successful probe (see _SERVER_DOWN_CONSECUTIVE_LIMIT). + # 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 @@ -108,38 +131,71 @@ def watch_job( # 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 _server_alive(h, p): - consecutive_down = 0 - terminal = _poll_local_once(state, host=host, port=port) - else: + if _probe_local_server(h, p) == _PROBE_UNREACHABLE: consecutive_down += 1 - if consecutive_down >= _SERVER_DOWN_CONSECUTIVE_LIMIT: - # A state that is already terminal keeps its own verdict — - # a dead server can't invalidate a completed/failed job. - if not state.is_terminal: + 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/`, 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.status = "error" - state.error = { - "code": "server_died", - "message": ( + 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": { + 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 - # Unreachable but still under the limit — a blip or a restart. - # Skip the poll (it can only fail against a server we just - # failed to reach) and re-probe on the next cycle. - terminal = False + 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: @@ -199,23 +255,76 @@ def _resolve_watch_target(state: jobs_state.JobState, host: str | None, port: in return h, p -def _server_alive(host: str, port: int) -> bool: - """Is the local ComfyUI server reachable? +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). - The same helper ``comfy jobs`` uses for its own up/down check, so the - watcher and the CLI never disagree. An unexpected error counts as *down* - rather than propagating — a probe must never crash the watcher. + ``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. """ - from comfy_cli.env_checker import check_comfy_server_running + import requests try: - return bool(check_comfy_server_running(port, host)) + # ``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 False + 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) -> bool: - """Update ``state`` in-place from a local ComfyUI server. Return True if terminal.""" +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) @@ -223,11 +332,11 @@ def _poll_local_once(state: jobs_state.JobState, *, host: str | None, port: int 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 @@ -235,22 +344,22 @@ def _poll_local_once(state: jobs_state.JobState, *, host: str | None, port: int 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 = { diff --git a/comfy_cli/error_codes.py b/comfy_cli/error_codes.py index 5a153393..af9706fb 100644 --- a/comfy_cli/error_codes.py +++ b/comfy_cli/error_codes.py @@ -336,8 +336,8 @@ class ErrorCode: ), ErrorCode( "server_died", - "The local ComfyUI server became unreachable while the job was in flight — " - "it likely crashed or was killed (e.g. an out-of-memory allocation).", + "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( diff --git a/tests/comfy_cli/jobs/test_jobs.py b/tests/comfy_cli/jobs/test_jobs.py index ef8ba04e..ba7ad26e 100644 --- a/tests/comfy_cli/jobs/test_jobs.py +++ b/tests/comfy_cli/jobs/test_jobs.py @@ -11,8 +11,10 @@ import subprocess import sys from pathlib import Path +from types import SimpleNamespace import pytest +import requests from comfy_cli.command import jobs as jobs_mod @@ -668,8 +670,8 @@ def test_snapshot_maps_interrupted_to_cancelled(monkeypatch): def test_poll_local_once_treats_cancelled_as_terminal(monkeypatch): - """_poll_local_once must return True (terminal) and set state.status='cancelled' - when _snapshot reports status='cancelled'.""" + """_poll_local_once must report terminal (and a record it saw) and set + state.status='cancelled' when _snapshot reports status='cancelled'.""" from comfy_cli import jobs_state from comfy_cli.command import job_watcher @@ -678,7 +680,7 @@ def test_poll_local_once_treats_cancelled_as_terminal(monkeypatch): lambda h, p, pid: {"prompt_id": pid, "status": "cancelled", "outputs": []}, ) state = jobs_state.new(prompt_id="pid", client_id="c", workflow="w", where="local") - assert job_watcher._poll_local_once(state, host=None, port=None) is True + assert job_watcher._poll_local_once(state, host=None, port=None) == (True, True) assert state.status == "cancelled" @@ -1383,16 +1385,23 @@ def _watched_local_job(monkeypatch, *, status="queued"): def _fake_probe(monkeypatch, results): """Point the watcher's liveness probe at a scripted sequence of verdicts. - The last verdict repeats forever so a test can't hang on a short script. + Each entry is a ``_PROBE_*`` constant, or a bool as shorthand for + alive/unreachable. The last verdict repeats forever so a test can't hang on + a short script. """ - seq = list(results) + from comfy_cli.command import job_watcher + + seq = [ + (job_watcher._PROBE_ALIVE if r else job_watcher._PROBE_UNREACHABLE) if isinstance(r, bool) else r + for r in results + ] calls = [] - def probe(port, host, *a, **kw): + def probe(host, port): calls.append((host, port)) return seq[min(len(calls) - 1, len(seq) - 1)] - monkeypatch.setattr("comfy_cli.env_checker.check_comfy_server_running", probe) + monkeypatch.setattr(job_watcher, "_probe_local_server", probe) return calls @@ -1406,10 +1415,12 @@ def test_watcher_records_server_died_after_consecutive_failed_probes(monkeypatch _watched_local_job(monkeypatch) calls = _fake_probe(monkeypatch, [False]) - # The poll must never run against a server we just failed to probe. + # Only the one last-chance poll at the limit runs; the earlier cycles must + # not poll a port that just refused them. + polls = [] monkeypatch.setattr( "comfy_cli.command.jobs._snapshot", - lambda h, p, pid: pytest.fail("polled a server that failed its liveness probe"), + lambda h, p, pid: polls.append((h, p)) or None, ) notified = [] monkeypatch.setattr(job_watcher, "_notify", notified.append) @@ -1417,6 +1428,7 @@ def test_watcher_records_server_died_after_consecutive_failed_probes(monkeypatch job_watcher.watch_job("pid-dead", where="local") assert len(calls) == job_watcher._SERVER_DOWN_CONSECUTIVE_LIMIT + assert polls == [("127.0.0.1", 8188)] on_disk = jobs_state.read("pid-dead") assert on_disk is not None assert on_disk.status == "error" @@ -1465,6 +1477,8 @@ def test_watcher_does_not_overwrite_a_terminal_state_when_server_dies(monkeypatc state.outputs = ["done.png"] jobs_state.write(state) _fake_probe(monkeypatch, [False]) + # The dead server has nothing left to report to the last-chance poll. + monkeypatch.setattr("comfy_cli.command.jobs._snapshot", lambda h, p, pid: None) monkeypatch.setattr(job_watcher, "_notify", lambda s: None) job_watcher.watch_job("pid-dead", where="local") @@ -1494,21 +1508,194 @@ def test_watcher_probe_targets_the_same_address_the_poll_uses(monkeypatch): jobs_state.write(state) monkeypatch.setattr(job_watcher.time, "sleep", lambda s: None) calls = _fake_probe(monkeypatch, [False]) + polls = [] + monkeypatch.setattr("comfy_cli.command.jobs._snapshot", lambda h, p, pid: polls.append((h, p)) or None) monkeypatch.setattr(job_watcher, "_notify", lambda s: None) job_watcher.watch_job("pid-v6", where="local") assert job_watcher._resolve_watch_target(state, None, None) == ("[::1]", 9999) assert set(calls) == {("[::1]", 9999)} + assert set(polls) == {("[::1]", 9999)} + + +@pytest.mark.parametrize( + "exc, expected", + [ + # Nothing listening — the only signal that means "dead". + (requests.exceptions.ConnectionError("refused"), "unreachable"), + # Slow, not dead. ConnectTimeout subclasses ConnectionError as well as + # Timeout, so it would be misread as a death if the except order ever + # regressed — which is exactly the false server_died this guards. + (requests.exceptions.ReadTimeout("slow"), "unresponsive"), + (requests.exceptions.ConnectTimeout("slow"), "unresponsive"), + # A probe must never crash the watcher, and never invent a death. + (RuntimeError("probe exploded"), "unresponsive"), + ], +) +def test_probe_classifies_failures_without_inventing_a_death(monkeypatch, exc, expected): + """Only a refused connection counts as unreachable — a busy server (loading + a model into VRAM) times out and must stay merely 'unresponsive'.""" + from comfy_cli.command import job_watcher + + urls = [] + + def fake_get(url, **kw): + urls.append(url) + raise exc + + monkeypatch.setattr(requests, "get", fake_get) + + assert job_watcher._probe_local_server("127.0.0.1", 8188) == expected + # ...and the probe stays cheap: the unbounded /history body grows without + # limit on a long-lived server, which would time the probe out by itself. + assert urls == ["http://127.0.0.1:8188/history?max_items=1"] + + +@pytest.mark.parametrize("status, expected", [(200, "alive"), (404, "unresponsive"), (500, "unresponsive")]) +def test_probe_treats_only_http_200_as_alive(monkeypatch, status, expected): + from comfy_cli.command import job_watcher + + monkeypatch.setattr(requests, "get", lambda url, **kw: SimpleNamespace(status_code=status)) + assert job_watcher._probe_local_server("127.0.0.1", 8188) == expected + + +def test_watcher_keeps_polling_an_alive_but_unresponsive_server(monkeypatch): + """A server that is up but too slow to answer the probe must never be + declared dead — it is still polled, and its job still completes.""" + from comfy_cli import jobs_state + from comfy_cli.command import job_watcher + + _watched_local_job(monkeypatch) + # Unresponsive forever: if it counted toward the death streak, the watcher + # would file server_died instead of reading the completion below. + _fake_probe(monkeypatch, [job_watcher._PROBE_UNRESPONSIVE]) + monkeypatch.setattr( + "comfy_cli.command.jobs._snapshot", + lambda h, p, pid: {"prompt_id": pid, "status": "completed", "outputs": ["slow.png"]}, + ) + monkeypatch.setattr(job_watcher, "_notify", lambda s: None) + job_watcher.watch_job("pid-dead", where="local") + + on_disk = jobs_state.read("pid-dead") + assert on_disk is not None + assert on_disk.status == "completed" + assert on_disk.error is None -def test_server_alive_probe_never_raises(monkeypatch): - """An unexpected probe failure counts as 'down' rather than crashing the - watcher (which would strand the state file at its last status).""" + +def test_watcher_recovers_a_job_that_finished_during_the_down_streak(monkeypatch): + """The last-chance poll at the limit wins over the server_died verdict: a + job that completed as the probes started failing must not be reported as a + failure with empty outputs.""" + from comfy_cli import jobs_state from comfy_cli.command import job_watcher - def boom(port, host, *a, **kw): - raise RuntimeError("probe exploded") + _watched_local_job(monkeypatch) + _fake_probe(monkeypatch, [False]) + monkeypatch.setattr( + "comfy_cli.command.jobs._snapshot", + lambda h, p, pid: {"prompt_id": pid, "status": "completed", "outputs": ["late.png"]}, + ) + monkeypatch.setattr(job_watcher, "_notify", lambda s: None) + + job_watcher.watch_job("pid-dead", where="local") + + on_disk = jobs_state.read("pid-dead") + assert on_disk is not None + assert on_disk.status == "completed" + assert on_disk.outputs == ["late.png"] + assert on_disk.error is None + - monkeypatch.setattr("comfy_cli.env_checker.check_comfy_server_running", boom) - assert job_watcher._server_alive("127.0.0.1", 8188) is False +def test_watcher_records_server_died_when_a_restart_loses_the_job(monkeypatch): + """A server OOM-killed and restarted inside the detection window makes the + next probe succeed, but the fresh process has no record of the prompt. That + must be recorded as a death rather than polled until the 6h ceiling.""" + from comfy_cli import jobs_state + from comfy_cli.command import job_watcher + + _watched_local_job(monkeypatch, status="running") + # One outage, then back up — the streak resets well short of the limit. + _fake_probe(monkeypatch, [False, True]) + # ...but the restarted server never heard of this prompt. + monkeypatch.setattr("comfy_cli.command.jobs._snapshot", lambda h, p, pid: None) + monkeypatch.setattr(job_watcher, "_LOST_AFTER_RESTART_S", 0.0) + monkeypatch.setattr(job_watcher, "_notify", lambda s: None) + + job_watcher.watch_job("pid-dead", where="local") + + on_disk = jobs_state.read("pid-dead") + assert on_disk is not None + assert on_disk.status == "error" + assert on_disk.error["code"] == "server_died" + assert on_disk.error["details"]["restarted"] is True + assert on_disk.error["details"]["last_status"] == "running" + + +def test_watcher_does_not_call_a_missing_record_a_restart_without_an_outage(monkeypatch): + """No outage seen → a record the server hasn't published yet is just a slow + start, not a death. The restart guard must stay latched off.""" + from comfy_cli import jobs_state + from comfy_cli.command import job_watcher + + _watched_local_job(monkeypatch) + _fake_probe(monkeypatch, [True]) + monkeypatch.setattr(job_watcher, "_LOST_AFTER_RESTART_S", 0.0) + monkeypatch.setattr(job_watcher, "_notify", lambda s: None) + + seen = [] + + def snapshot(h, p, pid): + seen.append(pid) + # Missing for a while, then it shows up and completes. + if len(seen) < 4: + return None + return {"prompt_id": pid, "status": "completed", "outputs": ["ok.png"]} + + monkeypatch.setattr("comfy_cli.command.jobs._snapshot", snapshot) + + job_watcher.watch_job("pid-dead", where="local") + + on_disk = jobs_state.read("pid-dead") + assert on_disk is not None + assert on_disk.status == "completed" + assert on_disk.error is None + + +def test_watcher_does_not_clobber_a_concurrent_cancel_with_server_died(monkeypatch): + """`comfy jobs cancel` writing a verdict while the watcher is failing probes + must survive: the verdict is re-read from disk, not taken from the + watcher's stale in-memory copy.""" + from comfy_cli import jobs_state + from comfy_cli.command import job_watcher + + _watched_local_job(monkeypatch) + _fake_probe(monkeypatch, [False]) + monkeypatch.setattr("comfy_cli.command.jobs._snapshot", lambda h, p, pid: None) + + def poll_and_cancel_out_of_band(state, **kw): + """The last-chance poll, racing a concurrent `comfy jobs cancel`. + + The cancel lands on disk while the watcher still holds its stale + pre-cancel copy of the state — the exact window the verdict re-read + has to cover. + """ + other = jobs_state.read("pid-dead") + other.status = "cancelled" + other.error = {"code": "cancelled", "message": "user cancelled", "details": {}} + jobs_state.write(other) + return False, False + + monkeypatch.setattr(job_watcher, "_poll_local_once", poll_and_cancel_out_of_band) + notified = [] + monkeypatch.setattr(job_watcher, "_notify", notified.append) + + job_watcher.watch_job("pid-dead", where="local") + + on_disk = jobs_state.read("pid-dead") + assert on_disk is not None + assert on_disk.status == "cancelled" + assert on_disk.error["code"] == "cancelled" + # The notification reports the verdict that actually stands. + assert notified[0].status == "cancelled"