diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 72bcb7f9dee..7bd37aa1aa4 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -177,8 +177,8 @@ Requirements and semantics: - **Redis is required.** Session payloads, per-model warmth stamps, and a per-router session index live in Redis so all pods share them and a single pod (via a Redis cron lock) runs the replays. Without Redis, warming logs a warning once and no-ops; requests are unaffected. Sessions are tracked through the index rather than keyspace scans, so Redis Cluster is supported. - **Prompt retention consent is a prerequisite.** Warming persists full request payloads (messages, system, tools) in Redis, so capture requires `store_prompts_in_spend_logs: true` and respects message redaction: with the flag off, or with `turn_off_message_logging` (globally or via the per-request redaction header) active, capture warns once and skips. - **Only Anthropic and Bedrock models that support prompt caching are warmed.** Other models in the tier pools are left alone. Requests must carry a `metadata.session_id` and exceed the warm set's minimum cacheable token count (`prompt_cache_min_tokens`, default 1024) to be captured. -- **A replay is admitted like a request.** Each replay is assembled as a request body and put through the proxy's own admission entry points before it is dispatched: the budget reservation (the same call the auth layer makes after `common_checks`, so the key, team, user, end-user, organization and tag counters all apply and are reconciled to the replay's actual cost) and then `ProxyLogging.pre_call_hook` (so RPM, TPM, max-parallel and every configured guardrail apply, on the proxy's own shared counters rather than a private copy). A rejection skips that one replay and is retried on the next tick; the failure hook returns whatever the rejected replay had already reserved. Warming stops for keys that are deleted, blocked, or expired; key state is verified fresh each tick and when the database is unreachable the tick is skipped, so warming pauses until it is reachable again (caches re-warm on the next successful tick). -- **Warming cost is visible where the customer already looks for cost.** Warming writes no spend logs of its own, so the replay rows are the only record of warming cost that exists, and they carry the same identity block a real request on that key carries: key, team, user and end-user attribution, the key's and team's own `tags` and `spend_logs_metadata`, and the `litellm_cache_warming` tag alongside them so warming is both included in per-tag chargeback and filterable out of it. Replays also fan out to the key-scoped and team-scoped logging callbacks, so a team pointing its traffic at its own Langfuse sees warming there too rather than only in the proxy-wide logs. A `max_tokens=1` replay of a warm prefix bills roughly 10% of the input cost. Key-level cache controls, `disable_fallbacks` and the global-guardrail opt-outs are applied the same way, which means a key that declares its own `cache` controls overrides warming's response-cache bypass exactly as it overrides a caller's. +- **A replay is admitted like a request.** Each replay is assembled as a request body and put through the proxy's own admission entry points before it is dispatched: the budget reservation (the same call the auth layer makes after `common_checks`, so the key, team, user, end-user, organization and tag counters all apply and are reconciled to the replay's actual cost) and then `ProxyLogging.pre_call_hook` (so RPM, TPM, max-parallel and every configured guardrail apply, on the proxy's own shared counters rather than a private copy). A rejection skips that one replay; the failure hook returns whatever the rejected replay had already reserved, and the attempt is recorded as attempted-but-cold, so the routing pick keeps treating that model as cold while the attempt still paces the retry. Warming stops for keys that are deleted, blocked, or expired; key state is verified fresh each tick and when the database is unreachable the tick is skipped, so warming pauses until it is reachable again (caches re-warm on the next successful tick). +- **Warming cost is visible where the customer already looks for cost.** Warming writes no spend logs of its own, so the replay rows are the only record of warming cost that exists, and they carry the same identity block a real request on that key carries: key, team, user and end-user attribution, the key's and team's own `tags` and `spend_logs_metadata`, plus a `litellm_cache_warming` marker in `spend_logs_metadata` so warming rows are identifiable in per-tag chargeback without putting warming's own tag on `metadata.tags`, which is an input to deployment selection and which operators can forbid requests from carrying. Replays also fan out to the key-scoped and team-scoped logging callbacks, so a team pointing its traffic at its own Langfuse sees warming there too rather than only in the proxy-wide logs. A `max_tokens=1` replay of a warm prefix bills roughly 10% of the input cost. Key-level cache controls, `disable_fallbacks` and the global-guardrail opt-outs are applied the same way, which means a key that declares its own `cache` controls overrides warming's response-cache bypass exactly as it overrides a caller's. - **Guardrails run on replays.** Because a replay goes through `pre_call_hook`, the guardrails configured for the model group (globally and on the deployment) also run on its warming replays. A blocking guardrail costs one skipped warm; a content-rewriting guardrail makes the warm ineffective rather than wrong, because a rewritten prefix simply is not the prefix real traffic sends. Sessions on a key that declares `max_iterations` are skipped instead of warmed, because that limiter counts every request on a `session_id` and cannot be consulted without incrementing it, so warming would consume the caller's own iteration budget. - **Multi-deployment groups require deployment affinity.** A replay routes by group name, so with several deployments it warms one member's cache while real traffic spreads across all of them: paying the cache-write premium against 1/N routing odds is worse than not warming, so such a group is skipped (with a one-time warning naming the group and both remedies) unless `DeploymentAffinityCheck` is active for it with the session_id mode, enabled globally via `router_settings.optional_pre_call_checks: ["session_affinity"]` or per group via `router_settings.model_group_affinity_config`. When active, replays carry the session's `session_id` (and the originating key hash), so the affinity check pins warming and real traffic to the same deployment. Single-deployment groups warm regardless. More than one deployment is used as the conservative stand-in for "more than one provider cache domain" - **`max_sessions`** caps concurrently warmed sessions per auto-router, enforced atomically at capture; once reached, new sessions are not admitted until existing ones expire. diff --git a/litellm/router_strategy/complexity_router/cache_warming/refresher.py b/litellm/router_strategy/complexity_router/cache_warming/refresher.py index d509ead716f..6bb80233c4d 100644 --- a/litellm/router_strategy/complexity_router/cache_warming/refresher.py +++ b/litellm/router_strategy/complexity_router/cache_warming/refresher.py @@ -580,7 +580,9 @@ class CacheWarmingRefresher: warmth = await store.get_warmth(session_key, warm_models) now = time.time() due_models = tuple( - model for model in warm_models if needs_rewarming(warmth.get(model, 0.0), now, refresh_interval_seconds) + model + for model in warm_models + if needs_rewarming(warmth[model].at if model in warmth else 0.0, now, refresh_interval_seconds) ) if not due_models: return @@ -589,38 +591,68 @@ class CacheWarmingRefresher: async with semaphore: if lease_lost.is_set(): return - admitted = await self._admit_replay( + attempted_at = time.time() + warmed = await self._replay_once( llm_router=llm_router, payload=payload, record=record, + session_key=session_key, model_group=model_group, key_state=key_state, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) - if admitted is None: - continue - data, principal = admitted - attempted_at = time.time() - try: - await ( - llm_router.aanthropic_messages(**data) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # factory-generated router surface is legacy-untyped - if payload.call_surface == "anthropic_messages" - else llm_router.acompletion(**data) # pyright: ignore[reportUnknownMemberType, reportCallIssue, reportUnknownVariableType, reportArgumentType] # router overloads are legacy-untyped - ) - await proxy_logging_obj.update_request_status( - litellm_call_id=str(data.get("litellm_call_id") or ""), status="success" - ) - except Exception as exc: # noqa: BLE001 # one failing replay must not abort the tick - verbose_router_logger.warning( - "cache_warming replay failed for session %s model %s", session_key, model_group, exc_info=True - ) - await self._report_rejection( - data=data, principal=principal, exc=exc, proxy_logging_obj=proxy_logging_obj - ) - finally: - await store.mark_warm_attempt(session_key, model_group, attempted_at, session_ttl_seconds) + await store.mark_warm_attempt( + session_key, model_group, attempted_at, session_ttl_seconds, warmed=warmed + ) + + async def _replay_once( + self, + *, + llm_router: "Router", + payload: CacheWarmingPayload, + record: CacheWarmingRecord, + session_key: str, + model_group: str, + key_state: "UserAPIKeyAuth | None", + prisma_client: "PrismaClient | None", + user_api_key_cache: "DualCache | None", + proxy_logging_obj: "ProxyLogging", + ) -> bool: + """True only when the provider accepted the replay, because a warmth stamp is the claim that the + provider now holds this session's prefix on this model. A refused or failed replay leaves it cold, so + it is stamped as attempted-but-cold: the pick keeps treating the model as cold, while the attempt + still paces the next tick instead of retrying a failing model every 30 seconds.""" + admitted = await self._admit_replay( + llm_router=llm_router, + payload=payload, + record=record, + model_group=model_group, + key_state=key_state, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if admitted is None: + return False + data, principal = admitted + try: + await ( + llm_router.aanthropic_messages(**data) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # factory-generated router surface is legacy-untyped + if payload.call_surface == "anthropic_messages" + else llm_router.acompletion(**data) # pyright: ignore[reportUnknownMemberType, reportCallIssue, reportUnknownVariableType, reportArgumentType] # router overloads are legacy-untyped + ) + except Exception as exc: # noqa: BLE001 # one failing replay must not abort the tick + verbose_router_logger.warning( + "cache_warming replay failed for session %s model %s", session_key, model_group, exc_info=True + ) + await self._report_rejection(data=data, principal=principal, exc=exc, proxy_logging_obj=proxy_logging_obj) + return False + await proxy_logging_obj.update_request_status( + litellm_call_id=str(data.get("litellm_call_id") or ""), status="success" + ) + return True async def _admit_replay( self, @@ -675,7 +707,6 @@ class CacheWarmingRefresher: await self._report_rejection(data=data, principal=principal, exc=exc, proxy_logging_obj=proxy_logging_obj) return None - @staticmethod @staticmethod async def _authorize_and_reserve( *, diff --git a/litellm/router_strategy/complexity_router/cache_warming/store.py b/litellm/router_strategy/complexity_router/cache_warming/store.py index 1a77bbedd9b..a141a97994b 100644 --- a/litellm/router_strategy/complexity_router/cache_warming/store.py +++ b/litellm/router_strategy/complexity_router/cache_warming/store.py @@ -11,6 +11,7 @@ from litellm.router_strategy.complexity_router.cache_warming.types import ( CACHE_WARMING_RECORD_SCHEMA_VERSION, CacheWarmingAttribution, CacheWarmingRecord, + WarmthStamp, ) _WARMTH_KEY_PREFIX = "complexity_router_cache_warmth:v1" @@ -82,15 +83,13 @@ def _parse_record(raw: object) -> CacheWarmingRecord | None: return record -def _parse_warmth(raw: object) -> float | None: - if isinstance(raw, (int, float)): - return float(raw) - if isinstance(raw, (str, bytes)): - try: - return float(raw) - except ValueError: - return None - return None +def _parse_warmth(raw: object) -> WarmthStamp | None: + try: + if isinstance(raw, (str, bytes)): + return WarmthStamp.model_validate_json(raw) + return WarmthStamp.model_validate(raw) + except ValidationError: + return None class CacheWarmingStore: @@ -188,17 +187,21 @@ class CacheWarmingStore: if admitted != 1: _warn_session_cap_reached(self.auto_router_model_name) return - await self.mark_warm_attempt(key, served_model, attempted_at=now, ttl_seconds=ttl_seconds) + await self.mark_warm_attempt(key, served_model, attempted_at=now, ttl_seconds=ttl_seconds, warmed=True) - async def mark_warm_attempt(self, key: str, model_group: str, attempted_at: float, ttl_seconds: int) -> None: + async def mark_warm_attempt( + self, key: str, model_group: str, attempted_at: float, ttl_seconds: int, warmed: bool + ) -> None: redis_cache = self._require_redis() if redis_cache is None: return await redis_cache.async_set_cache( # pyright: ignore[reportUnknownMemberType] # RedisCache is legacy-untyped - key=self.warmth_key(key, model_group), value=attempted_at, ttl=ttl_seconds + key=self.warmth_key(key, model_group), + value=WarmthStamp(at=attempted_at, warmed=warmed).model_dump(), + ttl=ttl_seconds, ) - async def get_warmth(self, key: str, model_groups: tuple[str, ...]) -> Mapping[str, float]: + async def get_warmth(self, key: str, model_groups: tuple[str, ...]) -> Mapping[str, WarmthStamp]: redis_cache = self._require_redis() if redis_cache is None: return {} # mutable-ok: fresh per-call result, not shared state diff --git a/litellm/router_strategy/complexity_router/cache_warming/types.py b/litellm/router_strategy/complexity_router/cache_warming/types.py index 7308495f514..74f8207534d 100644 --- a/litellm/router_strategy/complexity_router/cache_warming/types.py +++ b/litellm/router_strategy/complexity_router/cache_warming/types.py @@ -36,6 +36,17 @@ def needs_rewarming(warmed_at: float, now: float, refresh_interval_seconds: int) ) +class WarmthStamp(BaseModel): + """When a model group was last replayed for a session, and whether that replay actually landed. A failed + replay leaves the provider cache cold, so ``warmed`` keeps the pick from preferring it while ``at`` still + paces the next attempt.""" + + model_config = ConfigDict(extra="forbid") + + at: float + warmed: bool + + class CacheWarmingPayload(BaseModel): model_config = ConfigDict(extra="forbid") diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 592cd75c686..d129b2a50f7 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -671,7 +671,7 @@ class ComplexityRouter(CustomLogger): return None warmth = await store.get_warmth(record_key, tuple(pool)) now = time.time() - warmed = frozenset(model for model, warmed_at in warmth.items() if is_cache_fresh(warmed_at, now)) + warmed = frozenset(model for model, stamp in warmth.items() if stamp.warmed and is_cache_fresh(stamp.at, now)) served = frozenset((record.served_model,)) if is_cache_fresh(record.last_activity, now) else frozenset[str]() candidates = tuple(model for model in pool if model in warmed | served) if not candidates: diff --git a/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_capture.py b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_capture.py index 968d178e4c1..d189b060a76 100644 --- a/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_capture.py +++ b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_capture.py @@ -13,6 +13,7 @@ from litellm.router_strategy.complexity_router.cache_warming.capture import ( from litellm.router_strategy.complexity_router.cache_warming.store import CacheWarmingStore from litellm.router_strategy.complexity_router.cache_warming.types import ( CACHE_WARMING_REPLAY_MARKER_KEY, + WarmthStamp, decompress_payload, ) from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter @@ -100,13 +101,15 @@ async def test_second_turn_overwrites_payload_and_preserves_other_model_warmth() await router._capture_session(_kwargs(), MESSAGES, "claude-sonnet-4-5") key = CacheWarmingStore.record_key("smart-router", "hash-1", "sess-1") first = json.loads(redis.hashes[SESSIONS_KEY][key]) - redis.data[CacheWarmingStore.warmth_key(key, "gpt-5-mini")] = json.dumps(123.0) + other_stamp = WarmthStamp(at=123.0, warmed=True).model_dump() + redis.data[CacheWarmingStore.warmth_key(key, "gpt-5-mini")] = json.dumps(other_stamp) await router._capture_session( _kwargs(), MESSAGES + [{"role": "user", "content": "and rule 8?"}], "claude-sonnet-4-5" ) second = json.loads(redis.hashes[SESSIONS_KEY][key]) - assert json.loads(redis.data[CacheWarmingStore.warmth_key(key, "gpt-5-mini")]) == 123.0 - assert json.loads(redis.data[CacheWarmingStore.warmth_key(key, "claude-sonnet-4-5")]) > 0 + assert json.loads(redis.data[CacheWarmingStore.warmth_key(key, "gpt-5-mini")]) == other_stamp + served_stamp = WarmthStamp.model_validate_json(redis.data[CacheWarmingStore.warmth_key(key, "claude-sonnet-4-5")]) + assert served_stamp.at > 0 and served_stamp.warmed is True assert second["payload_sha256"] != first["payload_sha256"] diff --git a/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_refresher.py b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_refresher.py index a3c28d56de1..1770eabd0dc 100644 --- a/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_refresher.py +++ b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_refresher.py @@ -219,7 +219,42 @@ async def test_a_due_session_is_replayed_on_its_own_surface_and_stamped_warm(sur if surface == "anthropic_messages": assert call["system"] == "You are a policy assistant" stamp = warmth_stamp(redis, record_key, "smart-claude") - assert stamp is not None and stamp > 0 + assert stamp is not None and stamp.at > 0 and stamp.warmed is True + + +@pytest.mark.asyncio +async def test_a_replay_the_provider_rejected_is_stamped_cold_and_still_paced(): + """A stamp is the claim that the provider holds this session's prefix on that model, so a replay that + failed must not write a warm one: the pick would then send real traffic to a cold model for the provider's + whole cache TTL, which is the outcome warming exists to avoid. The attempt is still recorded, so a model + failing every replay is retried on the refresh interval rather than on every tick.""" + llm_router, redis = warming_rig(redis=FakeRedisCache()) + llm_router.failing_message_marker = "deployment policy" + record_key = seed_session(redis, warmth={"fast-claude": time.time()}) + await tick(llm_router) + assert llm_router.completion_calls == [] and len(llm_router.failed_calls) == 1 + stamp = warmth_stamp(redis, record_key, "smart-claude") + assert stamp is not None and stamp.at > 0 and stamp.warmed is False + + await tick(llm_router) + assert len(llm_router.failed_calls) == 1 + + +@pytest.mark.asyncio +async def test_a_replay_refused_by_admission_is_stamped_cold_rather_than_warm(): + """Same claim, for the replay that never reached the provider at all: a refusal from the request path's + own ceilings leaves the cache exactly as cold as a provider failure does.""" + limiter, counters = real_limiter() + llm_router, redis = warming_rig(redis=FakeRedisCache()) + record_key = seed_session(redis, user_api_key="k", warmth={"fast-claude": time.time()}) + keys = FakeKeyDirectory({"k": key_state(token="k", rpm_limit=1)}) + await counters.async_set_cache(key="{api_key:k}:window", value=str(int(time.time()))) + await counters.async_set_cache(key="{api_key:k}:requests", value=1) + with registered_callbacks(limiter): + await tick(llm_router, active=refresher(keys=keys, limiter=limiter)) + assert llm_router.completion_calls == [] + stamp = warmth_stamp(redis, record_key, "smart-claude") + assert stamp is not None and stamp.warmed is False @pytest.mark.asyncio diff --git a/tests/test_litellm/router_strategy/complexity_router/cache_warming/warming_rig.py b/tests/test_litellm/router_strategy/complexity_router/cache_warming/warming_rig.py index d56e6a85efd..d67709a4d5e 100644 --- a/tests/test_litellm/router_strategy/complexity_router/cache_warming/warming_rig.py +++ b/tests/test_litellm/router_strategy/complexity_router/cache_warming/warming_rig.py @@ -21,6 +21,7 @@ from litellm.router_strategy.complexity_router.cache_warming.types import ( CacheWarmingAttribution, CacheWarmingPayload, CacheWarmingRecord, + WarmthStamp, compress_payload, ) from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter @@ -61,6 +62,7 @@ class ReplayRouter(Router): ) self.cache.redis_cache = redis self.completion_calls: list[dict] = [] + self.failed_calls: list[dict] = [] self.anthropic_calls: list[dict] = [] self.replay_delay = replay_delay self.failing_message_marker: str | None = None @@ -76,6 +78,7 @@ class ReplayRouter(Router): async def acompletion(self, **kwargs: object): # pyright: ignore[reportIncompatibleMethodOverride] # test double narrows the overloads marker = self.failing_message_marker if marker is not None and marker in json.dumps(kwargs.get("messages"), default=str): + self.failed_calls.append(kwargs) raise RuntimeError("provider down") self._in_flight += 1 self.max_concurrent = max(self.max_concurrent, self._in_flight) @@ -247,13 +250,15 @@ def seed_session( redis.hashes.setdefault(store.sessions_key(), {})[record_key] = json.dumps(record.model_dump()) redis.zsets.setdefault(store.index_key(), {})[record_key] = time.time() + 3600 for model_group, stamp in (warmth or {}).items(): - redis.data[CacheWarmingStore.warmth_key(record_key, model_group)] = json.dumps(stamp) + redis.data[CacheWarmingStore.warmth_key(record_key, model_group)] = json.dumps( + WarmthStamp(at=stamp, warmed=True).model_dump() + ) return record_key -def warmth_stamp(redis: FakeRedisCache, record_key: str, model_group: str) -> float | None: +def warmth_stamp(redis: FakeRedisCache, record_key: str, model_group: str) -> WarmthStamp | None: raw = redis.data.get(CacheWarmingStore.warmth_key(record_key, model_group)) - return json.loads(raw) if raw is not None else None + return WarmthStamp.model_validate_json(raw) if raw is not None else None def proxy_logging_with_hooks(limiter: object | None = None) -> ProxyLogging: diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index d0458452c4a..c3598e8f408 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -4214,7 +4214,7 @@ class TestWarmAwarePick: ) @staticmethod - def _seed(redis, warmth, last_activity=None, served_model="fast-claude"): + def _seed(redis, warmth, last_activity=None, served_model="fast-claude", warmed=True): import json import time as time_module @@ -4224,6 +4224,7 @@ class TestWarmAwarePick: CacheWarmingAttribution, CacheWarmingPayload, CacheWarmingRecord, + WarmthStamp, compress_payload, ) @@ -4247,7 +4248,9 @@ class TestWarmAwarePick: key = store.record_key("warm-router", "hash-w", "warm-sess") redis.hashes.setdefault(store.sessions_key(), {})[key] = json.dumps(record.model_dump()) for model_group, stamp in warmth.items(): - redis.data[CacheWarmingStore.warmth_key(key, model_group)] = json.dumps(stamp) + redis.data[CacheWarmingStore.warmth_key(key, model_group)] = json.dumps( + WarmthStamp(at=stamp, warmed=warmed).model_dump() + ) @staticmethod def _kwargs(): @@ -4272,6 +4275,20 @@ class TestWarmAwarePick: assert "cold-model" not in picks assert picks <= {"fast-claude", "smart-claude"} + @pytest.mark.asyncio + async def test_a_failed_replay_does_not_make_a_model_look_warm(self, mock_router_instance): + """A replay that the provider refused leaves that model's cache cold, so its stamp must not steer the + pick onto it: the pick falls back to the served model, whose cache the session's own traffic wrote.""" + import time as time_module + + redis = self._fresh_redis() + self._seed(redis, warmth={"smart-claude": time_module.time()}, served_model="fast-claude", warmed=False) + router = self._router(mock_router_instance, redis) + picks = { + await router._pick_model_for_tier(ComplexityTier.SIMPLE, None, None, self._kwargs()) for _ in range(20) + } + assert picks == {"fast-claude"} + @pytest.mark.asyncio async def test_stale_warm_entries_fall_back(self, mock_router_instance): import time as time_module