feat(shadow_eval)!: gate the per-key budget on dollar spend instead of turns

litellm-proxy-extras/litellm_proxy_extras/migrations/20260819000000_shadow_eval_max_budget/migration.sql

@@ -0,0 +1,5 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_ShadowEvalJob" ADD COLUMN     "max_budget" DOUBLE PRECISION;
+
+-- AlterTable
+ALTER TABLE "LiteLLM_ShadowEvalAttempt" ADD COLUMN     "shadow_cost" DOUBLE PRECISION NOT NULL DEFAULT 0;

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -1502,7 +1502,8 @@ model LiteLLM_ShadowEvalJob {
   baseline_model    String?  // reverse only: the fixed model the router is judged against
   judge_model       String
   shadow_percentage Float
-  max_turns         Int      // this key's sample budget: judge at most this many turns
+  max_turns         Int      // sample-count ceiling: the whole budget on pre-max_budget jobs, the error-loop valve otherwise
+  max_budget        Float?   // per-key USD cap on the eval's own shadow + judge spend; null on jobs from before spend budgets
   created_at        DateTime @default(now())
   created_by        String?
   ends_at           DateTime
@@ -1525,6 +1526,7 @@ model LiteLLM_ShadowEvalAttempt {
   shadow_model String?
   confidence   Float?
   judge_cost   Float    @default(0)
+  shadow_cost  Float    @default(0)
   error        String?
   created_at   DateTime @default(now())
 

litellm/integrations/shadow_eval_logger.py

@@ -10,7 +10,7 @@
 import hashlib
 import random
 import traceback
-from collections.abc import Callable, Mapping, Sequence
+from collections.abc import Awaitable, Callable, Mapping, Sequence
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from itertools import groupby
@@ -42,8 +42,9 @@
     from litellm.router import Router
     from litellm.types.utils import StandardLoggingPayload
 
-# A job starting, stopping, or hitting its turn budget propagates to sampling within one
-# TTL; the turn budget can overshoot by at most one TTL of in-flight samples per pod.
+# A job starting, stopping, or hitting a budget propagates to sampling within one TTL;
+# the spend gate re-checks the cross-pod counter at pipeline entry, so it overshoots
+# only by the samples already in flight when the cap is crossed.
 _JOBS_CACHE_TTL_SECONDS: Final = 10
 
 # Concurrent shadow+judge pipelines per pod: a traffic spike turns into skipped samples
@@ -340,13 +341,24 @@ def _failure_detail(e: BaseException) -> str:
     return f"{type(e).__name__}{location}: {e}"
 
 
-def _judge_call_cost(response: object) -> float:
-    """Price a judge call, treating an unmapped judge model as free rather than fatal."""
+def _call_cost(response: object) -> float:
+    """Price one eval-arm call with the figure the spend pipeline bills: the router client
+    stamps _hidden_params.response_cost from the deployment's own pricing, which the public
+    price map lookup below cannot see (it reads 0 for deployment-priced models)."""
+    getter: Final = getattr(getattr(response, "_hidden_params", None), "get", None)
+    stamped: Final = getter("response_cost") if callable(getter) else None
+    if isinstance(stamped, (int, float)):
+        return float(stamped)
+    return _price_map_cost(response)
+
+
+def _price_map_cost(response: object) -> float:
+    """Public price map fallback, treating an unmapped model as free rather than fatal."""
     import litellm
 
     try:
         return litellm.completion_cost(completion_response=response) or 0.0
-    except Exception:  # noqa: BLE001  # unmapped judge model: the verdict still counts, cost stays 0
+    except Exception:  # noqa: BLE001  # unmapped model: the attempt still counts, cost stays 0
         return 0.0
 
 
@@ -374,6 +386,32 @@ def _judge_user_prompt(conversation: str, response_a: str, response_b: str) -> s
     )
 
 
+def _job_spend_counter_key(job_id: str) -> str:
+    return f"spend:shadow_eval:{job_id}"
+
+
+async def _job_spend_from_counter(counter_key: str, fallback_spend: float, max_budget: float) -> float:
+    """The leg's spend through the cross-pod counter the key budget gates read. The owner
+    degrades internally to the fill-time DB floor and raises only under fail-closed
+    enforcement, which the caller honors by skipping the sample."""
+    from litellm.proxy.proxy_server import get_current_spend
+
+    return await get_current_spend(counter_key=counter_key, fallback_spend=fallback_spend, max_budget=max_budget)
+
+
+async def _add_job_spend_to_counter(counter_key: str, cost: float) -> None:
+    """Advance the counter the moment a cost is known, so even a lost row closes the gate.
+    Known failure mode: a Redis outage freezes the counter (the owner invalidates it), the
+    gate degrades to the fill floor, and overshoot grows to in-flight plus one TTL of
+    samples, the same degradation the key budget counters accept."""
+    try:
+        from litellm.proxy.proxy_server import increment_spend_counter
+
+        await increment_spend_counter(counter_key=counter_key, increment=cost)
+    except Exception as e:  # noqa: BLE001  # attempt recording must proceed; the row stays truth and the fill floor gates
+        verbose_logger.warning("shadow_eval: spend counter increment failed for %s: %s", counter_key, e)
+
+
 async def _key_or_team_is_over_budget(metadata: Mapping[str, object]) -> bool:
     """Whether the shadowed key or its team is over budget, decided by the same owners
     the request path uses, so counter keys and thresholds can never drift from auth's.
@@ -438,8 +476,8 @@ def _request_was_routed_by(request_metadata: Mapping[str, object], router_name:
 
 @dataclass(frozen=True, slots=True)
 class _CallFailure:
-    """A shadow or judge call that produced no usable response. cost carries any judge
-    spend the failed attempt still billed, so job-level judge_spend never undercounts."""
+    """A shadow or judge call that produced no usable response. cost carries any spend
+    the failed call still billed, so job-level spend figures never undercount."""
 
     error: str
     cost: float = 0.0
@@ -452,6 +490,7 @@ class _ShadowResponse:
     text: str
     model: str
     tier: str | None
+    cost: float
 
 
 @dataclass(frozen=True, slots=True)
@@ -478,8 +517,10 @@ class ActiveShadowEvalJob(BaseModel):
     shadow_percentage: float
     judge_model: str
     max_turns: int
+    max_budget: float | None = None
     ends_at: datetime
     attempts: int = 0
+    spend: float = 0.0
 
     @field_validator("ends_at")
     @classmethod
@@ -500,7 +541,7 @@ def shadow_target(self) -> str:
         return self.baseline_model or self.router_name
 
 
-def _as_active_job(record: object, attempts: int) -> ActiveShadowEvalJob | None:
+def _as_active_job(record: object, attempts: int, spend: float) -> ActiveShadowEvalJob | None:
     """The sampling path's view of one job row, or None for a row it cannot sample: an
     unknown direction, or a reverse job with no baseline model to duplicate against.
     Failing closed here is what keeps the dispatch path total."""
@@ -509,7 +550,7 @@ def _as_active_job(record: object, attempts: int) -> ActiveShadowEvalJob | None:
     except ValidationError as e:
         verbose_logger.debug("shadow_eval: skipping unsamplable job row: %s", e)
         return None
-    return job.model_copy(update={"attempts": attempts})
+    return job.model_copy(update={"attempts": attempts, "spend": spend})  # mutable-ok: pydantic update payload
 
 
 _jobs_cache: Final = InMemoryCache(max_size_in_memory=4, default_ttl=_JOBS_CACHE_TTL_SECONDS)
@@ -524,12 +565,17 @@ def __init__(
         router_provider: Callable[[], "Router | None"] | None = None,
         prisma_provider: Callable[[], "PrismaClient | None"] | None = None,
         jobs_cache: InMemoryCache | None = None,
+        job_spend_reader: Callable[[str, float, float], Awaitable[float]] | None = None,
+        job_spend_writer: Callable[[str, float], Awaitable[None]] | None = None,
     ) -> None:
         """Providers are callables so the proxy's lazily-initialized globals are resolved
-        at call time, not at logger construction."""
+        at call time, not at logger construction. The spend reader and writer wrap the
+        proxy's cross-pod spend counter; tests inject a plain in-memory pair."""
         self._router_provider = router_provider or default_router_provider
         self._prisma_provider = prisma_provider or _default_prisma_provider
         self._jobs_cache = jobs_cache or _jobs_cache
+        self._read_job_spend = job_spend_reader or _job_spend_from_counter
+        self._write_job_spend = job_spend_writer or _add_job_spend_to_counter
         self._inflight_shadow_tasks: int = 0
         # Starts per job since the last cache fill, never decremented within a
         # generation; the refill absorbs written rows and resets.
@@ -556,18 +602,26 @@ async def _active_jobs(self) -> Mapping[str, tuple[ActiveShadowEvalJob, ...]]:
                 await prisma.db.litellm_shadowevalattempt.group_by(
                     by=["job_id"],
                     count=True,
+                    sum={"judge_cost": True, "shadow_cost": True},  # mutable-ok: Prisma aggregate spec
                     where={"job_id": {"in": [str(record.id) for record in records]}},  # mutable-ok: Prisma filter
                 )
                 if records
                 else ()
             )
-            attempt_counts: Final = {str(row["job_id"]): int(row["_count"]["_all"]) for row in grouped or []}
+            attempt_stats: Final = {  # mutable-ok: frozen snapshot of the grouped read
+                str(row["job_id"]): (
+                    int(row["_count"]["_all"]),
+                    float((row["_sum"] or {}).get("judge_cost") or 0.0)
+                    + float((row["_sum"] or {}).get("shadow_cost") or 0.0),
+                )
+                for row in grouped or []
+            }
             by_key: Final = tuple(
                 sorted(
                     (
                         (str(record.api_key_id), job)
                         for record in records or []
-                        if (job := _as_active_job(record, attempt_counts.get(str(record.id), 0))) is not None
+                        if (job := _as_active_job(record, *attempt_stats.get(str(record.id), (0, 0.0)))) is not None
                     ),
                     key=itemgetter(0),
                 )
@@ -624,6 +678,7 @@ async def async_log_success_event(
                 for job in (await self._active_jobs()).get(str(api_key_hash), ())
                 if datetime.now(timezone.utc) < job.ends_at
                 and job.attempts + self._job_starts.get(job.id, 0) < job.max_turns
+                and (job.max_budget is None or job.spend < job.max_budget)
                 and _sample_hits(request_id, job.id, job.shadow_percentage)
                 and _request_was_routed_by(request_metadata, job.router_name) == (job.direction == "reverse")
             )
@@ -684,12 +739,28 @@ async def _run_shadow_eval(
                 return
             if await _key_or_team_is_over_budget(parent_metadata):
                 return
-
+            if job.max_budget is not None:
+                try:
+                    spend: Final = await self._read_job_spend(_job_spend_counter_key(job.id), job.spend, job.max_budget)
+                except Exception as e:  # noqa: BLE001  # unverifiable budget: skip the sample rather than spend on it
+                    verbose_logger.warning("shadow_eval: budget unverifiable for %s, sample skipped: %s", job.id, e)
+                    return
+                if spend >= job.max_budget:
+                    return
             shadow: Final = await self._call_router_shadow(job.shadow_target, messages, shadow_params, parent_metadata)
-            if isinstance(shadow, _CallFailure):
-                await self._record_attempt(prisma, job, request_id, control_tier, outcome="error", error=shadow.error)
-                return
-
+        except Exception as e:  # noqa: BLE001  # detached task: nothing billed yet, record and never raise
+            verbose_logger.debug("shadow_eval: pipeline failed for %s: %s", request_id, e)
+            await self._record_attempt(
+                prisma, job, request_id, control_tier, outcome="error", error=f"pipeline error: {e}"
+            )
+            return
+        if isinstance(shadow, _CallFailure):
+            await self._record_attempt(
+                prisma, job, request_id, control_tier, outcome="error", error=shadow.error, shadow_cost=shadow.cost
+            )
+            return
+        # From here the shadow call has billed, so every exit records its cost.
+        try:
             verdict: Final = await self._call_judge(
                 judge_model=job.judge_model,
                 messages=messages,
@@ -707,6 +778,7 @@ async def _run_shadow_eval(
                     error=verdict.error,
                     shadow=shadow,
                     judge_cost=verdict.cost,
+                    shadow_cost=shadow.cost,
                 )
                 return
             await self._record_attempt(
@@ -719,15 +791,23 @@ async def _run_shadow_eval(
                 real_model=real_model,
                 confidence=verdict.confidence,
                 judge_cost=verdict.cost,
+                shadow_cost=shadow.cost,
             )
-        except Exception as e:  # noqa: BLE001  # detached task: record what happened, never raise
+        except Exception as e:  # noqa: BLE001  # detached task: the shadow call billed, record its cost, never raise
             verbose_logger.debug("shadow_eval: pipeline failed for %s: %s", request_id, e)
             await self._record_attempt(
-                prisma, job, request_id, control_tier, outcome="error", error=f"pipeline error: {e}"
+                prisma,
+                job,
+                request_id,
+                control_tier,
+                outcome="error",
+                error=f"pipeline error: {e}",
+                shadow=shadow,
+                shadow_cost=shadow.cost,
             )
 
-    @staticmethod
     async def _record_attempt(
+        self,
         prisma: "PrismaClient | None",
         job: ActiveShadowEvalJob,
         request_id: str,
@@ -738,8 +818,11 @@ async def _record_attempt(
         real_model: str = "",
         confidence: float | None = None,
         judge_cost: float = 0.0,
+        shadow_cost: float = 0.0,
         error: str | None = None,
     ) -> None:
+        if judge_cost + shadow_cost > 0:
+            await self._write_job_spend(_job_spend_counter_key(job.id), judge_cost + shadow_cost)
         if prisma is None:
             return
         try:
@@ -753,6 +836,7 @@ async def _record_attempt(
                     "shadow_model": shadow.model if shadow else None,
                     "confidence": confidence,
                     "judge_cost": judge_cost,
+                    "shadow_cost": shadow_cost,
                     "error": error[:_MAX_ERROR_CHARS] if error else None,
                 }
             )
@@ -792,11 +876,12 @@ async def _call_router_shadow(
             return _CallFailure(f"shadow router call failed: {_failure_detail(e)}")
         text: Final = _chat_final_text(response)
         if not text:
-            return _CallFailure("shadow router returned an empty response")
+            return _CallFailure("shadow router returned an empty response", cost=_call_cost(response))
         return _ShadowResponse(
             text=text,
             model=str(getattr(response, "model", None) or _routing_decision(shadow_metadata).get("routed_model") or ""),
             tier=_routed_tier(shadow_metadata),
+            cost=_call_cost(response),
         )
 
     async def _call_judge(
@@ -843,11 +928,11 @@ async def _call_judge(
             verdict: Final = PairwiseVerdict.model_validate(parse_json_verdict(raw))
         except Exception as e:  # noqa: BLE001  # malformed verdicts become error rows
             verbose_logger.debug("shadow_eval: unparseable judge verdict: %s", e)
-            return _CallFailure(f"unparseable judge verdict: {e}", cost=_judge_call_cost(response))
+            return _CallFailure(f"unparseable judge verdict: {e}", cost=_call_cost(response))
         return _JudgeVerdict(
             preference=_unmask_preference(verdict.preference, real_is_a),
             confidence=max(0.0, min(1.0, verdict.confidence)),
-            cost=_judge_call_cost(response),
+            cost=_call_cost(response),
         )