From 71ff91251f85ba56bbcde68bd53420086b04ae25 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 20 Aug 2026 15:57:50 -0700 Subject: [PATCH] fix(shadow_eval): split slot release from stops so stopped_at means exactly one thing --- .../migration.sql | 13 ++++ .../litellm_proxy_extras/schema.prisma | 1 + litellm/integrations/shadow_eval_logger.py | 1 + .../auto_router_endpoints.py | 11 +-- litellm/proxy/schema.prisma | 1 + .../auto_router_endpoints.py | 18 +++-- schema.prisma | 1 + .../test_auto_router_endpoints.py | 74 ++++++++++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 13 ++-- 9 files changed, 111 insertions(+), 22 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260820150000_shadow_eval_released_at/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260820150000_shadow_eval_released_at/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260820150000_shadow_eval_released_at/migration.sql new file mode 100644 index 00000000000..8d408701882 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260820150000_shadow_eval_released_at/migration.sql @@ -0,0 +1,13 @@ +-- AlterTable +ALTER TABLE "LiteLLM_ShadowEvalJob" ADD COLUMN IF NOT EXISTS "released_at" TIMESTAMP(3); + +UPDATE "LiteLLM_ShadowEvalJob" SET released_at = stopped_at +WHERE stopped_at IS NOT NULL AND released_at IS NULL; + +UPDATE "LiteLLM_ShadowEvalJob" SET stopped_at = NULL +WHERE stopped_at IS NOT NULL AND stopped_by IS NULL; + +DROP INDEX IF EXISTS "LiteLLM_ShadowEvalJob_one_active_per_key_direction"; + +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalJob_one_active_per_key_direction" + ON "LiteLLM_ShadowEvalJob"("api_key_id", "direction") WHERE "released_at" IS NULL; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index d9959677116..4612921c712 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1509,6 +1509,7 @@ model LiteLLM_ShadowEvalJob { ends_at DateTime stopped_at DateTime? stopped_by String? // operator who stopped it early; null when it ended on its own + released_at DateTime? // frees the key's slot in the partial unique index; stops and sweeps both set it, only stops set stopped_at @@index([group_id]) @@index([api_key_id]) diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index 5f4e7c71395..b557716b300 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -595,6 +595,7 @@ class ShadowEvalLogger(CustomLogger): records: Final = await prisma.db.litellm_shadowevaljob.find_many( where={ # mutable-ok: Prisma filter "stopped_at": None, + "released_at": None, "ends_at": {"gt": datetime.now(timezone.utc)}, # mutable-ok: Prisma filter }, ) diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 46aac82473c..581a6b06379 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -567,7 +567,7 @@ async def get_auto_router_benchmarks( # --------------------------------------------------------------------------- # Shadow eval: pre-adoption evaluation of an auto-router against live traffic. -# The job row is immutable config plus stopped_at; status, counts, spend, and errors +# The job row is immutable config plus the stop record and released_at slot marker; status, counts, spend, and errors # are derived from the append-only attempt rows, so reads here are aggregations # bounded by each job's max_turns through the attempt table's job_id index. # --------------------------------------------------------------------------- @@ -667,8 +667,8 @@ _ATTEMPT_AGG_BY_LEG_SQL: Final = "SELECT job_id AS grp," + _ATTEMPT_AGG_SELECT # reads the live counter, so admission can stop before a row-based guard would fire (safe # direction, and mid-deploy rows from old pods price as judge-only until the deploy ends). _SWEEP_FINISHED_JOBS_SQL: Final = """ -UPDATE "LiteLLM_ShadowEvalJob" j SET stopped_at = (NOW() AT TIME ZONE 'utc') -WHERE j.api_key_id = ANY($1::text[]) AND j.stopped_at IS NULL +UPDATE "LiteLLM_ShadowEvalJob" j SET released_at = (NOW() AT TIME ZONE 'utc') +WHERE j.api_key_id = ANY($1::text[]) AND j.released_at IS NULL AND ( j.ends_at <= (NOW() AT TIME ZONE 'utc') OR (SELECT COUNT(*) FROM "LiteLLM_ShadowEvalAttempt" a WHERE a.job_id = j.id) >= j.max_turns @@ -698,7 +698,8 @@ GROUP BY a.job_id _STOP_JOB_SQL: Final = """ UPDATE "LiteLLM_ShadowEvalJob" -SET stopped_by = $2, stopped_at = COALESCE(stopped_at, $3::timestamp) +SET stopped_by = $2, stopped_at = COALESCE(stopped_at, $3::timestamp), + released_at = COALESCE(released_at, $3::timestamp) WHERE group_id = $1 AND stopped_by IS NULL AND ends_at > (NOW() AT TIME ZONE 'utc') AND EXISTS ( @@ -981,7 +982,7 @@ async def start_shadow_eval( where={ # mutable-ok: Prisma filter "api_key_id": {"in": requested}, # mutable-ok: Prisma filter "direction": data.direction, - "stopped_at": None, + "released_at": None, }, ) if claimed: diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index d9959677116..4612921c712 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1509,6 +1509,7 @@ model LiteLLM_ShadowEvalJob { ends_at DateTime stopped_at DateTime? stopped_by String? // operator who stopped it early; null when it ended on its own + released_at DateTime? // frees the key's slot in the partial unique index; stops and sweeps both set it, only stops set stopped_at @@index([group_id]) @@index([api_key_id]) diff --git a/litellm/types/management_endpoints/auto_router_endpoints.py b/litellm/types/management_endpoints/auto_router_endpoints.py index e2469d4c78f..89f0979df44 100644 --- a/litellm/types/management_endpoints/auto_router_endpoints.py +++ b/litellm/types/management_endpoints/auto_router_endpoints.py @@ -386,7 +386,8 @@ class ShadowEvalJobResponse(BaseModel): description=( "The operator who stopped the job early, recorded by the stop endpoint; 'unknown' backfilled " "by migration for jobs that displayed stopped when the column arrived; None when the job " - "ended on its own. Its presence is what makes a job read stopped rather than completed" + "ended on its own. Budget and window completion never write the stop record: slot " + "bookkeeping lives in an internal column instead" ), ) @@ -399,19 +400,20 @@ class ShadowEvalJobResponse(BaseModel): @computed_field @property def status(self) -> ShadowEvalStatus: - """Three recorded facts, no history-guessing: a stop is stopped_by (the migration - backfills it for every job that displayed stopped when the column arrived, so the - pre-column population is closed), completion is the window passing or every key - spending its budget, and anything else is running. The all-keys-stamped fallback - covers only stops written by pre-column pods during a rolling deploy.""" + """Recorded facts only, no history-guessing. A key's stopped_at now means exactly + one thing, a stop: budget and window completion free the key's slot through the + internal released_at column and never touch it, so every key stamped, even by an + actor-less pre-released_at pod mid-deploy, reads stopped and no spend or count + arithmetic can reclassify it. Completion is the window passing or every key + spending its budget; spend is only ever consulted for unstamped keys.""" if self.stopped_by is not None: return "stopped" if datetime.now(timezone.utc) >= ( self.ends_at if self.ends_at.tzinfo else self.ends_at.replace(tzinfo=timezone.utc) ): return "completed" - if all(key.budget_spent for key in self.keys): - return "completed" if all(key.stopped_at is not None for key in self.keys): return "stopped" + if all(key.budget_spent for key in self.keys): + return "completed" return "running" diff --git a/schema.prisma b/schema.prisma index d9959677116..4612921c712 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1509,6 +1509,7 @@ model LiteLLM_ShadowEvalJob { ends_at DateTime stopped_at DateTime? stopped_by String? // operator who stopped it early; null when it ended on its own + released_at DateTime? // frees the key's slot in the partial unique index; stops and sweeps both set it, only stops set stopped_at @@index([group_id]) @@index([api_key_id]) diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 6165e869989..1acd7259eaa 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -525,6 +525,7 @@ def _leg_record(**overrides: object) -> MagicMock: "ends_at": datetime.now(timezone.utc) + timedelta(days=7), "stopped_at": None, "stopped_by": None, + "released_at": None, } fields = {**defaults, **overrides} record = MagicMock(spec=list(fields)) @@ -586,6 +587,8 @@ def _shadow_prisma(legs=(), agg_rows=None, by_leg_rows=None, known_keys=("key-ha current = [row for row in current if row.direction == w["direction"]] if "stopped_at" in w: current = [row for row in current if row.stopped_at is w["stopped_at"]] + if "released_at" in w: + current = [row for row in current if row.released_at is w["released_at"]] if "group_id" in w: wanted = w["group_id"]["in"] if isinstance(w["group_id"], dict) else [w["group_id"]] current = [row for row in current if row.group_id in wanted] @@ -615,6 +618,7 @@ def _shadow_prisma(legs=(), agg_rows=None, by_leg_rows=None, known_keys=("key-ha "ends_at", "stopped_at", "stopped_by", + "released_at", ) return {field: getattr(row, field) for field in fields} @@ -668,9 +672,10 @@ async def test_start_shadow_eval_writes_one_leg_per_key_in_one_statement(monkeyp response = await start_shadow_eval(_start_request(api_key_ids=("key-hash", "key-hash-2")), ADMIN) sweep_sql, sweep_keys = prisma.db.execute_raw.call_args.args - assert "stopped_at IS NULL" in sweep_sql + assert "released_at IS NULL" in sweep_sql + assert "SET released_at = (NOW() AT TIME ZONE 'utc')" in sweep_sql + assert "SET stopped_at" not in sweep_sql assert "j.ends_at <= (NOW() AT TIME ZONE 'utc')" in sweep_sql - assert "SET stopped_at = (NOW() AT TIME ZONE 'utc')" in sweep_sql assert ">= j.max_turns" in sweep_sql assert "j.max_budget IS NOT NULL" in sweep_sql assert ">= j.max_budget" in sweep_sql @@ -760,7 +765,16 @@ async def test_start_shadow_eval_reuses_a_key_whose_previous_job_already_stopped that forgets that would strand every key that has ever finished a job.""" import litellm.proxy.proxy_server as proxy_server - prisma = _shadow_prisma(legs=[_leg_record(group_id="job-7", stopped_at=datetime.now(timezone.utc))]) + prisma = _shadow_prisma( + legs=[ + _leg_record( + group_id="job-7", + stopped_at=datetime.now(timezone.utc), + stopped_by="admin", + released_at=datetime.now(timezone.utc), + ) + ] + ) monkeypatch.setattr(proxy_server, "prisma_client", prisma) monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) @@ -1129,6 +1143,59 @@ def test_a_start_request_still_sending_max_turns_is_rejected_not_silently_defaul _start_request(max_turns=200) +@pytest.mark.asyncio +async def test_sweep_released_job_reads_completed_not_stopped(monkeypatch: pytest.MonkeyPatch): + """The sweep frees a finished key's slot through released_at and never writes + stopped_at, so a budget-completed job stays completed after the sweep runs.""" + import litellm.proxy.proxy_server as proxy_server + + prisma = _shadow_prisma( + legs=[_leg_record(max_budget=1.0, released_at=datetime.now(timezone.utc))] + ) + prisma.attempt_rows = [{"job_id": "leg-1", "attempt_count": 3, "spend": 1.2}] + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + jobs = await list_shadow_eval_jobs(VIEWER, api_key_id=None, limit=50) + assert jobs[0].status == "completed" + assert all(key.stopped_at is None for key in jobs[0].keys) + + +@pytest.mark.asyncio +async def test_a_stamped_key_reads_stopped_even_over_budget(monkeypatch: pytest.MonkeyPatch): + """stopped_at means a stop, unconditionally: only pre-released_at pods stamp it + without an actor, and no spend arithmetic may reclassify their stop as completion.""" + import litellm.proxy.proxy_server as proxy_server + + prisma = _shadow_prisma( + legs=[_leg_record(max_budget=1.0, stopped_at=datetime.now(timezone.utc))] + ) + prisma.attempt_rows = [{"job_id": "leg-1", "attempt_count": 3, "spend": 1.2}] + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + jobs = await list_shadow_eval_jobs(VIEWER, api_key_id=None, limit=50) + assert jobs[0].status == "stopped" + + +def test_released_at_migration_moves_slot_stamps_off_stopped_at(): + """Sweep stamps written before released_at existed were slot bookkeeping, not stops: + the migration must move them to the new column, clear actor-less stopped_at, and + re-anchor the partial unique index, or every swept job would flip to stopped.""" + import litellm_proxy_extras + + sql = ( + Path(litellm_proxy_extras.__file__).parent + / "migrations" + / "20260820150000_shadow_eval_released_at" + / "migration.sql" + ).read_text() + assert 'ADD COLUMN IF NOT EXISTS "released_at" TIMESTAMP(3)' in sql + assert "SET released_at = stopped_at" in sql + assert "WHERE stopped_at IS NOT NULL AND released_at IS NULL" in sql + assert "SET stopped_at = NULL" in sql + assert "WHERE stopped_at IS NOT NULL AND stopped_by IS NULL" in sql + assert 'WHERE "released_at" IS NULL' in sql + + def test_max_budget_migration_is_additive_and_leaves_legacy_rows_null(): """max_budget stays NULL on pre-migration rows so they keep the turn budget they were configured with, and shadow_cost defaults to 0 so old rows price as judge-only.""" @@ -1270,6 +1337,7 @@ async def test_stop_shadow_eval_stops_every_unstopped_leg_and_rejects_non_runnin assert stopped.stopped_by == "admin" stop_sql, stop_group, stop_operator, stop_stamp = prisma.db.execute_raw.call_args.args assert "SET stopped_by = $2, stopped_at = COALESCE(stopped_at, $3::timestamp)" in stop_sql + assert "released_at = COALESCE(released_at, $3::timestamp)" in stop_sql assert "WHERE group_id = $1 AND stopped_by IS NULL" in stop_sql assert "ends_at > (NOW() AT TIME ZONE 'utc')" in stop_sql assert ") < k.max_turns" in stop_sql diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 847fe65a5cc..44a2b3e689c 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -33390,17 +33390,18 @@ export interface components { shadow_percentage: number; /** * Status - * @description Three recorded facts, no history-guessing: a stop is stopped_by (the migration - * backfills it for every job that displayed stopped when the column arrived, so the - * pre-column population is closed), completion is the window passing or every key - * spending its budget, and anything else is running. The all-keys-stamped fallback - * covers only stops written by pre-column pods during a rolling deploy. + * @description Recorded facts only, no history-guessing. A key's stopped_at now means exactly + * one thing, a stop: budget and window completion free the key's slot through the + * internal released_at column and never touch it, so every key stamped, even by an + * actor-less pre-released_at pod mid-deploy, reads stopped and no spend or count + * arithmetic can reclassify it. Completion is the window passing or every key + * spending its budget; spend is only ever consulted for unstamped keys. * @enum {string} */ readonly status: "running" | "completed" | "stopped"; /** * Stopped By - * @description The operator who stopped the job early, recorded by the stop endpoint; 'unknown' backfilled by migration for jobs that displayed stopped when the column arrived; None when the job ended on its own. Its presence is what makes a job read stopped rather than completed + * @description The operator who stopped the job early, recorded by the stop endpoint; 'unknown' backfilled by migration for jobs that displayed stopped when the column arrived; None when the job ended on its own. Budget and window completion never write the stop record: slot bookkeeping lives in an internal column instead */ stopped_by?: string | null; };