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