From d0f22dc80688bdc91321c3c0efc71efbb6c3195e Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Wed, 26 Aug 2026 21:17:11 -0700 Subject: [PATCH 1/2] feat: collect a Router deadline 504 under the same Idempotency-Key by default MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The router contract blesses exactly one same-key resend and the SDK was not making it. `spec/router-openapi.yaml`'s `deadline_exceeded` bucket says to "retry it with the SAME `Idempotency-Key`: when the provider had already accepted the generation, the retry collects that generation rather than dispatching another", and pins the `Retry-After` it carries to "seconds to wait before retrying the SAME request with the SAME `Idempotency-Key`" — sent "only when Comfy holds a handle to a generation the provider is still running". The default policy sorted that 504 with every other 5xx as an unknown outcome and declined it, so the designed replay path went unused and a caller got an error for a generation they are billed for either way. Adds a fourth failure class, "collectable", ahead of the unknown-outcome one: - a `deadline_exceeded` 504 carrying `Retry-After`, and - a 409 carrying `Retry-After` — the idempotency layer's "still in progress", which is what the collect retry meets before the generation finishes. Both gates are load-bearing. The 504 must name its bucket, because `deadline_exceeded` shares that status with `provider_timeout` (where nothing blesses the resend) and a header-less 504 from an intermediary reads as exactly that. The pace must be present, because without a handle to collect there is nothing to collect and a resend would dispatch new work. A 409 with no pace stays the deterministic refusal it was. `is_collectable` is the predicate; `error_bucket_of` reads the bucket off either layer's name for it (`error_type` on a typed router error, `code` on the protocol `ApiError` this route raises today) so the rule is not a silent no-op on the route it was written for. `retry_possibly_in_flight` is untouched and keeps its meaning: the class no contract characterises — a 5xx that named no collectable pace, and the client-side timeout. The new class is the one where the server itself said the resend is safe. Raises the default `max_elapsed` from 60s to 600s, one server deadline window and the same number as `MODEL_RUN_TIMEOUT`. A collect loop has to be able to outlast the deadline it is collecting after; 60s cannot. The cost lands on the other classes and is stated in the docs: an unreachable server now spends up to ten minutes connecting and backing off before it raises. `RetryPolicy(max_elapsed=60.0)` restores the old bound. No feature flag: `retry_collectable` is an ordinary `RetryPolicy` field (declared last so it cannot change what a positional `RetryPolicy(...)` means), and `Comfy(retry=NO_RETRY)` still disables retrying entirely. Corrects the `models.run()` docstring, which stated the v2 jobs API's single-use, reject-on-duplicate rule as though it governed this route. It is scoped to `submit()` now, and the router's replay-on-duplicate behaviour — same body collects, different body is rejected 422 — stated in its place. The stub server grows `model_run_collects_after_deadline`, the narrower carry the router describes (only the 504's reservation survives), kept distinct from `model_run_replays_idempotency_key`, the whole-deployment replay the opt-in is for — so the new tests assert against the contract rather than against a stub made permissive to fit them. --- README.md | 62 ++++--- src/comfy_sdk/models.py | 25 ++- src/comfy_sdk/retry.py | 258 ++++++++++++++++++++++------- src/comfy_sdk/router_exceptions.py | 31 ++-- tests/conftest.py | 31 +++- tests/test_models_run_retry.py | 249 ++++++++++++++++++++++++++-- tests/test_router_exceptions.py | 11 +- 7 files changed, 543 insertions(+), 124 deletions(-) diff --git a/README.md b/README.md index c33814a..ba91496 100644 --- a/README.md +++ b/README.md @@ -413,15 +413,19 @@ sends the same `Idempotency-Key`**, and a new call mints a new one — that is what lets a server tell a retry apart from a second order, on a surface where one call is a billed generation. -That one key is also what decides *which* failures are worth retrying. The API -contract makes `Idempotency-Key` **single-use, reject-on-duplicate, with no -response replay**: the first request to present a key is processed, and a later -one presenting the same key is rejected `422 idempotency_key_reuse` rather than -re-run. The contract also says when the key is released instead of claimed — a -request that definitively failed without starting work frees its key, while one -whose outcome the server could not characterise (a 5xx, an upstream timeout) -keeps it. Retrying under one key is only safe where that key is still unspent, -so that is exactly what the default policy retries. +That one key is also what decides *which* failures are worth retrying. The v2 +jobs contract makes its `Idempotency-Key` **single-use, reject-on-duplicate, +with no response replay**: the first request to present a key is processed, and +a later one presenting the same key is rejected `422 idempotency_key_reuse` +rather than re-run. (That rule governs `submit()`. On the router surface a +resend of the same key with the same body *collects* the generation the first +request started instead of dispatching another; a resend with a different body +is the one that is rejected `422`.) The contract also says when the key is +released instead of claimed — a request that definitively failed without +starting work frees its key, while one whose outcome the server could not +characterise (a 5xx, an upstream timeout) keeps it. Retrying under one key is +only safe where that key is still spendable, so that is exactly what the default +policy retries. The default policy: @@ -429,10 +433,11 @@ The default policy: |---|---| | Retried | connect-phase transport failures (connection refused, connect timeout, no pooled connection, proxy error) — the request never reached the server, so the key was never claimed | | Retried, at the server's pace | a `429` carrying `Retry-After` (queue full, out of credits, a concurrency limit) — a reject that started no work, so the key is released. The delay is the one the server named, not a guess | -| Not retried | every other 4xx — `400`/`content_policy_violation`, `404`, `409`, `422`, `401`, `402` — because asking again cannot change a deterministic refusal. A `429` with no `Retry-After` is not asking to be asked again either | -| Not retried by default | anything whose outcome is unknown: a **5xx response** (including the router's `service_unavailable` `503`, which asks a caller to retry with backoff but says nothing about the key), and a client-side timeout where the server may still be generating. The key stays claimed for these, so a same-key retry comes back `422 idempotency_key_reuse` and hides the real error — while a fresh-key retry is the second billed generation the one-key rule exists to prevent | -| Budget | 60 seconds of **total elapsed time** from the first attempt, not a number of attempts | -| Backoff | 0.5s doubling to a 15s ceiling, with full jitter (each wait is drawn from `[0, ceiling]`), clamped to whatever is left of the budget | +| Retried, at the server's pace | the answers that pace a resend of the *same* key for work already running: a `deadline_exceeded` `504` carrying `Retry-After` (Comfy stopped holding the connection at its own bound; the contract says to retry with the same key, which collects that generation rather than dispatching another), and an in-progress `409` carrying `Retry-After` (the same key, asked for again before the generation finished). One `run()` rides that loop to the finished result | +| Not retried | every other 4xx — `400`/`content_policy_violation`, `404`, a `409` that named no pace, `422`, `401`, `402` — because asking again cannot change a deterministic refusal. A `429` with no `Retry-After` is not asking to be asked again either, and neither is a `504` with none (the router sends it only when it holds a generation to collect) or a `504` that is `provider_timeout` rather than `deadline_exceeded` | +| Not retried by default | anything whose outcome is unknown: any **other 5xx response** (including the router's `service_unavailable` `503`, which asks a caller to retry with backoff but says nothing about the key), and a client-side timeout where the server may still be generating. The key stays claimed for these, so a same-key retry comes back `422 idempotency_key_reuse` and hides the real error — while a fresh-key retry is the second billed generation the one-key rule exists to prevent | +| Budget | 600 seconds of **total elapsed time** from the first attempt, not a number of attempts — one server deadline window, so a collect loop can outlast the deadline that started it | +| Backoff | 0.5s doubling to a 15s ceiling, with full jitter (each wait is drawn from `[0, ceiling]`), clamped to whatever is left of the budget. A `Retry-After` the server named is used as given instead | The budget bounds when the *last* attempt may **start**; an attempt already running is never interrupted by it, so a slow generation is never abandoned @@ -444,26 +449,33 @@ Tune or disable it per client: ```python from comfy_sdk import Comfy, NO_RETRY, RetryPolicy -Comfy(retry=NO_RETRY) # exactly one attempt, ever -Comfy(retry=RetryPolicy(max_elapsed=300.0)) # keep trying for five minutes +Comfy(retry=NO_RETRY) # exactly one attempt, ever +Comfy(retry=RetryPolicy(max_elapsed=60.0)) # give up after a minute +Comfy(retry=RetryPolicy(retry_collectable=False)) # raise the 504/409 instead -client.models.retry # the policy in force, read-only +client.models.retry # the policy in force, read-only ``` -5xx responses and client-side timeouts are the cases left out by default, and -for the same reason. `run` holds the connection open while the server -generates, so neither one tells you whether the generation happened — retrying -either starts a second generation unless the server replays the repeated key -rather than re-running it, and against a server that *rejects* it instead the -retry simply cannot succeed. Against a deployment that does replay, opt in: +The larger default budget is worth knowing about in the other direction too: a +genuinely unreachable server now spends up to ten minutes connecting and backing +off before it raises, where the old 60-second budget spent one. `max_elapsed` +buys that back. + +Other 5xx responses and client-side timeouts are the cases left out by default, +and for the same reason. `run` holds the connection open while the server +generates, so neither one tells you whether the generation happened, and no +contract says the key survives them — retrying either starts a second generation +unless the server replays the repeated key rather than re-running it, and +against a server that *rejects* it instead the retry simply cannot succeed. +Against a deployment that does replay, opt in: ```python Comfy(retry=RetryPolicy(max_elapsed=1200.0, retry_possibly_in_flight=True)) ``` -Raise `max_elapsed` when you do: one full-length client timeout on a run can -spend the default 60-second budget on its own, leaving no room for the retry -you just asked for. +Raise `max_elapsed` when you do: one full-length client timeout on a run spends +the whole default budget on its own, leaving no room for the retry you just +asked for. `retry` governs `client.models` only. `submit()`/`run()` on the client keep their own 429 handling, which follows the server's `Retry-After`. diff --git a/src/comfy_sdk/models.py b/src/comfy_sdk/models.py index 116db98..18c68cc 100644 --- a/src/comfy_sdk/models.py +++ b/src/comfy_sdk/models.py @@ -134,18 +134,29 @@ def run( An ``Idempotency-Key`` is sent on every run; a fresh one is minted per call unless ``idempotency_key`` is given, so an accidental exact resend - is the server's to reject rather than a second charged generation. + is the server's to deduplicate rather than a second charged generation. A failed attempt is retried under the client's policy, backed off with jitter and bounded by total elapsed time. **Every attempt of this one call reuses the one key**, which is what stops a retry from being billed as a second generation; calling ``run`` again is a new call and - mints a new key. Because the key is single-use and reject-on-duplicate, - only failures that leave it unclaimed are retried by default — - connect-phase failures, and a ``429`` that names a ``Retry-After``. A - completed 5xx or a client-side timeout leaves the outcome unknown, and - with it the key claimed, so those need - ``RetryPolicy(retry_possibly_in_flight=True)``. See + mints a new key. A key presented twice is *not* re-run: on the router + surface a resend of the same key with the same body collects the + generation the first request started rather than dispatching another, + and a resend with a *different* body is rejected ``422`` + ``idempotency_key_reuse`` — which is why the body is snapshotted before + the first attempt. (The single-use, reject-on-duplicate rule + ``spec/openapi.yaml`` states is the **v2 jobs API**'s, and governs + ``submit()``, not this route.) + + Retried by default: connect-phase failures, a ``429`` that names a + ``Retry-After``, and the answers that name a pace for collecting work + already running — a ``deadline_exceeded`` ``504`` and an in-progress + ``409``, each carrying ``Retry-After``. One ``run()`` can therefore ride + the collect loop through a server-side deadline to the finished + generation. Not retried by default: a completed 5xx that named no such + pace, and a client-side timeout — those leave the outcome genuinely + unknown and need ``RetryPolicy(retry_possibly_in_flight=True)``. See :mod:`comfy_sdk.retry`, and ``Comfy(retry=NO_RETRY)`` to switch it off. """ low = cast(ComfyLow, self._low) diff --git a/src/comfy_sdk/retry.py b/src/comfy_sdk/retry.py index 4d80e6a..309f737 100644 --- a/src/comfy_sdk/retry.py +++ b/src/comfy_sdk/retry.py @@ -15,24 +15,28 @@ unconditionally — there is no configuration that turns it off — because a retry without it is exactly the double-charge above. -**Only failures the key survives are retried by default.** Every documented -statement this repo makes about ``Idempotency-Key`` says it is *single-use, -reject-on-duplicate, with no response replay* — ``spec/openapi.yaml``'s shared -``IdempotencyKey`` parameter, :class:`~comfy_sdk.exceptions.IdempotencyKeyReuse` -and the README all agree that a later request presenting the same key is -rejected ``422`` ``idempotency_key_reuse`` rather than re-run or replayed. The -spec is equally explicit about when a key is *released* instead of claimed: a -request that "definitively fails without creating a job (a validation error, or -an upstream reject such as out-of-credits or queue-full)" frees it, while one -whose outcome the server cannot characterise ("an upstream timeout or 5xx where -the job may or may not have been created") keeps it claimed. - -``POST /models/run`` is not itself in that spec, so its key semantics are -undocumented — which is the point. Retrying under one key is only safe against a -server that *replays* a claimed key, and nothing here says this route does; the -alternatives are a ``422`` that replaces the real error, or a second billed -generation. So the SDK retries by default only where the key is provably still -unspent, and leaves the rest to an explicit opt-in. +**Only failures the key survives are retried by default.** The v2 jobs contract +in ``spec/openapi.yaml`` makes its shared ``IdempotencyKey`` parameter +*single-use, reject-on-duplicate, with no response replay*: the first request to +present a key is processed and any later one presenting it is rejected ``422`` +``idempotency_key_reuse`` rather than re-run or replayed (see +:class:`~comfy_sdk.exceptions.IdempotencyKeyReuse`). That spec is equally +explicit about when a key is *released* instead of claimed: a request that +"definitively fails without creating a job (a validation error, or an upstream +reject such as out-of-credits or queue-full)" frees it, while one whose outcome +the server cannot characterise ("an upstream timeout or 5xx where the job may or +may not have been created") keeps it claimed. + +``POST /models/run`` is not itself in that spec, so its key semantics are not +that spec's to state — and for one failure the *router* contract states them +directly. ``spec/router-openapi.yaml``'s ``deadline_exceeded`` bucket says to +"retry it with the SAME ``Idempotency-Key``: when the provider had already +accepted the generation, the retry collects that generation rather than +dispatching another", and pins the ``Retry-After`` it carries to "seconds to +wait before retrying the SAME request with the SAME ``Idempotency-Key``". So +the sorting rule is unchanged — retry where the one key is documented to be +spendable again — and one class is added to it, for the case where the server +has already said so on the wire. That, not a guess about the network, is what sorts the failures: @@ -46,16 +50,33 @@ out-of-credits, a concurrency limit) and tells clients to treat "429 + ``Retry-After``" as back-off-and-retry, so it is retried by default at the pace the server asked for rather than a blind backoff of our own. -3. **Outcome unknown** — a completed 5xx, or a transport failure that may have - delivered the request in full (a read timeout on a run is the important - member: the generation-sized client timeout expired with no answer, which is - precisely when the server is most likely still generating). This is the - class the spec keeps the key claimed for, so a same-key retry here is a - ``422`` that replaces the real error — and a *fresh*-key retry is the second - billed generation this module exists to prevent. Not retried by default. - :attr:`RetryPolicy.retry_possibly_in_flight` opts in, and it is correct - exactly when a deployment replays a repeated key instead of rejecting it. -4. **Everything else** — every other 4xx is the server's considered answer +3. **Collectable** — the server answered that the work it already holds is not + finished, *and* named the pace at which to ask the same key again for it: a + router ``deadline_exceeded`` ``504`` carrying ``Retry-After``, and a ``409`` + carrying ``Retry-After`` ("still in progress") from the idempotency layer on + the retry that follows it. This is the one class where the *server* has + stated the same-key resend is safe, and the pace it names is its own poll + interval — so it is retried by default, at that pace, and one ``run()`` call + rides the collect loop to the finished generation instead of handing the + caller a ``504`` for work that is still running. + :func:`is_collectable` is the predicate and + :attr:`RetryPolicy.retry_collectable` switches it off. Two gates keep it + narrow: the ``504`` must name the ``deadline_exceeded`` bucket (a bucket-less + ``504`` reads as ``provider_timeout``, where no contract blesses the resend), + and the ``Retry-After`` must be there (the router sends it only when it holds + a handle to a generation to collect — absent it, there is nothing to collect + and the ``504`` falls back to class 4). +4. **Outcome unknown** — a completed 5xx that named no collectable pace, or a + transport failure that may have delivered the request in full (a read timeout + on a run is the important member: the generation-sized client timeout expired + with no answer, which is precisely when the server is most likely still + generating). Nothing on the wire says the key survives this, so a same-key + retry may be a ``422`` that replaces the real error — and a *fresh*-key retry + is the second billed generation this module exists to prevent. Not retried by + default. :attr:`RetryPolicy.retry_possibly_in_flight` opts in, and it is + correct exactly when a deployment replays a repeated key instead of rejecting + it. +5. **Everything else** — every other 4xx is the server's considered answer about *this* request, and asking again spends money to be refused again. Never retried. @@ -67,16 +88,17 @@ class the spec keeps the key claimed for, so a same-key retry here is a Neither vendored contract says a ``503`` releases the key; the v2 contract says the opposite for the whole 5xx class ("an upstream timeout or 5xx where the job may or may not have been created" keeps it claimed), and the router spec -documents a same-key retry for exactly one bucket, ``deadline_exceeded``, and is -silent about this one. Retrying it by default would therefore trade a +documents a same-key retry for exactly one bucket, ``deadline_exceeded`` — which +is why that one is class 3 above and this one is not. The spec is silent about +what a ``503`` does to the key. Retrying it by default would therefore trade a diagnosable ``503`` for a ``422 idempotency_key_reuse`` on every deployment that -rejects a repeated key. So it stays in class 3 above, where +rejects a repeated key. So it stays in class 4 above, where :attr:`RetryPolicy.retry_possibly_in_flight` opts in — and that opt-in is the route to the contract's advice, because it keeps the one key across the retry. Know what the opt-in costs, though: it is a property of the *policy*, not of one bucket, so switching it on to get the blessed ``503`` retry also opts into -retrying every other completed 5xx and the client-side read timeout — class 3 +retrying every other completed 5xx and the client-side read timeout — class 4 entire, including the case this module calls the dangerous one. There is no per-bucket switch, deliberately: which failures a deployment's key survives is a fact about the deployment, not about the bucket, and a policy that claimed @@ -124,13 +146,28 @@ class the spec keeps the key claimed for, so a same-key retry here is a *last* attempt may start, so the worst case is that bound plus one per-attempt timeout, and the number of attempts falls out of the backoff schedule. +**The default budget is one server deadline window**, ten minutes — the same +number as :data:`~comfy_low.transport.MODEL_RUN_TIMEOUT`, deliberately, because +that is how long this surface is already willing to wait for one attempt. A +collect loop that outlives the deadline it is collecting after is the whole +point of class 3: the ``504`` says the server stopped holding *this* connection +at its own bound while the generation ran on, so a budget shorter than that bound +gives up mid-generation and hands the caller an error for work it will still be +charged for. The old 60-second budget was sized for the fast classes alone (a +connect failure, a paced ``429``) and could not outlast a single deadline +window. The cost of the larger default is paid in the *other* classes, and it is +worth naming: a genuinely unreachable server now spends up to ten minutes in +connect-and-back-off before it raises, where before it spent one. Pass +``RetryPolicy(max_elapsed=60.0)`` to get the old bound back. + Retry is on by default. ``Comfy(retry=NO_RETRY)`` turns it off; any other policy is a :class:`RetryPolicy` you construct:: from comfy_sdk import Comfy, NO_RETRY, RetryPolicy Comfy(retry=NO_RETRY) # exactly one attempt - Comfy(retry=RetryPolicy(max_elapsed=300.0)) # keep trying for 5 min + Comfy(retry=RetryPolicy(max_elapsed=60.0)) # give up after a minute + Comfy(retry=RetryPolicy(retry_collectable=False)) # no 504/409 collect loop """ from __future__ import annotations @@ -168,9 +205,24 @@ class the spec keeps the key claimed for, so a same-key retry here is a httpx.RemoteProtocolError, ) -#: The one 4xx another attempt can change, and only when it carries a pace. +#: The one 4xx another attempt can change on its own, and only when it carries +#: a pace. _TOO_MANY_REQUESTS = 429 +#: The 4xx the idempotency layer answers while the work the key already names is +#: still running. Retryable only when it carries a pace — see :func:`is_collectable`. +_CONFLICT = 409 + +#: Shared by the router's ``deadline_exceeded`` and ``provider_timeout`` buckets, +#: which is why the collect rule keys on the bucket and not on this. +_GATEWAY_TIMEOUT = 504 + +#: The one router bucket whose contract blesses a same-key resend: "retry it with +#: the SAME ``Idempotency-Key``: when the provider had already accepted the +#: generation, the retry collects that generation rather than dispatching +#: another" (``spec/router-openapi.yaml``). +_DEADLINE_EXCEEDED = "deadline_exceeded" + #: Policy fields that must be real numbers for the arithmetic below to mean #: anything. Kept beside the fields themselves so a numeric one added later is #: added here too. @@ -193,6 +245,13 @@ def is_unknown_outcome_status(status: int) -> bool: A ``502``/``504`` from an intermediary belongs here for the same reason: the proxy's response completed, which says nothing about whether the origin behind it stopped generating. + + This answers the *status* only, so it is still true of the one 5xx the + router contract does characterise — a ``deadline_exceeded`` ``504`` naming a + pace. :meth:`RetryPolicy.should_retry` consults :func:`is_collectable` + first, which is where that response stops being "unknown"; the outcome of + the *request* is genuinely still unknown there, and what the contract adds + is that asking again with the same key is how you find it out. """ return 500 <= status <= 599 @@ -218,6 +277,62 @@ def retry_after_of(exc: BaseException) -> float | None: return seconds +def error_bucket_of(exc: BaseException) -> str | None: + """The failure bucket the server named, or ``None`` if it named none. + + Read by attribute rather than by type, for the same reason + :func:`retry_after_of` is: one failure reaches this module modelled by two + different layers. A typed router error carries the wire ``error_type`` + (``RouterError.error_type``), while ``POST /models/run`` today raises the + protocol :class:`~comfy_low.errors.ApiError`, whose envelope names the same + thing ``code``. Reading only ``error_type`` would make every bucket-keyed + rule below unreachable on the route those rules were written for — a retry + that is a silent no-op with no test failing, which is the failure mode this + module has already been bitten by twice. + + Both attribute names are namespaced enough for that to be safe: the buckets + the rules below key on do not exist as v2 envelope codes, so a ``code`` that + reads as a router bucket *is* one. + """ + for attribute in ("error_type", "code"): + raw = getattr(exc, attribute, None) + if isinstance(raw, str) and raw.strip(): + return raw.strip() + return None + + +def is_collectable(exc: BaseException) -> bool: + """Whether the server said this failure's own work can be collected. + + True for the two answers that mean "the work your key already names is not + finished; ask again for it", each of which the server pairs with the pace to + ask at: + + * a router ``deadline_exceeded`` ``504`` carrying ``Retry-After`` — Comfy + stopped holding the connection at its own bound while the generation ran + on, and the contract says to "retry it with the SAME ``Idempotency-Key``", + which "collects that generation rather than dispatching another"; + * a ``409`` carrying ``Retry-After`` — the idempotency layer's answer that + the request under this key is still in progress, which is what the collect + retry above meets when it arrives before the generation finishes. + + Both gates are load-bearing. The ``504`` must name its bucket because + ``deadline_exceeded`` shares that status with ``provider_timeout``, where no + contract blesses the resend and a header-less ``504`` from an intermediary is + read as exactly that. The ``Retry-After`` must be present because the router + sends it on a ``deadline_exceeded`` "only when Comfy holds a handle to a + generation the provider is still running" — without it there is nothing to + collect, and a resend would be dispatching new work rather than gathering + old. A ``409`` with no pace is an ordinary conflict and stays a refusal. + """ + status = getattr(exc, "http_status", None) + if not isinstance(status, int) or retry_after_of(exc) is None: + return False + if status == _CONFLICT: + return True + return status == _GATEWAY_TIMEOUT and error_bucket_of(exc) == _DEADLINE_EXCEEDED + + @dataclass(frozen=True) class RetryPolicy: """When and how often to retry one logical call. @@ -230,8 +345,12 @@ class RetryPolicy: #: Seconds from the first attempt after which no *new* attempt is started. #: An attempt already running is never interrupted by it, so the worst-case #: wall clock for a call is this plus one per-attempt ``timeout``. Zero - #: disables retrying entirely (see :data:`NO_RETRY`). - max_elapsed: float = 60.0 + #: disables retrying entirely (see :data:`NO_RETRY`). The default is one + #: server deadline window — the same ten minutes as + #: :data:`~comfy_low.transport.MODEL_RUN_TIMEOUT` — so that a collect loop + #: after a ``deadline_exceeded`` ``504`` can outlast the bound that produced + #: it; see this module's docstring for what that costs the other classes. + max_elapsed: float = 600.0 #: Ceiling on the delay before the first retry. With ``jitter`` on — the #: default — the actual delay is drawn from ``[0, this]``. initial_backoff: float = 0.5 @@ -247,18 +366,31 @@ class RetryPolicy: #: together would otherwise retry together, converting one outage into a #: synchronised thundering herd on the recovering server. jitter: bool = True - #: Also retry when the request's outcome is unknown — a completed 5xx, or a - #: transport failure that may have delivered the request, so the server may - #: still be generating. Off by default, and the reason is the contract - #: rather than caution: every documented statement about - #: ``Idempotency-Key`` makes it single-use with no replay, and the spec - #: keeps it *claimed* across exactly this class of failure — so a same-key - #: retry here comes back ``422`` ``idempotency_key_reuse`` and hides the - #: real error. Turn it on for a deployment that replays a repeated key - #: instead of rejecting it — and raise ``max_elapsed`` when you do, since - #: one full-length client timeout on a run can spend the default budget on + #: Also retry when the request's outcome is *unknown* — a completed 5xx that + #: named no collectable pace, or a transport failure that may have delivered + #: the request, so the server may still be generating. This is the class no + #: contract characterises, which is what keeps it distinct from + #: :attr:`retry_collectable`: there the server named the resend safe, here + #: nothing did. Off by default, and the reason is the contract rather than + #: caution: the v2 jobs contract makes ``Idempotency-Key`` single-use with no + #: replay and keeps it *claimed* across exactly this class of failure — so a + #: same-key retry here can come back ``422`` ``idempotency_key_reuse`` and + #: hide the real error. Turn it on for a deployment that replays a repeated + #: key instead of rejecting it — and raise ``max_elapsed`` when you do, since + #: one full-length client timeout on a run spends the whole default budget on #: its own. retry_possibly_in_flight: bool = False + #: Retry the failures the server itself paced for a same-key resend — a + #: router ``deadline_exceeded`` ``504`` and an in-progress ``409``, each + #: carrying ``Retry-After``. See :func:`is_collectable` for the exact gates. + #: **On by default**, because this is the one class where the contract says + #: the resend collects the generation already running rather than + #: dispatching a second one, and the ``Retry-After`` is the server's own + #: poll interval. Switch it off to have those answers raised to the caller + #: instead — the whole retry surface goes away with ``NO_RETRY``. It is + #: declared last, out of reading order, so that inserting it cannot change + #: what an existing positional ``RetryPolicy(...)`` means. + retry_collectable: bool = True def __post_init__(self) -> None: for name in _NUMERIC_FIELDS: @@ -289,10 +421,12 @@ def should_retry(self, exc: BaseException) -> bool: An exception carrying an ``http_status`` is a response the server actually sent, whatever layer modelled it (the protocol ``ApiError``, - an SDK exception, a router one) — so the status decides. Everything - else is a transport failure, decided by whether the request can still - be executing server-side. See this module's docstring for the four - classes and why the key's contract, not the network, sorts them. + an SDK exception, a router one) — so the status decides, with the + failure bucket breaking the one tie a status cannot (the two ``504`` + buckets). Everything else is a transport failure, decided by whether the + request can still be executing server-side. See this module's docstring + for the five classes and why the key's contract, not the network, sorts + them. """ status = getattr(exc, "http_status", None) if isinstance(status, int): @@ -302,14 +436,24 @@ def should_retry(self, exc: BaseException) -> bool: # attached: the spec's retry signal is "429 + Retry-After", # and a 429 without one is not asking to be asked again. return retry_after_of(exc) is not None + if self.retry_collectable and is_collectable(exc): + # Checked before both branches below, because it overrides + # both: a `deadline_exceeded` 504 is a 5xx the router contract + # nevertheless blesses a same-key resend for, and an + # in-progress 409 is a 4xx that does become true on a later + # ask. Everything narrowing it to those two answers lives in + # `is_collectable`. + return True if is_unknown_outcome_status(status): - # Every 5xx, including the router's `service_unavailable` 503: - # the bucket says the condition clears on its own, but nothing - # says the Idempotency-Key does. See this module's docstring. + # Every other 5xx, including the router's `service_unavailable` + # 503: the bucket says the condition clears on its own, but + # nothing says the Idempotency-Key does. See this module's + # docstring. return self.retry_possibly_in_flight - # Every other 4xx is deterministic — 404 (no such model), 409, 422 - # (invalid input), a content-policy refusal — and none of them - # become true on the second ask. + # Every other 4xx is deterministic — 404 (no such model), a + # conflict that named no pace, 422 (invalid input), a + # content-policy refusal — and none of them become true on the + # second ask. return False if isinstance(exc, _NEVER_DELIVERED): return True @@ -418,6 +562,8 @@ def delay_before_retry(self, exc: BaseException) -> float | None: "NO_RETRY", "Retrier", "RetryPolicy", + "error_bucket_of", + "is_collectable", "is_unknown_outcome_status", "retry_after_of", ] diff --git a/src/comfy_sdk/router_exceptions.py b/src/comfy_sdk/router_exceptions.py index 26941aa..1453647 100644 --- a/src/comfy_sdk/router_exceptions.py +++ b/src/comfy_sdk/router_exceptions.py @@ -92,8 +92,9 @@ class docstrings below reproduce. ``tests/test_router_spec_contract.py`` reads #: :class:`ConcurrencyLimitExceeded`) and on a ``deadline_exceeded`` ``504``. #: It is read here rather than dropped because #: :func:`comfy_sdk.retry.retry_after_of` is what -#: :meth:`~comfy_sdk.retry.RetryPolicy.should_retry` keys its ``429`` branch on: -#: a router error that arrived without this attribute would make that branch +#: :meth:`~comfy_sdk.retry.RetryPolicy.should_retry` keys both of its +#: default-on branches on -- the ``429`` and the collectable ``504``/``409``: +#: a router error that arrived without this attribute would make those branches #: unreachable and turn the retry the default policy is built to make into a #: silent no-op. RETRY_AFTER_HEADER = "Retry-After" @@ -281,18 +282,22 @@ class DeadlineExceeded(RouterError): generation, the retry collects that generation rather than dispatching another, and a ``Retry-After`` on the ``504`` says when to ask. - That is the *contract's* advice, and as with :class:`ServiceUnavailable` the - SDK does **not** follow it by default. The bucket arrives on a ``504``, so - :func:`comfy_sdk.retry.is_unknown_outcome_status` sorts it with every other - 5xx and the default policy declines the retry; - ``RetryPolicy(retry_possibly_in_flight=True)`` opts in, and does keep the - one key across it. The policy cannot special-case this bucket on the status - alone, either -- a ``504`` reaching the SDK without a bucket header is read - as :class:`ProviderTimeout`, where a same-key retry is *not* blessed. What - the opt-in does honour is the pace: ``error_from_response`` preserves the - ``Retry-After`` on the exception, and + That is the *contract's* advice, and unlike :class:`ServiceUnavailable` the + SDK **does** follow it by default -- this is the one bucket a contract says + the same key survives, so ``client.models.run`` makes the retry rather than + describing it. :func:`comfy_sdk.retry.is_collectable` is the gate and it + keys on this bucket, never on the status alone: a ``504`` reaching the SDK + without a bucket header is read as :class:`ProviderTimeout`, where a + same-key retry is *not* blessed. It also requires the ``Retry-After``, which + the router sends only when it holds a handle to a generation still running + -- without one there is nothing to collect, so a bare + ``deadline_exceeded`` ``504`` is raised to the caller like any other 5xx. + The pace is honoured as given: ``error_from_response`` preserves the header + on the exception, and :meth:`~comfy_sdk.retry.Retrier.delay_before_retry` prefers a named pace - over its own jittered backoff. + over its own jittered backoff. ``RetryPolicy(retry_collectable=False)`` + switches the behaviour off; ``Comfy(retry=NO_RETRY)`` switches off retrying + entirely. """ error_type = "deadline_exceeded" diff --git a/tests/conftest.py b/tests/conftest.py index 1a393e2..8b198d9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -114,9 +114,19 @@ class ServerState: # is the vendored contract (reject-on-duplicate, no replay), under which a # same-key retry after a 5xx can only come back 422. model_run_replays_idempotency_key: bool = False - # Seconds sent as Retry-After alongside `model_run_error`, when that error - # is the 429 the policy is allowed to pace itself against. `None` sends no - # header at all, which is the 429 the policy must *not* retry. + # The narrower carry the router contract describes for `deadline_exceeded`: + # the reservation survives the 504 with a handle to the generation stamped + # on it, so the SAME key presented again *collects* that generation instead + # of being rejected 422. Only the 504 behaves that way — every other 5xx + # still claims the key — which is what makes this distinct from + # `model_run_replays_idempotency_key`, the whole-deployment replay the + # `retry_possibly_in_flight` opt-in exists for. + model_run_collects_after_deadline: bool = False + # Seconds sent as Retry-After alongside `model_run_error` / + # `model_run_transient_error`, for the failures the policy is allowed to + # pace itself against (a 429, a `deadline_exceeded` 504, an in-progress + # 409). `None` sends no header at all, which is the same failure the policy + # must *not* retry. model_run_retry_after: str | None = None # --- counters the tests assert on --- @@ -467,10 +477,17 @@ def _post_model_run(self) -> None: def claim_if_outcome_unknown(status: int) -> None: # The contract releases a key for a request that definitively # failed without starting work (a 4xx) and keeps it claimed - # when the outcome is unknown (a 5xx). A replaying deployment - # is the exception, and is what the opt-in is for. - if key and status >= 500 and not state.model_run_replays_idempotency_key: - state.model_run_idempotency[key] = "claimed" + # when the outcome is unknown (a 5xx). Two exceptions: a + # replaying deployment (what the opt-in is for), and the + # router's `deadline_exceeded` 504, whose reservation survives + # so the same key can collect the generation still running. + if not key or status < 500: + return + if state.model_run_replays_idempotency_key: + return + if state.model_run_collects_after_deadline and status == 504: + return + state.model_run_idempotency[key] = "claimed" def fail(status: int, code: str, message: str) -> None: claim_if_outcome_unknown(status) diff --git a/tests/test_models_run_retry.py b/tests/test_models_run_retry.py index 3a8edbf..235d352 100644 --- a/tests/test_models_run_retry.py +++ b/tests/test_models_run_retry.py @@ -8,21 +8,27 @@ Everything else here is the policy that decides *whether* to make that second attempt, and the thing that decides it is the key's own contract rather than a -guess about the network. ``spec/openapi.yaml`` makes ``Idempotency-Key`` -single-use, reject-on-duplicate, with no response replay, and says which -failures release the key (a definitive reject that started no work) and which -keep it claimed (a 5xx or upstream timeout, where the outcome is unknown). So -the default policy retries only what the one key survives — a connect-phase -failure, and a ``429`` that names its own ``Retry-After`` — and everything in -the unknown-outcome class sits behind ``retry_possibly_in_flight``, for the -deployment that replays a repeated key instead of rejecting it. +guess about the network. ``spec/openapi.yaml`` makes the v2 jobs API's +``Idempotency-Key`` single-use, reject-on-duplicate, with no response replay, +and says which failures release the key (a definitive reject that started no +work) and which keep it claimed (a 5xx or upstream timeout, where the outcome is +unknown). ``spec/router-openapi.yaml`` states the one exception: a +``deadline_exceeded`` ``504`` carrying ``Retry-After`` is to be retried "with +the SAME ``Idempotency-Key``", which collects the generation already running. +So the default policy retries what the one key survives — a connect-phase +failure, a ``429`` that names its own ``Retry-After``, and the collectable +answers that name one (that ``504``, and the in-progress ``409`` a collect +attempt meets) — and everything in the unknown-outcome class sits behind +``retry_possibly_in_flight``, for the deployment that replays a repeated key +instead of rejecting it. The stub server in ``conftest.py`` drives the wire half and now enforces that same reject-on-duplicate rule (``model_run_replays_idempotency_key`` switches -it to the replaying deployment), so a retry design the real server would reject -cannot pass here. The never-delivered class is asserted against ``_FlakyLow`` -instead: a stub HTTP server that is up cannot refuse a connection. The schedule -and deadline arithmetic is asserted directly against +it to the replaying deployment, ``model_run_collects_after_deadline`` to the +narrower carry the router describes for its own deadline), so a retry design the +real server would reject cannot pass here. The never-delivered class is asserted +against ``_FlakyLow`` instead: a stub HTTP server that is up cannot refuse a +connection. The schedule and deadline arithmetic is asserted directly against :class:`~comfy_sdk.retry.RetryPolicy` and :class:`~comfy_sdk.retry.Retrier`, where a fake clock makes it exact instead of timing-dependent. """ @@ -36,11 +42,26 @@ import httpx import pytest +from comfy_low.errors import ApiError +from comfy_low.transport import MODEL_RUN_TIMEOUT from comfy_sdk import DEFAULT_RETRY, NO_RETRY, AsyncComfy, Comfy, RetryPolicy from comfy_sdk.exceptions import ComfyError, IdempotencyKeyReuse from comfy_sdk.models import AsyncModels, Models -from comfy_sdk.retry import Retrier, is_unknown_outcome_status, retry_after_of -from comfy_sdk.router_exceptions import ContentPolicyViolation, InternalError, ServiceUnavailable +from comfy_sdk.retry import ( + Retrier, + error_bucket_of, + is_collectable, + is_unknown_outcome_status, + retry_after_of, +) +from comfy_sdk.router_exceptions import ( + ContentPolicyViolation, + DeadlineExceeded, + InternalError, + ProviderTimeout, + RouterError, + ServiceUnavailable, +) class _FakeClock: @@ -726,3 +747,203 @@ def test_a_local_failure_is_never_retryable(exc: Exception) -> None: opted_in = RetryPolicy(retry_possibly_in_flight=True) assert not DEFAULT_RETRY.should_retry(exc) assert not opted_in.should_retry(exc) + + +# --- the collect loop: the failures the server paced for a same-key resend --- +# +# The router contract blesses exactly one same-key resend: `deadline_exceeded`, +# where "the retry collects that generation rather than dispatching another" +# and the `Retry-After` on the 504 "says when to ask". These assert that the +# default policy makes that retry, that both gates on it hold (the bucket and +# the pace), and that the 409 the idempotency layer answers mid-collect is +# waited out the same way -- so one `run()` rides the loop to the result. + + +def test_a_deadline_504_that_names_a_pace_is_collected_under_the_one_key(server) -> None: + server.state.model_run_collects_after_deadline = True + server.state.model_run_transient_error = (504, "deadline_exceeded") + server.state.model_run_retry_after = "0" + server.state.model_run_fail_times = 1 + with Comfy(retry=FAST) as client: + assert client.models.run(MODEL, ARGS) == server.state.model_run_result + assert server.state.model_run_count == 2 + # Spelled out rather than `len(set(...)) == 1`, which is also true of + # `[None, None]`: the property is that a key was sent AND that the collect + # attempt presented the SAME one, because that is what makes it a + # collection of the first generation rather than a second order. + first, second = server.state.model_run_idempotency_keys + assert first is not None + assert second == first + + +async def test_a_deadline_504_is_collected_on_the_async_client(server) -> None: + server.state.model_run_collects_after_deadline = True + server.state.model_run_transient_error = (504, "deadline_exceeded") + server.state.model_run_retry_after = "0" + server.state.model_run_fail_times = 1 + async with AsyncComfy(retry=FAST) as client: + assert await client.models.run(MODEL, ARGS) == server.state.model_run_result + assert server.state.model_run_count == 2 + first, second = server.state.model_run_idempotency_keys + assert first is not None + assert second == first + + +def test_a_deadline_504_without_a_pace_is_not_retried(server) -> None: + # The router sends `Retry-After` on this bucket only when it holds a handle + # to a generation still running. Without one there is nothing to collect, + # so the resend would be dispatching new work -- back to the unknown-outcome + # class, where the caller sees the 504. + server.state.model_run_error = (504, "deadline_exceeded") + with Comfy(retry=FAST) as client: + with pytest.raises(ComfyError) as excinfo: + client.models.run(MODEL, ARGS) + assert server.state.model_run_count == 1 + assert excinfo.value.http_status == 504 + + +def test_a_504_that_is_not_the_deadline_bucket_is_not_collected(server) -> None: + # The gate that cannot be dropped: `deadline_exceeded` shares 504 with + # `provider_timeout`, and a header-less 504 from an intermediary reads as + # the latter. Only the bucket the contract blesses is resent, however + # helpfully the response is paced. + server.state.model_run_error = (504, "provider_timeout") + server.state.model_run_retry_after = "0" + with Comfy(retry=FAST) as client: + with pytest.raises(ComfyError) as excinfo: + client.models.run(MODEL, ARGS) + assert server.state.model_run_count == 1 + assert excinfo.value.http_status == 504 + + +def test_an_in_progress_409_that_names_a_pace_is_waited_out_and_retried(server) -> None: + # The answer the collect retry meets when it arrives before the generation + # has finished: the key is recognised, the work is still running, come back + # in `Retry-After`. One `run()` rides that to the 200. + server.state.model_run_transient_error = (409, "generation_in_progress") + server.state.model_run_retry_after = "0" + server.state.model_run_fail_times = 2 + with Comfy(retry=FAST) as client: + assert client.models.run(MODEL, ARGS) == server.state.model_run_result + assert server.state.model_run_count == 3 + keys = server.state.model_run_idempotency_keys + assert keys[0] is not None + assert keys == [keys[0]] * 3 + + +def test_a_409_that_names_no_pace_is_still_a_refusal(server) -> None: + # Unchanged: an ordinary conflict is deterministic. The pace is the whole + # signal that this one is "not yet" rather than "no". + server.state.model_run_error = (409, "generation_in_progress") + with Comfy(retry=FAST) as client: + with pytest.raises(ComfyError): + client.models.run(MODEL, ARGS) + assert server.state.model_run_count == 1 + + +def test_no_retry_disables_the_collect_loop(server) -> None: + server.state.model_run_collects_after_deadline = True + server.state.model_run_error = (504, "deadline_exceeded") + server.state.model_run_retry_after = "0" + with Comfy(retry=NO_RETRY) as client: + with pytest.raises(ComfyError): + client.models.run(MODEL, ARGS) + assert server.state.model_run_count == 1 + + +def test_the_collect_loop_can_be_switched_off_in_the_policy(server) -> None: + # No feature flag: the new default is a `RetryPolicy` knob like every other + # one, so a caller who wants the 504 raised keeps the rest of the policy. + server.state.model_run_collects_after_deadline = True + server.state.model_run_error = (504, "deadline_exceeded") + server.state.model_run_retry_after = "0" + policy = RetryPolicy( + max_elapsed=5.0, initial_backoff=0.01, max_backoff=0.02, retry_collectable=False + ) + with Comfy(retry=policy) as client: + with pytest.raises(ComfyError): + client.models.run(MODEL, ARGS) + assert server.state.model_run_count == 1 + + +def test_the_collect_loop_gives_up_at_the_elapsed_budget(server) -> None: + # A generation that outlives the budget still ends the call: `max_elapsed` + # bounds the loop however patiently the server keeps saying "not yet". + server.state.model_run_collects_after_deadline = True + server.state.model_run_error = (504, "deadline_exceeded") + server.state.model_run_retry_after = "1" + policy = RetryPolicy(max_elapsed=0.4, initial_backoff=0.01, max_backoff=0.02) + started = time.monotonic() + with Comfy(retry=policy) as client: + with pytest.raises(ComfyError) as excinfo: + client.models.run(MODEL, ARGS) + assert time.monotonic() - started < 5.0 + assert excinfo.value.http_status == 504 + assert server.state.model_run_count >= 2 + assert len(set(server.state.model_run_idempotency_keys)) == 1 + + +def test_the_budget_bounds_the_loop_even_at_the_server_named_pace() -> None: + # The same bound, exactly, against a fake clock: three full 30s waits fit + # in a 100s budget and the fourth is clamped to the 10s left rather than + # overshooting it, after which no further attempt starts. + clock = _FakeClock() + retrier = Retrier(RetryPolicy(max_elapsed=100.0), now=clock, rng=lambda: 1.0) + failure = DeadlineExceeded("still generating", http_status=504, retry_after=30) + delays = [] + while (delay := retrier.delay_before_retry(failure)) is not None: + delays.append(delay) + clock.advance(delay) + assert delays == [30.0, 30.0, 30.0, 10.0] + + +def test_the_default_budget_is_one_server_deadline_window() -> None: + # The decision this asserts, so it cannot drift silently: the budget is the + # same ten minutes `run` is already willing to spend on ONE attempt. A + # shorter one cannot outlast the deadline it is collecting after -- the 504 + # says the server stopped holding this connection at its own bound while + # the generation ran on, and the caller is billed for that generation + # whether or not the SDK waits to collect it. + assert DEFAULT_RETRY.max_elapsed == MODEL_RUN_TIMEOUT.read == 600.0 + + +@pytest.mark.parametrize( + "exc,expected", + [ + (DeadlineExceeded("deadline", http_status=504, retry_after=3), True), + (DeadlineExceeded("deadline", http_status=504), False), + (ProviderTimeout("upstream", http_status=504, retry_after=3), False), + (RouterError("in progress", http_status=409, retry_after=3), True), + (RouterError("conflict", http_status=409), False), + (ServiceUnavailable("later", http_status=503, retry_after=3), False), + (InternalError("boom", http_status=500, retry_after=3), False), + (httpx.ConnectError("refused"), False), + ], + ids=lambda v: str(v), +) +def test_what_counts_as_collectable(exc: BaseException, expected: bool) -> None: + assert is_collectable(exc) is expected + # And the default policy acts on exactly that, without the opt-in. + if expected: + assert DEFAULT_RETRY.should_retry(exc) + + +def test_the_bucket_is_read_from_either_layers_name_for_it() -> None: + # A typed router error carries `error_type`; the protocol `ApiError` that + # `POST /models/run` raises today names the same thing `code`. Reading only + # the first would make the whole collect rule unreachable on the route it + # was written for. + assert error_bucket_of(DeadlineExceeded("deadline", http_status=504)) == "deadline_exceeded" + assert error_bucket_of(ApiError("deadline", code="deadline_exceeded", http_status=504)) == ( + "deadline_exceeded" + ) + assert error_bucket_of(ApiError("boom", http_status=500)) == "error" + assert error_bucket_of(httpx.ConnectError("refused")) is None + + +def test_a_deadline_504_raised_as_a_protocol_error_is_still_collected() -> None: + # The wire test above proves it end to end; this pins the reason, which is + # that the bucket survives the layer boundary. `_CANDIDATE_FAILURES` exists + # for the same class of silent no-op. + exc = ApiError("deadline", code="deadline_exceeded", http_status=504, retry_after=2) + assert DEFAULT_RETRY.should_retry(exc) diff --git a/tests/test_router_exceptions.py b/tests/test_router_exceptions.py index e232080..99b17f3 100644 --- a/tests/test_router_exceptions.py +++ b/tests/test_router_exceptions.py @@ -487,12 +487,19 @@ def test_retry_after_is_preserved_on_the_throttled_buckets(error_type: str) -> N def test_retry_after_is_preserved_on_a_deadline_exceeded_504() -> None: - # The bucket whose docstring tells the caller a Retry-After says when to - # ask. Under `retry_possibly_in_flight=True` the policy honours it. + # The bucket whose contract tells the caller to retry with the SAME key and + # whose Retry-After says when. The default policy makes that retry, so the + # header being dropped here would silently disable it. + from comfy_sdk.retry import RetryPolicy + status, headers, body = stub_error_response("deadline_exceeded", 504) exc = error_from_response(status, {**headers, "Retry-After": "5"}, body) assert isinstance(exc, DeadlineExceeded) assert exc.retry_after == 5 + assert RetryPolicy().should_retry(exc) is True + # ...and without the pace there is no generation to collect, so it is not. + unpaced = error_from_response(status, headers, body) + assert RetryPolicy().should_retry(unpaced) is False def test_the_retry_after_header_is_matched_case_insensitively() -> None: From 8d0d883d47ad03a3404643eb622e5583a1ee47d2 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Wed, 26 Aug 2026 22:31:26 -0700 Subject: [PATCH 2/2] fix: close the gates the collect loop is only safe behind Review follow-up on the collect-by-default change. Eight findings from the review panel, and the first two are the ones that mattered: **The collect rule was a no-op against a real Router 504.** `POST /api/v2/models/run` is fronted by Router, whose error body is `{detail, error_type}` with the bucket repeated on `X-Comfy-Error-Type` -- not the v2 `{error: {code}}` envelope `error_from_envelope` read. Every real 504 therefore reached `is_collectable` with the status-derived `"error"` bucket and fell through, while every test passed because the stub only ever spoke the envelope. `error_from_envelope` now reads the header and the body's `error_type` -- but only after `_CODE_BY_STATUS`, so a Router 429 still maps to `queue_full` and nothing integrators already catch gets retyped; what is left is the 5xx range, where the status determines nothing. The stub grew a Router-shape mode and the tests to go with it. **`Retry-After: 0` was a busy loop.** `RouterRetryAfterHeader` is pinned to `minimum: 1`, so zero is a server answering outside its own contract -- and it was honoured verbatim, giving a zero-delay resend loop of full model-run POSTs for the whole budget, entirely server-controlled. Zero now names no usable pace and the jittered backoff answers instead. The header's *presence* still decides collectability: that is the server saying it holds a generation, which a nonsense value does not retract. The rest: - The 409 branch accepted any paced 409. Every 409 either vendored spec documents is a deterministic refusal -- and `spec/openapi.yaml` gives its `hash_mismatch` 409 a `Retry-After` outright -- so a permanent refusal or a bucket-less proxy conflict was resent for the whole budget. Gated on `generation_in_progress` now, fail-closed, the way the 504 is on its bucket. - The budget is split in two. `max_elapsed` goes back to 60s and governs the fast classes, so a blackholed host stops pinning a caller's thread for ten minutes to buy the collect loop its room. `collect_max_elapsed` is 1200s and governs the collect class alone -- a `deadline_exceeded` 504 arrives *at* the server's own 600s bound, so the previous one-window budget was already spent when it landed and the collect attempt it was sized for never started. - `retry_collectable=False` did not hold with `retry_possibly_in_flight=True`: the 504 fell through to the unknown-outcome branch and was retried anyway. Collectability is classified first now and answers `retry_collectable`. - The body snapshot was shallow, so a caller mutating a nested value during the retry window sent a different body under the one key -- the 422 the snapshot exists to prevent. Deep now. - The stub's key-claim carve-out keyed on status 504 alone, making it more permissive than the contract it stands in for; it takes the bucket too. - README's "not retried" row listed the two 504 cases that *are* retried under the in-flight opt-in. Coverage added for the non-collecting deployment the default now bets against: the resend comes back 422 `idempotency_key_reuse` in place of the real 504, and `retry_collectable=False` is the way out. That trade is stated in the README rather than left implicit. --- README.md | 41 +++-- src/comfy_low/errors.py | 44 +++++- src/comfy_low/transport.py | 11 +- src/comfy_sdk/models.py | 31 ++-- src/comfy_sdk/retry.py | 201 +++++++++++++++++-------- tests/conftest.py | 48 +++++- tests/test_error_mapping.py | 81 +++++++++- tests/test_models_run_retry.py | 267 +++++++++++++++++++++++++++++++-- 8 files changed, 615 insertions(+), 109 deletions(-) diff --git a/README.md b/README.md index ba91496..dba5df1 100644 --- a/README.md +++ b/README.md @@ -433,10 +433,10 @@ The default policy: |---|---| | Retried | connect-phase transport failures (connection refused, connect timeout, no pooled connection, proxy error) — the request never reached the server, so the key was never claimed | | Retried, at the server's pace | a `429` carrying `Retry-After` (queue full, out of credits, a concurrency limit) — a reject that started no work, so the key is released. The delay is the one the server named, not a guess | -| Retried, at the server's pace | the answers that pace a resend of the *same* key for work already running: a `deadline_exceeded` `504` carrying `Retry-After` (Comfy stopped holding the connection at its own bound; the contract says to retry with the same key, which collects that generation rather than dispatching another), and an in-progress `409` carrying `Retry-After` (the same key, asked for again before the generation finished). One `run()` rides that loop to the finished result | -| Not retried | every other 4xx — `400`/`content_policy_violation`, `404`, a `409` that named no pace, `422`, `401`, `402` — because asking again cannot change a deterministic refusal. A `429` with no `Retry-After` is not asking to be asked again either, and neither is a `504` with none (the router sends it only when it holds a generation to collect) or a `504` that is `provider_timeout` rather than `deadline_exceeded` | -| Not retried by default | anything whose outcome is unknown: any **other 5xx response** (including the router's `service_unavailable` `503`, which asks a caller to retry with backoff but says nothing about the key), and a client-side timeout where the server may still be generating. The key stays claimed for these, so a same-key retry comes back `422 idempotency_key_reuse` and hides the real error — while a fresh-key retry is the second billed generation the one-key rule exists to prevent | -| Budget | 600 seconds of **total elapsed time** from the first attempt, not a number of attempts — one server deadline window, so a collect loop can outlast the deadline that started it | +| Retried, at the server's pace | the answers that pace a resend of the *same* key for work already running: a `deadline_exceeded` `504` carrying `Retry-After` (Comfy stopped holding the connection at its own bound; the contract says to retry with the same key, which collects that generation rather than dispatching another), and a `generation_in_progress` `409` carrying `Retry-After` (the same key, asked for again before the generation finished). One `run()` rides that loop to the finished result | +| Not retried | every other 4xx — `400`/`content_policy_violation`, `404`, any `409` that is not the paced `generation_in_progress` one above (`hash_mismatch` carries a `Retry-After` and is still deterministic), `422`, `401`, `402` — because asking again cannot change a deterministic refusal. A `429` with no `Retry-After` is not asking to be asked again either | +| Not retried by default | anything whose outcome is unknown: any **other 5xx response** — including the router's `service_unavailable` `503` (which asks a caller to retry with backoff but says nothing about the key), a `504` carrying no `Retry-After` (the router sends it only when it holds a generation to collect), and a `504` that is `provider_timeout` rather than `deadline_exceeded` — and a client-side timeout where the server may still be generating. The key stays claimed for these, so a same-key retry comes back `422 idempotency_key_reuse` and hides the real error — while a fresh-key retry is the second billed generation the one-key rule exists to prevent | +| Budget | 60 seconds of **total elapsed time** from the first attempt, not a number of attempts. The collect loop gets its own, longer budget: 1200 seconds, two server deadline windows, so it can outlast the deadline that started it | | Backoff | 0.5s doubling to a 15s ceiling, with full jitter (each wait is drawn from `[0, ceiling]`), clamped to whatever is left of the budget. A `Retry-After` the server named is used as given instead | The budget bounds when the *last* attempt may **start**; an attempt already @@ -449,17 +449,29 @@ Tune or disable it per client: ```python from comfy_sdk import Comfy, NO_RETRY, RetryPolicy -Comfy(retry=NO_RETRY) # exactly one attempt, ever -Comfy(retry=RetryPolicy(max_elapsed=60.0)) # give up after a minute -Comfy(retry=RetryPolicy(retry_collectable=False)) # raise the 504/409 instead +Comfy(retry=NO_RETRY) # exactly one attempt, ever +Comfy(retry=RetryPolicy(max_elapsed=300.0)) # fast classes: five minutes +Comfy(retry=RetryPolicy(collect_max_elapsed=60.0)) # bound the collect loop +Comfy(retry=RetryPolicy(retry_collectable=False)) # raise the 504/409 instead -client.models.retry # the policy in force, read-only +client.models.retry # the policy in force, read-only ``` -The larger default budget is worth knowing about in the other direction too: a -genuinely unreachable server now spends up to ten minutes connecting and backing -off before it raises, where the old 60-second budget spent one. `max_elapsed` -buys that back. +There are two budgets because the classes have two shapes. `max_elapsed` governs +the fast ones — a connect failure, a paced `429` — which resolve in seconds or +not at all, so an unreachable host gives up in a minute rather than pinning a +caller (a whole thread, on the sync client) for longer. `collect_max_elapsed` +governs the collect loop alone, and is twenty minutes because that is the one +class that has to outlast a *server-side* bound: a `deadline_exceeded` `504` +arrives at Comfy's own ten-minute deadline, so a single-window budget would +already be spent when it lands and the collect attempt it exists for would never +start. Nothing else pays for that room. + +A note on what the default trades: `POST /models/run` is in neither vendored +spec, so a deployment may apply the v2 rule instead and keep the key claimed +across the `504`. There the collect resend comes back `422 +idempotency_key_reuse` in place of the real `504`. Set +`retry_collectable=False` on such a deployment. Other 5xx responses and client-side timeouts are the cases left out by default, and for the same reason. `run` holds the connection open while the server @@ -474,8 +486,9 @@ Comfy(retry=RetryPolicy(max_elapsed=1200.0, retry_possibly_in_flight=True)) ``` Raise `max_elapsed` when you do: one full-length client timeout on a run spends -the whole default budget on its own, leaving no room for the retry you just -asked for. +many times the default 60-second budget on its own, leaving no room for the +retry you just asked for. `collect_max_elapsed` does not help here — that budget +is the collect class's alone. `retry` governs `client.models` only. `submit()`/`run()` on the client keep their own 429 handling, which follows the server's `Retry-After`. diff --git a/src/comfy_low/errors.py b/src/comfy_low/errors.py index 8a3a12f..1736dcb 100644 --- a/src/comfy_low/errors.py +++ b/src/comfy_low/errors.py @@ -98,17 +98,50 @@ class Forbidden(ApiError): } +def _clean(value: Any) -> str | None: + """``value`` as a non-empty string, or ``None``. + + Anything else — a missing key, a number, Router's ``detail[]`` list form — + reads as absent rather than being coerced, so a malformed body degrades to + the status-derived default instead of producing a nonsense code. + """ + if not isinstance(value, str): + return None + stripped = value.strip() + return stripped or None + + def error_from_envelope( http_status: int, body: dict[str, Any] | None, *, retry_after: int | None = None, + error_type: str | None = None, ) -> ApiError: """Build the typed exception for an error response. Falls back to a status-derived code when the body is missing or not a well-formed envelope (so a bare ``401`` with no JSON still maps to ``Unauthorized``). + + Not every route answers in the envelope shape. ``POST /api/v2/models/run`` + is fronted by Router, whose error body is ``{detail, error_type}`` and which + repeats the same coarse bucket on the ``X-Comfy-Error-Type`` header + (``spec/router-openapi.yaml``). Reading only ``error["code"]`` would collapse + every one of those to the status-derived default — for a ``504`` that + default is the meaningless ``"error"``, and ``comfy_sdk.retry`` keys its + default-on collect rule on the bucket, so the rule would be a silent no-op + against every real Router ``504``. ``error_type`` (the header, passed by the + caller) and the body's own top-level ``error_type`` are read to prevent that. + + They are consulted in one narrow place, though: *after* the envelope's + ``code``, which always wins, and *after* :data:`_CODE_BY_STATUS`, which + already determines a code for every status where this API has a documented + one. So a Router ``429`` still maps to ``queue_full`` and a Router ``401`` + still to ``unauthorized`` — reordering those would silently retype + exceptions that integrators already catch. What is left is exactly the 5xx + range, where the status determines nothing and the bucket is the only name + the response has. """ err = (body or {}).get("error") if isinstance(body, dict) else None code = (err or {}).get("code") if isinstance(err, dict) else None @@ -116,7 +149,16 @@ def error_from_envelope( details = (err or {}).get("details") if isinstance(err, dict) else None if code is None: - code = _CODE_BY_STATUS.get(http_status, "error") + code = _CODE_BY_STATUS.get(http_status) + if code is None: + code = _clean(error_type) or _clean( + (body or {}).get("error_type") if isinstance(body, dict) else None + ) + if code is None: + code = "error" + if not message: + # Router names its human-readable string `detail`, not `error.message`. + message = _clean((body or {}).get("detail") if isinstance(body, dict) else None) if not message: message = f"HTTP {http_status}" diff --git a/src/comfy_low/transport.py b/src/comfy_low/transport.py index beb06e9..8a5e56c 100644 --- a/src/comfy_low/transport.py +++ b/src/comfy_low/transport.py @@ -220,7 +220,16 @@ def parse_or_raise(self, resp: httpx.Response, ok: tuple[int, ...]) -> dict[str, body = resp.json() except Exception: body = None - raise error_from_envelope(resp.status_code, body, retry_after=_retry_after(resp)) + raise error_from_envelope( + resp.status_code, + body, + retry_after=_retry_after(resp), + # Router repeats its coarse bucket on this header, and the model-run + # route answers in Router's body shape rather than the envelope's. + # Passing it here is what keeps the bucket alive across the layer + # boundary; see `error_from_envelope`. + error_type=resp.headers.get("X-Comfy-Error-Type"), + ) def parse_expiry(url: str) -> datetime | None: diff --git a/src/comfy_sdk/models.py b/src/comfy_sdk/models.py index 18c68cc..59f3f3b 100644 --- a/src/comfy_sdk/models.py +++ b/src/comfy_sdk/models.py @@ -35,6 +35,7 @@ import asyncio import time from collections.abc import Mapping +from copy import deepcopy from typing import Any, cast import httpx @@ -151,12 +152,13 @@ def run( Retried by default: connect-phase failures, a ``429`` that names a ``Retry-After``, and the answers that name a pace for collecting work - already running — a ``deadline_exceeded`` ``504`` and an in-progress - ``409``, each carrying ``Retry-After``. One ``run()`` can therefore ride - the collect loop through a server-side deadline to the finished - generation. Not retried by default: a completed 5xx that named no such - pace, and a client-side timeout — those leave the outcome genuinely - unknown and need ``RetryPolicy(retry_possibly_in_flight=True)``. See + already running — a ``deadline_exceeded`` ``504`` and a + ``generation_in_progress`` ``409``, each carrying ``Retry-After``. One + ``run()`` can therefore ride the collect loop through a server-side + deadline to the finished generation. Not retried by default: a completed + 5xx that named no such pace, and a client-side timeout — those leave the + outcome genuinely unknown and need + ``RetryPolicy(retry_possibly_in_flight=True)``. See :mod:`comfy_sdk.retry`, and ``Comfy(retry=NO_RETRY)`` to switch it off. """ low = cast(ComfyLow, self._low) @@ -167,7 +169,13 @@ def run( # mapping inside each attempt would let a mutation between attempts # send a different body under the *same* key, which is precisely the # same-key-different-body case the contract rejects outright. - payload = dict(arguments) + # + # Deep, not shallow: a shallow copy leaves every nested list and dict + # shared with the caller, so mutating `arguments["config"]["steps"]` + # during the retry window would still change the body under the one key + # and earn the 422 this snapshot exists to prevent. The body is JSON on + # the wire, so everything legal in it is deep-copyable. + payload = deepcopy(dict(arguments)) retrier = Retrier(self._retry, now=_now) with translating(): while True: @@ -203,10 +211,11 @@ async def run( """ low = cast(AsyncComfyLow, self._low) key = idempotency_key or new_idempotency_key() - # Snapshotted before the first attempt — see :meth:`Models.run`. The - # window is wider here: the caller's coroutine can mutate `arguments` - # while the retry sleeps. - payload = dict(arguments) + # Snapshotted deeply before the first attempt — see :meth:`Models.run`. + # The window is wider here: the retry sleeps inside the caller's own + # event loop, so another task is free to run and mutate `arguments`, + # nested values included. + payload = deepcopy(dict(arguments)) retrier = Retrier(self._retry, now=_now) with translating(): while True: diff --git a/src/comfy_sdk/retry.py b/src/comfy_sdk/retry.py index 309f737..5829a1b 100644 --- a/src/comfy_sdk/retry.py +++ b/src/comfy_sdk/retry.py @@ -52,20 +52,23 @@ pace the server asked for rather than a blind backoff of our own. 3. **Collectable** — the server answered that the work it already holds is not finished, *and* named the pace at which to ask the same key again for it: a - router ``deadline_exceeded`` ``504`` carrying ``Retry-After``, and a ``409`` - carrying ``Retry-After`` ("still in progress") from the idempotency layer on - the retry that follows it. This is the one class where the *server* has + router ``deadline_exceeded`` ``504`` carrying ``Retry-After``, and a + ``generation_in_progress`` ``409`` carrying ``Retry-After`` from the + idempotency layer on the retry that follows it. This is the one class where the *server* has stated the same-key resend is safe, and the pace it names is its own poll interval — so it is retried by default, at that pace, and one ``run()`` call rides the collect loop to the finished generation instead of handing the caller a ``504`` for work that is still running. :func:`is_collectable` is the predicate and - :attr:`RetryPolicy.retry_collectable` switches it off. Two gates keep it - narrow: the ``504`` must name the ``deadline_exceeded`` bucket (a bucket-less - ``504`` reads as ``provider_timeout``, where no contract blesses the resend), - and the ``Retry-After`` must be there (the router sends it only when it holds - a handle to a generation to collect — absent it, there is nothing to collect - and the ``504`` falls back to class 4). + :attr:`RetryPolicy.retry_collectable` switches it off, and + :attr:`RetryPolicy.collect_max_elapsed` bounds it. Two gates keep it narrow, + and both apply to each status: the response must name the bucket (a + bucket-less ``504`` reads as ``provider_timeout``, where no contract blesses + the resend; a bucket-less or ``hash_mismatch`` ``409`` is a deterministic + refusal that a pace does not soften), and the ``Retry-After`` must be there + (the router sends it only when it holds a handle to a generation to collect — + absent it, there is nothing to collect and the ``504`` falls back to class + 4). 4. **Outcome unknown** — a completed 5xx that named no collectable pace, or a transport failure that may have delivered the request in full (a read timeout on a run is the important member: the generation-sized client timeout expired @@ -142,23 +145,23 @@ The budget is **total elapsed time**, never an attempt count: a per-attempt budget multiplies out, and N attempts each allowed their own timeout stack into -a wait far longer than anything anybody chose. ``max_elapsed`` bounds when the -*last* attempt may start, so the worst case is that bound plus one per-attempt -timeout, and the number of attempts falls out of the backoff schedule. - -**The default budget is one server deadline window**, ten minutes — the same -number as :data:`~comfy_low.transport.MODEL_RUN_TIMEOUT`, deliberately, because -that is how long this surface is already willing to wait for one attempt. A -collect loop that outlives the deadline it is collecting after is the whole -point of class 3: the ``504`` says the server stopped holding *this* connection -at its own bound while the generation ran on, so a budget shorter than that bound -gives up mid-generation and hands the caller an error for work it will still be -charged for. The old 60-second budget was sized for the fast classes alone (a -connect failure, a paced ``429``) and could not outlast a single deadline -window. The cost of the larger default is paid in the *other* classes, and it is -worth naming: a genuinely unreachable server now spends up to ten minutes in -connect-and-back-off before it raises, where before it spent one. Pass -``RetryPolicy(max_elapsed=60.0)`` to get the old bound back. +a wait far longer than anything anybody chose. It bounds when the *last* attempt +may start, so the worst case is that bound plus one per-attempt timeout, and the +number of attempts falls out of the backoff schedule. + +**There are two budgets, because the classes above have two different shapes.** +``max_elapsed`` stays at one minute and governs the fast classes — a connect +failure, a paced ``429`` — which resolve in seconds or not at all; a blackholed +host has no business pinning a caller (a whole thread, on the sync client) for +longer than that. ``collect_max_elapsed`` is twenty minutes and governs class 3 +alone, because it is the one class whose budget has to outlast a *server-side* +bound: a ``deadline_exceeded`` ``504`` arrives at Comfy's own deadline, the same +ten minutes as :data:`~comfy_low.transport.MODEL_RUN_TIMEOUT`, so a budget of one +deadline window is already spent when the ``504`` lands and the collect attempt +it exists for never starts. Two windows is one to reach the ``504`` and one to +collect what it left running — and nothing else pays for it. Both are measured +from the same origin, the construction of the :class:`Retrier`; which one applies +is decided per failure, by :func:`is_collectable`. Retry is on by default. ``Comfy(retry=NO_RETRY)`` turns it off; any other policy is a :class:`RetryPolicy` you construct:: @@ -166,7 +169,8 @@ from comfy_sdk import Comfy, NO_RETRY, RetryPolicy Comfy(retry=NO_RETRY) # exactly one attempt - Comfy(retry=RetryPolicy(max_elapsed=60.0)) # give up after a minute + Comfy(retry=RetryPolicy(max_elapsed=300.0)) # fast classes: 5 minutes + Comfy(retry=RetryPolicy(collect_max_elapsed=60.0)) # bound the collect loop Comfy(retry=RetryPolicy(retry_collectable=False)) # no 504/409 collect loop """ @@ -223,10 +227,25 @@ #: another" (``spec/router-openapi.yaml``). _DEADLINE_EXCEEDED = "deadline_exceeded" +#: The bucket the idempotency layer answers a mid-collect ``409`` with: the key +#: is recognised and the generation it names has not finished. Neither vendored +#: spec contracts this bucket, which is exactly why the gate names it rather +#: than accepting any paced ``409``: the ``409`` s the specs *do* document are +#: deterministic refusals (``hash_mismatch``, which ``spec/openapi.yaml`` gives a +#: ``Retry-After``, and ``asset_in_use``), and a proxy or WAF conflict carries no +#: bucket at all. Fail closed — an unrecognised ``409`` stays a refusal. +_GENERATION_IN_PROGRESS = "generation_in_progress" + #: Policy fields that must be real numbers for the arithmetic below to mean #: anything. Kept beside the fields themselves so a numeric one added later is #: added here too. -_NUMERIC_FIELDS = ("max_elapsed", "initial_backoff", "backoff_factor", "max_backoff") +_NUMERIC_FIELDS = ( + "max_elapsed", + "collect_max_elapsed", + "initial_backoff", + "backoff_factor", + "max_backoff", +) def is_unknown_outcome_status(status: int) -> bool: @@ -312,25 +331,32 @@ def is_collectable(exc: BaseException) -> bool: stopped holding the connection at its own bound while the generation ran on, and the contract says to "retry it with the SAME ``Idempotency-Key``", which "collects that generation rather than dispatching another"; - * a ``409`` carrying ``Retry-After`` — the idempotency layer's answer that - the request under this key is still in progress, which is what the collect - retry above meets when it arrives before the generation finishes. + * a ``generation_in_progress`` ``409`` carrying ``Retry-After`` — the + idempotency layer's answer that the request under this key is still in + progress, which is what the collect retry above meets when it arrives + before the generation finishes. - Both gates are load-bearing. The ``504`` must name its bucket because + Every gate is load-bearing, and each status is gated on *both* its bucket + and its pace for the same reason. The ``504`` must name its bucket because ``deadline_exceeded`` shares that status with ``provider_timeout``, where no contract blesses the resend and a header-less ``504`` from an intermediary is - read as exactly that. The ``Retry-After`` must be present because the router - sends it on a ``deadline_exceeded`` "only when Comfy holds a handle to a - generation the provider is still running" — without it there is nothing to - collect, and a resend would be dispatching new work rather than gathering - old. A ``409`` with no pace is an ordinary conflict and stays a refusal. + read as exactly that. The ``409`` must name its bucket because every ``409`` + either vendored spec documents is a *deterministic* refusal that a pace does + not soften — ``spec/openapi.yaml`` gives its ``hash_mismatch`` ``409`` a + ``Retry-After`` outright — so accepting the status plus a pace alone would + resend a permanent refusal, or a bucket-less proxy conflict, for the whole + budget. The ``Retry-After`` must be present because the router sends it on a + ``deadline_exceeded`` "only when Comfy holds a handle to a generation the + provider is still running" — without it there is nothing to collect, and a + resend would be dispatching new work rather than gathering old. """ status = getattr(exc, "http_status", None) if not isinstance(status, int) or retry_after_of(exc) is None: return False + bucket = error_bucket_of(exc) if status == _CONFLICT: - return True - return status == _GATEWAY_TIMEOUT and error_bucket_of(exc) == _DEADLINE_EXCEEDED + return bucket == _GENERATION_IN_PROGRESS + return status == _GATEWAY_TIMEOUT and bucket == _DEADLINE_EXCEEDED @dataclass(frozen=True) @@ -345,12 +371,14 @@ class RetryPolicy: #: Seconds from the first attempt after which no *new* attempt is started. #: An attempt already running is never interrupted by it, so the worst-case #: wall clock for a call is this plus one per-attempt ``timeout``. Zero - #: disables retrying entirely (see :data:`NO_RETRY`). The default is one - #: server deadline window — the same ten minutes as - #: :data:`~comfy_low.transport.MODEL_RUN_TIMEOUT` — so that a collect loop - #: after a ``deadline_exceeded`` ``504`` can outlast the bound that produced - #: it; see this module's docstring for what that costs the other classes. - max_elapsed: float = 600.0 + #: disables retrying entirely (see :data:`NO_RETRY`), collect included. + #: This is the budget for the *fast* classes — a connect failure, a paced + #: ``429`` — which is why it stays at a minute: those resolve in seconds or + #: not at all, and a blackholed host should not pin a caller for longer than + #: the failure needs. The collect class gets its own, longer budget in + #: :attr:`collect_max_elapsed`, because it is the one class that has to + #: outlast a server-side deadline. + max_elapsed: float = 60.0 #: Ceiling on the delay before the first retry. With ``jitter`` on — the #: default — the actual delay is drawn from ``[0, this]``. initial_backoff: float = 0.5 @@ -381,8 +409,8 @@ class RetryPolicy: #: its own. retry_possibly_in_flight: bool = False #: Retry the failures the server itself paced for a same-key resend — a - #: router ``deadline_exceeded`` ``504`` and an in-progress ``409``, each - #: carrying ``Retry-After``. See :func:`is_collectable` for the exact gates. + #: router ``deadline_exceeded`` ``504`` and a ``generation_in_progress`` + #: ``409``, each carrying ``Retry-After``. See :func:`is_collectable` for the exact gates. #: **On by default**, because this is the one class where the contract says #: the resend collects the generation already running rather than #: dispatching a second one, and the ``Retry-After`` is the server's own @@ -391,6 +419,23 @@ class RetryPolicy: #: declared last, out of reading order, so that inserting it cannot change #: what an existing positional ``RetryPolicy(...)`` means. retry_collectable: bool = True + #: Seconds from the first attempt for the *collect* class specifically — the + #: failures :func:`is_collectable` recognises. Measured from the same origin + #: as :attr:`max_elapsed` and used in its place once the server has said the + #: work is still running. + #: + #: It is separate, and twenty minutes, because the collect loop is the one + #: class whose budget has to outlast a server-side bound. A + #: ``deadline_exceeded`` ``504`` arrives *at* Comfy's own deadline — the same + #: ten minutes as :data:`~comfy_low.transport.MODEL_RUN_TIMEOUT` — so a + #: budget of one deadline window is already spent by the time the ``504`` + #: lands and the collect attempt it was sized for never starts. Two windows + #: is one to reach the ``504`` and one to collect what it left running. The + #: cost is not charged to anything else: a refused connection still gives up + #: at :attr:`max_elapsed`. Ignored entirely when ``max_elapsed`` is zero, so + #: :data:`NO_RETRY` is still exactly one attempt. Declared last for the same + #: positional-compatibility reason as :attr:`retry_collectable`. + collect_max_elapsed: float = 1200.0 def __post_init__(self) -> None: for name in _NUMERIC_FIELDS: @@ -404,6 +449,8 @@ def __post_init__(self) -> None: raise ValueError(f"{name} must be a finite number") if self.max_elapsed < 0: raise ValueError("max_elapsed must not be negative") + if self.collect_max_elapsed < 0: + raise ValueError("collect_max_elapsed must not be negative") if self.initial_backoff <= 0: raise ValueError("initial_backoff must be positive") if self.backoff_factor < 1: @@ -436,14 +483,23 @@ def should_retry(self, exc: BaseException) -> bool: # attached: the spec's retry signal is "429 + Retry-After", # and a 429 without one is not asking to be asked again. return retry_after_of(exc) is not None - if self.retry_collectable and is_collectable(exc): - # Checked before both branches below, because it overrides + if is_collectable(exc): + # Classified before both branches below, because it overrides # both: a `deadline_exceeded` 504 is a 5xx the router contract # nevertheless blesses a same-key resend for, and an # in-progress 409 is a 4xx that does become true on a later # ask. Everything narrowing it to those two answers lives in # `is_collectable`. - return True + # + # The answer is `retry_collectable` rather than `True`, and the + # switch is read *after* the classification rather than as part + # of it, so that turning the collect loop off is an answer and + # not a fall-through: gated the other way round, a paced + # `deadline_exceeded` 504 with `retry_collectable=False` would + # drop into `is_unknown_outcome_status` below and be retried + # anyway under `retry_possibly_in_flight=True` — an opt-out that + # silently did nothing for the very status it is named for. + return self.retry_collectable if is_unknown_outcome_status(status): # Every other 5xx, including the router's `service_unavailable` # 503: the bucket says the condition clears on its own, but @@ -499,7 +555,7 @@ def _ceiling(self, exponent: int) -> float: #: Retrying turned off: one logical call is exactly one attempt. Pass it as #: ``Comfy(retry=NO_RETRY)``. The ``Idempotency-Key`` is still sent — it is not #: part of the retry policy, and a caller who retries by hand needs it. -NO_RETRY = RetryPolicy(max_elapsed=0.0) +NO_RETRY = RetryPolicy(max_elapsed=0.0, collect_max_elapsed=0.0) class Retrier: @@ -522,10 +578,11 @@ def __init__( self._now = now self._rng = rng self._attempts = 0 - # The whole call's budget, fixed here so every later decision measures - # against one origin — elapsed time across all attempts, not time - # granted afresh to each of them. - self._deadline = now() + policy.max_elapsed + # The one origin every later decision measures against — elapsed time + # across all attempts, not time granted afresh to each of them. The + # *budget* laid over it depends on which class the failure turns out to + # be (see `delay_before_retry`); the origin never does. + self._started = now() @property def attempts(self) -> int: @@ -541,19 +598,43 @@ def delay_before_retry(self, exc: BaseException) -> float | None: refusal — with full jitter one unlucky draw would otherwise end a call that still had seconds of budget, making identical calls take a randomly varying number of attempts. ``client.py::_retry_delay`` bounds - the workflow surface's 429 wait the same way. So ``max_elapsed`` is - when the last attempt may *start*, inclusive. + the workflow surface's 429 wait the same way. So the budget is when the + last attempt may *start*, inclusive. + + *Which* budget depends on the class: a collectable failure is measured + against :attr:`RetryPolicy.collect_max_elapsed` and everything else + against :attr:`RetryPolicy.max_elapsed`, both from this ``Retrier``'s + construction. See ``collect_max_elapsed`` for why the collect class + needs its own -- and why the other classes must not pay for it. """ failed_at = self._now() self._attempts += 1 if not self._policy.enabled or not self._policy.should_retry(exc): return None - remaining = self._deadline - failed_at + budget = ( + self._policy.collect_max_elapsed + if self._policy.retry_collectable and is_collectable(exc) + else self._policy.max_elapsed + ) + remaining = self._started + budget - failed_at if remaining <= 0: return None - # A pace the server named beats the schedule guessed here. + # A pace the server named beats the schedule guessed here -- but only a + # usable one. `Retry-After: 0` is outside what the router contract can + # send (`RouterRetryAfterHeader` is pinned to `minimum: 1`), and taking + # it verbatim would turn a server or intermediary that keeps answering + # it into a zero-delay resend loop of full model-run POSTs, bounded only + # by the budget. A pace of zero names no pace: fall back to the jittered + # backoff, which is also what stops such a loop synchronising across + # clients. Whether the header was *present* still decides collectability + # over in `is_collectable` -- that is the server saying it holds a handle + # to a generation, which a nonsense value does not retract. named = retry_after_of(exc) - delay = named if named is not None else self._policy.backoff(self._attempts, rng=self._rng) + delay = ( + named + if named is not None and named > 0 + else self._policy.backoff(self._attempts, rng=self._rng) + ) return min(delay, remaining) diff --git a/tests/conftest.py b/tests/conftest.py index 8b198d9..8889f90 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -128,6 +128,12 @@ class ServerState: # 409). `None` sends no header at all, which is the same failure the policy # must *not* retry. model_run_retry_after: str | None = None + # Answer model-run failures in Router's own error shape -- the coarse bucket + # on `X-Comfy-Error-Type` plus a `{detail, error_type}` body -- instead of + # the v2 `{error: {code, message}}` envelope. `POST /api/v2/models/run` is + # fronted by Router, so this is the shape a real deployment's 504 arrives + # in, and the bucket-keyed collect rule has to read it. + model_run_router_error_shape: bool = False # --- counters the tests assert on --- upload_count: int = 0 @@ -474,29 +480,55 @@ def _post_model_run(self) -> None: # The server holding the connection while it polls upstream. time.sleep(state.model_run_delay) - def claim_if_outcome_unknown(status: int) -> None: + def claim_if_outcome_unknown(status: int, code: str) -> None: # The contract releases a key for a request that definitively # failed without starting work (a 4xx) and keeps it claimed # when the outcome is unknown (a 5xx). Two exceptions: a # replaying deployment (what the opt-in is for), and the # router's `deadline_exceeded` 504, whose reservation survives # so the same key can collect the generation still running. + # + # The bucket is part of that second exception and not an + # afterthought: the carve-out the router documents is for + # `deadline_exceeded` specifically, not for the status, which it + # shares with `provider_timeout`. Keying on the status alone + # would make the stub more permissive than the contract it + # stands in for, and an SDK regression that resent a + # `provider_timeout` 504 would pass here while meeting a 422 on + # a real server. if not key or status < 500: return if state.model_run_replays_idempotency_key: return - if state.model_run_collects_after_deadline and status == 504: + if ( + state.model_run_collects_after_deadline + and status == 504 + and code == "deadline_exceeded" + ): return state.model_run_idempotency[key] = "claimed" def fail(status: int, code: str, message: str) -> None: - claim_if_outcome_unknown(status) - headers = ( - {"Retry-After": state.model_run_retry_after} - if state.model_run_retry_after is not None - else None + claim_if_outcome_unknown(status, code) + headers = {} + if state.model_run_retry_after is not None: + headers["Retry-After"] = state.model_run_retry_after + if state.model_run_router_error_shape: + # What a real Router failure looks like: the coarse bucket + # on `X-Comfy-Error-Type` and a `{detail, error_type}` body, + # with no v2 `error.code` anywhere. The SDK's bucket-keyed + # retry rules have to survive this shape too, and nothing + # exercised it while the stub only ever spoke the envelope. + headers["X-Comfy-Error-Type"] = code + self._json( + status, {"detail": message, "error_type": code}, headers=headers or None + ) + return + self._json( + status, + {"error": {"code": code, "message": message}}, + headers=headers or None, ) - self._json(status, {"error": {"code": code, "message": message}}, headers=headers) if state.model_run_fail_times > 0: state.model_run_fail_times -= 1 diff --git a/tests/test_error_mapping.py b/tests/test_error_mapping.py index 265f8ca..17a732d 100644 --- a/tests/test_error_mapping.py +++ b/tests/test_error_mapping.py @@ -1,10 +1,14 @@ -"""to_sdk_error must map the server's 404 codes to the typed NotFound.""" +"""How an error response on the wire becomes a typed exception. + +`to_sdk_error` mapping the server's 404 codes to the typed `NotFound`, and +`error_from_envelope` reading the two body shapes this API answers in. +""" from __future__ import annotations import pytest -from comfy_low.errors import ApiError +from comfy_low.errors import ApiError, QueueFull, error_from_envelope from comfy_sdk.exceptions import NotFound, to_sdk_error @@ -15,3 +19,76 @@ def test_404_codes_map_to_notfound(code: str) -> None: err = to_sdk_error(ApiError("not found", code=code, http_status=404)) assert isinstance(err, NotFound) assert err.code == code + + +# --- the two error shapes one client has to read --- +# +# `POST /api/v2/models/run` is fronted by Router, whose error body is +# `{detail, error_type}` with the coarse bucket repeated on `X-Comfy-Error-Type` +# -- not the v2 `{error: {code, message}}` envelope every other route answers in. +# `error_from_envelope` has to read both, because `comfy_sdk.retry` keys its +# default-on collect rule on that bucket and would otherwise be a silent no-op +# against every real Router 504. + + +def test_the_v2_envelope_code_still_wins() -> None: + err = error_from_envelope( + 504, + {"error": {"code": "generation_in_progress", "message": "still running"}}, + error_type="deadline_exceeded", + ) + assert err.code == "generation_in_progress" + assert err.message == "still running" + + +def test_routers_body_error_type_is_read_as_the_code() -> None: + err = error_from_envelope( + 504, {"detail": "upstream is slow", "error_type": "deadline_exceeded"} + ) + assert err.code == "deadline_exceeded" + assert err.message == "upstream is slow" + + +def test_routers_error_type_header_is_read_when_the_body_carries_none() -> None: + # The header is the bucket's other home, and `spec/router-openapi.yaml` + # marks it required on every Router error response. + err = error_from_envelope(504, None, error_type="deadline_exceeded") + assert err.code == "deadline_exceeded" + # A body-less error response is still diagnosable by status. + assert err.message == "HTTP 504" + + +def test_a_documented_status_code_is_not_retyped_by_the_bucket() -> None: + # The guard on the fallback's placement. Router's 429 bucket is + # `rate_limited` / `concurrency_limit_exceeded`, but this API's 429 has + # meant `queue_full` -- and `QueueFull` -- since before Router fronted + # anything. Letting the bucket win here would silently retype an exception + # integrators already catch, so the status-derived code is consulted first + # and only the 5xx range, where the status determines nothing, is left to + # the bucket. + err = error_from_envelope(429, None, error_type="concurrency_limit_exceeded", retry_after=3) + assert err.code == "queue_full" + assert isinstance(err, QueueFull) + # And the pace survives, which is what `comfy_sdk.retry` keys the 429 on. + assert err.retry_after == 3 + + +def test_a_router_validation_body_degrades_rather_than_coercing_its_detail() -> None: + # Router's per-field validation body is the fal/FastAPI `detail[]` shape. A + # list is not a message: stringifying it would put a Python repr in front of + # a caller, so the status-derived message answers instead. + err = error_from_envelope( + 500, + {"detail": [{"loc": ["body", "steps"], "msg": "too large", "type": "value_error"}]}, + error_type="internal_error", + ) + assert err.code == "internal_error" + assert err.message == "HTTP 500" + + +@pytest.mark.parametrize("body", [None, {}, {"error": None}, {"error_type": " "}, {"detail": 7}]) +def test_a_body_that_names_no_bucket_still_falls_back_to_the_status(body: object) -> None: + # Degrading to the status-derived code is what keeps a malformed error + # response diagnosable rather than replacing it with a decoding failure. + err = error_from_envelope(401, body) # type: ignore[arg-type] + assert err.code == "unauthorized" diff --git a/tests/test_models_run_retry.py b/tests/test_models_run_retry.py index 235d352..851d199 100644 --- a/tests/test_models_run_retry.py +++ b/tests/test_models_run_retry.py @@ -563,12 +563,20 @@ def test_jitter_can_be_turned_off_for_a_deterministic_schedule() -> None: "kwargs", [ {"max_elapsed": -1.0}, + {"collect_max_elapsed": -1.0}, {"initial_backoff": 0.0}, {"initial_backoff": -1.0}, {"backoff_factor": 0.5}, {"max_backoff": 0.1}, ], - ids=["negative-budget", "zero-backoff", "negative-backoff", "shrinking", "ceiling-below-floor"], + ids=[ + "negative-budget", + "negative-collect-budget", + "zero-backoff", + "negative-backoff", + "shrinking", + "ceiling-below-floor", + ], ) def test_an_incoherent_policy_is_rejected_at_construction(kwargs: dict) -> None: with pytest.raises(ValueError): @@ -585,6 +593,8 @@ def test_an_incoherent_policy_is_rejected_at_construction(kwargs: dict) -> None: # a NaN budget reads as "retrying disabled", a NaN backoff makes the # caller's `sleep` raise in place of the real error. {"max_elapsed": float("nan")}, + {"collect_max_elapsed": float("inf")}, + {"collect_max_elapsed": float("nan")}, {"initial_backoff": float("nan")}, {"backoff_factor": float("nan")}, {"max_backoff": float("nan")}, @@ -872,7 +882,11 @@ def test_the_collect_loop_gives_up_at_the_elapsed_budget(server) -> None: server.state.model_run_collects_after_deadline = True server.state.model_run_error = (504, "deadline_exceeded") server.state.model_run_retry_after = "1" - policy = RetryPolicy(max_elapsed=0.4, initial_backoff=0.01, max_backoff=0.02) + # `collect_max_elapsed`, not `max_elapsed`: the collect class has its own + # budget, and this is the bound that ends this loop. + policy = RetryPolicy( + max_elapsed=0.4, collect_max_elapsed=0.4, initial_backoff=0.01, max_backoff=0.02 + ) started = time.monotonic() with Comfy(retry=policy) as client: with pytest.raises(ComfyError) as excinfo: @@ -888,7 +902,7 @@ def test_the_budget_bounds_the_loop_even_at_the_server_named_pace() -> None: # in a 100s budget and the fourth is clamped to the 10s left rather than # overshooting it, after which no further attempt starts. clock = _FakeClock() - retrier = Retrier(RetryPolicy(max_elapsed=100.0), now=clock, rng=lambda: 1.0) + retrier = Retrier(RetryPolicy(collect_max_elapsed=100.0), now=clock, rng=lambda: 1.0) failure = DeadlineExceeded("still generating", http_status=504, retry_after=30) delays = [] while (delay := retrier.delay_before_retry(failure)) is not None: @@ -897,14 +911,48 @@ def test_the_budget_bounds_the_loop_even_at_the_server_named_pace() -> None: assert delays == [30.0, 30.0, 30.0, 10.0] -def test_the_default_budget_is_one_server_deadline_window() -> None: - # The decision this asserts, so it cannot drift silently: the budget is the - # same ten minutes `run` is already willing to spend on ONE attempt. A - # shorter one cannot outlast the deadline it is collecting after -- the 504 - # says the server stopped holding this connection at its own bound while - # the generation ran on, and the caller is billed for that generation - # whether or not the SDK waits to collect it. - assert DEFAULT_RETRY.max_elapsed == MODEL_RUN_TIMEOUT.read == 600.0 +def test_the_collect_budget_outlasts_a_server_deadline_window() -> None: + # The decision this asserts, so it cannot drift silently. A `deadline_exceeded` + # 504 arrives AT the server's own bound -- the same ten minutes `run` is + # already willing to spend on one attempt -- so a collect budget of exactly + # one window is spent by the time the 504 lands and the collect attempt it + # was sized for never starts. It has to be strictly more than one window, + # with room for a collect attempt of its own. + deadline_window = MODEL_RUN_TIMEOUT.read + assert deadline_window == 600.0 + assert DEFAULT_RETRY.collect_max_elapsed >= 2 * deadline_window + + # And the fast classes do not pay for it: a refused connection still gives + # up in a minute rather than sitting on a caller's thread for twenty. + assert DEFAULT_RETRY.max_elapsed == 60.0 + + +def test_the_collect_budget_applies_only_to_the_collect_class() -> None: + # Two budgets, one origin. The same `Retrier` answers a collectable failure + # against the long budget and everything else against the short one, so the + # collect loop's room is not a policy-wide regression for the fast classes. + policy = RetryPolicy(max_elapsed=10.0, collect_max_elapsed=100.0, jitter=False) + collectable = DeadlineExceeded("still generating", http_status=504, retry_after=5) + refused = httpx.ConnectError("refused") + + clock = _FakeClock() + retrier = Retrier(policy, now=clock, rng=lambda: 1.0) + # 40s in: past `max_elapsed`, well inside `collect_max_elapsed`. + clock.advance(40.0) + assert retrier.delay_before_retry(collectable) == 5.0 + assert retrier.delay_before_retry(refused) is None + + +def test_a_deadline_504_at_the_server_bound_still_gets_a_collect_attempt() -> None: + # The regression the split budget exists for, end to end against a fake + # clock: the 504 arrives after a full ten-minute attempt, which a + # one-window budget would have spent entirely. The collect attempt has to + # start anyway, because that generation is running and billed either way. + clock = _FakeClock() + retrier = Retrier(DEFAULT_RETRY, now=clock, rng=lambda: 1.0) + clock.advance(600.0) # one full server deadline window; see the test above + failure = DeadlineExceeded("still generating", http_status=504, retry_after=2) + assert retrier.delay_before_retry(failure) == 2.0 @pytest.mark.parametrize( @@ -913,7 +961,20 @@ def test_the_default_budget_is_one_server_deadline_window() -> None: (DeadlineExceeded("deadline", http_status=504, retry_after=3), True), (DeadlineExceeded("deadline", http_status=504), False), (ProviderTimeout("upstream", http_status=504, retry_after=3), False), - (RouterError("in progress", http_status=409, retry_after=3), True), + (RouterError("in progress", http_status=409, retry_after=3), False), + ( + ApiError( + "in progress", + code="generation_in_progress", + http_status=409, + retry_after=3, + ), + True, + ), + # `spec/openapi.yaml` gives its `hash_mismatch` 409 a `Retry-After`, and + # it is still a deterministic refusal: the pace does not make asking + # again change the answer. The bucket gate is what keeps it one. + (ApiError("bad hash", code="hash_mismatch", http_status=409, retry_after=3), False), (RouterError("conflict", http_status=409), False), (ServiceUnavailable("later", http_status=503, retry_after=3), False), (InternalError("boom", http_status=500, retry_after=3), False), @@ -947,3 +1008,185 @@ def test_a_deadline_504_raised_as_a_protocol_error_is_still_collected() -> None: # for the same class of silent no-op. exc = ApiError("deadline", code="deadline_exceeded", http_status=504, retry_after=2) assert DEFAULT_RETRY.should_retry(exc) + + +# --- the gates the collect loop is only safe behind --- + + +def test_a_zero_retry_after_does_not_become_a_zero_delay_resend_loop() -> None: + # `RouterRetryAfterHeader` is pinned to `minimum: 1`, so `Retry-After: 0` is + # a server or intermediary answering outside its own contract -- and taking + # it verbatim would resend a full model-run POST with no wait at all, for + # the whole budget, entirely under that server's control. Zero names no + # usable pace: the local backoff schedule answers instead. + clock = _FakeClock() + policy = RetryPolicy(initial_backoff=1.0, max_backoff=8.0, jitter=False) + retrier = Retrier(policy, now=clock, rng=lambda: 1.0) + failure = DeadlineExceeded("still generating", http_status=504, retry_after=0) + delays = [retrier.delay_before_retry(failure) for _ in range(3)] + assert delays == [1.0, 2.0, 4.0] + + +def test_a_zero_retry_after_still_counts_as_the_server_holding_a_handle() -> None: + # The value is unusable; the header's *presence* is not. It is the router + # saying it holds a generation to collect, which a nonsense number does not + # retract -- so the response stays collectable and only its pace is ignored. + failure = DeadlineExceeded("still generating", http_status=504, retry_after=0) + assert is_collectable(failure) + assert DEFAULT_RETRY.should_retry(failure) + + +def test_turning_the_collect_loop_off_holds_even_with_the_in_flight_opt_in() -> None: + # `retry_collectable=False` has to be an answer, not a fall-through. A + # paced `deadline_exceeded` 504 is also a 5xx, so classified the other way + # round it would drop into the unknown-outcome branch and be retried anyway + # under `retry_possibly_in_flight=True` -- an opt-out that silently did + # nothing for the one status it is named for. + policy = RetryPolicy(retry_collectable=False, retry_possibly_in_flight=True) + failure = DeadlineExceeded("still generating", http_status=504, retry_after=2) + assert is_collectable(failure) + assert not policy.should_retry(failure) + # And the 409 half of the class, which the fall-through happened to get + # right (a 4xx has no unknown-outcome branch to fall into) -- pinned so it + # stays right for the stated reason rather than by accident. + in_progress = ApiError( + "in progress", code="generation_in_progress", http_status=409, retry_after=2 + ) + assert is_collectable(in_progress) + assert not policy.should_retry(in_progress) + # The rest of the policy is untouched: this is one knob, not a kill switch. + assert policy.should_retry(InternalError("boom", http_status=500)) + + +def test_a_paced_hash_mismatch_409_is_still_a_refusal(server) -> None: + # The 409 the vendored spec actually documents, and it documents a + # `Retry-After` on it. It is a deterministic refusal all the same: the + # client-computed hash will not match on the second ask either. Without the + # bucket gate the collect loop would resend it for the whole budget -- and + # under the v2 rule that releases the key for a 4xx that started no work, + # each resend is a fresh billed dispatch rather than a collection. + server.state.model_run_error = (409, "hash_mismatch") + server.state.model_run_retry_after = "0" + with Comfy(retry=FAST) as client: + with pytest.raises(ComfyError) as excinfo: + client.models.run(MODEL, ARGS) + assert server.state.model_run_count == 1 + assert excinfo.value.http_status == 409 + + +# --- the wire shape a real Router failure arrives in --- + + +def test_a_deadline_504_in_routers_own_error_shape_is_still_collected(server) -> None: + # `POST /api/v2/models/run` is fronted by Router, whose error body is + # `{detail, error_type}` with the bucket repeated on `X-Comfy-Error-Type` -- + # not the v2 `{error: {code}}` envelope. Reading only the envelope would + # collapse every real 504 to the status-derived default, and the whole + # bucket-keyed collect rule would be a silent no-op in production while + # every test above passed. + server.state.model_run_router_error_shape = True + server.state.model_run_collects_after_deadline = True + server.state.model_run_transient_error = (504, "deadline_exceeded") + server.state.model_run_retry_after = "0" + server.state.model_run_fail_times = 1 + with Comfy(retry=FAST) as client: + assert client.models.run(MODEL, ARGS) == server.state.model_run_result + assert server.state.model_run_count == 2 + first, second = server.state.model_run_idempotency_keys + assert first is not None + assert second == first + + +def test_the_bucket_gate_holds_in_routers_own_error_shape_too(server) -> None: + # The gate has to survive the shape change in both directions: reading the + # header must not turn every Router 504 into a collectable one. + server.state.model_run_router_error_shape = True + server.state.model_run_error = (504, "provider_timeout") + server.state.model_run_retry_after = "0" + with Comfy(retry=FAST) as client: + with pytest.raises(ComfyError) as excinfo: + client.models.run(MODEL, ARGS) + assert server.state.model_run_count == 1 + assert excinfo.value.http_status == 504 + + +def test_a_router_error_body_keeps_its_detail_as_the_message(server) -> None: + # Router names its human-readable string `detail`; dropping it for a bare + # "HTTP 504" would make the shape that reaches real callers the least + # diagnosable one. + server.state.model_run_router_error_shape = True + server.state.model_run_error = (500, "internal_error") + with Comfy(retry=NO_RETRY) as client: + with pytest.raises(ComfyError) as excinfo: + client.models.run(MODEL, ARGS) + assert "internal_error" in str(excinfo.value) + + +# --- the deployment that does NOT collect --- + + +def test_the_default_collect_loop_against_a_non_collecting_deployment(server) -> None: + # The cost of defaulting this on, asserted rather than assumed. `POST + # /models/run` is in neither vendored spec, so a deployment may well apply + # the v2 rule instead -- key stays claimed across an unknown-outcome 5xx, + # and the resend is rejected. The caller then sees `422 + # idempotency_key_reuse` in place of the real 504. That is the trade the + # default makes; `retry_collectable=False` is the way out of it, and this + # pins both halves so neither can change silently. + server.state.model_run_collects_after_deadline = False + server.state.model_run_error = (504, "deadline_exceeded") + server.state.model_run_retry_after = "0" + with Comfy(retry=FAST) as client: + with pytest.raises(IdempotencyKeyReuse): + client.models.run(MODEL, ARGS) + assert server.state.model_run_count == 2 + + server.state.model_run_count = 0 + server.state.model_run_idempotency.clear() + policy = RetryPolicy( + max_elapsed=5.0, initial_backoff=0.01, max_backoff=0.02, retry_collectable=False + ) + with Comfy(retry=policy) as client: + with pytest.raises(ComfyError) as excinfo: + client.models.run(MODEL, ARGS) + assert server.state.model_run_count == 1 + assert excinfo.value.http_status == 504 + + +# --- the body snapshot the same-key rule rests on --- + + +def test_the_body_is_snapshotted_deeply_before_the_first_attempt(server) -> None: + # Same key, same body is what makes the resend a collection. A shallow copy + # leaves nested values shared with the caller, so a mutation during the + # retry window -- up to twenty minutes on the collect path -- would send a + # different body under the one key and earn the 422 the snapshot exists to + # prevent. + server.state.model_run_collects_after_deadline = True + server.state.model_run_transient_error = (504, "deadline_exceeded") + server.state.model_run_retry_after = "0" + server.state.model_run_fail_times = 1 + arguments: dict[str, Any] = {"prompt": "a cat", "config": {"steps": 4}} + + class _MutatingLow: + def __init__(self, inner: Any) -> None: + self._inner = inner + + def __getattr__(self, name: str) -> Any: + return getattr(self._inner, name) + + def post_model_run(self, model: str, args: Mapping[str, Any], **kw: Any) -> Any: + try: + return self._inner.post_model_run(model, args, **kw) + except BaseException: + # The caller mutating its own nested dict *between* attempts, + # which is the only window that matters: the collect resend is + # what has to carry the same body as the first attempt. + arguments["config"]["steps"] = 999 + raise + + with Comfy(retry=FAST) as client: + models = Models(cast(Any, _MutatingLow(client._low)), FAST) + assert models.run(MODEL, arguments) == server.state.model_run_result + assert server.state.model_run_count == 2 + assert server.state.last_model_run_body["arguments"]["config"] == {"steps": 4}