fix(cost-map): retry transient boot fetch failures and recover config deployments dropped by a stale cost map (#39230)

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-01 17:49:06 -07:00 committed by GitHub
parent 69029c139e
commit 04a25083a6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 291 additions and 44 deletions

View file

@ -12,6 +12,7 @@ import asyncio
import json
import os
import random
import time
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from datetime import datetime, timezone
@ -154,18 +155,6 @@ class GetModelCostMap:
return True
@staticmethod
def fetch_remote_model_cost_map(url: str, timeout: int = 5) -> dict:
"""
Fetch the model cost map from a remote URL.
Returns the parsed JSON dict. Raises on network/parse errors
(caller is expected to handle).
"""
response: Final = httpx.get(url, timeout=timeout)
response.raise_for_status()
return response.json()
RETRYABLE_FETCH_STATUS_CODES: Final = frozenset({429, 500, 502, 503, 504})
MODEL_COST_MAP_FETCH_MAX_ATTEMPTS: Final = 3
@ -212,6 +201,13 @@ class _AsyncGetClient(Protocol):
def get(self, url: str, *, timeout: float | None = None) -> Awaitable[httpx.Response]: ...
class _SyncGetClient(Protocol):
def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: ...
_FetchAttemptOutcome = ModelCostMapReloaded | ModelCostMapReloadUnavailable | _FetchAttemptRetryable
def _default_reload_client() -> _AsyncGetClient:
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.llms.custom_http import httpxSpecialProvider
@ -219,13 +215,30 @@ def _default_reload_client() -> _AsyncGetClient:
return get_async_httpx_client(llm_provider=httpxSpecialProvider.ModelCostMap)
async def _attempt_fetch(
client: _AsyncGetClient, url: str, timeout: int
) -> ModelCostMapReloaded | ModelCostMapReloadUnavailable | _FetchAttemptRetryable:
def _classify_fetch_error(error: httpx.HTTPError | httpx.InvalidURL, url: str) -> _FetchAttemptOutcome:
reason: Final = f"{type(error).__name__} fetching {url}: {error}"
if isinstance(error, (httpx.InvalidURL, httpx.UnsupportedProtocol)):
return ModelCostMapReloadUnavailable(reason=reason)
return _FetchAttemptRetryable(reason=reason, retry_after_seconds=None)
async def _attempt_fetch(client: _AsyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome:
try:
response: Final = await client.get(url, timeout=timeout)
except httpx.HTTPError as e:
return _FetchAttemptRetryable(reason=f"{type(e).__name__} fetching {url}: {e}", retry_after_seconds=None)
except (httpx.HTTPError, httpx.InvalidURL) as e:
return _classify_fetch_error(e, url)
return _classify_fetch_response(response, url)
def _attempt_fetch_sync(client: _SyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome:
try:
response: Final = client.get(url, timeout=timeout)
except (httpx.HTTPError, httpx.InvalidURL) as e:
return _classify_fetch_error(e, url)
return _classify_fetch_response(response, url)
def _classify_fetch_response(response: httpx.Response, url: str) -> _FetchAttemptOutcome:
if response.status_code in RETRYABLE_FETCH_STATUS_CODES:
return _FetchAttemptRetryable(
reason=f"HTTP {response.status_code} from {url}",
@ -242,6 +255,22 @@ async def _attempt_fetch(
return ModelCostMapReloaded(model_cost_map=parsed)
def _next_retry_wait(
outcome: _FetchAttemptRetryable, attempt: int, max_attempts: int, rng: random.Random
) -> float | ModelCostMapReloadUnavailable:
if attempt == max_attempts:
return ModelCostMapReloadUnavailable(reason=f"{outcome.reason} (after {max_attempts} attempts)")
wait_seconds: Final = _retry_wait_seconds(outcome=outcome, attempt=attempt, rng=rng)
verbose_logger.warning(
"LiteLLM: model cost map fetch attempt %d/%d failed (%s); retrying in %.1fs",
attempt,
max_attempts,
outcome.reason,
wait_seconds,
)
return wait_seconds
async def _fetch_remote_model_cost_map_with_retry(
url: str,
timeout: int,
@ -254,20 +283,32 @@ async def _fetch_remote_model_cost_map_with_retry(
outcome = await _attempt_fetch(client=client, url=url, timeout=timeout)
if not isinstance(outcome, _FetchAttemptRetryable):
return outcome
if attempt == max_attempts:
return ModelCostMapReloadUnavailable(reason=f"{outcome.reason} (after {max_attempts} attempts)")
wait_seconds = _retry_wait_seconds(outcome=outcome, attempt=attempt, rng=rng)
verbose_logger.warning(
"LiteLLM: model cost map fetch attempt %d/%d failed (%s); retrying in %.1fs",
attempt,
max_attempts,
outcome.reason,
wait_seconds,
)
wait_seconds = _next_retry_wait(outcome=outcome, attempt=attempt, max_attempts=max_attempts, rng=rng)
if isinstance(wait_seconds, ModelCostMapReloadUnavailable):
return wait_seconds
await sleep(wait_seconds)
return ModelCostMapReloadUnavailable(reason="model cost map fetch failed")
def _fetch_remote_model_cost_map_with_retry_sync(
url: str,
timeout: int,
max_attempts: int,
sleep: Callable[[float], None],
rng: random.Random,
client: _SyncGetClient,
) -> ModelCostMapReloadResult:
for attempt in range(1, max_attempts + 1):
outcome = _attempt_fetch_sync(client=client, url=url, timeout=timeout)
if not isinstance(outcome, _FetchAttemptRetryable):
return outcome
wait_seconds = _next_retry_wait(outcome=outcome, attempt=attempt, max_attempts=max_attempts, rng=rng)
if isinstance(wait_seconds, ModelCostMapReloadUnavailable):
return wait_seconds
sleep(wait_seconds)
return ModelCostMapReloadUnavailable(reason="model cost map fetch failed")
async def refetch_model_cost_map(
url: str,
timeout: int = 5,
@ -423,13 +464,21 @@ def _finalize_model_cost_map(model_cost: dict) -> dict:
return _expand_model_aliases(model_cost)
def get_model_cost_map(url: str) -> dict:
def get_model_cost_map(
url: str,
timeout: int = 5,
max_attempts: int = MODEL_COST_MAP_FETCH_MAX_ATTEMPTS,
sleep: Callable[[float], None] = time.sleep,
rng: random.Random | None = None,
client: "_SyncGetClient | None" = None,
) -> dict:
"""
Public entry point returns the model cost map dict.
1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set, uses the local backup only.
2. Otherwise fetches from ``url``, validates integrity, and falls back
to the local backup on any failure.
2. Otherwise fetches from ``url``, retrying transient HTTP errors
(429/5xx/transport) with Retry-After-aware backoff, validates
integrity, and falls back to the local backup on any failure.
Only the backup model count is cached (a single int) for validation.
The full backup dict is only parsed when it must be *returned* as a
@ -448,17 +497,24 @@ def get_model_cost_map(url: str) -> dict:
_cost_map_source_info.url = url
_cost_map_source_info.is_env_forced = False
try:
content: Final = GetModelCostMap.fetch_remote_model_cost_map(url)
except Exception as e:
result: Final = _fetch_remote_model_cost_map_with_retry_sync(
url=url,
timeout=timeout,
max_attempts=max_attempts,
sleep=sleep,
rng=rng if rng is not None else random.Random(),
client=client if client is not None else httpx,
)
if isinstance(result, ModelCostMapReloadUnavailable):
verbose_logger.warning(
"LiteLLM: Failed to fetch remote model cost map from %s: %s. Falling back to local backup.",
url,
str(e),
result.reason,
)
_cost_map_source_info.source = "local"
_cost_map_source_info.fallback_reason = f"Remote fetch failed: {e}"
_cost_map_source_info.fallback_reason = f"Remote fetch failed: {result.reason}"
return _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map())
content: Final = result.model_cost_map
# Validate using cached count (cheap int comparison, no file I/O)
if not GetModelCostMap.validate_model_cost_map(

View file

@ -22,7 +22,7 @@ import traceback
import weakref
from collections import defaultdict
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Mapping, Sequence
from functools import lru_cache
from functools import lru_cache, partial
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast
@ -825,6 +825,7 @@ class Router:
self._zero_cost_cache: dict[str, bool] = {}
self._routing_group_rows: tuple[DeploymentTypedDict, ...] | None = None
self._init_routing_groups(None)
self._provider_unresolved_deployments: tuple[Callable[[], Deployment | None], ...] = ()
self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds
self.model_group_affinity_config = model_group_affinity_config
@ -8472,6 +8473,19 @@ class Router:
return deployment
except Exception as e:
if self.ignore_invalid_deployments:
if isinstance(e, litellm.BadRequestError):
self._provider_unresolved_deployments = (
*self._provider_unresolved_deployments,
partial(
self._create_deployment,
deployment_info=deployment_info,
_model_name=_model_name,
_litellm_params=_litellm_params,
_model_info=_model_info,
declared_id=declared_id,
duplicate_ids=duplicate_ids,
),
)
verbose_router_logger.exception(
"Error creating deployment: %s, ignoring and continuing with other deployments.", e
)
@ -8901,6 +8915,7 @@ class Router:
self.quality_routers = {}
self.complexity_routers = {}
self.auto_routers = {}
self._provider_unresolved_deployments = ()
self._invalidate_model_group_info_cache()
self._invalidate_access_groups_cache()
# we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works
@ -9523,8 +9538,12 @@ class Router:
"""Re-assert this router's deployments onto a freshly fetched catalog.
Reads ``model_list`` at call time, so only deployments the router still
serves are restored.
serves are restored, plus any config deployment the fresh catalog now resolves.
"""
provider_unresolved: Final = self._provider_unresolved_deployments
self._provider_unresolved_deployments = ()
for create_deployment in provider_unresolved:
create_deployment()
for entry in tuple(self.model_list):
try:
deployment = entry if isinstance(entry, Deployment) else Deployment(**entry)

View file

@ -256,14 +256,12 @@ def test_get_model_cost_map_stamps_loaded_at(monkeypatch):
from litellm.litellm_core_utils import get_model_cost_map as module
monkeypatch.setattr(module._cost_map_source_info, "loaded_at", None)
monkeypatch.setattr(
module.GetModelCostMap,
"fetch_remote_model_cost_map",
staticmethod(lambda url, timeout=5: _load_root_cost_map()),
client, _calls = _mock_client(
[httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client
)
before = datetime.now(timezone.utc)
module.get_model_cost_map(url="https://example.invalid/cost_map.json")
module.get_model_cost_map(url="https://example.invalid/cost_map.json", client=client)
loaded_at = module.get_model_cost_map_loaded_at()
assert loaded_at is not None
@ -308,7 +306,7 @@ def _unset_local_cost_map_env(monkeypatch):
monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False)
def _mock_client(outcomes):
def _mock_client(outcomes, client_cls=httpx.AsyncClient):
"""httpx client over a MockTransport serving one outcome per request; an exception instance is raised."""
calls = {"count": 0}
@ -320,7 +318,7 @@ def _mock_client(outcomes):
raise outcome
return outcome
return httpx.AsyncClient(transport=httpx.MockTransport(handler)), calls
return client_cls(transport=httpx.MockTransport(handler)), calls
@pytest.mark.asyncio
@ -450,3 +448,97 @@ async def test_refetch_respects_local_env_override(monkeypatch):
)
assert isinstance(result, ModelCostMapReloaded)
assert len(result.model_cost_map) > 100
# ---------------------------------------------------------------------------
# get_model_cost_map: the boot-time load retries transient failures like a reload does
# ---------------------------------------------------------------------------
from litellm.litellm_core_utils.get_model_cost_map import (
get_model_cost_map,
get_model_cost_map_source_info,
)
class _SyncSleepRecorder:
"""Injected in place of time.sleep so the boot path's waits are asserted without delay."""
def __init__(self):
self.waits = []
def __call__(self, seconds: float) -> None:
self.waits.append(seconds)
def test_boot_load_retries_transient_failures_instead_of_falling_back():
"""A refused connection then a 503 at pod boot used to pin the process to the bundled
backup for its lifetime; both are transient and must be retried before giving up."""
client, calls = _mock_client(
[
httpx.ConnectError("connection refused"),
httpx.Response(503),
httpx.Response(200, content=_real_map_bytes()),
],
client_cls=httpx.Client,
)
sleeper = _SyncSleepRecorder()
cost_map = get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
assert calls["count"] == 3
assert len(sleeper.waits) == 2
assert 2.0 <= sleeper.waits[0] < 3.0
assert 4.0 <= sleeper.waits[1] < 5.0
source = get_model_cost_map_source_info()
assert source["source"] == "remote"
assert source["fallback_reason"] is None
assert cost_map.keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY}
def test_boot_load_honors_retry_after_then_falls_back_after_max_attempts():
"""An outage longer than the retry budget still ends on the bundled backup, and the
recorded fallback reason says how many attempts were spent so operators can tell."""
client, calls = _mock_client(
[httpx.Response(429, headers={"Retry-After": "7"})], client_cls=httpx.Client
)
sleeper = _SyncSleepRecorder()
cost_map = get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
assert calls["count"] == 3
assert sleeper.waits == [7.0, 7.0]
source = get_model_cost_map_source_info()
assert source["source"] == "local"
assert "after 3 attempts" in source["fallback_reason"]
assert len(cost_map) > 100
def test_boot_load_does_not_retry_permanent_failures():
"""A 404 or a malformed URL cannot heal by waiting: one attempt, no sleeps, backup."""
client, calls = _mock_client([httpx.Response(404)], client_cls=httpx.Client)
sleeper = _SyncSleepRecorder()
get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
assert calls["count"] == 1
assert sleeper.waits == []
assert get_model_cost_map_source_info()["source"] == "local"
get_model_cost_map(url="not a url", sleep=sleeper, rng=random.Random(0))
assert sleeper.waits == []
assert get_model_cost_map_source_info()["source"] == "local"
def test_boot_load_respects_local_env_override(monkeypatch):
"""LITELLM_LOCAL_MODEL_COST_MAP=True still short-circuits to the backup with zero HTTP."""
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
def _fail(request):
raise AssertionError("no HTTP request should be made when local map is forced")
cost_map = get_model_cost_map(
url=_URL,
sleep=_SyncSleepRecorder(),
client=httpx.Client(transport=httpx.MockTransport(_fail)),
)
assert len(cost_map) > 100
assert get_model_cost_map_source_info()["is_env_forced"] is True

View file

@ -2281,3 +2281,83 @@ def test_every_declaring_deployment_is_named(caplog):
assert "azure-ptu-east" in warnings[0]
assert "azure-ptu-west" in warnings[0]
assert "plain-gpt-4o" not in warnings[0]
def _simulate_price_data_reload_with_provider_sets(monkeypatch, fetched_catalog):
"""Like `_simulate_price_data_reload`, plus the provider model-set refresh the proxy's
`_swap_in_model_cost_map` does before replaying, so bare names in the new catalog resolve."""
monkeypatch.setattr(litellm, "model_cost", fetched_catalog)
_invalidate_model_cost_lowercase_map()
litellm.add_known_models(model_cost_map=fetched_catalog)
reapply_runtime_model_cost_registrations()
def test_a_config_deployment_dropped_by_a_stale_cost_map_comes_back_on_reload(monkeypatch):
"""
Booting on the bundled backup, a bare model that only the remote catalog knows
cannot be provider-resolved, so the proxy router (ignore_invalid_deployments) drops
it. Once a reload brings in a catalog that knows the model, the deployment must be
served again with its access groups, and exactly once however many reloads follow.
"""
backend = "lit-5766-only-in-remote-catalog"
try:
router = Router(
model_list=[
{
"model_name": "new-model",
"litellm_params": {"model": backend, "api_key": "k"},
"model_info": {"id": "new-id", "access_groups": ["team-models"]},
},
{
"model_name": "control-model",
"litellm_params": {"model": "hosted_vllm/control-backend", "api_key": "k"},
"model_info": {"id": "control-id", "access_groups": ["team-models"]},
},
],
ignore_invalid_deployments=True,
)
assert router.get_model_names() == ["control-model"]
assert router.get_model_access_groups(model_name="new-model") == {}
fresh_catalog = {**litellm.model_cost, backend: {"litellm_provider": "openai", "mode": "chat"}}
_simulate_price_data_reload_with_provider_sets(monkeypatch, fresh_catalog)
_simulate_price_data_reload_with_provider_sets(monkeypatch, fresh_catalog)
assert sorted(router.get_model_names()) == ["control-model", "new-model"]
assert router.get_model_access_groups(model_name="new-model") == {"team-models": ["new-model"]}
assert [d["model_info"]["id"] for d in router.model_list] == ["control-id", "new-id"]
assert "new-id" in litellm.model_cost
finally:
litellm.open_ai_chat_completion_models.discard(backend)
litellm.models_by_provider["openai"].discard(backend)
def test_a_config_deployment_dropped_for_a_permanent_reason_is_not_retried_on_reload(monkeypatch):
"""
Only provider-resolution drops can be healed by a fresh catalog. A deployment that
fails after its provider resolved (here a pass-through vertex entry with no project)
has already touched router state, so replaying it on every reload would leak into
`deployment_names` each time.
"""
router = Router(
model_list=[
{
"model_name": "vertex-passthrough",
"litellm_params": {"model": "vertex_ai/gemini-2.5-flash", "use_in_pass_through": True},
"model_info": {"id": "vertex-id"},
},
{
"model_name": "control-model",
"litellm_params": {"model": "hosted_vllm/control-backend", "api_key": "k"},
"model_info": {"id": "control-id"},
},
],
ignore_invalid_deployments=True,
)
assert router.get_model_names() == ["control-model"]
names_after_boot = list(router.deployment_names)
_simulate_price_data_reload_with_provider_sets(monkeypatch, dict(litellm.model_cost))
assert router.get_model_names() == ["control-model"]
assert router.deployment_names == names_after_boot