From 897e29bb51bb1affc361a28937bb0e833f9b5baa Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 30 Sep 2026 18:52:17 -0700 Subject: [PATCH 01/18] perf(cost-map): fetch the remote cost map in the background instead of blocking import import litellm used to wait on a GitHub download of the model cost map, about 115 ms on a normal network and up to 5 s when GitHub is unreachable. It now returns the bundled map right away and swaps in the remote map from a background thread once it passes the integrity check, unless the caller has already replaced litellm.model_cost. --- litellm/__init__.py | 8 +- .../litellm_core_utils/get_model_cost_map.py | 137 +++++++++++++++--- .../test_get_model_cost_map.py | 136 ++++++++++++++++- 3 files changed, 253 insertions(+), 28 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 58827b60a98..d41b7d71126 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -558,9 +558,13 @@ _key_management_system: Optional["KeyManagementSystem"] = None #### PII MASKING #### output_parse_pii: bool = False ############################################# -from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map, mark_litellm_import_complete +from litellm.litellm_core_utils.get_model_cost_map import ( + get_model_cost_map, + load_model_cost_map_at_import, + mark_litellm_import_complete, +) -model_cost = get_model_cost_map(url=model_cost_map_url) +model_cost = load_model_cost_map_at_import(url=model_cost_map_url) cost_discount_config: Dict[str, float] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount cost_margin_config: Dict[ str, Union[float, Dict[str, float]] diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 5471fe50d5f..ffcd4f68353 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -567,6 +567,54 @@ def adopt_model_cost_map( return fetched_model_count +def _adopt_remote_model_cost_map( + url: str, + timeout: int, + attempts: range, + sleep: Callable[[float], None], + rng: random.Random, + client: _SyncGetClient, + replaces: Mapping[str, object] | None, +) -> None: + result: Final = _fetch_remote_model_cost_map_with_retry_sync( + url=url, + timeout=timeout, + attempts=attempts, + sleep=sleep, + rng=rng, + client=client, + ) + if isinstance(result, ModelCostMapReloadUnavailable): + verbose_logger.warning( + "LiteLLM: Failed to fetch remote model cost map from %s (%s); keeping local backup", + url, + result.reason, + ) + _cost_map_source_info.fallback_reason = f"Remote fetch failed: {result.reason}" + return + _litellm_import_complete.wait() + 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, + ) + _cost_map_source_info.fallback_reason = "Remote data failed integrity validation" + return + import litellm + + if replaces is not None and litellm.model_cost is not replaces: + verbose_logger.debug("LiteLLM: litellm.model_cost was replaced during import; not adopting the remote map") + return + 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) + + def _retry_remote_fetch_in_background( url: str, timeout: int, @@ -581,40 +629,83 @@ def _retry_remote_fetch_in_background( if isinstance(first_wait, ModelCostMapReloadUnavailable): return sleep(first_wait) - result: Final = _fetch_remote_model_cost_map_with_retry_sync( + _adopt_remote_model_cost_map( url=url, timeout=timeout, attempts=range(2, max_attempts + 1), sleep=sleep, rng=rng, client=client, + replaces=None, ) - if isinstance(result, ModelCostMapReloadUnavailable): - verbose_logger.warning( - "LiteLLM: Failed to fetch remote model cost map from %s after %d attempts; keeping local backup", - url, - max_attempts, - ) - return - _litellm_import_complete.wait() - 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 - 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: # noqa: BLE001 # a failed background retry must not kill the task; the backup stays verbose_logger.warning("LiteLLM: Background model cost map retry failed: %s", e) +def _fetch_remote_in_background_after_import( + url: str, + timeout: int, + max_attempts: int, + sleep: Callable[[float], None], + rng: random.Random, + client: _SyncGetClient, + replaces: Mapping[str, object], +) -> None: + try: + _adopt_remote_model_cost_map( + url=url, + timeout=timeout, + attempts=range(1, max_attempts + 1), + sleep=sleep, + rng=rng, + client=client, + replaces=replaces, + ) + except Exception as e: # noqa: BLE001 # a failed background fetch must not kill the thread; the backup stays + verbose_logger.warning("LiteLLM: Background model cost map fetch failed: %s", e) + + +def load_model_cost_map_at_import( + 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: + """ + Return the bundled map right away and fetch the remote one in a background thread, + so ``import litellm`` never waits on the network. The remote map replaces + ``litellm.model_cost`` once it passes the integrity check, unless the caller has + already replaced ``litellm.model_cost`` with their own map. + """ + if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true" or _is_cli_process(): + return get_model_cost_map(url=url, timeout=timeout) + _cost_map_source_info.loaded_at = datetime.now(timezone.utc) + _cost_map_source_info.source = "local" + _cost_map_source_info.url = url + _cost_map_source_info.is_env_forced = False + _cost_map_source_info.fallback_reason = "Remote fetch in progress" + bundled: Final = _finalize_loaded_model_cost_map( + GetModelCostMap.load_local_model_cost_map_with_revision() + ).model_cost_map + threading.Thread( + target=_fetch_remote_in_background_after_import, + kwargs={ # mutable-ok: threading requires a mutable keyword-arguments mapping + "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, + "replaces": bundled, + }, + name="litellm-model-cost-map-fetch", + daemon=True, + ).start() + return bundled + + def get_model_cost_map( url: str, timeout: int = 5, diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index 00977d9c3ee..2b1aef0e37f 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -559,13 +559,13 @@ def test_boot_load_success_does_not_start_background_retry(): assert get_model_cost_map_source_info()["source"] == "remote" -def test_boot_load_transient_failure_returns_local_then_background_retry_adopts_remote(monkeypatch): +@pytest.fixture +def isolated_litellm_cost_state(monkeypatch): import litellm from litellm import utils as litellm_utils from litellm.litellm_core_utils import get_model_cost_map as module - original_model_cost = litellm.model_cost - monkeypatch.setattr(litellm, "model_cost", dict(original_model_cost)) + monkeypatch.setattr(litellm, "model_cost", dict(litellm.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)) @@ -579,6 +579,10 @@ def test_boot_load_transient_failure_returns_local_then_background_retry_adopts_ for name in ("source", "url", "is_env_forced", "fallback_reason", "loaded_at", "source_revision", "etag"): monkeypatch.setattr(source_info, name, getattr(source_info, name)) + +def test_boot_load_transient_failure_returns_local_then_background_retry_adopts_remote(isolated_litellm_cost_state): + import litellm + remote_map = _load_root_cost_map() remote_map["claude-remote-only-test"] = {"litellm_provider": "anthropic", "mode": "chat"} client, calls = _mock_client( @@ -734,3 +738,129 @@ def test_boot_load_skips_remote_fetch_for_cli_processes( assert source["source"] == "local" else: assert source["source"] == "remote" + + +def _fetch_threads(): + return [thread for thread in threading.enumerate() if thread.name == "litellm-model-cost-map-fetch"] + + +def _held_client(outcome): + """Serves `outcome` only after `release` is set, so a test can observe the caller before the fetch finishes.""" + started = threading.Event() + release = threading.Event() + + def handler(request): + started.set() + release.wait(timeout=10) + return outcome + + return httpx.Client(transport=httpx.MockTransport(handler)), started, release + + +def _finish_fetch(release): + release.set() + threads = _fetch_threads() + for thread in threads: + thread.join(timeout=10) + assert all(not thread.is_alive() for thread in threads) + + +def _remote_map_with(model: str) -> bytes: + remote_map = _load_root_cost_map() + remote_map[model] = {"litellm_provider": "anthropic", "mode": "chat"} + return json.dumps(remote_map).encode() + + +def test_import_load_returns_the_bundled_map_before_the_remote_fetch_answers_then_adopts_it( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) + try: + cost_map = load_model_cost_map_at_import(url=_URL, client=client) + assert started.wait(timeout=10) + 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"] == "Remote fetch in progress" + monkeypatch.setattr(litellm, "model_cost", cost_map) + finally: + _finish_fetch(release) + + assert "claude-remote-only-test" in litellm.model_cost + assert "claude-remote-only-test" in litellm.anthropic_models + source = get_model_cost_map_source_info() + assert source["source"] == "remote" + assert source["fallback_reason"] is None + + +def test_import_load_never_overwrites_a_model_cost_the_caller_replaced(isolated_litellm_cost_state, monkeypatch): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + callers_map = {"my-model": {"litellm_provider": "openai", "mode": "chat"}} + monkeypatch.setattr(litellm, "model_cost", callers_map) + finally: + _finish_fetch(release) + + assert litellm.model_cost is callers_map + assert get_model_cost_map_source_info()["source"] == "local" + + +def test_import_load_reports_why_the_remote_fetch_failed(isolated_litellm_cost_state): + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, _, release = _held_client(httpx.Response(404)) + try: + load_model_cost_map_at_import(url=_URL, client=client) + finally: + _finish_fetch(release) + + source = get_model_cost_map_source_info() + assert source["source"] == "local" + assert source["fallback_reason"].startswith("Remote fetch failed: HTTP 404") + + +def test_import_load_with_the_local_env_override_starts_no_fetch(monkeypatch): + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + 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 = load_model_cost_map_at_import(url=_URL, client=httpx.Client(transport=httpx.MockTransport(_fail))) + + assert cost_map + assert _fetch_threads() == [] + assert get_model_cost_map_source_info()["is_env_forced"] is True + + +def test_import_litellm_returns_while_the_cost_map_server_has_not_answered(tmp_path): + import socket + import subprocess + + with socket.socket() as server: + server.bind(("127.0.0.1", 0)) + server.listen() + url = f"http://127.0.0.1:{server.getsockname()[1]}/model_prices.json" + env = {key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"} + env["LITELLM_MODEL_COST_MAP_URL"] = url + probe = ( + "import litellm\n" + "from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map_source_info\n" + "info = get_model_cost_map_source_info()\n" + "print(info['source'], info['fallback_reason'], len(litellm.model_cost) > 100)\n" + ) + result = subprocess.run( + [sys.executable, "-P", "-c", probe], capture_output=True, text=True, env=env, cwd=tmp_path, timeout=120 + ) + + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "local Remote fetch in progress True" From b0f77f568a65c0efc797aa9e84180604cee228fe Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 30 Sep 2026 19:44:27 -0700 Subject: [PATCH 02/18] fix(cost-map): keep edits made since import, resume the fetch in forked workers Address review feedback on the background cost map fetch: - entries written to litellm.model_cost after import (per-request pricing, direct edits) are carried over when the remote map is adopted - a process forked before the fetch finishes, like a preloaded gunicorn worker, starts its own fetch instead of keeping the bundled map forever - once litellm.model_cost has been replaced, the import fetch no longer writes its failure reason over the newer map's status - the import test no longer runs a listening socket --- .../litellm_core_utils/get_model_cost_map.py | 119 ++++++++++++------ .../test_get_model_cost_map.py | 113 +++++++++++++---- 2 files changed, 167 insertions(+), 65 deletions(-) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index ffcd4f68353..b64c96c294e 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -567,6 +567,33 @@ def adopt_model_cost_map( return fetched_model_count +@dataclass(frozen=True, slots=True) +class _ImportFetch: + url: str + timeout: int + max_attempts: int + sleep: Callable[[float], None] + rng: random.Random + client: _SyncGetClient + bundled: dict # mutable-ok: the exact object handed out as litellm.model_cost, compared by identity + bundled_as_loaded: Mapping[str, Mapping[str, object]] + + +class _PendingImportFetch: + fetch: _ImportFetch | None = None + + +def _model_cost_replaced_since_import(fetch: _ImportFetch | None) -> bool: + import litellm + + return fetch is not None and litellm.model_cost is not fetch.bundled + + +def _entries_changed_since_import(fetch: _ImportFetch) -> dict[str, object]: + current: Final = fetch.bundled.copy() + return {key: value for key, value in current.items() if fetch.bundled_as_loaded.get(key) != value} + + def _adopt_remote_model_cost_map( url: str, timeout: int, @@ -574,7 +601,7 @@ def _adopt_remote_model_cost_map( sleep: Callable[[float], None], rng: random.Random, client: _SyncGetClient, - replaces: Mapping[str, object] | None, + import_fetch: _ImportFetch | None, ) -> None: result: Final = _fetch_remote_model_cost_map_with_retry_sync( url=url, @@ -590,7 +617,8 @@ def _adopt_remote_model_cost_map( url, result.reason, ) - _cost_map_source_info.fallback_reason = f"Remote fetch failed: {result.reason}" + if not _model_cost_replaced_since_import(import_fetch): + _cost_map_source_info.fallback_reason = f"Remote fetch failed: {result.reason}" return _litellm_import_complete.wait() if not GetModelCostMap.validate_model_cost_map( @@ -601,18 +629,18 @@ def _adopt_remote_model_cost_map( "LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s", url, ) - _cost_map_source_info.fallback_reason = "Remote data failed integrity validation" + if not _model_cost_replaced_since_import(import_fetch): + _cost_map_source_info.fallback_reason = "Remote data failed integrity validation" return - import litellm - - if replaces is not None and litellm.model_cost is not replaces: - verbose_logger.debug("LiteLLM: litellm.model_cost was replaced during import; not adopting the remote map") + if _model_cost_replaced_since_import(import_fetch): + verbose_logger.debug("LiteLLM: litellm.model_cost was replaced after import; not adopting the remote map") return + edits: Final = _entries_changed_since_import(import_fetch) if import_fetch is not None else {} 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) + adopt_model_cost_map({**finalized, **edits}) def _retry_remote_fetch_in_background( @@ -636,33 +664,43 @@ def _retry_remote_fetch_in_background( sleep=sleep, rng=rng, client=client, - replaces=None, + import_fetch=None, ) except Exception as e: # noqa: BLE001 # a failed background retry must not kill the task; the backup stays verbose_logger.warning("LiteLLM: Background model cost map retry failed: %s", e) -def _fetch_remote_in_background_after_import( - url: str, - timeout: int, - max_attempts: int, - sleep: Callable[[float], None], - rng: random.Random, - client: _SyncGetClient, - replaces: Mapping[str, object], -) -> None: +def _run_import_fetch(fetch: _ImportFetch) -> None: try: _adopt_remote_model_cost_map( - url=url, - timeout=timeout, - attempts=range(1, max_attempts + 1), - sleep=sleep, - rng=rng, - client=client, - replaces=replaces, + url=fetch.url, + timeout=fetch.timeout, + attempts=range(1, fetch.max_attempts + 1), + sleep=fetch.sleep, + rng=fetch.rng, + client=fetch.client, + import_fetch=fetch, ) except Exception as e: # noqa: BLE001 # a failed background fetch must not kill the thread; the backup stays verbose_logger.warning("LiteLLM: Background model cost map fetch failed: %s", e) + finally: + if _PendingImportFetch.fetch is fetch: + _PendingImportFetch.fetch = None + + +def _start_import_fetch(fetch: _ImportFetch) -> None: + _PendingImportFetch.fetch = fetch + threading.Thread(target=_run_import_fetch, args=(fetch,), name="litellm-model-cost-map-fetch", daemon=True).start() + + +def _resume_import_fetch_after_fork() -> None: + pending: Final = _PendingImportFetch.fetch + if pending is not None: + _start_import_fetch(pending) + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_resume_import_fetch_after_fork) def load_model_cost_map_at_import( @@ -676,8 +714,9 @@ def load_model_cost_map_at_import( """ Return the bundled map right away and fetch the remote one in a background thread, so ``import litellm`` never waits on the network. The remote map replaces - ``litellm.model_cost`` once it passes the integrity check, unless the caller has - already replaced ``litellm.model_cost`` with their own map. + ``litellm.model_cost`` once it passes the integrity check, keeping every entry + written since import, unless the caller has replaced ``litellm.model_cost`` outright. + A process forked before the fetch finishes starts its own fetch. """ if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true" or _is_cli_process(): return get_model_cost_map(url=url, timeout=timeout) @@ -689,20 +728,18 @@ def load_model_cost_map_at_import( bundled: Final = _finalize_loaded_model_cost_map( GetModelCostMap.load_local_model_cost_map_with_revision() ).model_cost_map - threading.Thread( - target=_fetch_remote_in_background_after_import, - kwargs={ # mutable-ok: threading requires a mutable keyword-arguments mapping - "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, - "replaces": bundled, - }, - name="litellm-model-cost-map-fetch", - daemon=True, - ).start() + _start_import_fetch( + _ImportFetch( + 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, + bundled=bundled, + bundled_as_loaded=GetModelCostMap.loaded_model_cost_map(), + ) + ) return bundled diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index 2b1aef0e37f..0be62129899 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -796,7 +796,9 @@ def test_import_load_returns_the_bundled_map_before_the_remote_fetch_answers_the assert source["fallback_reason"] is None -def test_import_load_never_overwrites_a_model_cost_the_caller_replaced(isolated_litellm_cost_state, monkeypatch): +def test_import_load_keeps_custom_pricing_and_edits_made_before_the_remote_map_lands( + isolated_litellm_cost_state, monkeypatch +): import litellm from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import @@ -804,21 +806,56 @@ def test_import_load_never_overwrites_a_model_cost_the_caller_replaced(isolated_ try: monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) assert started.wait(timeout=10) - callers_map = {"my-model": {"litellm_provider": "openai", "mode": "chat"}} - monkeypatch.setattr(litellm, "model_cost", callers_map) + litellm.register_model( + {"gpt-4o": {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 42.0}}, + persist_across_reloads=False, + ) + litellm.model_cost["my-direct-edit"] = { + "litellm_provider": "openai", + "mode": "chat", + } # test-quality-ok: the edit is what is under test, and the fixture restores model_cost finally: _finish_fetch(release) - assert litellm.model_cost is callers_map - assert get_model_cost_map_source_info()["source"] == "local" + assert "claude-remote-only-test" in litellm.model_cost + assert litellm.model_cost["gpt-4o"]["input_cost_per_token"] == 42.0 + assert "my-direct-edit" in litellm.model_cost + assert get_model_cost_map_source_info()["source"] == "remote" -def test_import_load_reports_why_the_remote_fetch_failed(isolated_litellm_cost_state): +@pytest.mark.parametrize("remote_response", [httpx.Response(200, content=b""), httpx.Response(404)], ids=["ok", "404"]) +def test_import_load_leaves_a_model_cost_the_caller_replaced_and_its_status_alone( + isolated_litellm_cost_state, monkeypatch, remote_response +): + import litellm + from litellm.litellm_core_utils import get_model_cost_map as module + + if remote_response.status_code == 200: + remote_response = httpx.Response(200, content=_remote_map_with("claude-remote-only-test")) + client, started, release = _held_client(remote_response) + try: + monkeypatch.setattr(litellm, "model_cost", module.load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + reloaded_map = {"my-model": {"litellm_provider": "openai", "mode": "chat"}} + monkeypatch.setattr(litellm, "model_cost", reloaded_map) + monkeypatch.setattr(module._cost_map_source_info, "source", "remote") + monkeypatch.setattr(module._cost_map_source_info, "fallback_reason", None) + finally: + _finish_fetch(release) + + assert litellm.model_cost is reloaded_map + source = get_model_cost_map_source_info() + assert source["source"] == "remote" + assert source["fallback_reason"] is None + + +def test_import_load_reports_why_the_remote_fetch_failed(isolated_litellm_cost_state, monkeypatch): + import litellm from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import client, _, release = _held_client(httpx.Response(404)) try: - load_model_cost_map_at_import(url=_URL, client=client) + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) finally: _finish_fetch(release) @@ -827,6 +864,39 @@ def test_import_load_reports_why_the_remote_fetch_failed(isolated_litellm_cost_s assert source["fallback_reason"].startswith("Remote fetch failed: HTTP 404") +@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") +def test_a_process_forked_before_the_remote_map_lands_fetches_it_itself(isolated_litellm_cost_state, monkeypatch): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + parent_started = threading.Event() + release = threading.Event() + body = _remote_map_with("claude-remote-only-test") + requests_served = {"count": 0} + + def handler(request): + requests_served["count"] += 1 + if requests_served["count"] == 1: + parent_started.set() + release.wait(timeout=10) + return httpx.Response(200, content=body) + + client = httpx.Client(transport=httpx.MockTransport(handler)) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert parent_started.wait(timeout=10) + pid = os.fork() + if pid == 0: + for thread in _fetch_threads(): + thread.join(timeout=10) + os._exit(0 if "claude-remote-only-test" in litellm.model_cost else 1) + _, status = os.waitpid(pid, 0) + finally: + _finish_fetch(release) + + assert os.waitstatus_to_exitcode(status) == 0 + + def test_import_load_with_the_local_env_override_starts_no_fetch(monkeypatch): from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import @@ -842,25 +912,20 @@ def test_import_load_with_the_local_env_override_starts_no_fetch(monkeypatch): assert get_model_cost_map_source_info()["is_env_forced"] is True -def test_import_litellm_returns_while_the_cost_map_server_has_not_answered(tmp_path): - import socket +def test_import_litellm_returns_before_the_remote_fetch_finishes(tmp_path): import subprocess - with socket.socket() as server: - server.bind(("127.0.0.1", 0)) - server.listen() - url = f"http://127.0.0.1:{server.getsockname()[1]}/model_prices.json" - env = {key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"} - env["LITELLM_MODEL_COST_MAP_URL"] = url - probe = ( - "import litellm\n" - "from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map_source_info\n" - "info = get_model_cost_map_source_info()\n" - "print(info['source'], info['fallback_reason'], len(litellm.model_cost) > 100)\n" - ) - result = subprocess.run( - [sys.executable, "-P", "-c", probe], capture_output=True, text=True, env=env, cwd=tmp_path, timeout=120 - ) + env = {key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"} + env["LITELLM_MODEL_COST_MAP_URL"] = "http://127.0.0.1:9/model_prices.json" + probe = ( + "import litellm\n" + "from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map_source_info\n" + "info = get_model_cost_map_source_info()\n" + "print(info['source'], info['fallback_reason'], len(litellm.model_cost) > 100)\n" + ) + result = subprocess.run( + [sys.executable, "-P", "-c", probe], capture_output=True, text=True, env=env, cwd=tmp_path, timeout=120 + ) assert result.returncode == 0, result.stderr assert result.stdout.strip() == "local Remote fetch in progress True" From 0e5ce6a0d7f07d89856357535681dcb22b6dedb6 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 30 Sep 2026 19:50:41 -0700 Subject: [PATCH 03/18] refactor(cost-map): make backup_model_count public and keep the edit overlay immutable --- .../litellm_core_utils/get_model_cost_map.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index b64c96c294e..701d9b58509 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -104,7 +104,7 @@ class GetModelCostMap: cls._loaded_catalog = MappingProxyType({key: MappingProxyType(entry) for key, entry in raw.items()}) @classmethod - def _get_backup_model_count(cls) -> int: + def backup_model_count(cls) -> int: """Return the number of models in the local backup (cached int).""" if cls._backup_model_count < 0: backup: Final = cls.load_local_model_cost_map() @@ -401,7 +401,7 @@ async def refetch_model_cost_map( return result if not GetModelCostMap.validate_model_cost_map( fetched_map=result.model_cost_map, - backup_model_count=GetModelCostMap._get_backup_model_count(), + backup_model_count=GetModelCostMap.backup_model_count(), ): return ModelCostMapReloadUnavailable(reason=f"model cost map from {url} failed integrity validation") _cost_map_source_info.loaded_at = datetime.now(timezone.utc) @@ -589,9 +589,9 @@ def _model_cost_replaced_since_import(fetch: _ImportFetch | None) -> bool: return fetch is not None and litellm.model_cost is not fetch.bundled -def _entries_changed_since_import(fetch: _ImportFetch) -> dict[str, object]: - current: Final = fetch.bundled.copy() - return {key: value for key, value in current.items() if fetch.bundled_as_loaded.get(key) != value} +def _entries_changed_since_import(fetch: _ImportFetch) -> Mapping[str, object]: + current: Final = tuple(fetch.bundled.items()) + return MappingProxyType({key: value for key, value in current if fetch.bundled_as_loaded.get(key) != value}) def _adopt_remote_model_cost_map( @@ -623,7 +623,7 @@ def _adopt_remote_model_cost_map( _litellm_import_complete.wait() 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 + backup_model_count=GetModelCostMap.backup_model_count(), ): verbose_logger.warning( "LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s", @@ -635,12 +635,12 @@ def _adopt_remote_model_cost_map( if _model_cost_replaced_since_import(import_fetch): verbose_logger.debug("LiteLLM: litellm.model_cost was replaced after import; not adopting the remote map") return - edits: Final = _entries_changed_since_import(import_fetch) if import_fetch is not None else {} + edits: Final = _entries_changed_since_import(import_fetch) if import_fetch is not None else MappingProxyType({}) 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, **edits}) + adopt_model_cost_map({**finalized, **edits}) # mutable-ok: litellm.model_cost is a plain mutable dict by contract def _retry_remote_fetch_in_background( @@ -810,7 +810,7 @@ def get_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(), + backup_model_count=GetModelCostMap.backup_model_count(), ): verbose_logger.warning( "LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s", From 8979167087bc7aa4753d0dac503ebf79ee667dd9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 9 Oct 2026 20:49:35 -0700 Subject: [PATCH 04/18] fix(cost-map): wait for the remote map on a lookup miss and keep edits made during adoption A model that only the remote cost map knows failed get_llm_provider and get_model_info right after import, because the bundled map was still in place. A lookup that misses now waits for the in-flight import fetch once and retries, bounded by the fetch timeout. Edits written through a reference to the bundled dict while the remote map is being adopted are carried over after the swap. The fork and fresh-import tests move to tests/integration/sdk so the unit file stays in-process. --- litellm/__init__.py | 2 +- .../get_llm_provider_logic.py | 48 +++-- .../litellm_core_utils/get_model_cost_map.py | 99 ++++++++-- litellm/utils.py | 15 ++ .../sdk/test_cost_map_background_fetch_sdk.py | 146 ++++++++++++++ .../test_get_model_cost_map.py | 187 ++++++++++++++---- 6 files changed, 427 insertions(+), 70 deletions(-) create mode 100644 tests/integration/sdk/test_cost_map_background_fetch_sdk.py diff --git a/litellm/__init__.py b/litellm/__init__.py index a3cf2b91914..c589b5b0614 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -563,7 +563,7 @@ _key_management_system: Optional["KeyManagementSystem"] = None output_parse_pii: bool = False ############################################# from litellm.litellm_core_utils.get_model_cost_map import ( - get_model_cost_map, + get_model_cost_map as get_model_cost_map, load_model_cost_map_at_import, mark_litellm_import_complete, ) diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index c9e5167e7a9..a52a5ca65b1 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -10,6 +10,7 @@ from litellm.constants import ( from litellm.litellm_core_utils.fallback_generalizations import ( match_routing_generalization, ) +from litellm.litellm_core_utils.get_model_cost_map import wait_for_import_model_cost_map from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.secret_managers.main import get_secret, get_secret_str @@ -174,6 +175,37 @@ def get_llm_provider( Return model, custom_llm_provider, dynamic_api_key, api_base """ + resolved: Final = _resolve_llm_provider(model, custom_llm_provider, api_base, api_key, litellm_params) + if resolved is not None: + return resolved + retried: Final = ( + _resolve_llm_provider(model, custom_llm_provider, api_base, api_key, litellm_params) + if wait_for_import_model_cost_map() + else None + ) + if retried is not None: + return retried + if litellm.suppress_debug_info is False: + print() # noqa: T201 + print( # noqa: T201 + "\033[1;31mProvider List: https://docs.litellm.ai/docs/providers\033[0m" + ) + print() # noqa: T201 + raise litellm.exceptions.BadRequestError( + message=f"LLM Provider NOT provided. Pass in the LLM provider you are trying to call. You passed model={model}\n Pass model as E.g. For 'Huggingface' inference endpoints pass in `completion(model='huggingface/starcoder',..)` Learn more: https://docs.litellm.ai/docs/providers", + model=model, + response=None, + llm_provider="", + ) + + +def _resolve_llm_provider( + model: str, + custom_llm_provider: str | None, + api_base: str | None, + api_key: str | None, + litellm_params: GenericLiteLLMParams | None, +) -> tuple[str, str, str | None, str | None] | None: try: # Early validation - model is required if model is None: @@ -533,20 +565,7 @@ def get_llm_provider( custom_llm_provider = match_routing_generalization(model) if not custom_llm_provider: - if litellm.suppress_debug_info is False: - print() # noqa: T201 - print( # noqa: T201 - "\033[1;31mProvider List: https://docs.litellm.ai/docs/providers\033[0m" - ) - print() # noqa: T201 - error_str = f"LLM Provider NOT provided. Pass in the LLM provider you are trying to call. You passed model={model}\n Pass model as E.g. For 'Huggingface' inference endpoints pass in `completion(model='huggingface/starcoder',..)` Learn more: https://docs.litellm.ai/docs/providers" - # maps to openai.NotFoundError, this is raised when openai does not recognize the llm - raise litellm.exceptions.BadRequestError( - message=error_str, - model=model, - response=None, - llm_provider="", - ) + return None if api_base is not None and not isinstance(api_base, str): raise Exception(f"api base needs to be a string. api_base={api_base}") if dynamic_api_key is not None and not isinstance(dynamic_api_key, str): @@ -556,7 +575,6 @@ def get_llm_provider( if isinstance(e, litellm.exceptions.BadRequestError): raise e else: - error_str = f"GetLLMProvider Exception - {e}\n\noriginal model: {model}" raise litellm.exceptions.BadRequestError( message=f"GetLLMProvider Exception - {e}\n\noriginal model: {model}", model=model, diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 5ceb08d33c8..7c890b49085 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -18,7 +18,7 @@ import sys import threading import time from collections.abc import Awaitable, Callable, Mapping -from dataclasses import dataclass, replace +from dataclasses import dataclass, field, replace from datetime import datetime, timezone from importlib.resources import files from pathlib import Path @@ -42,6 +42,7 @@ FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations" _CATALOG_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) _BUNDLED_CATALOG_ADAPTER: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) _CLI_ENTRYPOINT_NAMES: Final = frozenset({"lite", "litellm-proxy"}) +_NO_EDITS: Final[Mapping[str, object]] = MappingProxyType({}) def _is_cli_process() -> bool: @@ -53,7 +54,7 @@ def _is_cli_process() -> bool: RESERVED_TOP_LEVEL_KEYS: Final = frozenset({"sample_spec", FALLBACK_GENERALIZATIONS_KEY}) -def _count_model_entries(model_cost: dict) -> int: +def _count_model_entries(model_cost: Mapping[str, object]) -> int: """Count actual model entries, excluding reserved meta keys.""" return sum(1 for key in model_cost if key not in RESERVED_TOP_LEVEL_KEYS) @@ -108,7 +109,7 @@ class GetModelCostMap: def get_backup_model_count(cls) -> int: """Return the number of models in the local backup (cached int).""" if cls._backup_model_count < 0: - backup: Final = cls.load_local_model_cost_map() + backup: Final = cls.load_local_model_cost_map_with_revision().model_cost_map cls._backup_model_count = _count_model_entries(backup) return cls._backup_model_count @@ -135,7 +136,7 @@ class GetModelCostMap: @classmethod def _check_model_count_not_reduced( cls, - fetched_map: dict, + fetched_map: dict[str, object], backup_model_count: int, min_model_count: int = MODEL_COST_MAP_MIN_MODEL_COUNT, max_shrink_ratio: float = MODEL_COST_MAP_MAX_SHRINK_RATIO, @@ -170,7 +171,7 @@ class GetModelCostMap: @classmethod def validate_model_cost_map( cls, - fetched_map: dict, + fetched_map: dict[str, object], backup_model_count: int, min_model_count: int = MODEL_COST_MAP_MIN_MODEL_COUNT, max_shrink_ratio: float = MODEL_COST_MAP_MAX_SHRINK_RATIO, @@ -213,7 +214,7 @@ def mark_litellm_import_complete() -> None: @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 + model_cost_map: dict[str, object] # mutable-ok: becomes litellm.model_cost, a plain mutable dict by contract revision: str | None = None etag: str | None = None @@ -305,7 +306,9 @@ def _classify_fetch_response(response: httpx.Response, url: str) -> _FetchAttemp if not isinstance(parsed, dict): return ModelCostMapReloadUnavailable(reason=f"expected a JSON object from {url}, got {type(parsed).__name__}") return ModelCostMapReloaded( - model_cost_map=parsed, revision=git_blob_id(response.content), etag=response.headers.get("etag") + model_cost_map=_BUNDLED_CATALOG_ADAPTER.validate_python(parsed), + revision=git_blob_id(response.content), + etag=response.headers.get("etag"), ) @@ -578,18 +581,24 @@ class _ImportFetch: sleep: Callable[[float], None] rng: random.Random client: _SyncGetClient - bundled: dict # mutable-ok: the exact object handed out as litellm.model_cost, compared by identity + bundled: Mapping[str, object] bundled_as_loaded: Mapping[str, Mapping[str, object]] + first_attempt_settled: threading.Event = field(default_factory=threading.Event) class _PendingImportFetch: fetch: _ImportFetch | None = None + thread: threading.Thread | None = None + + +def _litellm_model_cost_is(expected: Mapping[str, object]) -> bool: + import litellm + + return litellm.model_cost is expected def _model_cost_replaced_since_import(fetch: _ImportFetch | None) -> bool: - import litellm - - return fetch is not None and litellm.model_cost is not fetch.bundled + return fetch is not None and not _litellm_model_cost_is(fetch.bundled) def _entries_changed_since_import(fetch: _ImportFetch) -> Mapping[str, object]: @@ -597,6 +606,26 @@ def _entries_changed_since_import(fetch: _ImportFetch) -> Mapping[str, object]: return MappingProxyType({key: value for key, value in current if fetch.bundled_as_loaded.get(key) != value}) +def _carry_edits_made_during_adoption( + fetch: _ImportFetch, + adopted: dict[str, object], # mutable-ok: the dict just adopted as litellm.model_cost, updated in place + carried: Mapping[str, object], + finalized: Mapping[str, object], +) -> None: + changed: Final = tuple(_entries_changed_since_import(fetch).items()) + late: Final = { + key: value + for key, value in changed + if value is not carried.get(key) and adopted.get(key) is carried.get(key, finalized.get(key)) + } + if not late or not _litellm_model_cost_is(adopted): + return + from litellm import utils + + adopted.update(late) + utils._invalidate_model_cost_lowercase_map() # pyright: ignore[reportPrivateUsage] # required cache invalidation + + def _adopt_remote_model_cost_map( url: str, timeout: int, @@ -638,12 +667,15 @@ def _adopt_remote_model_cost_map( if _model_cost_replaced_since_import(import_fetch): verbose_logger.debug("LiteLLM: litellm.model_cost was replaced after import; not adopting the remote map") return - edits: Final = _entries_changed_since_import(import_fetch) if import_fetch is not None else MappingProxyType({}) finalized: Final = _finalize_loaded_model_cost_map(result).model_cost_map + edits: Final = _entries_changed_since_import(import_fetch) if import_fetch is not None else _NO_EDITS + adopted: Final = {**finalized, **edits} _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, **edits}) # mutable-ok: litellm.model_cost is a plain mutable dict by contract + adopt_model_cost_map(adopted) + if import_fetch is not None: + _carry_edits_made_during_adoption(import_fetch, adopted=adopted, carried=edits, finalized=finalized) def _retry_remote_fetch_in_background( @@ -674,12 +706,16 @@ def _retry_remote_fetch_in_background( def _run_import_fetch(fetch: _ImportFetch) -> None: + def settle_then_sleep(seconds: float) -> None: + fetch.first_attempt_settled.set() + fetch.sleep(seconds) + try: _adopt_remote_model_cost_map( url=fetch.url, timeout=fetch.timeout, attempts=range(1, fetch.max_attempts + 1), - sleep=fetch.sleep, + sleep=settle_then_sleep, rng=fetch.rng, client=fetch.client, import_fetch=fetch, @@ -687,19 +723,45 @@ def _run_import_fetch(fetch: _ImportFetch) -> None: except Exception as e: # noqa: BLE001 # a failed background fetch must not kill the thread; the backup stays verbose_logger.warning("LiteLLM: Background model cost map fetch failed: %s", e) finally: + fetch.first_attempt_settled.set() if _PendingImportFetch.fetch is fetch: _PendingImportFetch.fetch = None def _start_import_fetch(fetch: _ImportFetch) -> None: + thread: Final = threading.Thread( + target=_run_import_fetch, args=(fetch,), name="litellm-model-cost-map-fetch", daemon=True + ) _PendingImportFetch.fetch = fetch - threading.Thread(target=_run_import_fetch, args=(fetch,), name="litellm-model-cost-map-fetch", daemon=True).start() + _PendingImportFetch.thread = thread + thread.start() def _resume_import_fetch_after_fork() -> None: pending: Final = _PendingImportFetch.fetch if pending is not None: - _start_import_fetch(pending) + _start_import_fetch(replace(pending, first_attempt_settled=threading.Event())) + + +def wait_for_import_model_cost_map() -> bool: + """ + Called by a lookup that missed the cost map while the remote map fetched at import is + still in flight. Blocks until that fetch's first attempt settles, at most one fetch + timeout, so a model the bundled map lacks resolves the way it did when import waited + for the fetch. True means the lookup should be tried once more. + """ + pending: Final = _PendingImportFetch.fetch + if ( + pending is None + or not _litellm_import_complete.is_set() + or threading.current_thread() is _PendingImportFetch.thread + ): + return False + if not pending.first_attempt_settled.is_set(): + verbose_logger.debug( + "LiteLLM: lookup missed the bundled cost map; waiting for the remote map fetched at import" + ) + return pending.first_attempt_settled.wait(timeout=pending.timeout) if hasattr(os, "register_at_fork"): @@ -713,13 +775,14 @@ def load_model_cost_map_at_import( sleep: Callable[[float], None] = time.sleep, rng: random.Random | None = None, client: "_SyncGetClient | None" = None, -) -> dict: +) -> dict: # mutable-ok: becomes litellm.model_cost, a plain mutable dict by contract """ Return the bundled map right away and fetch the remote one in a background thread, so ``import litellm`` never waits on the network. The remote map replaces ``litellm.model_cost`` once it passes the integrity check, keeping every entry written since import, unless the caller has replaced ``litellm.model_cost`` outright. - A process forked before the fetch finishes starts its own fetch. + A lookup that misses the bundled map while the fetch is in flight waits for it, see + ``wait_for_import_model_cost_map``. A process forked before the fetch finishes starts its own fetch. """ if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true" or _is_cli_process(): return get_model_cost_map(url=url, timeout=timeout) diff --git a/litellm/utils.py b/litellm/utils.py index e7a171f56c6..1baacd01ae8 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -90,6 +90,7 @@ from litellm.litellm_core_utils.fallback_generalizations import ( match_capability_generalizations, match_fill_missing_generalizations, ) +from litellm.litellm_core_utils.get_model_cost_map import wait_for_import_model_cost_map from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.litellm_core_utils.tokenizer import strip_special_tokens from litellm.rust_bridge import tokenizer as tokenizer_dispatch @@ -6036,6 +6037,20 @@ def get_model_info_helper( """ Helper for 'get_model_info'. Separated out to avoid infinite loop caused by returning 'supported_openai_param's """ + try: + return _model_info_from_cost_map(model, custom_llm_provider, api_base, api_key) + except ModelNotMappedError: + if not wait_for_import_model_cost_map(): + raise + return _model_info_from_cost_map(model, custom_llm_provider, api_base, api_key) + + +def _model_info_from_cost_map( + model: str, + custom_llm_provider: str | None, + api_base: str | None, + api_key: str | None, +) -> ModelInfoBase: from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider try: diff --git a/tests/integration/sdk/test_cost_map_background_fetch_sdk.py b/tests/integration/sdk/test_cost_map_background_fetch_sdk.py new file mode 100644 index 00000000000..6483f9a3a7a --- /dev/null +++ b/tests/integration/sdk/test_cost_map_background_fetch_sdk.py @@ -0,0 +1,146 @@ +import json +import os +import subprocess +import sys +import textwrap +import threading +from collections.abc import Iterator +from dataclasses import dataclass +from importlib.resources import files +from pathlib import Path +from typing import Final + +import pytest +from integration._support.wire import Reply, Request, Wire, wire_server + +pytestmark: Final = pytest.mark.timeout(300) + +_REMOTE_ONLY_MODEL: Final = "zz-remote-only-sdk-model" +_CHILD_SECONDS: Final = 120 + +_SCRIPT: Final = textwrap.dedent( + """ + import os, sys, threading + import litellm + from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map_source_info + + variant, model = sys.argv[1:3] + + def join_fetch(): + for thread in threading.enumerate(): + if thread.name == "litellm-model-cost-map-fetch": + thread.join(timeout=60) + + info = get_model_cost_map_source_info() + print(info["source"], info["fallback_reason"], model in litellm.model_cost, flush=True) + if variant == "adopt": + sys.stdin.readline() + join_fetch() + print(get_model_cost_map_source_info()["source"], model in litellm.model_cost, flush=True) + elif variant == "get_llm_provider": + print(litellm.get_llm_provider(model)[1], flush=True) + elif variant == "get_model_info": + print(litellm.get_model_info(model, custom_llm_provider="openai")["litellm_provider"], flush=True) + elif variant == "fork": + pid = os.fork() + if pid == 0: + join_fetch() + os._exit(0 if model in litellm.model_cost else 1) + print(os.waitstatus_to_exitcode(os.waitpid(pid, 0)[1]), flush=True) + else: + raise SystemExit(f"unknown variant {variant}") + """ +) + + +def _remote_map() -> bytes: + bundled: Final = json.loads(files("litellm").joinpath("model_prices_and_context_window_backup.json").read_text()) + return json.dumps({**bundled, _REMOTE_ONLY_MODEL: {"litellm_provider": "openai", "mode": "chat"}}).encode() + + +@dataclass(frozen=True, slots=True) +class _CostMapHost: + wire: Wire + first_request_answered: threading.Event + + +@pytest.fixture +def cost_map_host() -> Iterator[_CostMapHost]: + body: Final = _remote_map() + gate: Final = threading.Event() + first_seen: Final = threading.Event() + + def respond(request: Request) -> Reply: + if not first_seen.is_set(): + first_seen.set() + gate.wait(timeout=_CHILD_SECONDS) + return Reply(body=body) + + with wire_server(respond) as wire: + try: + yield _CostMapHost(wire=wire, first_request_answered=gate) + finally: + gate.set() + + +@dataclass(frozen=True, slots=True) +class _Child: + after_import: str + rest: tuple[str, ...] + + +def _run_child(host: _CostMapHost, variant: str, tmp_path: Path, answer_after_import: bool) -> _Child: + stderr_path: Final = tmp_path / "stderr.log" + with ( + stderr_path.open("w") as stderr, + subprocess.Popen( + [sys.executable, "-P", "-c", _SCRIPT, variant, _REMOTE_ONLY_MODEL], + env={ + **os.environ, + "LITELLM_LOCAL_MODEL_COST_MAP": "False", + "NO_PROXY": "*", + "LITELLM_MODEL_COST_MAP_URL": f"{host.wire.url}/model_prices_and_context_window.json", + }, + cwd=tmp_path, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=stderr, + text=True, + ) as child, + ): + assert child.stdout is not None + after_import: Final = child.stdout.readline().strip() + if answer_after_import: + host.first_request_answered.set() + rest: Final = child.communicate(input="\n", timeout=_CHILD_SECONDS)[0] + assert child.returncode == 0, stderr_path.read_text() + return _Child(after_import=after_import, rest=tuple(rest.split())) + + +def test_import_returns_before_the_remote_map_arrives_and_adopts_it_after( + cost_map_host: _CostMapHost, tmp_path: Path +) -> None: + child: Final = _run_child(cost_map_host, "adopt", tmp_path, answer_after_import=True) + + assert child.after_import == "local Remote fetch in progress False" + assert child.rest == ("remote", "True") + + +@pytest.mark.parametrize("lookup", ["get_llm_provider", "get_model_info"]) +def test_a_lookup_right_after_import_finds_a_model_only_the_remote_map_has( + cost_map_host: _CostMapHost, tmp_path: Path, lookup: str +) -> None: + child: Final = _run_child(cost_map_host, lookup, tmp_path, answer_after_import=True) + + assert child.after_import == "local Remote fetch in progress False" + assert child.rest == ("openai",) + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") +def test_a_worker_forked_while_the_remote_map_is_in_flight_fetches_it_itself( + cost_map_host: _CostMapHost, tmp_path: Path +) -> None: + child: Final = _run_child(cost_map_host, "fork", tmp_path, answer_after_import=False) + + assert child.after_import == "local Remote fetch in progress False" + assert child.rest == ("0",) diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index faeb732569b..56ed3ad2bd6 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -5,6 +5,7 @@ count actual model entries, not reserved meta keys) and the extraction of the """ import asyncio, importlib, importlib.resources as importlib_get_model, json, litellm +import logging import os import sys import threading @@ -844,37 +845,155 @@ def test_import_load_reports_why_the_remote_fetch_failed(isolated_litellm_cost_s assert source["fallback_reason"].startswith("Remote fetch failed: HTTP 404") -@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") -def test_a_process_forked_before_the_remote_map_lands_fetches_it_itself(isolated_litellm_cost_state, monkeypatch): +def test_a_fetch_resumed_after_fork_adopts_the_remote_map_while_the_inherited_one_hangs( + isolated_litellm_cost_state, monkeypatch +): import litellm - from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + from litellm.litellm_core_utils import get_model_cost_map as module - parent_started = threading.Event() + inherited_started = threading.Event() release = threading.Event() body = _remote_map_with("claude-remote-only-test") - requests_served = {"count": 0} + served = [] def handler(request): - requests_served["count"] += 1 - if requests_served["count"] == 1: - parent_started.set() + served.append(request) + if len(served) == 1: + inherited_started.set() release.wait(timeout=10) return httpx.Response(200, content=body) client = httpx.Client(transport=httpx.MockTransport(handler)) try: - monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) - assert parent_started.wait(timeout=10) - pid = os.fork() - if pid == 0: - for thread in _fetch_threads(): - thread.join(timeout=10) - os._exit(0 if "claude-remote-only-test" in litellm.model_cost else 1) - _, status = os.waitpid(pid, 0) + monkeypatch.setattr(litellm, "model_cost", module.load_model_cost_map_at_import(url=_URL, client=client)) + assert inherited_started.wait(timeout=10) + inherited_thread = module._PendingImportFetch.thread + module._resume_import_fetch_after_fork() + resumed_thread = module._PendingImportFetch.thread + assert resumed_thread is not inherited_thread + resumed_thread.join(timeout=10) + assert not resumed_thread.is_alive() + assert "claude-remote-only-test" in litellm.model_cost + assert inherited_thread.is_alive() finally: _finish_fetch(release) - assert os.waitstatus_to_exitcode(status) == 0 + assert len(served) == 2 + + +def test_an_edit_written_through_a_stale_reference_during_adoption_survives_it( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm import utils as litellm_utils + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) + monkeypatch.setattr(litellm_utils, "_runtime_registered_model_cost", {}) + bundled = load_model_cost_map_at_import(url=_URL, client=client) + monkeypatch.setattr(litellm, "model_cost", bundled) + original_replay = litellm_utils._LiveDeploymentReplay.callback + late_entry = {"litellm_provider": "openai", "mode": "chat"} + + def write_through_the_reference_read_before_the_swap(): + bundled["my-late-edit"] = late_entry + bundled["gpt-4o-mini"] = {**bundled["gpt-4o-mini"], "input_cost_per_token": 42.0} + bundled["gpt-4o"] = {**bundled["gpt-4o"], "input_cost_per_token": 42.0} + litellm.model_cost["gpt-4o"] = {**litellm.model_cost["gpt-4o"], "input_cost_per_token": 7.0} + + litellm_utils.set_live_deployment_replay(write_through_the_reference_read_before_the_swap) + try: + assert started.wait(timeout=10) + finally: + _finish_fetch(release) + litellm_utils.set_live_deployment_replay(original_replay) + + assert litellm.model_cost is not bundled + assert "claude-remote-only-test" in litellm.model_cost + assert litellm.model_cost["my-late-edit"] is late_entry + assert litellm.model_cost["gpt-4o-mini"]["input_cost_per_token"] == 42.0 + assert litellm.model_cost["gpt-4o"]["input_cost_per_token"] == 7.0 + + +class _ReleaseWhenALookupWaits(logging.Handler): + """Lets the held fetch answer only once a lookup is blocked waiting on it, so the lookup provably missed first.""" + + def __init__(self, release: threading.Event): + super().__init__(level=logging.DEBUG) + self.release = release + self.waited = threading.Event() + + def emit(self, record): + if "waiting for the remote map fetched at import" in record.getMessage(): + self.waited.set() + self.release.set() + + +@pytest.fixture() +def release_when_a_lookup_waits(): + from litellm import verbose_logger + + installed = [] + + def install(release: threading.Event) -> _ReleaseWhenALookupWaits: + handler = _ReleaseWhenALookupWaits(release) + verbose_logger.addHandler(handler) + installed.append(handler) + return handler + + previous_level = verbose_logger.level + verbose_logger.setLevel(logging.DEBUG) + yield install + verbose_logger.setLevel(previous_level) + for handler in installed: + verbose_logger.removeHandler(handler) + + +@pytest.mark.parametrize( + "lookup", + [ + pytest.param(lambda: litellm.get_llm_provider("zz-remote-only-test")[1], id="get_llm_provider"), + pytest.param( + lambda: litellm.get_model_info("zz-remote-only-test", custom_llm_provider="anthropic")["litellm_provider"], + id="get_model_info", + ), + ], +) +def test_a_lookup_that_misses_the_bundled_map_waits_for_the_remote_map_fetched_at_import( + isolated_litellm_cost_state, monkeypatch, release_when_a_lookup_waits, lookup +): + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("zz-remote-only-test"))) + waits = release_when_a_lookup_waits(release) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + assert "zz-remote-only-test" not in litellm.model_cost + provider = lookup() + finally: + _finish_fetch(release) + + assert waits.waited.is_set() + assert provider == "anthropic" + + +def test_a_lookup_that_misses_both_maps_still_raises_after_waiting( + isolated_litellm_cost_state, monkeypatch, release_when_a_lookup_waits +): + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("zz-remote-only-test"))) + waits = release_when_a_lookup_waits(release) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + with pytest.raises(litellm.exceptions.BadRequestError, match="LLM Provider NOT provided"): + litellm.get_llm_provider("zz-in-neither-map") + finally: + _finish_fetch(release) + + assert waits.waited.is_set() def test_import_load_with_the_local_env_override_starts_no_fetch(monkeypatch): @@ -892,31 +1011,13 @@ def test_import_load_with_the_local_env_override_starts_no_fetch(monkeypatch): assert get_model_cost_map_source_info()["is_env_forced"] is True -def test_import_litellm_returns_before_the_remote_fetch_finishes(tmp_path): - import subprocess - - env = {key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"} - env["LITELLM_MODEL_COST_MAP_URL"] = "http://127.0.0.1:9/model_prices.json" - probe = ( - "import litellm\n" - "from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map_source_info\n" - "info = get_model_cost_map_source_info()\n" - "print(info['source'], info['fallback_reason'], len(litellm.model_cost) > 100)\n" - ) - result = subprocess.run( - [sys.executable, "-P", "-c", probe], capture_output=True, text=True, env=env, cwd=tmp_path, timeout=120 - ) - - assert result.returncode == 0, result.stderr - assert result.stdout.strip() == "local Remote fetch in progress True" - - @pytest.fixture() def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="session") def event_loop(): try: @@ -926,6 +1027,7 @@ def event_loop(): yield loop loop.close() + @pytest.fixture(scope="function") def setup_and_teardown(event_loop): import litellm @@ -959,6 +1061,7 @@ def setup_and_teardown(event_loop): if pending: event_loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) + _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "set_verbose": getattr(litellm, "set_verbose", False), @@ -973,6 +1076,7 @@ _SCALAR_DEFAULTS = { "cohere_key": getattr(litellm, "cohere_key", None), } + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestCheckIsValidDict: """Unit tests for _check_is_valid_dict.""" @@ -997,6 +1101,7 @@ class TestCheckIsValidDict: """Non-empty dict should pass.""" assert GetModelCostMap._check_is_valid_dict({"model": {}}) is True + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestCheckModelCountNotReduced: """Unit tests for _check_model_count_not_reduced.""" @@ -1051,6 +1156,7 @@ class TestCheckModelCountNotReduced: is True ) + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestValidateModelCostMap: """Unit tests for validate_model_cost_map (combines both checks).""" @@ -1087,6 +1193,7 @@ class TestValidateModelCostMap: is True ) + def _fetch_with_single_outcome(monkeypatch, outcome): from litellm.litellm_core_utils import get_model_cost_map as module @@ -1154,6 +1261,7 @@ class TestGetModelCostMapFallback: assert isinstance(result, dict) assert len(result) > 0 + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestBackupModelCostMapExists: """Validates the local backup file is always present and valid.""" @@ -1169,6 +1277,7 @@ class TestBackupModelCostMapExists: backup = GetModelCostMap.load_local_model_cost_map() assert len(backup) > 100, f"Backup has only {len(backup)} models, expected > 100" + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestBadHostedModelCostMap: """ @@ -1233,12 +1342,14 @@ class TestBadHostedModelCostMap: finally: litellm.model_cost = original + @pytest.fixture() def _vcr_outcome_gate_local_testing(request, vcr): install_live_call_probe(request, vcr) yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def isolate_litellm_state(): """ @@ -1291,6 +1402,7 @@ def isolate_litellm_state(): setattr(litellm, attr, original_value) _invalidate_model_cost_lowercase_map() + _SCALAR_DEFAULTS_local_testing = { "num_retries": getattr(litellm, "num_retries", None), "num_retries_per_request": getattr(litellm, "num_retries_per_request", None), @@ -1311,6 +1423,7 @@ _SCALAR_DEFAULTS_local_testing = { "api_key": getattr(litellm, "api_key", None), } + @pytest.fixture(scope="module") def setup_and_teardown_local_testing(): """ @@ -1333,6 +1446,7 @@ def setup_and_teardown_local_testing(): litellm.in_memory_llm_clients_cache.flush_cache() yield + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -1344,6 +1458,7 @@ def test_get_model_cost_map(): except Exception as e: pytest.fail(f"An exception occurred: {e}") + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", From e5a61d1e0ef1ddfe198b44dc83a51eb582ff8ea1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 9 Oct 2026 22:15:29 -0700 Subject: [PATCH 05/18] fix(cost-map): retry a missed lookup whenever the remote map was adopted after it A lookup now reads an adoption counter before it looks in the map, so a fetch that finishes between the miss and the wait still triggers the retry, and a fetch that failed no longer causes a pointless second lookup. The lookup-wait test handler stored its event as `release`, which shadows logging.Handler.release on Python 3.12 and failed core-utils in CI. --- .../get_llm_provider_logic.py | 8 ++++-- .../litellm_core_utils/get_model_cost_map.py | 27 ++++++++++++++----- litellm/utils.py | 8 ++++-- .../test_get_model_cost_map.py | 24 +++++++++++++++-- 4 files changed, 54 insertions(+), 13 deletions(-) diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index a52a5ca65b1..a8b4b077e93 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -10,7 +10,10 @@ from litellm.constants import ( from litellm.litellm_core_utils.fallback_generalizations import ( match_routing_generalization, ) -from litellm.litellm_core_utils.get_model_cost_map import wait_for_import_model_cost_map +from litellm.litellm_core_utils.get_model_cost_map import ( + import_model_cost_map_generation, + wait_for_import_model_cost_map, +) from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.secret_managers.main import get_secret, get_secret_str @@ -175,12 +178,13 @@ def get_llm_provider( Return model, custom_llm_provider, dynamic_api_key, api_base """ + generation: Final = import_model_cost_map_generation() resolved: Final = _resolve_llm_provider(model, custom_llm_provider, api_base, api_key, litellm_params) if resolved is not None: return resolved retried: Final = ( _resolve_llm_provider(model, custom_llm_provider, api_base, api_key, litellm_params) - if wait_for_import_model_cost_map() + if wait_for_import_model_cost_map(since=generation) else None ) if retried is not None: diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 7c890b49085..39580ed47ae 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -589,6 +589,7 @@ class _ImportFetch: class _PendingImportFetch: fetch: _ImportFetch | None = None thread: threading.Thread | None = None + adoptions: int = 0 def _litellm_model_cost_is(expected: Mapping[str, object]) -> bool: @@ -676,6 +677,7 @@ def _adopt_remote_model_cost_map( adopt_model_cost_map(adopted) if import_fetch is not None: _carry_edits_made_during_adoption(import_fetch, adopted=adopted, carried=edits, finalized=finalized) + _PendingImportFetch.adoptions += 1 def _retry_remote_fetch_in_background( @@ -743,25 +745,36 @@ def _resume_import_fetch_after_fork() -> None: _start_import_fetch(replace(pending, first_attempt_settled=threading.Event())) -def wait_for_import_model_cost_map() -> bool: +def import_model_cost_map_generation() -> int: + """How many times the remote map fetched at import was adopted; read it before a lookup that may miss.""" + return _PendingImportFetch.adoptions + + +def wait_for_import_model_cost_map(since: int) -> bool: """ - Called by a lookup that missed the cost map while the remote map fetched at import is - still in flight. Blocks until that fetch's first attempt settles, at most one fetch - timeout, so a model the bundled map lacks resolves the way it did when import waited - for the fetch. True means the lookup should be tried once more. + Called by a lookup that missed the cost map, with the generation it read before looking up. + While the remote map fetched at import is still in flight, blocks until that fetch's first + attempt settles, at most one fetch timeout, so a model the bundled map lacks resolves the way + it did when import waited for the fetch. True means the remote map was adopted since + ``since`` and the lookup should be tried once more. """ + _wait_for_pending_import_fetch() + return _PendingImportFetch.adoptions != since + + +def _wait_for_pending_import_fetch() -> None: pending: Final = _PendingImportFetch.fetch if ( pending is None or not _litellm_import_complete.is_set() or threading.current_thread() is _PendingImportFetch.thread ): - return False + return if not pending.first_attempt_settled.is_set(): verbose_logger.debug( "LiteLLM: lookup missed the bundled cost map; waiting for the remote map fetched at import" ) - return pending.first_attempt_settled.wait(timeout=pending.timeout) + pending.first_attempt_settled.wait(timeout=pending.timeout) if hasattr(os, "register_at_fork"): diff --git a/litellm/utils.py b/litellm/utils.py index 1baacd01ae8..efd7c725d50 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -90,7 +90,10 @@ from litellm.litellm_core_utils.fallback_generalizations import ( match_capability_generalizations, match_fill_missing_generalizations, ) -from litellm.litellm_core_utils.get_model_cost_map import wait_for_import_model_cost_map +from litellm.litellm_core_utils.get_model_cost_map import ( + import_model_cost_map_generation, + wait_for_import_model_cost_map, +) from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.litellm_core_utils.tokenizer import strip_special_tokens from litellm.rust_bridge import tokenizer as tokenizer_dispatch @@ -6037,10 +6040,11 @@ def get_model_info_helper( """ Helper for 'get_model_info'. Separated out to avoid infinite loop caused by returning 'supported_openai_param's """ + generation: Final = import_model_cost_map_generation() try: return _model_info_from_cost_map(model, custom_llm_provider, api_base, api_key) except ModelNotMappedError: - if not wait_for_import_model_cost_map(): + if not wait_for_import_model_cost_map(since=generation): raise return _model_info_from_cost_map(model, custom_llm_provider, api_base, api_key) diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index 56ed3ad2bd6..8f2b90eb898 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -920,13 +920,13 @@ class _ReleaseWhenALookupWaits(logging.Handler): def __init__(self, release: threading.Event): super().__init__(level=logging.DEBUG) - self.release = release + self.fetch_release = release self.waited = threading.Event() def emit(self, record): if "waiting for the remote map fetched at import" in record.getMessage(): self.waited.set() - self.release.set() + self.fetch_release.set() @pytest.fixture() @@ -996,6 +996,26 @@ def test_a_lookup_that_misses_both_maps_still_raises_after_waiting( assert waits.waited.is_set() +@pytest.mark.parametrize("remote_status", [200, 404], ids=["adopted", "failed"]) +def test_a_lookup_that_missed_before_the_fetch_settled_retries_only_once_the_remote_map_was_adopted( + isolated_litellm_cost_state, monkeypatch, remote_status +): + from litellm.litellm_core_utils import get_model_cost_map as module + + body = _remote_map_with("zz-remote-only-test") if remote_status == 200 else b"" + client, started, release = _held_client(httpx.Response(remote_status, content=body)) + try: + monkeypatch.setattr(litellm, "model_cost", module.load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + generation_at_the_miss = module.import_model_cost_map_generation() + finally: + _finish_fetch(release) + + assert module._PendingImportFetch.fetch is None + assert module.wait_for_import_model_cost_map(since=generation_at_the_miss) is (remote_status == 200) + assert module.wait_for_import_model_cost_map(since=module.import_model_cost_map_generation()) is False + + def test_import_load_with_the_local_env_override_starts_no_fetch(monkeypatch): from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import From ef895bba48bc35f86b8f916568393ac482446e31 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 15:40:50 -0700 Subject: [PATCH 06/18] fix(cost-map): carry only the fields an edit changed when the remote map lands A Router built while the import fetch was still retrying registers its models on the bundled map, and adoption carried those whole bundled entries over the remote ones, so the models it serves kept bundled prices. Adoption now lays only the fields an edit added, changed, or removed over the remote entry, and keeps aliases sharing one entry. --- .../litellm_core_utils/get_model_cost_map.py | 47 ++++++- .../test_get_model_cost_map.py | 116 ++++++++++++++++++ 2 files changed, 158 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index d600f7134f5..a07a01808c9 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -43,8 +43,10 @@ from litellm.litellm_core_utils.fallback_generalizations import ( FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations" _CATALOG_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) _BUNDLED_CATALOG_ADAPTER: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) +_COST_ENTRY_ADAPTER: Final = TypeAdapter(dict[str, object]) _CLI_ENTRYPOINT_NAMES: Final = frozenset({"lite", "litellm-proxy"}) _NO_EDITS: Final[Mapping[str, object]] = MappingProxyType({}) +_ABSENT: Final = object() def _is_cli_process() -> bool: @@ -699,23 +701,55 @@ def _entries_changed_since_import(fetch: _ImportFetch) -> Mapping[str, object]: return MappingProxyType({key: value for key, value in current if fetch.bundled_as_loaded.get(key) != value}) +def _edit_overlaid_on( + remote: Mapping[str, object], edited: Mapping[str, object], as_loaded: Mapping[str, object] +) -> dict[str, object]: # mutable-ok: becomes a litellm.model_cost entry, a plain mutable dict by contract + removed: Final = as_loaded.keys() - edited.keys() + changed: Final = {name: value for name, value in edited.items() if as_loaded.get(name, _ABSENT) != value} + return {name: value for name, value in {**remote, **changed}.items() if name not in removed} + + +def _edit_writes( + onto: Mapping[str, object], key: str, edited: object, as_loaded: Mapping[str, object] | None +) -> Mapping[str, object]: + remote: Final = onto.get(key) + if as_loaded is None or not isinstance(remote, Mapping) or not isinstance(edited, Mapping): + return MappingProxyType({key: edited}) + overlaid: Final = _edit_overlaid_on( + _COST_ENTRY_ADAPTER.validate_python(remote), _COST_ENTRY_ADAPTER.validate_python(edited), as_loaded + ) + if not isinstance(remote, AliasedCostEntry): + return MappingProxyType({key: overlaid}) + shared: Final = AliasedCostEntry(overlaid, remote.names) + return MappingProxyType(dict.fromkeys((name for name in (key, *remote.names) if onto.get(name) is remote), shared)) + + +def _edits_overlaid_on( + fetch: _ImportFetch, onto: Mapping[str, object], edits: Mapping[str, object] +) -> Mapping[str, object]: + writes: Final = ( + _edit_writes(onto, key, edited, fetch.bundled_as_loaded.get(key)).items() for key, edited in edits.items() + ) + return MappingProxyType(dict(chain.from_iterable(writes))) + + def _carry_edits_made_during_adoption( fetch: _ImportFetch, adopted: dict[str, object], # mutable-ok: the dict just adopted as litellm.model_cost, updated in place carried: Mapping[str, object], - finalized: Mapping[str, object], + as_adopted: Mapping[str, object], ) -> None: changed: Final = tuple(_entries_changed_since_import(fetch).items()) late: Final = { key: value for key, value in changed - if value is not carried.get(key) and adopted.get(key) is carried.get(key, finalized.get(key)) + if value is not carried.get(key) and adopted.get(key) is as_adopted.get(key) } if not late or not _litellm_model_cost_is(adopted): return from litellm import utils - adopted.update(late) + adopted.update(_edits_overlaid_on(fetch, as_adopted, late)) utils._invalidate_model_cost_lowercase_map() # pyright: ignore[reportPrivateUsage] # required cache invalidation @@ -762,13 +796,16 @@ def _adopt_remote_model_cost_map( return finalized: Final = _finalize_loaded_model_cost_map(result).model_cost_map edits: Final = _entries_changed_since_import(import_fetch) if import_fetch is not None else _NO_EDITS - adopted: Final = {**finalized, **edits} + adopted: Final = ( + {**finalized, **_edits_overlaid_on(import_fetch, finalized, edits)} if import_fetch is not None else finalized + ) + as_adopted: Final = MappingProxyType(dict(adopted)) _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(adopted) if import_fetch is not None: - _carry_edits_made_during_adoption(import_fetch, adopted=adopted, carried=edits, finalized=finalized) + _carry_edits_made_during_adoption(import_fetch, adopted=adopted, carried=edits, as_adopted=as_adopted) _PendingImportFetch.adoptions += 1 diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index e5a746d55f0..81f0c5356ac 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -817,6 +817,22 @@ def _held_client(outcome): return httpx.Client(transport=httpx.MockTransport(handler)), started, release +def _rate_limited_then_held_client(body: bytes): + retry_started = threading.Event() + release = threading.Event() + served = [] + + def handler(request): + served.append(request) + if len(served) == 1: + return httpx.Response(429) + retry_started.set() + release.wait(timeout=10) + return httpx.Response(200, content=body) + + return httpx.Client(transport=httpx.MockTransport(handler)), retry_started, release + + def _finish_fetch(release): release.set() threads = _fetch_threads() @@ -994,6 +1010,106 @@ def test_an_edit_written_through_a_stale_reference_during_adoption_survives_it( assert litellm.model_cost["gpt-4o"]["input_cost_per_token"] == 7.0 +def _remote_map_repricing_gpt_4o() -> tuple[bytes, float]: + remote_map = _load_root_cost_map() + remote_price = remote_map["gpt-4o"]["input_cost_per_token"] * 10 + remote_map["gpt-4o"] = {**remote_map["gpt-4o"], "input_cost_per_token": remote_price} + return json.dumps(remote_map).encode(), remote_price + + +def test_a_router_built_before_the_remote_map_lands_prices_its_models_from_the_remote_map( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + body, remote_price = _remote_map_repricing_gpt_4o() + client, retry_started, release = _rate_limited_then_held_client(body) + try: + monkeypatch.setattr( + litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client, sleep=lambda _: None) + ) + assert retry_started.wait(timeout=10) + litellm.Router( + model_list=[ + { + "model_name": "gpt-4o", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-test"}, + } + ] + ) + finally: + _finish_fetch(release) + + assert litellm.model_cost["gpt-4o"]["input_cost_per_token"] == remote_price + assert litellm.get_model_info("openai/gpt-4o")["input_cost_per_token"] == remote_price + + +def test_a_partial_edit_made_before_the_remote_map_lands_keeps_only_the_fields_it_changed( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + body, remote_price = _remote_map_repricing_gpt_4o() + client, started, release = _held_client(httpx.Response(200, content=body)) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + edited = { + name: value for name, value in litellm.model_cost["gpt-4o"].items() if name != "output_cost_per_token" + } + litellm.model_cost["gpt-4o"] = { + **edited, + "max_tokens": 7, + } # test-quality-ok: the edit is what is under test, and the fixture restores model_cost + finally: + _finish_fetch(release) + + assert litellm.model_cost["gpt-4o"]["input_cost_per_token"] == remote_price + assert litellm.model_cost["gpt-4o"]["max_tokens"] == 7 + assert "output_cost_per_token" not in litellm.model_cost["gpt-4o"] + + +def _cost_map_with_aliased_model(input_cost_per_token: float) -> bytes: + cost_map = _load_root_cost_map() + cost_map["my-aliased-model"] = { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": input_cost_per_token, + "output_cost_per_token": 1.0, + "aliases": ["my-aliased-model-alias"], + } + return json.dumps(cost_map).encode() + + +def test_an_edit_to_an_aliased_model_made_before_the_remote_map_lands_stays_shared_by_its_aliases( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + monkeypatch.setattr(GetModelCostMap, "_loaded_catalog", GetModelCostMap._loaded_catalog) + monkeypatch.setattr(GetModelCostMap, "_loaded_catalog_fields", GetModelCostMap._loaded_catalog_fields) + monkeypatch.setattr( + GetModelCostMap, "read_local_model_cost_map_bytes", staticmethod(lambda: _cost_map_with_aliased_model(1.0)) + ) + client, retry_started, release = _rate_limited_then_held_client(_cost_map_with_aliased_model(5.0)) + try: + monkeypatch.setattr( + litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client, sleep=lambda _: None) + ) + assert retry_started.wait(timeout=10) + litellm.register_model({"my-aliased-model": {"output_cost_per_token": 9.0}}, persist_across_reloads=False) + finally: + _finish_fetch(release) + + entry = litellm.model_cost["my-aliased-model"] + assert litellm.model_cost["my-aliased-model-alias"] is entry + assert entry["input_cost_per_token"] == 5.0 + assert entry["output_cost_per_token"] == 9.0 + + class _ReleaseWhenALookupWaits(logging.Handler): """Lets the held fetch answer only once a lookup is blocked waiting on it, so the lookup provably missed first.""" From bf30bae894bc3d9edad43e8210f45871009a4ac0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 18:10:59 -0700 Subject: [PATCH 07/18] fix(cost-map): keep every edit through the background swap and wait for it where costs and forks need it Edits made since import are now found against a fresh parse of the bundled map, so nested edits and deleted models survive the swap. Cost calculations, get_valid_models, and a fork taken while the first download attempt is in flight wait for that attempt, and adoption adds remote-only models to litellm.model_list. Source info no longer stays "Remote fetch in progress" after the caller replaced litellm.model_cost. --- litellm/__init__.py | 192 +++++++++--------- litellm/cost_calculator.py | 3 + .../litellm_core_utils/get_model_cost_map.py | 112 +++++++--- litellm/utils.py | 3 +- .../test_get_model_cost_map.py | 192 +++++++++++++++++- 5 files changed, 371 insertions(+), 131 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 804e1b5f4e7..d26a329a562 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1014,6 +1014,9 @@ def add_known_models(model_cost_map: Optional[Dict] = None): """ _populate_provider_model_sets(model_cost_map if model_cost_map is not None else model_cost) models_by_provider.update(_build_models_by_provider()) + added: Final = _known_model_names() - model_list_set + model_list.extend(added) + model_list_set.update(added) _populate_provider_model_sets(model_cost) @@ -1046,98 +1049,103 @@ ollama_models = ["llama2"] maritalk_models = ["maritalk"] -model_list = list( - open_ai_chat_completion_models - | open_ai_text_completion_models - | cohere_models - | cohere_chat_models - | anthropic_models - | set(replicate_models) - | openrouter_models - | datarobot_models - | set(huggingface_models) - | vertex_chat_models - | vertex_text_models - | ai21_models - | ai21_chat_models - | set(together_ai_models) - | set(baseten_models) - | aleph_alpha_models - | nlp_cloud_models - | set(ollama_models) - | bedrock_models - | deepinfra_models - | perplexity_models - | set(maritalk_models) - | runwayml_models - | vertex_language_models - | watsonx_models - | gemini_models - | text_completion_codestral_models - | text_completion_inception_models - | xai_models - | zai_models - | fal_ai_models - | deepseek_models - | modelscope_models - | azure_ai_models - | voyage_models - | infinity_models - | databricks_models - | cloudflare_models - | codestral_models - | friendliai_models - | palm_models - | groq_models - | azure_models - | azure_anthropic_models - | anyscale_models - | cerebras_models - | nadir_models - | galadriel_models - | nvidia_nim_models - | nvidia_riva_models - | soniox_models - | sambanova_models - | scaledown_models - | azure_text_models - | novita_models - | assemblyai_models - | jina_ai_models - | snowflake_models - | gradient_ai_models - | llama_models - | featherless_ai_models - | nscale_models - | deepgram_models - | elevenlabs_models - | dashscope_models - | qwencloud_models - | qwen_ai_platform_models - | moonshot_models - | publicai_models - | darkbloom_models - | v0_models - | morph_models - | lambda_ai_models - | inception_models - | black_forest_labs_models - | recraft_models - | cometapi_models - | oci_models - | heroku_models - | vercel_ai_gateway_models - | edenai_models - | volcengine_models - | wandb_models - | ovhcloud_models - | lemonade_models - | docker_model_runner_models - | reducto_models - | bedrock_mantle_models - | coralbricks_models - | set(clarifai_models) -) + +def _known_model_names() -> Set[str]: + return ( + open_ai_chat_completion_models + | open_ai_text_completion_models + | cohere_models + | cohere_chat_models + | anthropic_models + | set(replicate_models) + | openrouter_models + | datarobot_models + | set(huggingface_models) + | vertex_chat_models + | vertex_text_models + | ai21_models + | ai21_chat_models + | set(together_ai_models) + | set(baseten_models) + | aleph_alpha_models + | nlp_cloud_models + | set(ollama_models) + | bedrock_models + | deepinfra_models + | perplexity_models + | set(maritalk_models) + | runwayml_models + | vertex_language_models + | watsonx_models + | gemini_models + | text_completion_codestral_models + | text_completion_inception_models + | xai_models + | zai_models + | fal_ai_models + | deepseek_models + | modelscope_models + | azure_ai_models + | voyage_models + | infinity_models + | databricks_models + | cloudflare_models + | codestral_models + | friendliai_models + | palm_models + | groq_models + | azure_models + | azure_anthropic_models + | anyscale_models + | cerebras_models + | nadir_models + | galadriel_models + | nvidia_nim_models + | nvidia_riva_models + | soniox_models + | sambanova_models + | scaledown_models + | azure_text_models + | novita_models + | assemblyai_models + | jina_ai_models + | snowflake_models + | gradient_ai_models + | llama_models + | featherless_ai_models + | nscale_models + | deepgram_models + | elevenlabs_models + | dashscope_models + | qwencloud_models + | qwen_ai_platform_models + | moonshot_models + | publicai_models + | darkbloom_models + | v0_models + | morph_models + | lambda_ai_models + | inception_models + | black_forest_labs_models + | recraft_models + | cometapi_models + | oci_models + | heroku_models + | vercel_ai_gateway_models + | edenai_models + | volcengine_models + | wandb_models + | ovhcloud_models + | lemonade_models + | docker_model_runner_models + | reducto_models + | bedrock_mantle_models + | coralbricks_models + | set(clarifai_models) + ) + + +model_list = list(_known_model_names()) model_list_set = set(model_list) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 3976612c6e1..7622f78b1a2 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -25,6 +25,7 @@ from litellm.constants import ( from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.cost_discount import resolve_cost_discount from litellm.litellm_core_utils.decisions_pricing import decisions_pricing_fields, is_decisions_call_type +from litellm.litellm_core_utils.get_model_cost_map import wait_for_pending_import_model_cost_map from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, @@ -459,6 +460,7 @@ def cost_per_token( if model is None: raise Exception("Invalid arg. Model cannot be none.") + wait_for_pending_import_model_cost_map() ## RECONSTRUCT USAGE BLOCK ## if usage_object is not None: @@ -1459,6 +1461,7 @@ def completion_cost( Raises if no candidate model name resolves to pricing. Register the model via `litellm.register_model` or add it to model_prices_and_context_window.json. """ + wait_for_pending_import_model_cost_map() try: call_type = infer_call_type(call_type, completion_response) or "completion" diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index a07a01808c9..d5a33279829 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -45,8 +45,8 @@ _CATALOG_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) _BUNDLED_CATALOG_ADAPTER: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) _COST_ENTRY_ADAPTER: Final = TypeAdapter(dict[str, object]) _CLI_ENTRYPOINT_NAMES: Final = frozenset({"lite", "litellm-proxy"}) -_NO_EDITS: Final[Mapping[str, object]] = MappingProxyType({}) _ABSENT: Final = object() +_IMPORT_FETCH_IN_PROGRESS: Final = "Remote fetch in progress" def _is_cli_process() -> bool: @@ -127,7 +127,7 @@ class GetModelCostMap: _get_backup_model_count = get_backup_model_count @staticmethod - def _check_is_valid_dict(fetched_map: dict) -> bool: + def _check_is_valid_dict(fetched_map: object) -> bool: """Check 1: fetched map is a non-empty dict.""" if not isinstance(fetched_map, dict): verbose_logger.warning( @@ -676,7 +676,7 @@ class _ImportFetch: rng: random.Random client: _SyncGetClient bundled: Mapping[str, object] - bundled_as_loaded: Mapping[str, Mapping[str, object]] + local: Callable[[], ModelCostMapReloaded] first_attempt_settled: threading.Event = field(default_factory=threading.Event) @@ -696,9 +696,23 @@ def _model_cost_replaced_since_import(fetch: _ImportFetch | None) -> bool: return fetch is not None and not _litellm_model_cost_is(fetch.bundled) -def _entries_changed_since_import(fetch: _ImportFetch) -> Mapping[str, object]: +def _bundled_as_loaded(fetch: _ImportFetch) -> Mapping[str, Mapping[str, object]]: + raw: Final = fetch.local().model_cost_map + catalog: Final = _CATALOG_ADAPTER.validate_python( + {key: value for key, value in raw.items() if key != FALLBACK_GENERALIZATIONS_KEY} + ) + return MappingProxyType(_expand_model_aliases(catalog)) + + +def _entries_changed_since_import( + fetch: _ImportFetch, as_loaded: Mapping[str, Mapping[str, object]] +) -> Mapping[str, object]: current: Final = tuple(fetch.bundled.items()) - return MappingProxyType({key: value for key, value in current if fetch.bundled_as_loaded.get(key) != value}) + return MappingProxyType({key: value for key, value in current if as_loaded.get(key) != value}) + + +def _keys_removed_since_import(fetch: _ImportFetch, as_loaded: Mapping[str, Mapping[str, object]]) -> frozenset[str]: + return frozenset(key for key in as_loaded if key not in fetch.bundled) def _edit_overlaid_on( @@ -725,31 +739,36 @@ def _edit_writes( def _edits_overlaid_on( - fetch: _ImportFetch, onto: Mapping[str, object], edits: Mapping[str, object] + as_loaded: Mapping[str, Mapping[str, object]], onto: Mapping[str, object], edits: Mapping[str, object] ) -> Mapping[str, object]: - writes: Final = ( - _edit_writes(onto, key, edited, fetch.bundled_as_loaded.get(key)).items() for key, edited in edits.items() - ) + writes: Final = (_edit_writes(onto, key, edited, as_loaded.get(key)).items() for key, edited in edits.items()) return MappingProxyType(dict(chain.from_iterable(writes))) def _carry_edits_made_during_adoption( fetch: _ImportFetch, + as_loaded: Mapping[str, Mapping[str, object]], adopted: dict[str, object], # mutable-ok: the dict just adopted as litellm.model_cost, updated in place carried: Mapping[str, object], + removed: frozenset[str], as_adopted: Mapping[str, object], ) -> None: - changed: Final = tuple(_entries_changed_since_import(fetch).items()) + changed: Final = tuple(_entries_changed_since_import(fetch, as_loaded).items()) late: Final = { key: value for key, value in changed if value is not carried.get(key) and adopted.get(key) is as_adopted.get(key) } - if not late or not _litellm_model_cost_is(adopted): + late_removed: Final = tuple( + key for key in _keys_removed_since_import(fetch, as_loaded) - removed if adopted.get(key) is as_adopted.get(key) + ) + if not (late or late_removed) or not _litellm_model_cost_is(adopted): return from litellm import utils - adopted.update(_edits_overlaid_on(fetch, as_adopted, late)) + adopted.update(_edits_overlaid_on(as_loaded, as_adopted, late)) + for key in late_removed: + adopted.pop(key, None) utils._invalidate_model_cost_lowercase_map() # pyright: ignore[reportPrivateUsage] # required cache invalidation @@ -795,18 +814,35 @@ def _adopt_remote_model_cost_map( verbose_logger.debug("LiteLLM: litellm.model_cost was replaced after import; not adopting the remote map") return finalized: Final = _finalize_loaded_model_cost_map(result).model_cost_map - edits: Final = _entries_changed_since_import(import_fetch) if import_fetch is not None else _NO_EDITS - adopted: Final = ( - {**finalized, **_edits_overlaid_on(import_fetch, finalized, edits)} if import_fetch is not None else finalized - ) - as_adopted: Final = MappingProxyType(dict(adopted)) + if import_fetch is None: + _record_remote_adoption() + adopt_model_cost_map(finalized) + return + _adopt_keeping_edits_since_import(import_fetch, finalized) + + +def _record_remote_adoption() -> None: _cost_map_source_info.source = "remote" _cost_map_source_info.fallback_reason = None _cost_map_source_info.loaded_at = datetime.now(timezone.utc) + + +def _adopt_keeping_edits_since_import(fetch: _ImportFetch, finalized: Mapping[str, object]) -> None: + as_loaded: Final = _bundled_as_loaded(fetch) + edits: Final = _entries_changed_since_import(fetch, as_loaded) + removed: Final = _keys_removed_since_import(fetch, as_loaded) + adopted: Final = { + key: value + for key, value in {**finalized, **_edits_overlaid_on(as_loaded, finalized, edits)}.items() + if key not in removed + } + as_adopted: Final = MappingProxyType(dict(adopted)) + _record_remote_adoption() adopt_model_cost_map(adopted) - if import_fetch is not None: - _carry_edits_made_during_adoption(import_fetch, adopted=adopted, carried=edits, as_adopted=as_adopted) - _PendingImportFetch.adoptions += 1 + _carry_edits_made_during_adoption( + fetch, as_loaded, adopted=adopted, carried=edits, removed=removed, as_adopted=as_adopted + ) + _PendingImportFetch.adoptions += 1 def _retry_remote_fetch_in_background( @@ -854,9 +890,14 @@ def _run_import_fetch(fetch: _ImportFetch) -> None: except Exception as e: # noqa: BLE001 # a failed background fetch must not kill the thread; the backup stays verbose_logger.warning("LiteLLM: Background model cost map fetch failed: %s", e) finally: - fetch.first_attempt_settled.set() if _PendingImportFetch.fetch is fetch: _PendingImportFetch.fetch = None + if ( + _model_cost_replaced_since_import(fetch) + and _cost_map_source_info.fallback_reason == _IMPORT_FETCH_IN_PROGRESS + ): + _cost_map_source_info.fallback_reason = "litellm.model_cost was replaced after import" + fetch.first_attempt_settled.set() def _start_import_fetch(fetch: _ImportFetch) -> None: @@ -887,11 +928,11 @@ def wait_for_import_model_cost_map(since: int) -> bool: it did when import waited for the fetch. True means the remote map was adopted since ``since`` and the lookup should be tried once more. """ - _wait_for_pending_import_fetch() + wait_for_pending_import_model_cost_map() return _PendingImportFetch.adoptions != since -def _wait_for_pending_import_fetch() -> None: +def wait_for_pending_import_model_cost_map() -> None: pending: Final = _PendingImportFetch.fetch if ( pending is None @@ -900,14 +941,22 @@ def _wait_for_pending_import_fetch() -> None: ): return if not pending.first_attempt_settled.is_set(): - verbose_logger.debug( - "LiteLLM: lookup missed the bundled cost map; waiting for the remote map fetched at import" - ) + verbose_logger.debug("LiteLLM: waiting for the remote map fetched at import") pending.first_attempt_settled.wait(timeout=pending.timeout) +def _settle_import_fetch_before_fork() -> None: + pending: Final = _PendingImportFetch.fetch + thread: Final = _PendingImportFetch.thread + if pending is None or thread is None or thread is threading.current_thread(): + return + wait_for_pending_import_model_cost_map() + if _PendingImportFetch.fetch is None: + thread.join(timeout=pending.timeout) + + if hasattr(os, "register_at_fork"): - os.register_at_fork(after_in_child=_resume_import_fetch_after_fork) + os.register_at_fork(before=_settle_import_fetch_before_fork, after_in_child=_resume_import_fetch_after_fork) def load_model_cost_map_at_import( @@ -917,6 +966,7 @@ def load_model_cost_map_at_import( sleep: Callable[[float], None] = time.sleep, rng: random.Random | None = None, client: "_SyncGetClient | None" = None, + local: Callable[[], ModelCostMapReloaded] = GetModelCostMap.load_local_model_cost_map_with_revision, ) -> dict: # mutable-ok: becomes litellm.model_cost, a plain mutable dict by contract """ Return the bundled map right away and fetch the remote one in a background thread, @@ -932,10 +982,8 @@ def load_model_cost_map_at_import( _cost_map_source_info.source = "local" _cost_map_source_info.url = url _cost_map_source_info.is_env_forced = False - _cost_map_source_info.fallback_reason = "Remote fetch in progress" - bundled: Final = _finalize_loaded_model_cost_map( - GetModelCostMap.load_local_model_cost_map_with_revision() - ).model_cost_map + _cost_map_source_info.fallback_reason = _IMPORT_FETCH_IN_PROGRESS + bundled: Final = _finalize_loaded_model_cost_map(local()).model_cost_map _start_import_fetch( _ImportFetch( url=url, @@ -945,7 +993,7 @@ def load_model_cost_map_at_import( rng=rng if rng is not None else random.Random(), client=client if client is not None else httpx, bundled=bundled, - bundled_as_loaded=GetModelCostMap.loaded_model_cost_map(), + local=local, ) ) return bundled diff --git a/litellm/utils.py b/litellm/utils.py index 6367f97cdb4..87bf8159a22 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -93,6 +93,7 @@ from litellm.litellm_core_utils.get_model_cost_map import ( AliasedCostEntry, import_model_cost_map_generation, wait_for_import_model_cost_map, + wait_for_pending_import_model_cost_map, ) from litellm.litellm_core_utils.model_param_helper import peek_reasoning_summary_aliases from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload @@ -8206,7 +8207,7 @@ def get_valid_models( Returns: A list of valid LLMs """ - + wait_for_pending_import_model_cost_map() try: ################################ # init litellm_params diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index 81f0c5356ac..f6efcd9cd9e 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -611,6 +611,8 @@ def isolated_litellm_cost_state(monkeypatch): 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, "model_list", list(litellm.model_list)) + monkeypatch.setattr(litellm, "model_list_set", set(litellm.model_list_set)) monkeypatch.setattr( litellm_utils, "_runtime_registered_model_cost", @@ -619,6 +621,8 @@ def isolated_litellm_cost_state(monkeypatch): 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)) + for name in ("_loaded_catalog", "_loaded_catalog_fields"): + monkeypatch.setattr(GetModelCostMap, name, getattr(GetModelCostMap, name)) def test_boot_load_transient_failure_returns_local_then_background_retry_adopts_remote(isolated_litellm_cost_state): @@ -867,6 +871,8 @@ def test_import_load_returns_the_bundled_map_before_the_remote_fetch_answers_the assert "claude-remote-only-test" in litellm.model_cost assert "claude-remote-only-test" in litellm.anthropic_models + assert "claude-remote-only-test" in litellm.model_list + assert "claude-remote-only-test" in litellm.model_list_set source = get_model_cost_map_source_info() assert source["source"] == "remote" assert source["fallback_reason"] is None @@ -989,9 +995,11 @@ def test_an_edit_written_through_a_stale_reference_during_adoption_survives_it( monkeypatch.setattr(litellm, "model_cost", bundled) original_replay = litellm_utils._LiveDeploymentReplay.callback late_entry = {"litellm_provider": "openai", "mode": "chat"} + late_removal = next(key for key in bundled if key not in ("gpt-4o", "gpt-4o-mini")) def write_through_the_reference_read_before_the_swap(): bundled["my-late-edit"] = late_entry + del bundled[late_removal] bundled["gpt-4o-mini"] = {**bundled["gpt-4o-mini"], "input_cost_per_token": 42.0} bundled["gpt-4o"] = {**bundled["gpt-4o"], "input_cost_per_token": 42.0} litellm.model_cost["gpt-4o"] = {**litellm.model_cost["gpt-4o"], "input_cost_per_token": 7.0} @@ -1006,6 +1014,7 @@ def test_an_edit_written_through_a_stale_reference_during_adoption_survives_it( assert litellm.model_cost is not bundled assert "claude-remote-only-test" in litellm.model_cost assert litellm.model_cost["my-late-edit"] is late_entry + assert late_removal not in litellm.model_cost assert litellm.model_cost["gpt-4o-mini"]["input_cost_per_token"] == 42.0 assert litellm.model_cost["gpt-4o"]["input_cost_per_token"] == 7.0 @@ -1071,6 +1080,68 @@ def test_a_partial_edit_made_before_the_remote_map_lands_keeps_only_the_fields_i assert "output_cost_per_token" not in litellm.model_cost["gpt-4o"] +def test_a_model_deleted_before_the_remote_map_lands_stays_deleted(isolated_litellm_cost_state, monkeypatch): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + del litellm.model_cost["gpt-4o-mini"] # test-quality-ok: the deletion is under test; the fixture restores it + finally: + _finish_fetch(release) + + assert "claude-remote-only-test" in litellm.model_cost + assert "gpt-4o-mini" not in litellm.model_cost + + +def _nested_fields(cost_map): + for model, entry in cost_map.items(): + for field, value in entry.items(): + if isinstance(value, dict): + yield model, field + + +def test_a_nested_edit_made_before_the_remote_map_lands_survives_it(isolated_litellm_cost_state, monkeypatch): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + model, field = next(_nested_fields(litellm.model_cost)) + litellm.model_cost[model][field][ + "nested-edit-test" + ] = 99.0 # test-quality-ok: the edit is what is under test, and the fixture restores model_cost + finally: + _finish_fetch(release) + + assert "claude-remote-only-test" in litellm.model_cost + assert litellm.model_cost[model][field]["nested-edit-test"] == 99.0 + + +def test_import_load_reports_that_the_caller_replaced_the_model_cost_it_returned( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + monkeypatch.setattr(litellm, "model_cost", {"my-model": {"litellm_provider": "openai", "mode": "chat"}}) + finally: + _finish_fetch(release) + + assert "claude-remote-only-test" not in litellm.model_cost + source = get_model_cost_map_source_info() + assert source["source"] == "local" + assert source["fallback_reason"] == "litellm.model_cost was replaced after import" + + def _cost_map_with_aliased_model(input_cost_per_token: float) -> bytes: cost_map = _load_root_cost_map() cost_map["my-aliased-model"] = { @@ -1089,15 +1160,17 @@ def test_an_edit_to_an_aliased_model_made_before_the_remote_map_lands_stays_shar import litellm from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import - monkeypatch.setattr(GetModelCostMap, "_loaded_catalog", GetModelCostMap._loaded_catalog) - monkeypatch.setattr(GetModelCostMap, "_loaded_catalog_fields", GetModelCostMap._loaded_catalog_fields) - monkeypatch.setattr( - GetModelCostMap, "read_local_model_cost_map_bytes", staticmethod(lambda: _cost_map_with_aliased_model(1.0)) - ) client, retry_started, release = _rate_limited_then_held_client(_cost_map_with_aliased_model(5.0)) try: monkeypatch.setattr( - litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client, sleep=lambda _: None) + litellm, + "model_cost", + load_model_cost_map_at_import( + url=_URL, + client=client, + sleep=lambda _: None, + local=lambda: ModelCostMapReloaded(model_cost_map=json.loads(_cost_map_with_aliased_model(1.0))), + ), ) assert retry_started.wait(timeout=10) litellm.register_model({"my-aliased-model": {"output_cost_per_token": 9.0}}, persist_across_reloads=False) @@ -1211,6 +1284,113 @@ def test_a_lookup_that_missed_before_the_fetch_settled_retries_only_once_the_rem assert module.wait_for_import_model_cost_map(since=module.import_model_cost_map_generation()) is False +@pytest.mark.parametrize( + "price_a_thousand_prompt_tokens", + [ + pytest.param( + lambda: litellm.cost_per_token(model="gpt-4o", prompt_tokens=1000, completion_tokens=0)[0], + id="cost_per_token", + ), + pytest.param( + lambda: litellm.completion_cost( + completion_response=litellm.ModelResponse( + model="gpt-4o", usage=litellm.Usage(prompt_tokens=1000, completion_tokens=0, total_tokens=1000) + ), + model="gpt-4o", + ), + id="completion_cost", + ), + ], +) +def test_a_cost_computed_while_the_remote_map_is_in_flight_uses_its_price( + isolated_litellm_cost_state, monkeypatch, release_when_a_lookup_waits, price_a_thousand_prompt_tokens +): + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + body, remote_price = _remote_map_repricing_gpt_4o() + client, started, release = _held_client(httpx.Response(200, content=body)) + waits = release_when_a_lookup_waits(release) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + cost = price_a_thousand_prompt_tokens() + finally: + _finish_fetch(release) + + assert waits.waited.is_set() + assert cost == pytest.approx(1000 * remote_price) + + +def test_an_image_priced_while_the_remote_map_is_in_flight_finds_a_model_only_the_remote_map_has( + isolated_litellm_cost_state, monkeypatch, release_when_a_lookup_waits +): + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + from litellm.types.utils import ImageObject, ImageResponse + + remote_map = _load_root_cost_map() + remote_map["zz-remote-only-image-test"] = { + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_image": 0.25, + } + client, started, release = _held_client(httpx.Response(200, content=json.dumps(remote_map).encode())) + waits = release_when_a_lookup_waits(release) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + cost = litellm.completion_cost( + completion_response=ImageResponse(created=1, data=[ImageObject(url="https://example.com/image.png")]), + model="zz-remote-only-image-test", + custom_llm_provider="openai", + call_type="image_generation", + ) + finally: + _finish_fetch(release) + + assert waits.waited.is_set() + assert cost == pytest.approx(0.25) + + +def test_models_listed_while_the_remote_map_is_in_flight_include_a_model_only_the_remote_map_has( + isolated_litellm_cost_state, monkeypatch, release_when_a_lookup_waits +): + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) + waits = release_when_a_lookup_waits(release) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + valid_models = litellm.get_valid_models(custom_llm_provider="anthropic") + finally: + _finish_fetch(release) + + assert waits.waited.is_set() + assert "claude-remote-only-test" in valid_models + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") +def test_a_process_forked_while_the_remote_map_is_in_flight_inherits_it_once_adopted( + isolated_litellm_cost_state, monkeypatch, release_when_a_lookup_waits +): + from litellm.litellm_core_utils import get_model_cost_map as module + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) + release_when_a_lookup_waits(release) + try: + monkeypatch.setattr(litellm, "model_cost", module.load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + child = os.fork() + if child == 0: + inherited = "claude-remote-only-test" in litellm.model_cost and module._PendingImportFetch.fetch is None + os._exit(0 if inherited else 1) + _, wait_status = os.waitpid(child, 0) + finally: + _finish_fetch(release) + + assert os.waitstatus_to_exitcode(wait_status) == 0 + + def test_import_load_with_the_local_env_override_starts_no_fetch(monkeypatch): from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import From 0e2bf0b41aed31b1be4a3a52264898687366aecb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 20:39:31 -0700 Subject: [PATCH 08/18] fix(cost-map): merge the remote map into the import map in place and wait for it on model info, costs, and forks References taken at import now see the remote map, pins equal to bundled values survive, the cost wait has no timeout, a forked child keeps the parent's retry delay, and register_model no longer lists custom deployment names --- litellm/__init__.py | 8 +- .../litellm_core_utils/get_model_cost_map.py | 205 ++++++++++-------- litellm/utils.py | 23 +- .../test_get_model_cost_map.py | 192 ++++++++++++---- tests/unit/test_utils.py | 18 ++ 5 files changed, 295 insertions(+), 151 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index d26a329a562..92588afb0a1 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1014,7 +1014,13 @@ def add_known_models(model_cost_map: Optional[Dict] = None): """ _populate_provider_model_sets(model_cost_map if model_cost_map is not None else model_cost) models_by_provider.update(_build_models_by_provider()) - added: Final = _known_model_names() - model_list_set + + +def add_known_models_to_model_list() -> None: + """``add_known_models``, plus listing every model name it newly made known in ``model_list``.""" + known_before: Final = _known_model_names() + add_known_models() + added: Final = _known_model_names() - known_before - model_list_set model_list.extend(added) model_list_set.update(added) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index d5a33279829..707fbfcad2f 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -29,7 +29,12 @@ from typing import Final, Protocol, cast # noqa: TID251 # adapts boto3 and bot import httpx from pydantic import ConfigDict, TypeAdapter -from typing_extensions import NotRequired, ReadOnly, TypedDict +from typing_extensions import ( + NotRequired, + ReadOnly, + TypedDict, + TypeIs, # noqa: TID251 # narrows live cost map entries without copying them +) from litellm import verbose_logger from litellm.constants import ( @@ -43,7 +48,6 @@ from litellm.litellm_core_utils.fallback_generalizations import ( FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations" _CATALOG_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) _BUNDLED_CATALOG_ADAPTER: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) -_COST_ENTRY_ADAPTER: Final = TypeAdapter(dict[str, object]) _CLI_ENTRYPOINT_NAMES: Final = frozenset({"lite", "litellm-proxy"}) _ABSENT: Final = object() _IMPORT_FETCH_IN_PROGRESS: Final = "Remote fetch in progress" @@ -675,17 +679,30 @@ class _ImportFetch: sleep: Callable[[float], None] rng: random.Random client: _SyncGetClient - bundled: Mapping[str, object] + bundled: dict[str, object] # mutable-ok: litellm.model_cost itself, which the remote map is merged into + loaded: Mapping[str, Mapping[str, object]] local: Callable[[], ModelCostMapReloaded] + first_attempt: int = 1 + first_attempt_at: float = 0.0 first_attempt_settled: threading.Event = field(default_factory=threading.Event) +@dataclass(frozen=True, slots=True) +class _ScheduledAttempt: + attempt: int + at: float + + class _PendingImportFetch: fetch: _ImportFetch | None = None thread: threading.Thread | None = None + next_attempt: _ScheduledAttempt | None = None adoptions: int = 0 +_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) + + def _litellm_model_cost_is(expected: Mapping[str, object]) -> bool: import litellm @@ -696,6 +713,12 @@ def _model_cost_replaced_since_import(fetch: _ImportFetch | None) -> bool: return fetch is not None and not _litellm_model_cost_is(fetch.bundled) +def _is_cost_entry( + value: object, +) -> TypeIs[dict[str, object]]: # guard-ok: cost map keys are str # mutable-ok: entries are merged in place + return isinstance(value, dict) + + def _bundled_as_loaded(fetch: _ImportFetch) -> Mapping[str, Mapping[str, object]]: raw: Final = fetch.local().model_cost_map catalog: Final = _CATALOG_ADAPTER.validate_python( @@ -704,72 +727,59 @@ def _bundled_as_loaded(fetch: _ImportFetch) -> Mapping[str, Mapping[str, object] return MappingProxyType(_expand_model_aliases(catalog)) -def _entries_changed_since_import( - fetch: _ImportFetch, as_loaded: Mapping[str, Mapping[str, object]] -) -> Mapping[str, object]: - current: Final = tuple(fetch.bundled.items()) - return MappingProxyType({key: value for key, value in current if as_loaded.get(key) != value}) - - -def _keys_removed_since_import(fetch: _ImportFetch, as_loaded: Mapping[str, Mapping[str, object]]) -> frozenset[str]: - return frozenset(key for key in as_loaded if key not in fetch.bundled) - - -def _edit_overlaid_on( - remote: Mapping[str, object], edited: Mapping[str, object], as_loaded: Mapping[str, object] -) -> dict[str, object]: # mutable-ok: becomes a litellm.model_cost entry, a plain mutable dict by contract - removed: Final = as_loaded.keys() - edited.keys() - changed: Final = {name: value for name, value in edited.items() if as_loaded.get(name, _ABSENT) != value} - return {name: value for name, value in {**remote, **changed}.items() if name not in removed} - - -def _edit_writes( - onto: Mapping[str, object], key: str, edited: object, as_loaded: Mapping[str, object] | None -) -> Mapping[str, object]: - remote: Final = onto.get(key) - if as_loaded is None or not isinstance(remote, Mapping) or not isinstance(edited, Mapping): - return MappingProxyType({key: edited}) - overlaid: Final = _edit_overlaid_on( - _COST_ENTRY_ADAPTER.validate_python(remote), _COST_ENTRY_ADAPTER.validate_python(edited), as_loaded +def _fields_written_since_import( + current: Mapping[str, object], loaded: Mapping[str, object], parsed: Mapping[str, object] +) -> frozenset[str]: + rewritten: Final = ( + name + for name, value in tuple(current.items()) + if loaded.get(name, _ABSENT) is not value or (isinstance(value, (dict, list)) and value != parsed.get(name)) ) - if not isinstance(remote, AliasedCostEntry): - return MappingProxyType({key: overlaid}) - shared: Final = AliasedCostEntry(overlaid, remote.names) - return MappingProxyType(dict.fromkeys((name for name in (key, *remote.names) if onto.get(name) is remote), shared)) + return frozenset(chain(rewritten, loaded.keys() - current.keys())) -def _edits_overlaid_on( - as_loaded: Mapping[str, Mapping[str, object]], onto: Mapping[str, object], edits: Mapping[str, object] -) -> Mapping[str, object]: - writes: Final = (_edit_writes(onto, key, edited, as_loaded.get(key)).items() for key, edited in edits.items()) - return MappingProxyType(dict(chain.from_iterable(writes))) - - -def _carry_edits_made_during_adoption( - fetch: _ImportFetch, - as_loaded: Mapping[str, Mapping[str, object]], - adopted: dict[str, object], # mutable-ok: the dict just adopted as litellm.model_cost, updated in place - carried: Mapping[str, object], - removed: frozenset[str], - as_adopted: Mapping[str, object], +def _merge_remote_fields( + current: dict[str, object], # mutable-ok: a live litellm.model_cost entry, so references to it see the remote map + remote: Mapping[str, object], + written: frozenset[str], ) -> None: - changed: Final = tuple(_entries_changed_since_import(fetch, as_loaded).items()) - late: Final = { - key: value - for key, value in changed - if value is not carried.get(key) and adopted.get(key) is as_adopted.get(key) - } - late_removed: Final = tuple( - key for key in _keys_removed_since_import(fetch, as_loaded) - removed if adopted.get(key) is as_adopted.get(key) - ) - if not (late or late_removed) or not _litellm_model_cost_is(adopted): - return + current.update({name: value for name, value in remote.items() if name not in written}) + for name in current.keys() - remote.keys() - written: + current.pop(name, None) + + +def _untouched_since_import(current: object, loaded: Mapping[str, object], parsed: Mapping[str, object]) -> bool: + return _is_cost_entry(current) and not _fields_written_since_import(current, loaded, parsed) + + +def _merge_into_import_map(fetch: _ImportFetch, finalized: Mapping[str, object]) -> None: + import litellm from litellm import utils - adopted.update(_edits_overlaid_on(as_loaded, as_adopted, late)) - for key in late_removed: - adopted.pop(key, None) + parsed: Final = _bundled_as_loaded(fetch) + remote: Final = tuple(finalized.items()) + shared_once: Final = { + id(current): (current, key, entry) + for key, entry in remote + if key in fetch.loaded and _is_cost_entry(current := fetch.bundled.get(key)) and _is_cost_entry(entry) + } + for current, key, entry in shared_once.values(): + _merge_remote_fields( + current, entry, _fields_written_since_import(current, fetch.loaded[key], parsed.get(key, _NO_FIELDS)) + ) + fetch.bundled.update({key: entry for key, entry in remote if key not in fetch.loaded and key not in fetch.bundled}) + dropped: Final = tuple( + key + for key in fetch.loaded.keys() - finalized.keys() + if _untouched_since_import(fetch.bundled.get(key), fetch.loaded[key], parsed.get(key, _NO_FIELDS)) + ) + for key in dropped: + fetch.bundled.pop(key, None) + _record_remote_adoption() utils._invalidate_model_cost_lowercase_map() # pyright: ignore[reportPrivateUsage] # required cache invalidation + litellm.add_known_models_to_model_list() + utils.reapply_runtime_model_cost_registrations() + _PendingImportFetch.adoptions += 1 def _adopt_remote_model_cost_map( @@ -818,7 +828,7 @@ def _adopt_remote_model_cost_map( _record_remote_adoption() adopt_model_cost_map(finalized) return - _adopt_keeping_edits_since_import(import_fetch, finalized) + _merge_into_import_map(import_fetch, finalized) def _record_remote_adoption() -> None: @@ -827,24 +837,6 @@ def _record_remote_adoption() -> None: _cost_map_source_info.loaded_at = datetime.now(timezone.utc) -def _adopt_keeping_edits_since_import(fetch: _ImportFetch, finalized: Mapping[str, object]) -> None: - as_loaded: Final = _bundled_as_loaded(fetch) - edits: Final = _entries_changed_since_import(fetch, as_loaded) - removed: Final = _keys_removed_since_import(fetch, as_loaded) - adopted: Final = { - key: value - for key, value in {**finalized, **_edits_overlaid_on(as_loaded, finalized, edits)}.items() - if key not in removed - } - as_adopted: Final = MappingProxyType(dict(adopted)) - _record_remote_adoption() - adopt_model_cost_map(adopted) - _carry_edits_made_during_adoption( - fetch, as_loaded, adopted=adopted, carried=edits, removed=removed, as_adopted=as_adopted - ) - _PendingImportFetch.adoptions += 1 - - def _retry_remote_fetch_in_background( url: str, timeout: int, @@ -874,14 +866,21 @@ def _retry_remote_fetch_in_background( def _run_import_fetch(fetch: _ImportFetch) -> None: def settle_then_sleep(seconds: float) -> None: + scheduled: Final = _PendingImportFetch.next_attempt + _PendingImportFetch.next_attempt = _ScheduledAttempt( + attempt=(scheduled.attempt if scheduled is not None else fetch.first_attempt) + 1, + at=time.monotonic() + seconds, + ) fetch.first_attempt_settled.set() fetch.sleep(seconds) try: + if fetch.first_attempt > 1: + fetch.sleep(max(0.0, fetch.first_attempt_at - time.monotonic())) _adopt_remote_model_cost_map( url=fetch.url, timeout=fetch.timeout, - attempts=range(1, fetch.max_attempts + 1), + attempts=range(fetch.first_attempt, fetch.max_attempts + 1), sleep=settle_then_sleep, rng=fetch.rng, client=fetch.client, @@ -909,10 +908,28 @@ def _start_import_fetch(fetch: _ImportFetch) -> None: thread.start() +def _settled_event() -> threading.Event: + settled: Final = threading.Event() + settled.set() + return settled + + def _resume_import_fetch_after_fork() -> None: pending: Final = _PendingImportFetch.fetch - if pending is not None: + scheduled: Final = _PendingImportFetch.next_attempt + if pending is None: + return + if scheduled is None: _start_import_fetch(replace(pending, first_attempt_settled=threading.Event())) + return + _start_import_fetch( + replace( + pending, + first_attempt=scheduled.attempt, + first_attempt_at=scheduled.at, + first_attempt_settled=_settled_event(), + ) + ) def import_model_cost_map_generation() -> int: @@ -924,9 +941,9 @@ def wait_for_import_model_cost_map(since: int) -> bool: """ Called by a lookup that missed the cost map, with the generation it read before looking up. While the remote map fetched at import is still in flight, blocks until that fetch's first - attempt settles, at most one fetch timeout, so a model the bundled map lacks resolves the way - it did when import waited for the fetch. True means the remote map was adopted since - ``since`` and the lookup should be tried once more. + attempt settles, so a model the bundled map lacks resolves the way it did when import waited + for the fetch. True means the remote map was adopted since ``since`` and the lookup should be + tried once more. """ wait_for_pending_import_model_cost_map() return _PendingImportFetch.adoptions != since @@ -942,7 +959,7 @@ def wait_for_pending_import_model_cost_map() -> None: return if not pending.first_attempt_settled.is_set(): verbose_logger.debug("LiteLLM: waiting for the remote map fetched at import") - pending.first_attempt_settled.wait(timeout=pending.timeout) + pending.first_attempt_settled.wait() def _settle_import_fetch_before_fork() -> None: @@ -970,11 +987,13 @@ def load_model_cost_map_at_import( ) -> dict: # mutable-ok: becomes litellm.model_cost, a plain mutable dict by contract """ Return the bundled map right away and fetch the remote one in a background thread, - so ``import litellm`` never waits on the network. The remote map replaces - ``litellm.model_cost`` once it passes the integrity check, keeping every entry - written since import, unless the caller has replaced ``litellm.model_cost`` outright. - A lookup that misses the bundled map while the fetch is in flight waits for it, see - ``wait_for_import_model_cost_map``. A process forked before the fetch finishes starts its own fetch. + so ``import litellm`` never waits on the network. Once the remote map passes the + integrity check it is merged into the returned dict in place, so references taken + at import see it too, and every field written since import keeps its value. A caller + that replaced ``litellm.model_cost`` outright keeps its own map. Costs and model info + looked up while the first attempt is in flight wait for it, see + ``wait_for_pending_import_model_cost_map``. A process forked during a retry wait + resumes the retry schedule on its own. """ if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true" or _is_cli_process(): return get_model_cost_map(url=url, timeout=timeout) @@ -984,6 +1003,7 @@ def load_model_cost_map_at_import( _cost_map_source_info.is_env_forced = False _cost_map_source_info.fallback_reason = _IMPORT_FETCH_IN_PROGRESS bundled: Final = _finalize_loaded_model_cost_map(local()).model_cost_map + _PendingImportFetch.next_attempt = None _start_import_fetch( _ImportFetch( url=url, @@ -993,6 +1013,7 @@ def load_model_cost_map_at_import( rng=rng if rng is not None else random.Random(), client=client if client is not None else httpx, bundled=bundled, + loaded=GetModelCostMap.loaded_model_cost_map(), local=local, ) ) diff --git a/litellm/utils.py b/litellm/utils.py index 87bf8159a22..e83a442e199 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -89,12 +89,7 @@ from litellm.litellm_core_utils.fallback_generalizations import ( match_capability_generalizations, match_fill_missing_generalizations, ) -from litellm.litellm_core_utils.get_model_cost_map import ( - AliasedCostEntry, - import_model_cost_map_generation, - wait_for_import_model_cost_map, - wait_for_pending_import_model_cost_map, -) +from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry, wait_for_pending_import_model_cost_map from litellm.litellm_core_utils.model_param_helper import peek_reasoning_summary_aliases from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.litellm_core_utils.tokenizer import strip_special_tokens @@ -6307,21 +6302,7 @@ def get_model_info_helper( """ Helper for 'get_model_info'. Separated out to avoid infinite loop caused by returning 'supported_openai_param's """ - generation: Final = import_model_cost_map_generation() - try: - return _model_info_from_cost_map(model, custom_llm_provider, api_base, api_key) - except ModelNotMappedError: - if not wait_for_import_model_cost_map(since=generation): - raise - return _model_info_from_cost_map(model, custom_llm_provider, api_base, api_key) - - -def _model_info_from_cost_map( - model: str, - custom_llm_provider: str | None, - api_base: str | None, - api_key: str | None, -) -> ModelInfoBase: + wait_for_pending_import_model_cost_map() from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider try: diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index f6efcd9cd9e..7843b1a65da 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -982,43 +982,6 @@ def test_a_fetch_resumed_after_fork_adopts_the_remote_map_while_the_inherited_on assert len(served) == 2 -def test_an_edit_written_through_a_stale_reference_during_adoption_survives_it( - isolated_litellm_cost_state, monkeypatch -): - import litellm - from litellm import utils as litellm_utils - from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import - - client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) - monkeypatch.setattr(litellm_utils, "_runtime_registered_model_cost", {}) - bundled = load_model_cost_map_at_import(url=_URL, client=client) - monkeypatch.setattr(litellm, "model_cost", bundled) - original_replay = litellm_utils._LiveDeploymentReplay.callback - late_entry = {"litellm_provider": "openai", "mode": "chat"} - late_removal = next(key for key in bundled if key not in ("gpt-4o", "gpt-4o-mini")) - - def write_through_the_reference_read_before_the_swap(): - bundled["my-late-edit"] = late_entry - del bundled[late_removal] - bundled["gpt-4o-mini"] = {**bundled["gpt-4o-mini"], "input_cost_per_token": 42.0} - bundled["gpt-4o"] = {**bundled["gpt-4o"], "input_cost_per_token": 42.0} - litellm.model_cost["gpt-4o"] = {**litellm.model_cost["gpt-4o"], "input_cost_per_token": 7.0} - - litellm_utils.set_live_deployment_replay(write_through_the_reference_read_before_the_swap) - try: - assert started.wait(timeout=10) - finally: - _finish_fetch(release) - litellm_utils.set_live_deployment_replay(original_replay) - - assert litellm.model_cost is not bundled - assert "claude-remote-only-test" in litellm.model_cost - assert litellm.model_cost["my-late-edit"] is late_entry - assert late_removal not in litellm.model_cost - assert litellm.model_cost["gpt-4o-mini"]["input_cost_per_token"] == 42.0 - assert litellm.model_cost["gpt-4o"]["input_cost_per_token"] == 7.0 - - def _remote_map_repricing_gpt_4o() -> tuple[bytes, float]: remote_map = _load_root_cost_map() remote_price = remote_map["gpt-4o"]["input_cost_per_token"] * 10 @@ -1026,6 +989,138 @@ def _remote_map_repricing_gpt_4o() -> tuple[bytes, float]: return json.dumps(remote_map).encode(), remote_price +def _remote_map_repricing_gpt_4o_and_adding(model: str) -> tuple[bytes, float]: + body, remote_price = _remote_map_repricing_gpt_4o() + remote_map = json.loads(body) + remote_map[model] = {"litellm_provider": "anthropic", "mode": "chat"} + return json.dumps(remote_map).encode(), remote_price + + +def test_references_to_the_map_taken_at_import_see_the_remote_map_once_it_lands( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + body, remote_price = _remote_map_repricing_gpt_4o_and_adding("claude-remote-only-test") + client, started, release = _held_client(httpx.Response(200, content=body)) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + held_map = litellm.model_cost + held_entry = litellm.model_cost["gpt-4o"] + assert started.wait(timeout=10) + finally: + _finish_fetch(release) + + held_map["my-direct-edit"] = {"litellm_provider": "openai", "mode": "chat"} + assert held_map["gpt-4o"]["input_cost_per_token"] == remote_price + assert held_entry["input_cost_per_token"] == remote_price + assert "claude-remote-only-test" in held_map + assert "my-direct-edit" in litellm.model_cost + + +def test_a_price_pinned_to_its_bundled_value_before_the_remote_map_lands_keeps_it( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + body, remote_price = _remote_map_repricing_gpt_4o() + client, started, release = _held_client(httpx.Response(200, content=body)) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + bundled_price = litellm.model_cost["gpt-4o"]["input_cost_per_token"] + litellm.model_cost["gpt-4o"]["input_cost_per_token"] = float(repr(bundled_price)) + finally: + _finish_fetch(release) + + assert remote_price != bundled_price + assert litellm.model_cost["gpt-4o"]["input_cost_per_token"] == bundled_price + + +def test_a_model_cost_the_caller_replaces_while_the_remote_map_is_merged_stays_theirs( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + caller_map = {"my-model": {"litellm_provider": "openai", "mode": "chat"}} + local_loads = [] + + def local_replacing_the_map_once_the_merge_starts(): + local_loads.append(None) + if len(local_loads) == 2: + monkeypatch.setattr(litellm, "model_cost", caller_map) + return GetModelCostMap.load_local_model_cost_map_with_revision() + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) + try: + monkeypatch.setattr( + litellm, + "model_cost", + load_model_cost_map_at_import(url=_URL, client=client, local=local_replacing_the_map_once_the_merge_starts), + ) + assert started.wait(timeout=10) + finally: + _finish_fetch(release) + + assert len(local_loads) == 2 + assert litellm.model_cost is caller_map + assert "claude-remote-only-test" not in litellm.model_cost + + +class _BlockingSleeps: + def __init__(self): + self.waits = [] + self.called = threading.Semaphore(0) + self.release = threading.Event() + + def __call__(self, seconds: float) -> None: + self.waits.append(seconds) + self.called.release() + self.release.wait(timeout=10) + + +def test_a_fetch_resumed_after_fork_during_a_retry_wait_waits_out_the_rest_of_it( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils import get_model_cost_map as module + + served = [] + + def handler(request): + served.append(request) + if len(served) == 1: + return httpx.Response(429, headers={"Retry-After": "30"}) + return httpx.Response(200, content=_remote_map_with("claude-remote-only-test")) + + sleeps = _BlockingSleeps() + try: + monkeypatch.setattr( + litellm, + "model_cost", + module.load_model_cost_map_at_import( + url=_URL, + client=httpx.Client(transport=httpx.MockTransport(handler)), + sleep=sleeps, + rng=random.Random(0), + ), + ) + assert sleeps.called.acquire(timeout=10) + module._resume_import_fetch_after_fork() + assert sleeps.called.acquire(timeout=10) + requests_during_the_wait = len(served) + finally: + _finish_fetch(sleeps.release) + + parent_wait, resumed_wait = sleeps.waits[:2] + assert requests_during_the_wait == 1 + assert 0 < resumed_wait <= parent_wait + assert "claude-remote-only-test" in litellm.model_cost + + def test_a_router_built_before_the_remote_map_lands_prices_its_models_from_the_remote_map( isolated_litellm_cost_state, monkeypatch ): @@ -1300,6 +1395,10 @@ def test_a_lookup_that_missed_before_the_fetch_settled_retries_only_once_the_rem ), id="completion_cost", ), + pytest.param( + lambda: 1000 * litellm.get_model_info("gpt-4o")["input_cost_per_token"], + id="get_model_info", + ), ], ) def test_a_cost_computed_while_the_remote_map_is_in_flight_uses_its_price( @@ -1321,6 +1420,25 @@ def test_a_cost_computed_while_the_remote_map_is_in_flight_uses_its_price( assert cost == pytest.approx(1000 * remote_price) +def test_a_cost_computed_while_the_first_attempt_outlasts_the_fetch_timeout_still_uses_the_remote_price( + isolated_litellm_cost_state, monkeypatch, release_when_a_lookup_waits +): + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + body, remote_price = _remote_map_repricing_gpt_4o() + client, started, release = _held_client(httpx.Response(200, content=body)) + waits = release_when_a_lookup_waits(release) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client, timeout=0)) + assert started.wait(timeout=10) + cost = litellm.cost_per_token(model="gpt-4o", prompt_tokens=1000, completion_tokens=0)[0] + finally: + _finish_fetch(release) + + assert waits.waited.is_set() + assert cost == pytest.approx(1000 * remote_price) + + def test_an_image_priced_while_the_remote_map_is_in_flight_finds_a_model_only_the_remote_map_has( isolated_litellm_cost_state, monkeypatch, release_when_a_lookup_waits ): diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index a8e90fded83..8d537c974eb 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -2471,6 +2471,24 @@ def test_registered_chat_model_lists_under_its_provider_wildcard( litellm.add_known_models(model_cost_map={}) +def test_a_registered_azure_o_series_deployment_name_keeps_reasoning_effort_supported(): + model_name: Final = "my-o3-deployment-regression" + try: + litellm.register_model( + {model_name: {"litellm_provider": "azure", "mode": "chat"}}, persist_across_reloads=False + ) + + supported: Final = litellm.get_supported_openai_params(model=model_name, custom_llm_provider="azure") + + assert supported is not None + assert "reasoning_effort" in supported + finally: + litellm.model_cost.pop(model_name, None) + litellm.azure_models.discard(model_name) + litellm.model_list_set.discard(model_name) + litellm.add_known_models(model_cost_map={}) + + def test_priced_router_deployment_id_stays_out_of_wildcard_model_listing(): from litellm.proxy.auth.model_checks import get_known_models_from_wildcard From 94823d5503c06b5f2d838a905fac4939e8fd43f4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 21:59:41 -0700 Subject: [PATCH 09/18] fix(cost-map): replace changed entries instead of resizing them when the remote map lands A thread reading a model's fields while the remote map was merged could hit a dictionary size change. Changed entries are now new objects set at every name that shared them, and an alias the remote map no longer lists stops resolving --- .../litellm_core_utils/get_model_cost_map.py | 35 +++++++------ .../test_get_model_cost_map.py | 50 ++++++++++++++++++- 2 files changed, 68 insertions(+), 17 deletions(-) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 707fbfcad2f..a7b9b2ee9bd 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -713,10 +713,8 @@ def _model_cost_replaced_since_import(fetch: _ImportFetch | None) -> bool: return fetch is not None and not _litellm_model_cost_is(fetch.bundled) -def _is_cost_entry( - value: object, -) -> TypeIs[dict[str, object]]: # guard-ok: cost map keys are str # mutable-ok: entries are merged in place - return isinstance(value, dict) +def _is_cost_entry(value: object) -> TypeIs[Mapping[str, object]]: # guard-ok: cost map keys are str + return isinstance(value, Mapping) def _bundled_as_loaded(fetch: _ImportFetch) -> Mapping[str, Mapping[str, object]]: @@ -738,14 +736,16 @@ def _fields_written_since_import( return frozenset(chain(rewritten, loaded.keys() - current.keys())) -def _merge_remote_fields( - current: dict[str, object], # mutable-ok: a live litellm.model_cost entry, so references to it see the remote map - remote: Mapping[str, object], - written: frozenset[str], -) -> None: - current.update({name: value for name, value in remote.items() if name not in written}) - for name in current.keys() - remote.keys() - written: - current.pop(name, None) +def _merged_entry( + current: Mapping[str, object], remote: Mapping[str, object], written: frozenset[str] +) -> dict[str, object]: # mutable-ok: cost map entries are dicts callers may edit + fields: Final = { + **{name: value for name, value in remote.items() if name not in written}, + **{name: value for name, value in current.items() if name in written}, + } + if isinstance(current, AliasedCostEntry): + return AliasedCostEntry(fields, current.names) + return fields def _untouched_since_import(current: object, loaded: Mapping[str, object], parsed: Mapping[str, object]) -> bool: @@ -763,16 +763,21 @@ def _merge_into_import_map(fetch: _ImportFetch, finalized: Mapping[str, object]) for key, entry in remote if key in fetch.loaded and _is_cost_entry(current := fetch.bundled.get(key)) and _is_cost_entry(entry) } - for current, key, entry in shared_once.values(): - _merge_remote_fields( + merged: Final = { + entry_id: _merged_entry( current, entry, _fields_written_since_import(current, fetch.loaded[key], parsed.get(key, _NO_FIELDS)) ) - fetch.bundled.update({key: entry for key, entry in remote if key not in fetch.loaded and key not in fetch.bundled}) + for entry_id, (current, key, entry) in shared_once.items() + } + changed: Final = {entry_id: entry for entry_id, entry in merged.items() if entry != shared_once[entry_id][0]} + replaced: Final = {key: changed[id(value)] for key, value in tuple(fetch.bundled.items()) if id(value) in changed} + added: Final = {key: entry for key, entry in remote if key not in fetch.loaded and key not in fetch.bundled} dropped: Final = tuple( key for key in fetch.loaded.keys() - finalized.keys() if _untouched_since_import(fetch.bundled.get(key), fetch.loaded[key], parsed.get(key, _NO_FIELDS)) ) + fetch.bundled.update({**replaced, **added}) for key in dropped: fetch.bundled.pop(key, None) _record_remote_adoption() diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index 7843b1a65da..97bc3921936 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -1007,18 +1007,64 @@ def test_references_to_the_map_taken_at_import_see_the_remote_map_once_it_lands( try: monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) held_map = litellm.model_cost - held_entry = litellm.model_cost["gpt-4o"] assert started.wait(timeout=10) finally: _finish_fetch(release) held_map["my-direct-edit"] = {"litellm_provider": "openai", "mode": "chat"} assert held_map["gpt-4o"]["input_cost_per_token"] == remote_price - assert held_entry["input_cost_per_token"] == remote_price assert "claude-remote-only-test" in held_map assert "my-direct-edit" in litellm.model_cost +def test_a_model_entry_being_read_field_by_field_when_the_remote_map_lands_reads_to_the_end( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + remote_map = _load_root_cost_map() + remote_map["gpt-4o"] = {**remote_map["gpt-4o"], "supports_remote_only_test_field": True} + client, started, release = _held_client(httpx.Response(200, content=json.dumps(remote_map).encode())) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + bundled_fields = len(litellm.model_cost["gpt-4o"]) + fields = iter(litellm.model_cost["gpt-4o"].items()) + first = next(fields) + finally: + _finish_fetch(release) + + assert len([first, *fields]) == bundled_fields + assert litellm.model_cost["gpt-4o"]["supports_remote_only_test_field"] is True + + +def test_an_alias_the_remote_map_no_longer_lists_stops_resolving(isolated_litellm_cost_state, monkeypatch): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + remote_map = json.loads(_cost_map_with_aliased_model(5.0)) + del remote_map["my-aliased-model"]["aliases"] + client, started, release = _held_client(httpx.Response(200, content=json.dumps(remote_map).encode())) + try: + monkeypatch.setattr( + litellm, + "model_cost", + load_model_cost_map_at_import( + url=_URL, + client=client, + local=lambda: ModelCostMapReloaded(model_cost_map=json.loads(_cost_map_with_aliased_model(1.0))), + ), + ) + assert started.wait(timeout=10) + assert "my-aliased-model-alias" in litellm.model_cost + finally: + _finish_fetch(release) + + assert litellm.model_cost["my-aliased-model"]["input_cost_per_token"] == 5.0 + assert "my-aliased-model-alias" not in litellm.model_cost + + def test_a_price_pinned_to_its_bundled_value_before_the_remote_map_lands_keeps_it( isolated_litellm_cost_state, monkeypatch ): From ce36b2c1d97fa6071908d6df2f5fe4bcf4138ceb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 18 Oct 2026 00:31:39 -0700 Subject: [PATCH 10/18] fix(cost-map): point an alias the remote map adds at its model's merged entry An alias that only the remote map lists now shares the entry its model kept, so prices set on the model before the remote map landed reach the alias too, and later register_model calls on the model update it --- .../litellm_core_utils/get_model_cost_map.py | 23 ++++++++++++++-- .../test_get_model_cost_map.py | 26 +++++++++++++++++++ 2 files changed, 47 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index a7b9b2ee9bd..c3077ece3d9 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -743,11 +743,21 @@ def _merged_entry( **{name: value for name, value in remote.items() if name not in written}, **{name: value for name, value in current.items() if name in written}, } + if isinstance(remote, AliasedCostEntry): + return AliasedCostEntry(fields, remote.names) if isinstance(current, AliasedCostEntry): return AliasedCostEntry(fields, current.names) return fields +def _alias_names(entry: Mapping[str, object]) -> tuple[str, ...] | None: + return entry.names if isinstance(entry, AliasedCostEntry) else None + + +def _same_entry(merged: Mapping[str, object], current: Mapping[str, object]) -> bool: + return merged == current and _alias_names(merged) == _alias_names(current) + + def _untouched_since_import(current: object, loaded: Mapping[str, object], parsed: Mapping[str, object]) -> bool: return _is_cost_entry(current) and not _fields_written_since_import(current, loaded, parsed) @@ -769,9 +779,18 @@ def _merge_into_import_map(fetch: _ImportFetch, finalized: Mapping[str, object]) ) for entry_id, (current, key, entry) in shared_once.items() } - changed: Final = {entry_id: entry for entry_id, entry in merged.items() if entry != shared_once[entry_id][0]} + changed: Final = { + entry_id: entry for entry_id, entry in merged.items() if not _same_entry(entry, shared_once[entry_id][0]) + } replaced: Final = {key: changed[id(value)] for key, value in tuple(fetch.bundled.items()) if id(value) in changed} - added: Final = {key: entry for key, entry in remote if key not in fetch.loaded and key not in fetch.bundled} + kept_for_remote: Final = { + id(entry): changed.get(entry_id, current) for entry_id, (current, _, entry) in shared_once.items() + } + added: Final = { + key: kept_for_remote.get(id(entry), entry) + for key, entry in remote + if key not in fetch.loaded and key not in fetch.bundled + } dropped: Final = tuple( key for key in fetch.loaded.keys() - finalized.keys() diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index 97bc3921936..67655243d2e 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -1324,6 +1324,32 @@ def test_an_edit_to_an_aliased_model_made_before_the_remote_map_lands_stays_shar assert entry["output_cost_per_token"] == 9.0 +def test_an_alias_the_remote_map_adds_shares_the_prices_its_model_kept(isolated_litellm_cost_state, monkeypatch): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + bundled_map = json.loads(_cost_map_with_aliased_model(1.0)) + del bundled_map["my-aliased-model"]["aliases"] + client, started, release = _held_client(httpx.Response(200, content=_cost_map_with_aliased_model(5.0))) + try: + monkeypatch.setattr( + litellm, + "model_cost", + load_model_cost_map_at_import( + url=_URL, client=client, local=lambda: ModelCostMapReloaded(model_cost_map=bundled_map) + ), + ) + assert started.wait(timeout=10) + litellm.register_model({"my-aliased-model": {"output_cost_per_token": 9.0}}, persist_across_reloads=False) + finally: + _finish_fetch(release) + + assert litellm.model_cost["my-aliased-model-alias"]["output_cost_per_token"] == 9.0 + assert litellm.model_cost["my-aliased-model-alias"]["input_cost_per_token"] == 5.0 + litellm.register_model({"my-aliased-model": {"input_cost_per_token": 7.0}}, persist_across_reloads=False) + assert litellm.model_cost["my-aliased-model-alias"]["input_cost_per_token"] == 7.0 + + class _ReleaseWhenALookupWaits(logging.Handler): """Lets the held fetch answer only once a lookup is blocked waiting on it, so the lookup provably missed first.""" From 041835a0991b65daa39a78bffa87770b2eee6712 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 18 Oct 2026 01:11:51 -0700 Subject: [PATCH 11/18] fix(cost-map): resume the fetch after a C-level fork, keep trimmed entries, update same-shape entries in place, report retry reasons, share snapshot values A process forked without the fork hooks, as uWSGI does, resumes the pending import fetch on its own instead of waiting forever. An entry that lost fields since import keeps its own field set, entries whose fields and aliases stay the same are updated in place so held references see remote prices, the fallback reason names the failed attempt while a retry waits and is only reported once import finishes, and the loaded catalog snapshot reuses live entries equal to the remote ones --- .../litellm_core_utils/get_model_cost_map.py | 105 ++++++++++++----- .../test_get_model_cost_map.py | 108 ++++++++++++++++++ 2 files changed, 187 insertions(+), 26 deletions(-) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index c3077ece3d9..c37ec742d30 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -18,7 +18,7 @@ import sys import threading import time import urllib.parse -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping, MutableMapping from dataclasses import dataclass, field, replace from datetime import datetime, timezone from importlib.resources import files @@ -431,7 +431,7 @@ def _fetch_remote_model_cost_map_with_retry_sync( url: str, timeout: int, attempts: range, - sleep: Callable[[float], None], + wait_to_retry: Callable[[_FetchAttemptRetryable, float], None], rng: random.Random, client: _SyncGetClient, ) -> ModelCostMapReloadResult: @@ -443,7 +443,7 @@ def _fetch_remote_model_cost_map_with_retry_sync( 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) + wait_to_retry(outcome, wait_seconds) return ModelCostMapReloadUnavailable(reason="model cost map fetch failed") @@ -649,12 +649,18 @@ def _finalize_model_cost_map(model_cost: dict) -> dict: return _expand_model_aliases(model_cost) -def _finalize_loaded_model_cost_map(loaded: ModelCostMapReloaded) -> ModelCostMapReloaded: +def _record_and_finalize(loaded: ModelCostMapReloaded) -> ModelCostMapReloaded: _cost_map_source_info.source_revision = loaded.revision _cost_map_source_info.etag = loaded.etag - finalized: Final = _finalize_model_cost_map(loaded.model_cost_map) - GetModelCostMap._snapshot_loaded_catalog(finalized) # pyright: ignore[reportPrivateUsage] # same module - return replace(loaded, model_cost_map=finalized) + return replace(loaded, model_cost_map=_finalize_model_cost_map(loaded.model_cost_map)) + + +def _finalize_loaded_model_cost_map(loaded: ModelCostMapReloaded) -> ModelCostMapReloaded: + finalized: Final = _record_and_finalize(loaded) + GetModelCostMap._snapshot_loaded_catalog( # pyright: ignore[reportPrivateUsage] # same module + finalized.model_cost_map + ) + return finalized def adopt_model_cost_map( @@ -696,6 +702,7 @@ class _ScheduledAttempt: class _PendingImportFetch: fetch: _ImportFetch | None = None thread: threading.Thread | None = None + pid: int = 0 next_attempt: _ScheduledAttempt | None = None adoptions: int = 0 @@ -713,6 +720,12 @@ def _model_cost_replaced_since_import(fetch: _ImportFetch | None) -> bool: return fetch is not None and not _litellm_model_cost_is(fetch.bundled) +def _report_import_fetch_failure(fetch: _ImportFetch | None, reason: str) -> None: + _litellm_import_complete.wait() + if not _model_cost_replaced_since_import(fetch): + _cost_map_source_info.fallback_reason = reason + + def _is_cost_entry(value: object) -> TypeIs[Mapping[str, object]]: # guard-ok: cost map keys are str return isinstance(value, Mapping) @@ -739,8 +752,9 @@ def _fields_written_since_import( def _merged_entry( current: Mapping[str, object], remote: Mapping[str, object], written: frozenset[str] ) -> dict[str, object]: # mutable-ok: cost map entries are dicts callers may edit + trimmed: Final = any(name not in current for name in written) fields: Final = { - **{name: value for name, value in remote.items() if name not in written}, + **{name: value for name, value in remote.items() if name not in written and (name in current or not trimmed)}, **{name: value for name, value in current.items() if name in written}, } if isinstance(remote, AliasedCostEntry): @@ -754,10 +768,22 @@ def _alias_names(entry: Mapping[str, object]) -> tuple[str, ...] | None: return entry.names if isinstance(entry, AliasedCostEntry) else None +def _updatable_in_place(merged: Mapping[str, object], current: Mapping[str, object]) -> bool: + return ( + isinstance(current, MutableMapping) + and merged.keys() == current.keys() + and _alias_names(merged) == _alias_names(current) + ) + + def _same_entry(merged: Mapping[str, object], current: Mapping[str, object]) -> bool: return merged == current and _alias_names(merged) == _alias_names(current) +def _live_entry_if_equal(live: object, remote: object) -> object: + return live if live == remote else remote + + def _untouched_since_import(current: object, loaded: Mapping[str, object], parsed: Mapping[str, object]) -> bool: return _is_cost_entry(current) and not _fields_written_since_import(current, loaded, parsed) @@ -782,9 +808,19 @@ def _merge_into_import_map(fetch: _ImportFetch, finalized: Mapping[str, object]) changed: Final = { entry_id: entry for entry_id, entry in merged.items() if not _same_entry(entry, shared_once[entry_id][0]) } - replaced: Final = {key: changed[id(value)] for key, value in tuple(fetch.bundled.items()) if id(value) in changed} + reshaped: Final = { + entry_id: entry + for entry_id, entry in changed.items() + if not _updatable_in_place(entry, shared_once[entry_id][0]) + } + updated_in_place: Final = tuple( + (live, entry) + for entry_id, entry in changed.items() + if entry_id not in reshaped and isinstance(live := shared_once[entry_id][0], MutableMapping) + ) + replaced: Final = {key: reshaped[id(value)] for key, value in tuple(fetch.bundled.items()) if id(value) in reshaped} kept_for_remote: Final = { - id(entry): changed.get(entry_id, current) for entry_id, (current, _, entry) in shared_once.items() + id(entry): reshaped.get(entry_id, current) for entry_id, (current, _, entry) in shared_once.items() } added: Final = { key: kept_for_remote.get(id(entry), entry) @@ -796,9 +832,14 @@ def _merge_into_import_map(fetch: _ImportFetch, finalized: Mapping[str, object]) for key in fetch.loaded.keys() - finalized.keys() if _untouched_since_import(fetch.bundled.get(key), fetch.loaded[key], parsed.get(key, _NO_FIELDS)) ) + for live_entry, entry in updated_in_place: + live_entry.update(entry) fetch.bundled.update({**replaced, **added}) for key in dropped: fetch.bundled.pop(key, None) + GetModelCostMap._snapshot_loaded_catalog( # pyright: ignore[reportPrivateUsage] # same module + {key: _live_entry_if_equal(fetch.bundled.get(key), entry) for key, entry in remote} + ) _record_remote_adoption() utils._invalidate_model_cost_lowercase_map() # pyright: ignore[reportPrivateUsage] # required cache invalidation litellm.add_known_models_to_model_list() @@ -810,7 +851,7 @@ def _adopt_remote_model_cost_map( url: str, timeout: int, attempts: range, - sleep: Callable[[float], None], + wait_to_retry: Callable[[_FetchAttemptRetryable, float], None], rng: random.Random, client: _SyncGetClient, import_fetch: _ImportFetch | None, @@ -819,7 +860,7 @@ def _adopt_remote_model_cost_map( url=url, timeout=timeout, attempts=attempts, - sleep=sleep, + wait_to_retry=wait_to_retry, rng=rng, client=client, ) @@ -829,8 +870,7 @@ def _adopt_remote_model_cost_map( url, result.reason, ) - if not _model_cost_replaced_since_import(import_fetch): - _cost_map_source_info.fallback_reason = f"Remote fetch failed: {result.reason}" + _report_import_fetch_failure(import_fetch, f"Remote fetch failed: {result.reason}") return _litellm_import_complete.wait() if not GetModelCostMap.validate_model_cost_map( @@ -841,18 +881,17 @@ def _adopt_remote_model_cost_map( "LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s", url, ) - if not _model_cost_replaced_since_import(import_fetch): - _cost_map_source_info.fallback_reason = "Remote data failed integrity validation" + _report_import_fetch_failure(import_fetch, "Remote data failed integrity validation") return if _model_cost_replaced_since_import(import_fetch): verbose_logger.debug("LiteLLM: litellm.model_cost was replaced after import; not adopting the remote map") return - finalized: Final = _finalize_loaded_model_cost_map(result).model_cost_map if import_fetch is None: + finalized: Final = _finalize_loaded_model_cost_map(result).model_cost_map _record_remote_adoption() adopt_model_cost_map(finalized) return - _merge_into_import_map(import_fetch, finalized) + _merge_into_import_map(import_fetch, _record_and_finalize(result).model_cost_map) def _record_remote_adoption() -> None: @@ -879,7 +918,7 @@ def _retry_remote_fetch_in_background( url=url, timeout=timeout, attempts=range(2, max_attempts + 1), - sleep=sleep, + wait_to_retry=lambda _, seconds: sleep(seconds), rng=rng, client=client, import_fetch=None, @@ -889,12 +928,13 @@ def _retry_remote_fetch_in_background( def _run_import_fetch(fetch: _ImportFetch) -> None: - def settle_then_sleep(seconds: float) -> None: + def settle_then_sleep(failed: _FetchAttemptRetryable, seconds: float) -> None: scheduled: Final = _PendingImportFetch.next_attempt _PendingImportFetch.next_attempt = _ScheduledAttempt( attempt=(scheduled.attempt if scheduled is not None else fetch.first_attempt) + 1, at=time.monotonic() + seconds, ) + _report_import_fetch_failure(fetch, f"Remote fetch failed: {failed.reason}") fetch.first_attempt_settled.set() fetch.sleep(seconds) @@ -905,7 +945,7 @@ def _run_import_fetch(fetch: _ImportFetch) -> None: url=fetch.url, timeout=fetch.timeout, attempts=range(fetch.first_attempt, fetch.max_attempts + 1), - sleep=settle_then_sleep, + wait_to_retry=settle_then_sleep, rng=fetch.rng, client=fetch.client, import_fetch=fetch, @@ -929,6 +969,7 @@ def _start_import_fetch(fetch: _ImportFetch) -> None: ) _PendingImportFetch.fetch = fetch _PendingImportFetch.thread = thread + _PendingImportFetch.pid = os.getpid() thread.start() @@ -973,7 +1014,18 @@ def wait_for_import_model_cost_map(since: int) -> bool: return _PendingImportFetch.adoptions != since +_resume_after_unhooked_fork: Final = threading.Lock() + + +def _resume_import_fetch_after_unhooked_fork() -> None: + with _resume_after_unhooked_fork: + if _PendingImportFetch.fetch is not None and _PendingImportFetch.pid != os.getpid(): + _resume_import_fetch_after_fork() + + def wait_for_pending_import_model_cost_map() -> None: + if _PendingImportFetch.fetch is not None and _PendingImportFetch.pid != os.getpid(): + _resume_import_fetch_after_unhooked_fork() pending: Final = _PendingImportFetch.fetch if ( pending is None @@ -1013,11 +1065,12 @@ def load_model_cost_map_at_import( Return the bundled map right away and fetch the remote one in a background thread, so ``import litellm`` never waits on the network. Once the remote map passes the integrity check it is merged into the returned dict in place, so references taken - at import see it too, and every field written since import keeps its value. A caller - that replaced ``litellm.model_cost`` outright keeps its own map. Costs and model info - looked up while the first attempt is in flight wait for it, see - ``wait_for_pending_import_model_cost_map``. A process forked during a retry wait - resumes the retry schedule on its own. + at import see it too, and every field written since import keeps its value. An entry + that lost fields since import keeps its own set of fields. A caller that replaced + ``litellm.model_cost`` outright keeps its own map. Costs and model info looked up + while the first attempt is in flight wait for it, see + ``wait_for_pending_import_model_cost_map``. A process forked while the fetch is + pending, by ``os.fork`` or by a server forking at the C level, resumes it on its own. """ if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true" or _is_cli_process(): return get_model_cost_map(url=url, timeout=timeout) diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index 67655243d2e..a4030d9b274 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -5,8 +5,10 @@ count actual model entries, not reserved meta keys) and the extraction of the """ import asyncio, importlib, importlib.resources as importlib_get_model, json, litellm +import ctypes import logging import os +import signal import sys import threading from typing import Final @@ -946,6 +948,27 @@ def test_import_load_reports_why_the_remote_fetch_failed(isolated_litellm_cost_s assert source["fallback_reason"].startswith("Remote fetch failed: HTTP 404") +def test_import_load_reports_why_the_last_attempt_failed_while_a_retry_waits(isolated_litellm_cost_state, monkeypatch): + import litellm + from litellm.litellm_core_utils import get_model_cost_map as module + + monkeypatch.setattr(module, "_litellm_import_complete", threading.Event()) + sleeps = _BlockingSleeps() + client = httpx.Client(transport=httpx.MockTransport(lambda request: httpx.Response(429))) + try: + bundled = module.load_model_cost_map_at_import(url=_URL, client=client, sleep=sleeps, rng=random.Random(0)) + retry_wait_began_during_import = sleeps.called.acquire(timeout=1) + monkeypatch.setattr(litellm, "model_cost", bundled) + module.mark_litellm_import_complete() + assert retry_wait_began_during_import or sleeps.called.acquire(timeout=10) + source = get_model_cost_map_source_info() + finally: + _finish_fetch(sleeps.release) + + assert source["source"] == "local" + assert source["fallback_reason"].startswith("Remote fetch failed: HTTP 429") + + def test_a_fetch_resumed_after_fork_adopts_the_remote_map_while_the_inherited_one_hangs( isolated_litellm_cost_state, monkeypatch ): @@ -1017,6 +1040,26 @@ def test_references_to_the_map_taken_at_import_see_the_remote_map_once_it_lands( assert "my-direct-edit" in litellm.model_cost +def test_a_model_entry_held_from_import_sees_the_remote_price_and_edits_through_it_reach_the_map( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + body, remote_price = _remote_map_repricing_gpt_4o() + client, started, release = _held_client(httpx.Response(200, content=body)) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + held_entry = litellm.model_cost["gpt-4o"] + assert started.wait(timeout=10) + finally: + _finish_fetch(release) + + held_entry["output_cost_per_token"] = 9.0 + assert held_entry["input_cost_per_token"] == remote_price + assert litellm.model_cost["gpt-4o"]["output_cost_per_token"] == 9.0 + + def test_a_model_entry_being_read_field_by_field_when_the_remote_map_lands_reads_to_the_end( isolated_litellm_cost_state, monkeypatch ): @@ -1221,6 +1264,33 @@ def test_a_partial_edit_made_before_the_remote_map_lands_keeps_only_the_fields_i assert "output_cost_per_token" not in litellm.model_cost["gpt-4o"] +def test_a_model_entry_replaced_by_one_with_fewer_fields_before_the_remote_map_lands_stays_as_written( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + remote_map = _load_root_cost_map() + remote_map["gpt-4o"] = {**remote_map["gpt-4o"], "output_cost_per_reasoning_token": 1.0} + client, started, release = _held_client(httpx.Response(200, content=json.dumps(remote_map).encode())) + replacement = { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + } + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + litellm.model_cost["gpt-4o"] = dict( + replacement + ) # test-quality-ok: the replacement is what is under test, and the fixture restores model_cost + finally: + _finish_fetch(release) + + assert litellm.model_cost["gpt-4o"] == replacement + + def test_a_model_deleted_before_the_remote_map_lands_stays_deleted(isolated_litellm_cost_state, monkeypatch): import litellm from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import @@ -1581,6 +1651,44 @@ def test_a_process_forked_while_the_remote_map_is_in_flight_inherits_it_once_ado assert os.waitstatus_to_exitcode(wait_status) == 0 +@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs fork") +def test_a_process_forked_without_the_fork_hooks_while_the_remote_map_is_in_flight_prices_from_it( + isolated_litellm_cost_state, monkeypatch +): + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + body, remote_price = _remote_map_repricing_gpt_4o() + first_request_started = threading.Event() + release = threading.Event() + served = [] + + def handler(request): + served.append(request) + if len(served) == 1: + first_request_started.set() + release.wait(timeout=10) + return httpx.Response(200, content=body) + + client = httpx.Client(transport=httpx.MockTransport(handler)) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert first_request_started.wait(timeout=10) + child = ctypes.PyDLL(None).fork() + if child == 0: + signal.signal(signal.SIGALRM, signal.SIG_DFL) + signal.alarm(10) + try: + cost = litellm.cost_per_token(model="gpt-4o", prompt_tokens=1000, completion_tokens=0)[0] + os._exit(0 if cost == pytest.approx(1000 * remote_price) else 1) + finally: + os._exit(2) + _, wait_status = os.waitpid(child, 0) + finally: + _finish_fetch(release) + + assert os.waitstatus_to_exitcode(wait_status) == 0 + + def test_import_load_with_the_local_env_override_starts_no_fetch(monkeypatch): from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import From d3589d3f8d8b777ec0e8cf1cccc0bc8e0858c088 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 18 Oct 2026 02:35:47 -0700 Subject: [PATCH 12/18] fix(cost-map): fetch the remote map before returning when running under uWSGI uWSGI forks workers at the C level after the app loads. A worker forked while the fetch thread holds the GIL deadlocks before it serves a request, so under uWSGI import fetches the remote map before returning, as it did before this change --- litellm/litellm_core_utils/get_model_cost_map.py | 13 +++++++++++-- .../litellm_core_utils/test_get_model_cost_map.py | 14 ++++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index c37ec742d30..beff2f4a71b 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -57,6 +57,10 @@ def _is_cli_process() -> bool: return Path(sys.argv[0]).stem in _CLI_ENTRYPOINT_NAMES +def _is_uwsgi_process() -> bool: + return "uwsgi" in sys.modules + + # Reserved top-level keys that are not model entries. They must be excluded # from the model-count integrity check so a real upstream shrink can't be masked. RESERVED_TOP_LEVEL_KEYS: Final = frozenset({"sample_spec", FALLBACK_GENERALIZATIONS_KEY}) @@ -1060,6 +1064,7 @@ def load_model_cost_map_at_import( rng: random.Random | None = None, client: "_SyncGetClient | None" = None, local: Callable[[], ModelCostMapReloaded] = GetModelCostMap.load_local_model_cost_map_with_revision, + fetch_before_returning: Callable[[], bool] = _is_uwsgi_process, ) -> dict: # mutable-ok: becomes litellm.model_cost, a plain mutable dict by contract """ Return the bundled map right away and fetch the remote one in a background thread, @@ -1071,9 +1076,13 @@ def load_model_cost_map_at_import( while the first attempt is in flight wait for it, see ``wait_for_pending_import_model_cost_map``. A process forked while the fetch is pending, by ``os.fork`` or by a server forking at the C level, resumes it on its own. + Under uWSGI, whose workers forked while a thread holds the GIL deadlock, the remote + map is fetched before returning instead. """ - if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true" or _is_cli_process(): - return get_model_cost_map(url=url, timeout=timeout) + if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true" or _is_cli_process() or fetch_before_returning(): + return get_model_cost_map( + url=url, timeout=timeout, max_attempts=max_attempts, sleep=sleep, rng=rng, client=client + ) _cost_map_source_info.loaded_at = datetime.now(timezone.utc) _cost_map_source_info.source = "local" _cost_map_source_info.url = url diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index a4030d9b274..e4815de0e97 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -948,6 +948,20 @@ def test_import_load_reports_why_the_remote_fetch_failed(isolated_litellm_cost_s assert source["fallback_reason"].startswith("Remote fetch failed: HTTP 404") +def test_import_load_under_uwsgi_returns_the_remote_map_once_fetched(isolated_litellm_cost_state): + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, _, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) + threading.Timer(0.2, release.set).start() + try: + loaded = load_model_cost_map_at_import(url=_URL, client=client, fetch_before_returning=lambda: True) + finally: + _finish_fetch(release) + + assert "claude-remote-only-test" in loaded + assert get_model_cost_map_source_info()["source"] == "remote" + + def test_import_load_reports_why_the_last_attempt_failed_while_a_retry_waits(isolated_litellm_cost_state, monkeypatch): import litellm from litellm.litellm_core_utils import get_model_cost_map as module From c0631a87629c3fd0d2c866fdfe6941357d02f757 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 18 Oct 2026 06:00:18 -0700 Subject: [PATCH 13/18] fix(cost-map): keep loops over the map, the event loop, and forked retries safe while the remote map lands Loops over litellm.model_cost walk a snapshot, so the background merge no longer raises 'dictionary changed size during iteration'. Async calls wait for the in-flight first attempt in a worker thread instead of blocking the event loop. Workers forked during a retry wait add their own jitter so they do not retry the map host in step. --- .../litellm_core_utils/get_model_cost_map.py | 46 +++++++++++-- litellm/utils.py | 7 +- .../sdk/test_cost_map_background_fetch_sdk.py | 41 ++++++++++- .../test_get_model_cost_map.py | 68 ++++++++++++++++++- 4 files changed, 155 insertions(+), 7 deletions(-) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index beff2f4a71b..e3145b1078e 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -18,14 +18,19 @@ import sys import threading import time import urllib.parse -from collections.abc import Awaitable, Callable, Mapping, MutableMapping +from collections.abc import Awaitable, Callable, Iterator, Mapping, MutableMapping from dataclasses import dataclass, field, replace from datetime import datetime, timezone from importlib.resources import files from itertools import chain from pathlib import Path from types import MappingProxyType -from typing import Final, Protocol, cast # noqa: TID251 # adapts boto3 and botocore's dynamic client responses +from typing import ( + TYPE_CHECKING, + Final, + Protocol, + cast, # noqa: TID251 # adapts boto3 and botocore's dynamic client responses +) import httpx from pydantic import ConfigDict, TypeAdapter @@ -45,6 +50,9 @@ from litellm.litellm_core_utils.fallback_generalizations import ( set_fallback_generalizations, ) +if TYPE_CHECKING: + from _collections_abc import dict_items, dict_keys, dict_values + FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations" _CATALOG_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) _BUNDLED_CATALOG_ADAPTER: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) @@ -576,6 +584,22 @@ class AliasedCostEntry(dict[str, object]): self.names: tuple[str, ...] = names +class _SnapshotIteratingDict(dict[str, object]): + __slots__ = () + + def __iter__(self) -> Iterator[str]: + return iter(tuple(super().keys())) + + def keys(self) -> "dict_keys[str, object]": + return dict(super().items()).keys() + + def values(self) -> "dict_values[str, object]": + return dict(super().items()).values() + + def items(self) -> "dict_items[str, object]": + return dict(super().items()).items() + + def _alias_groups(alias_owners: Mapping[str, str]) -> Mapping[str, tuple[str, ...]]: return MappingProxyType( { @@ -991,11 +1015,13 @@ def _resume_import_fetch_after_fork() -> None: if scheduled is None: _start_import_fetch(replace(pending, first_attempt_settled=threading.Event())) return + rng: Final = random.Random() _start_import_fetch( replace( pending, + rng=rng, first_attempt=scheduled.attempt, - first_attempt_at=scheduled.at, + first_attempt_at=scheduled.at + rng.uniform(0.0, 1.0), first_attempt_settled=_settled_event(), ) ) @@ -1042,6 +1068,18 @@ def wait_for_pending_import_model_cost_map() -> None: pending.first_attempt_settled.wait() +def _import_fetch_may_block() -> bool: + pending: Final = _PendingImportFetch.fetch + return pending is not None and ( + _PendingImportFetch.pid != os.getpid() or not pending.first_attempt_settled.is_set() + ) + + +async def await_pending_import_model_cost_map() -> None: + if _import_fetch_may_block(): + await asyncio.to_thread(wait_for_pending_import_model_cost_map) + + def _settle_import_fetch_before_fork() -> None: pending: Final = _PendingImportFetch.fetch thread: Final = _PendingImportFetch.thread @@ -1088,7 +1126,7 @@ def load_model_cost_map_at_import( _cost_map_source_info.url = url _cost_map_source_info.is_env_forced = False _cost_map_source_info.fallback_reason = _IMPORT_FETCH_IN_PROGRESS - bundled: Final = _finalize_loaded_model_cost_map(local()).model_cost_map + bundled: Final = _SnapshotIteratingDict(_finalize_loaded_model_cost_map(local()).model_cost_map) _PendingImportFetch.next_attempt = None _start_import_fetch( _ImportFetch( diff --git a/litellm/utils.py b/litellm/utils.py index f7a8a92bdb6..e87dd541160 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -89,7 +89,11 @@ from litellm.litellm_core_utils.fallback_generalizations import ( match_capability_generalizations, match_fill_missing_generalizations, ) -from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry, wait_for_pending_import_model_cost_map +from litellm.litellm_core_utils.get_model_cost_map import ( + AliasedCostEntry, + await_pending_import_model_cost_map, + wait_for_pending_import_model_cost_map, +) from litellm.litellm_core_utils.model_param_helper import peek_reasoning_summary_aliases from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.litellm_core_utils.tokenizer import strip_special_tokens @@ -2043,6 +2047,7 @@ def client(original_function): @wraps(original_function) async def wrapper_async(*args, **kwargs): + await await_pending_import_model_cost_map() print_args_passed_to_litellm(original_function, args, kwargs) start_time: Final = datetime.datetime.now() call_type: Final = _get_call_type_from_public_name(original_function.__name__) diff --git a/tests/integration/sdk/test_cost_map_background_fetch_sdk.py b/tests/integration/sdk/test_cost_map_background_fetch_sdk.py index 6483f9a3a7a..6e4fba0f7a4 100644 --- a/tests/integration/sdk/test_cost_map_background_fetch_sdk.py +++ b/tests/integration/sdk/test_cost_map_background_fetch_sdk.py @@ -4,7 +4,9 @@ import subprocess import sys import textwrap import threading +import time from collections.abc import Iterator +from queue import SimpleQueue from dataclasses import dataclass from importlib.resources import files from pathlib import Path @@ -20,7 +22,7 @@ _CHILD_SECONDS: Final = 120 _SCRIPT: Final = textwrap.dedent( """ - import os, sys, threading + import os, sys, threading, time import litellm from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map_source_info @@ -47,6 +49,18 @@ _SCRIPT: Final = textwrap.dedent( join_fetch() os._exit(0 if model in litellm.model_cost else 1) print(os.waitstatus_to_exitcode(os.waitpid(pid, 0)[1]), flush=True) + elif variant == "forks_during_retry_wait": + while get_model_cost_map_source_info()["fallback_reason"] == "Remote fetch in progress": + time.sleep(0.01) + children = [] + for _ in range(6): + pid = os.fork() + if pid == 0: + join_fetch() + os._exit(0 if model in litellm.model_cost else 1) + children.append(pid) + join_fetch() + print(*(os.waitstatus_to_exitcode(os.waitpid(pid, 0)[1]) for pid in children), flush=True) else: raise SystemExit(f"unknown variant {variant}") """ @@ -144,3 +158,28 @@ def test_a_worker_forked_while_the_remote_map_is_in_flight_fetches_it_itself( assert child.after_import == "local Remote fetch in progress False" assert child.rest == ("0",) + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") +def test_workers_forked_during_a_rate_limited_retry_wait_spread_their_retries(tmp_path: Path) -> None: + body: Final = _remote_map() + arrivals: Final[SimpleQueue[float]] = SimpleQueue() + + def respond(request: Request) -> Reply: + arrivals.put(time.monotonic()) + if arrivals.qsize() == 1: + return Reply(status=429, headers={"Retry-After": "2"}) + return Reply(body=body) + + with wire_server(respond) as wire: + child: Final = _run_child( + _CostMapHost(wire=wire, first_request_answered=threading.Event()), + "forks_during_retry_wait", + tmp_path, + answer_after_import=False, + ) + retries: Final = sorted(tuple(arrivals.get_nowait() for _ in range(arrivals.qsize()))[1:]) + + assert child.rest == ("0",) * 6 + assert len(retries) == 7 + assert retries[5] - retries[0] > 0.1 diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index e4815de0e97..541b971150e 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -11,6 +11,7 @@ import os import signal import sys import threading +import time from typing import Final import boto3 @@ -1096,6 +1097,35 @@ def test_a_model_entry_being_read_field_by_field_when_the_remote_map_lands_reads assert litellm.model_cost["gpt-4o"]["supports_remote_only_test_field"] is True +@pytest.mark.parametrize( + "start_loop", + [ + pytest.param(iter, id="models"), + pytest.param(lambda model_cost: iter(model_cost.keys()), id="keys"), + pytest.param(lambda model_cost: iter(model_cost.values()), id="values"), + pytest.param(lambda model_cost: iter(model_cost.items()), id="items"), + ], +) +def test_a_loop_over_the_map_running_when_the_remote_map_lands_reads_to_the_end( + isolated_litellm_cost_state, monkeypatch, start_loop +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + bundled_models = len(litellm.model_cost) + loop = start_loop(litellm.model_cost) + first = next(loop) + finally: + _finish_fetch(release) + + assert len([first, *loop]) == bundled_models + assert "claude-remote-only-test" in litellm.model_cost + + def test_an_alias_the_remote_map_no_longer_lists_stops_resolving(isolated_litellm_cost_state, monkeypatch): import litellm from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import @@ -1220,7 +1250,7 @@ def test_a_fetch_resumed_after_fork_during_a_retry_wait_waits_out_the_rest_of_it parent_wait, resumed_wait = sleeps.waits[:2] assert requests_during_the_wait == 1 - assert 0 < resumed_wait <= parent_wait + assert 0 < resumed_wait <= parent_wait + 1.0 assert "claude-remote-only-test" in litellm.model_cost @@ -1576,6 +1606,42 @@ def test_a_cost_computed_while_the_remote_map_is_in_flight_uses_its_price( assert cost == pytest.approx(1000 * remote_price) +async def test_an_async_call_priced_while_the_remote_map_is_in_flight_keeps_the_event_loop_running( + isolated_litellm_cost_state, monkeypatch +): + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + remote_map = _load_root_cost_map() + remote_price = remote_map["text-embedding-3-small"]["input_cost_per_token"] * 10 + remote_map["text-embedding-3-small"] = {**remote_map["text-embedding-3-small"], "input_cost_per_token": remote_price} + client, started, release = _held_client(httpx.Response(200, content=json.dumps(remote_map).encode())) + + def embed(): + return litellm.aembedding(model="text-embedding-3-small", input=["hi"], mock_response=[0.1]) + + await embed() + + async def loop_gaps_while(call): + task = asyncio.ensure_future(call) + ticks = [time.monotonic()] + while not task.done(): + await asyncio.sleep(0.01) + ticks.append(time.monotonic()) + return await task, tuple(later - earlier for earlier, later in zip(ticks, ticks[1:])) + + try: + monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + assert started.wait(timeout=10) + threading.Timer(0.5, release.set).start() + response, gaps = await loop_gaps_while(embed()) + finally: + _finish_fetch(release) + + assert max(gaps) < 0.25 + assert sum(gaps) >= 0.5 + assert response._hidden_params["response_cost"] == pytest.approx(10 * remote_price) + + def test_a_cost_computed_while_the_first_attempt_outlasts_the_fetch_timeout_still_uses_the_remote_price( isolated_litellm_cost_state, monkeypatch, release_when_a_lookup_waits ): From 6f7f96960977c1f69987dcd56051e524261fb04a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 18 Oct 2026 06:21:26 -0700 Subject: [PATCH 14/18] fix(cost-map): keep copies of the map whole while the remote map lands --- litellm/litellm_core_utils/get_model_cost_map.py | 5 +---- tests/unit/litellm_core_utils/test_get_model_cost_map.py | 1 - 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index e3145b1078e..916370a0df9 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -51,7 +51,7 @@ from litellm.litellm_core_utils.fallback_generalizations import ( ) if TYPE_CHECKING: - from _collections_abc import dict_items, dict_keys, dict_values + from _collections_abc import dict_items, dict_values FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations" _CATALOG_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) @@ -590,9 +590,6 @@ class _SnapshotIteratingDict(dict[str, object]): def __iter__(self) -> Iterator[str]: return iter(tuple(super().keys())) - def keys(self) -> "dict_keys[str, object]": - return dict(super().items()).keys() - def values(self) -> "dict_values[str, object]": return dict(super().items()).values() diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index 541b971150e..ffe7154bd0d 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -1101,7 +1101,6 @@ def test_a_model_entry_being_read_field_by_field_when_the_remote_map_lands_reads "start_loop", [ pytest.param(iter, id="models"), - pytest.param(lambda model_cost: iter(model_cost.keys()), id="keys"), pytest.param(lambda model_cost: iter(model_cost.values()), id="values"), pytest.param(lambda model_cost: iter(model_cost.items()), id="items"), ], From 6e418bce25c342ecb1d078ceb2683dc1f13410c6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 18 Oct 2026 06:35:50 -0700 Subject: [PATCH 15/18] fix(cost-map): start a retry wait from when the retry was scheduled, not from when import finished --- .../litellm_core_utils/get_model_cost_map.py | 5 ++-- .../test_get_model_cost_map.py | 23 +++++++++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 916370a0df9..0fc2be58b2d 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -955,13 +955,14 @@ def _retry_remote_fetch_in_background( def _run_import_fetch(fetch: _ImportFetch) -> None: def settle_then_sleep(failed: _FetchAttemptRetryable, seconds: float) -> None: scheduled: Final = _PendingImportFetch.next_attempt - _PendingImportFetch.next_attempt = _ScheduledAttempt( + retry: Final = _ScheduledAttempt( attempt=(scheduled.attempt if scheduled is not None else fetch.first_attempt) + 1, at=time.monotonic() + seconds, ) + _PendingImportFetch.next_attempt = retry _report_import_fetch_failure(fetch, f"Remote fetch failed: {failed.reason}") fetch.first_attempt_settled.set() - fetch.sleep(seconds) + fetch.sleep(max(0.0, retry.at - time.monotonic())) try: if fetch.first_attempt > 1: diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index ffe7154bd0d..d618266209c 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -984,6 +984,29 @@ def test_import_load_reports_why_the_last_attempt_failed_while_a_retry_waits(iso assert source["fallback_reason"].startswith("Remote fetch failed: HTTP 429") +def test_a_retry_scheduled_while_import_is_still_running_is_not_pushed_back_by_the_import( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils import get_model_cost_map as module + + monkeypatch.setattr(module, "_litellm_import_complete", threading.Event()) + sleeps = _BlockingSleeps() + client = httpx.Client( + transport=httpx.MockTransport(lambda request: httpx.Response(429, headers={"Retry-After": "2"})) + ) + try: + bundled = module.load_model_cost_map_at_import(url=_URL, client=client, sleep=sleeps, rng=random.Random(0)) + monkeypatch.setattr(litellm, "model_cost", bundled) + time.sleep(0.5) + module.mark_litellm_import_complete() + assert sleeps.called.acquire(timeout=10) + finally: + _finish_fetch(sleeps.release) + + assert sleeps.waits[0] < 1.6 + + def test_a_fetch_resumed_after_fork_adopts_the_remote_map_while_the_inherited_one_hangs( isolated_litellm_cost_state, monkeypatch ): From eb66bc14d94dc841bbe9bc1c4584ab7505a13456 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 18 Oct 2026 08:41:05 -0700 Subject: [PATCH 16/18] fix(cost-map): fill a model registered before the remote map lands from the remote entry A Router built while the first download is still failing writes an empty cost map entry for a bare model the bundled map does not know. The import merge skipped every key written after import, so that model never got its remote provider and its deployment stayed unresolved after the retry landed. Fields written since import still win, and the remote entry fills the rest. --- .../litellm_core_utils/get_model_cost_map.py | 7 ++++- .../test_get_model_cost_map.py | 30 +++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 0fc2be58b2d..69a19ffe6bf 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -852,6 +852,11 @@ def _merge_into_import_map(fetch: _ImportFetch, finalized: Mapping[str, object]) for key, entry in remote if key not in fetch.loaded and key not in fetch.bundled } + added_since_import: Final = { + key: _merged_entry(current, entry, frozenset(current.keys())) + for key, entry in remote + if key not in fetch.loaded and _is_cost_entry(current := fetch.bundled.get(key)) and _is_cost_entry(entry) + } dropped: Final = tuple( key for key in fetch.loaded.keys() - finalized.keys() @@ -859,7 +864,7 @@ def _merge_into_import_map(fetch: _ImportFetch, finalized: Mapping[str, object]) ) for live_entry, entry in updated_in_place: live_entry.update(entry) - fetch.bundled.update({**replaced, **added}) + fetch.bundled.update({**replaced, **added, **added_since_import}) for key in dropped: fetch.bundled.pop(key, None) GetModelCostMap._snapshot_loaded_catalog( # pyright: ignore[reportPrivateUsage] # same module diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index d618266209c..cc62fbf4eae 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -1304,6 +1304,36 @@ def test_a_router_built_before_the_remote_map_lands_prices_its_models_from_the_r assert litellm.get_model_info("openai/gpt-4o")["input_cost_per_token"] == remote_price +def test_a_router_built_before_the_remote_map_lands_serves_a_bare_model_only_the_remote_map_knows( + isolated_litellm_cost_state, monkeypatch +): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, retry_started, release = _rate_limited_then_held_client(_remote_map_with("claude-remote-only-test")) + try: + monkeypatch.setattr( + litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client, sleep=lambda _: None) + ) + assert retry_started.wait(timeout=10) + router = litellm.Router( + model_list=[ + { + "model_name": "remote-only", + "litellm_params": {"model": "claude-remote-only-test", "api_key": "sk-test"}, + } + ], + ignore_invalid_deployments=True, + ) + served_before_the_remote_map = len(router.model_list) + finally: + _finish_fetch(release) + + assert served_before_the_remote_map == 0 + assert [deployment["model_name"] for deployment in router.model_list] == ["remote-only"] + assert litellm.get_model_info("claude-remote-only-test")["litellm_provider"] == "anthropic" + + def test_a_partial_edit_made_before_the_remote_map_lands_keeps_only_the_fields_it_changed( isolated_litellm_cost_state, monkeypatch ): From 8590cad7ed2b46b36b3aeb102cabdfa52a807d6e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 18 Oct 2026 11:18:54 -0700 Subject: [PATCH 17/18] fix(cost-map): move a model's provider membership when the remote map changes or drops it The import merge only ever added provider set memberships, so a bundled model the remote map moved to another provider kept its old one, and get_llm_provider still answered the old provider for its bare name. A model the remote map dropped also stayed in its bundled set, which sent an azure_ai deployment of a dropped OpenAI model down the azure route. The merge now removes the membership the old entry gave a model once the new one is in place, so lookups match a process that loaded the remote map at import. The provider set chain becomes a function that names a model's set, and the sets are typed as string sets. The raw-fork test moves to the integration suite, since unit tests run without child processes, and the other fork test was already covered there. --- litellm/__init__.py | 771 +++++++++--------- .../litellm_core_utils/get_model_cost_map.py | 10 + .../sdk/test_cost_map_background_fetch_sdk.py | 40 +- .../test_get_model_cost_map.py | 142 ++-- .../test_model_prices_and_context_window.py | 2 +- 5 files changed, 524 insertions(+), 441 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index a7e7ff4ae8a..d7b6319fd2c 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -8,7 +8,7 @@ _apply_process_defaults() import threading import os -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from typing import ( Any, Callable, @@ -592,125 +592,125 @@ vertex_ai_safety_settings: Optional[dict] = None ####### COMPLETION MODELS ################### from typing import Set -open_ai_chat_completion_models: Set = set() -open_ai_text_completion_models: Set = set() -cohere_models: Set = set() -cohere_chat_models: Set = set() -mistral_chat_models: Set = set() -text_completion_codestral_models: Set = set() -text_completion_inception_models: Set = set() -anthropic_models: Set = set() -openrouter_models: Set = set() -datarobot_models: Set = set() -vertex_language_models: Set = set() -vertex_vision_models: Set = set() -vertex_chat_models: Set = set() -vertex_code_chat_models: Set = set() -vertex_ai_image_models: Set = set() -vertex_ai_video_models: Set = set() -vertex_text_models: Set = set() -vertex_code_text_models: Set = set() -vertex_embedding_models: Set = set() -vertex_anthropic_models: Set = set() -vertex_llama3_models: Set = set() -vertex_deepseek_models: Set = set() -vertex_ai_ai21_models: Set = set() -vertex_mistral_models: Set = set() -vertex_openai_models: Set = set() -vertex_minimax_models: Set = set() -vertex_moonshot_models: Set = set() -vertex_zai_models: Set = set() -ai21_models: Set = set() -ai21_chat_models: Set = set() -nlp_cloud_models: Set = set() -aleph_alpha_models: Set = set() -bedrock_models: Set = set() -bedrock_converse_models: Set = set(BEDROCK_CONVERSE_MODELS) -fal_ai_models: Set = set() -fireworks_ai_models: Set = set() -fireworks_ai_embedding_models: Set = set() -deepinfra_models: Set = set() -perplexity_models: Set = set() -watsonx_models: Set = set() -gemini_models: Set = set() -xai_models: Set = set() -zai_models: Set = set() -deepseek_models: Set = set() -tencent_models: Set = set() -runwayml_models: Set = set() -azure_ai_models: Set = set() -jina_ai_models: Set = set() -voyage_models: Set = set() -infinity_models: Set = set() -heroku_models: Set = set() -databricks_models: Set = set() -cloudflare_models: Set = set() -codestral_models: Set = set() -friendliai_models: Set = set() -featherless_ai_models: Set = set() -palm_models: Set = set() -groq_models: Set = set() -azure_models: Set = set() -azure_anthropic_models: Set = set() -azure_text_models: Set = set() -anyscale_models: Set = set() -cerebras_models: Set = set() -nadir_models: Set = set() -galadriel_models: Set = set() -nvidia_nim_models: Set = set() -nvidia_riva_models: Set = set() -soniox_models: Set = set() -sambanova_models: Set = set() +open_ai_chat_completion_models: Set[str] = set() +open_ai_text_completion_models: Set[str] = set() +cohere_models: Set[str] = set() +cohere_chat_models: Set[str] = set() +mistral_chat_models: Set[str] = set() +text_completion_codestral_models: Set[str] = set() +text_completion_inception_models: Set[str] = set() +anthropic_models: Set[str] = set() +openrouter_models: Set[str] = set() +datarobot_models: Set[str] = set() +vertex_language_models: Set[str] = set() +vertex_vision_models: Set[str] = set() +vertex_chat_models: Set[str] = set() +vertex_code_chat_models: Set[str] = set() +vertex_ai_image_models: Set[str] = set() +vertex_ai_video_models: Set[str] = set() +vertex_text_models: Set[str] = set() +vertex_code_text_models: Set[str] = set() +vertex_embedding_models: Set[str] = set() +vertex_anthropic_models: Set[str] = set() +vertex_llama3_models: Set[str] = set() +vertex_deepseek_models: Set[str] = set() +vertex_ai_ai21_models: Set[str] = set() +vertex_mistral_models: Set[str] = set() +vertex_openai_models: Set[str] = set() +vertex_minimax_models: Set[str] = set() +vertex_moonshot_models: Set[str] = set() +vertex_zai_models: Set[str] = set() +ai21_models: Set[str] = set() +ai21_chat_models: Set[str] = set() +nlp_cloud_models: Set[str] = set() +aleph_alpha_models: Set[str] = set() +bedrock_models: Set[str] = set() +bedrock_converse_models: Set[str] = set(BEDROCK_CONVERSE_MODELS) +fal_ai_models: Set[str] = set() +fireworks_ai_models: Set[str] = set() +fireworks_ai_embedding_models: Set[str] = set() +deepinfra_models: Set[str] = set() +perplexity_models: Set[str] = set() +watsonx_models: Set[str] = set() +gemini_models: Set[str] = set() +xai_models: Set[str] = set() +zai_models: Set[str] = set() +deepseek_models: Set[str] = set() +tencent_models: Set[str] = set() +runwayml_models: Set[str] = set() +azure_ai_models: Set[str] = set() +jina_ai_models: Set[str] = set() +voyage_models: Set[str] = set() +infinity_models: Set[str] = set() +heroku_models: Set[str] = set() +databricks_models: Set[str] = set() +cloudflare_models: Set[str] = set() +codestral_models: Set[str] = set() +friendliai_models: Set[str] = set() +featherless_ai_models: Set[str] = set() +palm_models: Set[str] = set() +groq_models: Set[str] = set() +azure_models: Set[str] = set() +azure_anthropic_models: Set[str] = set() +azure_text_models: Set[str] = set() +anyscale_models: Set[str] = set() +cerebras_models: Set[str] = set() +nadir_models: Set[str] = set() +galadriel_models: Set[str] = set() +nvidia_nim_models: Set[str] = set() +nvidia_riva_models: Set[str] = set() +soniox_models: Set[str] = set() +sambanova_models: Set[str] = set() scaledown_models: Final[Set[str]] = set() # Price reloads update this registry through existing references. -sambanova_embedding_models: Set = set() -novita_models: Set = set() -assemblyai_models: Set = set() -snowflake_models: Set = set() -gradient_ai_models: Set = set() -llama_models: Set = set() -nscale_models: Set = set() -nebius_models: Set = set() -nebius_embedding_models: Set = set() -aiml_models: Set = set() -deepgram_models: Set = set() -elevenlabs_models: Set = set() -dashscope_models: Set = set() -qwencloud_models: Set = set() -qwen_ai_platform_models: Set = set() -moonshot_models: Set = set() -publicai_models: Set = set() -darkbloom_models: Set = set() +sambanova_embedding_models: Set[str] = set() +novita_models: Set[str] = set() +assemblyai_models: Set[str] = set() +snowflake_models: Set[str] = set() +gradient_ai_models: Set[str] = set() +llama_models: Set[str] = set() +nscale_models: Set[str] = set() +nebius_models: Set[str] = set() +nebius_embedding_models: Set[str] = set() +aiml_models: Set[str] = set() +deepgram_models: Set[str] = set() +elevenlabs_models: Set[str] = set() +dashscope_models: Set[str] = set() +qwencloud_models: Set[str] = set() +qwen_ai_platform_models: Set[str] = set() +moonshot_models: Set[str] = set() +publicai_models: Set[str] = set() +darkbloom_models: Set[str] = set() openinfer_models: Final[Set[str]] = set() -v0_models: Set = set() -morph_models: Set = set() -lambda_ai_models: Set = set() -inception_models: Set = set() -hyperbolic_models: Set = set() -black_forest_labs_models: Set = set() -recraft_models: Set = set() -cometapi_models: Set = set() -oci_models: Set = set() -vercel_ai_gateway_models: Set = set() -edenai_models: Set = set() -volcengine_models: Set = set() -wandb_models: Set = set(WANDB_MODELS) -ovhcloud_models: Set = set() -ovhcloud_embedding_models: Set = set() -lemonade_models: Set = set() -docker_model_runner_models: Set = set() -amazon_nova_models: Set = set() -stability_models: Set = set() -github_copilot_models: Set = set() +v0_models: Set[str] = set() +morph_models: Set[str] = set() +lambda_ai_models: Set[str] = set() +inception_models: Set[str] = set() +hyperbolic_models: Set[str] = set() +black_forest_labs_models: Set[str] = set() +recraft_models: Set[str] = set() +cometapi_models: Set[str] = set() +oci_models: Set[str] = set() +vercel_ai_gateway_models: Set[str] = set() +edenai_models: Set[str] = set() +volcengine_models: Set[str] = set() +wandb_models: Set[str] = set(WANDB_MODELS) +ovhcloud_models: Set[str] = set() +ovhcloud_embedding_models: Set[str] = set() +lemonade_models: Set[str] = set() +docker_model_runner_models: Set[str] = set() +amazon_nova_models: Set[str] = set() +stability_models: Set[str] = set() +github_copilot_models: Set[str] = set() microsoft_365_copilot_models: Set[str] = set() -chatgpt_models: Set = set() -minimax_models: Set = set() -aws_polly_models: Set = set() -transcribe_models: Set = set() -gigachat_models: Set = set() -llamagate_models: Set = set() +chatgpt_models: Set[str] = set() +minimax_models: Set[str] = set() +aws_polly_models: Set[str] = set() +transcribe_models: Set[str] = set() +gigachat_models: Set[str] = set() +llamagate_models: Set[str] = set() coralbricks_models: Set[str] = set() -reducto_models: Set = set() -bedrock_mantle_models: Set = set() +reducto_models: Set[str] = set() +bedrock_mantle_models: Set[str] = set() def is_bedrock_pricing_only_model(key: str) -> bool: @@ -746,268 +746,264 @@ def is_openai_finetune_model(key: str) -> bool: return key.startswith("ft:") and not key.count(":") > 1 +def _provider_model_set(key: str, value: Mapping[str, object]) -> Optional[Tuple[Set[str], str]]: + provider: Final = value.get("litellm_provider") + if provider == "openai" and not is_openai_finetune_model(key): + return open_ai_chat_completion_models, key + if provider == "text-completion-openai": + return open_ai_text_completion_models, key + if provider == "azure_text": + return azure_text_models, key + if provider == "cohere": + return cohere_models, key + if provider == "cohere_chat": + return cohere_chat_models, key + if provider == "mistral": + return mistral_chat_models, key + if provider == "anthropic": + return anthropic_models, key + if provider == "empower": + return empower_models, key + if provider == "openrouter": + return openrouter_models, key + if provider == "vercel_ai_gateway": + return vercel_ai_gateway_models, key + if provider == "edenai": + return edenai_models, key + if provider == "datarobot": + return datarobot_models, key + if provider == "vertex_ai-text-models": + return vertex_text_models, key + if provider == "vertex_ai-code-text-models": + return vertex_code_text_models, key + if provider == "vertex_ai-language-models": + return vertex_language_models, key + if provider == "vertex_ai-vision-models": + return vertex_vision_models, key + if provider == "vertex_ai-chat-models": + return vertex_chat_models, key + if provider == "vertex_ai-code-chat-models": + return vertex_code_chat_models, key + if provider == "vertex_ai-embedding-models": + return vertex_embedding_models, key + if provider == "vertex_ai-anthropic_models": + return vertex_anthropic_models, key.replace("vertex_ai/", "") + if provider == "vertex_ai-llama_models": + return vertex_llama3_models, key.replace("vertex_ai/", "") + if provider == "vertex_ai-deepseek_models": + return vertex_deepseek_models, key.replace("vertex_ai/", "") + if provider == "vertex_ai-mistral_models": + return vertex_mistral_models, key.replace("vertex_ai/", "") + if provider == "vertex_ai-ai21_models": + return vertex_ai_ai21_models, key.replace("vertex_ai/", "") + if provider == "vertex_ai-image-models": + return vertex_ai_image_models, key.replace("vertex_ai/", "") + if provider == "vertex_ai-video-models": + return vertex_ai_video_models, key.replace("vertex_ai/", "") + if provider == "vertex_ai-openai_models": + return vertex_openai_models, key.replace("vertex_ai/", "") + if provider == "vertex_ai-minimax_models": + return vertex_minimax_models, key.replace("vertex_ai/", "") + if provider == "vertex_ai-moonshot_models": + return vertex_moonshot_models, key.replace("vertex_ai/", "") + if provider == "vertex_ai-zai_models": + return vertex_zai_models, key.replace("vertex_ai/", "") + if provider == "ai21": + if value.get("mode") == "chat": + return ai21_chat_models, key + else: + return ai21_models, key + if provider == "nlp_cloud": + return nlp_cloud_models, key + if provider == "aleph_alpha": + return aleph_alpha_models, key + if provider == "bedrock" and value.get("mode") == "guardrail": + return None + if provider == "bedrock" and not is_bedrock_pricing_only_model(key): + return bedrock_models, key + if provider == "bedrock_converse": + return bedrock_converse_models, key + if provider == "deepinfra": + return deepinfra_models, key + if provider == "perplexity": + return perplexity_models, key + if provider == "watsonx": + return watsonx_models, key + if provider == "gemini": + return gemini_models, key + if provider == "fireworks_ai": + # ignore the 'up-to', '-to-' model names -> not real models. just for cost tracking based on model params. + if "-to-" not in key and "fireworks-ai-default" not in key: + return fireworks_ai_models, key + if provider == "fireworks_ai-embedding-models": + # ignore the 'up-to', '-to-' model names -> not real models. just for cost tracking based on model params. + if "-to-" not in key: + return fireworks_ai_embedding_models, key + if provider == "text-completion-codestral": + return text_completion_codestral_models, key + if provider == "text-completion-inception": + return text_completion_inception_models, key + if provider == "xai": + return xai_models, key + if provider == "zai": + return zai_models, key + if provider == "fal_ai": + return fal_ai_models, key + if provider == "deepseek": + return deepseek_models, key + if provider == "tencent": + return tencent_models, key + if provider == "runwayml": + return runwayml_models, key + if provider == "meta_llama": + return llama_models, key + if provider == "nscale": + return nscale_models, key + if provider == "azure_ai": + return azure_ai_models, key + if provider == "voyage": + return voyage_models, key + if provider == "infinity": + return infinity_models, key + if provider == "databricks": + return databricks_models, key + if provider == "cloudflare": + return cloudflare_models, key + if provider == "codestral": + return codestral_models, key + if provider == "friendliai": + return friendliai_models, key + if provider == "palm": + return palm_models, key + if provider == "groq": + return groq_models, key + if provider == "azure": + return azure_models, key + if provider == "azure_anthropic": + return azure_anthropic_models, key + if provider == "anyscale": + return anyscale_models, key + if provider == "cerebras": + return cerebras_models, key + if provider == "nadir": + return nadir_models, key + if provider == "galadriel": + return galadriel_models, key + if provider == "nvidia_nim": + return nvidia_nim_models, key + if provider == "nvidia_riva": + return nvidia_riva_models, key + if provider == "soniox": + return soniox_models, key + if provider == "sambanova": + return sambanova_models, key + if provider == "scaledown": + return scaledown_models, key + if provider == "sambanova-embedding-models": + return sambanova_embedding_models, key + if provider == "novita": + return novita_models, key + if provider == "nebius-chat-models": + return nebius_models, key + if provider == "nebius-embedding-models": + return nebius_embedding_models, key + if provider == "aiml": + return aiml_models, key + if provider == "assemblyai": + return assemblyai_models, key + if provider == "jina_ai": + return jina_ai_models, key + if provider == "snowflake": + return snowflake_models, key + if provider == "gradient_ai": + return gradient_ai_models, key + if provider == "featherless_ai": + return featherless_ai_models, key + if provider == "deepgram": + return deepgram_models, key + if provider == "elevenlabs": + return elevenlabs_models, key + if provider == "heroku": + return heroku_models, key + if provider == "dashscope": + return dashscope_models, key + if provider == "qwencloud": + return qwencloud_models, key + if provider == "qwen_ai_platform": + return qwen_ai_platform_models, key + if provider == "modelscope": + return modelscope_models, key + if provider == "moonshot": + return moonshot_models, key + if provider == "publicai": + return publicai_models, key + if provider == "darkbloom": + return darkbloom_models, key + if provider == "openinfer": + return openinfer_models, key + if provider == "v0": + return v0_models, key + if provider == "morph": + return morph_models, key + if provider == "lambda_ai": + return lambda_ai_models, key + if provider == "inception": + return inception_models, key + if provider == "hyperbolic": + return hyperbolic_models, key + if provider == "black_forest_labs": + return black_forest_labs_models, key + if provider == "recraft": + return recraft_models, key + if provider == "cometapi": + return cometapi_models, key + if provider == "oci": + return oci_models, key + if provider == "volcengine": + return volcengine_models, key + if provider == "wandb": + return wandb_models, key + if provider == "ovhcloud": + return ovhcloud_models, key + if provider == "ovhcloud-embedding-models": + return ovhcloud_embedding_models, key + if provider == "lemonade": + return lemonade_models, key + if provider == "docker_model_runner": + return docker_model_runner_models, key + if provider == "amazon_nova": + return amazon_nova_models, key + if provider == "stability": + return stability_models, key + if provider == "github_copilot": + return github_copilot_models, key + if provider == "microsoft_365_copilot": + return microsoft_365_copilot_models, key + if provider == "chatgpt": + return chatgpt_models, key + if provider == "minimax": + return minimax_models, key + if provider == "aws_polly": + return aws_polly_models, key + if provider == "transcribe": + return transcribe_models, key + if provider == "gigachat": + return gigachat_models, key + if provider == "llamagate": + return llamagate_models, key + if provider == "coralbricks": + return coralbricks_models, key + if provider == "reducto": + return reducto_models, key + if provider == "bedrock_mantle": + return bedrock_mantle_models, key + return None + + def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, object]]) -> None: for key, value in tuple(model_cost_map.items()): - if value.get("litellm_provider") == "openai" and not is_openai_finetune_model(key): - open_ai_chat_completion_models.add(key) - elif value.get("litellm_provider") == "text-completion-openai": - open_ai_text_completion_models.add(key) - elif value.get("litellm_provider") == "azure_text": - azure_text_models.add(key) - elif value.get("litellm_provider") == "cohere": - cohere_models.add(key) - elif value.get("litellm_provider") == "cohere_chat": - cohere_chat_models.add(key) - elif value.get("litellm_provider") == "mistral": - mistral_chat_models.add(key) - elif value.get("litellm_provider") == "anthropic": - anthropic_models.add(key) - elif value.get("litellm_provider") == "empower": - empower_models.add(key) - elif value.get("litellm_provider") == "openrouter": - openrouter_models.add(key) - elif value.get("litellm_provider") == "vercel_ai_gateway": - vercel_ai_gateway_models.add(key) - elif value.get("litellm_provider") == "edenai": - edenai_models.add(key) - elif value.get("litellm_provider") == "datarobot": - datarobot_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-text-models": - vertex_text_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-code-text-models": - vertex_code_text_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-language-models": - vertex_language_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-vision-models": - vertex_vision_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-chat-models": - vertex_chat_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-code-chat-models": - vertex_code_chat_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-embedding-models": - vertex_embedding_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-anthropic_models": - key = key.replace("vertex_ai/", "") - vertex_anthropic_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-llama_models": - key = key.replace("vertex_ai/", "") - vertex_llama3_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-deepseek_models": - key = key.replace("vertex_ai/", "") - vertex_deepseek_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-mistral_models": - key = key.replace("vertex_ai/", "") - vertex_mistral_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-ai21_models": - key = key.replace("vertex_ai/", "") - vertex_ai_ai21_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-image-models": - key = key.replace("vertex_ai/", "") - vertex_ai_image_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-video-models": - key = key.replace("vertex_ai/", "") - vertex_ai_video_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-openai_models": - key = key.replace("vertex_ai/", "") - vertex_openai_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-minimax_models": - key = key.replace("vertex_ai/", "") - vertex_minimax_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-moonshot_models": - key = key.replace("vertex_ai/", "") - vertex_moonshot_models.add(key) - elif value.get("litellm_provider") == "vertex_ai-zai_models": - key = key.replace("vertex_ai/", "") - vertex_zai_models.add(key) - elif value.get("litellm_provider") == "ai21": - if value.get("mode") == "chat": - ai21_chat_models.add(key) - else: - ai21_models.add(key) - elif value.get("litellm_provider") == "nlp_cloud": - nlp_cloud_models.add(key) - elif value.get("litellm_provider") == "aleph_alpha": - aleph_alpha_models.add(key) - elif value.get("litellm_provider") == "bedrock" and value.get("mode") == "guardrail": - pass - elif value.get("litellm_provider") == "bedrock" and not is_bedrock_pricing_only_model(key): - bedrock_models.add(key) - elif value.get("litellm_provider") == "bedrock_converse": - bedrock_converse_models.add(key) - elif value.get("litellm_provider") == "deepinfra": - deepinfra_models.add(key) - elif value.get("litellm_provider") == "perplexity": - perplexity_models.add(key) - elif value.get("litellm_provider") == "watsonx": - watsonx_models.add(key) - elif value.get("litellm_provider") == "gemini": - gemini_models.add(key) - elif value.get("litellm_provider") == "fireworks_ai": - # ignore the 'up-to', '-to-' model names -> not real models. just for cost tracking based on model params. - if "-to-" not in key and "fireworks-ai-default" not in key: - fireworks_ai_models.add(key) - elif value.get("litellm_provider") == "fireworks_ai-embedding-models": - # ignore the 'up-to', '-to-' model names -> not real models. just for cost tracking based on model params. - if "-to-" not in key: - fireworks_ai_embedding_models.add(key) - elif value.get("litellm_provider") == "text-completion-codestral": - text_completion_codestral_models.add(key) - elif value.get("litellm_provider") == "text-completion-inception": - text_completion_inception_models.add(key) - elif value.get("litellm_provider") == "xai": - xai_models.add(key) - elif value.get("litellm_provider") == "zai": - zai_models.add(key) - elif value.get("litellm_provider") == "fal_ai": - fal_ai_models.add(key) - elif value.get("litellm_provider") == "deepseek": - deepseek_models.add(key) - elif value.get("litellm_provider") == "tencent": - tencent_models.add(key) - elif value.get("litellm_provider") == "runwayml": - runwayml_models.add(key) - elif value.get("litellm_provider") == "meta_llama": - llama_models.add(key) - elif value.get("litellm_provider") == "nscale": - nscale_models.add(key) - elif value.get("litellm_provider") == "azure_ai": - azure_ai_models.add(key) - elif value.get("litellm_provider") == "voyage": - voyage_models.add(key) - elif value.get("litellm_provider") == "infinity": - infinity_models.add(key) - elif value.get("litellm_provider") == "databricks": - databricks_models.add(key) - elif value.get("litellm_provider") == "cloudflare": - cloudflare_models.add(key) - elif value.get("litellm_provider") == "codestral": - codestral_models.add(key) - elif value.get("litellm_provider") == "friendliai": - friendliai_models.add(key) - elif value.get("litellm_provider") == "palm": - palm_models.add(key) - elif value.get("litellm_provider") == "groq": - groq_models.add(key) - elif value.get("litellm_provider") == "azure": - azure_models.add(key) - elif value.get("litellm_provider") == "azure_anthropic": - azure_anthropic_models.add(key) - elif value.get("litellm_provider") == "anyscale": - anyscale_models.add(key) - elif value.get("litellm_provider") == "cerebras": - cerebras_models.add(key) - elif value.get("litellm_provider") == "nadir": - nadir_models.add(key) - elif value.get("litellm_provider") == "galadriel": - galadriel_models.add(key) - elif value.get("litellm_provider") == "nvidia_nim": - nvidia_nim_models.add(key) - elif value.get("litellm_provider") == "nvidia_riva": - nvidia_riva_models.add(key) - elif value.get("litellm_provider") == "soniox": - soniox_models.add(key) - elif value.get("litellm_provider") == "sambanova": - sambanova_models.add(key) - elif value.get("litellm_provider") == "scaledown": - scaledown_models.add(key) - elif value.get("litellm_provider") == "sambanova-embedding-models": - sambanova_embedding_models.add(key) - elif value.get("litellm_provider") == "novita": - novita_models.add(key) - elif value.get("litellm_provider") == "nebius-chat-models": - nebius_models.add(key) - elif value.get("litellm_provider") == "nebius-embedding-models": - nebius_embedding_models.add(key) - elif value.get("litellm_provider") == "aiml": - aiml_models.add(key) - elif value.get("litellm_provider") == "assemblyai": - assemblyai_models.add(key) - elif value.get("litellm_provider") == "jina_ai": - jina_ai_models.add(key) - elif value.get("litellm_provider") == "snowflake": - snowflake_models.add(key) - elif value.get("litellm_provider") == "gradient_ai": - gradient_ai_models.add(key) - elif value.get("litellm_provider") == "featherless_ai": - featherless_ai_models.add(key) - elif value.get("litellm_provider") == "deepgram": - deepgram_models.add(key) - elif value.get("litellm_provider") == "elevenlabs": - elevenlabs_models.add(key) - elif value.get("litellm_provider") == "heroku": - heroku_models.add(key) - elif value.get("litellm_provider") == "dashscope": - dashscope_models.add(key) - elif value.get("litellm_provider") == "qwencloud": - qwencloud_models.add(key) - elif value.get("litellm_provider") == "qwen_ai_platform": - qwen_ai_platform_models.add(key) - elif value.get("litellm_provider") == "modelscope": - modelscope_models.add(key) - elif value.get("litellm_provider") == "moonshot": - moonshot_models.add(key) - elif value.get("litellm_provider") == "publicai": - publicai_models.add(key) - elif value.get("litellm_provider") == "darkbloom": - darkbloom_models.add(key) - elif value.get("litellm_provider") == "openinfer": - openinfer_models.add(key) - elif value.get("litellm_provider") == "v0": - v0_models.add(key) - elif value.get("litellm_provider") == "morph": - morph_models.add(key) - elif value.get("litellm_provider") == "lambda_ai": - lambda_ai_models.add(key) - elif value.get("litellm_provider") == "inception": - inception_models.add(key) - elif value.get("litellm_provider") == "hyperbolic": - hyperbolic_models.add(key) - elif value.get("litellm_provider") == "black_forest_labs": - black_forest_labs_models.add(key) - elif value.get("litellm_provider") == "recraft": - recraft_models.add(key) - elif value.get("litellm_provider") == "cometapi": - cometapi_models.add(key) - elif value.get("litellm_provider") == "oci": - oci_models.add(key) - elif value.get("litellm_provider") == "volcengine": - volcengine_models.add(key) - elif value.get("litellm_provider") == "wandb": - wandb_models.add(key) - elif value.get("litellm_provider") == "ovhcloud": - ovhcloud_models.add(key) - elif value.get("litellm_provider") == "ovhcloud-embedding-models": - ovhcloud_embedding_models.add(key) - elif value.get("litellm_provider") == "lemonade": - lemonade_models.add(key) - elif value.get("litellm_provider") == "docker_model_runner": - docker_model_runner_models.add(key) - elif value.get("litellm_provider") == "amazon_nova": - amazon_nova_models.add(key) - elif value.get("litellm_provider") == "stability": - stability_models.add(key) - elif value.get("litellm_provider") == "github_copilot": - github_copilot_models.add(key) - elif value.get("litellm_provider") == "microsoft_365_copilot": - microsoft_365_copilot_models.add(key) - elif value.get("litellm_provider") == "chatgpt": - chatgpt_models.add(key) - elif value.get("litellm_provider") == "minimax": - minimax_models.add(key) - elif value.get("litellm_provider") == "aws_polly": - aws_polly_models.add(key) - elif value.get("litellm_provider") == "transcribe": - transcribe_models.add(key) - elif value.get("litellm_provider") == "gigachat": - gigachat_models.add(key) - elif value.get("litellm_provider") == "llamagate": - llamagate_models.add(key) - elif value.get("litellm_provider") == "coralbricks": - coralbricks_models.add(key) - elif value.get("litellm_provider") == "reducto": - reducto_models.add(key) - elif value.get("litellm_provider") == "bedrock_mantle": - bedrock_mantle_models.add(key) + if (membership := _provider_model_set(key, value)) is not None: + membership[0].add(membership[1]) def add_known_models(model_cost_map: Optional[Dict] = None): @@ -1028,6 +1024,29 @@ def add_known_models_to_model_list() -> None: model_list_set.update(added) +def _provider_model_set_left( + key: str, before: Mapping[str, object], after: Mapping[str, object] +) -> Optional[Tuple[Set[str], str]]: + old: Final = _provider_model_set(key, before) + new: Final = _provider_model_set(key, after) + if old is None or (new is not None and old[0] is new[0] and old[1] == new[1]): + return None + return old + + +def drop_stale_known_models(changes: Iterable[Tuple[str, Mapping[str, object], Mapping[str, object]]]) -> None: + left: Final = tuple( + membership + for key, before, after in changes + if (membership := _provider_model_set_left(key, before, after)) is not None + ) + if not left: + return + for model_set, name in left: + model_set.discard(name) + models_by_provider.update(_build_models_by_provider()) + + _populate_provider_model_sets(model_cost) # known openai compatible endpoints - we'll eventually move this list to the model_prices_and_context_window.json dictionary diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 69a19ffe6bf..4dcdf6e75cb 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -844,6 +844,11 @@ def _merge_into_import_map(fetch: _ImportFetch, finalized: Mapping[str, object]) if entry_id not in reshaped and isinstance(live := shared_once[entry_id][0], MutableMapping) ) replaced: Final = {key: reshaped[id(value)] for key, value in tuple(fetch.bundled.items()) if id(value) in reshaped} + moved: Final = tuple( + (key, MappingProxyType(dict(value)), changed[id(value)]) + for key, value in tuple(fetch.bundled.items()) + if id(value) in changed and _is_cost_entry(value) + ) kept_for_remote: Final = { id(entry): reshaped.get(entry_id, current) for entry_id, (current, _, entry) in shared_once.items() } @@ -862,6 +867,10 @@ def _merge_into_import_map(fetch: _ImportFetch, finalized: Mapping[str, object]) for key in fetch.loaded.keys() - finalized.keys() if _untouched_since_import(fetch.bundled.get(key), fetch.loaded[key], parsed.get(key, _NO_FIELDS)) ) + regrouped: Final = ( + *moved, + *((key, entry, _NO_FIELDS) for key in dropped if _is_cost_entry(entry := fetch.bundled.get(key))), + ) for live_entry, entry in updated_in_place: live_entry.update(entry) fetch.bundled.update({**replaced, **added, **added_since_import}) @@ -873,6 +882,7 @@ def _merge_into_import_map(fetch: _ImportFetch, finalized: Mapping[str, object]) _record_remote_adoption() utils._invalidate_model_cost_lowercase_map() # pyright: ignore[reportPrivateUsage] # required cache invalidation litellm.add_known_models_to_model_list() + litellm.drop_stale_known_models(regrouped) utils.reapply_runtime_model_cost_registrations() _PendingImportFetch.adoptions += 1 diff --git a/tests/integration/sdk/test_cost_map_background_fetch_sdk.py b/tests/integration/sdk/test_cost_map_background_fetch_sdk.py index 6e4fba0f7a4..f2b295c0289 100644 --- a/tests/integration/sdk/test_cost_map_background_fetch_sdk.py +++ b/tests/integration/sdk/test_cost_map_background_fetch_sdk.py @@ -49,6 +49,17 @@ _SCRIPT: Final = textwrap.dedent( join_fetch() os._exit(0 if model in litellm.model_cost else 1) print(os.waitstatus_to_exitcode(os.waitpid(pid, 0)[1]), flush=True) + elif variant == "raw_fork": + import ctypes, signal + pid = ctypes.PyDLL(None).fork() + if pid == 0: + signal.alarm(60) + try: + print(litellm.cost_per_token(model="gpt-4o", prompt_tokens=1000, completion_tokens=0)[0], flush=True) + os._exit(0) + finally: + os._exit(1) + print(os.waitstatus_to_exitcode(os.waitpid(pid, 0)[1]), flush=True) elif variant == "forks_during_retry_wait": while get_model_cost_map_source_info()["fallback_reason"] == "Remote fetch in progress": time.sleep(0.01) @@ -67,9 +78,23 @@ _SCRIPT: Final = textwrap.dedent( ) +def _bundled_map() -> dict: + return json.loads(files("litellm").joinpath("model_prices_and_context_window_backup.json").read_text()) + + +def _remote_gpt_4o_input_cost() -> float: + return 3 * _bundled_map()["gpt-4o"]["input_cost_per_token"] + + def _remote_map() -> bytes: - bundled: Final = json.loads(files("litellm").joinpath("model_prices_and_context_window_backup.json").read_text()) - return json.dumps({**bundled, _REMOTE_ONLY_MODEL: {"litellm_provider": "openai", "mode": "chat"}}).encode() + bundled: Final = _bundled_map() + return json.dumps( + { + **bundled, + "gpt-4o": {**bundled["gpt-4o"], "input_cost_per_token": _remote_gpt_4o_input_cost()}, + _REMOTE_ONLY_MODEL: {"litellm_provider": "openai", "mode": "chat"}, + } + ).encode() @dataclass(frozen=True, slots=True) @@ -160,6 +185,17 @@ def test_a_worker_forked_while_the_remote_map_is_in_flight_fetches_it_itself( assert child.rest == ("0",) +@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") +def test_a_worker_forked_without_the_fork_hooks_while_the_remote_map_is_in_flight_prices_from_it( + cost_map_host: _CostMapHost, tmp_path: Path +) -> None: + child: Final = _run_child(cost_map_host, "raw_fork", tmp_path, answer_after_import=False) + + assert child.after_import == "local Remote fetch in progress False" + assert float(child.rest[0]) == pytest.approx(1000 * _remote_gpt_4o_input_cost()) + assert child.rest[1:] == ("0",) + + @pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") def test_workers_forked_during_a_rate_limited_retry_wait_spread_their_retries(tmp_path: Path) -> None: body: Final = _remote_map() diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index cc62fbf4eae..eae34fd8011 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -5,10 +5,8 @@ count actual model entries, not reserved meta keys) and the extraction of the """ import asyncio, importlib, importlib.resources as importlib_get_model, json, litellm -import ctypes import logging import os -import signal import sys import threading import time @@ -1334,6 +1332,86 @@ def test_a_router_built_before_the_remote_map_lands_serves_a_bare_model_only_the assert litellm.get_model_info("claude-remote-only-test")["litellm_provider"] == "anthropic" +def _bundled_chat_model_of(provider: str) -> str: + return min( + key + for key, entry in _load_root_cost_map().items() + if entry.get("litellm_provider") == provider and entry.get("mode") == "chat" + ) + + +def _remote_map_editing(model: str, fields: dict) -> bytes: + remote_map = _load_root_cost_map() + remote_map[model] = {**remote_map[model], **fields} + return json.dumps(remote_map).encode() + + +def _land_remote_map(monkeypatch, body: bytes) -> None: + import litellm + from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import + + client, retry_started, release = _rate_limited_then_held_client(body) + try: + monkeypatch.setattr( + litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client, sleep=lambda _: None) + ) + assert retry_started.wait(timeout=10) + finally: + _finish_fetch(release) + + +@pytest.mark.parametrize( + "remote_fields, provider_after", + ( + ({"litellm_provider": "text-completion-openai"}, "text-completion-openai"), + ({"input_cost_per_token": 1.0}, "openai"), + ), + ids=("moved", "repriced"), +) +def test_a_bare_model_the_remote_map_edits_resolves_to_the_remote_provider( + isolated_litellm_cost_state, monkeypatch, remote_fields, provider_after +): + import litellm + + model = _bundled_chat_model_of("openai") + assert litellm.get_llm_provider(model)[1] == "openai" + + _land_remote_map(monkeypatch, _remote_map_editing(model, remote_fields)) + + assert litellm.get_llm_provider(model)[1] == provider_after + + +def test_an_azure_ai_deployment_of_a_model_the_remote_map_drops_stays_on_azure_ai( + isolated_litellm_cost_state, monkeypatch +): + import litellm + + model = _bundled_chat_model_of("openai") + deployment = f"azure_ai/{model}" + api_base = "https://example.invalid" + assert litellm.get_llm_provider(deployment, api_base=api_base)[1] == "azure" + remote_map = _load_root_cost_map() + del remote_map[model] + + _land_remote_map(monkeypatch, json.dumps(remote_map).encode()) + + assert litellm.get_llm_provider(deployment, api_base=api_base)[1] == "azure_ai" + + +def test_a_model_the_remote_map_moves_to_another_provider_leaves_the_old_provider_wildcard_list( + isolated_litellm_cost_state, monkeypatch +): + import litellm + + model = _bundled_chat_model_of("azure") + assert model in litellm.models_by_provider["azure"] + + _land_remote_map(monkeypatch, _remote_map_editing(model, {"litellm_provider": "openai"})) + + assert model not in litellm.models_by_provider["azure"] + assert model in litellm.models_by_provider["openai"] + + def test_a_partial_edit_made_before_the_remote_map_lands_keeps_only_the_fields_it_changed( isolated_litellm_cost_state, monkeypatch ): @@ -1761,66 +1839,6 @@ def test_models_listed_while_the_remote_map_is_in_flight_include_a_model_only_th assert "claude-remote-only-test" in valid_models -@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") -def test_a_process_forked_while_the_remote_map_is_in_flight_inherits_it_once_adopted( - isolated_litellm_cost_state, monkeypatch, release_when_a_lookup_waits -): - from litellm.litellm_core_utils import get_model_cost_map as module - - client, started, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) - release_when_a_lookup_waits(release) - try: - monkeypatch.setattr(litellm, "model_cost", module.load_model_cost_map_at_import(url=_URL, client=client)) - assert started.wait(timeout=10) - child = os.fork() - if child == 0: - inherited = "claude-remote-only-test" in litellm.model_cost and module._PendingImportFetch.fetch is None - os._exit(0 if inherited else 1) - _, wait_status = os.waitpid(child, 0) - finally: - _finish_fetch(release) - - assert os.waitstatus_to_exitcode(wait_status) == 0 - - -@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs fork") -def test_a_process_forked_without_the_fork_hooks_while_the_remote_map_is_in_flight_prices_from_it( - isolated_litellm_cost_state, monkeypatch -): - from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import - - body, remote_price = _remote_map_repricing_gpt_4o() - first_request_started = threading.Event() - release = threading.Event() - served = [] - - def handler(request): - served.append(request) - if len(served) == 1: - first_request_started.set() - release.wait(timeout=10) - return httpx.Response(200, content=body) - - client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) - assert first_request_started.wait(timeout=10) - child = ctypes.PyDLL(None).fork() - if child == 0: - signal.signal(signal.SIGALRM, signal.SIG_DFL) - signal.alarm(10) - try: - cost = litellm.cost_per_token(model="gpt-4o", prompt_tokens=1000, completion_tokens=0)[0] - os._exit(0 if cost == pytest.approx(1000 * remote_price) else 1) - finally: - os._exit(2) - _, wait_status = os.waitpid(child, 0) - finally: - _finish_fetch(release) - - assert os.waitstatus_to_exitcode(wait_status) == 0 - - def test_import_load_with_the_local_env_override_starts_no_fetch(monkeypatch): from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index 395fe8c9750..3775d35d8c2 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -697,7 +697,7 @@ def unregistered_providers(rows: Mapping[str, object]) -> list[str]: def test_every_cost_map_provider_is_registered(path: Path): assert unregistered_providers(json.loads(path.read_text())) == [], ( f"{path.name} carries a litellm_provider whose models a `/*` grant does not list. A new provider " - "needs a `_models` set in litellm/__init__.py, filled in _populate_provider_model_sets and listed " + "needs a `_models` set in litellm/__init__.py, returned by _provider_model_set and listed " "in _build_models_by_provider. A new `-` label needs its rows added to a set that " "`models_by_provider[]` includes" ) From bd32f663d12d9f26eb85c1e8c6f2c4e5d26f2452 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 18 Oct 2026 13:12:15 -0700 Subject: [PATCH 18/18] refactor(cost-map): inject the import fetch clock and test retries without real time The retry, uWSGI, and event-loop tests for the import fetch no longer sleep or start timers. A stepped clock drives the retry schedule, the uWSGI case checks the fetch ran on the caller's thread, and the event-loop case releases the fetch from the loop itself. Adds an integration cell where a bare model only the remote map knows is served once a rate-limited first download retries. --- .../litellm_core_utils/get_model_cost_map.py | 9 ++- .../routing/test_stale_cost_map_boot.py | 51 ++++++++++++++- .../test_get_model_cost_map.py | 63 +++++++++++-------- 3 files changed, 93 insertions(+), 30 deletions(-) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 4dcdf6e75cb..85952d6b6a2 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -713,6 +713,7 @@ class _ImportFetch: bundled: dict[str, object] # mutable-ok: litellm.model_cost itself, which the remote map is merged into loaded: Mapping[str, Mapping[str, object]] local: Callable[[], ModelCostMapReloaded] + clock: Callable[[], float] first_attempt: int = 1 first_attempt_at: float = 0.0 first_attempt_settled: threading.Event = field(default_factory=threading.Event) @@ -972,16 +973,16 @@ def _run_import_fetch(fetch: _ImportFetch) -> None: scheduled: Final = _PendingImportFetch.next_attempt retry: Final = _ScheduledAttempt( attempt=(scheduled.attempt if scheduled is not None else fetch.first_attempt) + 1, - at=time.monotonic() + seconds, + at=fetch.clock() + seconds, ) _PendingImportFetch.next_attempt = retry _report_import_fetch_failure(fetch, f"Remote fetch failed: {failed.reason}") fetch.first_attempt_settled.set() - fetch.sleep(max(0.0, retry.at - time.monotonic())) + fetch.sleep(max(0.0, retry.at - fetch.clock())) try: if fetch.first_attempt > 1: - fetch.sleep(max(0.0, fetch.first_attempt_at - time.monotonic())) + fetch.sleep(max(0.0, fetch.first_attempt_at - fetch.clock())) _adopt_remote_model_cost_map( url=fetch.url, timeout=fetch.timeout, @@ -1116,6 +1117,7 @@ def load_model_cost_map_at_import( client: "_SyncGetClient | None" = None, local: Callable[[], ModelCostMapReloaded] = GetModelCostMap.load_local_model_cost_map_with_revision, fetch_before_returning: Callable[[], bool] = _is_uwsgi_process, + clock: Callable[[], float] = time.monotonic, ) -> dict: # mutable-ok: becomes litellm.model_cost, a plain mutable dict by contract """ Return the bundled map right away and fetch the remote one in a background thread, @@ -1152,6 +1154,7 @@ def load_model_cost_map_at_import( bundled=bundled, loaded=GetModelCostMap.loaded_model_cost_map(), local=local, + clock=clock, ) ) return bundled diff --git a/tests/integration/routing/test_stale_cost_map_boot.py b/tests/integration/routing/test_stale_cost_map_boot.py index 13b3d777587..78bae63a6a4 100644 --- a/tests/integration/routing/test_stale_cost_map_boot.py +++ b/tests/integration/routing/test_stale_cost_map_boot.py @@ -1,5 +1,7 @@ +import functools import json import threading +import time import uuid from pathlib import Path from typing import Final @@ -7,7 +9,7 @@ from typing import Final import httpx import pytest from integration._support.client import Gateway, eventually -from integration._support.process import owned_proxy +from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.wire import Reply, Request, wire_server @@ -86,3 +88,50 @@ def test_config_deployment_dropped_by_stale_boot_cost_map_is_restored_after_relo "body": {"model": model, "messages": [{"role": "user", "content": "stale cost map control"}]}, } ] + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_a_bare_model_only_the_remote_map_knows_is_served_once_a_rate_limited_first_download_retries( + gateway: Gateway, tmp_path: Path +) -> None: + model: Final = "integration-remote-only-" + uuid.uuid4().hex + remote_map: Final = json.dumps( + {model: {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 0, "output_cost_per_token": 0}} + ).encode() + + @functools.cache + def first_download_at() -> float: + return time.monotonic() + + def respond(request: Request) -> Reply: + assert request.target == "/model_prices.json", request + if time.monotonic() < first_download_at() + 55: + return Reply(status=429, headers={"Retry-After": "30"}) + return Reply(body=remote_map) + + overrides: Final = {"MODEL_COST_MAP_MIN_MODEL_COUNT": "1", "MODEL_COST_MAP_MAX_SHRINK_RATIO": "0"} + with ( + wire_server(respond) as peer, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + owned_proxy( + gateway, + tmp_path, + {**overrides, "LITELLM_MODEL_COST_MAP_URL": peer.url + "/model_prices.json"}, + config=_proxy_config(tmp_path, model, gateway.upstream_url), + remove_environment=("LITELLM_LOCAL_MODEL_COST_MAP",), + workers=2, + ) as candidate, + ): + upstream.get("/__observations").raise_for_status() + body: Final = {"model": model, "messages": [{"role": "user", "content": "rate-limited cost map control"}]} + eventually( + lambda: candidate.request("POST", "/v1/chat/completions", body).status_code, + lambda status: status == 200, + seconds=90, + ) + statuses: Final = tuple(candidate.request("POST", "/v1/chat/completions", body).status_code for _ in range(6)) + assert statuses == (200,) * 6 + forwarded: Final = tuple( + (observed["path"], observed["body"]) for observed in upstream.get("/__observations").json()["requests"] + ) + assert forwarded == (("/v1/chat/completions", body),) * 7 diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index eae34fd8011..942c0c5d1cd 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -9,7 +9,6 @@ import logging import os import sys import threading -import time from typing import Final import boto3 @@ -816,8 +815,7 @@ def _held_client(outcome): def handler(request): started.set() - release.wait(timeout=10) - return outcome + return outcome if release.wait(timeout=10) else httpx.Response(503) return httpx.Client(transport=httpx.MockTransport(handler)), started, release @@ -950,13 +948,16 @@ def test_import_load_reports_why_the_remote_fetch_failed(isolated_litellm_cost_s def test_import_load_under_uwsgi_returns_the_remote_map_once_fetched(isolated_litellm_cost_state): from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import - client, _, release = _held_client(httpx.Response(200, content=_remote_map_with("claude-remote-only-test"))) - threading.Timer(0.2, release.set).start() - try: - loaded = load_model_cost_map_at_import(url=_URL, client=client, fetch_before_returning=lambda: True) - finally: - _finish_fetch(release) + fetched_on = [] + def handler(request): + fetched_on.append(threading.current_thread()) + return httpx.Response(200, content=_remote_map_with("claude-remote-only-test")) + + client = httpx.Client(transport=httpx.MockTransport(handler)) + loaded = load_model_cost_map_at_import(url=_URL, client=client, fetch_before_returning=lambda: True) + + assert fetched_on == [threading.current_thread()] assert "claude-remote-only-test" in loaded assert get_model_cost_map_source_info()["source"] == "remote" @@ -990,19 +991,34 @@ def test_a_retry_scheduled_while_import_is_still_running_is_not_pushed_back_by_t monkeypatch.setattr(module, "_litellm_import_complete", threading.Event()) sleeps = _BlockingSleeps() + clock = _SteppedClock() client = httpx.Client( transport=httpx.MockTransport(lambda request: httpx.Response(429, headers={"Retry-After": "2"})) ) try: - bundled = module.load_model_cost_map_at_import(url=_URL, client=client, sleep=sleeps, rng=random.Random(0)) + bundled = module.load_model_cost_map_at_import( + url=_URL, client=client, sleep=sleeps, rng=random.Random(0), clock=clock + ) monkeypatch.setattr(litellm, "model_cost", bundled) - time.sleep(0.5) + assert clock.read.wait(timeout=10) + clock.now = 0.5 module.mark_litellm_import_complete() assert sleeps.called.acquire(timeout=10) finally: _finish_fetch(sleeps.release) - assert sleeps.waits[0] < 1.6 + assert sleeps.waits[0] == pytest.approx(1.5) + + +class _SteppedClock: + def __init__(self): + self.now = 0.0 + self.read = threading.Event() + + def __call__(self) -> float: + now = self.now + self.read.set() + return now def test_a_fetch_resumed_after_fork_adopts_the_remote_map_while_the_inherited_one_hangs( @@ -1737,7 +1753,7 @@ def test_a_cost_computed_while_the_remote_map_is_in_flight_uses_its_price( async def test_an_async_call_priced_while_the_remote_map_is_in_flight_keeps_the_event_loop_running( - isolated_litellm_cost_state, monkeypatch + isolated_litellm_cost_state, monkeypatch, release_when_a_lookup_waits ): from litellm.litellm_core_utils.get_model_cost_map import load_model_cost_map_at_import @@ -1751,24 +1767,19 @@ async def test_an_async_call_priced_while_the_remote_map_is_in_flight_keeps_the_ await embed() - async def loop_gaps_while(call): - task = asyncio.ensure_future(call) - ticks = [time.monotonic()] - while not task.done(): - await asyncio.sleep(0.01) - ticks.append(time.monotonic()) - return await task, tuple(later - earlier for earlier, later in zip(ticks, ticks[1:])) - try: - monkeypatch.setattr(litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client)) + monkeypatch.setattr( + litellm, "model_cost", load_model_cost_map_at_import(url=_URL, client=client, max_attempts=1) + ) assert started.wait(timeout=10) - threading.Timer(0.5, release.set).start() - response, gaps = await loop_gaps_while(embed()) + lookup = release_when_a_lookup_waits(threading.Event()) + priced = asyncio.ensure_future(embed()) + assert await asyncio.to_thread(lookup.waited.wait, 10) + release.set() + response = await priced finally: _finish_fetch(release) - assert max(gaps) < 0.25 - assert sum(gaps) >= 0.5 assert response._hidden_params["response_cost"] == pytest.approx(10 * remote_price)