From 4d03724dcf22a11314f599b1063beebf5eafa8ce Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Wed, 9 Sep 2026 15:27:16 -0700 Subject: [PATCH 1/3] feat(models): queue a model request with submit, subscribe and handle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `models.run` holds one connection open until the generation is finished. This adds the queued form of the same request for a caller who cannot wait that long: `client.models.submit(model, arguments)` returns a `RequestHandle` as soon as the server accepts the request, and the generation is collected later — including from another process, since `client.models.handle(model, request_id)` rebuilds the handle from nothing but the two ids that address it. The handle carries `status()`, `get()`, `cancel()` and `iter_events()`, and `models.subscribe()` is submit plus polling plus collection in one call, with an `on_queue_update=` callback for progress and a client-side `timeout=` that makes a best-effort cancel before it raises. The shape is the one `comfy_sdk.jobs` already uses for workflow jobs rather than a second idiom, and the queue itself stays server-owned: ordering, admission, retries, timeouts, billing and expiry are not reimplemented here. Three properties are contract rather than implementation detail: - Polling is authoritative and paced by the server's own `Retry-After` when it names one, by an adaptive backoff when it does not. There is no stream on this surface, so `iter_events` is the poll loop with its updates exposed. - A `COMPLETED` status carrying an `error_type` — how the server reports a failed *and* a cancelled request — raises the typed `RouterError` subclass. A 200 with an error payload is never handed back as a result. - Each `submit` *call* mints one fresh `Idempotency-Key`, so two deliberate submits are two requests while a transport retry inside one call replays the original rather than queueing a second generation. `AsyncComfy` awaits the same method names with the same arguments in the same order; the existing introspection parity test discovers the new pair and compares it. There is no `submit_async`, for the same reason there is no `run_async`. The four routes are hand-bound and confined to the `_MODEL_REQUEST*` constants in `comfy_low.transport`: the contract declaring them is authored but held, so the vendored Router spec does not carry them yet and there is nothing for the contract test to pin them against. That is the one part of this change a spec sync is expected to correct. `client.models.run` is unchanged in behaviour and signature. --- CHANGELOG.md | 28 ++ README.md | 114 ++++- src/comfy_low/transport.py | 233 +++++++++- src/comfy_sdk/__init__.py | 6 + src/comfy_sdk/client.py | 5 +- src/comfy_sdk/exceptions.py | 3 +- src/comfy_sdk/model_requests.py | 537 ++++++++++++++++++++++ src/comfy_sdk/models.py | 268 ++++++++++- src/comfy_sdk/router_exceptions.py | 58 +++ tests/conftest.py | 209 ++++++++- tests/test_models_queue.py | 687 +++++++++++++++++++++++++++++ tests/test_sync_async_parity.py | 11 +- 12 files changed, 2145 insertions(+), 14 deletions(-) create mode 100644 src/comfy_sdk/model_requests.py create mode 100644 tests/test_models_queue.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 5f4968d..80fca15 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,34 @@ notes for each version. ### Added +- **The queued model surface** — `client.models.submit()`, `client.models.subscribe()` + and `client.models.handle()`, plus the `RequestHandle` they hand back + (`status()`, `get()`, `cancel()`, `iter_events()`). `models.run` holds one + connection open until the generation is finished; `submit` returns as soon as + the server accepts the request, and the generation is collected later — + including from another process, since `client.models.handle(model, + request_id)` rebuilds the handle from nothing but the two ids. `subscribe` is + submit plus polling plus collection in one call, with an `on_queue_update=` + callback for progress and a client-side `timeout=` that makes a best-effort + cancel before it raises. `AsyncComfy` awaits the same method names with the + same arguments in the same order — there is no `submit_async`, for the reason + there is no `run_async`. Three properties are contract rather than + implementation: polling is authoritative and paced by the server's own + `Retry-After` when it names one; a `COMPLETED` status carrying an + `error_type` — which is how the server reports a failed *or* a cancelled + request — raises the typed `RouterError` subclass rather than being returned + as a successful result; and each `submit` **call** mints one fresh + `Idempotency-Key`, so two deliberate submits are two requests while a + transport retry inside one call replays the original. The queue itself stays + entirely server-owned: ordering, admission, retries, timeouts, billing and + expiry are not reimplemented here. The surface is gated server side — a + caller it is not switched on for is answered `403 not_enabled`, which arrives + as `comfy_sdk.router_exceptions.NotEnabled`. `client.models.run` is unchanged + in behaviour and signature. +- `comfy_sdk.router_exceptions.error_from_completion()` — the typed exception a + completed-but-failed queued request reports, or `None`. Public because the + rule it encodes ("a `200` is not the same thing as a success on this + surface") is one a caller reading a raw payload has to apply too. - Every exception `client.models.run()` raises **for a failed call** now carries the `Idempotency-Key` it was made under, on `.idempotency_key` — the typed `RouterError` buckets, a `RouterError` whose `error_type` this version diff --git a/README.md b/README.md index ac04659..ed7d8ea 100644 --- a/README.md +++ b/README.md @@ -395,7 +395,9 @@ client's `COMFY_BASE_URL`. See two variables. `base_url` and `timeout` are a read-only view of that configuration; model -operations are added to this namespace as they land. +operations are added to this namespace as they land. There are two ways to run +a model on it — `run`, which waits, and `submit`, which queues — and they send +the same request. ### `models.run` — one call, one result @@ -445,6 +447,103 @@ indefinitely. Each call also sends a fresh `Idempotency-Key`, so an accidental exact resend is rejected by the server instead of billing a second generation; pass `idempotency_key=` to choose the value yourself. +### `models.submit` — queue it, collect it later + +`run` holds one connection open until the generation is finished. When the +caller cannot wait that long — a web request that has to return now, a worker +that submits in one process and collects in another, a batch that should be in +flight all at once — submit it to the queue instead: + +```python +handle = client.models.submit("fal-ai/flux-pro", {"prompt": "a cat"}) + +handle.request_id # 'req_...' — all another process needs +handle.status().status # 'IN_QUEUE' / 'IN_PROGRESS' / 'COMPLETED' +result = handle.get() # blocks until complete, returns the provider payload +``` + +`submit` sends the same request `run` does — the same model id, the same native +body — and returns as soon as the server has **accepted** it. The queue is the +server's: ordering, admission, retries, timeouts, billing and expiry are all +decided there, and this SDK adds polling and ergonomics on top of it and +nothing else. + +The handle carries four operations: + +| | | +|---|---| +| `handle.status()` | one authoritative poll, returned as a `QueueUpdate` (`status`, `queue_position`, `error_type`, `retry_after`, `raw`) | +| `handle.get(timeout=None)` | poll to completion, then return the provider's own payload — the same value `run` would have returned | +| `handle.cancel()` | ask the server to cancel. A request, not a guarantee: a request that already completed stays completed | +| `handle.iter_events(timeout=None)` | the poll loop with its updates exposed — yields the first observation, every change of status or queue position, and the completion | + +Polling is **poll-authoritative**: there is no stream to reconcile against on +this surface, and `iter_events` is the poll loop rather than SSE. It backs off +adaptively, and a `Retry-After` the server names on a poll beats that schedule +— the server knows its own pace. + +**A `200` is not the same thing as a success here.** The server reports a +failed *and* a cancelled request as `COMPLETED` carrying an `error_type`, so +`get()` raises the matching typed exception from +[`comfy_sdk.router_exceptions`](#typed-errors) rather than handing the failure +back as a result. `iter_events` deliberately does not raise — it is a view of +the queue's progress, and `get()` is the one that collects. + +Rehydrate a handle in another process from the two ids that address the +request, with no call made: + +```python +handle = client.models.handle("fal-ai/flux-pro", request_id) +result = handle.get() +``` + +Both ids are needed because both address the route +(`/v2/models/{provider}/{model}/requests/{request_id}`), and both are validated +locally before anything is sent. + +### `models.subscribe` — submit, follow, collect + +```python +def on_update(update): + print(update.status, update.queue_position) + +result = client.models.subscribe( + "fal-ai/flux-pro", {"prompt": "a cat"}, on_queue_update=on_update, timeout=300 +) +``` + +`submit` + poll + `get`, in one call, for a caller who does want to wait but +also wants to show progress. `timeout=` is a **client-side** bound with no +server-side meaning; when it runs out, `subscribe` makes a best-effort +`cancel()` — so a caller who has stopped waiting is not still paying for a +generation nobody will collect — and then raises `TimeoutError`. Use `submit` +when the request should outlive the caller's patience. + +Each `submit` **call** mints one fresh `Idempotency-Key`: two deliberate +submits of the same input are two requests, while a transport-level retry +inside one call keeps the one key and replays the original rather than queueing +a second generation. Pass `idempotency_key=` to choose it yourself — the case +that earns it is a lost response, where the request may have been accepted and +its id lost with the reply; every exception carries the key on +`.idempotency_key` for exactly that. + +The awaitable form is the async client, with the identical names, arguments and +argument order — there is no `submit_async`, for the same reason there is no +`run_async`: + +```python +async with AsyncComfy(api_key="comfyui-...") as client: + handle = await client.models.submit("fal-ai/flux-pro", {"prompt": "a cat"}) + async for update in handle.iter_events(): + print(update.status) + result = await handle.get() +``` + +This surface is **gated server side**. A caller the queue is not switched on +for is answered `403 not_enabled`, which arrives as +`comfy_sdk.router_exceptions.NotEnabled` — nothing about the request is wrong, +and it is terminal: do not retry it. + ### Retrying a run without paying for it twice A failed `models.run` is retried automatically. **Every attempt of one call @@ -585,7 +684,8 @@ All three read as `None` rather than raising on any exception `models.run` raises, so a handler never has to guard the attribute access itself. `idempotency_key` is `None` on errors from *other* surfaces, though — it is -`models.run` that records it, and `submit()` sends a key without stamping one. +`models.run` and `models.submit` that record it, and the workflow surface's +`client.submit()` sends a key without stamping one. So `None` means "this SDK did not record a key for you", **not** "no key was sent, resend freely": check for it before replaying, as the snippet above does, rather than passing it straight back into `idempotency_key=` where `None` means @@ -706,3 +806,13 @@ python scripts/check_drift.py # same check CI runs; fails if committed models Releases are published to PyPI from a GitHub Release (tag `vX.Y.Z`) by [`.github/workflows/publish.yml`](.github/workflows/publish.yml), using PyPI's Trusted Publishing (OIDC) — no API token is stored in this repo. + +**Pre-releases** (a release candidate) go out the same way, with one thing to +get right: the workflow validates the tag as **SemVer**, where a pre-release is +a `-`-separated segment. So tag it `v0.2.0-rc1` (or `v0.2.0-rc.1`) — **not** +`v0.2.0rc1`, which is PEP 440's own spelling and is rejected by that check. The +build normalises the accepted form to the PEP 440 version anyway, so +`v0.2.0-rc1` produces `comfy_sdk-0.2.0rc1-py3-none-any.whl`, which `pip install +comfy-sdk` will not pick up without `--pre`. That last part is the point of +shipping one: a pre-release reaches the people who ask for it and nobody +else. diff --git a/src/comfy_low/transport.py b/src/comfy_low/transport.py index 8b564d7..ec8af11 100644 --- a/src/comfy_low/transport.py +++ b/src/comfy_low/transport.py @@ -10,9 +10,10 @@ * **per-request timeout / abort** — every method takes ``timeout`` and the raw httpx cancellation applies. -One binding is *not* backed by an ``operationId`` *in this module's sense*: -``post_model_run``. It targets a different surface — Comfy Router, on its own -host (:data:`ROUTER_BASE_URL`) — rather than the ``/api/v2`` deployment the rest +One family of bindings is *not* backed by an ``operationId`` *in this module's +sense*: the model bindings. ``post_model_run`` targets a different surface — +Comfy Router, on its own host +(:data:`ROUTER_BASE_URL`) — rather than the ``/api/v2`` deployment the rest of these methods speak to, and it is declared by a *second* vendored contract, ``spec/router-openapi.yaml`` (``operationId: runRouterModel``, path ``/v2/models/{provider}/{model}``). Nothing is generated from that second file @@ -23,6 +24,15 @@ — and ``tests/test_router_spec_contract.py`` plus ``scripts/check_drift.py`` fail if that constant and the vendored path disagree. +The four ``post_model_submit`` / ``get_model_request_status`` / +``get_model_request_result`` / ``post_model_request_cancel`` bindings are the +same story one step earlier: they are the *queued* form of that one operation, +and the contract declaring them is authored but held, so the vendored Router +spec does not carry them yet and there is nothing for the contract test to pin +them against. Their routes are confined to the ``_MODEL_REQUEST*`` constants +for exactly the reason the run path was, and they are the one part of this +change a spec sync is expected to correct. + This layer contains no orchestration, retries, hashing, or reconnection — those live in ``comfy_sdk``. """ @@ -89,6 +99,26 @@ #: constant does not follow it. _MODEL_RUN_PATH_TEMPLATE = "/v2/models/{provider}/{model}" +#: Routes for the *queued* form of a model request — submit, poll, collect, +#: cancel. They extend :data:`_MODEL_RUN_PATH_TEMPLATE` with a ``requests`` +#: collection under the same model-ID-addressed prefix, because a queued +#: request is the same operation on the same model, reached without holding the +#: connection open for it. +#: +#: **These four are not in the vendored contract yet.** The queue operations +#: are authored upstream but held, and the one-way sync into +#: ``spec/router-openapi.yaml`` strips a held operation — so unlike +#: :data:`_MODEL_RUN_PATH_TEMPLATE`, which +#: ``tests/test_router_spec_contract.py`` pins against the vendored file, these +#: are hand-bound with nothing to pin them to. They are gathered here, in one +#: place and nowhere else, precisely so the sync that publishes them is a +#: four-line diff plus the assertion that pins them — exactly as the run path +#: was hand-bound before its own spec arrived. +_MODEL_REQUESTS_PATH_TEMPLATE = _MODEL_RUN_PATH_TEMPLATE + "/requests" +_MODEL_REQUEST_PATH_TEMPLATE = _MODEL_REQUESTS_PATH_TEMPLATE + "/{request_id}" +_MODEL_REQUEST_STATUS_PATH_TEMPLATE = _MODEL_REQUEST_PATH_TEMPLATE + "/status" +_MODEL_REQUEST_CANCEL_PATH_TEMPLATE = _MODEL_REQUEST_PATH_TEMPLATE + "/cancel" + _DEFAULT_PORTS = {"http": 80, "https": 443} @@ -170,6 +200,75 @@ def model_run_request( return path, body, headers +def parse_request_id(request_id: str) -> str: + """``request_id`` unchanged, or an error for one that cannot address a route. + + A queued request's id is the last segment of + ``/v2/models/{provider}/{model}/requests/{request_id}``, so it is subject to + exactly the discipline :func:`parse_model_id` applies to the two segments + before it: a wrong *type* raises ``TypeError``, a wrong *value* raises + ``ValueError``, and both fail locally rather than being pasted into a URL + and answered by whatever route they land on. ``.``/``..`` are refused rather + than encoded for the same reason they are there — ``quote`` leaves ``.`` + alone, so a dot segment would survive into the path and walk the route on + any intermediary that normalizes it. + """ + if not isinstance(request_id, str): + raise TypeError(f"request id must be a str, got {type(request_id).__name__}") + if not request_id: + raise ValueError("request id must not be empty") + if "/" in request_id: + raise ValueError( + f"request id must be a single path segment — it addresses " + f"{_MODEL_REQUEST_PATH_TEMPLATE}; got {request_id!r}" + ) + if request_id in (".", ".."): + raise ValueError( + f"request id must not be '.' or '..' — it would traverse the request path " + f"rather than name a request; got {request_id!r}" + ) + return request_id + + +def model_submit_request( + model: str, + arguments: Mapping[str, Any], + idempotency_key: str | None, +) -> tuple[str, dict[str, Any], dict[str, str]]: + """Sans-IO ``(path, json_body, headers)`` for one *queued* model request. + + Identical in shape to :func:`model_run_request` — the model id addresses the + request and the body is the partner model's own native JSON input, verbatim + — differing only in the route it targets. That sameness is deliberate: a + caller moves between the awaited and the queued form by choosing a method, + not by rewriting the request. + """ + provider, name = parse_model_id(model) + path = _MODEL_REQUESTS_PATH_TEMPLATE.format( + provider=quote(provider, safe=""), model=quote(name, safe="") + ) + body: dict[str, Any] = dict(arguments) + headers = {"Idempotency-Key": idempotency_key} if idempotency_key else {} + return path, body, headers + + +def model_request_path(model: str, request_id: str, template: str) -> str: + """Sans-IO path for one queued request, from ``template``. + + ``template`` is one of the ``_MODEL_REQUEST_*`` constants; passing it in + rather than branching on an operation name keeps every route this family + reaches spelled in exactly one place. Each of the three segments is + percent-encoded with ``safe=""`` so nothing in it can add a path segment, a + query or a fragment. + """ + provider, name = parse_model_id(model) + return template.format( + provider=quote(provider, safe=""), + model=quote(name, safe=""), + request_id=quote(parse_request_id(request_id), safe=""), + ) + + def _build_user_agent(client_info: str | None) -> str: """SDK identity sent on every request. This is request metadata (not telemetry — no phone-home), so adoption is measurable server-side from @@ -794,6 +893,89 @@ def post_model_run( resp = self.raw_request("POST", url, headers=headers, json=body, timeout=timeout) return self._p.parse_or_raise(resp, (200, 201)) + # -- models: the queued form ------------------------------------------ + # + # These four return the response HEADERS alongside the decoded body, which + # nothing else in this transport does. The reason is specific to the queue + # rather than a change of house style: a poll's whole job is to say when to + # ask again, and the server says it on `Retry-After` — a header, on a + # SUCCESS response, which `parse_or_raise` has no way to hand back. Reading + # it is what makes the layer above pace itself against the server's own + # answer instead of only against a local backoff schedule. The same channel + # carries `X-Comfy-Request-Id`, which is what a failure reported inside a + # 200 body has to be attributable by. + def post_model_submit( + self, + model: str, + arguments: Mapping[str, Any], + *, + idempotency_key: str | None = None, + timeout: Any = _UNSET, + ) -> tuple[dict[str, Any], httpx.Headers]: + """POST ``{router_base_url}/v2/models/{provider}/{model}/requests`` — queued. + + The queued sibling of :meth:`post_model_run`: same host, same + model-ID-addressed prefix, same verbatim body, but the server answers as + soon as the request is *accepted* rather than holding the connection + until the generation is finished. The response names the request id + every later call in this family is addressed by. + + The timeout is therefore the client's ordinary default rather than + :data:`MODEL_RUN_TIMEOUT` — nothing here waits on a generation. + + Raises ``TypeError``/``ValueError`` from :func:`parse_model_id` before + any request when ``model`` is not a ``{provider}/{model}`` id. + """ + path, body, headers = model_submit_request(model, arguments, idempotency_key) + url = self._p.router_base_url + path + resp = self.raw_request("POST", url, headers=headers, json=body, timeout=timeout) + return self._p.parse_or_raise(resp, (200, 201, 202)), resp.headers + + def get_model_request_status( + self, model: str, request_id: str, *, timeout: Any = _UNSET + ) -> tuple[dict[str, Any], httpx.Headers]: + """GET the queue status of one submitted request — the authoritative read. + + This is the source of truth for how far a queued request has got, and + the only one: there is no stream to reconcile against on this surface. + """ + url = self._p.router_base_url + model_request_path( + model, request_id, _MODEL_REQUEST_STATUS_PATH_TEMPLATE + ) + resp = self.raw_request("GET", url, timeout=timeout) + return self._p.parse_or_raise(resp, (200,)), resp.headers + + def get_model_request_result( + self, model: str, request_id: str, *, timeout: Any = _UNSET + ) -> tuple[dict[str, Any], httpx.Headers]: + """GET the finished result of one submitted request. + + The body is the provider's own payload, exactly as + :meth:`post_model_run` returns it — this route is where a queued + request's result is collected, not a differently-shaped one. + """ + url = self._p.router_base_url + model_request_path( + model, request_id, _MODEL_REQUEST_PATH_TEMPLATE + ) + resp = self.raw_request("GET", url, timeout=timeout) + return self._p.parse_or_raise(resp, (200,)), resp.headers + + def post_model_request_cancel( + self, model: str, request_id: str, *, timeout: Any = _UNSET + ) -> tuple[dict[str, Any], httpx.Headers]: + """POST a cancellation for one submitted request. + + A request, not a guarantee — a deployment that answers ``204`` gives an + empty body, which ``parse_or_raise`` returns as ``{}``. The + authoritative state is whatever :meth:`get_model_request_status` says + next, exactly as it is for a job. + """ + url = self._p.router_base_url + model_request_path( + model, request_id, _MODEL_REQUEST_CANCEL_PATH_TEMPLATE + ) + resp = self.raw_request("POST", url, timeout=timeout) + return self._p.parse_or_raise(resp, (200, 202, 204)), resp.headers + class AsyncComfyLow: """Asynchronous protocol bindings — mirrors :class:`ComfyLow`.""" @@ -1109,6 +1291,51 @@ async def post_model_run( resp = await self.raw_request("POST", url, headers=headers, json=body, timeout=timeout) return self._p.parse_or_raise(resp, (200, 201)) + # -- models: the queued form ------------------------------------------ + async def post_model_submit( + self, + model: str, + arguments: Mapping[str, Any], + *, + idempotency_key: str | None = None, + timeout: Any = _UNSET, + ) -> tuple[dict[str, Any], httpx.Headers]: + """Async :meth:`ComfyLow.post_model_submit`.""" + path, body, headers = model_submit_request(model, arguments, idempotency_key) + url = self._p.router_base_url + path + resp = await self.raw_request("POST", url, headers=headers, json=body, timeout=timeout) + return self._p.parse_or_raise(resp, (200, 201, 202)), resp.headers + + async def get_model_request_status( + self, model: str, request_id: str, *, timeout: Any = _UNSET + ) -> tuple[dict[str, Any], httpx.Headers]: + """Async :meth:`ComfyLow.get_model_request_status`.""" + url = self._p.router_base_url + model_request_path( + model, request_id, _MODEL_REQUEST_STATUS_PATH_TEMPLATE + ) + resp = await self.raw_request("GET", url, timeout=timeout) + return self._p.parse_or_raise(resp, (200,)), resp.headers + + async def get_model_request_result( + self, model: str, request_id: str, *, timeout: Any = _UNSET + ) -> tuple[dict[str, Any], httpx.Headers]: + """Async :meth:`ComfyLow.get_model_request_result`.""" + url = self._p.router_base_url + model_request_path( + model, request_id, _MODEL_REQUEST_PATH_TEMPLATE + ) + resp = await self.raw_request("GET", url, timeout=timeout) + return self._p.parse_or_raise(resp, (200,)), resp.headers + + async def post_model_request_cancel( + self, model: str, request_id: str, *, timeout: Any = _UNSET + ) -> tuple[dict[str, Any], httpx.Headers]: + """Async :meth:`ComfyLow.post_model_request_cancel`.""" + url = self._p.router_base_url + model_request_path( + model, request_id, _MODEL_REQUEST_CANCEL_PATH_TEMPLATE + ) + resp = await self.raw_request("POST", url, timeout=timeout) + return self._p.parse_or_raise(resp, (200, 202, 204)), resp.headers + def _looks_like_path(s: str) -> bool: return s.startswith("http") or s.startswith("/") diff --git a/src/comfy_sdk/__init__.py b/src/comfy_sdk/__init__.py index b2ea25d..7bf3dde 100644 --- a/src/comfy_sdk/__init__.py +++ b/src/comfy_sdk/__init__.py @@ -66,6 +66,7 @@ WorkflowFormatUi, ) from .jobs import AsyncJob, Job, JobWorkflow +from .model_requests import COMPLETED, AsyncRequestHandle, QueueUpdate, RequestHandle from .outputs import AsyncOutput, DownloadUrl, Output from .retry import DEFAULT_RETRY, NO_RETRY, RetryPolicy from .workflows import Workflow, WorkflowFactory @@ -97,6 +98,11 @@ "Job", "AsyncJob", "JobWorkflow", + # queued model requests + "RequestHandle", + "AsyncRequestHandle", + "QueueUpdate", + "COMPLETED", "Output", "AsyncOutput", "DownloadUrl", diff --git a/src/comfy_sdk/client.py b/src/comfy_sdk/client.py index a7aa403..74417c5 100644 --- a/src/comfy_sdk/client.py +++ b/src/comfy_sdk/client.py @@ -227,8 +227,9 @@ class Comfy: ``retry`` is the policy ``client.models`` calls fail under — :data:`~comfy_sdk.retry.DEFAULT_RETRY` unless replaced, and :data:`~comfy_sdk.retry.NO_RETRY` to make every call exactly one attempt. - It does not govern ``submit``/``run``, whose 429 handling follows the - server's own ``Retry-After`` instead. + It does not govern this client's own ``submit``/``run`` (the workflow + surface), whose 429 handling follows the server's own ``Retry-After`` + instead. """ def __init__( diff --git a/src/comfy_sdk/exceptions.py b/src/comfy_sdk/exceptions.py index be2986a..40396d0 100644 --- a/src/comfy_sdk/exceptions.py +++ b/src/comfy_sdk/exceptions.py @@ -23,7 +23,8 @@ class ComfyError(Exception): """Base for every SDK-level error.""" #: The ``Idempotency-Key`` the failed call was made under. Populated by - #: :meth:`comfy_sdk.models.Models.run` and its async twin, which are the + #: :meth:`comfy_sdk.models.Models.run` and :meth:`comfy_sdk.models.Models.submit` + #: (and their async twins), which are the #: operations that pass a key to :func:`translating`; ``None`` everywhere #: else — including on operations that *do* send a key but do not stamp it #: (``Comfy.submit()``), and on an exception constructed by hand. So diff --git a/src/comfy_sdk/model_requests.py b/src/comfy_sdk/model_requests.py new file mode 100644 index 0000000..ade0392 --- /dev/null +++ b/src/comfy_sdk/model_requests.py @@ -0,0 +1,537 @@ +"""Handles for *queued* model requests — ``client.models.submit`` and friends. + +The queued form of a model run. ``models.run`` holds one connection open until +the generation is finished; ``models.submit`` hands back a +:class:`RequestHandle` the moment the server accepts the request, and the +generation is collected later — from another coroutine, another process, or +another machine, since a handle is rehydratable from nothing but the model id +and the request id (``client.models.handle``). + +**The server owns the queue.** Ordering, admission, retries, timeouts, billing +and expiry are all decided server side; this module adds polling and ergonomics +and nothing else. Anything here that looked like queue *behaviour* — a local +position estimate, a client-side retry of a rejected submit, an expiry clock — +would be a second, disagreeing implementation of a decision that has already +been made somewhere authoritative. + +The shape is deliberately the one :mod:`comfy_sdk.jobs` already uses for +workflow jobs — submit, hold a handle, poll to a terminal state with adaptive +backoff, collect or cancel — rather than a second idiom for the same thing. Two +differences follow from the surface rather than from taste: + +* **Terminal means ``COMPLETED``, and a failure is a completion.** The server + reports a failed or cancelled request as ``COMPLETED`` carrying an + ``error_type``, so a ``200`` is not the same thing as a success. Every path + that reads a completion runs it through + :func:`~comfy_sdk.router_exceptions.error_from_completion` and raises the + typed router exception, which is what keeps a failed generation from being + returned as a result. A status this SDK version has never heard of is treated + as *not yet terminal* — the set grows on the server's release cycle, and + guessing that an unknown state is finished would collect a result that does + not exist yet. + +* **There is no event stream.** :meth:`RequestHandle.iter_events` is the poll + loop with its updates exposed, not SSE — the queue's streaming surfaces are + not part of this. Polls are paced by the server's own ``Retry-After`` when it + names one and by an adaptive backoff when it does not. +""" + +from __future__ import annotations + +import asyncio +import time +from collections.abc import AsyncGenerator, Awaitable, Callable, Generator, Iterator, Mapping +from dataclasses import dataclass, field +from typing import Any + +import httpx + +from comfy_low.errors import ApiError +from comfy_low.transport import AsyncComfyLow, ComfyLow, parse_request_id + +from . import _core +from .exceptions import ComfyError, translating +from .retry import DEFAULT_RETRY, Retrier, RetryPolicy +from .router_exceptions import RouterError, error_from_completion + +#: The one terminal queue status. Deliberately a single value rather than a +#: :data:`comfy_sdk._core.TERMINAL`-style set: the server does not express a +#: cancel or a failure as its own status, it expresses them as this status plus +#: an ``error_type``. Adding "CANCELED" here on the assumption it exists would +#: strand a caller whose request really did reach ``COMPLETED``. +COMPLETED = "COMPLETED" + +#: Failures a poll is retried through, matching ``models.run``'s tuple exactly. +#: Being listed is not "retryable" — :meth:`~comfy_sdk.retry.RetryPolicy.should_retry` +#: decides that. What it buys is that anything else propagates untouched. +_CANDIDATE_FAILURES = (ApiError, RouterError, httpx.TransportError) + +#: What a best-effort cancel is allowed to fail with. Narrow on purpose: this +#: is only ever used to keep a cancel from masking the ``TimeoutError`` that +#: prompted it, and swallowing a bare ``Exception`` there would hide a bug in +#: this SDK just as readily as it hides an unreachable server. +_CANCEL_FAILURES = (ComfyError, httpx.HTTPError) + +_now = time.monotonic + + +@dataclass(frozen=True) +class QueueUpdate: + """One observation of a queued request's place in the queue. + + What :meth:`RequestHandle.status` returns, what + :meth:`RequestHandle.iter_events` yields, and what ``subscribe``'s + ``on_queue_update`` callback is handed. + + ``status`` is an **open string**, not an enum: a status added server side + has to reach the caller as itself rather than as a decoding failure. Compare + it against :data:`COMPLETED`, or read :attr:`is_completed`. + """ + + #: The request id this update is about — the same id + #: ``client.models.handle`` rehydrates from. + request_id: str + #: The server's status for the request, verbatim. ``""`` when the response + #: named none at all (a cancel answered ``204``, say). + status: str + #: Position in the queue when the server reported one, else ``None``. It is + #: the server's number, never computed here. + queue_position: int | None = None + #: The failure bucket a completion carries, when it carries one. Present + #: here as data; the raising is done by the methods that collect a result. + error_type: str | None = None + #: Seconds the server asked the caller to wait before polling again, from + #: ``Retry-After``. ``None`` when it named no pace, in which case the + #: adaptive backoff decides. + retry_after: int | None = None + #: The decoded response body, unmodified — the escape hatch for a field + #: this dataclass does not model yet. + raw: Mapping[str, Any] = field(default_factory=dict) + + @property + def is_completed(self) -> bool: + """Whether the request has reached the queue's one terminal status.""" + return self.status == COMPLETED + + def __repr__(self) -> str: + position = "" if self.queue_position is None else f", queue_position={self.queue_position}" + error = "" if self.error_type is None else f", error_type={self.error_type!r}" + return ( + f"QueueUpdate(request_id={self.request_id!r}, status={self.status!r}{position}{error})" + ) + + +def _retry_after_seconds(headers: httpx.Headers) -> int | None: + """``Retry-After`` as a positive whole number of seconds, or ``None``. + + Non-positive and unparseable values are dropped rather than honoured, for + the reason :class:`~comfy_sdk.retry.Retrier` gives for the same check: a + pace of zero names no pace, and taking it verbatim turns a server that + keeps answering it into a zero-delay poll loop. + """ + raw = headers.get("Retry-After") + if raw is None: + return None + try: + seconds = int(raw) + except ValueError: + return None + return seconds if seconds > 0 else None + + +def _update_from( + payload: Mapping[str, Any], headers: httpx.Headers, *, request_id: str +) -> QueueUpdate: + """Build a :class:`QueueUpdate` from one response. + + ``request_id`` is the id the call was addressed by, used when the body does + not repeat it — which is what keeps an update from a body-less ``204`` + cancel still identifying the request it is about. + """ + body_id = payload.get("request_id") + status = payload.get("status") + position = payload.get("queue_position") + return QueueUpdate( + request_id=body_id if isinstance(body_id, str) and body_id else request_id, + status=status if isinstance(status, str) else "", + queue_position=position + if isinstance(position, int) and not isinstance(position, bool) + else None, + error_type=payload.get("error_type") + if isinstance(payload.get("error_type"), str) + else None, + retry_after=_retry_after_seconds(headers), + raw=dict(payload), + ) + + +def _request_id_of(payload: Mapping[str, Any]) -> str: + """The request id a submit response names, or a :class:`ComfyError`. + + A submit whose response carries no usable id is unusable in the specific + way that matters here: the work may well have been accepted and billed, and + the caller has been left with no way to reach it. That is a failure of the + call, so it is raised rather than papered over with a placeholder id that + would 404 on the first poll. + """ + request_id = payload.get("request_id") + if not isinstance(request_id, str) or not request_id: + raise ComfyError( + "the queue accepted the request but its response named no request_id, so the " + "request cannot be polled, collected or cancelled", + code="invalid_response", + ) + try: + # The id becomes a path segment on every later call, so it is held to + # the same rule a caller-supplied one is. Failing here rather than on + # the first poll keeps the failure next to the response that caused it + # — and next to the `Idempotency-Key` that can recover the request. + return parse_request_id(request_id) + except ValueError as exc: + raise ComfyError( + f"the queue named a request_id that cannot address a route: {exc}", + code="invalid_response", + ) from exc + + +def _raise_for_completion( + payload: Mapping[str, Any], *, request_id: str, envelope_only: bool = False +) -> None: + """Raise the typed router exception a completion reports, if it reports one. + + The gate behind "a ``200`` with an error payload is never returned as + success". It runs on the terminal status read *and* on the collected + result, because either can be the response that carries the ``error_type`` + and checking only one leaves the other handing a failure back as data. + + ``envelope_only`` is for the result route, and it is a deliberate + narrowing rather than caution. That route's body is the **provider's own + payload**, forwarded verbatim — a partner model is free to have a field + called ``error_type`` in its native output, and turning one of those into a + raised exception would fail a generation that succeeded. So on that body an + ``error_type`` only counts when it arrives inside the queue's own envelope, + which is what a ``COMPLETED`` status alongside it identifies. The status + read has no such ambiguity and is checked unconditionally, so the failure + the server reports where it is authoritative is never the one that gets + missed. + """ + if envelope_only and payload.get("status") != COMPLETED: + return + error = error_from_completion(payload, request_id=request_id) + if error is not None: + raise error + + +def _changed(previous: QueueUpdate | None, current: QueueUpdate) -> bool: + """Whether ``current`` is worth reporting given ``previous``. + + The first observation always is. After that, only a change in the two + fields a caller renders — the status and the queue position — counts, so a + progress bar is not redrawn once a second for a request that has not moved. + """ + if previous is None: + return True + return (previous.status, previous.queue_position) != (current.status, current.queue_position) + + +def _pace(update: QueueUpdate, backoff: Iterator[float]) -> float: + """Seconds to wait before the next poll. + + A pace the server named beats the schedule guessed here — that is the whole + point of ``Retry-After`` — and the adaptive backoff carries the interval + when it named none. The backoff is advanced either way so the schedule does + not restart from its floor the moment the server stops naming a pace. + """ + scheduled = next(backoff) + return float(update.retry_after) if update.retry_after is not None else scheduled + + +def _remaining(deadline: float | None) -> float | None: + """Seconds left before ``deadline``, or ``None`` when there is no deadline.""" + return None if deadline is None else deadline - _now() + + +def _completed(update: QueueUpdate | None) -> QueueUpdate: + """The terminal update the poll loop ended on. + + ``iter_events`` always yields the completion — the status it carries + differs from every update before it, and the loop only returns once it has + seen one — so ``None`` here is unreachable rather than a state to handle. It + is still checked, because the alternative is a ``None`` dereference in the + middle of collecting a result if that ever stops being true. + """ + if update is None: # pragma: no cover - unreachable; see the docstring + raise ComfyError( + "the poll loop ended without observing a completion", code="invalid_response" + ) + return update + + +def _last(updates: Iterator[QueueUpdate]) -> QueueUpdate: + """Drain ``updates`` and return the final one — the completion.""" + final: QueueUpdate | None = None + for update in updates: + final = update + return _completed(final) + + +def _timed_out(request_id: str, timeout: float | None, update: QueueUpdate) -> TimeoutError: + return TimeoutError( + f"model request {request_id} not complete after {timeout}s (status={update.status!r})" + ) + + +class _RequestHandleBase: + """State and read-only views shared by the sync and async handles.""" + + _model: str + _request_id: str + _retry: RetryPolicy + + @property + def model(self) -> str: + """The canonical ``{provider}/{model}`` id this request was submitted to. + + Part of the handle's identity rather than a convenience: every route in + this family is addressed by the model id *and* the request id, which is + why ``client.models.handle`` takes both. + """ + return self._model + + @property + def request_id(self) -> str: + """The server-minted id for this request — all a rehydration needs.""" + return self._request_id + + def __repr__(self) -> str: + return f"{type(self).__name__}(model={self._model!r}, request_id={self._request_id!r})" + + +class RequestHandle(_RequestHandleBase): + """A queued model request on :class:`~comfy_sdk.client.Comfy`. + + Built by ``client.models.submit`` and by ``client.models.handle``; there is + nothing to construct by hand, and nothing in it that a second process + cannot rebuild from :attr:`model` and :attr:`request_id`. + """ + + def __init__( + self, + low: ComfyLow, + model: str, + request_id: str, + retry: RetryPolicy = DEFAULT_RETRY, + ) -> None: + self._low = low + self._model = model + self._request_id = request_id + self._retry = retry + + # -- polling (authoritative) ------------------------------------------ + def status(self) -> QueueUpdate: + """Poll the queue once and return what it said. + + One request, no waiting — the queued surface's counterpart to + :meth:`comfy_sdk.jobs.Job.refresh`. It reports a completion carrying an + ``error_type`` as data on :attr:`QueueUpdate.error_type` rather than + raising: this is the read a caller uses to *look*, and the raising + belongs to :meth:`get`, which is the one that hands back a result. + """ + with translating(): + payload, headers = self._call( + lambda: self._low.get_model_request_status(self._model, self._request_id) + ) + return _update_from(payload, headers, request_id=self._request_id) + + def iter_events(self, timeout: float | None = None) -> Generator[QueueUpdate, None, None]: + """Poll to completion, yielding an update whenever the queue moves. + + The first observation is always yielded; after that only a change in + status or queue position is. The final yield is the completion itself, + after which the iterator stops — it does *not* raise for a completion + carrying an ``error_type``, because this is a view of the queue's + progress and a caller who wants the result calls :meth:`get`, which + does raise. + + ``timeout`` is a client-side bound in seconds on the whole loop; ``None`` + polls until the server says the request is done. Exceeding it raises + ``TimeoutError`` and leaves the request running — the queue is the + server's, so a local clock running out says nothing about it. Cancelling + on the way out is ``models.subscribe``'s behaviour, deliberately not + this one's: an iterator that cancelled the work it was iterating would + make a ``for`` loop with a ``break`` destructive. + """ + deadline = None if timeout is None else _now() + timeout + backoff = _core.backoff_schedule() + previous: QueueUpdate | None = None + while True: + update = self.status() + if _changed(previous, update): + yield update + previous = update + if update.is_completed: + return + remaining = _remaining(deadline) + if remaining is not None and remaining <= 0: + raise _timed_out(self._request_id, timeout, update) + delay = _pace(update, backoff) + time.sleep(delay if remaining is None else min(delay, remaining)) + + def get(self, timeout: float | None = None) -> dict[str, Any]: + """Wait for the request to complete and return the provider's payload. + + The result is the partner model's own output, decoded from JSON and + handed back as-is — the same value ``models.run`` returns for the same + model and arguments. + + Raises the typed router exception + (:mod:`comfy_sdk.router_exceptions`) when the completion carries an + ``error_type``, which is how the server reports a failed *or* cancelled + request. ``timeout`` bounds the wait exactly as it does on + :meth:`iter_events`, and raises ``TimeoutError`` without cancelling. + + Calling it on a request that has already completed is one status poll + and one fetch, so collecting a result twice — or from a second process + — costs no more than the first time. + """ + return self._collect(_last(self.iter_events(timeout=timeout))) + + def _collect(self, completion: QueueUpdate) -> dict[str, Any]: + """Turn an observed completion into a result, or into the typed error. + + Split out of :meth:`get` so ``models.subscribe`` — which has already + polled its way to the completion — can collect from the update it is + holding instead of spending one more status request re-discovering it. + """ + _raise_for_completion(completion.raw, request_id=self._request_id) + with translating(): + payload, _headers = self._call( + lambda: self._low.get_model_request_result(self._model, self._request_id) + ) + # Checked again on the result body: which of the two responses carries + # the `error_type` is the server's choice, and reading only one of them + # is how a failure gets returned as a result. + _raise_for_completion(payload, request_id=self._request_id, envelope_only=True) + return payload + + def cancel(self) -> QueueUpdate: + """Ask the server to cancel this request. + + A request, not a guarantee — exactly as it is for a workflow job. A + request that has already completed stays completed, so read the returned + update's :attr:`~QueueUpdate.status`, or poll :meth:`status`, rather + than assuming the work stopped. A deployment that answers with no body + gives an update whose ``status`` is ``""``; the authoritative state is + the next :meth:`status`. + """ + with translating(): + payload, headers = self._call( + lambda: self._low.post_model_request_cancel(self._model, self._request_id) + ) + return _update_from(payload, headers, request_id=self._request_id) + + def _call( + self, send: Callable[[], tuple[dict[str, Any], httpx.Headers]] + ) -> tuple[dict[str, Any], httpx.Headers]: + """Run one queue call under the client's retry policy. + + The same ``Retrier`` ``models.run`` uses, constructed per call because + its budget runs from its construction: a poll loop that shared one + would spend the whole budget on its first hour of polling and then + surface the next blip as a hard failure. What it buys here is that a + ``429`` naming a ``Retry-After`` paces the poll instead of ending it. + """ + retrier = Retrier(self._retry, now=_now) + while True: + try: + return send() + except _CANDIDATE_FAILURES as exc: + delay = retrier.delay_before_retry(exc) + if delay is None: + raise + time.sleep(delay) + + +class AsyncRequestHandle(_RequestHandleBase): + """A queued model request on :class:`~comfy_sdk.client.AsyncComfy` — mirrors + :class:`RequestHandle`.""" + + def __init__( + self, + low: AsyncComfyLow, + model: str, + request_id: str, + retry: RetryPolicy = DEFAULT_RETRY, + ) -> None: + self._low = low + self._model = model + self._request_id = request_id + self._retry = retry + + async def status(self) -> QueueUpdate: + """Awaitable :meth:`RequestHandle.status` — one authoritative poll.""" + with translating(): + payload, headers = await self._call( + lambda: self._low.get_model_request_status(self._model, self._request_id) + ) + return _update_from(payload, headers, request_id=self._request_id) + + async def iter_events(self, timeout: float | None = None) -> AsyncGenerator[QueueUpdate, None]: + """Async :meth:`RequestHandle.iter_events` — ``async for`` over the updates.""" + deadline = None if timeout is None else _now() + timeout + backoff = _core.backoff_schedule() + previous: QueueUpdate | None = None + while True: + update = await self.status() + if _changed(previous, update): + yield update + previous = update + if update.is_completed: + return + remaining = _remaining(deadline) + if remaining is not None and remaining <= 0: + raise _timed_out(self._request_id, timeout, update) + delay = _pace(update, backoff) + await asyncio.sleep(delay if remaining is None else min(delay, remaining)) + + async def get(self, timeout: float | None = None) -> dict[str, Any]: + """Async :meth:`RequestHandle.get` — wait, then collect or raise.""" + completion: QueueUpdate | None = None + async for update in self.iter_events(timeout=timeout): + completion = update + return await self._collect(_completed(completion)) + + async def _collect(self, completion: QueueUpdate) -> dict[str, Any]: + """Async :meth:`RequestHandle._collect`.""" + _raise_for_completion(completion.raw, request_id=self._request_id) + with translating(): + payload, _headers = await self._call( + lambda: self._low.get_model_request_result(self._model, self._request_id) + ) + _raise_for_completion(payload, request_id=self._request_id, envelope_only=True) + return payload + + async def cancel(self) -> QueueUpdate: + """Async :meth:`RequestHandle.cancel` — a request, not a guarantee.""" + with translating(): + payload, headers = await self._call( + lambda: self._low.post_model_request_cancel(self._model, self._request_id) + ) + return _update_from(payload, headers, request_id=self._request_id) + + async def _call( + self, send: Callable[[], Awaitable[tuple[dict[str, Any], httpx.Headers]]] + ) -> tuple[dict[str, Any], httpx.Headers]: + """Async :meth:`RequestHandle._call` — one queue call under the retry policy.""" + retrier = Retrier(self._retry, now=_now) + while True: + try: + return await send() + except _CANDIDATE_FAILURES as exc: + delay = retrier.delay_before_retry(exc) + if delay is None: + raise + await asyncio.sleep(delay) + + +__all__ = ["COMPLETED", "AsyncRequestHandle", "QueueUpdate", "RequestHandle"] diff --git a/src/comfy_sdk/models.py b/src/comfy_sdk/models.py index d79b9d9..27b461a 100644 --- a/src/comfy_sdk/models.py +++ b/src/comfy_sdk/models.py @@ -41,18 +41,33 @@ from __future__ import annotations import asyncio +import contextlib import time -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from copy import deepcopy from typing import Any, cast import httpx from comfy_low.errors import ApiError -from comfy_low.transport import MODEL_RUN_TIMEOUT, AsyncComfyLow, ComfyLow +from comfy_low.transport import ( + MODEL_RUN_TIMEOUT, + AsyncComfyLow, + ComfyLow, + parse_model_id, + parse_request_id, +) from ._core import new_idempotency_key, validate_idempotency_key from .exceptions import translating +from .model_requests import ( + _CANCEL_FAILURES, + AsyncRequestHandle, + QueueUpdate, + RequestHandle, + _completed, + _request_id_of, +) from .retry import DEFAULT_RETRY, Retrier, RetryPolicy from .router_exceptions import RouterError @@ -257,6 +272,165 @@ def run( raise time.sleep(delay) + # -- the queued form: submit, hold a handle, collect ------------------ + def submit( + self, + model: str, + arguments: Mapping[str, Any], + *, + idempotency_key: str | None = None, + ) -> RequestHandle: + """Queue ``model`` with ``arguments`` and return a handle to the request. + + The queued counterpart of :meth:`run`, and the same request either way: + ``model`` is the canonical ``{provider}/{model}`` id and ``arguments`` + is the partner model's own native JSON input, forwarded unchanged. The + difference is when the server answers — here, as soon as the request is + *accepted*, with the generation collected later through the returned + :class:`~comfy_sdk.model_requests.RequestHandle`. + + Reach for it over :meth:`run` when the caller cannot hold a connection + for the length of a generation: a web request that has to return now, a + worker that submits in one process and collects in another, or a batch + where the submits should all be in flight at once. + + **A fresh** ``Idempotency-Key`` **is minted per call**, which is what + makes two deliberate submits of the same input two requests rather than + one deduplicated request, while a transport-level retry *inside* this + one call keeps the one key and so replays the original rather than + queueing a second generation. Pass ``idempotency_key`` to choose the + key yourself — the case that earns it is a lost response: the request + may have been accepted and its id lost with the response, and resending + under the same key is the only way back to it. The uniqueness rules are + :meth:`run`'s, unchanged: the keyspace is the whole workspace's, so + mint keys with real entropy and never from a guessable label. + + Every exception this raises carries that key on ``.idempotency_key``, + for exactly that recovery. + + The surface is gated server side: a caller the queue is not switched on + for is answered ``403`` ``not_enabled``, which arrives here as + :class:`~comfy_sdk.router_exceptions.NotEnabled`. Nothing about the + request is wrong in that case, and it is terminal — do not retry it. + """ + low = cast(ComfyLow, self._low) + key = ( + validate_idempotency_key(idempotency_key) + if idempotency_key is not None + else new_idempotency_key() + ) + # Snapshotted deeply before the first attempt, for the reason `run` + # gives: a mutation between attempts would send a different body under + # the one key, which is the same-key-different-body case the contract + # refuses outright. + payload = deepcopy(dict(arguments)) + retrier = Retrier(self._retry, now=_now) + with translating(idempotency_key=key): + while True: + try: + body, _headers = low.post_model_submit(model, payload, idempotency_key=key) + break + except _CANDIDATE_FAILURES as exc: + delay = retrier.delay_before_retry(exc) + if delay is None: + raise + time.sleep(delay) + request_id = _request_id_of(body) + return RequestHandle(low, model, request_id, self._retry) + + def subscribe( + self, + model: str, + arguments: Mapping[str, Any], + *, + on_queue_update: Callable[[QueueUpdate], Any] | None = None, + timeout: float | None = None, + idempotency_key: str | None = None, + ) -> dict[str, Any]: + """Queue a request, follow it to completion, and return its result. + + :meth:`submit` plus polling plus + :meth:`~comfy_sdk.model_requests.RequestHandle.get`, in one call — the + ergonomic form for a caller who does want to wait but also wants to + show progress while waiting. The return value is the provider's own + payload, identical to what :meth:`run` would have returned. + + ``on_queue_update`` is called with a + :class:`~comfy_sdk.model_requests.QueueUpdate` each time the queue + moves — first observation, every change of status or position, and the + completion. It is called from this thread, so keep it quick; an + exception it raises propagates and abandons the wait (the request keeps + running server side). + + ``timeout`` is a **client-side** bound in seconds on the whole wait, + with no server-side meaning: the queue's own timeouts are the server's. + When it runs out this makes a best-effort + :meth:`~comfy_sdk.model_requests.RequestHandle.cancel` — so a caller + that has stopped waiting is not also still paying for a generation + nobody will collect — and then raises ``TimeoutError``. Best-effort is + literal: a cancel that itself fails is swallowed, because the timeout + is the failure worth reporting and a masked one would send the caller + looking in the wrong place. Use :meth:`submit` instead when the request + should outlive the caller's patience. + + A completion carrying an ``error_type`` — which is how the server + reports a failure *and* a cancellation — raises the typed router + exception rather than returning, so a ``200`` never comes back as a + successful result. + """ + handle = self.submit(model, arguments, idempotency_key=idempotency_key) + # ``closing`` so the poll generator is finalised on every exit — the + # completion, the timeout, and above all the one where the caller's + # callback raises, which otherwise leaves it suspended until the + # collector happens to reach it. + with contextlib.closing(handle.iter_events(timeout=timeout)) as updates: + completion: QueueUpdate | None = None + while True: + # Only the *iteration* is inside the ``except TimeoutError``. A + # callback is the caller's own code and may raise a + # ``TimeoutError`` of its own — from an HTTP call it makes to + # render progress, say — and catching that here would cancel a + # perfectly healthy request on the strength of a failure that + # had nothing to do with the wait. + try: + update = next(updates) + except StopIteration: + break + except TimeoutError: + try: + handle.cancel() + except _CANCEL_FAILURES: + # Best-effort is literal: the timeout is the failure + # worth reporting, and a masked one sends the caller + # looking in the wrong place. + pass + raise + completion = update + if on_queue_update is not None: + on_queue_update(update) + return handle._collect(_completed(completion)) + + def handle(self, model: str, request_id: str) -> RequestHandle: + """Rebuild the handle for a request submitted anywhere. + + Takes no state beyond the two ids that address the request, so a + process that never made the submit — a worker draining a queue of ids, + a retry after a restart — reaches the same + :class:`~comfy_sdk.model_requests.RequestHandle` the submitting process + held. Makes no request of its own: an id that names nothing surfaces on + the first :meth:`~comfy_sdk.model_requests.RequestHandle.status` or + :meth:`~comfy_sdk.model_requests.RequestHandle.get`, as the server's + own answer rather than as a guess made here. + + Both ids are validated locally — a malformed ``{provider}/{model}`` id + or a ``request_id`` that is not a single path segment raises + ``ValueError`` (a non-string raises ``TypeError``) rather than being + pasted into a URL. + """ + parse_model_id(model) + parse_request_id(request_id) + return RequestHandle(cast(ComfyLow, self._low), model, request_id, self._retry) + class AsyncModels(_ModelsBase): """``client.models`` on :class:`~comfy_sdk.client.AsyncComfy` — mirrors :class:`Models`.""" @@ -307,3 +481,93 @@ async def run( if delay is None: raise await asyncio.sleep(delay) + + # -- the queued form: submit, hold a handle, collect ------------------ + async def submit( + self, + model: str, + arguments: Mapping[str, Any], + *, + idempotency_key: str | None = None, + ) -> AsyncRequestHandle: + """Awaitable :meth:`Models.submit` — same arguments, an async handle. + + Including the fresh-key-per-call rule and the ``.idempotency_key`` every + exception carries for a lost-response recovery. See :meth:`Models.submit`. + """ + low = cast(AsyncComfyLow, self._low) + key = ( + validate_idempotency_key(idempotency_key) + if idempotency_key is not None + else new_idempotency_key() + ) + payload = deepcopy(dict(arguments)) + retrier = Retrier(self._retry, now=_now) + with translating(idempotency_key=key): + while True: + try: + body, _headers = await low.post_model_submit( + model, payload, idempotency_key=key + ) + break + except _CANDIDATE_FAILURES as exc: + delay = retrier.delay_before_retry(exc) + if delay is None: + raise + await asyncio.sleep(delay) + request_id = _request_id_of(body) + return AsyncRequestHandle(low, model, request_id, self._retry) + + async def subscribe( + self, + model: str, + arguments: Mapping[str, Any], + *, + on_queue_update: Callable[[QueueUpdate], Any] | None = None, + timeout: float | None = None, + idempotency_key: str | None = None, + ) -> dict[str, Any]: + """Awaitable :meth:`Models.subscribe` — same arguments, same result. + + ``on_queue_update`` may be a plain callable or a coroutine function; + an awaitable it returns is awaited before the next poll, so an async + callback does not need wrapping. See :meth:`Models.subscribe`. + """ + handle = await self.submit(model, arguments, idempotency_key=idempotency_key) + # ``aclosing`` for the reason ``Models.subscribe`` uses ``closing``, and + # a sharper one: an async generator left suspended is finalised by the + # event loop's own shutdown hook, well after this call returned. + async with contextlib.aclosing(handle.iter_events(timeout=timeout)) as updates: + completion: QueueUpdate | None = None + while True: + # Scoped to the iteration alone, for the reason + # `Models.subscribe` gives — and it bites harder here, where a + # callback that awaits anything under `asyncio.wait_for` raises + # `TimeoutError` natively. + try: + update = await anext(updates) + except StopAsyncIteration: + break + except TimeoutError: + try: + await handle.cancel() + except _CANCEL_FAILURES: + pass + raise + completion = update + if on_queue_update is not None: + outcome = on_queue_update(update) + if isinstance(outcome, Awaitable): + await outcome + return await handle._collect(_completed(completion)) + + async def handle(self, model: str, request_id: str) -> AsyncRequestHandle: + """Awaitable :meth:`Models.handle` — rebuild a handle from the two ids. + + Awaited for symmetry with the rest of the async client rather than + because it does any I/O; it makes no request, exactly as the sync form + makes none. See :meth:`Models.handle`. + """ + parse_model_id(model) + parse_request_id(request_id) + return AsyncRequestHandle(cast(AsyncComfyLow, self._low), model, request_id, self._retry) diff --git a/src/comfy_sdk/router_exceptions.py b/src/comfy_sdk/router_exceptions.py index 8cae4b1..3bccdd2 100644 --- a/src/comfy_sdk/router_exceptions.py +++ b/src/comfy_sdk/router_exceptions.py @@ -521,6 +521,63 @@ def error_from_response( ) +def error_from_completion( + payload: Any, + *, + request_id: str | None = None, + retry_after: int | None = None, +) -> RouterError | None: + """The typed exception a *completed* queued request reports, or ``None``. + + The queue expresses a failure and a cancellation the same way it expresses + a success: the request reaches ``COMPLETED``, and the failure rides in the + body as an ``error_type``. There is no error *status* to read — the poll + that discovered it was a ``200`` — so a client that only mapped HTTP status + codes would hand a caller a failed generation as a successful result. + + Returns ``None`` when the payload names no ``error_type``, which is the + ordinary success path; every caller has to treat that as "no error found" + rather than as "no error possible". + + ``http_status`` is left unset on what this builds, deliberately: there was + no failing status. That also keeps :mod:`comfy_sdk.retry` out of it — a + completion carrying an ``error_type`` is the server's final answer about a + request that already ran, not a transport condition another attempt could + survive. + + Like :func:`error_from_response`, this never raises: a malformed body + degrades to the least specific exception the payload still supports. + """ + if not isinstance(payload, Mapping): + return None + error_type = _clean(payload.get("error_type")) + if error_type is None: + return None + + errors: tuple[ValidationErrorDetail, ...] = () + detail: str | None = None + raw_detail = payload.get("detail") + if isinstance(raw_detail, str): + detail = raw_detail or None + elif isinstance(raw_detail, Sequence) and not isinstance(raw_detail, (str, bytes)): + errors = tuple(_detail_from(entry) for entry in raw_detail if isinstance(entry, Mapping)) + + if detail is None: + detail = _summarise(errors) or f"the request completed with error_type {error_type!r}" + + return exception_for(error_type)( + detail, + error_type=error_type, + # Filtered, not merely stripped, and by the same function + # `error_from_response` uses: a completion's id is just as + # server-controlled as a header's, and it lands in the same displayed, + # pasted-into-a-support-ticket place. + request_id=clean_request_id(request_id), + retry_after=retry_after, + errors=errors, + ) + + def _lowercase_headers(headers: Mapping[str, str] | None) -> dict[str, str]: """Header names are case-insensitive on the wire; normalise so a plain dict works as well as an ``httpx.Headers``.""" @@ -617,6 +674,7 @@ def _summarise(errors: Sequence[ValidationErrorDetail]) -> str: "ServiceUnavailable", "Unauthorized", "ValidationErrorDetail", + "error_from_completion", "error_from_response", "exception_for", ] diff --git a/tests/conftest.py b/tests/conftest.py index 799351a..5e48478 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -8,8 +8,10 @@ Both, because the SDK speaks to two surfaces: the ``/api/v2`` deployment (jobs, assets) and Comfy Router (``/v2/models/{provider}/{model}``), which is a -different host in production. This one stub answers both route families, so a -test that exercises either gets a single server — while a test that is *about* +different host in production. This one stub answers three route families — the +``/api/v2`` paths, the awaited model run, and the queued model routes under +``.../requests`` — so a test that exercises any of them gets a single server +— while a test that is *about* the two being separate points ``COMFY_ROUTER_BASE_URL`` at ``second_server``. """ @@ -165,6 +167,67 @@ class ServerState: # in, and the bucket-keyed collect rule has to read it. model_run_router_error_shape: bool = False + # --- the queued model surface (submit / status / result / cancel) --- + # POST .../requests answers this status with a body naming a request id. + queue_submit_status: int = 200 + # (status, code) answered instead of accepting the submit — permanent. + queue_submit_error: tuple[int, str] | None = None + # Submits that fail transiently before one is accepted, and the (status, + # code) each answers with. Checked *before* `queue_submit_error`, exactly as + # `model_run_fail_times` is checked before `model_run_error`. + queue_submit_fail_times: int = 0 + queue_submit_transient_error: tuple[int, str] = (429, "rate_limited") + queue_submit_transient_retry_after: str | None = "1" + # Answer the submit with a body carrying no `request_id` at all — accepted + # work the caller has no way to reach. + queue_submit_omits_request_id: bool = False + # The id the queue hands back, and the one every later route answers for. + queue_request_id: str = "req_stub_01" + # Status polls that report a non-terminal state before the request reaches + # COMPLETED. 0 means the very first poll is already complete. + queue_polls_to_complete: int = 2 + # Non-terminal status reported by those polls, and the queue position they + # report (decremented per poll, floored at 0). + queue_pending_status: str = "IN_QUEUE" + queue_start_position: int = 2 + # Sent as Retry-After on every *successful* status poll — the server naming + # its own poll pace, which the SDK honours over its local backoff. + queue_status_retry_after: str | None = None + # (status, code) answered by every status poll instead of the queue state + # — the permanent-failure knob. + queue_status_error: tuple[int, str] | None = None + # Status polls that fail transiently before answering normally, and the + # (status, code) each of those answers with. Checked *before* + # `queue_status_error`, exactly as `model_run_fail_times` is. + queue_status_fail_times: int = 0 + queue_status_transient_error: tuple[int, str] = (429, "rate_limited") + # Retry-After sent alongside a transient status failure. + queue_status_transient_retry_after: str | None = "1" + # `error_type` carried by the COMPLETED status — how the server reports a + # failed or cancelled request. `None` is the ordinary success path. + queue_error_type: str | None = None + queue_error_detail: str | None = "the model refused the request" + # `error_type` carried by the *result* body only, with the status reporting + # a clean completion — the other half of "a 200 is not a success". + queue_result_error_type: str | None = None + # Extra keys merged into the served result payload — for the case where the + # provider's OWN native output happens to carry a field the queue envelope + # also uses (`error_type`), with no queue envelope around it. + queue_result_extra: dict[str, Any] = field(default_factory=dict) + # The provider's native payload served by GET .../requests/{id}. + queue_result: dict[str, Any] = field( + default_factory=lambda: { + "images": [{"url": "http://example.invalid/queued.png"}], + "seed": 7, + } + ) + # Status code for the cancel response; 204 exercises the empty-body path. + queue_cancel_status: int = 200 + # The bucket a cancelled request's completion carries. + queue_cancel_error_type: str = "client_disconnected" + # Set by a cancel; makes every later status poll report the cancellation. + queue_canceled: bool = False + # --- counters the tests assert on --- upload_count: int = 0 from_hash_count: int = 0 @@ -216,6 +279,24 @@ class ServerState: # and does not increment this, which is what lets a test tell a real replay # apart from a second generation that merely returns an equal payload. model_run_generations: int = 0 + queue_submit_count: int = 0 + queue_status_count: int = 0 + # Status polls that were actually *answered with a queue state*, as + # distinct from polls that arrived (`queue_status_count`). A poll answered + # with a transient failure must not advance the request towards completion, + # or a retry test would silently shorten the queue it is testing. + queue_status_served: int = 0 + queue_result_count: int = 0 + queue_cancel_count: int = 0 + # Every Idempotency-Key seen on a queue submit, in arrival order. + queue_submit_idempotency_keys: list[str | None] = field(default_factory=list) + # The native body of the last queue submit, and the two decoded id segments. + last_queue_submit_body: dict[str, Any] | None = None + last_queue_provider: str | None = None + last_queue_model: str | None = None + # Every raw path the queue routes answered, in order — for the tests that + # are about the routes themselves rather than about what came back. + queue_paths: list[str] = field(default_factory=list) def _asset_json(asset_id: str, hash_: str, created_new: bool, size: int) -> dict: @@ -375,6 +456,18 @@ def do_GET(self) -> None: return self._json(200, _asset_json(m.group(1), state.server_hash, False, 33)) return + # Comfy Router's queued model routes — the status poll and the + # result collection. Matched before the two-segment run route + # patterns for the same reason they are anchored: a request id is + # a path segment, not a model name. + m = re.match(r"/v2/models/([^/]+)/([^/]+)/requests/([^/]+)/status$", self.path) + if m: + self._serve_queue_status(m.group(3)) + return + m = re.match(r"/v2/models/([^/]+)/([^/]+)/requests/([^/]+)$", self.path) + if m: + self._serve_queue_result(m.group(3)) + return m = re.match(r"/api/v2/jobs/([^/]+)/events$", self.path) if m: self._serve_events(m.group(1)) @@ -503,6 +596,15 @@ def do_POST(self) -> None: if m: self._post_model_run(m.group(1), m.group(2)) return + m = re.match(r"/v2/models/([^/]+)/([^/]+)/requests$", self.path) + if m: + self._post_queue_submit(m.group(1), m.group(2)) + return + m = re.match(r"/v2/models/([^/]+)/([^/]+)/requests/([^/]+)/cancel$", self.path) + if m: + self._read_body() + self._post_queue_cancel(m.group(3)) + return m = re.match(r"/api/v2/jobs/([^/]+)/cancel$", self.path) if m: self._json(200, _job_json(m.group(1), "canceling")) @@ -527,6 +629,109 @@ def _post_from_hash(self) -> None: else: self._err(404, "blob_not_found", "no such blob") + # -- the queued model surface -- + def _post_queue_submit(self, provider: str, model: str) -> None: + state.queue_submit_count += 1 + state.queue_paths.append(self.path) + state.last_queue_provider = unquote(provider) + state.last_queue_model = unquote(model) + state.last_queue_submit_body = json.loads(self._read_body() or b"{}") + state.queue_submit_idempotency_keys.append(self.headers.get("Idempotency-Key")) + if state.queue_submit_fail_times > 0: + state.queue_submit_fail_times -= 1 + status, code = state.queue_submit_transient_error + self._router_err(status, code, retry_after=state.queue_submit_transient_retry_after) + return + if state.queue_submit_error: + status, code = state.queue_submit_error + self._router_err(status, code) + return + body: dict[str, Any] = {"status": state.queue_pending_status} + if not state.queue_submit_omits_request_id: + body["request_id"] = state.queue_request_id + self._json(state.queue_submit_status, body) + + def _serve_queue_status(self, request_id: str) -> None: + state.queue_status_count += 1 + state.queue_paths.append(self.path) + if state.queue_status_fail_times > 0: + state.queue_status_fail_times -= 1 + status, code = state.queue_status_transient_error + self._router_err(status, code, retry_after=state.queue_status_transient_retry_after) + return + if state.queue_status_error: + status, code = state.queue_status_error + self._router_err(status, code) + return + headers = ( + {"Retry-After": state.queue_status_retry_after} + if state.queue_status_retry_after + else {} + ) + body: dict[str, Any] = {"request_id": unquote(request_id)} + if state.queue_canceled: + body["status"] = "COMPLETED" + body["error_type"] = state.queue_cancel_error_type + body["detail"] = "the request was cancelled" + self._json(200, body, headers=headers) + return + served = state.queue_status_served + state.queue_status_served += 1 + if served < state.queue_polls_to_complete: + body["status"] = state.queue_pending_status + body["queue_position"] = max(state.queue_start_position - served, 0) + self._json(200, body, headers=headers) + return + body["status"] = "COMPLETED" + if state.queue_error_type: + body["error_type"] = state.queue_error_type + if state.queue_error_detail is not None: + body["detail"] = state.queue_error_detail + self._json(200, body, headers=headers) + + def _serve_queue_result(self, request_id: str) -> None: + state.queue_result_count += 1 + state.queue_paths.append(self.path) + if state.queue_result_error_type: + self._json( + 200, + { + "request_id": unquote(request_id), + "status": "COMPLETED", + "error_type": state.queue_result_error_type, + "detail": "the result body carried the failure", + }, + ) + return + self._json(200, {**state.queue_result, **state.queue_result_extra}) + + def _post_queue_cancel(self, request_id: str) -> None: + state.queue_cancel_count += 1 + state.queue_paths.append(self.path) + state.queue_canceled = True + if state.queue_cancel_status == 204: + self.send_response(204) + self.send_header("Content-Length", "0") + self.end_headers() + return + self._json( + state.queue_cancel_status, + { + "request_id": unquote(request_id), + "status": "COMPLETED", + "error_type": state.queue_cancel_error_type, + }, + ) + + def _router_err( + self, status: int, code: str, message: str = "err", retry_after: str | None = None + ) -> None: + """Router's own error shape: the bucket on the header and in the body.""" + headers = {"X-Comfy-Error-Type": code} + if retry_after: + headers["Retry-After"] = retry_after + self._json(status, {"detail": message, "error_type": code}, headers=headers) + def _post_model_run(self, provider: str, model: str) -> None: state.model_run_count += 1 # Decoded, because the SDK percent-encodes each segment and a real diff --git a/tests/test_models_queue.py b/tests/test_models_queue.py new file mode 100644 index 0000000..4967e61 --- /dev/null +++ b/tests/test_models_queue.py @@ -0,0 +1,687 @@ +"""The queued model surface: ``models.submit`` / ``subscribe`` / ``handle``. + +Driven against the recorded stub in ``conftest.py``, never a live deployment — +``server.state`` is set to the queue scenario under test and the SDK is pointed +at the stub by the fixture. + +What each group here is for: + +* the **routes** the four operations address, and the encoding of the ids in + them, since a queued request is addressed by model id *and* request id; +* **poll-authoritative** completion — adaptive backoff, and the server's own + ``Retry-After`` beating it when it names one; +* the **completion-is-not-a-success** rule: a ``COMPLETED`` carrying an + ``error_type`` raises the typed router exception from every path that hands + back a result, whether the bucket arrives on the status or on the result; +* the **Idempotency-Key** contract — one fresh key per ``submit`` call; +* ``subscribe``'s client-side timeout, which cancels before it raises; +* and that ``models.run`` is untouched by all of it. +""" + +from __future__ import annotations + +import asyncio +import time +from typing import Any + +import pytest + +from comfy_low.transport import ( + _MODEL_REQUEST_CANCEL_PATH_TEMPLATE, + _MODEL_REQUEST_PATH_TEMPLATE, + _MODEL_REQUEST_STATUS_PATH_TEMPLATE, + _MODEL_REQUESTS_PATH_TEMPLATE, + _MODEL_RUN_PATH_TEMPLATE, +) +from comfy_sdk import AsyncComfy, Comfy, QueueUpdate +from comfy_sdk.exceptions import ComfyError +from comfy_sdk.model_requests import COMPLETED, AsyncRequestHandle, RequestHandle +from comfy_sdk.retry import NO_RETRY +from comfy_sdk.router_exceptions import ( + ContentPolicyViolation, + NotEnabled, + RouterError, + error_from_completion, +) + +MODEL = "acme/fast-sdxl" +ARGS = {"prompt": "a red bicycle"} + + +@pytest.fixture +def fast_poll(monkeypatch): + """Collapse the poll backoff so a multi-poll test is not a multi-second one. + + Patched on ``comfy_sdk._core`` — the one place the schedule is defined — + rather than on the handles, so both the sync and the async loop pick it up + and neither can drift onto a second schedule. + """ + import comfy_sdk._core as core + + monkeypatch.setattr(core, "backoff_schedule", lambda *a, **k: iter(lambda: 0.0, None)) + return None + + +def _client(**kw: Any) -> Comfy: + return Comfy(api_key="comfyui-test-key", **kw) + + +# --- routes ---------------------------------------------------------------- + + +def test_submit_posts_to_the_requests_route_under_the_model_id(server, fast_poll) -> None: + with _client() as client: + handle = client.models.submit(MODEL, ARGS) + + assert isinstance(handle, RequestHandle) + assert handle.request_id == server.state.queue_request_id + assert handle.model == MODEL + assert server.state.queue_paths == ["/v2/models/acme/fast-sdxl/requests"] + # The partner model's native input, with no Comfy-shaped envelope — the + # same body `models.run` sends. + assert server.state.last_queue_submit_body == ARGS + assert (server.state.last_queue_provider, server.state.last_queue_model) == ( + "acme", + "fast-sdxl", + ) + + +def test_every_queue_route_is_addressed_by_both_ids(server, fast_poll) -> None: + server.state.queue_polls_to_complete = 0 + with _client() as client: + handle = client.models.submit(MODEL, ARGS) + handle.status() + handle.get() + handle.cancel() + + prefix = "/v2/models/acme/fast-sdxl/requests" + rid = server.state.queue_request_id + assert server.state.queue_paths == [ + prefix, + f"{prefix}/{rid}/status", + f"{prefix}/{rid}/status", + f"{prefix}/{rid}", + f"{prefix}/{rid}/cancel", + ] + + +def test_the_ids_are_percent_encoded_into_the_path(server, fast_poll) -> None: + server.state.queue_request_id = "req id?x=1" + with _client() as client: + handle = client.models.submit("acme/model name", ARGS) + handle.status() + + assert server.state.queue_paths[0] == "/v2/models/acme/model%20name/requests" + # Nothing in either id can add a segment, a query or a fragment. + assert server.state.queue_paths[1].endswith("/requests/req%20id%3Fx%3D1/status") + + +def test_a_server_named_request_id_that_would_walk_the_path_is_refused(server) -> None: + """A hostile or broken id fails at the submit, not three calls later.""" + server.state.queue_request_id = "../escape" + with _client() as client: + with pytest.raises(ComfyError) as excinfo: + client.models.submit(MODEL, ARGS) + + assert "cannot address a route" in str(excinfo.value) + assert excinfo.value.idempotency_key, "the accepted request stays recoverable" + + +@pytest.mark.parametrize( + "model,expected", + [("one-segment", ValueError), ("a/b/c", ValueError), (object(), TypeError)], +) +def test_submit_rejects_a_malformed_model_id_before_any_request(server, model, expected) -> None: + with _client() as client: + with pytest.raises(expected): + client.models.submit(model, ARGS) + assert server.state.queue_submit_count == 0 + + +@pytest.mark.parametrize( + "request_id,expected", + [("", ValueError), ("a/b", ValueError), ("..", ValueError), (object(), TypeError)], +) +def test_handle_rejects_a_malformed_request_id_before_any_request( + server, request_id, expected +) -> None: + with _client() as client: + with pytest.raises(expected): + client.models.handle(MODEL, request_id) + assert server.state.queue_status_count == 0 + + +def test_handle_rehydrates_from_the_two_ids_without_a_request(server, fast_poll) -> None: + """The other-process case: no submit here, only the ids.""" + with _client() as client: + handle = client.models.handle(MODEL, "req_from_elsewhere") + assert (handle.model, handle.request_id) == (MODEL, "req_from_elsewhere") + # Constructing it made no call at all — the first one is the poll. + assert server.state.queue_status_count == 0 + update = handle.status() + + assert server.state.queue_status_count == 1 + assert update.request_id == "req_from_elsewhere" + + +def test_a_submit_whose_response_names_no_request_id_is_an_error(server) -> None: + server.state.queue_submit_omits_request_id = True + with _client() as client: + with pytest.raises(ComfyError) as excinfo: + client.models.submit(MODEL, ARGS) + assert "request_id" in str(excinfo.value) + # The key is still reachable, because the work may have been accepted. + assert excinfo.value.idempotency_key + + +# --- poll-authoritative completion ------------------------------------------ + + +def test_get_polls_to_completion_then_collects_the_result(server, fast_poll) -> None: + server.state.queue_polls_to_complete = 3 + with _client() as client: + result = client.models.submit(MODEL, ARGS).get() + + assert result == server.state.queue_result + # Three pending polls plus the completing one; the result is fetched once. + assert server.state.queue_status_count == 4 + assert server.state.queue_result_count == 1 + + +def test_iter_events_yields_the_first_state_every_change_and_the_completion( + server, fast_poll +) -> None: + server.state.queue_polls_to_complete = 3 + server.state.queue_start_position = 2 + with _client() as client: + updates = list(client.models.submit(MODEL, ARGS).iter_events()) + + assert [u.status for u in updates] == ["IN_QUEUE", "IN_QUEUE", "IN_QUEUE", COMPLETED] + # Positions 2, 1, 0 — the server's numbers, never computed locally. The + # third pending poll repeats position 0 and is therefore not re-reported. + assert [u.queue_position for u in updates] == [2, 1, 0, None] + assert updates[-1].is_completed + + +def test_an_unchanged_poll_is_not_re_reported(server, fast_poll) -> None: + """A queue that has not moved must not redraw the caller's progress bar.""" + server.state.queue_polls_to_complete = 4 + server.state.queue_start_position = 0 # every pending poll reports position 0 + with _client() as client: + updates = list(client.models.submit(MODEL, ARGS).iter_events()) + + assert [u.status for u in updates] == ["IN_QUEUE", COMPLETED] + assert server.state.queue_status_count == 5 + + +def test_an_unknown_status_is_not_treated_as_terminal(server, fast_poll) -> None: + """A status this version has never heard of keeps polling rather than + collecting a result that does not exist yet.""" + server.state.queue_pending_status = "SOME_FUTURE_STATE" + server.state.queue_polls_to_complete = 2 + with _client() as client: + result = client.models.submit(MODEL, ARGS).get() + + assert result == server.state.queue_result + assert server.state.queue_status_count == 3 + + +def test_a_server_named_retry_after_paces_the_poll(server, monkeypatch) -> None: + """``Retry-After`` on a *successful* poll beats the local backoff.""" + slept: list[float] = [] + monkeypatch.setattr("comfy_sdk.model_requests.time.sleep", slept.append) + server.state.queue_polls_to_complete = 2 + server.state.queue_status_retry_after = "7" + + with _client() as client: + client.models.submit(MODEL, ARGS).get() + + # One sleep per pending poll, each at the server's pace rather than the + # 0.5s the adaptive schedule would have started from. + assert slept == [7.0, 7.0] + + +def test_a_useless_retry_after_falls_back_to_the_backoff(server, monkeypatch) -> None: + """``Retry-After: 0`` names no pace and must not become a zero-delay loop.""" + slept: list[float] = [] + monkeypatch.setattr("comfy_sdk.model_requests.time.sleep", slept.append) + server.state.queue_polls_to_complete = 2 + server.state.queue_status_retry_after = "0" + + with _client() as client: + client.models.submit(MODEL, ARGS).get() + + assert slept and all(delay > 0 for delay in slept) + + +def test_a_throttled_poll_is_retried_under_the_client_policy( + server, fast_poll, monkeypatch +) -> None: + monkeypatch.setattr("comfy_sdk.model_requests.time.sleep", lambda _s: None) + server.state.queue_polls_to_complete = 0 + server.state.queue_status_fail_times = 2 # 429 rate_limited, twice + + with _client() as client: + result = client.models.submit(MODEL, ARGS).get() + + assert result == server.state.queue_result + # Two throttled polls plus the one that answered — the throttled ones did + # not advance the queue, and did not surface to the caller. + assert server.state.queue_status_count == 3 + assert server.state.queue_status_served == 1 + + +def test_a_throttled_poll_raises_when_the_client_does_not_retry(server, fast_poll) -> None: + server.state.queue_polls_to_complete = 0 + server.state.queue_status_fail_times = 1 + + with _client(retry=NO_RETRY) as client: + handle = client.models.submit(MODEL, ARGS) + with pytest.raises(RouterError) as excinfo: + handle.status() + assert excinfo.value.error_type == "rate_limited" + + +# --- a completion is not a success ------------------------------------------ + + +def test_a_completed_status_carrying_an_error_type_raises_the_typed_error( + server, fast_poll +) -> None: + server.state.queue_polls_to_complete = 1 + server.state.queue_error_type = "content_policy_violation" + server.state.queue_error_detail = "the prompt was refused" + + with _client() as client: + with pytest.raises(ContentPolicyViolation) as excinfo: + client.models.submit(MODEL, ARGS).get() + + assert excinfo.value.error_type == "content_policy_violation" + assert excinfo.value.detail == "the prompt was refused" + assert excinfo.value.request_id == server.state.queue_request_id + # The result was never collected: the failure was already known. + assert server.state.queue_result_count == 0 + + +def test_an_error_type_on_the_result_body_alone_still_raises(server, fast_poll) -> None: + """The other half of the rule — whichever response carries the bucket.""" + server.state.queue_polls_to_complete = 0 + server.state.queue_result_error_type = "provider_error" + + with _client() as client: + with pytest.raises(RouterError) as excinfo: + client.models.submit(MODEL, ARGS).get() + + assert excinfo.value.error_type == "provider_error" + assert server.state.queue_result_count == 1 + + +def test_a_providers_own_error_type_field_is_not_mistaken_for_a_failure(server, fast_poll) -> None: + """The result body is the provider's native output, forwarded verbatim. + + A partner model is free to have a field called ``error_type`` in its own + schema. Raising on one would fail a generation that succeeded, so on the + result route the bucket only counts inside the queue's own envelope — which + the ``COMPLETED`` status alongside it is what identifies. + """ + server.state.queue_polls_to_complete = 0 + server.state.queue_result_extra = {"error_type": "provider_error"} + + with _client() as client: + result = client.models.submit(MODEL, ARGS).get() + + assert result["error_type"] == "provider_error" + assert result["seed"] == server.state.queue_result["seed"] + + +def test_an_unknown_error_type_still_raises_the_base_router_error(server, fast_poll) -> None: + server.state.queue_polls_to_complete = 0 + server.state.queue_error_type = "a_bucket_from_the_future" + + with _client() as client: + with pytest.raises(RouterError) as excinfo: + client.models.submit(MODEL, ARGS).get() + + assert type(excinfo.value) is RouterError + assert excinfo.value.error_type == "a_bucket_from_the_future" + + +def test_iter_events_reports_the_failure_as_data_rather_than_raising(server, fast_poll) -> None: + """The split that makes ``iter_events`` a view and ``get`` the collector.""" + server.state.queue_polls_to_complete = 0 + server.state.queue_error_type = "provider_error" + + with _client() as client: + handle = client.models.submit(MODEL, ARGS) + updates = list(handle.iter_events()) + assert updates[-1].error_type == "provider_error" + with pytest.raises(RouterError): + handle.get() + + +def test_a_clean_completion_reports_no_error() -> None: + assert error_from_completion({"status": COMPLETED, "request_id": "r"}) is None + assert error_from_completion(None) is None + assert error_from_completion({"error_type": " "}) is None + + +def test_a_request_the_server_does_not_know_surfaces_as_the_servers_own_answer( + server, fast_poll +) -> None: + """``handle`` makes no call, so an id that names nothing fails on the poll.""" + server.state.queue_status_error = (404, "model_not_found") + + with _client() as client: + handle = client.models.handle(MODEL, "req_that_never_existed") + with pytest.raises(RouterError) as excinfo: + handle.status() + + assert excinfo.value.http_status == 404 + assert excinfo.value.error_type == "model_not_found" + + +def test_an_unflagged_caller_gets_the_servers_not_enabled_as_the_typed_error(server) -> None: + """The surface is gated server side; the SDK's job is to type the refusal.""" + server.state.queue_submit_error = (403, "not_enabled") + + with _client() as client: + with pytest.raises(NotEnabled) as excinfo: + client.models.submit(MODEL, ARGS) + + assert excinfo.value.error_type == "not_enabled" + assert excinfo.value.http_status == 403 + + +# --- the Idempotency-Key contract ------------------------------------------- + + +def test_each_submit_call_mints_a_fresh_key(server, fast_poll) -> None: + with _client() as client: + client.models.submit(MODEL, ARGS) + client.models.submit(MODEL, ARGS) + + keys = server.state.queue_submit_idempotency_keys + assert len(keys) == 2 + assert all(keys) + assert keys[0] != keys[1], "two deliberate submits must be two requests, not one" + + +def test_a_transport_retry_of_one_submit_keeps_the_one_key(server, monkeypatch) -> None: + """The other half of the rule: one *call* is one key, however many attempts. + + A fresh key per attempt would bill a retried submit as a second queued + generation, which is the failure the one-key rule exists to prevent. + """ + monkeypatch.setattr("comfy_sdk.models.time.sleep", lambda _s: None) + server.state.queue_submit_fail_times = 2 # two 429s, then accepted + + with _client() as client: + client.models.submit(MODEL, ARGS) + + keys = server.state.queue_submit_idempotency_keys + assert len(keys) == 3, "three attempts should have reached the server" + assert len(set(keys)) == 1, "every attempt of one call carries one key" + + +def test_an_explicit_key_is_used_verbatim(server, fast_poll) -> None: + with _client() as client: + client.models.submit(MODEL, ARGS, idempotency_key="my-own-key-01") + assert server.state.queue_submit_idempotency_keys == ["my-own-key-01"] + + +def test_an_unusable_explicit_key_is_refused_locally(server) -> None: + with _client() as client: + with pytest.raises(ValueError): + client.models.submit(MODEL, ARGS, idempotency_key="") + assert server.state.queue_submit_count == 0 + + +def test_the_key_rides_out_on_a_failure(server) -> None: + server.state.queue_submit_error = (400, "invalid_input") + with _client(retry=NO_RETRY) as client: + with pytest.raises(RouterError) as excinfo: + client.models.submit(MODEL, ARGS, idempotency_key="key-for-recovery") + assert excinfo.value.idempotency_key == "key-for-recovery" + + +# --- cancel and subscribe --------------------------------------------------- + + +def test_cancel_asks_the_server_and_reports_what_it_said(server, fast_poll) -> None: + with _client() as client: + handle = client.models.submit(MODEL, ARGS) + update = handle.cancel() + + assert server.state.queue_cancel_count == 1 + assert update.request_id == server.state.queue_request_id + assert update.error_type == server.state.queue_cancel_error_type + + +def test_a_cancel_answered_with_no_body_still_identifies_the_request(server, fast_poll) -> None: + server.state.queue_cancel_status = 204 + with _client() as client: + update = client.models.submit(MODEL, ARGS).cancel() + + assert update.request_id == server.state.queue_request_id + assert update.status == "" + + +def test_subscribe_reports_progress_and_returns_the_result(server, fast_poll) -> None: + server.state.queue_polls_to_complete = 2 + seen: list[QueueUpdate] = [] + + with _client() as client: + result = client.models.subscribe(MODEL, ARGS, on_queue_update=seen.append) + + assert result == server.state.queue_result + assert [u.status for u in seen] == ["IN_QUEUE", "IN_QUEUE", COMPLETED] + + +def test_subscribe_needs_no_callback(server, fast_poll) -> None: + server.state.queue_polls_to_complete = 1 + with _client() as client: + assert client.models.subscribe(MODEL, ARGS) == server.state.queue_result + + +def test_subscribe_raises_the_typed_error_for_a_failed_completion(server, fast_poll) -> None: + server.state.queue_polls_to_complete = 0 + server.state.queue_error_type = "content_policy_violation" + with _client() as client: + with pytest.raises(ContentPolicyViolation): + client.models.subscribe(MODEL, ARGS) + + +def test_subscribes_timeout_cancels_before_it_raises(server, monkeypatch) -> None: + """Acceptance: a caller who has stopped waiting is not still paying.""" + monkeypatch.setattr("comfy_sdk.model_requests.time.sleep", lambda _s: None) + # Never completes on its own. + server.state.queue_polls_to_complete = 10_000 + + with _client() as client: + with pytest.raises(TimeoutError): + client.models.subscribe(MODEL, ARGS, timeout=0.0) + + assert server.state.queue_cancel_count == 1 + + +def test_a_failing_cancel_does_not_mask_the_timeout(server, monkeypatch) -> None: + """Best-effort is literal: the timeout is the failure worth reporting.""" + monkeypatch.setattr("comfy_sdk.model_requests.time.sleep", lambda _s: None) + server.state.queue_polls_to_complete = 10_000 + + def _explode(self: Any) -> None: + raise ComfyError("cancel is unreachable") + + monkeypatch.setattr(RequestHandle, "cancel", _explode) + + with _client() as client: + with pytest.raises(TimeoutError): + client.models.subscribe(MODEL, ARGS, timeout=0.0) + + +def test_a_callbacks_own_timeout_error_does_not_cancel_the_request(server, monkeypatch) -> None: + """The caller's callback is the caller's code, and its failures are its own. + + A progress callback that makes its own HTTP call can raise ``TimeoutError`` + for reasons that have nothing to do with this wait; cancelling a healthy + queued request on the strength of it would be a charge thrown away. + """ + monkeypatch.setattr("comfy_sdk.model_requests.time.sleep", lambda _s: None) + server.state.queue_polls_to_complete = 3 + + def _explodes(_update: QueueUpdate) -> None: + raise TimeoutError("the callback's own HTTP call timed out") + + with _client() as client: + with pytest.raises(TimeoutError, match="callback"): + client.models.subscribe(MODEL, ARGS, on_queue_update=_explodes) + + assert server.state.queue_cancel_count == 0 + + +def test_subscribe_collects_without_a_second_status_poll(server, fast_poll) -> None: + """It has already polled its way to the completion; re-discovering it is + one request spent on something it is holding.""" + server.state.queue_polls_to_complete = 2 + + with _client() as client: + client.models.subscribe(MODEL, ARGS) + + # Two pending polls plus the completing one — and no fourth. + assert server.state.queue_status_count == 3 + assert server.state.queue_result_count == 1 + + +def test_iter_events_does_not_cancel_on_its_own_timeout(server, monkeypatch) -> None: + """A ``for`` loop over the queue must not be destructive.""" + monkeypatch.setattr("comfy_sdk.model_requests.time.sleep", lambda _s: None) + server.state.queue_polls_to_complete = 10_000 + + with _client() as client: + handle = client.models.submit(MODEL, ARGS) + with pytest.raises(TimeoutError): + list(handle.iter_events(timeout=0.0)) + + assert server.state.queue_cancel_count == 0 + + +def test_a_timeout_never_sleeps_past_the_deadline(server, monkeypatch) -> None: + """A long server-named pace must not outlive the caller's own bound.""" + slept: list[float] = [] + monkeypatch.setattr("comfy_sdk.model_requests.time.sleep", slept.append) + server.state.queue_polls_to_complete = 10_000 + server.state.queue_status_retry_after = "600" + + with _client() as client: + handle = client.models.submit(MODEL, ARGS) + started = time.monotonic() + with pytest.raises(TimeoutError): + list(handle.iter_events(timeout=0.25)) + + assert slept, "the loop should have paced at least one poll" + assert max(slept) <= 0.25, f"slept past the caller's deadline: {slept}" + assert time.monotonic() - started < 5 + + +# --- the routes as constants ------------------------------------------------ + + +def test_the_queue_routes_extend_the_run_route() -> None: + """The four queue routes hang off the model-run path, in one place. + + The Router contract that declares them is authored but held, so unlike + ``_MODEL_RUN_PATH_TEMPLATE`` there is no vendored spec to pin them against + (``tests/test_router_spec_contract.py`` does that for the run path). What + can be pinned is the relationship: they are the same model-ID-addressed + prefix plus a ``requests`` collection, and each fills exactly the segments + it declares. When the operations land in ``spec/router-openapi.yaml`` this + is the assertion to replace with a comparison against the file. + """ + assert _MODEL_REQUESTS_PATH_TEMPLATE == _MODEL_RUN_PATH_TEMPLATE + "/requests" + assert _MODEL_REQUEST_PATH_TEMPLATE == _MODEL_REQUESTS_PATH_TEMPLATE + "/{request_id}" + for template in ( + _MODEL_REQUEST_PATH_TEMPLATE, + _MODEL_REQUEST_STATUS_PATH_TEMPLATE, + _MODEL_REQUEST_CANCEL_PATH_TEMPLATE, + ): + assert template.count("{") == 3 + assert all(part in template for part in ("{provider}", "{model}", "{request_id}")) + + +# --- models.run is untouched ------------------------------------------------ + + +def test_run_still_posts_to_its_own_route_and_returns_the_payload(server) -> None: + """Acceptance: the queued surface must not have moved ``run`` an inch.""" + with _client() as client: + result = client.models.run(MODEL, ARGS) + + assert result == server.state.model_run_result + assert server.state.last_model_run_path == "/v2/models/acme/fast-sdxl" + assert server.state.queue_submit_count == 0 + assert server.state.queue_status_count == 0 + + +# --- the async client ------------------------------------------------------- + + +async def test_async_submit_and_get(server, fast_poll) -> None: + server.state.queue_polls_to_complete = 2 + async with AsyncComfy(api_key="comfyui-test-key") as client: + handle = await client.models.submit(MODEL, ARGS) + assert isinstance(handle, AsyncRequestHandle) + assert await handle.get() == server.state.queue_result + + +async def test_async_iter_events(server, fast_poll) -> None: + server.state.queue_polls_to_complete = 2 + async with AsyncComfy(api_key="comfyui-test-key") as client: + handle = await client.models.submit(MODEL, ARGS) + statuses = [update.status async for update in handle.iter_events()] + assert statuses == ["IN_QUEUE", "IN_QUEUE", COMPLETED] + + +async def test_async_handle_rehydrates_without_a_request(server, fast_poll) -> None: + async with AsyncComfy(api_key="comfyui-test-key") as client: + handle = await client.models.handle(MODEL, "req_elsewhere") + assert server.state.queue_status_count == 0 + update = await handle.status() + assert update.request_id == "req_elsewhere" + + +async def test_async_subscribe_awaits_an_async_callback(server, fast_poll) -> None: + server.state.queue_polls_to_complete = 1 + seen: list[str] = [] + + async def _record(update: QueueUpdate) -> None: + await asyncio.sleep(0) + seen.append(update.status) + + async with AsyncComfy(api_key="comfyui-test-key") as client: + result = await client.models.subscribe(MODEL, ARGS, on_queue_update=_record) + + assert result == server.state.queue_result + assert seen == ["IN_QUEUE", COMPLETED] + + +async def test_async_subscribe_timeout_cancels_before_raising(server, monkeypatch) -> None: + async def _no_sleep(_delay: float) -> None: + return None + + monkeypatch.setattr("comfy_sdk.model_requests.asyncio.sleep", _no_sleep) + server.state.queue_polls_to_complete = 10_000 + + async with AsyncComfy(api_key="comfyui-test-key") as client: + with pytest.raises(TimeoutError): + await client.models.subscribe(MODEL, ARGS, timeout=0.0) + + assert server.state.queue_cancel_count == 1 + + +async def test_async_completion_error_raises_the_typed_exception(server, fast_poll) -> None: + server.state.queue_polls_to_complete = 0 + server.state.queue_error_type = "content_policy_violation" + async with AsyncComfy(api_key="comfyui-test-key") as client: + handle = await client.models.submit(MODEL, ARGS) + with pytest.raises(ContentPolicyViolation): + await handle.get() diff --git a/tests/test_sync_async_parity.py b/tests/test_sync_async_parity.py index ed21e24..f241b06 100644 --- a/tests/test_sync_async_parity.py +++ b/tests/test_sync_async_parity.py @@ -586,8 +586,15 @@ def test_the_models_namespace_is_covered() -> None: models_pairs = [pair for pair in _PAIRS if pair[1] is _SYNC_NAMESPACES["models"]] assert models_pairs, "the models namespace produced no pair to compare" _label, sync_models, async_models = models_pairs[0] - assert "run" in _methods(sync_models), f"{sync_models.__name__}.run is not being compared" - assert "run" in _methods(async_models), f"{async_models.__name__}.run is not being compared" + # Every operation the namespace publishes, named outright. `run` is the one + # this test was written for; the queued trio joined it, and a namespace + # method that silently dropped out of the walk would otherwise leave the + # generic comparisons above passing on a smaller surface than they claim. + for name in ("run", "submit", "subscribe", "handle"): + assert name in _methods(sync_models), f"{sync_models.__name__}.{name} is not being compared" + assert name in _methods(async_models), ( + f"{async_models.__name__}.{name} is not being compared" + ) def test_introspection_is_not_vacuous() -> None: From a7a8e118e4a14c77be721d6ce1eb9c1ea590db17 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Thu, 10 Sep 2026 11:02:27 -0700 Subject: [PATCH 2/3] fix: address bot review comments on the queued model surface - bound a caller's timeout over the poll itself, its retries and the result fetch, not only the pause between polls (get/subscribe/iter_events) - start subscribe's clock before the submit, and issue its cleanup cancel under NO_RETRY with a short HTTP bound so a timeout is not followed by a minute of retrying the cancel - best-effort remote cancel when an async subscribe task is cancelled from outside, shielded and bounded - cap a server-named Retry-After at 60s before sleeping on it, so a huge value neither parks the caller nor overflows float() - guard non-object bodies: a status or submit body that is not a JSON object is an invalid_response, a result that is not one is handed back unchanged - require the authoritative status read to name a status, so a 200 {} cannot be polled forever - read a blank error_type as no error on an update, the same way the raising path already does - carry the addressed request id on an update rather than the body's copy, and bound a request id to 256 printable characters before it reaches a path or a message - send the cancel as PUT, which is the verb the contract declares - README: rehydration needs the model id and the request id, and the id is a UUID --- README.md | 2 +- src/comfy_low/transport.py | 27 +++- src/comfy_sdk/model_requests.py | 271 ++++++++++++++++++++++++++------ src/comfy_sdk/models.py | 84 ++++++---- tests/conftest.py | 45 +++++- tests/test_models_queue.py | 141 ++++++++++++++++- 6 files changed, 478 insertions(+), 92 deletions(-) diff --git a/README.md b/README.md index ed7d8ea..3855351 100644 --- a/README.md +++ b/README.md @@ -457,7 +457,7 @@ flight all at once — submit it to the queue instead: ```python handle = client.models.submit("fal-ai/flux-pro", {"prompt": "a cat"}) -handle.request_id # 'req_...' — all another process needs +handle.request_id # a UUID; with the model id, all another process needs handle.status().status # 'IN_QUEUE' / 'IN_PROGRESS' / 'COMPLETED' result = handle.get() # blocks until complete, returns the provider payload ``` diff --git a/src/comfy_low/transport.py b/src/comfy_low/transport.py index ec8af11..a1a0261 100644 --- a/src/comfy_low/transport.py +++ b/src/comfy_low/transport.py @@ -25,7 +25,7 @@ fail if that constant and the vendored path disagree. The four ``post_model_submit`` / ``get_model_request_status`` / -``get_model_request_result`` / ``post_model_request_cancel`` bindings are the +``get_model_request_result`` / ``put_model_request_cancel`` bindings are the same story one step earlier: they are the *queued* form of that one operation, and the contract declaring them is authored but held, so the vendored Router spec does not carry them yet and there is nothing for the contract test to pin @@ -119,6 +119,11 @@ _MODEL_REQUEST_STATUS_PATH_TEMPLATE = _MODEL_REQUEST_PATH_TEMPLATE + "/status" _MODEL_REQUEST_CANCEL_PATH_TEMPLATE = _MODEL_REQUEST_PATH_TEMPLATE + "/cancel" +#: Longest request id accepted into a path. The contract mints UUIDs (36 +#: characters); the bound exists so a server-controlled value that is NOT one +#: cannot reach the public handle, a log line or an exception message unbounded. +_MAX_REQUEST_ID_LENGTH = 256 + _DEFAULT_PORTS = {"http": 80, "https": 443} @@ -227,6 +232,14 @@ def parse_request_id(request_id: str) -> str: f"request id must not be '.' or '..' — it would traverse the request path " f"rather than name a request; got {request_id!r}" ) + if len(request_id) > _MAX_REQUEST_ID_LENGTH: + raise ValueError( + f"request id must be at most {_MAX_REQUEST_ID_LENGTH} characters; got {len(request_id)}" + ) + if not request_id.isprintable(): + # It is displayed and interpolated into exception messages as well as + # into the path, so a control character is refused rather than encoded. + raise ValueError(f"request id must not contain control characters; got {request_id!r}") return request_id @@ -960,10 +973,10 @@ def get_model_request_result( resp = self.raw_request("GET", url, timeout=timeout) return self._p.parse_or_raise(resp, (200,)), resp.headers - def post_model_request_cancel( + def put_model_request_cancel( self, model: str, request_id: str, *, timeout: Any = _UNSET ) -> tuple[dict[str, Any], httpx.Headers]: - """POST a cancellation for one submitted request. + """PUT a cancellation for one submitted request. A request, not a guarantee — a deployment that answers ``204`` gives an empty body, which ``parse_or_raise`` returns as ``{}``. The @@ -973,7 +986,7 @@ def post_model_request_cancel( url = self._p.router_base_url + model_request_path( model, request_id, _MODEL_REQUEST_CANCEL_PATH_TEMPLATE ) - resp = self.raw_request("POST", url, timeout=timeout) + resp = self.raw_request("PUT", url, timeout=timeout) return self._p.parse_or_raise(resp, (200, 202, 204)), resp.headers @@ -1326,14 +1339,14 @@ async def get_model_request_result( resp = await self.raw_request("GET", url, timeout=timeout) return self._p.parse_or_raise(resp, (200,)), resp.headers - async def post_model_request_cancel( + async def put_model_request_cancel( self, model: str, request_id: str, *, timeout: Any = _UNSET ) -> tuple[dict[str, Any], httpx.Headers]: - """Async :meth:`ComfyLow.post_model_request_cancel`.""" + """Async :meth:`ComfyLow.put_model_request_cancel`.""" url = self._p.router_base_url + model_request_path( model, request_id, _MODEL_REQUEST_CANCEL_PATH_TEMPLATE ) - resp = await self.raw_request("POST", url, timeout=timeout) + resp = await self.raw_request("PUT", url, timeout=timeout) return self._p.parse_or_raise(resp, (200, 202, 204)), resp.headers diff --git a/src/comfy_sdk/model_requests.py b/src/comfy_sdk/model_requests.py index ade0392..ae37fb8 100644 --- a/src/comfy_sdk/model_requests.py +++ b/src/comfy_sdk/model_requests.py @@ -41,7 +41,7 @@ import asyncio import time from collections.abc import AsyncGenerator, Awaitable, Callable, Generator, Iterator, Mapping -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import Any import httpx @@ -51,7 +51,7 @@ from . import _core from .exceptions import ComfyError, translating -from .retry import DEFAULT_RETRY, Retrier, RetryPolicy +from .retry import DEFAULT_RETRY, NO_RETRY, Retrier, RetryPolicy from .router_exceptions import RouterError, error_from_completion #: The one terminal queue status. Deliberately a single value rather than a @@ -72,6 +72,24 @@ #: this SDK just as readily as it hides an unreachable server. _CANCEL_FAILURES = (ComfyError, httpx.HTTPError) +#: Ceiling, in seconds, on a server-named ``Retry-After`` between two polls. +#: The header is honoured because the server knows its own pace, but it is a +#: hint and not a bound, and taking it verbatim would let one ``Retry-After: +#: 86400`` park a thread for a day -- or an absurd-but-parseable value overflow +#: ``float()``. A minute is long enough that a request told to wait longer is +#: still polled rarely, and short enough that nothing is parked. +_MAX_PACE = 60 + +#: Floor, in seconds, on the per-request HTTP timeout derived from a caller's +#: remaining deadline. A sub-second bound cannot complete a TLS handshake, so +#: without the floor the last poll before a deadline would be a certain +#: transport failure rather than an answer. +_MIN_HTTP_TIMEOUT = 1.0 + +#: HTTP timeout, in seconds, for the best-effort cleanup cancel ``subscribe`` +#: issues after its own timeout -- see ``RequestHandle._cancel_best_effort``. +_CANCEL_TIMEOUT = 10.0 + _now = time.monotonic @@ -139,33 +157,64 @@ def _retry_after_seconds(headers: httpx.Headers) -> int | None: return seconds if seconds > 0 else None +def _text(value: Any) -> str | None: + """A body field as a non-empty, stripped string, or ``None``. + + The same reading :func:`~comfy_sdk.router_exceptions.error_from_completion` + gives ``error_type``, so an update and the raising path cannot disagree + about whether a blank bucket is a failure. + """ + if not isinstance(value, str): + return None + stripped = value.strip() + return stripped or None + + def _update_from( - payload: Mapping[str, Any], headers: httpx.Headers, *, request_id: str + payload: Any, headers: httpx.Headers, *, request_id: str, require_status: bool = False ) -> QueueUpdate: """Build a :class:`QueueUpdate` from one response. - ``request_id`` is the id the call was addressed by, used when the body does - not repeat it — which is what keeps an update from a body-less ``204`` - cancel still identifying the request it is about. + ``request_id`` is the id the call was addressed by, and it is the one the + update carries: the body's copy is server-controlled and unvalidated, and + an update that named a different request from the one it was asked about + would be wrong in exactly the place a caller pastes into a support ticket. + It is also what keeps an update from a body-less ``204`` cancel still + identifying the request it is about. + + ``require_status`` is set by the authoritative status read, where a body + naming no status is not a state to poll again but a response this SDK + cannot act on -- treating it as "not yet terminal" would poll a ``200 {}`` + forever. It stays off for the cancel, whose ``204`` legitimately names none. """ - body_id = payload.get("request_id") + if not isinstance(payload, Mapping): + raise ComfyError( + "the queue answered with a body that is not a JSON object, so the request's " + "state cannot be read from it", + code="invalid_response", + ) status = payload.get("status") + if not isinstance(status, str): + status = "" + if require_status and not status.strip(): + raise ComfyError( + "the status read answered without naming a status, so the request's state is unknown", + code="invalid_response", + ) position = payload.get("queue_position") return QueueUpdate( - request_id=body_id if isinstance(body_id, str) and body_id else request_id, - status=status if isinstance(status, str) else "", + request_id=request_id, + status=status, queue_position=position if isinstance(position, int) and not isinstance(position, bool) else None, - error_type=payload.get("error_type") - if isinstance(payload.get("error_type"), str) - else None, + error_type=_text(payload.get("error_type")), retry_after=_retry_after_seconds(headers), raw=dict(payload), ) -def _request_id_of(payload: Mapping[str, Any]) -> str: +def _request_id_of(payload: Any) -> str: """The request id a submit response names, or a :class:`ComfyError`. A submit whose response carries no usable id is unusable in the specific @@ -174,6 +223,12 @@ def _request_id_of(payload: Mapping[str, Any]) -> str: call, so it is raised rather than papered over with a placeholder id that would 404 on the first poll. """ + if not isinstance(payload, Mapping): + raise ComfyError( + "the queue accepted the request but answered with a body that is not a JSON " + "object, so no request_id could be read from it", + code="invalid_response", + ) request_id = payload.get("request_id") if not isinstance(request_id, str) or not request_id: raise ComfyError( @@ -194,9 +249,7 @@ def _request_id_of(payload: Mapping[str, Any]) -> str: ) from exc -def _raise_for_completion( - payload: Mapping[str, Any], *, request_id: str, envelope_only: bool = False -) -> None: +def _raise_for_completion(payload: Any, *, request_id: str, envelope_only: bool = False) -> None: """Raise the typed router exception a completion reports, if it reports one. The gate behind "a ``200`` with an error payload is never returned as @@ -215,6 +268,11 @@ def _raise_for_completion( the server reports where it is authoritative is never the one that gets missed. """ + if not isinstance(payload, Mapping): + # A partner's native output is whatever JSON document the partner + # answers with -- an array or a bare value is a result, not an + # envelope, and there is nothing in it the queue could have reported. + return if envelope_only and payload.get("status") != COMPLETED: return error = error_from_completion(payload, request_id=request_id) @@ -240,10 +298,14 @@ def _pace(update: QueueUpdate, backoff: Iterator[float]) -> float: A pace the server named beats the schedule guessed here — that is the whole point of ``Retry-After`` — and the adaptive backoff carries the interval when it named none. The backoff is advanced either way so the schedule does - not restart from its floor the moment the server stops naming a pace. + not restart from its floor the moment the server stops naming a pace. The + named pace is capped at :data:`_MAX_PACE`: a hint, honoured, but not a bound + a single header can park the caller behind. """ scheduled = next(backoff) - return float(update.retry_after) if update.retry_after is not None else scheduled + if update.retry_after is None: + return scheduled + return float(min(update.retry_after, _MAX_PACE)) def _remaining(deadline: float | None) -> float | None: @@ -275,12 +337,41 @@ def _last(updates: Iterator[QueueUpdate]) -> QueueUpdate: return _completed(final) -def _timed_out(request_id: str, timeout: float | None, update: QueueUpdate) -> TimeoutError: +def _timed_out(request_id: str, timeout: float | None, update: QueueUpdate | None) -> TimeoutError: + status = "unknown" if update is None else update.status return TimeoutError( - f"model request {request_id} not complete after {timeout}s (status={update.status!r})" + f"model request {request_id} not complete after {timeout}s (status={status!r})" + ) + + +def _bounded(policy: RetryPolicy, budget: float | None) -> RetryPolicy: + """``policy`` with both of its elapsed budgets capped at ``budget`` seconds. + + How a caller's deadline reaches the retrier: a poll made with two seconds + left must not be allowed a minute of retries, and one made with nothing + left gets exactly one attempt (a zero ``max_elapsed`` is ``NO_RETRY``). + """ + if budget is None: + return policy + left = max(budget, 0.0) + return replace( + policy, + max_elapsed=min(policy.max_elapsed, left), + collect_max_elapsed=min(policy.collect_max_elapsed, left), ) +def _http_timeout(budget: float | None) -> dict[str, Any]: + """Keyword arguments bounding one low-level call's HTTP timeout by ``budget``. + + Empty when there is no budget, so the client's own timeout applies. Floored + at :data:`_MIN_HTTP_TIMEOUT` for the reason given there. + """ + if budget is None: + return {} + return {"timeout": max(budget, _MIN_HTTP_TIMEOUT)} + + class _RequestHandleBase: """State and read-only views shared by the sync and async handles.""" @@ -337,11 +428,23 @@ def status(self) -> QueueUpdate: raising: this is the read a caller uses to *look*, and the raising belongs to :meth:`get`, which is the one that hands back a result. """ + return self._status(budget=None) + + def _status(self, *, budget: float | None) -> QueueUpdate: + """One authoritative poll, bounded by ``budget`` seconds when one is given. + + The bound covers the whole call — the HTTP request and any retry of it + — so a caller's ``timeout`` on :meth:`iter_events` is a bound on the + loop and not only on the pauses between its polls. + """ with translating(): payload, headers = self._call( - lambda: self._low.get_model_request_status(self._model, self._request_id) + lambda: self._low.get_model_request_status( + self._model, self._request_id, **_http_timeout(budget) + ), + budget=budget, ) - return _update_from(payload, headers, request_id=self._request_id) + return _update_from(payload, headers, request_id=self._request_id, require_status=True) def iter_events(self, timeout: float | None = None) -> Generator[QueueUpdate, None, None]: """Poll to completion, yielding an update whenever the queue moves. @@ -353,19 +456,28 @@ def iter_events(self, timeout: float | None = None) -> Generator[QueueUpdate, No progress and a caller who wants the result calls :meth:`get`, which does raise. - ``timeout`` is a client-side bound in seconds on the whole loop; ``None`` - polls until the server says the request is done. Exceeding it raises - ``TimeoutError`` and leaves the request running — the queue is the - server's, so a local clock running out says nothing about it. Cancelling - on the way out is ``models.subscribe``'s behaviour, deliberately not - this one's: an iterator that cancelled the work it was iterating would - make a ``for`` loop with a ``break`` destructive. + ``timeout`` is a client-side bound in seconds on the whole loop — the + polls, their retries and the pauses between them, not the pauses alone + — and ``None`` polls until the server says the request is done. The + first poll is always made, so ``timeout=0`` reads "look once". Past the + bound no further poll is started, and the last one is held to what is + left of it (with a floor of :data:`_MIN_HTTP_TIMEOUT` so it can still + complete a handshake), so the loop overruns its bound by at most one + such request. Exceeding it raises ``TimeoutError`` and leaves the + request running — the queue is the server's, so a local clock running + out says nothing about it. Cancelling on the way out is + ``models.subscribe``'s behaviour, deliberately not this one's: an + iterator that cancelled the work it was iterating would make a ``for`` + loop with a ``break`` destructive. """ deadline = None if timeout is None else _now() + timeout backoff = _core.backoff_schedule() previous: QueueUpdate | None = None while True: - update = self.status() + remaining = _remaining(deadline) + if previous is not None and remaining is not None and remaining <= 0: + raise _timed_out(self._request_id, timeout, previous) + update = self._status(budget=remaining) if _changed(previous, update): yield update previous = update @@ -382,31 +494,42 @@ def get(self, timeout: float | None = None) -> dict[str, Any]: The result is the partner model's own output, decoded from JSON and handed back as-is — the same value ``models.run`` returns for the same - model and arguments. + model and arguments. It is a JSON object for every model Router serves + today; a partner whose native output were an array or a bare value + would be handed back unchanged too, since the payload is the partner's + and not this SDK's to reshape. Raises the typed router exception (:mod:`comfy_sdk.router_exceptions`) when the completion carries an ``error_type``, which is how the server reports a failed *or* cancelled request. ``timeout`` bounds the wait exactly as it does on - :meth:`iter_events`, and raises ``TimeoutError`` without cancelling. + :meth:`iter_events`, the result fetch included, and raises + ``TimeoutError`` without cancelling. Calling it on a request that has already completed is one status poll and one fetch, so collecting a result twice — or from a second process — costs no more than the first time. """ - return self._collect(_last(self.iter_events(timeout=timeout))) + deadline = None if timeout is None else _now() + timeout + completion = _last(self.iter_events(timeout=timeout)) + return self._collect(completion, budget=_remaining(deadline)) - def _collect(self, completion: QueueUpdate) -> dict[str, Any]: + def _collect(self, completion: QueueUpdate, *, budget: float | None = None) -> dict[str, Any]: """Turn an observed completion into a result, or into the typed error. Split out of :meth:`get` so ``models.subscribe`` — which has already polled its way to the completion — can collect from the update it is holding instead of spending one more status request re-discovering it. + ``budget`` is what is left of the caller's deadline, and bounds the + fetch the way :meth:`_status` bounds a poll. """ _raise_for_completion(completion.raw, request_id=self._request_id) with translating(): payload, _headers = self._call( - lambda: self._low.get_model_request_result(self._model, self._request_id) + lambda: self._low.get_model_request_result( + self._model, self._request_id, **_http_timeout(budget) + ), + budget=budget, ) # Checked again on the result body: which of the two responses carries # the `error_type` is the server's choice, and reading only one of them @@ -426,12 +549,34 @@ def cancel(self) -> QueueUpdate: """ with translating(): payload, headers = self._call( - lambda: self._low.post_model_request_cancel(self._model, self._request_id) + lambda: self._low.put_model_request_cancel(self._model, self._request_id) ) return _update_from(payload, headers, request_id=self._request_id) + def _cancel_best_effort(self) -> None: + """The cleanup cancel ``models.subscribe`` issues after its own timeout. + + One attempt under :data:`~comfy_sdk.retry.NO_RETRY` and a short HTTP + bound, because it runs inside the handling of a ``TimeoutError`` the + caller is about to see: a cancel that rode the client's full retry + policy could hold that caller for the whole of ``max_elapsed`` — or + ``collect_max_elapsed``, if the cancel were answered with a paced + ``429`` — after they had already stopped waiting. + """ + with translating(): + self._call( + lambda: self._low.put_model_request_cancel( + self._model, self._request_id, timeout=_CANCEL_TIMEOUT + ), + policy=NO_RETRY, + ) + def _call( - self, send: Callable[[], tuple[dict[str, Any], httpx.Headers]] + self, + send: Callable[[], tuple[dict[str, Any], httpx.Headers]], + *, + budget: float | None = None, + policy: RetryPolicy | None = None, ) -> tuple[dict[str, Any], httpx.Headers]: """Run one queue call under the client's retry policy. @@ -440,8 +585,10 @@ def _call( would spend the whole budget on its first hour of polling and then surface the next blip as a hard failure. What it buys here is that a ``429`` naming a ``Retry-After`` paces the poll instead of ending it. + ``budget`` caps that policy's elapsed budgets at what is left of the + caller's deadline; ``policy`` substitutes another policy outright. """ - retrier = Retrier(self._retry, now=_now) + retrier = Retrier(_bounded(policy or self._retry, budget), now=_now) while True: try: return send() @@ -470,11 +617,18 @@ def __init__( async def status(self) -> QueueUpdate: """Awaitable :meth:`RequestHandle.status` — one authoritative poll.""" + return await self._status(budget=None) + + async def _status(self, *, budget: float | None) -> QueueUpdate: + """Async :meth:`RequestHandle._status` — one poll, bounded by ``budget``.""" with translating(): payload, headers = await self._call( - lambda: self._low.get_model_request_status(self._model, self._request_id) + lambda: self._low.get_model_request_status( + self._model, self._request_id, **_http_timeout(budget) + ), + budget=budget, ) - return _update_from(payload, headers, request_id=self._request_id) + return _update_from(payload, headers, request_id=self._request_id, require_status=True) async def iter_events(self, timeout: float | None = None) -> AsyncGenerator[QueueUpdate, None]: """Async :meth:`RequestHandle.iter_events` — ``async for`` over the updates.""" @@ -482,7 +636,10 @@ async def iter_events(self, timeout: float | None = None) -> AsyncGenerator[Queu backoff = _core.backoff_schedule() previous: QueueUpdate | None = None while True: - update = await self.status() + remaining = _remaining(deadline) + if previous is not None and remaining is not None and remaining <= 0: + raise _timed_out(self._request_id, timeout, previous) + update = await self._status(budget=remaining) if _changed(previous, update): yield update previous = update @@ -496,17 +653,23 @@ async def iter_events(self, timeout: float | None = None) -> AsyncGenerator[Queu async def get(self, timeout: float | None = None) -> dict[str, Any]: """Async :meth:`RequestHandle.get` — wait, then collect or raise.""" + deadline = None if timeout is None else _now() + timeout completion: QueueUpdate | None = None async for update in self.iter_events(timeout=timeout): completion = update - return await self._collect(_completed(completion)) + return await self._collect(_completed(completion), budget=_remaining(deadline)) - async def _collect(self, completion: QueueUpdate) -> dict[str, Any]: + async def _collect( + self, completion: QueueUpdate, *, budget: float | None = None + ) -> dict[str, Any]: """Async :meth:`RequestHandle._collect`.""" _raise_for_completion(completion.raw, request_id=self._request_id) with translating(): payload, _headers = await self._call( - lambda: self._low.get_model_request_result(self._model, self._request_id) + lambda: self._low.get_model_request_result( + self._model, self._request_id, **_http_timeout(budget) + ), + budget=budget, ) _raise_for_completion(payload, request_id=self._request_id, envelope_only=True) return payload @@ -515,15 +678,29 @@ async def cancel(self) -> QueueUpdate: """Async :meth:`RequestHandle.cancel` — a request, not a guarantee.""" with translating(): payload, headers = await self._call( - lambda: self._low.post_model_request_cancel(self._model, self._request_id) + lambda: self._low.put_model_request_cancel(self._model, self._request_id) ) return _update_from(payload, headers, request_id=self._request_id) + async def _cancel_best_effort(self) -> None: + """Async :meth:`RequestHandle._cancel_best_effort` — one bounded attempt.""" + with translating(): + await self._call( + lambda: self._low.put_model_request_cancel( + self._model, self._request_id, timeout=_CANCEL_TIMEOUT + ), + policy=NO_RETRY, + ) + async def _call( - self, send: Callable[[], Awaitable[tuple[dict[str, Any], httpx.Headers]]] + self, + send: Callable[[], Awaitable[tuple[dict[str, Any], httpx.Headers]]], + *, + budget: float | None = None, + policy: RetryPolicy | None = None, ) -> tuple[dict[str, Any], httpx.Headers]: """Async :meth:`RequestHandle._call` — one queue call under the retry policy.""" - retrier = Retrier(self._retry, now=_now) + retrier = Retrier(_bounded(policy or self._retry, budget), now=_now) while True: try: return await send() diff --git a/src/comfy_sdk/models.py b/src/comfy_sdk/models.py index 27b461a..c7fa3c2 100644 --- a/src/comfy_sdk/models.py +++ b/src/comfy_sdk/models.py @@ -62,10 +62,12 @@ from .exceptions import translating from .model_requests import ( _CANCEL_FAILURES, + _CANCEL_TIMEOUT, AsyncRequestHandle, QueueUpdate, RequestHandle, _completed, + _remaining, _request_id_of, ) from .retry import DEFAULT_RETRY, Retrier, RetryPolicy @@ -362,8 +364,10 @@ def subscribe( exception it raises propagates and abandons the wait (the request keeps running server side). - ``timeout`` is a **client-side** bound in seconds on the whole wait, - with no server-side meaning: the queue's own timeouts are the server's. + ``timeout`` is a **client-side** bound in seconds on the whole call — + the submit's wait excluded only where its own retry policy is already + running, then every poll, retry and pause, and the result fetch — with + no server-side meaning: the queue's own timeouts are the server's. When it runs out this makes a best-effort :meth:`~comfy_sdk.model_requests.RequestHandle.cancel` — so a caller that has stopped waiting is not also still paying for a generation @@ -378,12 +382,15 @@ def subscribe( exception rather than returning, so a ``200`` never comes back as a successful result. """ + # The clock starts here, before the submit, so ``timeout`` bounds the + # whole call as documented and not only the polling after it. + deadline = None if timeout is None else _now() + timeout handle = self.submit(model, arguments, idempotency_key=idempotency_key) # ``closing`` so the poll generator is finalised on every exit — the # completion, the timeout, and above all the one where the caller's # callback raises, which otherwise leaves it suspended until the # collector happens to reach it. - with contextlib.closing(handle.iter_events(timeout=timeout)) as updates: + with contextlib.closing(handle.iter_events(timeout=_remaining(deadline))) as updates: completion: QueueUpdate | None = None while True: # Only the *iteration* is inside the ``except TimeoutError``. A @@ -398,7 +405,7 @@ def subscribe( break except TimeoutError: try: - handle.cancel() + handle._cancel_best_effort() except _CANCEL_FAILURES: # Best-effort is literal: the timeout is the failure # worth reporting, and a masked one sends the caller @@ -408,7 +415,7 @@ def subscribe( completion = update if on_queue_update is not None: on_queue_update(update) - return handle._collect(_completed(completion)) + return handle._collect(_completed(completion), budget=_remaining(deadline)) def handle(self, model: str, request_id: str) -> RequestHandle: """Rebuild the handle for a request submitted anywhere. @@ -533,33 +540,50 @@ async def subscribe( an awaitable it returns is awaited before the next poll, so an async callback does not need wrapping. See :meth:`Models.subscribe`. """ + deadline = None if timeout is None else _now() + timeout handle = await self.submit(model, arguments, idempotency_key=idempotency_key) - # ``aclosing`` for the reason ``Models.subscribe`` uses ``closing``, and - # a sharper one: an async generator left suspended is finalised by the - # event loop's own shutdown hook, well after this call returned. - async with contextlib.aclosing(handle.iter_events(timeout=timeout)) as updates: - completion: QueueUpdate | None = None - while True: - # Scoped to the iteration alone, for the reason - # `Models.subscribe` gives — and it bites harder here, where a - # callback that awaits anything under `asyncio.wait_for` raises - # `TimeoutError` natively. - try: - update = await anext(updates) - except StopAsyncIteration: - break - except TimeoutError: + completion: QueueUpdate | None = None + try: + # ``aclosing`` for the reason ``Models.subscribe`` uses ``closing``, + # and a sharper one: an async generator left suspended is finalised + # by the event loop's own shutdown hook, well after this call + # returned. + async with contextlib.aclosing( + handle.iter_events(timeout=_remaining(deadline)) + ) as updates: + while True: + # Scoped to the iteration alone, for the reason + # `Models.subscribe` gives — and it bites harder here, where + # a callback that awaits anything under `asyncio.wait_for` + # raises `TimeoutError` natively. try: - await handle.cancel() - except _CANCEL_FAILURES: - pass - raise - completion = update - if on_queue_update is not None: - outcome = on_queue_update(update) - if isinstance(outcome, Awaitable): - await outcome - return await handle._collect(_completed(completion)) + update = await anext(updates) + except StopAsyncIteration: + break + except TimeoutError: + try: + await handle._cancel_best_effort() + except _CANCEL_FAILURES: + pass + raise + completion = update + if on_queue_update is not None: + outcome = on_queue_update(update) + if isinstance(outcome, Awaitable): + await outcome + except asyncio.CancelledError: + # The task was cancelled from outside while the request is still + # queued or running. `subscribe` has exposed neither the handle nor + # its key, so the caller has no way back to a generation that would + # otherwise keep running — and keep billing — after they stopped + # waiting for it. One shielded, bounded best-effort cancel, exactly + # as on the timeout path, then the cancellation proceeds. + with contextlib.suppress(*_CANCEL_FAILURES, TimeoutError, asyncio.TimeoutError): + await asyncio.shield( + asyncio.wait_for(handle._cancel_best_effort(), _CANCEL_TIMEOUT + 1.0) + ) + raise + return await handle._collect(_completed(completion), budget=_remaining(deadline)) async def handle(self, model: str, request_id: str) -> AsyncRequestHandle: """Awaitable :meth:`Models.handle` — rebuild a handle from the two ids. diff --git a/tests/conftest.py b/tests/conftest.py index 5e48478..6321bd7 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -223,6 +223,17 @@ class ServerState: ) # Status code for the cancel response; 204 exercises the empty-body path. queue_cancel_status: int = 200 + # Cancels that answer a transient failure (status, code) before one is + # accepted — for proving the cleanup cancel after a timeout does not ride + # the client's retry policy. + queue_cancel_fail_times: int = 0 + queue_cancel_transient_error: tuple[int, str] = (429, "rate_limited") + queue_cancel_transient_retry_after: str | None = "1" + # When set, the result read answers this JSON document verbatim instead of + # `queue_result` — for a provider whose native output is not an object. + queue_result_raw: Any = None + # The status read answers a body naming no `status` at all. + queue_status_omits_status: bool = False # The bucket a cancelled request's completion carries. queue_cancel_error_type: str = "client_disconnected" # Set by a cancel; makes every later status poll report the cancellation. @@ -288,6 +299,8 @@ class ServerState: queue_status_served: int = 0 queue_result_count: int = 0 queue_cancel_count: int = 0 + # The HTTP method each cancel arrived with. + queue_cancel_methods: list[str] = field(default_factory=list) # Every Idempotency-Key seen on a queue submit, in arrival order. queue_submit_idempotency_keys: list[str | None] = field(default_factory=list) # The native body of the last queue submit, and the two decoded id segments. @@ -572,6 +585,19 @@ def frame(event: str, data: dict) -> None: frame("status", {"status": state.terminal_status}) # -- POST -- + def do_PUT(self) -> None: + # Comfy Router's queue cancel is a PUT (the contract's + # `cancelRouterModelRequest`), so it is served here and nowhere + # else: a POST to the same path is the wrong verb and gets a 404 + # like any other unrouted request. + m = re.match(r"/v2/models/([^/]+)/([^/]+)/requests/([^/]+)/cancel$", self.path) + if m: + self._read_body() + self._put_queue_cancel(m.group(3)) + return + self._read_body() + self._err(404, "not_found") + def do_POST(self) -> None: if not self._auth_ok(): self._read_body() @@ -600,11 +626,6 @@ def do_POST(self) -> None: if m: self._post_queue_submit(m.group(1), m.group(2)) return - m = re.match(r"/v2/models/([^/]+)/([^/]+)/requests/([^/]+)/cancel$", self.path) - if m: - self._read_body() - self._post_queue_cancel(m.group(3)) - return m = re.match(r"/api/v2/jobs/([^/]+)/cancel$", self.path) if m: self._json(200, _job_json(m.group(1), "canceling")) @@ -669,6 +690,9 @@ def _serve_queue_status(self, request_id: str) -> None: else {} ) body: dict[str, Any] = {"request_id": unquote(request_id)} + if state.queue_status_omits_status: + self._json(200, body, headers=headers) + return if state.queue_canceled: body["status"] = "COMPLETED" body["error_type"] = state.queue_cancel_error_type @@ -692,6 +716,9 @@ def _serve_queue_status(self, request_id: str) -> None: def _serve_queue_result(self, request_id: str) -> None: state.queue_result_count += 1 state.queue_paths.append(self.path) + if state.queue_result_raw is not None: + self._json(200, state.queue_result_raw) + return if state.queue_result_error_type: self._json( 200, @@ -705,9 +732,15 @@ def _serve_queue_result(self, request_id: str) -> None: return self._json(200, {**state.queue_result, **state.queue_result_extra}) - def _post_queue_cancel(self, request_id: str) -> None: + def _put_queue_cancel(self, request_id: str) -> None: state.queue_cancel_count += 1 state.queue_paths.append(self.path) + state.queue_cancel_methods.append(self.command) + if state.queue_cancel_fail_times > 0: + state.queue_cancel_fail_times -= 1 + status, code = state.queue_cancel_transient_error + self._router_err(status, code, retry_after=state.queue_cancel_transient_retry_after) + return state.queue_canceled = True if state.queue_cancel_status == 204: self.send_response(204) diff --git a/tests/test_models_queue.py b/tests/test_models_queue.py index 4967e61..0cb70fc 100644 --- a/tests/test_models_queue.py +++ b/tests/test_models_queue.py @@ -512,7 +512,7 @@ def test_a_failing_cancel_does_not_mask_the_timeout(server, monkeypatch) -> None def _explode(self: Any) -> None: raise ComfyError("cancel is unreachable") - monkeypatch.setattr(RequestHandle, "cancel", _explode) + monkeypatch.setattr(RequestHandle, "_cancel_best_effort", _explode) with _client() as client: with pytest.raises(TimeoutError): @@ -685,3 +685,142 @@ async def test_async_completion_error_raises_the_typed_exception(server, fast_po handle = await client.models.submit(MODEL, ARGS) with pytest.raises(ContentPolicyViolation): await handle.get() + + +# --- review follow-ups: bounded waits, malformed bodies, the cleanup cancel --- + + +def test_a_blank_error_type_reads_as_no_error_on_an_update() -> None: + """The update and the raising path read ``error_type`` the same way.""" + import httpx + + from comfy_sdk.model_requests import _update_from + + body = {"status": COMPLETED, "error_type": " "} + update = _update_from(body, httpx.Headers(), request_id="r") + + assert update.error_type is None + assert error_from_completion(body) is None + + +def test_an_update_carries_the_id_it_was_addressed_by() -> None: + import httpx + + from comfy_sdk.model_requests import _update_from + + body = {"request_id": "somebody-else\n", "status": "IN_QUEUE"} + update = _update_from(body, httpx.Headers(), request_id="mine") + + assert update.request_id == "mine" + assert update.raw == body + + +def test_a_result_that_is_not_a_json_object_is_returned_unchanged(server, fast_poll) -> None: + """The result is the partner's document, whatever shape the partner gave it.""" + server.state.queue_polls_to_complete = 0 + server.state.queue_result_raw = [{"url": "http://example.invalid/a.png"}] + + with _client() as client: + assert client.models.submit(MODEL, ARGS).get() == server.state.queue_result_raw + + +def test_a_status_read_naming_no_status_is_an_invalid_response(server, fast_poll) -> None: + """A ``200 {}`` from the authoritative read must not poll forever.""" + server.state.queue_status_omits_status = True + + with _client() as client: + handle = client.models.submit(MODEL, ARGS) + with pytest.raises(ComfyError) as excinfo: + handle.get() + + assert excinfo.value.code == "invalid_response" + + +def test_a_huge_retry_after_is_capped_before_it_is_slept(server, monkeypatch) -> None: + slept: list[float] = [] + monkeypatch.setattr("comfy_sdk.model_requests.time.sleep", slept.append) + server.state.queue_polls_to_complete = 1 + # Parses as an int; `float()` of it would overflow. + server.state.queue_status_retry_after = "9" * 400 + + with _client() as client: + client.models.submit(MODEL, ARGS).get() + + assert slept == [60.0] + + +def test_a_timeout_bounds_the_poll_and_its_retries_not_only_the_sleep(server, monkeypatch) -> None: + """``get(timeout=...)`` on a server that keeps throttling returns within the + bound instead of riding the retry policy's whole minute.""" + monkeypatch.setattr("comfy_sdk.model_requests.time.sleep", lambda _s: None) + server.state.queue_status_fail_times = 10_000 + + started = time.monotonic() + with _client() as client: + handle = client.models.submit(MODEL, ARGS) + with pytest.raises((RouterError, TimeoutError)): + handle.get(timeout=0.5) + + assert time.monotonic() - started < 5 + + +def test_the_cleanup_cancel_after_a_timeout_does_not_ride_the_retry_policy( + server, monkeypatch +) -> None: + monkeypatch.setattr("comfy_sdk.model_requests.time.sleep", lambda _s: None) + server.state.queue_polls_to_complete = 10_000 + # Every cancel is answered with a paced 429, which the full policy would + # retry for the whole of its budget. + server.state.queue_cancel_fail_times = 10_000 + + started = time.monotonic() + with _client() as client: + with pytest.raises(TimeoutError): + client.models.subscribe(MODEL, ARGS, timeout=0.0) + + assert server.state.queue_cancel_count == 1 + assert time.monotonic() - started < 5 + + +@pytest.mark.parametrize("request_id", ["abc\n", "with\x00nul", "x" * 257]) +def test_handle_refuses_an_unprintable_or_oversized_request_id(server, request_id) -> None: + with _client() as client: + with pytest.raises(ValueError): + client.models.handle(MODEL, request_id) + + +def test_cancel_is_a_put(server, fast_poll) -> None: + """The contract's cancel is ``PUT``; a POST would be the wrong verb.""" + with _client() as client: + client.models.submit(MODEL, ARGS).cancel() + + assert server.state.queue_cancel_methods == ["PUT"] + + +async def test_async_subscribe_cancellation_requests_a_remote_cancel(server, monkeypatch) -> None: + """A task cancelled from outside still asks the server to stop the run. + + The cancel is delivered while the loop is in its own pause between polls, + so the test exercises this SDK's handling of the cancellation rather than + the HTTP stack's: a ``Task.cancel()`` that lands inside an in-flight + request is the transport's to surface, and when it surfaces late the loop + simply reaches this same pause on its next iteration. + """ + server.state.queue_polls_to_complete = 10_000 + real_sleep = asyncio.sleep + pausing = asyncio.Event() + + async def _pause(_delay: float) -> None: + pausing.set() + await real_sleep(0.05) + + monkeypatch.setattr("comfy_sdk.model_requests.asyncio.sleep", _pause) + + async with AsyncComfy(api_key="comfyui-test-key") as client: + task = asyncio.ensure_future(client.models.subscribe(MODEL, ARGS)) + await pausing.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert server.state.queue_cancel_count == 1 From 0f5426dce20005c4e19fc8f73bcb113533c305ee Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Thu, 10 Sep 2026 11:38:54 -0700 Subject: [PATCH 3/3] fix: address bot review comments on the queue tests and result docs - conftest: do_PUT runs the same _auth_ok check as the other verbs, so the cancel route records its Authorization header and honours require_auth - test_models_queue: the asyncio.sleep patch is a pass-through that only reports the pause; every caller still waits the delay it asked for - RequestHandle.get: the docstring says the dict[str, Any] annotation is the contract (shared with models.run) and that an off-contract non-object payload is passed through as robustness, not as a second return type --- src/comfy_sdk/model_requests.py | 11 +++++++---- tests/conftest.py | 4 ++++ tests/test_models_queue.py | 8 ++++++-- 3 files changed, 17 insertions(+), 6 deletions(-) diff --git a/src/comfy_sdk/model_requests.py b/src/comfy_sdk/model_requests.py index ae37fb8..8913565 100644 --- a/src/comfy_sdk/model_requests.py +++ b/src/comfy_sdk/model_requests.py @@ -494,10 +494,13 @@ def get(self, timeout: float | None = None) -> dict[str, Any]: The result is the partner model's own output, decoded from JSON and handed back as-is — the same value ``models.run`` returns for the same - model and arguments. It is a JSON object for every model Router serves - today; a partner whose native output were an array or a bare value - would be handed back unchanged too, since the payload is the partner's - and not this SDK's to reshape. + model and arguments, under the same ``dict[str, Any]`` annotation. That + annotation is the contract: every model Router serves answers with a + JSON object. A partner whose native output were an array or a bare + value would still be handed back unchanged rather than rejected, since + the payload is the partner's and not this SDK's to reshape — but that + is robustness against an off-contract payload, not a second supported + return type. Raises the typed router exception (:mod:`comfy_sdk.router_exceptions`) when the completion carries an diff --git a/tests/conftest.py b/tests/conftest.py index 6321bd7..80c9e30 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -586,6 +586,10 @@ def frame(event: str, data: dict) -> None: # -- POST -- def do_PUT(self) -> None: + if not self._auth_ok(): + self._read_body() + self._err(401, "unauthorized", "no key") + return # Comfy Router's queue cancel is a PUT (the contract's # `cancelRouterModelRequest`), so it is served here and nowhere # else: a POST to the same path is the wrong verb and gets a 404 diff --git a/tests/test_models_queue.py b/tests/test_models_queue.py index 0cb70fc..04720e1 100644 --- a/tests/test_models_queue.py +++ b/tests/test_models_queue.py @@ -810,9 +810,13 @@ async def test_async_subscribe_cancellation_requests_a_remote_cancel(server, mon real_sleep = asyncio.sleep pausing = asyncio.Event() - async def _pause(_delay: float) -> None: + # `comfy_sdk.model_requests` imports the `asyncio` module itself, so this + # patch lands on the shared `asyncio.sleep` for the test's duration. It is + # a pass-through that only *reports* the pause: every caller still waits + # the delay it asked for, and the cancel below is what cuts the wait short. + async def _pause(delay: float) -> None: pausing.set() - await real_sleep(0.05) + await real_sleep(delay) monkeypatch.setattr("comfy_sdk.model_requests.asyncio.sleep", _pause)