fix(alerts): resume retirement scans across replicas

This commit is contained in:
moe-berri 2026-11-07 12:01:59 +08:00
parent d20ec80f90
commit a75d6335b5
4 changed files with 161 additions and 21 deletions

View file

@ -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()

View file

@ -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)

View file

@ -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

View file

@ -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