diff --git a/CHANGELOG.md b/CHANGELOG.md index 78ca920..a6047e5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,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. - `Asset.get_download_url()` / `AsyncAsset.get_download_url()` — a directly-fetchable URL for an *uploaded* asset's bytes, mirroring `Output.get_download_url()` (same `DownloadUrl`, commits the asset first if diff --git a/README.md b/README.md index 008e9f8..1fc5056 100644 --- a/README.md +++ b/README.md @@ -401,7 +401,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 @@ -507,6 +509,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 # 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 +``` + +`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 @@ -650,10 +749,10 @@ raises, so a handler never has to guard the attribute access itself. itself, and so on the submit phase of `run()`. A failure raised before the request exists (a UI-format workflow, an asset that would not upload) or while `run()` polls the job afterwards carries none. What the key is *good for* -differs, so read it with the surface in mind. `models.run` sends it to a -surface that **replays** a claimed key, which is what makes it a handle on a -generation you were already billed for. `POST /jobs` instead **rejects** a -reused key with `422 idempotency_key_reuse` (see the +differs, so read it with the surface in mind. `models.run` and `models.submit` +send it to a surface that **replays** a claimed key, which is what makes it a +handle on a generation you were already billed for. `POST /jobs` instead +**rejects** a reused key with `422 idempotency_key_reuse` (see the [`IdempotencyKeyReuse`](#typed-errors) bullet below): keys there are single-use and there is no replay. So on a `submit()` failure `exc.idempotency_key` is the key this attempt was made under, not a replay handle. Whether the server ever @@ -797,3 +896,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 536cde6..18fc864 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`` / ``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 +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,31 @@ #: 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" + +#: 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} @@ -170,6 +205,83 @@ 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}" + ) + 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 + + +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 @@ -827,6 +939,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 put_model_request_cancel( + self, model: str, request_id: str, *, timeout: Any = _UNSET + ) -> tuple[dict[str, Any], httpx.Headers]: + """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 + 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("PUT", url, timeout=timeout) + return self._p.parse_or_raise(resp, (200, 202, 204)), resp.headers + class AsyncComfyLow: """Asynchronous protocol bindings — mirrors :class:`ComfyLow`.""" @@ -1142,6 +1337,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 put_model_request_cancel( + self, model: str, request_id: str, *, timeout: Any = _UNSET + ) -> tuple[dict[str, Any], httpx.Headers]: + """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("PUT", 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 911b92a..3bc6b9f 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 a3535b9..d597ab7 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, and by + #: :meth:`comfy_sdk.models.Models.run` and + #: :meth:`comfy_sdk.models.Models.submit` and their async twins, and by #: :meth:`comfy_sdk.client.Comfy.submit` / #: :meth:`comfy_sdk.client.AsyncComfy.submit` on every failure of the #: ``POST /jobs`` attempt itself — and so by the submit phase of @@ -37,15 +38,15 @@ class ComfyError(Exception): #: a resend is safe. #: #: What the key is *good for* differs by surface, so read it with the - #: operation in mind: ``models.run`` sends it to a surface that replays a - #: claimed key, so the key is a handle on the generation you were already - #: billed for. ``POST /jobs`` instead *rejects* a reused key with - #: ``422 idempotency_key_reuse``, so on a ``submit`` failure the key is the - #: one this attempt was made under, not a replay handle: after an - #: ambiguous failure poll or list for the job the first attempt may have - #: created rather than resubmitting under it, while a failure the server - #: never saw (a connect failure, an exhausted ``QueueFull``) leaves the key - #: unclaimed. + #: operation in mind: ``models.run`` and ``models.submit`` send it to a + #: surface that replays a claimed key, so the key is a handle on the + #: generation you were already billed for. ``POST /jobs`` instead + #: *rejects* a reused key with ``422 idempotency_key_reuse``, so on a + #: ``Comfy.submit`` failure the key is the one this attempt was made + #: under, not a replay handle: after an ambiguous failure poll or list for + #: the job the first attempt may have created rather than resubmitting + #: under it, while a failure the server never saw (a connect failure, an + #: exhausted ``QueueFull``) leaves the key unclaimed. #: #: Declared on the base rather than set per subclass so that a bucket this #: SDK version has never heard of — which arrives as a bare diff --git a/src/comfy_sdk/model_requests.py b/src/comfy_sdk/model_requests.py new file mode 100644 index 0000000..8913565 --- /dev/null +++ b/src/comfy_sdk/model_requests.py @@ -0,0 +1,717 @@ +"""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, replace +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, NO_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) + +#: 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 + + +@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 _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: 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, 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. + """ + 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=request_id, + status=status, + queue_position=position + if isinstance(position, int) and not isinstance(position, bool) + else None, + error_type=_text(payload.get("error_type")), + retry_after=_retry_after_seconds(headers), + raw=dict(payload), + ) + + +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 + 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. + """ + 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( + "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: 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 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) + 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. 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) + if update.retry_after is None: + return scheduled + return float(min(update.retry_after, _MAX_PACE)) + + +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 | None) -> TimeoutError: + status = "unknown" if update is None else update.status + return TimeoutError( + 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.""" + + _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. + """ + 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, **_http_timeout(budget) + ), + budget=budget, + ) + 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. + + 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 — 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: + 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 + 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, 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 + ``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`, 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. + """ + 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, *, 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, **_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 + # 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.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]], + *, + budget: float | None = None, + policy: RetryPolicy | None = None, + ) -> 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. + ``budget`` caps that policy's elapsed budgets at what is left of the + caller's deadline; ``policy`` substitutes another policy outright. + """ + retrier = Retrier(_bounded(policy or self._retry, budget), 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.""" + 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, **_http_timeout(budget) + ), + budget=budget, + ) + 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.""" + deadline = None if timeout is None else _now() + timeout + backoff = _core.backoff_schedule() + previous: QueueUpdate | None = None + while True: + 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 + 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.""" + 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), budget=_remaining(deadline)) + + 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, **_http_timeout(budget) + ), + budget=budget, + ) + _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.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]]], + *, + 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(_bounded(policy or self._retry, budget), 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..c7fa3c2 100644 --- a/src/comfy_sdk/models.py +++ b/src/comfy_sdk/models.py @@ -41,18 +41,35 @@ 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, + _CANCEL_TIMEOUT, + AsyncRequestHandle, + QueueUpdate, + RequestHandle, + _completed, + _remaining, + _request_id_of, +) from .retry import DEFAULT_RETRY, Retrier, RetryPolicy from .router_exceptions import RouterError @@ -257,6 +274,170 @@ 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 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 + 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. + """ + # 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=_remaining(deadline))) 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_best_effort() + 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), budget=_remaining(deadline)) + + 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 +488,110 @@ 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`. + """ + deadline = None if timeout is None else _now() + timeout + handle = await self.submit(model, arguments, idempotency_key=idempotency_key) + 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: + 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. + + 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 ffd073e..d6363e8 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,78 @@ 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 + # 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. + queue_canceled: bool = False + # --- counters the tests assert on --- upload_count: int = 0 from_hash_count: int = 0 @@ -221,6 +295,26 @@ 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 + # 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. + 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: @@ -380,6 +474,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)) @@ -484,6 +590,23 @@ def frame(event: str, data: dict) -> None: frame("status", {"status": state.terminal_status}) # -- 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 + # 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() @@ -508,6 +631,10 @@ 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"/api/v2/jobs/([^/]+)/cancel$", self.path) if m: self._json(200, _job_json(m.group(1), "canceling")) @@ -532,6 +659,121 @@ 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_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 + 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_raw is not None: + self._json(200, state.queue_result_raw) + return + 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 _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) + 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..04720e1 --- /dev/null +++ b/tests/test_models_queue.py @@ -0,0 +1,830 @@ +"""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_best_effort", _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() + + +# --- 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() + + # `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(delay) + + 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 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: