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/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(

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__)

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

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