feat(ptu): add PTU reservation table and admin CRUD endpoints - #33130
feat(ptu): add PTU reservation table and admin CRUD endpoints#33130yucheng-berri wants to merge 4 commits into
Conversation
|
|
Greptile SummaryThis PR adds a PTU (Provisioned Throughput Unit) reservation table and four admin-only CRUD endpoints as a storage foundation for a future flat-cost attribution feature. The feature is gated by
Confidence Score: 5/5Safe to merge — purely additive behind a default-off feature flag with no changes to the hot request path. The change is strictly additive: new table, new repository, new endpoints, all gated by a default-off flag. The datetime timezone comparison issue flagged in the previous round has been properly fixed with idempotent UTC coercion at the Pydantic validation layer in both the request types and the domain model, and three targeted regression tests confirm the fix. No existing code paths are modified. No files require special attention — all changed files are new additions or trivial router registrations.
|
| Filename | Overview |
|---|---|
| litellm/types/proxy/management_endpoints/ptu_reservation.py | Request/response types with field validators that coerce naive datetimes to UTC at the Pydantic layer — the P1 fix from the prior review round. |
| litellm/models/ptu_reservation.py | Domain model with UTC-coercion validator and business-rule enforcement (cost_source invariants, effective_to > effective_from). Clean. |
| litellm/proxy/management_endpoints/ptu_reservation_endpoints.py | Four admin-only endpoints with feature-flag + role gates. Timezone comparison in close() is now safe after pydantic-layer UTC coercion fix. |
| litellm/repositories/ptu_reservation_repository.py | Repository with find_active and find_overlapping using correct half-open interval semantics for datetime window overlap detection. |
| litellm-proxy-extras/litellm_proxy_extras/migrations/20260713000000_add_litellm_ptu_reservation_table/migration.sql | Additive migration creating LiteLLM_PTUReservation with three correct indexes; uses TIMESTAMP(3) consistently with Prisma conventions. |
| tests/test_litellm/proxy/management_endpoints/test_ptu_reservation_endpoints.py | 28 mock-only tests covering create, list, info, close, flag-gating, RBAC, overlap detection, and three new regression tests for the naive-datetime TZ fix. |
Reviews (2): Last reviewed commit: "fix(ptu): coerce naive datetimes on rese..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| close_at = body.effective_to or datetime.now(timezone.utc) | ||
| if close_at <= row.effective_from: |
There was a problem hiding this comment.
body.effective_to is parsed by Pydantic as a naive datetime when the caller omits timezone info (e.g., "2026-08-15T00:00:00" with no Z/+00:00). row.effective_from returned by Prisma is a UTC-aware datetime. Python raises TypeError: can't compare offset-naive and offset-aware datetimes on the <= comparison, turning a valid-looking close request into a 500. The same naive input is then written to effective_to in the DB, storing a timezone-stripped value. Normalising close_at to UTC before the comparison (and before the update) fixes both: replace close_at = body.effective_to or datetime.now(timezone.utc) with an explicit replace(tzinfo=timezone.utc) guard on naive inputs.
| repo = PTUReservationRepository(prisma_client) | ||
| overlapping = await repo.find_overlapping( | ||
| team_id=validated_reservation["team_id"], | ||
| model=validated_reservation["model"], | ||
| effective_from=validated_reservation["effective_from"], | ||
| effective_to=validated_reservation["effective_to"], | ||
| ) | ||
| if overlapping: | ||
| raise HTTPException( | ||
| status_code=409, | ||
| detail={ | ||
| "error": "reservation overlaps existing active reservation(s) for the same (team, model)", | ||
| "overlapping_ids": [r.id for r in overlapping], | ||
| }, | ||
| ) | ||
|
|
||
| actor = user_api_key_dict.user_id or "admin" | ||
| create_data = { | ||
| **{k: v for k, v in validated_reservation.items() if k != "id" and v is not None}, | ||
| "created_by": actor, | ||
| "updated_by": actor, | ||
| } | ||
| return await repo.table.create(data=create_data) |
There was a problem hiding this comment.
find_overlapping and table.create are two separate round-trips to the DB with no transaction wrapping them. Two concurrent admin requests for the same (team, model) window will both pass the overlap check (since neither row exists yet) and both succeed, silently creating duplicate reservations. The fix requires either a serializable transaction around the check+insert pair or a unique partial index on (team_id, model) for open-ended rows. Low-risk given admin-only traffic, but worth addressing before the later stages start reading these rows.
| where: dict = {} | ||
| if team_id is not None: | ||
| where["team_id"] = team_id | ||
| if model is not None: | ||
| where["model"] = model | ||
| return await repo.table.find_many(where=where) |
There was a problem hiding this comment.
repo.table.find_many(where=where) with no take limit returns every matching row in a single query. As historical reservations accumulate over many teams and models, this will return unbounded result sets. The BaseRepository.find_many already accepts take/skip parameters; threading a limit query param through here (defaulting to e.g. 1 000) would be consistent with how other list endpoints in the repo behave.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
|
@greptileai please re-review — pushed P1 fix for the datetime-tz mismatch. Both /new and /close now coerce naive datetimes to UTC at the pydantic layer (field validators on both the request types and the domain model); 3 regression tests added that fail without the fix. |
|
bugbot run |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Want reviews to match your repository better? Bugbot Learning can learn team-specific rules from PR activity. A team admin can enable Learning in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 2a4b5a1. Configure here.
| ) | ||
|
|
||
| close_at = body.effective_to or datetime.now(timezone.utc) | ||
| if close_at <= row.effective_from: |
There was a problem hiding this comment.
Close compares naive and aware datetimes
High Severity
The close_ptu_reservation endpoint can return a 500 error due to a TypeError. This occurs because close_at (a timezone-aware datetime) is compared directly with row.effective_from, which is a naive datetime loaded from Prisma, preventing reservations from being closed.
Reviewed by Cursor Bugbot for commit 2a4b5a1. Configure here.
9e82ea4 to
5105676
Compare
Adds a new LiteLLM_PTUReservation table storing admin-registered PTU reservations for a (team, model) pair over a time window, plus admin CRUD endpoints (new, list, info, close). Feature-gated by enable_ptu_cost_attribution in general_settings; default off. Stage 1 storage + endpoints only. No spend behavior change: no daily rollup job, no writes to LiteLLM_DailyTeamSpend, no impact on the per-request cost tracking pipeline. - schema.prisma + migration for LiteLLM_PTUReservation - Pydantic domain models in litellm/models/ptu_reservation.py - Repository in litellm/repositories/ptu_reservation_repository.py - Endpoints in litellm/proxy/management_endpoints/ptu_reservation_endpoints.py - Request/response types under litellm/types/proxy/management_endpoints/ - 28 unit tests covering validation, feature-flag gate, admin-only auth, overlap detection, close/create semantics - Routes listed under management_routes in litellm/proxy/_types.py - Router mounted in proxy_server.py alongside budget_management_router
Pydantic accepted naive datetimes on effective_from and effective_to; comparing those against Prisma's UTC-aware row values raised TypeError, turning /ptu_reservation/close into a 500 on inputs like '2026-08-15T00:00:00'. Adds a field validator on both the request types and the domain model that stamps missing tzinfo as UTC and converts tz-aware values into UTC. Three regression tests pin the new contract (create with naive input, close with naive input, close with naive input before effective_from).
Three CI failures on the stage 1 branch: 1. schema_migration_check drifted: the LiteLLM_PTUReservation CREATE TABLE omitted DEFAULT CURRENT_TIMESTAMP on updated_at, so prisma migrate diff reported the column as needing an ALTER. Added the default to match every other table's convention. 2. proxy-infra test_gateway_plus_backend_covers_full_app: the four new /ptu_reservation/* routes weren't listed on either component's allowlist. Added /ptu_reservation/ as a backend prefix (admin API, never data-plane). 3. Verify schema.d.ts matches the proxy OpenAPI spec: adding endpoints changed the spec so the dashboard's generated types were stale. Regenerated with npm run gen:api.
5105676 to
da4e195
Compare
|
Closing as superseded. This is part of the v1 PTU design, which stored PTU config in a separate reservation table. The shipped design puts that config on the model deployment instead, merged as #35341, #35343, #35391, #35393 and #36829. This one has no successor: the reservation table was dropped entirely rather than reimplemented, so everything it added is obsolete. The branch is kept, so nothing here is lost. |
2 participants
Relevant issues
Linear ticket
Resolves LIT-1697 (stage 1 of 5)
Pre-Submission checklist
Screenshots / Proof of Fix
Feature gate is off by default so the endpoints 403 out of the box. Enable in general_settings, then
Type
New Feature
Changes
Adds admin-registered PTU (Provisioned Throughput Unit) reservations as a storage and CRUD-only foundation. A reservation records that a specific team owns N PTUs of a specific model at $X/PTU/month over a time window
[effective_from, effective_to). Later stages will read these rows and attribute prorated flat cost toLiteLLM_DailyTeamSpendand the Usage page; this PR intentionally does none of that.Feature-gated by
enable_ptu_cost_attributioningeneral_settings, default off. Every endpoint returns 403 when the flag is off; when the flag is on, every endpoint requires PROXY_ADMIN.Endpoints:
POST /ptu_reservation/new-> create a reservation. Overlapping active reservations for the same(team, model)are rejected with 409GET /ptu_reservation/list-> optionalteam_id,model,active_onlyfiltersGET /ptu_reservation/info?id=X-> single rowPOST /ptu_reservation/close-> seteffective_to; refuses to close an already-closed reservation or one whose neweffective_tois at or beforeeffective_from. Never deletes rowsEditing a reservation is modeled as closing the existing row and creating a new one. Historical rows stay queryable for audit.
The schema carries
cost_sourceandazure_resource_idcolumns so a future automated Azure Cost Management integration can land without a schema migration;cost_source="azure_billing"requests are rejected in this release.Deviation from the admin-entity pattern
Every other admin CRUD entity in the codebase ships a
/<entity>/updateroute (budget, team, tag, organization, model, customer). This PR intentionally does not. Every cost-affecting field on a reservation (team_id,model,ptu_count,cost_per_ptu,cost_source,azure_resource_id,effective_from) is history-load-bearing: mutating it in place would retroactively rewrite the flat cost the daily rollup (stage 2) has already attributed for past days, breaking the audit trail. The correct edit flow is/closethe old row and/newa fresh one; the resulting two-row history is the feature, not the tradeoff.effective_tois already mutable through/close.updated_by/updated_atare server-managed. That leaves nothing meaningful for/updateto touch, so shipping one would either always 400 or be a rename of/close.Behavior changes
None. This change is strictly additive.
_PROXY_track_cost_callback,db_spend_update_writer.update_database,_batch_database_updates,_commit_spend_updates_to_db,common_daily_activity.py,LiteLLM_TeamTable.spend, budget enforcement are all untouchedLiteLLM_DailyTeamSpendBoilerplate & extensibility note
The new files closely mirror the shape of
LiteLLM_BudgetTableand its friends (model inlitellm/models/, repository inlitellm/repositories/, endpoints inlitellm/proxy/management_endpoints/, request types underlitellm/types/proxy/management_endpoints/). Every new admin-owned CRUD entity in this codebase reproduces roughly the same scaffold (endpoint auth/flag/prisma/repo pattern, near-identical repository properties, form-shaped pydantic requests). Extracting a generic admin-entity router and repository base would remove several hundred lines per entity but touches Budget, Team, Organization, and Access Group at once, so it belongs in a separate refactor rather than inside a feature PR. Filing follow-up.Files changed
schema.prisma: newLiteLLM_PTUReservationmodel with three indexeslitellm-proxy-extras/litellm_proxy_extras/migrations/20260713000000_add_litellm_ptu_reservation_table/migration.sql: DDLlitellm/models/ptu_reservation.py:LiteLLM_PTUReservation+LiteLLM_PTUReservationFullwith pydantic validator enforcing manual/azure_billing invariants andeffective_to > effective_fromlitellm/repositories/ptu_reservation_repository.py:PTUReservationRepositorywithfind_activeandfind_overlapping(half-open interval semantics)litellm/proxy/management_endpoints/ptu_reservation_endpoints.py: four endpoints, feature-flag + admin-role gateslitellm/types/proxy/management_endpoints/ptu_reservation.py: request/response typeslitellm/proxy/_types.py: routes added tomanagement_routesallowlistlitellm/proxy/proxy_server.py: router import +app.include_routertests/test_litellm/proxy/management_endpoints/test_ptu_reservation_endpoints.py: 28 testsTests
28 tests. Suite covers:
cost_source="manual"ptu_count, withoutcost_per_ptu, with non-positiveptu_count, with negativecost_per_ptueffective_to <= effective_from(both strictly-less and equal)cost_source="azure_billing"in this releaseoverlapping_idsactive_onlyeffective_to(default: now UTC; explicit override respected)effective_to <= effective_fromNote
Medium Risk
Adds a DB migration and new admin-only APIs behind a feature flag; spend attribution is not wired yet, so runtime LLM cost behavior is unchanged when the flag is off.
Overview
Introduces admin-managed PTU reservations—time-bounded
(team, model)prepaid capacity with manualptu_countandcost_per_ptu—as storage and API only (stage 1; no spend rollup or usage UI yet).Adds
LiteLLM_PTUReservation(Prisma + migration), Pydantic models with manual vsazure_billingfield rules, and a repository withfind_active/find_overlappinghalf-open windows. Four endpoints (/new,/list,/info,/close) are wired on the proxy, listed in management routes and the UI backend allowlist (/ptu_reservation/), and documented in OpenAPI types. No/update—changes are modeled as close + new for auditability.Access is gated by
enable_ptu_cost_attribution(default off) and PROXY_ADMIN. Creating overlapping reservations for the same team/model returns 409;azure_billingis rejected for this release. 28 endpoint tests cover validation, overlap, auth, and the feature flag.Reviewed by Cursor Bugbot for commit 2a4b5a1. Bugbot is set up for automated code reviews on this repo. Configure here.