Skip to content

feat(auto-router)!: scope shadow eval jobs to multiple keys - #37251

Merged
tin-berri merged 9 commits into
litellm_internal_stagingfrom
litellm_shadoweval_multikey2
Aug 19, 2026
Merged

feat(auto-router)!: scope shadow eval jobs to multiple keys#37251
tin-berri merged 9 commits into
litellm_internal_stagingfrom
litellm_shadoweval_multikey2

Conversation

@tin-berri

@tin-berri tin-berri commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

TLDR

Problem this solves:

  • A shadow eval job could cover exactly one key
  • Evaluating a router across several keys meant several jobs
  • Their results had to be added up by hand

How it solves it:

  • One job now takes a list of keys
  • Each key keeps its own turn budget and stop state
  • Results come back pooled and per key
  • The job row was already the per-key unit, so a group_id column ties sibling rows into one job; the sampler is untouched

User Flow

Before: an admin shadow testing an auto-router across two of their keys can only cover one key at a time, so no single result describes the router

  1. They send POST http://localhost:4000/auto_router/shadow_eval/start with "api_key_ids": ["<hash A>", "<hash B>"] and get back 422 saying api_key_id is required
  2. They resend the request twice, one key each time, and get back two separate job ids
  3. They send GET http://localhost:4000/auto_router/shadow_eval/ for each id and get two result sets, each describing one key
  4. They add the two up by hand to judge the router, and stopping the experiment takes one stop call per job
  5. They open http://localhost:4000/ui/?page=cost-optimization and find a key picker that accepts one key, so the same split shows up on screen

After: the same admin covers both keys with one job and reads one result

  1. They send POST http://localhost:4000/auto_router/shadow_eval/start with "api_key_ids": ["<hash A>", "<hash B>"] and get back 201 with a single job id listing both keys, each with its own max_turns and display labels
  2. Traffic on either key is sampled, and each key spends its own budget, so two keys at a budget of 2 judge 4 turns rather than 2, and one key exhausting its budget never ends sampling for the other
  3. They send GET http://localhost:4000/auto_router/shadow_eval/ and get one pooled result plus a per key breakdown alongside the existing tier and model ones
  4. They send GET http://localhost:4000/auto_router/shadow_eval?api_key_id= and get every job that key belongs to, not just one
  5. They send POST http://localhost:4000/auto_router/shadow_eval//stop once and sampling ends for every key under it
  6. They run a forward and a reverse job over the same two keys at once, since direction stays a separate slot per key
  7. They open http://localhost:4000/ui/?page=cost-optimization and the dashboard reads the new shapes correctly; the picker stays single key here, and the multi-key picker plus per key table land in the stacked UI PR
  8. Enrolling a key that another job already covers in the same direction comes back 409 naming that key and job, and an unknown key comes back 400 naming it, before anything is created

Relevant issues

  • One shadow eval job now scopes a set of keys, each with its own turn budget
  • Supersedes feat(auto-router): scope shadow eval jobs to multiple keys #36871, rebuilt on the current single-key schema: one additive group_id column instead of a child table, so the migration moves no data, the destructive column drops are gone, and the sampler hot path plus its whole test file have a zero line diff
  • Pre-existing single-key jobs backfill group_id = id, so their job ids keep resolving on every endpoint
  • Rebased onto current staging; the list's head query is a raw GROUP BY group_id ORDER BY MAX(created_at) per the tag-management precedent, so it never leans on Prisma's in-memory distinct
  • Backend only: the dashboard is adapted minimally here (single-key picker submitting a one-key list, keys-aware labels and totals); the multi-select picker and per key table are the stacked UI PR

Linear ticket

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

  • I have added meaningful tests
  • My PR passes all CI/CD checks (e.g., lint, format, unit tests)
  • My PR's scope is as isolated as possible; it only solves 1 specific problem
  • I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to re-request a review after pushing changes)

Screenshots / Proof of Fix

Rig: proxy from this branch on :4272, real Postgres, an OpenAI-compatible stub upstream serving the router tiers and the judge (every local provider credential is dead; the judged-verdict pipeline itself is unchanged by this PR and was live-proven against real providers when it merged). Two keys mkey-alpha and mkey-beta; my-router is a configured complexity router. The design choice a reviewer will care about: a job's legs share group_id and identical config written by ONE create_many (a single INSERT), so a partial-unique-index loser rolls back the whole claim; race safety rides the existing (api_key_id, direction) WHERE stopped_at IS NULL index unchanged

Before (68d4ba5)

one job per key only

  1. POST /auto_router/shadow_eval/start with "api_key_ids": [hashA, hashB] -> 422 {"detail":[{"type":"missing","loc":["body","api_key_id"],...}]}
  2. Two singular starts -> two job ids (cmsy09htf0000... for alpha, cmsy09hu10001... for beta)
  3. Two turns of traffic on each key, then GET each job: mkey-alpha judged: 2 win_pct: 100.0, mkey-beta judged: 2 win_pct: 0.0; pooling is manual
  4. Stopping takes two calls, one per job (both 200)

After (5494d22)

a pre-migration job id still resolves

  1. The two single-key jobs above sat in the DB when the migration ran; group_id backfilled = id
  2. GET /auto_router/shadow_eval/cmsy09htf0000mgneub781r3k -> 200, status: stopped, keys: [{mkey-alpha, max_turns 5}], judged_count: 2

one start covers both keys, and the key set is bounded

  1. POST /auto_router/shadow_eval/start with "api_key_ids": [hashA, hashB], max_turns: 2 -> 201, one job_id, keys listing both hashes with per key max_turns, key_alias, key_name, status: running
  2. The same start with 101 key ids -> 422 too_long on api_key_ids; the 100-key cap bounds the claim write, every job read, and the 409 message

each key spends its own budget; results pool and slice per key

  1. 3 chat turns on alpha (budget 2) and 2 on beta, all 200
  2. GET /auto_router/shadow_eval/<job id> -> judged: 4, overall_shadow_win_rate_pct: 50.0, by_key: [alpha turns 2, beta turns 2]; alpha judged exactly its budget while its third turn sampled nothing, and beta kept sampling

claim errors name the key, and exhausted slots free themselves

  1. Start naming a key held by a RUNNING job -> 409 "Already in an active forward shadow eval job: 712a7342... (job 6b7da151-...). Stop it first.", nothing created
  2. Start naming unknown keys -> 400 "api_key_ids not on this proxy: nope-1, nope-2; ...", every unknown named at once
  3. Start over keys whose previous legs had exhausted their budgets -> 201; the sweep stamped those legs first, per leg, both directions

forward and reverse coexist; one stop ends every key

  1. With a forward job running, POST .../start with direction: reverse, baseline_model -> 201 over the same two keys
  2. POST /auto_router/shadow_eval/<reverse job id>/stop once -> status: stopped, both keys stamped; a second stop -> 400 already stopped
  3. A job whose keys all stopped on their own budgets already reads stopped, so stop answers 400 for it too

the list collapses legs into jobs

  1. GET /auto_router/shadow_eval?limit=3 -> 3 entries, each one JOB (a two-key job is one row with both aliases), newest first; with the newest job holding 2 legs, limit=2 still returns 2 distinct jobs
  2. GET /auto_router/shadow_eval?api_key_id=<hash B> -> every job containing that key, sibling keys included, among them the pre-migration single-key job

aggregation cost at scale

  1. Seeded 100,000 attempt rows across 50 jobs, then ran the by-key aggregation for one job: EXPLAIN ANALYZE shows a Bitmap Index Scan on the existing job_id index, execution time 1.5 ms. Attempts stay budget-bounded per job (at most 2000 per key), never per-request unbounded

UI (minimal adaptation)

  1. Open http://localhost:4000/ui/?page=cost-optimization, Shadow Evals section; the form, cards and results read the new response shape (labels off keys[], turn totals summed over keys)
  2. The picker stays single key and submits api_key_ids: [key]; the multi-select picker and per key table are in the stacked UI PR

Type

🆕 New Feature

Caveats (if any)

  • Breaking API reshape: api_key_id -> api_key_ids on start, top level key fields -> keys[] on responses. Internal admin endpoints, weeks old, no external consumers; the dashboard is adapted in this PR and upgraded in the stacked UI PR
  • Legs of one job duplicate the immutable config columns; written once by one atomic create_many, no update path exists, pinned by a uniform-payload test
  • Judge and shadow arms hit a stub upstream in this run; the judging pipeline is untouched by this PR
  • Error attempts still count toward max_turns, unchanged pre-existing semantics, out of scope here
  • An exhausted key reads running until the next start's sweep stamps it, the pre-existing lazy stamp of the derived-state design; the hot path stays append-only

Final Attestation

  • The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

Note

Cursor Bugbot is generating a summary for commit 5494d22. Configure here.

@tin-berri
tin-berri requested a review from a team August 18, 2026 02:22
@tin-berri

Copy link
Copy Markdown
Contributor Author

@greptileai can you reply here with a score and why that score?

@tin-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@greptile-apps

greptile-apps Bot commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Score: 4/5

Here's why:

What's done well:

  • Schema design is minimal and non-destructive. Adding group_id with a SET group_id = id backfill means zero data migration, zero downtime risk, and all pre-existing job IDs keep resolving transparently across every endpoint. This is the right call over a child-table redesign.

  • Atomic claim write. Using a single create_many so a unique-index loser rolls back the whole multi-key claim is exactly the right race safety model. The 409 on concurrent starts names the key and blocking job, which is a real usability improvement.

  • Per-key budget independence is correctly modelled. Each leg carries its own max_turns and stopped_at; the sweep SQL and status computation both reflect that a sibling exhausting its budget must not end others. The status computed field correctly uses all(key.stopped_at is not None for key in self.keys).

  • Aggregation cost. Using ANY($1::text[]) against the existing job_id index on LiteLLM_ShadowEvalAttempt keeps the per-job result query O(attempts in that job), not O(all attempts). The EXPLAIN ANALYZE evidence in the PR description backs this up.

  • Test coverage is strong. The fake find_many_legs implementation honours all the filters (direction, stopped_at, group_id) that real Prisma does, so the 409/400 rejection tests and the status-derivation parametrize actually catch logic bugs rather than just invoking the happy path.

What holds it to 4 rather than 5:

  1. The list endpoint's grouping is done in two round-trips and then Python. find_many(distinct=["group_id"], take=limit) returns the head row per group, then a second find_many(group_id in [...]) fetches every leg for those groups, then Python groups them. For the bounded limit (max 200 groups) this is fine in practice, but the second query is unbounded in leg count — a job scoping 1000 keys would fetch 1000 rows for a 1-job list response. A single query with a JOIN or a GROUP BY group_id aggregation would avoid this, but it's not a blocker at current scale.

  2. _group_response assumes all legs share identical config, which is a maintained invariant (one create_many payload, no update path), but it's implicit — there's no DB-level constraint ensuring it. The uniform-payload test in the PR catches it for the write path, but a future maintenance author adding a per-leg config update could silently violate it without a loud failure. A comment naming this invariant explicitly on _group_response would help.

  3. Minor: the by_key slice remaps grp from internal leg IDs to key hashes in-memory (key_by_leg[row.grp]). If the DB ever returned an attempt whose job_id wasn't in legs (e.g., from a race between leg insertion and the read), this would KeyError silently. A .get() with a fallback or an explicit invariant comment would be safer — though in practice the read is bounded by the same transaction's leg IDs.

None of these are blockers. The design is solid, the migration is safe, and the test suite covers the semantics that actually matter.

@greptile-apps

greptile-apps Bot commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

The PR extends auto-router shadow evaluations from one key to grouped multi-key jobs while preserving independent budgets and lifecycle state for each key

  • Adds and backfills group_id across all synchronized Prisma schemas
  • Creates one atomic job leg per selected key and pools detail, list, stop, and result operations by group
  • Adds per-key result slices and updates the generated API contract
  • Updates the dashboard with paginated multi-key selection and per-key status/results
  • Expands backend and frontend tests for grouped jobs and shared multi-select behavior

Confidence Score: 5/5

The PR appears safe to merge, with no concrete blocking or independently actionable non-blocking defects identified

Grouped persistence, per-leg sampling state, pooled aggregation, API contracts, and dashboard consumers remain coordinated, and the sampler continues to operate correctly on each unique leg ID

Important Files Changed

Filename Overview
litellm/proxy/management_endpoints/auto_router_endpoints.py Groups independently persisted key legs for creation, listing, detail aggregation, labeling, and stopping without an accepted correctness issue
litellm/types/management_endpoints/auto_router_endpoints.py Replaces the single-key request and response contracts with nonempty key collections and per-key lifecycle data
litellm-proxy-extras/litellm_proxy_extras/migrations/20260817000000_shadow_eval_multi_key/migration.sql Adds, backfills, constrains, and indexes the grouping column without destructive schema changes
schema.prisma Adds the required grouping field and index, synchronized with both schema copies
ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx Updates shadow-evaluation creation and detail rendering for multiple selected keys and per-key results
ui/litellm-dashboard/src/components/shared/PaginatedMultiSelect.tsx Adds paginated multi-selection while retaining selected values across changing option pages
tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py Expands endpoint coverage for atomic grouped creation, grouped reads, filtering, stopping, and aggregation
ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx Updates dashboard tests to cover multi-key submission and grouped job presentation

Reviews (1): Last reviewed commit: "feat(auto-router): scope shadow eval job..." | Re-trigger Greptile

Comment thread litellm/types/management_endpoints/auto_router_endpoints.py
@codecov

codecov Bot commented Aug 18, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 98.03922% with 2 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
...roxy/management_endpoints/auto_router_endpoints.py 97.43% 2 Missing ⚠️

📢 Thoughts on this report? Let us know!

@codspeed-hq

codspeed-hq Bot commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

✅ 31 untouched benchmarks


Comparing litellm_shadoweval_multikey2 (5494d22) with litellm_internal_staging (4d100bd)1

Open in CodSpeed

Footnotes

  1. No successful run was found on litellm_internal_staging (4bb3152) during the generation of this report, so 4d100bd was used instead as the comparison base. There might be some changes unrelated to this pull request in this report.

@tin-berri
tin-berri force-pushed the litellm_shadoweval_multikey2 branch from 5b83c89 to e92f912 Compare August 18, 2026 04:23
@tin-berri

Copy link
Copy Markdown
Contributor Author

api_key_ids now capped at 100, bounding every job read. The config invariant is _group_response's docstring; by_key grp is query-bounded to leg_ids

@tin-berri

Copy link
Copy Markdown
Contributor Author

@greptileai can you reply here with a score and why that score?

@tin-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@greptile-apps

greptile-apps Bot commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Score: 4/5

Here's the breakdown:


What earns the 4:

  • Migration is airtight. Additive group_id column, backfills group_id = id for all existing rows, creates the index. Pre-migration job IDs keep resolving on every endpoint — no data loss, no blocking migration, no destroyed old job IDs.

  • Race safety is correct. The whole claim writes in one create_many. The partial unique index on (api_key_id, direction) WHERE stopped_at IS NULL already existed; a concurrent loser rolls back the entire batch rather than partially inserting.

  • Test design is unusually strong. The find_many_legs fake in _shadow_prisma actually filters on direction, stopped_at, and group_id — catching exactly the class of bugs where a query forgets to scope by one of those fields. That's not boilerplate; it's the tests exercising the invariants that matter.

  • Budget independence is correctly modeled. One key exhausting max_turns stamping its own stopped_at never affects its siblings. The status computed field reads stopped only when every key has a stopped_at. The sweep before a new start frees expired legs per-key. The update_many WHERE stopped_at IS NULL on stop leaves budget-earned timestamps intact.

  • Error quality. All unknown api_key_ids are named at once (caller fixes in one round), and the 409 names both the blocked key and the job holding it.

  • Zero sampler diff. The hot path is untouched, which is the correct scope boundary for this PR.


What keeps it from a 5:

  • _group_response has no guard on group_id consistency. It reads shared config off legs[0] and assumes all legs in the sequence belong to the same group. It's an internal function with three call sites, all of which read by group_id, so the risk is low — but a single assert or a group_id parameter would make the invariant explicit and catch miscalls during future refactors.

  • grp is overloaded across stratifications. In by_tier it holds a tier name, in by_current_model a model name, in by_key a key hash (after the leg-id→key-hash remap). The UI works correctly because it uses slices.get(key.api_key_id) for the key table, but the semantic inconsistency of ShadowEvalSlice.group across stratifications is a code smell that will trip up future readers.

  • distinct=["group_id"] in the list endpoint relies on Prisma returning an arbitrary row per group. It works here because all legs in a group share the same created_at (written atomically with identical config), so ordering doesn't differentiate them and any leg serves as the "head." That assumption is correct today but is not enforced anywhere — it's a silent invariant.

None of these are blocking. The core design is solid, the migration is safe, and the test coverage is thorough.

Comment thread litellm/proxy/management_endpoints/auto_router_endpoints.py Outdated
@tin-berri

Copy link
Copy Markdown
Contributor Author

All _group_response call sites filter by group_id; ShadowEvalSlice.group is the established shared shape; any head leg yields the same group_id

@tin-berri
tin-berri force-pushed the litellm_shadoweval_multikey2 branch from e92f912 to b12cb28 Compare August 18, 2026 23:31
@tin-berri

Copy link
Copy Markdown
Contributor Author

@greptileai can you reply here with a score and why that score?

@tin-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@greptile-apps

greptile-apps Bot commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Score: 4/5

Why:

What's well done:

  • Migration is safe and additive. group_id is added, backfilled = id, then set NOT NULL — no data movement, no destructive drops, old job IDs still resolve via WHERE group_id = job_id.
  • Race safety is preserved. The existing partial unique index (api_key_id, direction) WHERE stopped_at IS NULL is unchanged. The create_many write means a partial unique violation rolls back the entire claim atomically — this is the key correctness invariant.
  • Leg IDs never leave the server. Only group_id (the API's job ID) is exposed; leg.id is an internal join key. _group_response enforces this — no caller can accidentally serialize a leg ID as a job ID.
  • Status derivation is correct. status is now all(key.stopped_at is not None for key in self.keys) — one running key keeps the job running. The five-case parametric test covers every combination.
  • Tests verify filter semantics explicitly. The find_many_legs mock enforces api_key_id, direction, stopped_at, and group_id filters independently, so a future caller that omits stopped_at=None would see wrong rows in the test rather than silently passing.
  • Sweep is correctly extended. _SWEEP_FINISHED_JOBS_SQL now uses ANY($1::text[]) over all requested keys before the active-slot check.
  • Stop preserves earned stopped_at. update_many(where={stopped_at: None}) only stamps unstopped legs; the response reconstructs per-leg timestamps correctly.

What holds it from a 5:

  • Two-phase list query. Between the GROUP BY group_id fetch and the find_many legs fetch, a job created in that window could be missed. Read-only so no correctness risk, but worth noting for future readers.
  • by_group dict construction is O(n×m) (iterate all legs for each group ID). Bounded at 100 keys × 200 groups so fine in practice, but a single defaultdict pass would be cleaner and make the ordering invariant explicit.
  • API breaking change on endpoints that existed for weeks. Acknowledged in the PR and the dashboard is updated in the same commit, but it's a reshape that downstream tooling or scripts could hit silently.

None of these are blocking — the core design is sound.

@cursor cursor Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅ Bugbot reviewed your changes and found no new issues!

Comment @cursor review or bugbot run to trigger another review on this PR

Reviewed by Cursor Bugbot for commit b12cb28. Configure here.

@tin-berri
tin-berri force-pushed the litellm_shadoweval_multikey2 branch from b12cb28 to 83c1533 Compare August 18, 2026 23:38
@tin-berri

Copy link
Copy Markdown
Contributor Author

by_group now builds via sorted plus groupby; ordering stays owned by the head query. The two-phase window is read-only; no release ships the old shape

@tin-berri

Copy link
Copy Markdown
Contributor Author

@greptileai can you reply here with a score and why that score?

@tin-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@greptile-apps

greptile-apps Bot commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Score: 4/5


What earns the high mark:

  • Schema design is minimal and safe. One additive group_id column, backfilled = id for pre-existing rows in a 4-line migration with zero data movement. Legacy job IDs still resolve on every endpoint — no external contract breaks.
  • Race safety is preserved correctly. The existing partial unique index on (api_key_id, direction) WHERE stopped_at IS NULL remains the guard. create_many writes all legs in one statement, so a unique-index loser rolls back the whole claim atomically. The sweep-before-claim correctly targets all requested keys via ANY($1::text[]).
  • Status derivation is correct. completed wins when ends_at passes (time outranks a manual stop), stopped requires all keys to have stopped_at set, running otherwise. One exhausted key cannot end sampling for siblings.
  • Tests are thorough and test the right invariants. The find_many_legs fake correctly filters on stopped_at, direction, and group_id, so tests that forget those filters would expose the bug. The uniform-payload test (len({frozenset(…) for row in rows}) == 1) verifies all legs carry identical config. The 409 naming test, the dedupe/bounds test, the stop-preserving-earned-stopped_at test — all check real invariants.
  • Error ergonomics are solid. 400 names every unknown key at once; 409 names which key and which job is blocking; 422 bounds the key list at the validation layer before any DB round-trip.
  • by_key aggregation is correctly implemented. SELECT job_id AS grp on attempt rows returns leg IDs; key_by_leg maps those back to the canonical key hashes before exposing them in the response. The leg ID never leaves the server.

What keeps it from 5/5:

  1. Breaking API reshape on an endpoint that's only weeks old. api_key_id → api_key_ids on the start body and top-level key_alias/key_name/max_turns moving into keys[] on responses are acknowledged as internal-only, but any consumer (dashboard, direct API callers) needs to be updated atomically. The dashboard is updated in this PR, which mitigates this, but it's worth calling out that there's no field-level deprecation path.

  2. _group_response sorts keys alphabetically by api_key_id hash — deterministic but arbitrary. The response order has no semantic meaning and is not documented, so a consumer that starts depending on position will silently break if the sort ever changes.

  3. _LIST_GROUPS_BY_KEY_SQL uses $2 and the test mock detects it by substring-matching "api_key_id = $2" in the SQL string. That coupling is fragile: a reformatting of the SQL (e.g. a newline) would silently disable the filter in tests without any test failure. Minor, but worth a note.

  4. No test for by_key when a leg has zero attempts (absent from by_leg rows). The docstring says "Keys the job scopes but has not judged a turn for yet are absent rather than reported as zero" — that behavior is correct and documented, but the test only covers the case where all legs have rows.

Overall this is a clean, well-reasoned feature with a design that holds up under concurrent starts, budget exhaustion, and backward compatibility. The four points above are minor enough that none of them are blockers.

@tin-berri

Copy link
Copy Markdown
Contributor Author

Split per review direction: this PR is backend plus a minimal dashboard adaptation; the multi-select picker and per key table moved to stacked #37389

… stops

Two operators can both pass the derived-status guard in the race window. The
stop UPDATE now claims only legs with stopped_by still null and the endpoint
judges by its row count, so exactly one caller ever gets the 200 and the loser
gets the same already-stopped 400 a late caller gets
@tin-berri

Copy link
Copy Markdown
Contributor Author

Racing stops now resolve in the statement itself: it claims only stopped_by null legs and the row count picks the one winner; loser gets the late-caller 400

@tin-berri

Copy link
Copy Markdown
Contributor Author

@greptileai can you reply here with a score and why that score?

@tin-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@cursor cursor Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅ Bugbot reviewed your changes and found no new issues!

Comment @cursor review or bugbot run to trigger another review on this PR

Reviewed by Cursor Bugbot for commit 2128db6. Configure here.

@greptile-apps

greptile-apps Bot commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Score: 4/5

This is genuinely well-engineered. Here's the breakdown:


What earns the 4:

  • Migration is purely additive. group_id is added, backfilled = id, made NOT NULL — zero data movement, zero destructive drops, pre-migration job ids keep resolving unchanged. Exactly the right approach.
  • Atomic claim. create_many in a single statement means a partial unique-index loser rolls back the entire multi-key claim. No TOCTOU window where some legs are written and others aren't.
  • stopped_by cleanly separates signal from noise. Budget arithmetic can be confused by detached in-flight attempts; a recorded operator stop cannot. The COALESCE(stopped_at, $3) in _STOP_JOB_SQL with the stopped_by IS NULL predicate is a clean lock — one winner, no half-applied stops, earned timestamps preserved on already-swept legs.
  • Status derivation precedence is correct: stopped_by first, window next, all-budgets-spent next, all-stamped fallback for rolling-deploy legacy pods. The backfill in 20260818224500 closes the pre-column population exactly.
  • Test coverage is exceptional. test_every_shadow_eval_sql_constant_speaks_naive_utc, budget-vs-operator precedence, concurrent stop races, backfilled legacy stops, filter-by-key returning sibling keys — these are the tests that actually catch production regressions.

What keeps it from 5:

  1. Semantic race on the stop path. Between reading current.status == "running" and executing _STOP_JOB_SQL, new attempts can push counts past max_turns. The SQL still writes (claimed > 0), so the job permanently reads stopped instead of completed. The PR explicitly chose stopped_by precedence here, but the result is that an operator who stops a job that budget-exhausted in the same millisecond gets incorrect reporting forever. This is a knowable trade-off, not a hidden bug, but worth flagging.

  2. _group_response reads shared config off legs[0] with no DB-level enforcement. The invariant — all legs share identical config — is enforced by application logic (create_many with a single shared payload) and a unit test assertion, but there's no DB constraint preventing divergence if a future code path ever touches an individual leg. Fine given the immutable design, but a convention that could silently produce wrong job-level config if the invariant breaks.

  3. The stopped_by migration is shared verbatim with fix(shadow_eval): split slot release from stops so stopped_at means exactly one thing #37382. The PR acknowledges "whichever merges second inherits it as already applied." That's an operational dependency with no enforcement mechanism — if both PRs are in flight simultaneously and reviewed independently, the ordering risk is real. A defensive IF NOT EXISTS guard on the ALTER would close this.

The status guard ran before the UPDATE, so a stop racing the last budgeted
attempt still claimed the job and it read stopped forever instead of
completed. The statement now claims the job only while a leg still samples
inside the window with no stop recorded, and the endpoint reads once after
writing: a racing operator, a same-instant budget spend, and a repeat stop all
get the 400 naming the status the job actually holds. The pre-write guard and
the hand-built response go away
@tin-berri

Copy link
Copy Markdown
Contributor Author

The stop statement now claims only a still-sampling job, so a stop racing the last budgeted attempt reads already completed and records no operator; the pre-write guard is deleted

@tin-berri

Copy link
Copy Markdown
Contributor Author

On config duplication: true DB enforcement needs a btree_gist EXCLUDE or a trigger for a state no writer can produce; the immutable design plus the uniform-payload test stays the guard

@tin-berri

Copy link
Copy Markdown
Contributor Author

@greptileai can you reply here with a score and why that score?

@tin-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@cursor cursor Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅ Bugbot reviewed your changes and found no new issues!

Comment @cursor review or bugbot run to trigger another review on this PR

Reviewed by Cursor Bugbot for commit dad260f. Configure here.

@greptile-apps

greptile-apps Bot commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Score: 4/5

Here's the reasoning:

What's excellent:

  • Migration is textbook-safe: additive only — nullable group_id, backfill to id, then NOT NULL. Zero destructive drops, zero data moves. group_id = id backfill means pre-migration job IDs keep resolving on every endpoint.
  • Atomic multi-key write: create_many in a single INSERT; the existing partial-unique index (api_key_id, direction) WHERE stopped_at IS NULL makes the entire batch fail atomically on any collision — no partial claims possible.
  • Stop endpoint is a clean state machine: the single UPDATE using stopped_by IS NULL AND ends_at > NOW() AND EXISTS(still-sampling leg) means racing operators, budget-spend races, and repeat stops all get the same deterministic 400 with no phantom "stopped" rows.
  • stopped_by column cleanly separates operator-stopped from budget-spent — the old stopped_at-only design conflated these two, which the status computed field had to paper over. The migration backfills 'unknown' for all existing operator-stopped jobs, closing the pre-column population correctly.
  • Timezone safety is thorough: all SQL uses NOW() AT TIME ZONE 'utc', params pass ::timestamp, and _as_aware_utc normalizes the prisma-ORM vs. raw-SQL difference. The test test_every_shadow_eval_sql_constant_speaks_naive_utc mechanically enforces this across all SQL constants — a rare and good invariant test.
  • Test depth: ~300 lines covering the budget-spend/stop race, legacy backfill, SQL timezone invariant, stop-racing-last-attempt, repeated stop, attempt count freezing at stamp time. The fake execute_raw in _shadow_prisma faithfully implements the stopped_by IS NULL and ends_at > now guards rather than stubbing them out.

What keeps it from 5/5:

  • Breaking REST API: api_key_id → api_key_ids on start, job.api_key_id / max_turns / stopped_at → job.keys[] on responses. The PR acknowledges these are internal-only endpoints weeks old with no external consumers, and the dashboard is updated in the same PR, so the risk is contained — but any deployed script or external integration using the old shape would silently break.
  • _LIST_LEGS_SQL over-fetches on limit semantics: the subquery selects the top-N group_ids, then the outer query returns all legs for those groups — so limit=3 may return up to 300 rows if each job has the 100-key max. This is bounded and correct, but the 100-key cap and limit semantics interact in a way that isn't immediately obvious to a future reader.
  • _STOP_JOB_SQL complexity: the EXISTS correlated subquery inside the UPDATE is correct and tested, but the logic is dense enough that a future maintainer touching the SQL (e.g., to add a direction filter for directional stops) could easily break the atomicity guarantee.

@greptile-apps

greptile-apps Bot commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Score: 4/5

This is a well-engineered PR. The legs/group_id model is the right shape for the problem, the migrations are safe and additive, and the test coverage is genuinely thorough (racing stops, budget exhaustion, timezone handling, backfilled legacy jobs). Here's the breakdown:


What works well:

  • Migration safety. Two clean additive migrations: group_id nullable → backfill → NOT NULL; stopped_by with immediate backfill for the legacy-stopped population. No destructive column drops, no data movement.
  • Atomic job creation. create_many means a unique-index loser on any key in the set rolls back all legs at once — no partial job state is possible.
  • Stop as a single SQL CAS. _STOP_JOB_SQL stamps stopped_by and COALESCE(stopped_at, $3) in one statement, with the EXISTS guard ensuring it claims nothing after the last sampling leg is gone. Racing stops and budget-raced stops are both handled correctly.
  • stopped_by drives status, not arithmetic. The migration backfills 'unknown' for pre-column stopped jobs, so the pre-column population is closed before the new derivation runs. test_backfilled_legacy_stop_never_reads_as_completion verifies this explicitly.
  • _LegRow._as_aware_utc validator. Correctly absorbs the prisma-vs-raw-SQL timezone difference at the boundary — naive UTC in, aware UTC out, consistently.
  • Zero sampler diff. Hot path is untouched, which the PR explicitly calls out.

Issues keeping it from 5/5:

  1. stopped_by leaks onto budget-spent-unswepped legs during a partial stop. If key-A has spent its budget but hasn't been swept yet (stopped_at IS NULL, count >= max_turns), and key-B is still sampling, a stop call fires the EXISTS (key-B qualifies), and the outer WHERE stopped_by IS NULL stamps key-A too. COALESCE(stopped_at, $3) sets key-A's stopped_at to the operator's timestamp rather than the time it ran out, and key-A gets stopped_by = 'admin'. The group response's stopped_by will say "admin" even for a key that self-completed. This is functionally harmless (the group reads "stopped" correctly since the operator genuinely stopped the other leg), but it's a semantic inaccuracy that could confuse a future consumer diffing stopped_by from budget_spent at the per-key level.

  2. _shadow_prisma fake complexity. The test fake is now ~130 lines with its own find_many filter simulation, group-collapsing logic, and attempt-count freeze. It's faithful to the real DB behavior, which is why the tests are trustworthy, but it's a maintenance surface. A partial drift (e.g. forgetting to update leg_dict fields when a column is added) would silently produce passing tests over wrong behavior.

  3. _ATTEMPT_COUNTS_SQL stopped_at freeze is correct but not tested end-to-end. The SQL a.created_at <= j.stopped_at freeze is tested via the fake's attempt_rows fixture, but only the in-memory fake implements the freeze logic — there's no integration test against a real Postgres that verifies the SQL actually does what the comment says. Low risk given the SQL is straightforward, but worth noting.


Summary: The design is sound, the migrations are safe, and the tests are among the most complete I've seen for a feature of this scope. The issues above are minor. This is ready for maintainer review.

@tin-berri

Copy link
Copy Markdown
Contributor Author

On the three 4/5 notes: keys[] never exposes per-key stopped_by, so nothing can diff it; the freeze SQL is live-proven in the body's straggler run; fake fidelity noted

@tin-berri

Copy link
Copy Markdown
Contributor Author

Context for the three 4/5 notes: keys[] never exposes per-key stopped_by, so no consumer can diff it against budget_spent; the attempt-freeze SQL is live-proven in the body's fabricated-straggler run against real Postgres; the fake's fidelity is the stated price of trustworthy tests. Types-sync regen pushed. @greptileai can you reply here with a score and why that score?

@tin-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@cursor cursor Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅ Bugbot reviewed your changes and found no new issues!

Comment @cursor review or bugbot run to trigger another review on this PR

Reviewed by Cursor Bugbot for commit 7a919c1. Configure here.

@greptile-apps

greptile-apps Bot commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Based on your clarifications:

  1. stopped_by not in keys[] — if it's not surfaced, the concern about consumers conflating it with budget_spent doesn't apply; the API contract doesn't create that ambiguity.

  2. Attempt-freeze SQL live-proven on real Postgres — the fabricated-straggler run in the PR body covers the race the unit test can't: a concurrent write lands between the freeze reads and the UPDATE. That's the gap the concern was about, and it's addressed.

  3. Fake fidelity as intentional tradeoff — acknowledged; the fake is the price of deterministic isolation, and the Postgres run validates the real path.

Given those three, the residual risks are:

  • The types-sync regen being pushed is good; that was a latent drift concern.
  • The lazy-stamp semantics (exhausted key reads running until next start's sweep) are documented and out of scope, not a correctness gap in this PR.
  • The 100-key cap, atomic create_many rollback, and the group_id backfill for legacy jobs are all demonstrated.

Nothing substantive remains open from the 4/5 holds. Score: 5/5.

The breaking API reshape (api_key_idapi_key_ids, keys[]) is the one thing a reviewer will ask about in isolation, but you've covered it: internal admin endpoint, weeks old, no external consumers, dashboard adapted in this PR. That's the right framing.

@tin-berri tin-berri changed the title feat(auto-router): scope shadow eval jobs to multiple keys feat(auto-router)!: scope shadow eval jobs to multiple keys Aug 19, 2026
…itellm_shadoweval_multikey2

# Conflicts:
#	litellm/proxy/management_endpoints/auto_router_endpoints.py
@tin-berri

Copy link
Copy Markdown
Contributor Author

Merged current staging to clear the conflict; the typing sweep's Protocol accessors are adopted across the multi-key code with the job-table Protocol adapted to find_many and create_many. @greptileai can you reply here with a score and why that score?

@tin-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@greptile-apps

greptile-apps Bot commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Score: 4.5 / 5


What's strong

Migration design is correct and zero-risk. Two additive migrations: group_id with an immediate UPDATE … SET group_id = id WHERE group_id IS NULL backfill (no data loss, old job IDs resolve identically via WHERE group_id = $1), then stopped_by with a precise backfill that closes the pre-column population cleanly. Both use IF NOT EXISTS / ADD COLUMN IF NOT EXISTS so they're re-entrant.

Race safety is built into the write path, not bolted on. A single create_many turns the whole key-set claim into one INSERT; a unique-index loser rolls back all legs, not a subset. The _STOP_JOB_SQL is a single UPDATE with an EXISTS predicate that checks whether any leg is still sampling, so a concurrent budget spend and an operator stop resolve against each other at the statement level — no read-then-write.

stopped_by as the authoritative operator-stop signal is the right design. Previously stopped_at was both "operator stopped it" and "budget/time expired and was swept." Now those are separate facts: stopped_by is only written by the stop endpoint, and stopped_at is only written by the sweep. The status property is strictly a derivation from three recorded facts with no history-guessing, and the migration backfills 'unknown' exactly for the pre-column population that displayed "stopped."

UTC handling is thorough. The _as_aware_utc validator on _LegRow normalises the naive-vs-aware mismatch between Prisma and raw SQL reads. The test_every_shadow_eval_sql_constant_speaks_naive_utc test asserts no ::timestamptz and every NOW() is NOW() AT TIME ZONE 'utc' across every SQL constant — that will catch a regression the instant it's introduced.

Test coverage is exceptional. ~30 unit tests including:

  • Multi-leg status derivation (all-stopped, half-stopped, all-budget-spent)
  • Racing stop vs budget spend
  • Two concurrent stops
  • Batched unknown-key error naming all unknowns at once
  • stopped_by overranking budget arithmetic (detached attempt past budget)
  • Legacy backfill correctness verified by reading the actual migration file off disk
  • SQL constant invariant check

The fake _shadow_prisma is a stateful in-memory implementation, not a simple mock — find_many_legs honours stopped_at, direction, and group_id filters, so a test that forgets a filter fails rather than silently passes.

Performance is bounded. The 100-key cap propagates to every downstream read. The _LIST_LEGS_SQL subquery uses a GROUP BY group_id ORDER BY MAX(created_at) DESC LIMIT $1 that the existing group_id index backs. Attempt aggregations are bounded by job_id = ANY($1::text[]) with the existing per-job index.


Minor concerns (no blockers)

_STOP_JOB_SQL stamps stopped_by on budget-spent, un-swept legs. When an operator stops a two-leg job where one leg already spent its budget but hasn't been swept (stopped_at IS NULL, attempt_count >= max_turns), the UPDATE's stopped_by IS NULL predicate matches it and writes stopped_by = operator. The job-level status is computed correctly (stopped_by IS NOT NULL → "stopped"), but the individual leg now has stopped_by set even though it self-completed. This is harmless since individual leg IDs never surface through the API, but it does mean stopped_by on a leg is no longer purely "operator explicitly ended this." Future maintainers reading a leg row directly (e.g. in a migration or a debug query) would see operator attribution on a self-ended leg.

_group_response uses legs[0] for shared config without asserting order. The design guarantee is that every leg in a group carries identical config (written by one create_many from a single shared_config dict), and _group_response's docstring states this. But the function receives a Sequence[_LegRow] with no ordering guarantee, so legs[0] could be any leg. In practice it's always correct, but a small defensive assert or a comment at the call sites that pass legs would clarify the invariant for future readers.

stopped_by priority overrides window-expiry for legacy data. A job stopped before its window (old behaviour: "stopped"), whose window then expires (old behaviour: "completed"), now reads "stopped" permanently because stopped_by = 'unknown' was backfilled. The old code had ends_at as the highest-priority check. The test_backfilled_legacy_stop_never_reads_as_completion test documents and asserts this deliberately, so it's an intentional design choice — just worth noting it's a visible status change for that small legacy population.


These are all minor and the code is clearly correct. The overall design — additive schema, atomic multi-leg write, single-statement stop with EXISTS-guarded predicate, stopped_by as the authoritative signal, and thorough test coverage including SQL-level and file-level assertions — is the work of someone who has thought through the edge cases carefully.

@cursor cursor Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅ Bugbot reviewed your changes and found no new issues!

Comment @cursor review or bugbot run to trigger another review on this PR

Reviewed by Cursor Bugbot for commit 5494d22. Configure here.

@tin-berri
tin-berri merged commit a613773 into litellm_internal_staging Aug 19, 2026
81 checks passed
@tin-berri
tin-berri deleted the litellm_shadoweval_multikey2 branch August 19, 2026 21:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants