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.
This commit is contained in:
Tin Chi Lo 2026-07-30 01:20:08 -07:00
parent f5014e99a9
commit a4a1573e2d
2 changed files with 41 additions and 3 deletions

View file

@ -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()

View file

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