[Partner Nodes] fix: respect Retry-After header

Signed-off-by: bigcat88 <bigcat88@icloud.com>
This commit is contained in:
bigcat88 2026-06-02 10:54:57 +03:00
parent 5955ddff52
commit b9cf3dd502
No known key found for this signature in database
GPG Key ID: 1F0BF0EC3CF22721
2 changed files with 31 additions and 0 deletions

View File

@ -4,6 +4,8 @@ import os
import re import re
import time import time
from collections.abc import Callable from collections.abc import Callable
from datetime import datetime, timezone
from email.utils import parsedate_to_datetime
from io import BytesIO from io import BytesIO
from yarl import URL from yarl import URL
@ -91,6 +93,32 @@ async def sleep_with_interrupt(
await asyncio.sleep(min(1.0, end - now)) await asyncio.sleep(min(1.0, end - now))
def _retry_after_wait(value: str | None, fallback: float, max_wait: float) -> float:
"""Delay before the next retry, honoring a server ``Retry-After`` header."""
seconds: float | None = None
if value is not None:
value = value.strip()
if value.isascii() and value.isdigit():
# delay-seconds form. The ASCII-digit guard keeps exotic Unicode "digit" characters away from float()
# an all-digit string always converts (huge values become inf, never raising).
seconds = float(value)
elif value:
# HTTP-date form. parsedate_to_datetime raises OverflowError (not a ValueError) on absurd years/offsets
try:
parsed = parsedate_to_datetime(value)
except (TypeError, ValueError, OverflowError):
parsed = None
if parsed is not None:
if parsed.tzinfo is None: # naive datetime: HTTP-date is UTC
parsed = parsed.replace(tzinfo=timezone.utc)
delta = (parsed - datetime.now(timezone.utc)).total_seconds()
seconds = delta if delta > 0 else 0.0
if seconds is None:
return fallback
return min(seconds, max_wait)
def mimetype_to_extension(mime_type: str) -> str: def mimetype_to_extension(mime_type: str) -> str:
"""Converts a MIME type to a file extension.""" """Converts a MIME type to a file extension."""
return mime_type.split("/")[-1].lower() return mime_type.split("/")[-1].lower()

View File

@ -21,6 +21,7 @@ from server import PromptServer
from . import request_logger from . import request_logger
from ._helpers import ( from ._helpers import (
_retry_after_wait,
default_base_url, default_base_url,
get_comfy_api_headers, get_comfy_api_headers,
get_node_id, get_node_id,
@ -82,6 +83,7 @@ class _PollUIState:
_RETRY_STATUS = {408, 500, 502, 503, 504} # status 429 is handled separately _RETRY_STATUS = {408, 500, 502, 503, 504} # status 429 is handled separately
_MAX_RETRY_AFTER_WAIT = 150.0 # Cap a server Retry-After at this many seconds so a large hint can't block execution
COMPLETED_STATUSES = ["succeeded", "succeed", "success", "completed", "finished", "done", "complete"] COMPLETED_STATUSES = ["succeeded", "succeed", "success", "completed", "finished", "done", "complete"]
FAILED_STATUSES = ["cancelled", "canceled", "canceling", "fail", "failed", "error"] FAILED_STATUSES = ["cancelled", "canceled", "canceling", "fail", "failed", "error"]
QUEUED_STATUSES = ["created", "queued", "queueing", "submitted", "initializing", "wait", "in_queue"] QUEUED_STATUSES = ["created", "queued", "queueing", "submitted", "initializing", "wait", "in_queue"]
@ -747,6 +749,7 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
should_retry = True should_retry = True
if should_retry: if should_retry:
wait_time = _retry_after_wait(resp.headers.get("Retry-After"), wait_time, _MAX_RETRY_AFTER_WAIT)
logging.warning( logging.warning(
"HTTP %s %s -> %s. Waiting %.2fs (%s).", "HTTP %s %s -> %s. Waiting %.2fs (%s).",
method, method,