mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
Merge pull request #40350 from BerriAI/litellm_cost_map_background_retries
fix(cost-map): keep first fetch blocking, run retries in background
This commit is contained in:
commit
bfcc6404d3
6 changed files with 240 additions and 114 deletions
|
|
@ -546,7 +546,7 @@ _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
|
||||
from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map, mark_litellm_import_complete
|
||||
|
||||
model_cost = get_model_cost_map(url=model_cost_map_url)
|
||||
cost_discount_config: Dict[str, float] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount
|
||||
|
|
@ -2405,3 +2405,5 @@ def __getattr__(name: str) -> Any:
|
|||
|
||||
|
||||
# ALL_LITELLM_RESPONSE_TYPES is lazy-loaded via __getattr__ to avoid loading utils at import time
|
||||
|
||||
mark_litellm_import_complete()
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ import hashlib
|
|||
import json
|
||||
import os
|
||||
import random
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass, replace
|
||||
|
|
@ -176,6 +177,11 @@ class GetModelCostMap:
|
|||
RETRYABLE_FETCH_STATUS_CODES: Final = frozenset({429, 500, 502, 503, 504})
|
||||
MODEL_COST_MAP_FETCH_MAX_ATTEMPTS: Final = 3
|
||||
MODEL_COST_MAP_FETCH_MAX_WAIT_SECONDS: Final = 30.0
|
||||
_litellm_import_complete = threading.Event()
|
||||
|
||||
|
||||
def mark_litellm_import_complete() -> None:
|
||||
_litellm_import_complete.set()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -314,12 +320,13 @@ async def _fetch_remote_model_cost_map_with_retry(
|
|||
def _fetch_remote_model_cost_map_with_retry_sync(
|
||||
url: str,
|
||||
timeout: int,
|
||||
max_attempts: int,
|
||||
attempts: range,
|
||||
sleep: Callable[[float], None],
|
||||
rng: random.Random,
|
||||
client: _SyncGetClient,
|
||||
) -> ModelCostMapReloadResult:
|
||||
for attempt in range(1, max_attempts + 1):
|
||||
max_attempts: Final = attempts.stop - 1
|
||||
for attempt in attempts:
|
||||
outcome = _attempt_fetch_sync(client=client, url=url, timeout=timeout)
|
||||
if not isinstance(outcome, _FetchAttemptRetryable):
|
||||
return outcome
|
||||
|
|
@ -520,6 +527,68 @@ def _finalize_loaded_model_cost_map(loaded: ModelCostMapReloaded) -> ModelCostMa
|
|||
return replace(loaded, model_cost_map=_finalize_model_cost_map(loaded.model_cost_map))
|
||||
|
||||
|
||||
def adopt_model_cost_map(
|
||||
new_model_cost_map: dict, # mutable-ok: public API preserves the mutable cost-map contract
|
||||
) -> int:
|
||||
import litellm
|
||||
from litellm import utils
|
||||
|
||||
litellm.model_cost = new_model_cost_map
|
||||
utils._invalidate_model_cost_lowercase_map() # pyright: ignore[reportPrivateUsage] # required cache invalidation
|
||||
litellm.add_known_models(model_cost_map=new_model_cost_map)
|
||||
fetched_model_count: Final = len(new_model_cost_map) if new_model_cost_map else 0
|
||||
utils.reapply_runtime_model_cost_registrations()
|
||||
return fetched_model_count
|
||||
|
||||
|
||||
def _retry_remote_fetch_in_background(
|
||||
url: str,
|
||||
timeout: int,
|
||||
max_attempts: int,
|
||||
sleep: Callable[[float], None],
|
||||
rng: random.Random,
|
||||
client: _SyncGetClient,
|
||||
first_outcome: _FetchAttemptRetryable,
|
||||
) -> None:
|
||||
try:
|
||||
first_wait: Final = _next_retry_wait(outcome=first_outcome, attempt=1, max_attempts=max_attempts, rng=rng)
|
||||
if isinstance(first_wait, ModelCostMapReloadUnavailable):
|
||||
return
|
||||
sleep(first_wait)
|
||||
result: Final = _fetch_remote_model_cost_map_with_retry_sync(
|
||||
url=url,
|
||||
timeout=timeout,
|
||||
attempts=range(2, max_attempts + 1),
|
||||
sleep=sleep,
|
||||
rng=rng,
|
||||
client=client,
|
||||
)
|
||||
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:
|
||||
verbose_logger.warning("LiteLLM: Background model cost map retry failed: %s", e)
|
||||
|
||||
|
||||
def get_model_cost_map(
|
||||
url: str,
|
||||
timeout: int = 5,
|
||||
|
|
@ -532,9 +601,7 @@ def get_model_cost_map(
|
|||
Public entry point — returns the model cost map dict.
|
||||
|
||||
1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set, uses the local backup only.
|
||||
2. Otherwise fetches from ``url``, retrying transient HTTP errors
|
||||
(429/5xx/transport) with Retry-After-aware backoff, validates
|
||||
integrity, and falls back to the local backup on any failure.
|
||||
2. Otherwise fetches from ``url``, retrying transient errors in a background thread.
|
||||
|
||||
Only the backup model count is cached (a single int) for validation.
|
||||
The full backup dict is only parsed when it must be *returned* as a
|
||||
|
|
@ -553,24 +620,34 @@ def get_model_cost_map(
|
|||
_cost_map_source_info.url = url
|
||||
_cost_map_source_info.is_env_forced = False
|
||||
|
||||
result: Final = _fetch_remote_model_cost_map_with_retry_sync(
|
||||
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,
|
||||
)
|
||||
if isinstance(result, ModelCostMapReloadUnavailable):
|
||||
fetch_client: Final = client if client is not None else httpx
|
||||
fetch_rng: Final = rng if rng is not None else random.Random()
|
||||
outcome: Final = _attempt_fetch_sync(client=fetch_client, url=url, timeout=timeout)
|
||||
if isinstance(outcome, _FetchAttemptRetryable) and max_attempts > 1:
|
||||
threading.Thread(
|
||||
target=_retry_remote_fetch_in_background,
|
||||
kwargs={ # mutable-ok: threading requires a mutable keyword-arguments mapping
|
||||
"url": url,
|
||||
"timeout": timeout,
|
||||
"max_attempts": max_attempts,
|
||||
"sleep": sleep,
|
||||
"rng": fetch_rng,
|
||||
"client": fetch_client,
|
||||
"first_outcome": outcome,
|
||||
},
|
||||
name="litellm-model-cost-map-retry",
|
||||
daemon=True,
|
||||
).start()
|
||||
if not isinstance(outcome, ModelCostMapReloaded):
|
||||
verbose_logger.warning(
|
||||
"LiteLLM: Failed to fetch remote model cost map from %s: %s. Falling back to local backup.",
|
||||
url,
|
||||
result.reason,
|
||||
outcome.reason,
|
||||
)
|
||||
_cost_map_source_info.source = "local"
|
||||
_cost_map_source_info.fallback_reason = f"Remote fetch failed: {result.reason}"
|
||||
_cost_map_source_info.fallback_reason = f"Remote fetch failed: {outcome.reason}"
|
||||
return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map
|
||||
content: Final = result.model_cost_map
|
||||
content: Final = outcome.model_cost_map
|
||||
|
||||
# Validate using cached count (cheap int comparison, no file I/O)
|
||||
if not GetModelCostMap.validate_model_cost_map(
|
||||
|
|
@ -587,4 +664,4 @@ def get_model_cost_map(
|
|||
|
||||
_cost_map_source_info.source = "remote"
|
||||
_cost_map_source_info.fallback_reason = None
|
||||
return _finalize_loaded_model_cost_map(result).model_cost_map
|
||||
return _finalize_loaded_model_cost_map(outcome).model_cost_map
|
||||
|
|
|
|||
|
|
@ -138,11 +138,7 @@ from litellm.types.utils import (
|
|||
TextCompletionResponse,
|
||||
TokenCountResponse,
|
||||
)
|
||||
from litellm.utils import (
|
||||
_invalidate_model_cost_lowercase_map,
|
||||
load_credentials_from_list,
|
||||
reapply_runtime_model_cost_registrations,
|
||||
)
|
||||
from litellm.utils import load_credentials_from_list
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from aiohttp import ClientSession
|
||||
|
|
@ -4436,20 +4432,9 @@ def resolve_classifier_plugin(
|
|||
|
||||
|
||||
def _swap_in_model_cost_map(new_model_cost_map: dict) -> int:
|
||||
"""Adopt a freshly fetched cost map into this process's litellm state, return the model count"""
|
||||
litellm.model_cost = new_model_cost_map
|
||||
# Invalidate case-insensitive lookup map since model_cost was replaced
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
# Repopulate provider model sets (e.g. litellm.anthropic_models) so that
|
||||
# wildcard patterns like "anthropic/*" include any newly added models.
|
||||
litellm.add_known_models(model_cost_map=new_model_cost_map)
|
||||
# Counted before the re-apply below, which writes into this same dict, so the
|
||||
# number reported describes the fetched price data alone.
|
||||
fetched_model_count: Final = len(new_model_cost_map) if new_model_cost_map else 0
|
||||
# The swap discards everything registered at runtime (deployment model_info,
|
||||
# register_model overrides), so put it back on top of the fresh catalog.
|
||||
reapply_runtime_model_cost_registrations()
|
||||
return fetched_model_count
|
||||
from litellm.litellm_core_utils.get_model_cost_map import adopt_model_cost_map
|
||||
|
||||
return adopt_model_cost_map(new_model_cost_map)
|
||||
|
||||
|
||||
def should_load_db_object(object_type: str | SupportedDBObjectType) -> bool:
|
||||
|
|
|
|||
|
|
@ -3079,7 +3079,7 @@ def register_model(
|
|||
# Convert stringified numbers to appropriate numeric types
|
||||
loaded_model_cost = model_cost
|
||||
elif isinstance(model_cost, str):
|
||||
loaded_model_cost = litellm.get_model_cost_map(url=model_cost)
|
||||
loaded_model_cost = litellm.get_model_cost_map(url=model_cost, max_attempts=1)
|
||||
|
||||
if persist_across_reloads:
|
||||
_registrations: Final[Mapping[str, Mapping[str, object]]] = loaded_model_cost
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ count actual model entries, not reserved meta keys) and the extraction of the
|
|||
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -26,9 +27,7 @@ from litellm.litellm_core_utils.get_model_cost_map import (
|
|||
|
||||
|
||||
def _load_root_cost_map() -> dict:
|
||||
path = os.path.join(
|
||||
os.path.dirname(__file__), "../../../model_prices_and_context_window.json"
|
||||
)
|
||||
path = os.path.join(os.path.dirname(__file__), "../../../model_prices_and_context_window.json")
|
||||
with open(path) as f:
|
||||
return json.load(f)
|
||||
|
||||
|
|
@ -44,9 +43,7 @@ def test_git_blob_id_is_what_git_hash_object_prints():
|
|||
|
||||
|
||||
def _make_models(n: int) -> dict:
|
||||
return {
|
||||
f"model-{i}": {"litellm_provider": "openai", "mode": "chat"} for i in range(n)
|
||||
}
|
||||
return {f"model-{i}": {"litellm_provider": "openai", "mode": "chat"} for i in range(n)}
|
||||
|
||||
|
||||
def test_count_model_entries_excludes_reserved_keys():
|
||||
|
|
@ -129,9 +126,7 @@ def test_finalize_pops_key_and_installs_rules():
|
|||
def test_finalize_with_no_block_clears_rules():
|
||||
previous = list(get_fallback_generalization_rules())
|
||||
try:
|
||||
set_fallback_generalizations(
|
||||
[{"name": "stale", "pattern": r"^x", "model_info": {"a": 1}}]
|
||||
)
|
||||
set_fallback_generalizations([{"name": "stale", "pattern": r"^x", "model_info": {"a": 1}}])
|
||||
_finalize_model_cost_map(_make_models(2))
|
||||
assert match_capability_generalizations("x-1") is None
|
||||
finally:
|
||||
|
|
@ -317,9 +312,7 @@ def test_get_model_cost_map_stamps_loaded_at():
|
|||
|
||||
from litellm.litellm_core_utils import get_model_cost_map as module
|
||||
|
||||
client, _calls = _mock_client(
|
||||
[httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client
|
||||
)
|
||||
client, _calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client)
|
||||
|
||||
before = datetime.now(timezone.utc)
|
||||
module.get_model_cost_map(url="https://example.invalid/cost_map.json", client=client)
|
||||
|
|
@ -328,6 +321,7 @@ def test_get_model_cost_map_stamps_loaded_at():
|
|||
assert loaded_at is not None
|
||||
assert before <= loaded_at <= datetime.now(timezone.utc)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# refetch_model_cost_map: retry/backoff behavior for runtime reloads
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -394,9 +388,7 @@ async def test_refetch_retries_429_honoring_retry_after():
|
|||
]
|
||||
)
|
||||
sleeper = _SleepRecorder()
|
||||
result = await refetch_model_cost_map(
|
||||
url=_URL, sleep=sleeper, rng=random.Random(0), client=client
|
||||
)
|
||||
result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
|
||||
assert isinstance(result, ModelCostMapReloaded)
|
||||
assert len(result.model_cost_map) > 100
|
||||
assert calls["count"] == 3
|
||||
|
|
@ -408,9 +400,7 @@ async def test_refetch_gives_up_after_max_attempts_with_exponential_backoff():
|
|||
"""All 429 without Retry-After: exponential backoff waits, then a failure value."""
|
||||
client, calls = _mock_client([httpx.Response(429)])
|
||||
sleeper = _SleepRecorder()
|
||||
result = await refetch_model_cost_map(
|
||||
url=_URL, sleep=sleeper, rng=random.Random(0), client=client
|
||||
)
|
||||
result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
|
||||
assert isinstance(result, ModelCostMapReloadUnavailable)
|
||||
assert "429" in result.reason
|
||||
assert "after 3 attempts" in result.reason
|
||||
|
|
@ -430,9 +420,7 @@ async def test_refetch_caps_retry_after_wait():
|
|||
]
|
||||
)
|
||||
sleeper = _SleepRecorder()
|
||||
result = await refetch_model_cost_map(
|
||||
url=_URL, sleep=sleeper, rng=random.Random(0), client=client
|
||||
)
|
||||
result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
|
||||
assert isinstance(result, ModelCostMapReloaded)
|
||||
assert sleeper.waits == [30.0]
|
||||
|
||||
|
|
@ -447,9 +435,7 @@ async def test_refetch_retries_transport_errors():
|
|||
]
|
||||
)
|
||||
sleeper = _SleepRecorder()
|
||||
result = await refetch_model_cost_map(
|
||||
url=_URL, sleep=sleeper, rng=random.Random(0), client=client
|
||||
)
|
||||
result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
|
||||
assert isinstance(result, ModelCostMapReloaded)
|
||||
assert calls["count"] == 2
|
||||
assert len(sleeper.waits) == 1
|
||||
|
|
@ -460,9 +446,7 @@ async def test_refetch_non_retryable_status_fails_immediately():
|
|||
"""A 404 is permanent: one attempt, no sleeps, failure value."""
|
||||
client, calls = _mock_client([httpx.Response(404)])
|
||||
sleeper = _SleepRecorder()
|
||||
result = await refetch_model_cost_map(
|
||||
url=_URL, sleep=sleeper, rng=random.Random(0), client=client
|
||||
)
|
||||
result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
|
||||
assert isinstance(result, ModelCostMapReloadUnavailable)
|
||||
assert "404" in result.reason
|
||||
assert calls["count"] == 1
|
||||
|
|
@ -473,9 +457,7 @@ async def test_refetch_non_retryable_status_fails_immediately():
|
|||
async def test_refetch_invalid_json_fails_immediately():
|
||||
client, calls = _mock_client([httpx.Response(200, content=b"not json")])
|
||||
sleeper = _SleepRecorder()
|
||||
result = await refetch_model_cost_map(
|
||||
url=_URL, sleep=sleeper, rng=random.Random(0), client=client
|
||||
)
|
||||
result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
|
||||
assert isinstance(result, ModelCostMapReloadUnavailable)
|
||||
assert "invalid JSON" in result.reason
|
||||
assert calls["count"] == 1
|
||||
|
|
@ -487,9 +469,7 @@ async def test_refetch_shrunk_map_fails_integrity_not_swapped_in():
|
|||
"""A drastically shrunk upstream file is rejected instead of being adopted."""
|
||||
tiny = json.dumps(_make_models(60)).encode()
|
||||
client, _calls = _mock_client([httpx.Response(200, content=tiny)])
|
||||
result = await refetch_model_cost_map(
|
||||
url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client
|
||||
)
|
||||
result = await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client)
|
||||
assert isinstance(result, ModelCostMapReloadUnavailable)
|
||||
assert "integrity validation" in result.reason
|
||||
|
||||
|
|
@ -586,69 +566,125 @@ from litellm.litellm_core_utils.get_model_cost_map import (
|
|||
class _SyncSleepRecorder:
|
||||
"""Injected in place of time.sleep so the boot path's waits are asserted without delay."""
|
||||
|
||||
def __init__(self):
|
||||
def __init__(self, block=False):
|
||||
self.waits = []
|
||||
self.block = block
|
||||
self.started = threading.Event()
|
||||
self.release = threading.Event()
|
||||
|
||||
def __call__(self, seconds: float) -> None:
|
||||
if self.block:
|
||||
self.started.set()
|
||||
self.release.wait(timeout=10)
|
||||
self.waits.append(seconds)
|
||||
|
||||
|
||||
def test_boot_load_retries_transient_failures_instead_of_falling_back():
|
||||
"""A refused connection then a 503 at pod boot used to pin the process to the bundled
|
||||
backup for its lifetime; both are transient and must be retried before giving up."""
|
||||
def _retry_threads():
|
||||
return [thread for thread in threading.enumerate() if thread.name == "litellm-model-cost-map-retry"]
|
||||
|
||||
|
||||
def test_boot_load_success_does_not_start_background_retry():
|
||||
client, calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client)
|
||||
sleeper = _SyncSleepRecorder()
|
||||
|
||||
cost_map = get_model_cost_map(
|
||||
url=_URL,
|
||||
sleep=sleeper,
|
||||
rng=random.Random(0),
|
||||
client=client,
|
||||
)
|
||||
assert calls["count"] == 1
|
||||
assert sleeper.waits == []
|
||||
assert _retry_threads() == []
|
||||
assert cost_map.keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY}
|
||||
assert get_model_cost_map_source_info()["source"] == "remote"
|
||||
|
||||
|
||||
def test_boot_load_transient_failure_returns_local_then_background_retry_adopts_remote(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))
|
||||
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))
|
||||
monkeypatch.setattr(litellm, "models_by_provider", dict(litellm.models_by_provider))
|
||||
monkeypatch.setattr(
|
||||
litellm_utils,
|
||||
"_runtime_registered_model_cost",
|
||||
dict(litellm_utils._runtime_registered_model_cost),
|
||||
)
|
||||
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))
|
||||
|
||||
remote_map = _load_root_cost_map()
|
||||
remote_map["claude-remote-only-test"] = {"litellm_provider": "anthropic", "mode": "chat"}
|
||||
client, calls = _mock_client(
|
||||
[
|
||||
httpx.ConnectError("connection refused"),
|
||||
httpx.Response(503),
|
||||
httpx.Response(200, content=_real_map_bytes()),
|
||||
httpx.Response(200, content=json.dumps(remote_map).encode()),
|
||||
],
|
||||
client_cls=httpx.Client,
|
||||
)
|
||||
sleeper = _SyncSleepRecorder()
|
||||
sleeper = _SyncSleepRecorder(block=True)
|
||||
litellm.register_model({"my-runtime-model": {"litellm_provider": "custom", "max_input_tokens": 4321}})
|
||||
|
||||
cost_map = get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
|
||||
|
||||
assert calls["count"] == 3
|
||||
assert len(sleeper.waits) == 2
|
||||
assert 2.0 <= sleeper.waits[0] < 3.0
|
||||
assert 4.0 <= sleeper.waits[1] < 5.0
|
||||
source = get_model_cost_map_source_info()
|
||||
assert source["source"] == "remote"
|
||||
assert source["fallback_reason"] is None
|
||||
assert cost_map.keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY}
|
||||
|
||||
|
||||
def test_boot_load_honors_retry_after_then_falls_back_after_max_attempts():
|
||||
"""An outage longer than the retry budget still ends on the bundled backup, and the
|
||||
recorded fallback reason says how many attempts were spent so operators can tell."""
|
||||
client, calls = _mock_client(
|
||||
[httpx.Response(429, headers={"Retry-After": "7"})], client_cls=httpx.Client
|
||||
cost_map = get_model_cost_map(
|
||||
url=_URL,
|
||||
max_attempts=3,
|
||||
sleep=sleeper,
|
||||
rng=random.Random(0),
|
||||
client=client,
|
||||
)
|
||||
sleeper = _SyncSleepRecorder()
|
||||
|
||||
cost_map = get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
|
||||
|
||||
assert calls["count"] == 3
|
||||
assert sleeper.waits == [7.0, 7.0]
|
||||
source = get_model_cost_map_source_info()
|
||||
assert source["source"] == "local"
|
||||
assert "after 3 attempts" in source["fallback_reason"]
|
||||
assert len(cost_map) > 100
|
||||
assert calls["count"] == 1
|
||||
assert sleeper.waits == []
|
||||
assert sleeper.started.wait(timeout=10)
|
||||
threads = _retry_threads()
|
||||
try:
|
||||
assert len(threads) == 1
|
||||
assert "claude-remote-only-test" not in cost_map
|
||||
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"].startswith("Remote fetch failed:")
|
||||
sleeper.release.set()
|
||||
for thread in threads:
|
||||
thread.join(timeout=10)
|
||||
assert all(not thread.is_alive() for thread in threads)
|
||||
assert sleeper.waits and 2.0 <= sleeper.waits[0] < 3.0
|
||||
assert calls["count"] == 2
|
||||
assert "claude-remote-only-test" in litellm.model_cost
|
||||
assert "claude-remote-only-test" in litellm.anthropic_models
|
||||
assert "my-runtime-model" in litellm.model_cost
|
||||
source = get_model_cost_map_source_info()
|
||||
assert source["source"] == "remote"
|
||||
assert source["fallback_reason"] is None
|
||||
finally:
|
||||
sleeper.release.set()
|
||||
for thread in _retry_threads():
|
||||
thread.join(timeout=10)
|
||||
|
||||
|
||||
def test_boot_load_does_not_retry_permanent_failures():
|
||||
"""A 404 or a malformed URL cannot heal by waiting: one attempt, no sleeps, backup."""
|
||||
def test_boot_load_does_not_retry_non_retryable_failure():
|
||||
client, calls = _mock_client([httpx.Response(404)], client_cls=httpx.Client)
|
||||
sleeper = _SyncSleepRecorder()
|
||||
|
||||
get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
|
||||
get_model_cost_map(
|
||||
url=_URL,
|
||||
sleep=sleeper,
|
||||
rng=random.Random(0),
|
||||
client=client,
|
||||
)
|
||||
assert calls["count"] == 1
|
||||
assert sleeper.waits == []
|
||||
assert get_model_cost_map_source_info()["source"] == "local"
|
||||
|
||||
get_model_cost_map(url="not a url", sleep=sleeper, rng=random.Random(0))
|
||||
assert sleeper.waits == []
|
||||
assert get_model_cost_map_source_info()["source"] == "local"
|
||||
assert _retry_threads() == []
|
||||
source = get_model_cost_map_source_info()
|
||||
assert source["source"] == "local"
|
||||
assert source["fallback_reason"] is not None
|
||||
|
||||
|
||||
def test_boot_load_respects_local_env_override(monkeypatch):
|
||||
|
|
@ -701,7 +737,9 @@ def test_boot_load_that_fails_the_integrity_check_reports_the_backup_not_the_rej
|
|||
)
|
||||
get_model_cost_map(url=_URL, sleep=_SyncSleepRecorder(), rng=random.Random(0), client=remote)
|
||||
shrunk_body = b'{"gpt-5.4-mini": {"mode": "chat", "input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06}}'
|
||||
shrunk, _ = _mock_client([httpx.Response(200, headers={"ETag": 'W/"shrunk"'}, content=shrunk_body)], client_cls=httpx.Client)
|
||||
shrunk, _ = _mock_client(
|
||||
[httpx.Response(200, headers={"ETag": 'W/"shrunk"'}, content=shrunk_body)], client_cls=httpx.Client
|
||||
)
|
||||
|
||||
get_model_cost_map(url=_URL, sleep=_SyncSleepRecorder(), rng=random.Random(0), client=shrunk)
|
||||
|
||||
|
|
|
|||
|
|
@ -2,10 +2,12 @@ import asyncio
|
|||
import json
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from jsonschema import validate
|
||||
|
|
@ -2382,6 +2384,28 @@ def test_register_model_with_scientific_notation():
|
|||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_register_model_url_fetch_uses_single_attempt(monkeypatch):
|
||||
monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False)
|
||||
monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost))
|
||||
before = dict(litellm.model_cost)
|
||||
threads_before = {thread.name for thread in threading.enumerate()}
|
||||
route = respx.get("https://example.invalid/custom_pricing.json").mock(
|
||||
return_value=httpx.Response(503)
|
||||
)
|
||||
|
||||
litellm.register_model(model_cost="https://example.invalid/custom_pricing.json")
|
||||
|
||||
threads_after = {thread.name for thread in threading.enumerate()}
|
||||
assert route.call_count == 1
|
||||
assert not (threads_after - threads_before) & {"litellm-model-cost-map-retry"}
|
||||
assert not any(
|
||||
thread.name == "litellm-model-cost-map-retry" and thread.is_alive()
|
||||
for thread in threading.enumerate()
|
||||
)
|
||||
assert litellm.model_cost.keys() >= before.keys()
|
||||
|
||||
|
||||
def test_register_model_openrouter_without_slash():
|
||||
"""
|
||||
Test that register_model handles openrouter models without '/' in the name.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue