fix(alerts): retry delivery and retain distinct retirement models

This commit is contained in:
moe-berri 2026-11-07 11:24:29 +08:00
parent d44a94ee00
commit 4c4f2bfc19
7 changed files with 237 additions and 65 deletions

View file

@ -414,7 +414,11 @@ def _scan_deprecation_model(
position: Final = positions[cursor.group_position]
candidate: Final = info(position)
advanced: Final = cursor.model_copy(update={"group_position": cursor.group_position + 1})
if candidate is None or candidate.info.deprecation_date != cursor.model.deprecation_date:
if candidate is None or (
candidate.info.deprecation_date,
candidate.info.litellm_model,
candidate.info.litellm_provider,
) != (cursor.model.deprecation_date, cursor.model.litellm_model, cursor.model.litellm_provider):
return advanced, None
if position < cursor.deployment:
return DeprecationScanCursor(deployment=cursor.deployment + 1), None

View file

@ -229,7 +229,7 @@ class PerformanceIssueProjector:
while asyncio.get_running_loop().time() < projection_deadline:
if (next_position := await self._projection_page(position, work, observed_at)) is None:
await self._store.complete_projection(work.day, position.version)
await self._notify(deadline)
await self.notify(deadline)
await self._store.prune()
if position.version != work.version:
await rules.save_projection_cursor(UsageAlertProjectionCursor(version=work.version, day=work.day))
@ -238,7 +238,7 @@ class PerformanceIssueProjector:
position = next_position
await rules.save_projection_cursor(position)
await asyncio.sleep(0.01)
await self._notify(deadline)
await self.notify(deadline)
return False
async def _projection_page(
@ -705,7 +705,7 @@ class PerformanceIssueProjector:
while batch := tuple(islice(bindings, 10)):
await self._store.queue_deliveries(batch)
async def _notify(self, stop_at: float | None = None) -> None:
async def notify(self, stop_at: float | None = None) -> None:
deadline: Final = stop_at if stop_at is not None else asyncio.get_running_loop().time() + 60
while asyncio.get_running_loop().time() + 2 < deadline and (
pending := await self._store.pending_deliveries(limit=10)

View file

@ -1,6 +1,7 @@
"""Once-daily evaluation, polled hourly by one database-elected worker."""
import asyncio
import re
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import date, datetime, time, timedelta, timezone
@ -59,7 +60,8 @@ def _has_activity_page(scope: UsageAlertScope, entity: str) -> bool:
if not entity.strip():
return False
if scope == "key":
return entity != "unattributed" and not entity.removeprefix("raw:").startswith("sk-")
key: Final = re.sub(r"^(?:raw:)+", "", entity)
return key not in ("", "unattributed") and not key.startswith("sk-")
if scope == "user":
return entity != "default_user_id"
if scope == "team":
@ -224,9 +226,9 @@ class UsageAlertJob:
)
if self._performance_projection is None:
await self._notify(deadline)
elif (
work := await self._store.pending_projection_work()
) is not None and await self._performance_projection.run_pending(now, work, self._store, deadline):
elif (work := await self._store.pending_projection_work()) is None:
await self._performance_projection.notify(deadline)
elif await self._performance_projection.run_pending(now, work, self._store, deadline):
await self._store.finish_projection_work(work)
await self._store.prune()

View file

@ -260,6 +260,105 @@ async def test_key_spend_issue_consolidates_traffic_and_uses_fresh_days(db: Mana
assert (await issues.list_issues(scope_type="key", scope_id=_KEY_B)).data[0].active
@pytest.mark.asyncio(loop_scope="session")
async def test_due_issue_delivery_retries_after_daily_projection_is_complete(db: ManagementDatabase) -> None:
async with _isolated(db) as prisma:
await _seed(prisma)
rules = UsageAlertsRepository.from_prisma(prisma)
rule = await rules.save_rule(
UsageAlertRuleInput(
name="Key spend notification",
scope_type="key",
scope_id=_KEY_A,
metric="spend",
direction="increase",
notification_channels=("slack",),
),
"qa-admin",
)
assert rule is not None
job = UsageAlertJob.from_prisma(prisma, SlackAlerting(alerting=[], alert_types=[]))
await job._run(_NOW)
assert await rules.pending_projection_work() is None
before = await prisma.writer_db.query_raw(
"SELECT id,(record->>'attempts')::int AS attempts,record->>'last_error' AS error "
'FROM "LiteLLM_IssueDelivery" WHERE NOT sent'
)
assert len(before) == 1
assert (before[0]["attempts"], before[0]["error"]) == (1, "transport_not_configured")
await prisma.writer_db.query_raw(
"UPDATE \"LiteLLM_IssueDelivery\" SET next_attempt=CURRENT_TIMESTAMP-interval '1 second' "
"WHERE id=$1 RETURNING id",
before[0]["id"],
)
await job._run(_NOW + timedelta(hours=1))
after = await prisma.writer_db.query_raw(
"SELECT (record->>'attempts')::int AS attempts FROM \"LiteLLM_IssueDelivery\" WHERE id=$1",
before[0]["id"],
)
assert after == [{"attempts": 2}], "due delivery was not retried after projection completed"
assert await rules.pending_projection_work() is None
@pytest.mark.asyncio(loop_scope="session")
async def test_key_evaluation_rejects_normalized_sentinels_and_secret_shaped_ids(db: ManagementDatabase) -> None:
keys = (
"unattributed",
"raw:unattributed",
"sk-synthetic-only",
"raw:sk-synthetic-only",
"raw:raw:sk-synthetic-only",
_KEY_A,
)
async with _isolated(db) as prisma:
await prisma.writer_db.query_raw(
"""INSERT INTO "LiteLLM_DailyUserSpend" (id,user_id,date,api_key,model,spend,api_requests,updated_at)
SELECT 'key-identity-'||n||'-'||d,'owner',($1::date-d)::text,key,'example-model',
CASE WHEN d=0 THEN 100 ELSE 1 END,1000,CURRENT_TIMESTAMP
FROM unnest($2::text[]) WITH ORDINALITY AS keys(key,n) CROSS JOIN generate_series(0,7) d RETURNING id""",
_DAY.isoformat(),
keys,
)
await _coverage(prisma, _DAY)
job = UsageAlertJob.from_prisma(prisma, SlackAlerting(alerting=[], alert_types=[]))
await job._run(_NOW)
alerts = await UsageAlertsRepository.from_prisma(prisma).list_alerts(scope_type="key")
assert alerts.total == 1, "normalized invalid keys created incidents"
assert alerts.data[0].scope_id == _KEY_A
@pytest.mark.asyncio(loop_scope="session")
async def test_retirement_detects_used_model_after_unused_model_in_the_same_alias(db: ManagementDatabase) -> None:
from litellm import Router
from litellm.proxy.common_utils.model_deprecation import RouterDeprecationSource
from litellm.proxy.spend_tracking.deprecation_issues import DeprecationIssueProjector
models = ("openai/unused-retirement-example", "openai/used-retirement-example")
async with _isolated(db) as prisma:
router = Router(
model_list=[
{
"model_name": "shared-retirement-alias",
"litellm_params": {"model": model},
"model_info": {"id": model, "deprecation_date": "2026-11-20"},
}
for model in models
]
)
await prisma.writer_db.query_raw(
"""INSERT INTO "LiteLLM_DailyUserSpend" (id,user_id,date,api_key,model,api_requests,updated_at)
VALUES('used-retirement','owner',$1,$2,$3,7,CURRENT_TIMESTAMP) RETURNING id""",
_DAY.isoformat(),
_KEY_A,
models[1],
)
await DeprecationIssueProjector.from_prisma(prisma, RouterDeprecationSource(lambda: router)).run(_NOW)
issues = await IssuesRepository.from_prisma(prisma).list_issues(kind="model_deprecation")
assert issues.total == 1, "retired model with actual usage was lost"
assert issues.data[0].cohort.scope_id == models[1]
assert issues.data[0].impact.requests == 7
@pytest.mark.asyncio(loop_scope="session")
async def test_single_key_reliability_collapses_matching_scopes_without_traffic_growth(db: ManagementDatabase) -> None:
async with _isolated(db) as prisma:

View file

@ -0,0 +1,103 @@
import asyncio
import smtplib
import threading
from contextlib import suppress
from datetime import date, datetime, timezone
from typing import Final
import pytest
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
from litellm.integrations.SlackAlerting.usage_alerts import UsageAlertDelivery
from litellm.types.proxy.usage_alerts import UsageAlert, UsageAlertRule
_NOW: Final = datetime(2026, 10, 20, tzinfo=timezone.utc)
_RULE: Final = UsageAlertRule(
name="Request spike",
id="rule",
scope_type="key",
scope_id="reporting-key",
metric="requests",
direction="increase",
created_at=_NOW,
updated_at=_NOW,
created_by=None,
notification_channels=("slack",),
)
_ALERT: Final = UsageAlert(
id="incident",
rule_id="rule",
rule_name="Request spike",
scope_type="key",
scope_id="reporting-key",
metric="requests",
direction="increase",
status="open",
day=date(2026, 10, 20),
observed_value=1000,
baseline_value=200,
threshold=600,
first_seen=_NOW,
last_seen=_NOW,
baseline_start=date(2026, 10, 13),
baseline_end=date(2026, 10, 19),
baseline_sample_days=7,
activity_path="/api-keys?key=reporting-key&detail_tab=activity",
)
@pytest.mark.asyncio
async def test_cancelled_smtp_attempt_retains_only_one_send(monkeypatch: pytest.MonkeyPatch) -> None:
connected: Final = asyncio.Event()
release: Final = asyncio.Event()
accepted: Final[list[asyncio.StreamWriter]] = []
async def stalled(_: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
accepted.append(writer)
connected.set()
await release.wait()
writer.close()
await writer.wait_closed()
server: Final = await asyncio.start_server(stalled, "127.0.0.1", 0)
monkeypatch.setenv("SMTP_HOST", "127.0.0.1")
monkeypatch.setenv("SMTP_PORT", str(server.sockets[0].getsockname()[1]))
monkeypatch.setenv("SMTP_SENDER_EMAIL", "alerts@example.com")
monkeypatch.setenv("SMTP_TLS", "False")
monkeypatch.setenv("SMTP_USE_SSL", "False")
monkeypatch.delenv("SMTP_USERNAME", raising=False)
monkeypatch.delenv("SMTP_PASSWORD", raising=False)
slack: Final = SlackAlerting(alerting=["email"], alert_types=[])
delivery: Final = UsageAlertDelivery(slack)
rule: Final = _RULE.model_copy(
update={"notification_channels": ("email",), "email_recipients": ("qa@example.com",)}
)
try:
sending = asyncio.create_task(delivery.send(_ALERT, rule, "email"))
await asyncio.wait_for(connected.wait(), timeout=10)
sending.cancel()
with pytest.raises(asyncio.CancelledError):
await sending
threads: Final = threading.active_count()
pending: Final = UsageAlertDelivery._email_task
assert pending is not None and not pending.done()
for _ in range(100):
with pytest.raises(RuntimeError, match="still in flight"):
await delivery.send(_ALERT, rule, "email")
assert UsageAlertDelivery._email_task is pending
assert len(accepted) == 1 and threading.active_count() == threads
with pytest.raises(smtplib.SMTPServerDisconnected, match="timed out"):
await asyncio.wait_for(asyncio.shield(pending), timeout=10)
assert pending.done()
finally:
release.set()
server.close()
await server.wait_closed()
for writer in accepted:
writer.close()
pending = UsageAlertDelivery._email_task
if pending is not None:
with suppress(smtplib.SMTPServerDisconnected):
await asyncio.wait_for(asyncio.shield(pending), timeout=10)
UsageAlertDelivery._email_task = None
UsageAlertDelivery._email_key = None

View file

@ -1,4 +1,3 @@
import smtplib
from datetime import date, datetime, timezone
from typing import Final
@ -164,57 +163,3 @@ async def test_slack_source_labels_cannot_inject_mentions_or_links() -> None:
assert b"&lt;@someone&gt;" in requests[0].content
finally:
await handler.close()
@pytest.mark.asyncio
async def test_cancelled_smtp_attempt_retains_only_one_send(monkeypatch: pytest.MonkeyPatch) -> None:
import asyncio
import threading
connected: Final = asyncio.Event()
release: Final = asyncio.Event()
accepted: Final[list[asyncio.StreamWriter]] = []
async def stalled(_: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
accepted.append(writer)
connected.set()
await release.wait()
writer.close()
await writer.wait_closed()
server: Final = await asyncio.start_server(stalled, "127.0.0.1", 0)
monkeypatch.setenv("SMTP_HOST", "127.0.0.1")
monkeypatch.setenv("SMTP_PORT", str(server.sockets[0].getsockname()[1]))
monkeypatch.setenv("SMTP_SENDER_EMAIL", "alerts@example.com")
monkeypatch.setenv("SMTP_TLS", "False")
monkeypatch.setenv("SMTP_USE_SSL", "False")
monkeypatch.delenv("SMTP_USERNAME", raising=False)
monkeypatch.delenv("SMTP_PASSWORD", raising=False)
slack: Final = SlackAlerting(alerting=["email"], alert_types=[])
delivery: Final = UsageAlertDelivery(slack)
rule: Final = _RULE.model_copy(
update={"notification_channels": ("email",), "email_recipients": ("qa@example.com",)}
)
try:
with pytest.raises(asyncio.TimeoutError):
await asyncio.wait_for(delivery.send(_ALERT, rule, "email"), timeout=0.1)
await asyncio.wait_for(connected.wait(), timeout=1)
threads: Final = threading.active_count()
pending: Final = UsageAlertDelivery._email_task
assert pending is not None and not pending.done()
for _ in range(100):
with pytest.raises(RuntimeError, match="still in flight"):
await delivery.send(_ALERT, rule, "email")
assert UsageAlertDelivery._email_task is pending
assert len(accepted) == 1 and threading.active_count() == threads
with pytest.raises(smtplib.SMTPServerDisconnected, match="timed out"):
await asyncio.wait_for(asyncio.shield(pending), timeout=3)
assert pending.done()
finally:
release.set()
server.close()
await server.wait_closed()
for writer in accepted:
writer.close()
UsageAlertDelivery._email_task = None
UsageAlertDelivery._email_key = None

View file

@ -762,7 +762,7 @@ def test_retirement_pages_resume_inside_large_alias_without_losing_configured_su
model_list=[
{
"model_name": "alias",
"litellm_params": {"model": f"openai/paged-model-{index}"},
"litellm_params": {"model": "openai/paged-model"},
"model_info": {
"id": f"paged-{index}",
"deprecation_date": "2026-11-20",
@ -790,10 +790,29 @@ def test_retirement_pages_resume_inside_large_alias_without_losing_configured_su
expected = collect_model_deprecations(router, today=today).imminent
assert models == tuple(expected)
assert len(models) == 1
assert models[0].litellm_model == "openai/paged-model-0"
assert models[0].litellm_model == "openai/paged-model"
assert models[0].successor_model == "replacement"
def test_retirement_pages_preserve_distinct_models_sharing_an_alias_and_date() -> None:
router = Router(
model_list=[
{
"model_name": "shared-alias",
"litellm_params": {"model": model},
"model_info": {"id": model, "deprecation_date": "2026-11-20"},
}
for model in ("openai/unused-retired-model", "openai/used-retired-model")
]
)
page = RouterDeprecationSource(lambda: router).page(date(2026, 11, 5), DeprecationScanCursor())
assert page.complete
assert tuple(model.litellm_model for model in page.models) == (
"openai/unused-retired-model",
"openai/used-retired-model",
)
def test_retirement_source_bounds_pages_and_restarts_identity_after_router_replacement() -> None:
router = Router(
model_list=[