mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
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:
parent
f5014e99a9
commit
a4a1573e2d
2 changed files with 41 additions and 3 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue