From a7015fb185537128fa3c2e293f8441559d045b75 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sun, 8 Nov 2026 15:33:17 +0800 Subject: [PATCH] fix(alerts): rotate bounded work across scope queues --- .../proxy/spend_tracking/usage_alert_job.py | 17 ++++---- litellm/types/proxy/usage_alerts.py | 15 ++++++- .../database/test_performance_issues.py | 41 +++++++++++++++++++ 3 files changed, 63 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/spend_tracking/usage_alert_job.py b/litellm/proxy/spend_tracking/usage_alert_job.py index 374d9ee433a..de9d81a2f58 100644 --- a/litellm/proxy/spend_tracking/usage_alert_job.py +++ b/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: diff --git a/litellm/types/proxy/usage_alerts.py b/litellm/types/proxy/usage_alerts.py index 727bee5770e..b5b3956fdc1 100644 --- a/litellm/types/proxy/usage_alerts.py +++ b/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): diff --git a/tests/integration/database/test_performance_issues.py b/tests/integration/database/test_performance_issues.py index 269a2a83942..addfafd2617 100644 --- a/tests/integration/database/test_performance_issues.py +++ b/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}]