mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(alerts): resume retirement scans across replicas
This commit is contained in:
parent
d20ec80f90
commit
a75d6335b5
4 changed files with 161 additions and 21 deletions
|
|
@ -1,5 +1,6 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
|
|
@ -14,7 +15,6 @@ from typing import (
|
|||
TypeAlias,
|
||||
cast, # noqa: TID251 # router exposes its live validated deployment list
|
||||
)
|
||||
from uuid import uuid4
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
|
@ -431,6 +431,40 @@ def _scan_deprecation_model(
|
|||
), None
|
||||
|
||||
|
||||
async def _retirement_fingerprint(
|
||||
deployment: Mapping[str, object], router: Router, today: date, catalog: Mapping[str, object]
|
||||
) -> bytes:
|
||||
info: Final = _build_info(
|
||||
deployment, router, today, DEFAULT_DEPRECATION_WARN_DAYS, catalog, logging.DEBUG, check_successor=False
|
||||
)
|
||||
model_name: Final = _as_model_key(deployment.get("model_name"))
|
||||
fields: Final = (
|
||||
(model_name,)
|
||||
if info is None
|
||||
else (
|
||||
model_name,
|
||||
info.info.model_name,
|
||||
info.info.team_public_model_name,
|
||||
info.info.litellm_model,
|
||||
info.info.litellm_provider,
|
||||
info.info.deprecation_date.isoformat(),
|
||||
str(info.info.days_until_deprecation),
|
||||
info.info.successor_model,
|
||||
info.info.successor_model_source,
|
||||
info.info.deprecation_source,
|
||||
)
|
||||
)
|
||||
digest: Final = hashlib.sha256()
|
||||
for value in fields:
|
||||
digest.update(b"\x00" if value is None else b"\x01")
|
||||
if value is not None:
|
||||
digest.update(len(value).to_bytes(8, "big"))
|
||||
for offset in range(0, len(value), 4096):
|
||||
digest.update(value[offset : offset + 4096].encode())
|
||||
await asyncio.sleep(0)
|
||||
return digest.digest()
|
||||
|
||||
|
||||
class RouterDeprecationSource:
|
||||
"""Read pages from router-owned views; never retain a second deployment catalog."""
|
||||
|
||||
|
|
@ -439,14 +473,48 @@ class RouterDeprecationSource:
|
|||
) -> None:
|
||||
self._router = router
|
||||
self._loaded_catalog = loaded_catalog
|
||||
self._instance = str(uuid4())
|
||||
self._generation: tuple[int, int, int, int] | None = None
|
||||
self._revision_day: date | None = None
|
||||
self._revision = ""
|
||||
self._revision_position = 0
|
||||
self._revision_digest = b""
|
||||
|
||||
@property
|
||||
def revision(self) -> str:
|
||||
def _local_generation(self) -> tuple[int, int, int, int]:
|
||||
router: Final = self._router()
|
||||
version: Final = router.registry().version if router is not None else 0
|
||||
model_cost_identity: Final = id(litellm.model_cost) # pyright: ignore[reportUnknownMemberType,reportUnknownArgumentType] # identity only, from the legacy untyped catalog
|
||||
return f"{self._instance}:{id(router)}:{version}:{model_cost_identity}:{id(self._loaded_catalog())}"
|
||||
return id(router), version, model_cost_identity, id(self._loaded_catalog())
|
||||
|
||||
@property
|
||||
def revision(self) -> str:
|
||||
return self._revision if self._generation == self._local_generation() else ""
|
||||
|
||||
async def prepare_revision(self, today: date) -> str | None:
|
||||
router: Final = self._router()
|
||||
if router is None:
|
||||
return None
|
||||
generation: Final = self._local_generation()
|
||||
if self._generation != generation or self._revision_day != today:
|
||||
self._generation = generation
|
||||
self._revision_day = today
|
||||
self._revision = ""
|
||||
self._revision_position = 0
|
||||
self._revision_digest = b""
|
||||
deployments: Final = cast( # cast-ok: router validated each deployment
|
||||
Sequence[Mapping[str, object]], router.registry().all()
|
||||
)
|
||||
catalog: Final = self._loaded_catalog()
|
||||
while self._local_generation() == generation:
|
||||
if self._revision_position >= len(deployments):
|
||||
self._revision = hashlib.sha256(self._revision_digest).hexdigest()
|
||||
return self._revision
|
||||
for position in range(self._revision_position, min(self._revision_position + 100, len(deployments))):
|
||||
self._revision_digest = hashlib.sha256(
|
||||
self._revision_digest + await _retirement_fingerprint(deployments[position], router, today, catalog)
|
||||
).digest()
|
||||
self._revision_position = position + 1
|
||||
await asyncio.sleep(0.01)
|
||||
return None
|
||||
|
||||
def page(self, today: date, cursor: DeprecationScanCursor) -> DeprecationModelPage:
|
||||
router: Final = self._router()
|
||||
|
|
|
|||
|
|
@ -195,15 +195,16 @@ class DeprecationIssueProjector:
|
|||
stored: Final = _ProjectionCheckpoint.model_validate(previous.param_value) if previous is not None else None
|
||||
if stored is not None and stored.complete and stored.hour >= hour:
|
||||
return
|
||||
revision: Final = await self._source.prepare_revision(hour.date())
|
||||
if revision is None:
|
||||
return
|
||||
state = ( # rebind-ok: durable continuation after each completed page
|
||||
stored
|
||||
if stored is not None
|
||||
and not stored.complete
|
||||
and stored.source_revision == self._source.revision
|
||||
and stored.source_revision == revision
|
||||
and stored.hour.date() == hour.date()
|
||||
else _ProjectionCheckpoint(
|
||||
hour=hour, observed_at=now, source_revision=self._source.revision, complete=False
|
||||
)
|
||||
else _ProjectionCheckpoint(hour=hour, observed_at=now, source_revision=revision, complete=False)
|
||||
)
|
||||
while self._source.revision == state.source_revision:
|
||||
page = self._source.page(state.hour.date(), state.cursor)
|
||||
|
|
|
|||
|
|
@ -253,7 +253,7 @@ async def test_retirement_scan_resumes_bounded_pages_without_premature_recovery(
|
|||
projector = DeprecationIssueProjector.from_prisma(db.prisma, source)
|
||||
|
||||
async def interrupt_after_page(
|
||||
projector: DeprecationIssueProjector, observed_at: datetime, revision: str
|
||||
projector: DeprecationIssueProjector, observed_at: datetime, revision: str, after: int = 0
|
||||
) -> _ProjectionCheckpoint:
|
||||
task = asyncio.create_task(projector.run(observed_at))
|
||||
try:
|
||||
|
|
@ -263,7 +263,7 @@ async def test_retirement_scan_resumes_bounded_pages_without_premature_recovery(
|
|||
value = await config.get_param(checkpoint_name)
|
||||
if value is not None:
|
||||
state = _ProjectionCheckpoint.model_validate(value.param_value)
|
||||
if state.source_revision == revision and 0 < state.cursor.deployment < 600:
|
||||
if state.source_revision == revision and after < state.cursor.deployment < 600:
|
||||
return state
|
||||
await asyncio.sleep(0)
|
||||
pytest.fail("projector finished before observing a partial scan")
|
||||
|
|
@ -274,18 +274,24 @@ async def test_retirement_scan_resumes_bounded_pages_without_premature_recovery(
|
|||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
partial = await interrupt_after_page(projector, now, source.revision)
|
||||
revision = await source.prepare_revision(now.date())
|
||||
assert revision is not None
|
||||
partial = await interrupt_after_page(projector, now, revision)
|
||||
assert partial.complete is False
|
||||
assert (await repository.get_issue(old.id)).active is True
|
||||
# A different process/router generation cannot trust the same numeric registry version.
|
||||
next_source = RouterDeprecationSource(lambda: router)
|
||||
replica = Router(model_list=list(router.registry().all()))
|
||||
next_source = RouterDeprecationSource(lambda: replica)
|
||||
assert await next_source.prepare_revision(now.date()) == revision
|
||||
replacement = DeprecationIssueProjector.from_prisma(db.prisma, next_source)
|
||||
restarted = await interrupt_after_page(replacement, now + timedelta(seconds=10), next_source.revision)
|
||||
assert restarted.source_revision != partial.source_revision
|
||||
assert restarted.complete is False
|
||||
resumed = await interrupt_after_page(
|
||||
replacement, now + timedelta(seconds=10), revision, partial.cursor.deployment
|
||||
)
|
||||
assert resumed.observed_at == partial.observed_at
|
||||
assert resumed.cursor.deployment > partial.cursor.deployment
|
||||
assert (await repository.get_issue(old.id)).active is True
|
||||
for iteration in range(10):
|
||||
await replacement.run(now + timedelta(seconds=20 + iteration))
|
||||
current = replacement if iteration % 2 == 0 else projector
|
||||
await current.run(now + timedelta(seconds=20 + iteration))
|
||||
complete = await config.get_param(checkpoint_name)
|
||||
if _ProjectionCheckpoint.model_validate(complete.param_value).complete:
|
||||
break
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ These tests focus on the helper itself — not on the proxy endpoint or
|
|||
Slack integration — so they can run without the full proxy stack.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from collections.abc import Mapping
|
||||
from datetime import date, datetime, timezone
|
||||
|
|
@ -813,7 +814,8 @@ def test_retirement_pages_preserve_distinct_models_sharing_an_alias_and_date() -
|
|||
)
|
||||
|
||||
|
||||
def test_retirement_source_bounds_pages_and_restarts_identity_after_router_replacement() -> None:
|
||||
@pytest.mark.asyncio
|
||||
async def test_retirement_source_bounds_pages_and_resumes_identical_router_replacements() -> None:
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
|
|
@ -832,10 +834,14 @@ def test_retirement_source_bounds_pages_and_restarts_identity_after_router_repla
|
|||
third = source.page(date(2026, 11, 5), second.cursor)
|
||||
assert len(third.models) == 5 and third.complete
|
||||
assert len({model.model_name for model in (*first.models, *second.models, *third.models)}) == 25
|
||||
revision = source.revision
|
||||
revision = await source.prepare_revision(date(2026, 11, 5))
|
||||
replica = Router(model_list=list(router.registry().all()))
|
||||
replacement = RouterDeprecationSource(lambda: replica)
|
||||
assert await replacement.prepare_revision(date(2026, 11, 5)) == revision
|
||||
assert replacement.page(date(2026, 11, 5), first.cursor) == second
|
||||
router = Router(model_list=[])
|
||||
assert source.revision != revision
|
||||
assert source.revision != RouterDeprecationSource(lambda: router).revision
|
||||
assert await source.prepare_revision(date(2026, 11, 5)) != revision
|
||||
|
||||
|
||||
def test_unavailable_retirement_router_is_not_a_completed_empty_catalog() -> None:
|
||||
|
|
@ -843,3 +849,62 @@ def test_unavailable_retirement_router_is_not_a_completed_empty_catalog() -> Non
|
|||
cursor = DeprecationScanCursor()
|
||||
page = source.page(date(2026, 11, 5), cursor)
|
||||
assert page.models == () and page.cursor == cursor and not page.complete
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retirement_revision_rejects_reordered_or_changed_model_metadata() -> None:
|
||||
models = [
|
||||
{
|
||||
"model_name": f"alias-{index}",
|
||||
"litellm_params": {"model": f"openai/retired-{index}"},
|
||||
"model_info": {"id": str(index), "deprecation_date": "2026-11-20"},
|
||||
}
|
||||
for index in range(3)
|
||||
]
|
||||
router = Router(model_list=models)
|
||||
source = RouterDeprecationSource(lambda: router)
|
||||
today = date(2026, 11, 5)
|
||||
revision = await source.prepare_revision(today)
|
||||
reordered = Router(model_list=list(reversed(models)))
|
||||
assert await RouterDeprecationSource(lambda: reordered).prepare_revision(today) != revision
|
||||
changed = Router(
|
||||
model_list=[
|
||||
{**model, "model_info": {**model["model_info"], "deprecation_date": "2026-11-21"}} for model in models
|
||||
]
|
||||
)
|
||||
assert await RouterDeprecationSource(lambda: changed).prepare_revision(today) != revision
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retirement_revision_resumes_cancelled_preparation_and_invalidates_changed_catalog(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": f"alias-{index}",
|
||||
"litellm_params": {"model": "openai/revision-model"},
|
||||
"model_info": {"id": str(index)},
|
||||
}
|
||||
for index in range(250)
|
||||
]
|
||||
)
|
||||
catalog = {"openai/revision-model": {"deprecation_date": "2026-11-20"}}
|
||||
monkeypatch.setattr(litellm, "model_cost", catalog)
|
||||
source = RouterDeprecationSource(lambda: router, lambda: catalog)
|
||||
today = date(2026, 11, 5)
|
||||
preparation = asyncio.create_task(source.prepare_revision(today))
|
||||
await asyncio.sleep(0)
|
||||
preparation.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await preparation
|
||||
revision = await source.prepare_revision(today)
|
||||
copied_catalog = dict(catalog)
|
||||
equivalent = RouterDeprecationSource(lambda: router, lambda: copied_catalog)
|
||||
assert revision is not None
|
||||
assert await equivalent.prepare_revision(today) == revision
|
||||
catalog = {"openai/revision-model": {"deprecation_date": "2026-11-21"}}
|
||||
monkeypatch.setattr(litellm, "model_cost", catalog)
|
||||
assert source.revision != revision
|
||||
changed = await source.prepare_revision(today)
|
||||
assert changed is not None and changed != revision
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue