mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-28 08:21: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
@@ -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