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

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

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)