From 7fc47bac5a5eb7c1dea342ac94cdb16f6b16771a Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 30 Jul 2026 10:32:39 -0700 Subject: [PATCH] fix(cache_warming): warm the model the session was actually served A tier may be a pool the router picks from at random, but the warm set resolved to one representative per tier, so for a pooled tier the member holding the session's cache was usually not warmed at all. That inverted the feature: the session's own cache expired at the provider TTL and its next turn on the same tier paid a full cache write, while warming spent on a pool member the session had never touched and might never be routed to. The warm set is now per session, leading with the record's served_model and then the tier representatives, so a session's own cache is always refreshed. Eligibility is still resolved once per tick over the union, so no extra model-list lookups happen per session. Single-model tiers, which every existing test and the live proof used, are unaffected: the served model is the tier representative there, which is why the gap did not surface earlier. --- .../cache_warming/refresher.py | 14 ++++++++++--- .../cache_warming/test_refresher.py | 21 +++++++++++++++++++ 2 files changed, 32 insertions(+), 3 deletions(-) diff --git a/litellm/router_strategy/complexity_router/cache_warming/refresher.py b/litellm/router_strategy/complexity_router/cache_warming/refresher.py index 72082d8623d..e302dbc6cb2 100644 --- a/litellm/router_strategy/complexity_router/cache_warming/refresher.py +++ b/litellm/router_strategy/complexity_router/cache_warming/refresher.py @@ -516,8 +516,14 @@ class CacheWarmingRefresher: ) if not active: return - warm_models = filter_cache_warmable(llm_router, resolve_warm_models(complexity_router.config)) - if not warm_models: + tier_models = resolve_warm_models(complexity_router.config) + warmable = frozenset( + filter_cache_warmable( + llm_router, + tuple(dict.fromkeys((*tier_models, *(record.served_model for _, record in active)))), + ) + ) + if not warmable: return attributed = frozenset( record.attribution.user_api_key @@ -547,7 +553,9 @@ class CacheWarmingRefresher: store=store, session_key=key, record=record, - warm_models=warm_models, + warm_models=tuple( + model for model in dict.fromkeys((record.served_model, *tier_models)) if model in warmable + ), refresh_interval_seconds=config.refresh_interval_seconds, session_ttl_seconds=config.session_ttl_seconds, semaphore=semaphore, 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 d3ed7f7f9ff..42cbc552c99 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 @@ -423,3 +423,24 @@ async def test_the_concurrency_bound_bounds_decompressed_payloads_not_just_repla assert len(llm_router.completion_calls) == 12, "every seeded session should warm both due models" assert llm_router.max_concurrent <= 2, "replays in flight must respect the bound" assert max(inflated_before_first_replay_completed) <= 2, "payloads inflated must respect the same bound" + + +@pytest.mark.asyncio +async def test_a_pooled_tier_warms_the_model_the_session_was_actually_served(): + """A tier may be a pool that the router picks from at random, so the member holding this session's cache + is not necessarily the one the tier resolves to for warming. Warming the tier representative alone left + the session's own cache to expire, which is the case the feature exists to prevent, while spending on a + member the session never touched.""" + pool_list = [ + {"model_name": "haiku-a", "litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "sk-t"}}, + {"model_name": "haiku-b", "litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "sk-t"}}, + {"model_name": "sonnet-a", "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-t"}}, + ] + llm_router, redis = warming_rig( + redis=FakeRedisCache(), + tiers={"SIMPLE": ["haiku-a", "haiku-b"], "COMPLEX": ["sonnet-a"]}, + model_list=pool_list, + ) + seed_session(redis, served_model="haiku-b") + await tick(llm_router) + assert "haiku-b" in replayed_models(llm_router), "the session's own cache must be kept warm"