fix(alerts): rotate bounded work across scope queues

litellm/proxy/spend_tracking/usage_alert_job.py

@@ -229,18 +229,19 @@ class UsageAlertJob:
     async def _evaluate(self, target: date, deadline: float) -> None:
         cursor = await self._store.work_cursor()  # rebind-ok: durable round-robin position
         exhausted: Final[set[int]] = set()  # mutable-ok: at most the 10 scope/config phases in this invocation
-        phase_progress = False  # rebind-ok: whether a page advanced in the current phase
+        progressed: Final[set[int]] = set()  # mutable-ok: at most one progress flag per scope/config phase
         while asyncio.get_running_loop().time() < deadline and len(exhausted) < 10:
             if cursor.phase in exhausted:
-                cursor = UsageAlertWorkCursor(phase=(cursor.phase + 1) % 10)
+                cursor = cursor.advance(cursor.after_id)
                 continue
             phase: int = cursor.phase
             cursor, advanced, finished = await self._evaluation_page(cursor, target)
-            phase_progress = phase_progress or advanced
+            if advanced:
+                progressed.add(phase)
             if finished:
-                if not phase_progress:
+                if phase not in progressed:
                     exhausted.add(phase)
-                phase_progress = False
+                progressed.discard(phase)
             await self._store.save_work_cursor(cursor)
             await asyncio.sleep(0.01)
 
@@ -252,7 +253,7 @@ class UsageAlertJob:
         limit: Final = 8 if wildcard else 10
         rules: Final = await self._store.enabled_rule_page(scope, wildcard, cursor.after_id, limit=limit)
         if not rules:
-            return UsageAlertWorkCursor(phase=(cursor.phase + 1) % 10), False, True
+            return cursor.advance(None), False, True
         # Keep one date/cursor per page. A newly edited rule cannot rewind another rule.
         position: Final = self._position(rules[0], target)
         selected: Final = tuple(
@@ -271,8 +272,8 @@ class UsageAlertJob:
         due: Final = tuple(rule for rule in selected if rule.id not in complete)
         advanced: Final = bool(due) and await self._attempt_scope_page(due, position.day, wildcard)
         if len(rules) < limit and len(selected) == len(rules):
-            return UsageAlertWorkCursor(phase=(cursor.phase + 1) % 10), advanced, True
-        return UsageAlertWorkCursor(phase=cursor.phase, after_id=selected[-1].id), advanced, False
+            return cursor.advance(None), advanced, True
+        return cursor.advance(selected[-1].id), advanced, False
 
     async def _attempt_scope_page(self, rules: tuple[UsageAlertRule, ...], target: date, wildcard: bool) -> bool:
         try:

litellm/types/proxy/usage_alerts.py

@@ -1,7 +1,7 @@
 """Daily usage alert rules and incidents shared by the API and evaluator."""
 
 from datetime import date, datetime
-from typing import Literal, TypeAlias
+from typing import Annotated, Final, Literal, TypeAlias
 
 from pydantic import BaseModel, ConfigDict, Field, model_validator
 from typing_extensions import Self
@@ -75,7 +75,18 @@ class UsageAlertProgress(BaseModel, frozen=True):
 
 class UsageAlertWorkCursor(BaseModel, frozen=True):
     phase: int = Field(default=0, ge=0, lt=10)
-    after_id: str | None = None
+    after_id: str | None = Field(default=None, max_length=200)
+    positions: tuple[Annotated[str, Field(max_length=200)] | None, ...] = Field(
+        default=(None,) * 10, min_length=10, max_length=10
+    )
+
+    def advance(self, after_id: str | None) -> "UsageAlertWorkCursor":
+        phase: Final = (self.phase + 1) % 10
+        return UsageAlertWorkCursor(
+            phase=phase,
+            after_id=self.positions[phase],
+            positions=tuple(after_id if index == self.phase else value for index, value in enumerate(self.positions)),
+        )
 
 
 class UsageAlertProjectionWork(BaseModel, frozen=True):

tests/integration/database/test_performance_issues.py

@@ -1046,3 +1046,44 @@ async def test_cancelled_email_receipt_survives_a_different_delivery(
         UsageAlertDelivery._email_task = None
         UsageAlertDelivery._email_key = None
         UsageAlertDelivery._email_store = None
+
+
+@pytest.mark.asyncio(loop_scope="session")
+async def test_large_custom_rule_scope_does_not_starve_other_scopes_on_resume(db: ManagementDatabase) -> None:
+    from litellm.types.proxy.usage_alerts import UsageAlertWorkCursor
+
+    async with _isolated(db) as prisma:
+        await _seed(prisma)
+        store: Final = UsageAlertsRepository.from_prisma(prisma)
+        key_rule: Final = await store.save_rule(
+            UsageAlertRuleInput(
+                name="Project spend", scope_type="key", scope_id=_KEY_A, metric="spend", direction="increase"
+            ),
+            "qa-user",
+        )
+        team_rule: Final = await store.save_rule(
+            UsageAlertRuleInput(
+                name="Team spend", scope_type="team", scope_id="team", metric="spend", direction="increase"
+            ),
+            "qa-admin",
+        )
+        assert key_rule is not None and team_rule is not None
+        await prisma.writer_db.query_raw(
+            """INSERT INTO "LiteLLM_UsageAlertRule" (id,scope_type,scope_id,enabled,record)
+            SELECT 'custom-'||lpad(n::text,6,'0'),scope_type,scope_id,enabled,
+                record||jsonb_build_object('id','custom-'||lpad(n::text,6,'0'))
+            FROM "LiteLLM_UsageAlertRule",generate_series(1,1000) n WHERE id=$1 RETURNING id""",
+            key_rule.id,
+        )
+        await store.save_work_cursor(UsageAlertWorkCursor(phase=3))
+        for _ in range(12):
+            job: Final = UsageAlertJob.from_prisma(prisma, SlackAlerting(alerting=[], alert_types=[]))
+            cursor, _, _ = await job._evaluation_page(await store.work_cursor(), _DAY)
+            await store.save_work_cursor(cursor)
+        assert await store.completed_rule_ids((team_rule,), _DAY) == frozenset((team_rule.id,))
+        completed: Final = await prisma.writer_db.query_raw(
+            """SELECT COUNT(*)::int AS count FROM "LiteLLM_UsageAlertRule"
+            WHERE scope_type='key' AND record->>'last_evaluated_day'=$1""",
+            _DAY.isoformat(),
+        )
+        assert completed == [{"count": 20}]