From a4a1573e2dc337b27b9044190ed61c8a672162bc Mon Sep 17 00:00:00 2001
From: Tin Chi Lo <tin@berri.ai>
Date: Thu, 30 Jul 2026 01:20:08 -0700
Subject: [PATCH] fix(cache_warming): let the concurrency bound also bound
 decompressed payloads

Every session in a tick is started at once, so anything a session materializes
before acquiring its replay slot scales with max_sessions rather than with the
concurrency setting. The payload was inflated above the semaphore and held for the
whole replay, and capture admits payloads up to eight times the compressed cap, so
one tick could hold a decompressed payload per active session instead of per replay
in flight.

Decompression moves inside the slot, which is now held across the session's whole
due set. The replay ceiling is unchanged, since a session's models are replayed in
sequence inside its slot, and the CPU burst of inflating is now serialized to the
same bound rather than running for every session at once.

The bound had no test at all; the rig tracked peak in-flight replays and nothing
asserted on it. One test now pins both halves through a real tick: replays in
flight and payloads inflated must both respect max_concurrent_replays. It fails on
the previous ordering with the payload count at the session count, not the bound.
---
 .../cache_warming/refresher.py                | 11 +++++--
 .../cache_warming/test_refresher.py           | 33 +++++++++++++++++++
 2 files changed, 41 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 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"