diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 67d06a0758e..f81ddbfee2e 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -184,10 +184,6 @@ def mark_litellm_import_complete() -> None: _litellm_import_complete.set() -def _start_daemon_thread(fn: Callable[[], None]) -> None: - threading.Thread(target=fn, name="litellm-model-cost-map-retry", daemon=True).start() - - @dataclass(frozen=True, slots=True) class ModelCostMapReloaded: model_cost_map: dict # mutable-ok: adopted as litellm.model_cost, whose consumer contract is a plain mutable dict @@ -531,32 +527,6 @@ def _finalize_loaded_model_cost_map(loaded: ModelCostMapReloaded) -> ModelCostMa return replace(loaded, model_cost_map=_finalize_model_cost_map(loaded.model_cost_map)) -def _use_local_backup(reason: str | None) -> dict: # mutable-ok: returns the mutable model-cost map contract - _cost_map_source_info.source = "local" - _cost_map_source_info.fallback_reason = reason - return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map - - -def _accept_remote( - result: ModelCostMapReloaded, url: str -) -> dict | None: # mutable-ok: returns the mutable model-cost map contract - if not GetModelCostMap.validate_model_cost_map( - fetched_map=result.model_cost_map, - backup_model_count=GetModelCostMap._get_backup_model_count(), # pyright: ignore[reportPrivateUsage] # integrity cache - ): - verbose_logger.warning( - "LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s", - url, - ) - return None - finalized: Final = _finalize_loaded_model_cost_map(result).model_cost_map - _cost_map_source_info.source = "remote" - _cost_map_source_info.url = url - _cost_map_source_info.is_env_forced = False - _cost_map_source_info.fallback_reason = None - return finalized - - def adopt_model_cost_map( new_model_cost_map: dict, # mutable-ok: public API preserves the mutable cost-map contract ) -> int: @@ -579,7 +549,6 @@ def _retry_remote_fetch_in_background( rng: random.Random, client: _SyncGetClient, first_outcome: _FetchAttemptRetryable, - apply: Callable[[dict], object], # mutable-ok: injected callback receives the mutable cost-map dict ) -> None: try: first_wait: Final = _next_retry_wait(outcome=first_outcome, attempt=1, max_attempts=max_attempts, rng=rng) @@ -602,12 +571,20 @@ def _retry_remote_fetch_in_background( ) return _litellm_import_complete.wait() - accepted: Final = _accept_remote(result, url) - if accepted is None: - _cost_map_source_info.fallback_reason = "Remote data failed integrity validation" + if not GetModelCostMap.validate_model_cost_map( + fetched_map=result.model_cost_map, + backup_model_count=GetModelCostMap._get_backup_model_count(), # pyright: ignore[reportPrivateUsage] # integrity cache + ): + verbose_logger.warning( + "LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s", + url, + ) return - apply(accepted) + finalized: Final = _finalize_loaded_model_cost_map(result).model_cost_map + _cost_map_source_info.source = "remote" + _cost_map_source_info.fallback_reason = None _cost_map_source_info.loaded_at = datetime.now(timezone.utc) + adopt_model_cost_map(finalized) except Exception as e: verbose_logger.warning("LiteLLM: Background model cost map retry failed: %s", e) @@ -619,19 +596,12 @@ def get_model_cost_map( sleep: Callable[[float], None] = time.sleep, rng: random.Random | None = None, client: "_SyncGetClient | None" = None, - start_background: Callable[[Callable[[], None]], None] = _start_daemon_thread, - apply: Callable[ # mutable-ok: injected callback receives the mutable cost-map dict - [dict], - object, - ] = adopt_model_cost_map, ) -> 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 the first response, and falls - back to the local backup while retrying transient HTTP errors in the - background. + 2. Otherwise fetches from ``url``, retrying transient errors in a background thread. 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 @@ -641,9 +611,11 @@ def get_model_cost_map( # Note: can't use get_secret_bool here — this runs during litellm.__init__ # before litellm._key_management_settings is set. if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true": + _cost_map_source_info.source = "local" _cost_map_source_info.url = None _cost_map_source_info.is_env_forced = True - return _use_local_backup(None) + _cost_map_source_info.fallback_reason = None + return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map _cost_map_source_info.url = url _cost_map_source_info.is_env_forced = False @@ -651,32 +623,45 @@ def get_model_cost_map( fetch_client: Final = client if client is not None else httpx fetch_rng: Final = rng if rng is not None else random.Random() outcome: Final = _attempt_fetch_sync(client=fetch_client, url=url, timeout=timeout) - if isinstance(outcome, ModelCostMapReloaded): - accepted: Final = _accept_remote(outcome, url) - return accepted if accepted is not None else _use_local_backup("Remote data failed integrity validation") if isinstance(outcome, _FetchAttemptRetryable) and max_attempts > 1: - verbose_logger.warning( - "LiteLLM: model cost map fetch attempt 1/%d failed: %s; " - "using local backup while retrying in the background", - max_attempts, - outcome.reason, - ) - start_background( - lambda: _retry_remote_fetch_in_background( - url=url, - timeout=timeout, - max_attempts=max_attempts, - sleep=sleep, - rng=fetch_rng, - client=fetch_client, - first_outcome=outcome, - apply=apply, - ) - ) - else: + threading.Thread( + target=_retry_remote_fetch_in_background, + kwargs={ # mutable-ok: threading requires a mutable keyword-arguments mapping + "url": url, + "timeout": timeout, + "max_attempts": max_attempts, + "sleep": sleep, + "rng": fetch_rng, + "client": fetch_client, + "first_outcome": outcome, + }, + name="litellm-model-cost-map-retry", + daemon=True, + ).start() + if not isinstance(outcome, ModelCostMapReloaded): verbose_logger.warning( "LiteLLM: Failed to fetch remote model cost map from %s: %s. Falling back to local backup.", url, outcome.reason, ) - return _use_local_backup(f"Remote fetch failed: {outcome.reason}") + _cost_map_source_info.source = "local" + _cost_map_source_info.fallback_reason = f"Remote fetch failed: {outcome.reason}" + return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map + content: Final = outcome.model_cost_map + + # Validate using cached count (cheap int comparison, no file I/O) + if not GetModelCostMap.validate_model_cost_map( + fetched_map=content, + backup_model_count=GetModelCostMap._get_backup_model_count(), + ): + verbose_logger.warning( + "LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s", + url, + ) + _cost_map_source_info.source = "local" + _cost_map_source_info.fallback_reason = "Remote data failed integrity validation" + return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map + + _cost_map_source_info.source = "remote" + _cost_map_source_info.fallback_reason = None + return _finalize_loaded_model_cost_map(outcome).model_cost_map diff --git a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py index e39658a8f37..266c2ca1465 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py @@ -6,6 +6,7 @@ count actual model entries, not reserved meta keys) and the extraction of the import json import os +import threading import pytest @@ -20,7 +21,6 @@ from litellm.litellm_core_utils.get_model_cost_map import ( GetModelCostMap, _count_model_entries, _finalize_model_cost_map, - adopt_model_cost_map, get_model_cost_map_provenance, git_blob_id, ) @@ -566,194 +566,125 @@ from litellm.litellm_core_utils.get_model_cost_map import ( class _SyncSleepRecorder: """Injected in place of time.sleep so the boot path's waits are asserted without delay.""" - def __init__(self): + def __init__(self, block=False): self.waits = [] + self.block = block + self.started = threading.Event() + self.release = threading.Event() def __call__(self, seconds: float) -> None: + if self.block: + self.started.set() + self.release.wait(timeout=10) self.waits.append(seconds) -class _BackgroundRecorder: - def __init__(self): - self.callbacks = [] - - def __call__(self, callback): - self.callbacks.append(callback) +def _retry_threads(): + return [thread for thread in threading.enumerate() if thread.name == "litellm-model-cost-map-retry"] -def test_boot_load_returns_local_map_and_schedules_transient_retry(): - client, calls = _mock_client( - [ - httpx.ConnectError("connection refused"), - ], - client_cls=httpx.Client, - ) - sleeper = _SyncSleepRecorder() - background = _BackgroundRecorder() - - cost_map = get_model_cost_map( - url=_URL, - sleep=sleeper, - rng=random.Random(0), - client=client, - start_background=background, - ) - assert calls["count"] == 1 - assert sleeper.waits == [] - assert len(background.callbacks) == 1 - assert len(cost_map) > 100 - source = get_model_cost_map_source_info() - assert source["source"] == "local" - assert source["fallback_reason"] is not None - - -def test_background_retry_adopts_valid_remote_map(): - client, calls = _mock_client( - [ - httpx.ConnectError("connection refused"), - httpx.Response(200, content=_real_map_bytes()), - ], - client_cls=httpx.Client, - ) - sleeper = _SyncSleepRecorder() - background = _BackgroundRecorder() - applied = [] - - get_model_cost_map( - url=_URL, - sleep=sleeper, - rng=random.Random(0), - client=client, - start_background=background, - apply=applied.append, - ) - assert calls["count"] == 1 - assert sleeper.waits == [] - assert len(background.callbacks) == 1 - - background.callbacks[0]() - - assert calls["count"] == 2 - assert len(sleeper.waits) == 1 - assert 2.0 <= sleeper.waits[0] < 3.0 - assert len(applied) == 1 - assert applied[0].keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY} - source = get_model_cost_map_source_info() - assert source["source"] == "remote" - assert source["fallback_reason"] is None - - -def test_background_retry_keeps_local_map_after_remaining_failures(): - client, calls = _mock_client( - [httpx.ConnectError("connection refused"), httpx.Response(503)], - client_cls=httpx.Client, - ) - sleeper = _SyncSleepRecorder() - background = _BackgroundRecorder() - applied = [] - - get_model_cost_map( - url=_URL, - sleep=sleeper, - rng=random.Random(0), - client=client, - start_background=background, - apply=applied.append, - ) - background.callbacks[0]() - - 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 - assert applied == [] - assert get_model_cost_map_source_info()["source"] == "local" - - -def test_boot_load_does_not_schedule_non_retryable_failure(): - client, calls = _mock_client([httpx.Response(404)], client_cls=httpx.Client) - sleeper = _SyncSleepRecorder() - background = _BackgroundRecorder() - - get_model_cost_map( - url=_URL, - sleep=sleeper, - rng=random.Random(0), - client=client, - start_background=background, - ) - - assert calls["count"] == 1 - assert sleeper.waits == [] - assert background.callbacks == [] - source = get_model_cost_map_source_info() - assert source["source"] == "local" - assert source["fallback_reason"] is not None - - -def test_boot_load_success_does_not_schedule_background_retry(): +def test_boot_load_success_does_not_start_background_retry(): client, calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client) sleeper = _SyncSleepRecorder() - background = _BackgroundRecorder() cost_map = get_model_cost_map( url=_URL, sleep=sleeper, rng=random.Random(0), client=client, - start_background=background, ) assert calls["count"] == 1 assert sleeper.waits == [] - assert background.callbacks == [] - assert get_model_cost_map_source_info()["source"] == "remote" + assert _retry_threads() == [] assert cost_map.keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY} + assert get_model_cost_map_source_info()["source"] == "remote" -def test_boot_load_with_one_attempt_does_not_schedule_background_retry(): - client, calls = _mock_client([httpx.ConnectError("connection refused")], client_cls=httpx.Client) - sleeper = _SyncSleepRecorder() - background = _BackgroundRecorder() +def test_boot_load_transient_failure_returns_local_then_background_retry_adopts_remote(monkeypatch): + import litellm + from litellm import utils as litellm_utils + from litellm.litellm_core_utils import get_model_cost_map as module - get_model_cost_map( + original_model_cost = litellm.model_cost + monkeypatch.setattr(litellm, "model_cost", dict(original_model_cost)) + for name, provider_models in tuple(vars(litellm).items()): + if name.endswith("_models") and isinstance(provider_models, set): + monkeypatch.setattr(litellm, name, set(provider_models)) + monkeypatch.setattr(litellm, "models_by_provider", dict(litellm.models_by_provider)) + monkeypatch.setattr( + litellm_utils, + "_runtime_registered_model_cost", + dict(litellm_utils._runtime_registered_model_cost), + ) + source_info = module._cost_map_source_info + for name in ("source", "url", "is_env_forced", "fallback_reason", "loaded_at", "source_revision", "etag"): + monkeypatch.setattr(source_info, name, getattr(source_info, name)) + + remote_map = _load_root_cost_map() + remote_map["claude-remote-only-test"] = {"litellm_provider": "anthropic", "mode": "chat"} + client, calls = _mock_client( + [ + httpx.ConnectError("connection refused"), + httpx.Response(200, content=json.dumps(remote_map).encode()), + ], + client_cls=httpx.Client, + ) + sleeper = _SyncSleepRecorder(block=True) + litellm.register_model({"my-runtime-model": {"litellm_provider": "custom", "max_input_tokens": 4321}}) + + cost_map = get_model_cost_map( url=_URL, - max_attempts=1, + max_attempts=3, sleep=sleeper, rng=random.Random(0), client=client, - start_background=background, ) assert calls["count"] == 1 assert sleeper.waits == [] - assert background.callbacks == [] - assert get_model_cost_map_source_info()["source"] == "local" - - -def test_adopt_model_cost_map_replays_runtime_registration_and_provider_models(): - import litellm - from litellm import utils as litellm_utils - - original_model_cost = litellm.model_cost - original_registry = dict(litellm_utils._runtime_registered_model_cost) - original_anthropic_models = set(litellm.anthropic_models) + assert sleeper.started.wait(timeout=10) + threads = _retry_threads() try: - litellm.register_model( - model_cost={"custom/deployment-model": {"litellm_provider": "custom", "max_input_tokens": 4321}} - ) - - models_count = adopt_model_cost_map({"anthropic/new-model": {"litellm_provider": "anthropic", "mode": "chat"}}) - - assert models_count == 1 - assert "anthropic/new-model" in litellm.anthropic_models - assert litellm.model_cost["custom/deployment-model"]["max_input_tokens"] == 4321 + assert len(threads) == 1 + assert "claude-remote-only-test" not in cost_map + assert cost_map.keys() == _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map()).keys() + source = get_model_cost_map_source_info() + assert source["source"] == "local" + assert source["fallback_reason"].startswith("Remote fetch failed:") + sleeper.release.set() + for thread in threads: + thread.join(timeout=10) + assert all(not thread.is_alive() for thread in threads) + assert sleeper.waits and 2.0 <= sleeper.waits[0] < 3.0 + assert calls["count"] == 2 + assert "claude-remote-only-test" in litellm.model_cost + assert "claude-remote-only-test" in litellm.anthropic_models + assert "my-runtime-model" in litellm.model_cost + source = get_model_cost_map_source_info() + assert source["source"] == "remote" + assert source["fallback_reason"] is None finally: - litellm.model_cost = original_model_cost # test-quality-ok: restore the module state changed by adoption - litellm_utils._runtime_registered_model_cost.clear() - litellm_utils._runtime_registered_model_cost.update(original_registry) - litellm.anthropic_models.clear() - litellm.anthropic_models.update(original_anthropic_models) - litellm_utils._invalidate_model_cost_lowercase_map() + sleeper.release.set() + for thread in _retry_threads(): + thread.join(timeout=10) + + +def test_boot_load_does_not_retry_non_retryable_failure(): + 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 _retry_threads() == [] + source = get_model_cost_map_source_info() + assert source["source"] == "local" + assert source["fallback_reason"] is not None def test_boot_load_respects_local_env_override(monkeypatch):