From a75d6335b580f8e9fc175b0adbc4732b4cae5de3 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 7 Nov 2026 12:01:59 +0800 Subject: [PATCH] fix(alerts): resume retirement scans across replicas --- .../proxy/common_utils/model_deprecation.py | 78 +++++++++++++++++-- .../spend_tracking/deprecation_issues.py | 9 ++- .../test_deprecation_issue_evaluation.py | 24 +++--- .../common_utils/test_model_deprecation.py | 71 ++++++++++++++++- 4 files changed, 161 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/common_utils/model_deprecation.py b/litellm/proxy/common_utils/model_deprecation.py index f84aa05977a..f548f21dc41 100644 --- a/litellm/proxy/common_utils/model_deprecation.py +++ b/litellm/proxy/common_utils/model_deprecation.py @@ -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() diff --git a/litellm/proxy/spend_tracking/deprecation_issues.py b/litellm/proxy/spend_tracking/deprecation_issues.py index 31b66058a8a..5f688357068 100644 --- a/litellm/proxy/spend_tracking/deprecation_issues.py +++ b/litellm/proxy/spend_tracking/deprecation_issues.py @@ -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) diff --git a/tests/integration/database/test_deprecation_issue_evaluation.py b/tests/integration/database/test_deprecation_issue_evaluation.py index 6e049a336d3..3417ca92ebe 100644 --- a/tests/integration/database/test_deprecation_issue_evaluation.py +++ b/tests/integration/database/test_deprecation_issue_evaluation.py @@ -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 diff --git a/tests/unit/proxy/common_utils/test_model_deprecation.py b/tests/unit/proxy/common_utils/test_model_deprecation.py index 9d526d37bdd..3e4b456658c 100644 --- a/tests/unit/proxy/common_utils/test_model_deprecation.py +++ b/tests/unit/proxy/common_utils/test_model_deprecation.py @@ -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