mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-28 00:11:27 -03:00
feat(services): add per-destination rate-limit gate for API traffic (#1085)
Implement Phase 1 of docs/plans/issue-1085-rate-limit-design.md: - New RateLimitCoordinator: per-host shared Retry-After gate with exponential backoff (30s base, 1800s cap), minimum inter-request pacing (default 0.75s), herd-free waiter serialization via per-destination locks, and a bounded wait (default 300s) that raises instead of parking. - Downloader.make_request: connectivity-guard fail-fast first, then gate pacing; on 429 register the cooldown and wait-and-resend (bounded); errors that passed through the gate are marked gate_handled. - FallbackMetadataProvider / MetadataSyncService: a network provider 429 no longer fails over to other network providers (stops the CivArchive flood); sqlite stays as local last resort. Rate-limited lookups now report "Rate limited" instead of "Model not found", so transient 429s no longer mark models civitai_deleted. - _RateLimitRetryHelper skips its own sleep for gate_handled errors, removing the double wait. - New settings: rate_limit_gate_enabled, rate_limit_max_wait_seconds, rate_limit_min_interval_seconds.
This commit is contained in:
Will Miao
@@ -1,7 +1,11 @@
|
||||
# Plan: Global Rate-Limit Abidance for Recipe Ingest & Metadata Fetching
|
||||
|
||||
**Issue:** [#1085 — Large Recipe Ingest Appears to not abide by vendor rate limits, possibly a few other errors?](https://github.com/willmiao/ComfyUI-Lora-Manager/issues/1085)
|
||||
**Status:** v1 — awaiting review
|
||||
**Status:** v2 — reviewed; decisions recorded in §10. **Phase 1 implemented**
|
||||
(2026-08-27): coordinator + downloader gate + Fix C failover semantics +
|
||||
helper double-wait fix + settings; full regression 2385 passed. Changes vs v1:
|
||||
Fix C moved to Phase 1, helper double-wait resolved in Phase 1, gate/guard
|
||||
ordering specified, WebSocket slowdown hint confirmed in scope (Phase 2).
|
||||
**Scope:** HTTP API traffic to CivitAI (`civitai.red`) and CivArchive (`civarchive.com`) from metadata fetching (bulk refresh, metadata sync, recipe analysis/enrichment, usage-control lookups). Large binary downloads (model files / preview images via `download_file`) are out of scope for *pacing* (they are already single-connection transfers) but their 429 responses should still be *registered*.
|
||||
|
||||
> Context: a first batch of fixes for this issue was already committed as
|
||||
@@ -187,12 +191,15 @@ API:
|
||||
|
||||
Enforcement points:
|
||||
|
||||
1. **`Downloader.make_request`** (`downloader.py:1102-1132`): before
|
||||
`session.request`, `await coordinator.wait_for_slot(destination)`. On 429:
|
||||
`coordinator.register_rate_limit(...)`, then *wait for the gate and
|
||||
re-send* (loop, bounded by `rate_limit_max_wait_seconds`, default 300;
|
||||
`retry_after ≥ cap` ⇒ fail immediately). After the loop, return the
|
||||
`RateLimitError` to the caller (unchanged contract). 200 path calls
|
||||
1. **`Downloader.make_request`** (`downloader.py:1102-1132`): ordering inside
|
||||
the method is **connectivity-guard fail-fast first** (offline short-circuit
|
||||
costs nothing to check), **then** `await coordinator.wait_for_slot(destination)`
|
||||
before `session.request`. On 429: `coordinator.register_rate_limit(...)`,
|
||||
then *wait for the gate and re-send* (loop, bounded by
|
||||
`rate_limit_max_wait_seconds`, default 300; `retry_after ≥ cap` ⇒ fail
|
||||
immediately). After the loop, return the `RateLimitError` to the caller
|
||||
(unchanged contract) **with `exc.gate_handled = True` set** so downstream
|
||||
retry helpers know the wait already happened. 200 path calls
|
||||
`register_success`.
|
||||
2. **`Downloader.download_to_memory` / `get_response_headers`** (phase 2):
|
||||
register 429s (so API calls queue); waiting only in `make_request`
|
||||
@@ -211,10 +218,12 @@ Enforcement points:
|
||||
`"rate limited (retry_after=…s); re-run the import later"` instead of
|
||||
`FAILED`, and surface a `rate_limited` flag in the WebSocket progress
|
||||
broadcast.
|
||||
6. **`_RateLimitRetryHelper` retries** (`model_metadata_provider.py`): with the
|
||||
gate at the downloader, the helper's `retry_after` sleeps become mostly
|
||||
redundant; demote/simplify in phase 2 (keep the wiring so a
|
||||
`RateLimitError` still propagates cleanly).
|
||||
6. **`_RateLimitRetryHelper` retries** (`model_metadata_provider.py`):
|
||||
**Phase 1** — when the raised `RateLimitError` carries `gate_handled = True`
|
||||
(set by the downloader after honoring the gate), the helper skips its own
|
||||
`retry_after` sleep and re-raises immediately, eliminating the double wait.
|
||||
The wiring stays so a `RateLimitError` still propagates cleanly; full
|
||||
demotion/removal can follow once the gate proves out.
|
||||
|
||||
Settings (`settings.json`, schema extension in `SettingsManager`):
|
||||
|
||||
@@ -230,7 +239,7 @@ Settings (`settings.json`, schema extension in `SettingsManager`):
|
||||
|---|---|
|
||||
| `py/services/rate_limit_coordinator.py` (new) | coordinator singleton + per-destination state + tests seam |
|
||||
| `py/services/downloader.py` | gate pre-check + 429 register/wait/retry loop + `register_success`; log the 429 notice at INFO once per cooldown, then DEBUG |
|
||||
| `py/services/model_metadata_provider.py` | `FallbackMetadataProvider`: stop network failover on `RateLimitError`; helper simplification |
|
||||
| `py/services/model_metadata_provider.py` | `FallbackMetadataProvider`: stop network failover on `RateLimitError`; helper skips its sleep when the error is marked `gate_handled` |
|
||||
| `py/services/metadata_sync_service.py` | `fetch_and_update_model`: same failover semantics; keep sqlite last resort |
|
||||
| `py/services/batch_import_service.py` | (phase 2) rate-limit failures → `SKIPPED` + `rate_limited` progress flag |
|
||||
| `py/services/settings_manager.py` | new settings keys + defaults |
|
||||
@@ -245,18 +254,24 @@ Settings (`settings.json`, schema extension in `SettingsManager`):
|
||||
up to the wait cap — UI actions that call the API (e.g. a model-details
|
||||
fetch) may take longer during cooldowns. Mitigation: bounded cap + INFO log
|
||||
+ the existing async request handling already tolerates slow responses.
|
||||
Open question: should interactive (non-batch) requests skip the wait and
|
||||
fail fast instead? (recommend: same wait — simpler, and cooldowns are short)
|
||||
**Decided (§10): interactive requests take the same bounded wait** — one
|
||||
behavior, no call-source plumbing; cooldowns are usually short.
|
||||
- **Gate waits occupy batch slots**: with the 1–5 batch semaphore, all slots
|
||||
can park on a gate simultaneously, freezing visible progress for up to one
|
||||
wait cap per wave. Bounded and acceptable; the phase-2 `SKIPPED` mapping +
|
||||
WebSocket `rate_limited` flag (both confirmed in scope, §10) make the stall
|
||||
visible and recoverable.
|
||||
- **Rate limit reality check**: CivitAI anonymous vs keyed limits, and whether
|
||||
`civitai.red` differs, is unverified. Default pacing `0.75 s/req` is a
|
||||
conservative guess (R6). Open question for maintainer: preferred default
|
||||
and whether an API-keyed ceiling should be higher.
|
||||
- **Long CivArchive windows**: `Retry-After ~1500 s` observed in code
|
||||
comments. A 300 s default cap means such lookups fail rather than wait.
|
||||
Open question: raise the default cap, or accept failure+skip semantics?
|
||||
- **Double waiting**: `_RateLimitRetryHelper` + gate could stack waits; the
|
||||
phase-2 simplification removes the helper's own sleeps for requests that go
|
||||
through the downloader.
|
||||
comments. **Decided (§10): keep the 300 s default cap** — such lookups
|
||||
fail/skip rather than park a request path for 25 minutes; batch import maps
|
||||
them to `SKIPPED` (phase 2) so the user can re-run later.
|
||||
- **Double waiting**: `_RateLimitRetryHelper` + gate could stack waits.
|
||||
**Resolved in Phase 1**: the downloader marks gate-honored errors with
|
||||
`gate_handled = True` and the helper skips its own sleep for those.
|
||||
- **Downloads**: `download_file` 429s return an error to download managers
|
||||
unchanged (already handled); only *registration* is proposed, so future
|
||||
API calls queue behind a large `Retry-After` from a download burst.
|
||||
@@ -289,22 +304,30 @@ Settings (`settings.json`, schema extension in `SettingsManager`):
|
||||
|
||||
## 9. Implementation Phases
|
||||
|
||||
- **Phase 1 (this plan, after review):** `RateLimitCoordinator` +
|
||||
`Downloader.make_request` integration (pre-check pacing + 429
|
||||
register/wait/retry loop + cap) + settings + coordinator/downloader tests.
|
||||
- **Phase 2:** failover semantics (`FallbackMetadataProvider`,
|
||||
`fetch_and_update_model`), helper simplification, batch-import
|
||||
`SKIPPED`-on-rate-limit + progress flag, `download_to_memory`/HEAD 429
|
||||
registration, provider/sync/batch tests.
|
||||
- **Phase 1 (this plan, reviewed):** `RateLimitCoordinator` +
|
||||
`Downloader.make_request` integration (guard fail-fast → gate pre-check
|
||||
pacing → 429 register/wait/retry loop with cap → `gate_handled` marking) +
|
||||
settings + **Fix C failover semantics** (`FallbackMetadataProvider`,
|
||||
`fetch_and_update_model` — moved up from phase 2: smallest diff, kills the
|
||||
CivArchive flood immediately, independent of coordinator correctness) +
|
||||
`_RateLimitRetryHelper` double-wait fix + coordinator/downloader/provider/
|
||||
sync tests.
|
||||
- **Phase 2:** batch-import `SKIPPED`-on-rate-limit + `rate_limited` WebSocket
|
||||
progress flag + slowdown hint (confirmed, §10),
|
||||
`download_to_memory`/HEAD 429 registration, batch tests.
|
||||
- **Phase 3:** full regression + docs + commit referencing `(#1085)`.
|
||||
|
||||
## 10. Review Checklist
|
||||
## 10. Review Checklist — Decisions (2026-08-27)
|
||||
|
||||
- [ ] Default pacing interval acceptable (`0.75 s`)? Prefer higher/lower?
|
||||
- [ ] Wait cap `300 s` acceptable, or should long-window CivArchive lookups
|
||||
wait longer?
|
||||
- [ ] OK that interactive API calls also wait (bounded) instead of failing
|
||||
fast?
|
||||
- [ ] Keep sqlite as last resort behind a network rate limit?
|
||||
- [ ] Add a UI hint ("rate limited — slowing down") surfaced via WebSocket,
|
||||
or is INFO logging enough?
|
||||
- [x] Default pacing interval `0.75 s` — **accepted** as conservative default;
|
||||
tunable via `rate_limit_min_interval_seconds`. Revisit if CivitAI
|
||||
publishes keyed/anonymous ceilings.
|
||||
- [x] Wait cap `300 s` — **accepted**; long-window CivArchive lookups fail →
|
||||
batch import marks them `SKIPPED` with a rate-limit reason (phase 2).
|
||||
- [x] Interactive API calls also wait (bounded) — **yes**, same behavior for
|
||||
all callers.
|
||||
- [x] Keep sqlite as last resort behind a network rate limit — **yes**
|
||||
(local-only, no vendor cost).
|
||||
- [x] UI hint — **yes**: WebSocket `rate_limited` flag + "rate limited —
|
||||
slowing down" hint in batch-import progress (phase 2); INFO logging
|
||||
regardless.
|
||||
@@ -32,6 +32,7 @@ from .connectivity_guard import (
|
||||
ConnectivityGuard,
|
||||
)
|
||||
from .errors import RateLimitError
|
||||
from .rate_limit_coordinator import RateLimitCoordinator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -1074,12 +1075,37 @@ class Downloader:
|
||||
|
||||
Returns:
|
||||
Tuple[bool, Union[Dict, str]]: (success, response data or error message)
|
||||
|
||||
When the rate-limit gate is enabled (``rate_limit_gate_enabled``),
|
||||
requests are paced per destination and 429 responses are honored by
|
||||
waiting out the ``Retry-After`` window (bounded by
|
||||
``rate_limit_max_wait_seconds``) before re-sending. A ``RateLimitError``
|
||||
returned after gate involvement is marked with ``gate_handled = True``
|
||||
so downstream retry helpers do not wait a second time.
|
||||
"""
|
||||
guard = await ConnectivityGuard.get_instance()
|
||||
destination = self._guard_destination(url)
|
||||
# Fail fast on transport-level outages before pacing: there is no
|
||||
# point waiting out a vendor cooldown while the network is down.
|
||||
if guard.should_block_request(destination):
|
||||
return False, OFFLINE_COOLDOWN_ERROR
|
||||
|
||||
coordinator = await RateLimitCoordinator.get_instance()
|
||||
gate_enabled = coordinator.enabled
|
||||
# Safety bound on the wait-and-resend loop; each 429 normally exits
|
||||
# via the wait cap in wait_for_slot, this covers pathological 429s
|
||||
# with tiny Retry-After values.
|
||||
max_resend_attempts = 5
|
||||
attempt = 0
|
||||
|
||||
while True:
|
||||
if gate_enabled:
|
||||
try:
|
||||
await coordinator.wait_for_slot(destination)
|
||||
except RateLimitError as exc:
|
||||
exc.gate_handled = True
|
||||
return False, exc
|
||||
|
||||
try:
|
||||
session = await self.session
|
||||
# Debug log for proxy mode at request time
|
||||
@@ -1104,6 +1130,8 @@ class Downloader:
|
||||
) as response:
|
||||
if response.status == 200:
|
||||
guard.register_success(destination)
|
||||
if gate_enabled:
|
||||
coordinator.register_success(destination)
|
||||
# Try to parse as JSON, fall back to text
|
||||
try:
|
||||
data = await response.json()
|
||||
@@ -1120,6 +1148,7 @@ class Downloader:
|
||||
elif response.status == 429:
|
||||
retry_after = self._extract_retry_after(response.headers)
|
||||
error_msg = "Request rate limited"
|
||||
if not gate_enabled:
|
||||
logger.warning(
|
||||
"Rate limit encountered for %s %s; retry_after=%s",
|
||||
method,
|
||||
@@ -1130,6 +1159,17 @@ class Downloader:
|
||||
error_msg,
|
||||
retry_after=retry_after,
|
||||
)
|
||||
# The coordinator logs the cooldown notice (INFO once
|
||||
# per window, DEBUG on extension).
|
||||
coordinator.register_rate_limit(destination, retry_after)
|
||||
attempt += 1
|
||||
if attempt >= max_resend_attempts:
|
||||
error = RateLimitError(error_msg, retry_after=retry_after)
|
||||
error.gate_handled = True
|
||||
return False, error
|
||||
# Loop back: wait_for_slot blocks until the cooldown
|
||||
# elapses (or raises once the wait exceeds the cap).
|
||||
continue
|
||||
else:
|
||||
return False, f"Request failed with status {response.status}"
|
||||
|
||||
|
||||
@@ -245,16 +245,23 @@ class MetadataSyncService:
|
||||
civitai_api_not_found = False
|
||||
any_rate_limited = False
|
||||
|
||||
skip_network_providers = False
|
||||
for provider_name, provider in provider_attempts:
|
||||
if skip_network_providers and provider_name != "sqlite":
|
||||
# A network provider was already rate-limited; failing
|
||||
# over to another network provider just spreads the flood
|
||||
# (#1085). The local sqlite archive stays as last resort.
|
||||
continue
|
||||
try:
|
||||
civitai_metadata_candidate, error = await provider.get_model_by_hash(sha256)
|
||||
except RateLimitError as exc:
|
||||
logger.warning(
|
||||
"Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
|
||||
"Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers",
|
||||
provider_name or provider.__class__.__name__,
|
||||
exc.retry_after or 0,
|
||||
)
|
||||
any_rate_limited = True
|
||||
skip_network_providers = True
|
||||
continue
|
||||
except Exception as exc: # pragma: no cover - defensive logging
|
||||
logger.error("Provider %s failed for hash %s: %s", provider_name, sha256, exc)
|
||||
|
||||
@@ -66,6 +66,14 @@ class _RateLimitRetryHelper:
|
||||
except RateLimitError as exc:
|
||||
attempt += 1
|
||||
|
||||
# The downloader's rate-limit gate already applied the wait
|
||||
# policy for this request (waited out the vendor window or
|
||||
# deliberately refused because it exceeds the cap). Sleeping
|
||||
# again here would double the wait — just propagate.
|
||||
if getattr(exc, "gate_handled", False):
|
||||
exc.provider = exc.provider or label
|
||||
raise
|
||||
|
||||
# Determine effective retry limit based on rate-limit magnitude
|
||||
effective_retry_limit = self._retry_limit # default: 3
|
||||
if exc.retry_after is not None and exc.retry_after >= 120.0:
|
||||
@@ -101,6 +109,12 @@ class _RateLimitRetryHelper:
|
||||
|
||||
return min(self._max_delay, max(0.0, base_delay))
|
||||
|
||||
|
||||
# Labels of providers that are free to consult even while a network provider
|
||||
# is rate-limited (local lookups, no vendor cost).
|
||||
_LOCAL_PROVIDER_LABELS = frozenset({"sqlite"})
|
||||
|
||||
|
||||
class ModelMetadataProvider(ABC):
|
||||
"""Base abstract class for all model metadata providers"""
|
||||
|
||||
@@ -451,7 +465,14 @@ class SQLiteModelMetadataProvider(ModelMetadataProvider):
|
||||
return None
|
||||
|
||||
class FallbackMetadataProvider(ModelMetadataProvider):
|
||||
"""Try providers in order, return first successful result."""
|
||||
"""Try providers in order, return first successful result.
|
||||
|
||||
Rate-limit policy (#1085): once a *network* provider raises
|
||||
``RateLimitError``, the chain stops consulting further network providers —
|
||||
failing over would just spread the flood to the next vendor. Local-only
|
||||
providers (see ``_LOCAL_PROVIDER_LABELS``) are still allowed as a last
|
||||
resort because they cost the vendor nothing.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -486,7 +507,10 @@ class FallbackMetadataProvider(ModelMetadataProvider):
|
||||
)
|
||||
|
||||
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
|
||||
rate_limited = False
|
||||
for provider, label in self._iter_providers():
|
||||
if rate_limited and label not in _LOCAL_PROVIDER_LABELS:
|
||||
continue
|
||||
try:
|
||||
result, error = await self._call_with_rate_limit(
|
||||
label,
|
||||
@@ -496,8 +520,9 @@ class FallbackMetadataProvider(ModelMetadataProvider):
|
||||
if result:
|
||||
return result, error
|
||||
except RateLimitError as exc:
|
||||
rate_limited = True
|
||||
logger.warning(
|
||||
"Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
|
||||
"Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers",
|
||||
label,
|
||||
exc.retry_after or 0,
|
||||
)
|
||||
@@ -505,11 +530,18 @@ class FallbackMetadataProvider(ModelMetadataProvider):
|
||||
except Exception as e:
|
||||
logger.debug("Provider %s failed for get_model_by_hash: %s", label, e)
|
||||
continue
|
||||
if rate_limited:
|
||||
# Distinct from "Model not found": callers must not mistake a
|
||||
# rate-limited lookup for a confirmed deletion.
|
||||
return None, "Rate limited"
|
||||
return None, "Model not found"
|
||||
|
||||
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
|
||||
not_found_confirmed = False
|
||||
rate_limited = False
|
||||
for provider, label in self._iter_providers():
|
||||
if rate_limited and label not in _LOCAL_PROVIDER_LABELS:
|
||||
continue
|
||||
try:
|
||||
result = await self._call_with_rate_limit(
|
||||
label,
|
||||
@@ -519,8 +551,9 @@ class FallbackMetadataProvider(ModelMetadataProvider):
|
||||
if result:
|
||||
return result
|
||||
except RateLimitError as exc:
|
||||
rate_limited = True
|
||||
logger.warning(
|
||||
"Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
|
||||
"Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers",
|
||||
label,
|
||||
exc.retry_after or 0,
|
||||
)
|
||||
@@ -539,7 +572,10 @@ class FallbackMetadataProvider(ModelMetadataProvider):
|
||||
return None
|
||||
|
||||
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
|
||||
rate_limited = False
|
||||
for provider, label in self._iter_providers():
|
||||
if rate_limited and label not in _LOCAL_PROVIDER_LABELS:
|
||||
continue
|
||||
try:
|
||||
result = await self._call_with_rate_limit(
|
||||
label,
|
||||
@@ -550,8 +586,9 @@ class FallbackMetadataProvider(ModelMetadataProvider):
|
||||
if result:
|
||||
return result
|
||||
except RateLimitError as exc:
|
||||
rate_limited = True
|
||||
logger.warning(
|
||||
"Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
|
||||
"Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers",
|
||||
label,
|
||||
exc.retry_after or 0,
|
||||
)
|
||||
@@ -562,7 +599,10 @@ class FallbackMetadataProvider(ModelMetadataProvider):
|
||||
return None
|
||||
|
||||
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
|
||||
rate_limited = False
|
||||
for provider, label in self._iter_providers():
|
||||
if rate_limited and label not in _LOCAL_PROVIDER_LABELS:
|
||||
continue
|
||||
try:
|
||||
result, error = await self._call_with_rate_limit(
|
||||
label,
|
||||
@@ -572,8 +612,9 @@ class FallbackMetadataProvider(ModelMetadataProvider):
|
||||
if result:
|
||||
return result, error
|
||||
except RateLimitError as exc:
|
||||
rate_limited = True
|
||||
logger.warning(
|
||||
"Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
|
||||
"Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers",
|
||||
label,
|
||||
exc.retry_after or 0,
|
||||
)
|
||||
@@ -581,12 +622,17 @@ class FallbackMetadataProvider(ModelMetadataProvider):
|
||||
except Exception as e:
|
||||
logger.debug("Provider %s failed for get_model_version_info: %s", label, e)
|
||||
continue
|
||||
if rate_limited:
|
||||
return None, "Rate limited"
|
||||
return None, "No provider could retrieve the data"
|
||||
|
||||
async def get_model_versions_by_hashes(
|
||||
self, hashes: List[str]
|
||||
) -> Optional[List[Dict[str, Any]]]:
|
||||
rate_limited = False
|
||||
for provider, label in self._iter_providers():
|
||||
if rate_limited and label not in _LOCAL_PROVIDER_LABELS:
|
||||
continue
|
||||
try:
|
||||
result = await self._call_with_rate_limit(
|
||||
label,
|
||||
@@ -598,8 +644,9 @@ class FallbackMetadataProvider(ModelMetadataProvider):
|
||||
except NotImplementedError:
|
||||
continue
|
||||
except RateLimitError as exc:
|
||||
rate_limited = True
|
||||
logger.warning(
|
||||
"Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
|
||||
"Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers",
|
||||
label,
|
||||
exc.retry_after or 0,
|
||||
)
|
||||
@@ -614,7 +661,10 @@ class FallbackMetadataProvider(ModelMetadataProvider):
|
||||
return None
|
||||
|
||||
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict[str, Any]]:
|
||||
rate_limited = False
|
||||
for provider, label in self._iter_providers():
|
||||
if rate_limited and label not in _LOCAL_PROVIDER_LABELS:
|
||||
continue
|
||||
try:
|
||||
result = await self._call_with_rate_limit(
|
||||
label,
|
||||
@@ -625,8 +675,9 @@ class FallbackMetadataProvider(ModelMetadataProvider):
|
||||
if result is not None:
|
||||
return result
|
||||
except RateLimitError as exc:
|
||||
rate_limited = True
|
||||
logger.warning(
|
||||
"Provider %s is rate-limited (retry_after=%.0fs); skipping to next provider",
|
||||
"Provider %s is rate-limited (retry_after=%.0fs); not failing over to other network providers",
|
||||
label,
|
||||
exc.retry_after or 0,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,213 @@
|
||||
"""Process-wide, per-destination rate-limit gate for outbound API traffic.
|
||||
|
||||
Implements the pacing/gating layer designed in
|
||||
``docs/plans/issue-1085-rate-limit-design.md``:
|
||||
|
||||
- **Reactive gate**: a 429 response arms ``next_allowed_send`` from the
|
||||
vendor's ``Retry-After`` (or exponential backoff when the header is
|
||||
missing); subsequent requests to the same destination wait out the window.
|
||||
- **Preemptive pacing**: a minimum inter-request interval per destination
|
||||
spaces consecutive sends so bursts never form in the first place.
|
||||
- **Herd-free**: waiters are serialized through a per-destination lock, so
|
||||
each one claims a distinct send slot instead of thousands of coroutines
|
||||
waking up together.
|
||||
- **Bounded**: waits longer than ``rate_limit_max_wait_seconds`` are refused
|
||||
by raising :class:`RateLimitError`, leaving the final decision to callers.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Dict, Optional
|
||||
|
||||
from .errors import RateLimitError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_MIN_INTERVAL_SECONDS = 0.75
|
||||
DEFAULT_MAX_WAIT_SECONDS = 300.0
|
||||
BASE_BACKOFF_SECONDS = 30.0
|
||||
MAX_BACKOFF_SECONDS = 1800.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class _DestinationState:
|
||||
"""Rate-limit bookkeeping for one destination (hostname)."""
|
||||
|
||||
next_allowed_send: float = 0.0 # time.monotonic() timestamp
|
||||
consecutive_429: int = 0
|
||||
last_send_at: float = 0.0 # time.monotonic() timestamp
|
||||
lock: asyncio.Lock = field(default_factory=asyncio.Lock)
|
||||
|
||||
|
||||
class RateLimitCoordinator:
|
||||
"""Coordinates outbound request pacing per destination.
|
||||
|
||||
Singleton mirroring :class:`ConnectivityGuard`'s pattern. All waits are
|
||||
bounded by the ``rate_limit_max_wait_seconds`` setting; when the required
|
||||
wait exceeds the cap, :meth:`wait_for_slot` raises :class:`RateLimitError`
|
||||
instead of parking the caller.
|
||||
"""
|
||||
|
||||
_instance: "RateLimitCoordinator | None" = None
|
||||
_instance_lock = asyncio.Lock()
|
||||
|
||||
@classmethod
|
||||
async def get_instance(cls) -> "RateLimitCoordinator":
|
||||
async with cls._instance_lock:
|
||||
if cls._instance is None:
|
||||
cls._instance = cls()
|
||||
return cls._instance
|
||||
|
||||
def __init__(self) -> None:
|
||||
if hasattr(self, "_initialized"):
|
||||
return
|
||||
self._initialized = True
|
||||
self._states: Dict[str, _DestinationState] = {}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Settings (read live so settings edits apply without a restart)
|
||||
|
||||
@staticmethod
|
||||
def _setting(key: str, default):
|
||||
try:
|
||||
from .settings_manager import get_settings_manager
|
||||
|
||||
return get_settings_manager().get(key, default)
|
||||
except Exception: # pragma: no cover - defensive: settings unavailable
|
||||
return default
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
return bool(self._setting("rate_limit_gate_enabled", True))
|
||||
|
||||
@property
|
||||
def min_interval_seconds(self) -> float:
|
||||
try:
|
||||
return max(0.0, float(self._setting("rate_limit_min_interval_seconds", DEFAULT_MIN_INTERVAL_SECONDS)))
|
||||
except (TypeError, ValueError):
|
||||
return DEFAULT_MIN_INTERVAL_SECONDS
|
||||
|
||||
@property
|
||||
def max_wait_seconds(self) -> float:
|
||||
try:
|
||||
return max(0.0, float(self._setting("rate_limit_max_wait_seconds", DEFAULT_MAX_WAIT_SECONDS)))
|
||||
except (TypeError, ValueError):
|
||||
return DEFAULT_MAX_WAIT_SECONDS
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# State helpers
|
||||
|
||||
@staticmethod
|
||||
def _normalize(destination: Optional[str]) -> str:
|
||||
if destination is None or not destination.strip():
|
||||
return "__global__"
|
||||
return destination.lower().strip()
|
||||
|
||||
def _state_for(self, destination: Optional[str]) -> _DestinationState:
|
||||
key = self._normalize(destination)
|
||||
if key not in self._states:
|
||||
self._states[key] = _DestinationState()
|
||||
return self._states[key]
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Drop all per-destination state. Test seam."""
|
||||
self._states.clear()
|
||||
|
||||
def in_cooldown(self, destination: Optional[str] = None) -> bool:
|
||||
return self.remaining_seconds(destination) > 0
|
||||
|
||||
def remaining_seconds(self, destination: Optional[str] = None) -> float:
|
||||
state = self._state_for(destination)
|
||||
return max(0.0, state.next_allowed_send - time.monotonic())
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Gate operations
|
||||
|
||||
async def wait_for_slot(self, destination: Optional[str] = None) -> None:
|
||||
"""Block until this caller may send the next request to *destination*.
|
||||
|
||||
Waits for both the rate-limit cooldown (``next_allowed_send``) and the
|
||||
minimum inter-request interval (``last_send_at + min_interval``).
|
||||
Waiters queue on the per-destination lock, so concurrent callers are
|
||||
spaced out instead of stampeding when a cooldown expires.
|
||||
|
||||
Raises:
|
||||
RateLimitError: when the required wait exceeds
|
||||
``rate_limit_max_wait_seconds``.
|
||||
"""
|
||||
state = self._state_for(destination)
|
||||
deadline = time.monotonic() + self.max_wait_seconds
|
||||
async with state.lock:
|
||||
now = time.monotonic()
|
||||
wake_at = max(
|
||||
state.next_allowed_send,
|
||||
state.last_send_at + self.min_interval_seconds,
|
||||
)
|
||||
if wake_at > deadline:
|
||||
raise RateLimitError(
|
||||
f"Rate limit wait for '{self._normalize(destination)}' "
|
||||
f"exceeds the {self.max_wait_seconds:.0f}s cap",
|
||||
retry_after=wake_at - now,
|
||||
)
|
||||
delay = wake_at - now
|
||||
if delay > 0:
|
||||
logger.debug(
|
||||
"Rate-limit gate: pacing request to '%s' by %.2fs",
|
||||
self._normalize(destination),
|
||||
delay,
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
state.last_send_at = time.monotonic()
|
||||
|
||||
def register_rate_limit(
|
||||
self,
|
||||
destination: Optional[str],
|
||||
retry_after: Optional[float] = None,
|
||||
) -> float:
|
||||
"""Record a 429 for *destination* and arm the cooldown window.
|
||||
|
||||
Honors the vendor's ``Retry-After`` when present; otherwise grows an
|
||||
exponential backoff (30s base, doubling per consecutive 429, capped at
|
||||
1800s). Returns the cooldown duration in seconds.
|
||||
"""
|
||||
state = self._state_for(destination)
|
||||
state.consecutive_429 += 1
|
||||
if retry_after is not None and retry_after > 0:
|
||||
backoff = min(MAX_BACKOFF_SECONDS, float(retry_after))
|
||||
else:
|
||||
backoff = min(
|
||||
MAX_BACKOFF_SECONDS,
|
||||
BASE_BACKOFF_SECONDS * (2 ** (state.consecutive_429 - 1)),
|
||||
)
|
||||
now = time.monotonic()
|
||||
already_cooling = state.next_allowed_send > now
|
||||
state.next_allowed_send = max(state.next_allowed_send, now + backoff)
|
||||
if already_cooling:
|
||||
logger.debug(
|
||||
"Rate-limit cooldown for '%s' extended by %.0fs (consecutive_429=%d)",
|
||||
self._normalize(destination),
|
||||
backoff,
|
||||
state.consecutive_429,
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
"Rate limited by '%s'; pausing requests for %.0fs",
|
||||
self._normalize(destination),
|
||||
backoff,
|
||||
)
|
||||
return backoff
|
||||
|
||||
def register_success(self, destination: Optional[str]) -> None:
|
||||
"""Reset rate-limit state after a successful request.
|
||||
|
||||
A 200 proves the vendor is accepting traffic again, so any armed
|
||||
cooldown window is cleared alongside the backoff counter (mirrors
|
||||
``ConnectivityGuard.register_success`` semantics).
|
||||
"""
|
||||
state = self._state_for(destination)
|
||||
state.consecutive_429 = 0
|
||||
state.next_allowed_send = 0.0
|
||||
@@ -70,6 +70,9 @@ DEFAULT_SETTINGS: Dict[str, Any] = {
|
||||
"enable_metadata_archive_db": False,
|
||||
"enable_civarchive_api": True,
|
||||
"metadata_provider_order": "civitai_archive_sqlite",
|
||||
"rate_limit_gate_enabled": True,
|
||||
"rate_limit_max_wait_seconds": 300,
|
||||
"rate_limit_min_interval_seconds": 0.75,
|
||||
"proxy_enabled": False,
|
||||
"proxy_host": "",
|
||||
"proxy_port": "",
|
||||
|
||||
@@ -812,3 +812,56 @@ async def test_fetch_and_update_model_does_not_overwrite_api_metadata_with_archi
|
||||
|
||||
helpers.metadata_manager.save_metadata.assert_awaited()
|
||||
update_cache.assert_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_and_update_model_keeps_sqlite_last_resort_after_civarchive_rate_limit(tmp_path):
|
||||
"""A CivArchive 429 must not block the local sqlite last resort (#1085)."""
|
||||
civarchive_provider = SimpleNamespace(
|
||||
get_model_by_hash=AsyncMock(
|
||||
side_effect=RateLimitError("limited", retry_after=30)
|
||||
),
|
||||
get_model_version=AsyncMock(),
|
||||
)
|
||||
sqlite_payload = {
|
||||
"source": "archive_db",
|
||||
"model": {"name": "Recovered", "description": "", "tags": []},
|
||||
"images": [],
|
||||
"baseModel": "sdxl",
|
||||
}
|
||||
sqlite_provider = SimpleNamespace(
|
||||
get_model_by_hash=AsyncMock(return_value=(sqlite_payload, None)),
|
||||
get_model_version=AsyncMock(),
|
||||
)
|
||||
|
||||
async def select_provider(name: str):
|
||||
if name == "civarchive_api":
|
||||
return civarchive_provider
|
||||
if name == "sqlite":
|
||||
return sqlite_provider
|
||||
raise AssertionError(f"unexpected provider request: {name}")
|
||||
|
||||
helpers = build_service(
|
||||
settings_values={"enable_metadata_archive_db": True},
|
||||
provider_selector=AsyncMock(side_effect=select_provider),
|
||||
)
|
||||
|
||||
model_path = tmp_path / "model.safetensors"
|
||||
model_data = {
|
||||
"civitai_deleted": True,
|
||||
"db_checked": False,
|
||||
"file_path": str(model_path),
|
||||
}
|
||||
update_cache = AsyncMock()
|
||||
|
||||
ok, error = await helpers.service.fetch_and_update_model(
|
||||
sha256="cafe",
|
||||
file_path=str(model_path),
|
||||
model_data=model_data,
|
||||
update_cache_func=update_cache,
|
||||
)
|
||||
|
||||
assert ok and error is None
|
||||
civarchive_provider.get_model_by_hash.assert_awaited_once()
|
||||
sqlite_provider.get_model_by_hash.assert_awaited_once()
|
||||
assert model_data["metadata_source"] == "archive_db"
|
||||
|
||||
@@ -101,7 +101,9 @@ async def test_fallback_retries_same_provider_on_rate_limit(monkeypatch):
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fallback_continues_to_next_provider_on_rate_limit(monkeypatch):
|
||||
"""After exhausting retries on primary, fallback should continue to secondary."""
|
||||
"""#1085: a rate-limited network provider no longer fails over to another
|
||||
network provider (that just spreads the flood); local providers such as
|
||||
sqlite remain as a last resort."""
|
||||
sleep_mock = AsyncMock()
|
||||
monkeypatch.setattr(provider_module.asyncio, "sleep", sleep_mock)
|
||||
monkeypatch.setattr(provider_module.random, "uniform", lambda *_: 0.0)
|
||||
@@ -114,13 +116,26 @@ async def test_fallback_continues_to_next_provider_on_rate_limit(monkeypatch):
|
||||
rate_limit_retry_limit=2,
|
||||
)
|
||||
|
||||
# After Change A: no longer raises; falls through to secondary
|
||||
result, error = await fallback.get_model_by_hash("abc")
|
||||
|
||||
# Secondary is a network provider: it must NOT be consulted after the 429.
|
||||
assert result is None
|
||||
assert error == "Rate limited"
|
||||
assert primary.calls == 2 # retry_limit exhausted on primary
|
||||
assert secondary.calls == 0 # no network failover
|
||||
|
||||
# A local sqlite provider behind the rate-limited one is still allowed.
|
||||
sqlite = TrackingProvider()
|
||||
fallback = FallbackMetadataProvider(
|
||||
[("primary", AlwaysRateLimitedProvider()), ("sqlite", sqlite)],
|
||||
rate_limit_retry_limit=2,
|
||||
)
|
||||
|
||||
result, error = await fallback.get_model_by_hash("abc")
|
||||
|
||||
assert error is None
|
||||
assert result == {"id": "secondary"}
|
||||
assert primary.calls == 2 # retry_limit exhausted on primary
|
||||
assert secondary.calls == 1 # secondary IS called now
|
||||
assert sqlite.calls == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -0,0 +1,400 @@
|
||||
"""Tests for the per-destination rate-limit gate (#1085).
|
||||
|
||||
Covers the RateLimitCoordinator itself, its integration into
|
||||
``Downloader.make_request``, the failover semantics change in
|
||||
``FallbackMetadataProvider``, and the ``_RateLimitRetryHelper`` double-wait
|
||||
fix.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from datetime import datetime
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, Dict, Optional
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from py.services.connectivity_guard import ConnectivityGuard
|
||||
from py.services.downloader import Downloader
|
||||
from py.services.errors import RateLimitError
|
||||
from py.services.model_metadata_provider import (
|
||||
FallbackMetadataProvider,
|
||||
_RateLimitRetryHelper,
|
||||
)
|
||||
from py.services.rate_limit_coordinator import RateLimitCoordinator
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_singletons():
|
||||
RateLimitCoordinator._instance = None
|
||||
ConnectivityGuard._instance = None
|
||||
yield
|
||||
RateLimitCoordinator._instance = None
|
||||
ConnectivityGuard._instance = None
|
||||
|
||||
|
||||
def _patch_gate_settings(monkeypatch, **overrides):
|
||||
"""Override the coordinator's settings reads for the test."""
|
||||
monkeypatch.setattr(
|
||||
RateLimitCoordinator,
|
||||
"_setting",
|
||||
staticmethod(lambda key, default: overrides.get(key, default)),
|
||||
)
|
||||
|
||||
|
||||
async def _make_coordinator(monkeypatch, **overrides) -> RateLimitCoordinator:
|
||||
_patch_gate_settings(monkeypatch, **overrides)
|
||||
return await RateLimitCoordinator.get_instance()
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Coordinator unit tests
|
||||
|
||||
|
||||
async def test_pacing_enforces_min_interval(monkeypatch):
|
||||
coordinator = await _make_coordinator(
|
||||
monkeypatch, rate_limit_min_interval_seconds=0.1
|
||||
)
|
||||
start = time.monotonic()
|
||||
await coordinator.wait_for_slot("example.com")
|
||||
await coordinator.wait_for_slot("example.com")
|
||||
elapsed = time.monotonic() - start
|
||||
assert elapsed >= 0.1
|
||||
|
||||
|
||||
async def test_pacing_is_per_destination(monkeypatch):
|
||||
coordinator = await _make_coordinator(
|
||||
monkeypatch, rate_limit_min_interval_seconds=0.2
|
||||
)
|
||||
await coordinator.wait_for_slot("a.example.com")
|
||||
start = time.monotonic()
|
||||
await coordinator.wait_for_slot("b.example.com")
|
||||
elapsed = time.monotonic() - start
|
||||
assert elapsed < 0.1
|
||||
|
||||
|
||||
async def test_register_rate_limit_arms_cooldown_and_waits(monkeypatch):
|
||||
coordinator = await _make_coordinator(
|
||||
monkeypatch,
|
||||
rate_limit_min_interval_seconds=0.0,
|
||||
rate_limit_max_wait_seconds=5.0,
|
||||
)
|
||||
coordinator.register_rate_limit("example.com", retry_after=0.15)
|
||||
assert coordinator.in_cooldown("example.com")
|
||||
assert 0.1 < coordinator.remaining_seconds("example.com") <= 0.15
|
||||
|
||||
start = time.monotonic()
|
||||
await coordinator.wait_for_slot("example.com")
|
||||
elapsed = time.monotonic() - start
|
||||
assert elapsed >= 0.14
|
||||
assert not coordinator.in_cooldown("example.com")
|
||||
|
||||
|
||||
async def test_concurrent_waiters_share_one_cooldown_window(monkeypatch):
|
||||
"""Herd test: N waiters wake after ~one window, not N windows."""
|
||||
coordinator = await _make_coordinator(
|
||||
monkeypatch,
|
||||
rate_limit_min_interval_seconds=0.0,
|
||||
rate_limit_max_wait_seconds=5.0,
|
||||
)
|
||||
coordinator.register_rate_limit("example.com", retry_after=0.2)
|
||||
|
||||
start = time.monotonic()
|
||||
await asyncio.gather(
|
||||
*(coordinator.wait_for_slot("example.com") for _ in range(4))
|
||||
)
|
||||
elapsed = time.monotonic() - start
|
||||
# 4 independent windows would take ~0.8s; a shared window is ~0.2s.
|
||||
assert 0.19 <= elapsed < 0.5
|
||||
|
||||
|
||||
async def test_backoff_grows_on_consecutive_429_and_resets_on_success(
|
||||
monkeypatch,
|
||||
):
|
||||
coordinator = await _make_coordinator(
|
||||
monkeypatch, rate_limit_min_interval_seconds=0.0
|
||||
)
|
||||
coordinator.register_rate_limit("example.com", retry_after=None)
|
||||
first = coordinator.remaining_seconds("example.com")
|
||||
assert 29.0 < first <= 30.0
|
||||
|
||||
coordinator.register_rate_limit("example.com", retry_after=None)
|
||||
second = coordinator.remaining_seconds("example.com")
|
||||
assert 59.0 < second <= 60.0
|
||||
|
||||
coordinator.register_success("example.com")
|
||||
coordinator.register_rate_limit("example.com", retry_after=None)
|
||||
third = coordinator.remaining_seconds("example.com")
|
||||
assert 29.0 < third <= 30.0
|
||||
|
||||
|
||||
async def test_wait_beyond_cap_raises_rate_limit_error(monkeypatch):
|
||||
coordinator = await _make_coordinator(
|
||||
monkeypatch,
|
||||
rate_limit_min_interval_seconds=0.0,
|
||||
rate_limit_max_wait_seconds=0.05,
|
||||
)
|
||||
coordinator.register_rate_limit("example.com", retry_after=30.0)
|
||||
|
||||
start = time.monotonic()
|
||||
with pytest.raises(RateLimitError) as excinfo:
|
||||
await coordinator.wait_for_slot("example.com")
|
||||
elapsed = time.monotonic() - start
|
||||
assert elapsed < 1.0 # refused immediately instead of parking
|
||||
assert excinfo.value.retry_after is not None
|
||||
assert excinfo.value.retry_after > 1.0
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Downloader integration tests
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(
|
||||
self,
|
||||
status: int,
|
||||
payload: Optional[Dict[str, Any]] = None,
|
||||
headers: Optional[Dict[str, str]] = None,
|
||||
):
|
||||
self.status = status
|
||||
self._payload = payload
|
||||
self.headers = headers or {}
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
async def json(self):
|
||||
if self._payload is None:
|
||||
raise ValueError("no json payload")
|
||||
return self._payload
|
||||
|
||||
async def text(self):
|
||||
return ""
|
||||
|
||||
|
||||
class _FakeSession:
|
||||
def __init__(self, responses):
|
||||
self._responses = list(responses)
|
||||
self.requests = []
|
||||
|
||||
def request(self, method, url, headers=None, **kwargs):
|
||||
self.requests.append({"method": method, "url": url})
|
||||
assert self._responses, "unexpected extra request"
|
||||
return self._responses.pop(0)
|
||||
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
|
||||
def _build_downloader(responses) -> Downloader:
|
||||
downloader = Downloader()
|
||||
fake_session = _FakeSession(responses)
|
||||
downloader._session = fake_session # pyright: ignore[reportAttributeAccessIssue]
|
||||
downloader._session_created_at = datetime.now()
|
||||
downloader._proxy_url = None
|
||||
|
||||
async def _noop_create_session():
|
||||
downloader._session = fake_session # pyright: ignore[reportAttributeAccessIssue]
|
||||
downloader._session_created_at = datetime.now()
|
||||
downloader._proxy_url = None
|
||||
|
||||
downloader._create_session = _noop_create_session # type: ignore[assignment]
|
||||
return downloader
|
||||
|
||||
|
||||
async def test_make_request_waits_out_429_then_resends(monkeypatch):
|
||||
_patch_gate_settings(
|
||||
monkeypatch,
|
||||
rate_limit_gate_enabled=True,
|
||||
rate_limit_min_interval_seconds=0.0,
|
||||
rate_limit_max_wait_seconds=5.0,
|
||||
)
|
||||
downloader = _build_downloader(
|
||||
[
|
||||
_FakeResponse(429, headers={"Retry-After": "1"}),
|
||||
_FakeResponse(200, payload={"ok": True}),
|
||||
]
|
||||
)
|
||||
|
||||
start = time.monotonic()
|
||||
success, payload = await downloader.make_request(
|
||||
"GET", "https://api.example.com/models/1"
|
||||
)
|
||||
elapsed = time.monotonic() - start
|
||||
|
||||
assert success is True
|
||||
assert payload == {"ok": True}
|
||||
assert len(downloader._session.requests) == 2
|
||||
assert elapsed >= 0.9
|
||||
|
||||
|
||||
async def test_make_request_paces_consecutive_calls(monkeypatch):
|
||||
_patch_gate_settings(
|
||||
monkeypatch,
|
||||
rate_limit_gate_enabled=True,
|
||||
rate_limit_min_interval_seconds=0.15,
|
||||
rate_limit_max_wait_seconds=5.0,
|
||||
)
|
||||
downloader = _build_downloader(
|
||||
[_FakeResponse(200, payload={}), _FakeResponse(200, payload={})]
|
||||
)
|
||||
|
||||
start = time.monotonic()
|
||||
await downloader.make_request("GET", "https://api.example.com/a")
|
||||
await downloader.make_request("GET", "https://api.example.com/b")
|
||||
elapsed = time.monotonic() - start
|
||||
assert elapsed >= 0.14
|
||||
|
||||
|
||||
async def test_make_request_gate_disabled_returns_429_immediately(monkeypatch):
|
||||
_patch_gate_settings(monkeypatch, rate_limit_gate_enabled=False)
|
||||
downloader = _build_downloader(
|
||||
[_FakeResponse(429, headers={"Retry-After": "30"})]
|
||||
)
|
||||
|
||||
start = time.monotonic()
|
||||
success, payload = await downloader.make_request(
|
||||
"GET", "https://api.example.com/models/1"
|
||||
)
|
||||
elapsed = time.monotonic() - start
|
||||
|
||||
assert success is False
|
||||
assert isinstance(payload, RateLimitError)
|
||||
assert payload.retry_after == 30.0
|
||||
# Gate was off: the error is NOT marked, so retry helpers keep their
|
||||
# legacy behavior.
|
||||
assert getattr(payload, "gate_handled", False) is False
|
||||
assert len(downloader._session.requests) == 1
|
||||
assert elapsed < 1.0
|
||||
|
||||
|
||||
async def test_make_request_refuses_wait_beyond_cap(monkeypatch):
|
||||
_patch_gate_settings(
|
||||
monkeypatch,
|
||||
rate_limit_gate_enabled=True,
|
||||
rate_limit_min_interval_seconds=0.0,
|
||||
rate_limit_max_wait_seconds=0.2,
|
||||
)
|
||||
downloader = _build_downloader(
|
||||
[_FakeResponse(429, headers={"Retry-After": "3600"})]
|
||||
)
|
||||
|
||||
start = time.monotonic()
|
||||
success, payload = await downloader.make_request(
|
||||
"GET", "https://api.example.com/models/1"
|
||||
)
|
||||
elapsed = time.monotonic() - start
|
||||
|
||||
assert success is False
|
||||
assert isinstance(payload, RateLimitError)
|
||||
assert payload.gate_handled is True
|
||||
assert len(downloader._session.requests) == 1
|
||||
assert elapsed < 1.0
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# FallbackMetadataProvider failover semantics (Fix C)
|
||||
|
||||
|
||||
def _stub_provider(*, result=None, error=None, exc: Exception | None = None):
|
||||
if exc is not None:
|
||||
call = AsyncMock(side_effect=exc)
|
||||
else:
|
||||
call = AsyncMock(return_value=(result, error))
|
||||
return SimpleNamespace(get_model_by_hash=call)
|
||||
|
||||
|
||||
async def test_fallback_does_not_fail_over_to_network_provider_on_429():
|
||||
civitai = _stub_provider(exc=RateLimitError("limited", retry_after=30))
|
||||
civarchive = _stub_provider(result={"id": 1}, error=None)
|
||||
sqlite = _stub_provider(result=None, error="not in archive")
|
||||
|
||||
fallback = FallbackMetadataProvider(
|
||||
[
|
||||
("civitai_api", civitai),
|
||||
("civarchive_api", civarchive),
|
||||
("sqlite", sqlite),
|
||||
]
|
||||
)
|
||||
|
||||
result, error = await fallback.get_model_by_hash("deadbeef")
|
||||
|
||||
assert result is None
|
||||
assert error == "Rate limited"
|
||||
civarchive.get_model_by_hash.assert_not_called() # no network failover
|
||||
sqlite.get_model_by_hash.assert_called_once() # local last resort kept
|
||||
|
||||
|
||||
async def test_fallback_still_fails_over_on_not_found():
|
||||
civitai = _stub_provider(result=None, error="Model not found")
|
||||
civarchive = _stub_provider(result={"id": 1}, error=None)
|
||||
|
||||
fallback = FallbackMetadataProvider(
|
||||
[("civitai_api", civitai), ("civarchive_api", civarchive)]
|
||||
)
|
||||
|
||||
result, _ = await fallback.get_model_by_hash("deadbeef")
|
||||
|
||||
assert result == {"id": 1}
|
||||
civarchive.get_model_by_hash.assert_called_once()
|
||||
|
||||
|
||||
async def test_fallback_404_failover_still_works_after_rate_limit_change():
|
||||
"""A 404 from the first network provider still reaches the second."""
|
||||
civitai = _stub_provider(result=None, error="Resource not found")
|
||||
civarchive = _stub_provider(result={"id": 2}, error=None)
|
||||
|
||||
fallback = FallbackMetadataProvider(
|
||||
[("civitai_api", civitai), ("civarchive_api", civarchive)]
|
||||
)
|
||||
|
||||
result, _ = await fallback.get_model_by_hash("deadbeef")
|
||||
assert result == {"id": 2}
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# _RateLimitRetryHelper double-wait fix
|
||||
|
||||
|
||||
async def test_retry_helper_does_not_sleep_for_gate_handled_errors():
|
||||
calls = 0
|
||||
|
||||
async def failing():
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
error = RateLimitError("limited", retry_after=30)
|
||||
error.gate_handled = True
|
||||
raise error
|
||||
|
||||
helper = _RateLimitRetryHelper()
|
||||
start = time.monotonic()
|
||||
with pytest.raises(RateLimitError) as excinfo:
|
||||
await helper.run("civitai_api", failing)
|
||||
elapsed = time.monotonic() - start
|
||||
|
||||
assert calls == 1 # propagated immediately, no retry loop
|
||||
assert elapsed < 1.0
|
||||
assert excinfo.value.provider == "civitai_api"
|
||||
|
||||
|
||||
async def test_retry_helper_keeps_legacy_retry_for_ungated_errors():
|
||||
calls = 0
|
||||
|
||||
async def failing():
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
raise RateLimitError("limited", retry_after=None)
|
||||
|
||||
helper = _RateLimitRetryHelper(
|
||||
retry_limit=2, base_delay=0.01, max_delay=0.05, jitter_ratio=0.0
|
||||
)
|
||||
with pytest.raises(RateLimitError):
|
||||
await helper.run("civitai_api", failing)
|
||||
|
||||
assert calls == 2 # legacy retry behavior unchanged
|
||||
Reference in New Issue
Block a user