scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

fix(cache_warming): warm the model the session was actually served

Browse source 
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.
This commit is contained in:
Tin Chi Lo 2026-07-30 10:32:39 -07:00
parent bb0d5684f7
commit ed39e60a07
2 changed files with 32 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm/router_strategy/complexity_router/cache_warming/refresher.py
View file

@ -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,

21
tests/test_litellm/router_strategy/complexity_router/cache_warming/test_refresher.py
View file

@ -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"