mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-04 07:27:51 +00:00
fix(rate-limiter): charge token counters of calls that skipped admission without waiting for a window
A proxy sub-call, such as the compaction summary or the safeguards classifier, never passes admission, so the scopes it charges have no open window and the sub-call precheck reads their counters directly. Gating those charges on a window left a key's summary-model TPM limit unenforced; they are now charged as before.
This commit is contained in:
parent
5b94aee8b3
commit
acdcad9769
3 changed files with 47 additions and 3 deletions
|
|
@ -5325,10 +5325,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
reserved_tokens=reserved_tokens,
|
||||
)
|
||||
)
|
||||
unreserved_targets: Final = [target for target in charged_targets if target not in reserved_scopes]
|
||||
# A call that never passed admission, such as a proxy sub-call, opened no window, and the sub-call
|
||||
# precheck reads its counters directly, so it charges them unconditionally
|
||||
pipeline_operations.extend(
|
||||
self._build_open_window_tpm_ops(
|
||||
targets=[target for target in charged_targets if target not in reserved_scopes],
|
||||
tokens=total_tokens,
|
||||
self._build_open_window_tpm_ops(targets=unreserved_targets, tokens=total_tokens)
|
||||
if stash is not None
|
||||
else self._build_reservation_aware_tpm_ops(
|
||||
targets=unreserved_targets,
|
||||
reserved_scopes=reserved_scopes,
|
||||
actual_tokens=total_tokens,
|
||||
reserved_tokens=0,
|
||||
)
|
||||
)
|
||||
pipeline_operations.extend(
|
||||
|
|
|
|||
|
|
@ -140,6 +140,20 @@ def test_caller_without_the_summary_model_sees_the_edit_error(compacting: Rig) -
|
|||
assert len(compacting.backend_calls(received)) == 1
|
||||
|
||||
|
||||
def test_caller_over_the_summary_models_tpm_limit_sees_the_edit_error(compacting: Rig) -> None:
|
||||
with compacting.gateway.scenario() as scenario:
|
||||
key: Final = scenario.key(model_tpm_limit={CLASSIFIER_MODEL: 1})
|
||||
within_nonce: Final = nonce()
|
||||
compacting.register(within_nonce, chat_backend(tool_id(within_nonce), stream=False), summarizing_or_classifying)
|
||||
within: Final = applied_edit(json_body(compacting.post(long_request(within_nonce), key=key)))
|
||||
assert within["type"] == "compact_20260112" and "error" not in within, within
|
||||
over_nonce: Final = nonce()
|
||||
compacting.register(over_nonce, chat_backend(tool_id(over_nonce), stream=False), summarizing_or_classifying)
|
||||
over: Final = json_body(compacting.post(long_request(over_nonce), key=key))
|
||||
assert applied_edit(over)["error"] == "summary_model_rate_limit_exceeded", over
|
||||
assert compacting.classifier_calls(compacting.received(over_nonce)) == ()
|
||||
|
||||
|
||||
def assert_summarized_and_classified_on_their_models(rig: Rig, request_nonce: str, body: dict[str, JsonValue]) -> str:
|
||||
received: Final = rig.received(request_nonce)
|
||||
targets: Final = [request.target for request in received]
|
||||
|
|
|
|||
|
|
@ -8842,3 +8842,26 @@ async def test_success_charges_a_numeric_body_user_to_the_window_its_admission_o
|
|||
)
|
||||
|
||||
assert _charged_windows(handler, _scoped_success_kwargs(user=123))["{end_user:123}:tokens"] == "{end_user:123}:window"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("parent_admitted", [True, False], ids=["sub_call_of_an_admitted_request", "no_admission"])
|
||||
async def test_success_of_a_call_that_skipped_admission_charges_every_scope_without_waiting_for_a_window(
|
||||
parent_admitted: bool,
|
||||
) -> None:
|
||||
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()))
|
||||
if parent_admitted:
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t"),
|
||||
cache=handler.internal_usage_cache.dual_cache,
|
||||
data={
|
||||
"model": "gpt-4o-mini",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"litellm_call_id": "parent",
|
||||
},
|
||||
call_type="acompletion",
|
||||
)
|
||||
|
||||
assert _charged_windows(handler, _scoped_success_kwargs(litellm_call_id="sub-call")) == dict.fromkeys(
|
||||
_SCOPE_TOKEN_KEYS
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue