fix(llm): add failure cooldown and lock for model catalog fetch

This commit is contained in:
Will Miao
2026-09-07 10:17:48 +08:00
parent 5ae4aef30e
commit a7995db009
2 changed files with 152 additions and 47 deletions
+82 -47
View File
@@ -11,6 +11,7 @@ from __future__ import annotations
import asyncio
import json
import logging
import time
from typing import Any, Dict, List, Optional
import aiohttp
@@ -32,6 +33,16 @@ _catalog_cache: Optional[Dict[str, List[str]]] = None
# ``{provider_id: {model_id: max_output_tokens}}``.
_model_output_limits: Dict[str, Dict[str, int]] = {}
# Monotonic timestamp of the last failed catalog fetch (None = no failure
# yet). Failed fetches are negatively cached: further calls return the
# empty fallback without hitting the network until the cooldown elapses,
# so users on broken networks don't stall on every settings-modal open.
_catalog_last_failure: Optional[float] = None
_CATALOG_FAILURE_COOLDOWN = 600.0 # seconds
# Serializes catalog fetches so concurrent callers don't duplicate requests.
_catalog_lock = asyncio.Lock()
_CATALOG_TIMEOUT = aiohttp.ClientTimeout(total=30)
# Cloudflare serves brotli when the client advertises it, and brotli is a
@@ -54,61 +65,85 @@ async def _load_model_catalog() -> Dict[str, List[str]]:
value has a ``models`` sub-dict keyed by model ID. The result is cached
in memory after the first successful fetch.
Subsequent calls return the cached data immediately.
Failed fetches are negatively cached: further calls return an empty
dict without hitting the network until ``_CATALOG_FAILURE_COOLDOWN``
has elapsed, so a broken network does not stall every settings-modal
open. Concurrent callers are serialized behind :data:`_catalog_lock`
so only one request is ever in flight.
"""
global _catalog_cache, _model_output_limits
global _catalog_cache, _model_output_limits, _catalog_last_failure
if _catalog_cache is not None:
return _catalog_cache
try:
async with aiohttp.ClientSession(timeout=_CATALOG_TIMEOUT) as session:
async with session.get(_MODEL_CATALOG_URL, headers=_NO_BROTLI_HEADERS) as resp:
if resp.status != 200:
logger.warning("Model catalog returned HTTP %s", resp.status)
return _catalog_cache or {}
data = await resp.json()
except (aiohttp.ClientError, asyncio.TimeoutError, json.JSONDecodeError, UnicodeDecodeError) as exc:
logger.warning("Failed to fetch model catalog: %s", exc)
return _catalog_cache or {}
async with _catalog_lock:
# Re-check under the lock: another caller may have fetched (or
# failed) while we were waiting.
if _catalog_cache is not None:
return _catalog_cache
if (
_catalog_last_failure is not None
and time.monotonic() - _catalog_last_failure < _CATALOG_FAILURE_COOLDOWN
):
logger.debug(
"Skipping model catalog fetch: last attempt failed %.0fs ago",
time.monotonic() - _catalog_last_failure,
)
return {}
if not isinstance(data, dict):
logger.warning("Model catalog is not a dict, got %s", type(data).__name__)
return _catalog_cache or {}
try:
async with aiohttp.ClientSession(timeout=_CATALOG_TIMEOUT) as session:
async with session.get(_MODEL_CATALOG_URL, headers=_NO_BROTLI_HEADERS) as resp:
if resp.status != 200:
logger.warning("Model catalog returned HTTP %s", resp.status)
_catalog_last_failure = time.monotonic()
return {}
data = await resp.json()
except (aiohttp.ClientError, asyncio.TimeoutError, json.JSONDecodeError, UnicodeDecodeError) as exc:
logger.warning("Failed to fetch model catalog: %s", exc)
_catalog_last_failure = time.monotonic()
return {}
result: Dict[str, List[str]] = {}
output_limits: Dict[str, Dict[str, int]] = {}
for provider_id, provider_info in data.items():
if not isinstance(provider_info, dict):
continue
models_dict = provider_info.get("models")
if not isinstance(models_dict, dict):
continue
model_ids: List[str] = []
provider_limits: Dict[str, int] = {}
for mid, model_info in models_dict.items():
if not isinstance(mid, str):
if not isinstance(data, dict):
logger.warning("Model catalog is not a dict, got %s", type(data).__name__)
_catalog_last_failure = time.monotonic()
return {}
result: Dict[str, List[str]] = {}
output_limits: Dict[str, Dict[str, int]] = {}
for provider_id, provider_info in data.items():
if not isinstance(provider_info, dict):
continue
model_ids.append(mid)
if isinstance(model_info, dict):
limit = model_info.get("limit")
if isinstance(limit, dict):
output = limit.get("output")
if isinstance(output, (int, float)) and output > 0:
provider_limits[mid] = int(output)
if model_ids:
result[provider_id] = model_ids
if provider_limits:
output_limits[provider_id] = provider_limits
models_dict = provider_info.get("models")
if not isinstance(models_dict, dict):
continue
model_ids: List[str] = []
provider_limits: Dict[str, int] = {}
for mid, model_info in models_dict.items():
if not isinstance(mid, str):
continue
model_ids.append(mid)
if isinstance(model_info, dict):
limit = model_info.get("limit")
if isinstance(limit, dict):
output = limit.get("output")
if isinstance(output, (int, float)) and output > 0:
provider_limits[mid] = int(output)
if model_ids:
result[provider_id] = model_ids
if provider_limits:
output_limits[provider_id] = provider_limits
_catalog_cache = result
_model_output_limits = output_limits
logger.debug(
"Loaded model catalog: %d providers, %d total models "
"(%d providers have output limits)",
len(result),
sum(len(m) for m in result.values()),
len(output_limits),
)
return result
_catalog_cache = result
_model_output_limits = output_limits
logger.debug(
"Loaded model catalog: %d providers, %d total models "
"(%d providers have output limits)",
len(result),
sum(len(m) for m in result.values()),
len(output_limits),
)
return result
def _get_model_max_output(provider: str, model: str) -> Optional[int]:
+70
View File
@@ -4,6 +4,7 @@ from __future__ import annotations
import asyncio
import json
import time
from unittest import mock
import pytest
@@ -344,6 +345,14 @@ class CorruptJsonResponse(MockResponse):
raise UnicodeDecodeError("utf-8", b"\x9a", 0, 1, "invalid start byte")
class SlowResponse(MockResponse):
"""Response whose body takes a moment to read, to force contention."""
async def json(self):
await asyncio.sleep(0.05)
return self._json_data
class TestModelCatalog:
"""Tests for _load_model_catalog / fetch_ollama_models error handling."""
@@ -352,9 +361,11 @@ class TestModelCatalog:
"""Reset the module-level catalog cache around each test."""
llm_module._catalog_cache = None
llm_module._model_output_limits = {}
llm_module._catalog_last_failure = None
yield
llm_module._catalog_cache = None
llm_module._model_output_limits = {}
llm_module._catalog_last_failure = None
@pytest.mark.asyncio
async def test_load_model_catalog_falls_back_on_unicode_decode_error(self):
@@ -401,3 +412,62 @@ class TestModelCatalog:
assert models == ["llama3"]
assert session.last_headers == {"Accept-Encoding": "gzip, deflate"}
@pytest.mark.asyncio
async def test_failed_fetch_is_negatively_cached(self):
"""A failed fetch is not retried until the cooldown elapses."""
created = []
def factory(*args, **kwargs):
session = MockGetSession(MockResponse(500, text_data="error"))
created.append(session)
return session
with mock.patch("aiohttp.ClientSession", side_effect=factory):
first = await llm_module._load_model_catalog()
second = await llm_module._load_model_catalog()
assert first == {}
assert second == {}
assert len(created) == 1
assert llm_module._catalog_last_failure is not None
@pytest.mark.asyncio
async def test_fetch_retries_after_cooldown(self):
"""Once the cooldown elapses, the next call fetches again."""
bad = MockGetSession(MockResponse(500, text_data="error"))
with mock.patch("aiohttp.ClientSession", return_value=bad):
assert await llm_module._load_model_catalog() == {}
# Simulate the cooldown having elapsed.
llm_module._catalog_last_failure = (
time.monotonic() - llm_module._CATALOG_FAILURE_COOLDOWN - 1
)
good = MockGetSession(
MockResponse(200, json_data={"openai": {"models": {"gpt-4o": {}}}})
)
with mock.patch("aiohttp.ClientSession", return_value=good):
catalog = await llm_module._load_model_catalog()
assert catalog == {"openai": ["gpt-4o"]}
@pytest.mark.asyncio
async def test_concurrent_fetches_are_deduplicated(self):
"""Concurrent callers share a single in-flight fetch."""
created = []
def factory(*args, **kwargs):
session = MockGetSession(
SlowResponse(200, json_data={"openai": {"models": {"gpt-4o": {}}}})
)
created.append(session)
return session
with mock.patch("aiohttp.ClientSession", side_effect=factory):
results = await asyncio.gather(
*(llm_module._load_model_catalog() for _ in range(3))
)
assert len(created) == 1
assert all(r == {"openai": ["gpt-4o"]} for r in results)