BerriAI/litellm PR #37539 — Enqueued-token rate limiting for batch submissions — PR #37539
BerriAI/litellm · pull request #37539 ·
Transcript
PlainEnglish
Welcome to PR 37539 for LiteLLM. This pull request introduces opt-in enqueued-token rate limiting for batch submissions, changing how the proxy governs long-running batch workloads. Let's walk through the change together.
PlainEnglish
The PR pursues two primary objectives. First, it implements a new enqueued-token reservation system for batch submissions, with automatic refunds when batches reach terminal states. Second, it ensures that only proxy admins can write the new limit, preventing self-service quota escalation. Along the way, it adds post-call hooks for refund logic and maintains dual-backend semantics across Redis and in-memory stores. Importantly, this is opt-in — the feature won't touch existing batch rate limiting unless an admin sets the metadata.
PlainEnglish
The proxy currently enforces RPM and TPM limits on batch submissions using per-minute rolling windows. That's a poor fit for batch workloads that can take hours or even days to complete. This PR introduces an opt-in alternative where administrators set a long-lived enqueued-token allowance on a key or team. When a batch is submitted, tokens are reserved against that allowance. When the batch finishes, fails, expires, or is cancelled, the reservation is refunded. This is safer for batch workloads and prevents queue-based starvation.
Architecture
We start by adding two configuration constants to litellm constants dot py. The first is a TTL of eight days, which ensures that even if a proxy crashes, a reservation won't leak tokens indefinitely. The second is the metadata field name itself, so it's centralized and consistent across the codebase.
Architecture
The heart of the change is a new module called batch enqueued tokens dot py. It defines a store class that manages per-scope counters — one for the API key, one for the team. The store offers a dual backend: Redis with one-key-at-a-time Lua scripts for cluster safety, or an in-memory fallback with asyncio locking. The interface is straightforward. Reserve tokens at submission, rolling back partial reserves on failure. Save the reservation keyed by batch ID with a clamped TTL. Pop the reservation when the batch reaches a terminal state. And refund the tokens back to the counters.
Architecture
In the batch rate limiter pre-call hook, we resolve any enqueued scopes from the user metadata. If scopes exist, we skip the standard per-minute rate limit checks and instead call a new method that reserves tokens against the enqueued allowance. If the reservation succeeds, the request proceeds. If it fails, we raise a four twenty-nine with a detailed error message that includes the remaining capacity. This is the fork point where the new allowance replaces the old rolling-window logic.
Architecture
Here's how the reservation flows through the request lifecycle. The pre-call hook reserves tokens and stashes the reservation. If the batch is successfully created, the success hook saves the reservation to Redis keyed by the batch ID with a clamped TTL. If the batch status is already terminal — completed, failed, expired, or cancelled — the hook immediately pops the reservation and refunds it. If the request fails before reaching the provider, the failure hook refunds the stashed reservation so tokens aren't leaked. This covers all the paths.
Architecture
Security is gated by a new enforcement function in auth utils. It compares the requested metadata value to the stored value. If they differ, the function checks that the user has the proxy admin role. If not, it raises a four oh three. If the values match, the request is allowed — that's idempotent resubmission, which keeps forms from breaking. This enforcement is called in every key and team mutation path: generation, update, regeneration, and bulk updates.
Architecture
Let's zoom out and see where the change lives. The core store module is a brand-new file in proxy hooks. The pre-call reservation logic goes into the batch rate limiter. The post-call save and refund hooks live in parallel request limiter v3. Admin enforcement is added to auth utils, and the gate is called from key and team management endpoints. The test coverage spans unit tests for the store, integration tests for the rate limiter, post-call hook tests, admin-gate tests, and end-to-end acceptance tests.
PlainEnglish
So what's the outcome? The proxy now offers administrators an opt-in batch enqueued token limit metadata field for keys and teams. When set, batch submissions reserve their estimated token count against the allowance instead of checking per-minute RPM or TPM. The reservation is refunded when the batch completes, fails, is cancelled, or expires. The implementation dual-backs Redis for distributed deployments and in-memory storage, with automatic TTL-based cleanup to prevent token leaks in a crash scenario. The feature is gated to proxy admins only, so non-admins cannot write the field and escalate their own quota.
CodeQuality
Three risks were identified during the review. The first is the most subtle: the pop operation handles an empty string tombstone to prevent double-refund of stale local records, but the invariant is delicate. The second is a TTL clamping edge case where elapsed time can exceed the configured TTL due to clock drift or test injection, resulting in a one-second TTL. And the third is that the post-call hook calls batch response view on every response, relying on defensive shape validation to skip non-batch calls. Let's look at each one.
CodeQuality
The first gotcha is the pop reservation method. When the Redis record is an empty string, that's a tombstone left by a failed save. The method returns None to signal that there's nothing to refund. The local ghost from the failed save can still be in memory cache, but pop local record deletes it only on a successful pop. This prevents double-refund of stale local records. The comment on line thirteen accurately explains the invariant, and the test coverage at line two seventy-eight through two ninety-seven of the test file confirms the behavior. It's correct design, just subtle.
CodeQuality
The second gotcha is TTL clamping. When saving a reservation to Redis, the code calculates the remaining TTL by subtracting elapsed time from the configured TTL. It then clamps the value to a minimum of one second so Redis EXPIRE gets a positive integer. If elapsed time exceeds the configured TTL — due to monotonic clock drift or test injection — the TTL becomes one second. This bounds the record lifetime but doesn't prevent it from being too short. No observed issue in tests, but it's worth noting as a latency-sensitive edge case.
CodeQuality
The third gotcha is in the post-call hook. The hook calls batch response view on every response, regardless of call type. If the provider returns a shape missing the required fields, the view returns None and the hook returns early. Non-batch objects like chat completions correctly leave the stash untouched, as confirmed by tests at line fifty-nine ninety-three through sixty oh eight. However, there's no active validation that responses come from the right provider. In practice, batch response view enforces a literal batch string on the object field, so the shape is unambiguous. A chat completion with an object field can't match because the literal enforces the exact string. Still, the open question is whether there's a cleaner way to skip this for non-batch calls to avoid the validation overhead.
PlainEnglish
This is a thoughtful, well-tested change that brings long-lived enqueued-token rate limiting to batch submissions. The admin gates are solid, the dual-backend semantics are preserved, and the test coverage is comprehensive. The gotchas are all documented and understood, with the post-call hook question being the only open design consideration. I'd recommend approving with a note on that question. Thanks for watching, and happy reviewing.
How this was made
Lenzon read BerriAI/litellm at pull request #37539 and generated this walkthrough automatically. The narration above is the transcript of what it says.
Explain a pull request from your own repo
Point Lenzon at a repo or a pull request and get a narrated walkthrough like this one.
Try it