mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(alerts): retry delivery and retain distinct retirement models
This commit is contained in:
parent
d44a94ee00
commit
4c4f2bfc19
7 changed files with 237 additions and 65 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
103
tests/integration/observability/test_usage_alert_email.py
Normal file
103
tests/integration/observability/test_usage_alert_email.py
Normal 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
|
||||
|
|
@ -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"<@someone>" 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
|
||||
|
|
|
|||
|
|
@ -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=[
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue