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] 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 ):