diff --git a/litellm/router_strategy/complexity_router/cache_warming/refresher.py b/litellm/router_strategy/complexity_router/cache_warming/refresher.py index 30295889288..72082d8623d 100644 --- a/litellm/router_strategy/complexity_router/cache_warming/refresher.py +++ b/litellm/router_strategy/complexity_router/cache_warming/refresher.py @@ -583,6 +583,11 @@ class CacheWarmingRefresher: proxy_logging_obj: "ProxyLogging", lease_lost: asyncio.Event, ) -> None: + """The semaphore is held across the session's whole due set, and decompression happens inside it, so + the concurrency bound also bounds decompressed residency. Every session is started at once, so + inflating above the semaphore let the tick hold one decompressed payload per active session rather + than per replay in flight, which the cap sizes at max_sessions times the uncompressed ceiling. The + replay ceiling is unchanged, since a session's models are replayed in sequence inside the slot.""" warmth = await store.get_warmth(session_key, warm_models) now = time.time() due_models = tuple( @@ -592,9 +597,9 @@ class CacheWarmingRefresher: ) if not due_models: return - payload = decompress_payload(record.payload_compressed) - for model_group in due_models: - async with semaphore: + async with semaphore: + payload = decompress_payload(record.payload_compressed) + for model_group in due_models: if lease_lost.is_set(): return attempted_at = time.time() 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 b46271135fb..f94ca7e63be 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 @@ -396,3 +396,36 @@ def test_scim_deactivation_is_one_predicate_shared_with_the_auth_paths(): assert user_is_scim_deactivated(user({"scim_active": True})) is False assert user_is_scim_deactivated(user({})) is False assert user_is_scim_deactivated(None) is False + + +@pytest.mark.asyncio +async def test_the_concurrency_bound_bounds_decompressed_payloads_not_just_replays(): + """Every session is started at once, so anything a session materializes before acquiring its slot scales + with max_sessions instead of with the concurrency setting. Payloads are held decompressed for the whole + replay, and capture admits them up to eight times the compressed cap, so inflating above the semaphore let + one tick hold a thousand of them. Pins both halves of the bound: replays in flight and payloads inflated.""" + from litellm.router_strategy.complexity_router.cache_warming import refresher as refresher_module + + llm_router, redis = warming_rig(redis=FakeRedisCache(), replay_delay=0.02) + for index in range(6): + seed_session(redis, session_id=f"sess-{index}", caller_scope=f"hash-{index}", user_api_key=f"hash-{index}") + real_decompress = refresher_module.decompress_payload + inflated_before_first_replay_completed: list[int] = [] + inflated = 0 + + def counting_decompress(blob): + nonlocal inflated + inflated += 1 + if not llm_router.completion_calls: + inflated_before_first_replay_completed.append(inflated) + return real_decompress(blob) + + refresher_module.decompress_payload = counting_decompress + try: + await tick(llm_router, active=refresher(max_concurrent_replays=2)) + finally: + refresher_module.decompress_payload = real_decompress + + 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"