feat(proxy): add project-level model_itpm_limit/model_otpm_limit - #35098
feat(proxy): add project-level model_itpm_limit/model_otpm_limit#35098shivijain2323 wants to merge 1 commit into
Conversation
Bedrock Mantle bills input and output tokens separately, but project-level quotas only supported a single combined model_tpm_limit. Adds model_itpm_limit/model_otpm_limit on the project, enforced via the existing project rate-limit hook's reservation/reconciliation path alongside (not instead of) the combined TPM limit, mirroring how the deployment-level itpm/otpm check already handles that overlap.
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| itpm_scopes = self._get_reserved_itpm_scopes_from_kwargs(kwargs=request_data) | ||
| pipeline_operations.extend( | ||
| self._build_reservation_aware_tpm_ops( | ||
| targets=list(itpm_scopes), | ||
| reserved_scopes=itpm_scopes, | ||
| actual_tokens=0, | ||
| reserved_tokens=itpm_reserved, | ||
| ) | ||
| ) | ||
| if otpm_reserved > 0: | ||
| otpm_scopes = self._get_reserved_otpm_scopes_from_kwargs(kwargs=request_data) | ||
| pipeline_operations.extend( | ||
| self._build_reservation_aware_tpm_ops( | ||
| targets=list(otpm_scopes), | ||
| reserved_scopes=otpm_scopes, | ||
| actual_tokens=0, | ||
| reserved_tokens=otpm_reserved, |
There was a problem hiding this comment.
When a client disconnects after provider dispatch but before a partial-success event is available, this cleanup reconciles both reservations with actual_tokens=0, fully refunding provider-billable input and output usage. Repeated early stream cancellations can therefore evade the project's ITPM and OTPM quotas.
How this was verified: The streaming cleanup invokes this fallback only when no success event owns reconciliation, and both counter updates subtract the full reservation as zero actual usage.
Rule Used: What: Fail any PR which may contains a security in... (source)
Knowledge Base Used: Proxy Server Request Flow
| dedicated = { | ||
| k: v | ||
| for k, v in { | ||
| "model_itpm_limit": self.model_itpm_limit, | ||
| "model_otpm_limit": self.model_otpm_limit, | ||
| "model_rpm_limit": self.model_rpm_limit, | ||
| "model_tpm_limit": self.model_tpm_limit, | ||
| }.items() | ||
| if v is not None |
There was a problem hiding this comment.
When an existing project stores a model TPM or RPM limit in metadata while its dedicated JSON column has the database default {}, this merge treats the empty object as an explicit value and overwrites the configured limit. The limiter then finds no matching model descriptor, so the existing project quota is silently disabled.
Knowledge Base Used:
Greptile SummaryAdds project-level input/output token quotas and integrates their persistence, authentication metadata, reservation, reconciliation, Responses API handling, and tests
Confidence Score: 2/5This PR should not merge until disconnect accounting preserves billable usage and empty dedicated columns no longer disable existing project limits The new limiter can remove reservations for dispatched requests on an early stream disconnect, and the metadata merge can replace configured legacy quotas with empty database defaults Files Needing Attention: litellm/proxy/hooks/parallel_request_limiter_v3.py, litellm/models/project.py
|
| Filename | Overview |
|---|---|
| litellm/proxy/hooks/parallel_request_limiter_v3.py | Adds separate ITPM/OTPM reservation and reconciliation, but the disconnect fallback can refund already-consumed usage |
| litellm/models/project.py | Adds merged project rate-limit metadata, but empty database-default columns can erase legacy metadata limits |
| litellm/proxy/auth/user_api_key_auth.py | Routes project metadata through the new merged view across authentication paths |
| litellm/proxy/_types.py | Adds project ITPM/OTPM fields to management request contracts and metadata-field routing |
| litellm-proxy-extras/litellm_proxy_extras/migrations/20260722200000_add_project_itpm_otpm_limit/migration.sql | Adds non-null JSONB columns for project input/output token limits |
| tests/test_litellm/proxy/hooks/test_tpm_concurrent.py | Substantially expands reservation, reconciliation, concurrency, and Responses API coverage |
Reviews (1): Last reviewed commit: "feat(proxy): add project-level model_itp..." | Re-trigger Greptile
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
| }.items() | ||
| if v is not None | ||
| } | ||
| return {**base, **dedicated} |
There was a problem hiding this comment.
Medium: Empty columns disable configured project limits
NewProjectRequest and UpdateProjectRequest move these fields into metadata, but every dedicated database column has a non-null {} default. After loading the row, this merge therefore replaces configured metadata with {}, disabling the new ITPM/OTPM limits and existing per-model RPM/TPM limits; a project-key holder can then exceed those quotas. Persist management fields into the dedicated columns, or prevent empty defaults from shadowing legacy metadata and backfill existing rows.
| explicit_max_tokens = next( | ||
| ( | ||
| value | ||
| for value in (data.get("max_tokens"), data.get("max_completion_tokens"), data.get("max_output_tokens")) |
There was a problem hiding this comment.
Medium: Conflicting Responses API caps bypass output reservation
A Responses caller can send max_tokens: 0 together with a large max_output_tokens. This selects zero for reservation and suppresses the implicit cap, while the Responses transformation only forwards max_output_tokens, allowing concurrent requests to generate output without reserving the corresponding TPM/OTPM budget. Select the cap canonical for the call type and reject conflicting output-cap fields.
| self._build_reservation_aware_tpm_ops( | ||
| targets=list(itpm_scopes), | ||
| reserved_scopes=itpm_scopes, | ||
| actual_tokens=0, |
There was a problem hiding this comment.
Low: Stream cancellation refunds processed tokens
This treats a disconnected request as consuming zero input and output tokens even after its prompt reached the provider or partial output was delivered. A caller can repeatedly submit large streaming requests and disconnect before the usage callback, recovering the full ITPM/OTPM reservation each time. Reconcile recovered partial usage when available and retain the reservation when usage is unknown rather than refunding it as zero.
PR overviewThis pull request adds project-level per-model input-token-per-minute (ITPM) and output-token-per-minute (OTPM) limits to the proxy. It updates project configuration storage and the parallel request limiter, including streaming and Responses API accounting. Three issues remain open, and none have yet been addressed. Project limits can be shadowed by empty database values, conflicting Responses API token caps can avoid output-budget reservation, and stream cancellation can incorrectly refund consumed capacity. These behaviors allow project-key holders to exceed configured token quotas, though the impact is limited to quota and resource-accounting enforcement. Open issues (3)
Fixed/addressed: 0 · PR risk: 6/10 |
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
1 participant
TLDR
Problem this solves:
How it solves it:
Relevant issues
Linear ticket
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaito re-request a review after pushing changes)Delays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
Screenshots / Proof of Fix
Type
🆕 New Feature
Changes
We already support model_tpm_limit/model_rpm_limit per project, but that treats input and output tokens as one combined bucket, which doesn't match how Bedrock Mantle actually meters usage. This adds model_itpm_limit/model_otpm_limit to LiteLLM_ProjectTable, NewProjectRequest, and UpdateProjectRequest, routed into project metadata the same way model_tpm_limit already is, with a matching Prisma migration.
There's already deployment-level itpm/otpm support in router_utils/pre_call_checks/io_token_rate_limit_check.py, but that code only ever deals with one deployment at a time, whereas the proxy hook (parallel_request_limiter_v3.py) has to juggle key/team/org/project limits together on the same request. Rather than bolt a second copy of the reservation logic on, this adds project-scoped ITPM/OTPM as two more descriptors into the same reservation path the hook already uses for TPM/RPM. If a project sets both the combined TPM limit and the new itpm/otpm limits on the same model, both get enforced, not one overriding the other, matching what the deployment-level check already does. Input and output are reserved separately up front, and if one side reserves fine but the other is over limit, the side that already went through gets rolled back so nothing is left over-counted. On success, cached prompt tokens get excluded from the ITPM count since that's how Bedrock bills it, but that's purely for the rate limit math; cost and usage logging still see the full token count.
Adds get_project_model_itpm_limit/get_project_model_otpm_limit to auth_utils.py for parity with the existing key/team/project rate-limit accessors.
Also fixes several correctness issues in the Responses API path of parallel_request_limiter_v3.py surfaced during review: the combined-TPM and OTPM output-cap checks classified any request with data["input"] set as an embedding, which also misclassifies the Responses API (it puts its prompt in "input" too, but does generate output) and skipped the output cap entirely for it; the implicit output cap was being written to data["max_tokens"], which the Responses-to-chat-completion transformation ignores (it only reads max_output_tokens), so the cap was silently dropped before provider dispatch; and an explicit max_output_tokens=0 was getting folded away by a truthy
orchain and then re-floored to 1 by the OTPM reservation, which could false-reject a genuine zero-output request against an already-exhausted OTPM bucket. Also routes Responses API input through the standard transform_responses_api_input_to_messages helper before token counting, since token_counter's text argument only joins plain strings in a list and was silently dropping input_image content blocks from the ITPM estimate.Out of scope for this PR: key/team/org-level model_itpm_limit/model_otpm_limit, and Admin UI support (only the auto-generated schema.d.ts types are updated; no form fields yet).
Final Attestation
Docs PR BerriAI/litellm-docs#638