From e79109c7e9683b2889db246b787acc8a75d4dfa7 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 30 Jul 2026 17:57:59 -0700 Subject: [PATCH] fix(proxy): identify manual reloads by revision instead of comparing timestamps Comparing a request timestamp against each pod's data age made correctness depend on clock resolution: Postgres stores TIMESTAMP(3) while Python stamps microseconds, and two events inside the same millisecond are indistinguishable no matter how the comparison is written. Replace reload_requested_at with a reload_revision counter the manual reload endpoint increments atomically in the database. Each pod records the revision it last applied and reloads whenever the row's differs, so a request reaches every pod exactly once regardless of clock skew or precision, and concurrent requests publish distinct revisions instead of overwriting one another. A pod adopts the current revision on its first poll, since data it loaded at boot already satisfies any earlier request. Interval reloads still key off the pod's own data age, where hour scale comparisons make precision irrelevant. --- .../migration.sql | 2 +- .../litellm_proxy_extras/schema.prisma | 2 +- .../common_utils/periodic_reload_schedule.py | 61 ++++---- litellm/proxy/proxy_server.py | 25 +-- litellm/proxy/schema.prisma | 2 +- schema.prisma | 2 +- .../test_periodic_reload_schedule.py | 145 ++++++++---------- .../test_routes_model_cost_map.py | 10 +- tests/test_litellm/proxy/test_proxy_server.py | 103 ++++++++++--- 9 files changed, 205 insertions(+), 147 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260729000000_add_reload_tracking_to_litellm_config/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260729000000_add_reload_tracking_to_litellm_config/migration.sql index 548582e6b52..2a8460a9b60 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260729000000_add_reload_tracking_to_litellm_config/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260729000000_add_reload_tracking_to_litellm_config/migration.sql @@ -1,3 +1,3 @@ -- AlterTable ALTER TABLE "LiteLLM_Config" ADD COLUMN IF NOT EXISTS "last_run_at" TIMESTAMP(3), -ADD COLUMN IF NOT EXISTS "reload_requested_at" TIMESTAMP(3); +ADD COLUMN IF NOT EXISTS "reload_revision" BIGINT NOT NULL DEFAULT 0; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index cca69741357..df872e96303 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -602,7 +602,7 @@ model LiteLLM_Config { param_name String @id param_value Json? last_run_at DateTime? - reload_requested_at DateTime? + reload_revision BigInt @default(0) } // View spend, model, api_key per request diff --git a/litellm/proxy/common_utils/periodic_reload_schedule.py b/litellm/proxy/common_utils/periodic_reload_schedule.py index 1712376fb14..bc353f11ddd 100644 --- a/litellm/proxy/common_utils/periodic_reload_schedule.py +++ b/litellm/proxy/common_utils/periodic_reload_schedule.py @@ -5,12 +5,12 @@ Persistence for the admin-configured periodic model cost map reload schedule sto Field ownership is split by writer so concurrent writers never overwrite each other: the schedule endpoints own the ``param_value`` JSON (``interval_hours``), while the reload job and the manual reload endpoints own the dedicated ``last_run_at`` / -``reload_requested_at`` columns. ``last_run_at`` lives in the row rather than process -memory so the Admin UI still reports the last execution after a restart and across pods. -``reload_requested_at`` fans a manual reload out to every pod without any pod having to -write the flag back down, so no pod can starve another of the request. Each pod tracks -only when its own in-memory copy was loaded (seeded at boot) and reloads whenever that -copy is older than the request or the configured interval. +``reload_revision`` columns. ``last_run_at`` lives in the row rather than process memory +so the Admin UI still reports the last execution after a restart and across pods. +``reload_revision`` is a monotonic counter a manual reload increments; each pod records +the revision it last applied and reloads whenever the row's differs, so a request reaches +every pod exactly once without any pod clearing it and without comparing clocks. Interval +reloads stay per-pod, driven by when that pod's own copy of the data was loaded. """ from collections.abc import Mapping @@ -36,11 +36,15 @@ if TYPE_CHECKING: MODEL_COST_MAP_RELOAD_PARAM_NAME = "model_cost_map_reload_config" +class _RevisionIncrement(TypedDict): + increment: int + + class _ConfigRowWrite(TypedDict, total=False): param_name: str param_value: str last_run_at: datetime - reload_requested_at: datetime + reload_revision: int | _RevisionIncrement class _ConfigUpsertData(TypedDict): @@ -63,7 +67,7 @@ def _config_table(prisma_client: PrismaClient) -> _ConfigTable: @dataclass(frozen=True, slots=True) class ReloadSchedule: interval_hours: int | None = None - reload_requested_at: datetime | None = None + reload_revision: int = 0 last_run_at: datetime | None = None @@ -80,17 +84,8 @@ class _IntervalConfig(BaseModel): interval_hours: int | None = None -def to_db_precision(value: datetime) -> datetime: - """ - Drop sub-millisecond digits, matching the TIMESTAMP(3) columns these values round-trip - through. Comparing a microsecond-precision stamp against its own truncated copy would - read as newer and skip the request it recorded - """ - return value.replace(microsecond=(value.microsecond // 1000) * 1000) - - def utc_now() -> datetime: - return to_db_precision(datetime.now(timezone.utc)) + return datetime.now(timezone.utc) def _parse_interval_hours(param_value: object) -> int | None: @@ -109,7 +104,7 @@ def _as_utc(value: datetime | None) -> datetime | None: def parse_reload_schedule(row: "LiteLLM_Config") -> ReloadSchedule: return ReloadSchedule( interval_hours=_parse_interval_hours(row.param_value), - reload_requested_at=_as_utc(row.reload_requested_at), + reload_revision=int(row.reload_revision or 0), last_run_at=_as_utc(row.last_run_at), ) @@ -135,16 +130,19 @@ def reload_schedule_status(schedule: ReloadSchedule | None) -> ReloadScheduleSta def pod_reload_is_due( *, schedule: ReloadSchedule, + pod_applied_revision: int | None, pod_data_loaded_at: datetime, current_time: datetime, description: str, ) -> bool: """ - Whether this pod should reload now: when its in-memory data (loaded at boot or at its - last reload) is older than a manual request or the configured interval. A schedule that - has never run anywhere fires immediately so the first run is not one interval away + Whether this pod should reload now. A revision it has not applied means a manual reload + it has not served; ``pod_applied_revision`` is None only until the pod's first poll, when + its data is boot-fresh and it simply adopts whatever revision it finds. Interval reloads + compare against this pod's own data, and a schedule that has never run anywhere fires + immediately rather than one interval later """ - if schedule.reload_requested_at is not None and schedule.reload_requested_at > pod_data_loaded_at: + if pod_applied_revision is not None and schedule.reload_revision != pod_applied_revision: verbose_proxy_logger.info("%s reload triggered by manual reload request", description) return True if schedule.interval_hours is None: @@ -194,14 +192,19 @@ async def record_reload_run(prisma_client: PrismaClient, param_name: str, ran_at await invalidate_config_param(param_name) -async def record_manual_reload(prisma_client: PrismaClient, param_name: str, ran_at: datetime) -> None: - """After a manual in-pod reload: stamps the shared last run and the request every other - pod compares against on its next poll""" - await _config_table(prisma_client).upsert( +async def record_manual_reload(prisma_client: PrismaClient, param_name: str, ran_at: datetime) -> int: + """ + After a manual in-pod reload: stamp the shared last run and bump the revision every other + pod compares against. The increment is atomic, so concurrent requests each publish a + distinct revision instead of overwriting one another. Returns the published revision so + the serving pod can adopt it rather than reloading again on its next poll + """ + row = await _config_table(prisma_client).upsert( where={"param_name": param_name}, data={ - "create": {"param_name": param_name, "last_run_at": ran_at, "reload_requested_at": ran_at}, - "update": {"last_run_at": ran_at, "reload_requested_at": ran_at}, + "create": {"param_name": param_name, "last_run_at": ran_at, "reload_revision": 1}, + "update": {"last_run_at": ran_at, "reload_revision": {"increment": 1}}, }, ) await invalidate_config_param(param_name) + return int(row.reload_revision) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4f75763061e..bf55d39e6ec 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -333,7 +333,6 @@ from litellm.proxy.common_utils.periodic_reload_schedule import ( record_manual_reload, record_reload_run, reload_schedule_status, - to_db_precision, utc_now, write_reload_interval, ) @@ -3863,10 +3862,10 @@ class ProxyConfig: get_model_cost_map_loaded_at, ) - cost_map_loaded_at = get_model_cost_map_loaded_at() - self.model_cost_map_loaded_at: datetime = ( - to_db_precision(cost_map_loaded_at) if cost_map_loaded_at is not None else utc_now() - ) + self.model_cost_map_loaded_at: datetime = get_model_cost_map_loaded_at() or utc_now() + # None until the first poll adopts whatever revision the row carries; this pod's + # data is boot-fresh, so a request published before it started is already satisfied + self.model_cost_map_applied_revision: int | None = None def is_yaml(self, config_file_path: str) -> bool: if not os.path.isfile(config_file_path): @@ -6451,12 +6450,17 @@ class ProxyConfig: return current_time = utc_now() - if not pod_reload_is_due( + is_due = pod_reload_is_due( schedule=schedule, + pod_applied_revision=self.model_cost_map_applied_revision, pod_data_loaded_at=self.model_cost_map_loaded_at, current_time=current_time, description="Model cost map", - ): + ) + # Adopt the revision we just compared against, whether or not it was the reason + # to reload; a request published after this read keeps a higher one for next tick + self.model_cost_map_applied_revision = schedule.reload_revision + if not is_due: return models_count = await asyncio.to_thread(_reload_model_cost_map_in_process) @@ -15643,8 +15647,11 @@ async def reload_model_cost_map( current_time = utc_now() proxy_config.model_cost_map_loaded_at = current_time - # Stamp the shared last run and the reload request so other pods reload too - await record_manual_reload(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME, current_time) + # Publish a new revision so every other pod reloads on its next poll; this pod has + # already served it, so adopt it here rather than reloading again a tick later + proxy_config.model_cost_map_applied_revision = await record_manual_reload( + prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME, current_time + ) verbose_proxy_logger.info(f"Model cost map reloaded successfully in current pod. Models count: {models_count}") diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index cca69741357..df872e96303 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -602,7 +602,7 @@ model LiteLLM_Config { param_name String @id param_value Json? last_run_at DateTime? - reload_requested_at DateTime? + reload_revision BigInt @default(0) } // View spend, model, api_key per request diff --git a/schema.prisma b/schema.prisma index cca69741357..df872e96303 100644 --- a/schema.prisma +++ b/schema.prisma @@ -602,7 +602,7 @@ model LiteLLM_Config { param_name String @id param_value Json? last_run_at DateTime? - reload_requested_at DateTime? + reload_revision BigInt @default(0) } // View spend, model, api_key per request diff --git a/tests/test_litellm/proxy/common_utils/test_periodic_reload_schedule.py b/tests/test_litellm/proxy/common_utils/test_periodic_reload_schedule.py index e282059eac2..bef5552f007 100644 --- a/tests/test_litellm/proxy/common_utils/test_periodic_reload_schedule.py +++ b/tests/test_litellm/proxy/common_utils/test_periodic_reload_schedule.py @@ -13,8 +13,6 @@ from litellm.proxy.common_utils.periodic_reload_schedule import ( record_manual_reload, record_reload_run, reload_schedule_status, - to_db_precision, - utc_now, write_reload_interval, ) @@ -22,37 +20,39 @@ LAST_RUN = datetime(2024, 1, 1, 6, 0, 0, tzinfo=timezone.utc) NOW = datetime(2024, 1, 1, 12, 0, tzinfo=timezone.utc) -def _row(param_value=None, reload_requested_at=None, last_run_at=None): +def _row(param_value=None, reload_revision=0, last_run_at=None): return SimpleNamespace( param_name="model_cost_map_reload_config", param_value=param_value, - reload_requested_at=reload_requested_at, + reload_revision=reload_revision, last_run_at=last_run_at, ) -def _mock_prisma(row=None): +def _mock_prisma(row=None, upserted_revision=1): prisma_client = MagicMock() prisma_client.db.litellm_config.find_unique = AsyncMock(return_value=row) - prisma_client.db.litellm_config.upsert = AsyncMock(return_value=None) + prisma_client.db.litellm_config.upsert = AsyncMock(return_value=_row(reload_revision=upserted_revision)) prisma_client.db.litellm_config.update_many = AsyncMock(return_value=1) return prisma_client def test_parse_reads_interval_from_json_and_state_from_columns(): - schedule = parse_reload_schedule( - _row(param_value={"interval_hours": 6}, reload_requested_at=NOW, last_run_at=LAST_RUN) - ) + schedule = parse_reload_schedule(_row(param_value={"interval_hours": 6}, reload_revision=7, last_run_at=LAST_RUN)) - assert schedule == ReloadSchedule(interval_hours=6, reload_requested_at=NOW, last_run_at=LAST_RUN) + assert schedule == ReloadSchedule(interval_hours=6, reload_revision=7, last_run_at=LAST_RUN) def test_parse_treats_naive_column_timestamps_as_utc(): - schedule = parse_reload_schedule( - _row(reload_requested_at=NOW.replace(tzinfo=None), last_run_at=LAST_RUN.replace(tzinfo=None)) - ) + schedule = parse_reload_schedule(_row(last_run_at=LAST_RUN.replace(tzinfo=None))) - assert (schedule.reload_requested_at, schedule.last_run_at) == (NOW, LAST_RUN) + assert schedule.last_run_at == LAST_RUN + + +def test_parse_defaults_revision_when_the_column_is_null(): + """Rows written before the column existed read back as NULL and must not crash the + comparison; nobody has applied revision 0, so treating it as 0 is a no-op""" + assert parse_reload_schedule(_row(reload_revision=None)).reload_revision == 0 @pytest.mark.parametrize( @@ -68,7 +68,7 @@ def test_parse_ignores_legacy_json_force_reload(): would re-trigger a reload every poll because nothing clears the JSON copy""" schedule = parse_reload_schedule(_row(param_value={"interval_hours": 6, "force_reload": True})) - assert schedule == ReloadSchedule(interval_hours=6, reload_requested_at=None, last_run_at=None) + assert schedule == ReloadSchedule(interval_hours=6, reload_revision=0, last_run_at=None) def test_status_reports_persisted_last_run_and_next_run(): @@ -98,36 +98,29 @@ def test_next_run_needs_both_interval_and_last_run(): @pytest.mark.parametrize( - "schedule, pod_data_loaded_at, expected", + "schedule, pod_applied_revision, pod_data_loaded_at, expected", [ - (ReloadSchedule(reload_requested_at=datetime(2024, 1, 1, 11, 0, tzinfo=timezone.utc)), LAST_RUN, True), - ( - ReloadSchedule( - interval_hours=6, - reload_requested_at=datetime(2024, 1, 1, 11, 30, tzinfo=timezone.utc), - last_run_at=datetime(2024, 1, 1, 11, 0, tzinfo=timezone.utc), - ), - datetime(2024, 1, 1, 11, 0, tzinfo=timezone.utc), - True, - ), - (ReloadSchedule(reload_requested_at=LAST_RUN), LAST_RUN, False), - (ReloadSchedule(reload_requested_at=datetime(2024, 1, 1, 5, 0, tzinfo=timezone.utc)), LAST_RUN, False), - (ReloadSchedule(reload_requested_at=LAST_RUN), datetime(2024, 1, 1, 11, 0, tzinfo=timezone.utc), False), - (ReloadSchedule(interval_hours=6, reload_requested_at=LAST_RUN, last_run_at=LAST_RUN), LAST_RUN, True), - (ReloadSchedule(), LAST_RUN, False), - (ReloadSchedule(interval_hours=6), datetime(2024, 1, 1, 11, 59, tzinfo=timezone.utc), True), + (ReloadSchedule(reload_revision=4), 3, NOW, True), + (ReloadSchedule(interval_hours=6, reload_revision=4, last_run_at=LAST_RUN), 3, NOW, True), + (ReloadSchedule(reload_revision=3), 3, LAST_RUN, False), + (ReloadSchedule(reload_revision=4), None, LAST_RUN, False), + (ReloadSchedule(interval_hours=6, reload_revision=4), None, NOW, True), + (ReloadSchedule(), 0, LAST_RUN, False), + (ReloadSchedule(interval_hours=6), 0, datetime(2024, 1, 1, 11, 59, tzinfo=timezone.utc), True), ( ReloadSchedule(interval_hours=6, last_run_at=datetime(2024, 1, 1, 6, 30, tzinfo=timezone.utc)), + 0, datetime(2024, 1, 1, 11, 0, tzinfo=timezone.utc), False, ), - (ReloadSchedule(interval_hours=6, last_run_at=LAST_RUN), LAST_RUN, True), + (ReloadSchedule(interval_hours=6, last_run_at=LAST_RUN), 0, LAST_RUN, True), ], ) -def test_pod_reload_is_due(schedule, pod_data_loaded_at, expected): +def test_pod_reload_is_due(schedule, pod_applied_revision, pod_data_loaded_at, expected): assert ( pod_reload_is_due( schedule=schedule, + pod_applied_revision=pod_applied_revision, pod_data_loaded_at=pod_data_loaded_at, current_time=NOW, description="test", @@ -136,44 +129,35 @@ def test_pod_reload_is_due(schedule, pod_data_loaded_at, expected): ) -def test_stamps_match_the_precision_they_are_stored_at(): - """Postgres stores TIMESTAMP(3), so a microsecond-precision pod clock would read as - newer than its own persisted copy and skip the request it just recorded""" - assert utc_now().microsecond % 1000 == 0 - assert to_db_precision(datetime(2024, 1, 1, 6, 0, 0, 123456, tzinfo=timezone.utc)) == datetime( - 2024, 1, 1, 6, 0, 0, 123000, tzinfo=timezone.utc +def test_manual_request_is_identified_not_ordered(): + """Comparing revisions for inequality rather than ordering timestamps: a pod applies a + request once and is not due again, no matter how the clocks or precisions line up""" + unapplied = pod_reload_is_due( + schedule=ReloadSchedule(reload_revision=9), + pod_applied_revision=8, + pod_data_loaded_at=NOW, + current_time=NOW, + description="test", + ) + applied = pod_reload_is_due( + schedule=ReloadSchedule(reload_revision=9), + pod_applied_revision=9, + pod_data_loaded_at=NOW, + current_time=NOW, + description="test", ) + assert (unapplied, applied) == (True, False) -def test_request_in_the_same_millisecond_is_not_lost_to_truncation(): - """A pod whose data loaded microseconds before a request must still see it: both sides - are floored to the stored resolution, so the request never reads as older than it is""" - loaded_at = to_db_precision(datetime(2024, 1, 1, 6, 0, 0, 123400, tzinfo=timezone.utc)) - requested_at = datetime(2024, 1, 1, 6, 0, 0, 124000, tzinfo=timezone.utc) +def test_fresh_pod_adopts_the_revision_without_reloading(): + """A pod that just booted holds data newer than any earlier request, so its first poll + must not re-serve one; the interval is the only thing that can make it due""" assert ( pod_reload_is_due( - schedule=ReloadSchedule(reload_requested_at=requested_at), - pod_data_loaded_at=loaded_at, - current_time=NOW, - description="test", - ) - is True - ) - - -def test_manual_request_reaches_every_pod_with_older_data(): - """Nothing clears reload_requested_at, so each pod reloads exactly once per request: - due while its data predates the request, not due once refreshed or booted after it""" - schedule = ReloadSchedule(reload_requested_at=datetime(2024, 1, 1, 11, 0, tzinfo=timezone.utc)) - - assert ( - pod_reload_is_due(schedule=schedule, pod_data_loaded_at=LAST_RUN, current_time=NOW, description="test") is True - ) - assert ( - pod_reload_is_due( - schedule=schedule, - pod_data_loaded_at=datetime(2024, 1, 1, 11, 0, 1, tzinfo=timezone.utc), + schedule=ReloadSchedule(reload_revision=12), + pod_applied_revision=None, + pod_data_loaded_at=NOW, current_time=NOW, description="test", ) @@ -187,6 +171,7 @@ def test_pod_reload_decision_ignores_persisted_last_run(): assert ( pod_reload_is_due( schedule=ReloadSchedule(interval_hours=6, last_run_at=datetime(2024, 1, 1, 11, 59, tzinfo=timezone.utc)), + pod_applied_revision=0, pod_data_loaded_at=LAST_RUN, current_time=NOW, description="test", @@ -201,6 +186,7 @@ def test_schedule_that_never_ran_fires_immediately(): assert ( pod_reload_is_due( schedule=ReloadSchedule(interval_hours=6, last_run_at=None), + pod_applied_revision=0, pod_data_loaded_at=datetime(2024, 1, 1, 11, 59, tzinfo=timezone.utc), current_time=NOW, description="test", @@ -215,14 +201,14 @@ async def test_read_reload_schedule_returns_none_for_missing_row(): @pytest.mark.asyncio -async def test_read_reload_schedule_surfaces_request_on_interval_less_row(): +async def test_read_reload_schedule_surfaces_revision_on_interval_less_row(): """A manual reload on a proxy with no schedule creates a row with only the columns - set; the request must still reach other pods""" - prisma_client = _mock_prisma(row=_row(param_value=None, reload_requested_at=NOW)) + set; the revision must still reach other pods""" + prisma_client = _mock_prisma(row=_row(param_value=None, reload_revision=3)) schedule = await read_reload_schedule(prisma_client, "model_cost_map_reload_config") - assert schedule == ReloadSchedule(interval_hours=None, reload_requested_at=NOW, last_run_at=None) + assert schedule == ReloadSchedule(interval_hours=None, reload_revision=3, last_run_at=None) @pytest.mark.asyncio @@ -237,9 +223,9 @@ async def test_write_reload_interval_touches_only_param_value(): @pytest.mark.asyncio -async def test_record_reload_run_updates_last_run_without_creating_or_clearing(): - """update_many so a schedule deleted mid-poll stays deleted, and the untouched - reload_requested_at keeps fanning the request out to pods that have not seen it""" +async def test_record_reload_run_updates_last_run_without_creating_or_bumping(): + """update_many so a schedule deleted mid-poll stays deleted, and the untouched revision + keeps fanning the request out to pods that have not applied it""" prisma_client = _mock_prisma() await record_reload_run(prisma_client, "model_cost_map_reload_config", LAST_RUN) @@ -250,15 +236,18 @@ async def test_record_reload_run_updates_last_run_without_creating_or_clearing() @pytest.mark.asyncio -async def test_record_manual_reload_stamps_request_and_last_run(): - prisma_client = _mock_prisma() +async def test_record_manual_reload_bumps_the_revision_atomically(): + """The increment must be delegated to the database: two concurrent requests that both + read then wrote a computed value would publish the same revision and one would be lost""" + prisma_client = _mock_prisma(upserted_revision=5) - await record_manual_reload(prisma_client, "model_cost_map_reload_config", LAST_RUN) + published = await record_manual_reload(prisma_client, "model_cost_map_reload_config", LAST_RUN) data = prisma_client.db.litellm_config.upsert.await_args.kwargs["data"] - assert data["update"] == {"last_run_at": LAST_RUN, "reload_requested_at": LAST_RUN} + assert data["update"] == {"last_run_at": LAST_RUN, "reload_revision": {"increment": 1}} assert data["create"] == { "param_name": "model_cost_map_reload_config", "last_run_at": LAST_RUN, - "reload_requested_at": LAST_RUN, + "reload_revision": 1, } + assert published == 5 diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_cost_map.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_cost_map.py index 680ae1b5663..08bf655de24 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_cost_map.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_cost_map.py @@ -83,8 +83,8 @@ def test_reload_model_cost_map_happy(client, auth_as, monkeypatch, mock_prisma): } assert table.upsert.await_count == 1 update_payload = table.upsert.await_args.kwargs["data"]["update"] - assert set(update_payload) == {"last_run_at", "reload_requested_at"} - assert update_payload["reload_requested_at"] == update_payload["last_run_at"] + assert set(update_payload) == {"last_run_at", "reload_revision"} + assert update_payload["reload_revision"] == {"increment": 1} def test_reload_model_cost_map_not_admin_forbidden(client, auth_as): @@ -255,7 +255,7 @@ def test_get_model_cost_map_reload_status_scheduled( table = _attach_litellm_config(mock_prisma) config_row = MagicMock() config_row.param_value = {"interval_hours": 12} - config_row.reload_requested_at = None + config_row.reload_revision = 0 config_row.last_run_at = None table.find_unique = AsyncMock(return_value=config_row) monkeypatch.setattr(ps, "prisma_client", mock_prisma) @@ -283,7 +283,7 @@ def test_get_model_cost_map_reload_status_reports_persisted_last_run( table = _attach_litellm_config(mock_prisma) config_row = MagicMock() config_row.param_value = {"interval_hours": 6} - config_row.reload_requested_at = None + config_row.reload_revision = 0 config_row.last_run_at = datetime(2024, 1, 1, 6, 0, 0, tzinfo=timezone.utc) table.find_unique = AsyncMock(return_value=config_row) monkeypatch.setattr(ps, "prisma_client", mock_prisma) @@ -311,7 +311,7 @@ def test_get_model_cost_map_reload_status_no_config_not_scheduled( table = _attach_litellm_config(mock_prisma) config_row = MagicMock() config_row.param_value = {"interval_hours": None} - config_row.reload_requested_at = datetime(2024, 1, 1, 6, 0, 0, tzinfo=timezone.utc) + config_row.reload_revision = 3 config_row.last_run_at = None table.find_unique = AsyncMock(return_value=config_row) monkeypatch.setattr(ps, "prisma_client", mock_prisma) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 552cab81c42..325f90e5f44 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -3639,13 +3639,13 @@ async def test_chat_completion_result_no_nested_none_values(): def _reload_schedule_row( param_value: dict, *, - reload_requested_at: datetime | None = None, + reload_revision: int = 0, last_run_at: datetime | None = None, ) -> types.SimpleNamespace: """LiteLLM_Config row shape: admin-owned interval in param_value, run state in dedicated columns""" return types.SimpleNamespace( param_value=param_value, - reload_requested_at=reload_requested_at, + reload_revision=reload_revision, last_run_at=last_run_at, ) @@ -3686,7 +3686,9 @@ class TestPriceDataReloadAPI: } # Mock the database connection with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: - mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) + mock_prisma.db.litellm_config.upsert = AsyncMock( + return_value=_reload_schedule_row({}, reload_revision=1) + ) response = client_with_auth.post("/reload/model_cost_map") @@ -3742,7 +3744,9 @@ class TestPriceDataReloadAPI: # Mock the database connection with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) - mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) + mock_prisma.db.litellm_config.upsert = AsyncMock( + return_value=_reload_schedule_row({}, reload_revision=1) + ) response = client_with_auth.post("/reload/model_cost_map") @@ -3756,7 +3760,9 @@ class TestPriceDataReloadAPI: """Admin schedule write owns param_value only, so it can't clobber the job-owned run columns""" with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: # Mock database upsert - mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) + mock_prisma.db.litellm_config.upsert = AsyncMock( + return_value=_reload_schedule_row({}, reload_revision=1) + ) response = client_with_auth.post("/schedule/model_cost_map_reload?hours=6") @@ -3888,7 +3894,7 @@ class TestPriceDataReloadAPI: mock_prisma.db.litellm_config.find_unique = AsyncMock( return_value=_reload_schedule_row( {"interval_hours": None}, - reload_requested_at=datetime(2024, 1, 1, 6, 0, tzinfo=timezone.utc), + reload_revision=3, ) ) @@ -3959,7 +3965,9 @@ class TestPriceDataReloadIntegration: # Mock the database connection with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: - mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) + mock_prisma.db.litellm_config.upsert = AsyncMock( + return_value=_reload_schedule_row({}, reload_revision=1) + ) # Test reload endpoint response = client_with_auth.post("/reload/model_cost_map") @@ -3988,14 +3996,16 @@ class TestPriceDataReloadIntegration: def test_distributed_reload_check_function(self): """ - A reload requested by another pod after this pod's last reload takes effect here - even one minute into a 6h interval; a missing row is a no-op + A revision this pod has not applied takes effect here even one minute into a 6h + interval; a missing row is a no-op """ from litellm.proxy.proxy_server import ProxyConfig proxy_config = ProxyConfig() mock_prisma = MagicMock() - mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) + mock_prisma.db.litellm_config.upsert = AsyncMock( + return_value=_reload_schedule_row({}, reload_revision=1) + ) mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None) mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) @@ -4006,15 +4016,15 @@ class TestPriceDataReloadIntegration: assert proxy_config.model_cost_map_loaded_at == boot_loaded_at frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc) - requested_at = datetime(2024, 1, 1, 6, 59, 30, tzinfo=timezone.utc) mock_prisma.db.litellm_config.find_unique = AsyncMock( return_value=_reload_schedule_row( {"interval_hours": 6}, - reload_requested_at=requested_at, - last_run_at=requested_at, + reload_revision=4, + last_run_at=datetime(2024, 1, 1, 6, 59, 30, tzinfo=timezone.utc), ) ) proxy_config.model_cost_map_loaded_at = frozen_now - timedelta(minutes=1) + proxy_config.model_cost_map_applied_revision = 3 original_model_cost = litellm.model_cost.copy() try: @@ -4039,14 +4049,15 @@ class TestPriceDataReloadIntegration: "where": {"param_name": "model_cost_map_reload_config"}, } mock_prisma.db.litellm_config.upsert.assert_not_called() + assert proxy_config.model_cost_map_applied_revision == 4 finally: litellm.model_cost = original_model_cost _invalidate_model_cost_lowercase_map() def test_distributed_reload_ignores_already_applied_request(self): """ - A reload request older than this pod's own last reload was already applied here, so - it must not re-trigger on every job tick for the rest of the interval + A revision this pod already applied must not re-trigger on every job tick for the + rest of the interval """ from litellm.proxy.proxy_server import ProxyConfig @@ -4056,13 +4067,14 @@ class TestPriceDataReloadIntegration: mock_prisma.db.litellm_config.find_unique = AsyncMock( return_value=_reload_schedule_row( {"interval_hours": 6}, - reload_requested_at=datetime(2024, 1, 1, 6, 0, tzinfo=timezone.utc), + reload_revision=4, last_run_at=datetime(2024, 1, 1, 6, 0, tzinfo=timezone.utc), ) ) frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc) pod_data_loaded_at = frozen_now - timedelta(minutes=1) proxy_config.model_cost_map_loaded_at = pod_data_loaded_at + proxy_config.model_cost_map_applied_revision = 4 original_model_cost = litellm.model_cost.copy() try: @@ -4136,6 +4148,42 @@ class TestPriceDataReloadIntegration: litellm.model_cost = original_model_cost _invalidate_model_cost_lowercase_map() + def test_every_pod_applies_a_manual_revision_exactly_once(self): + """The fleet property: no pod clears the revision, so each one reloads on the tick + after it is published and then stops, whatever order the pods poll in""" + from litellm.proxy.proxy_server import ProxyConfig + + pods = [ProxyConfig(), ProxyConfig(), ProxyConfig()] + frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc) + mock_prisma = MagicMock() + mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None) + mock_prisma.db.litellm_config.find_unique = AsyncMock( + return_value=_reload_schedule_row({}, reload_revision=1) + ) + for pod in pods: + pod.model_cost_map_applied_revision = 0 + pod.model_cost_map_loaded_at = frozen_now + + original_model_cost = litellm.model_cost.copy() + try: + with ( + patch( + "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map" + ) as mock_get_map, + patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), + ): + mock_get_map.return_value = {"gpt-4": {"input_cost_per_token": 0.001}} + + for _ in range(3): + for pod in pods: + asyncio.run(pod._check_and_reload_model_cost_map(mock_prisma)) + + assert mock_get_map.call_count == len(pods) + assert all(p.model_cost_map_applied_revision == 1 for p in pods) + finally: + litellm.model_cost = original_model_cost + _invalidate_model_cost_lowercase_map() + def test_distributed_reload_stamps_last_run_without_creating_row(self): """ Regression: the job's write carries neither param_value (which would clobber the @@ -4149,7 +4197,9 @@ class TestPriceDataReloadIntegration: mock_prisma.db.litellm_config.find_unique = AsyncMock( return_value=_reload_schedule_row({"interval_hours": 24}) ) - mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) + mock_prisma.db.litellm_config.upsert = AsyncMock( + return_value=_reload_schedule_row({}, reload_revision=1) + ) mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None) frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc) @@ -4204,7 +4254,9 @@ class TestPriceDataReloadIntegration: patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): mock_get_map.return_value = {"gpt-4": {"input_cost_per_token": 0.001}} - mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) + mock_prisma.db.litellm_config.upsert = AsyncMock( + return_value=_reload_schedule_row({}, reload_revision=9) + ) response = client.post("/reload/model_cost_map") assert response.status_code == 200 @@ -4213,17 +4265,20 @@ class TestPriceDataReloadIntegration: call_args = mock_prisma.db.litellm_config.upsert.call_args assert call_args[1]["data"]["update"] == { "last_run_at": frozen_now, - "reload_requested_at": frozen_now, + "reload_revision": {"increment": 1}, } assert call_args[1]["data"]["create"] == { "param_name": "model_cost_map_reload_config", "last_run_at": frozen_now, - "reload_requested_at": frozen_now, + "reload_revision": 1, } assert ( proxy_server_module.proxy_config.model_cost_map_loaded_at == frozen_now ) + assert ( + proxy_server_module.proxy_config.model_cost_map_applied_revision == 9 + ), "the serving pod must adopt the revision it published, not reload again" finally: litellm.model_cost = original_model_cost _invalidate_model_cost_lowercase_map() @@ -4247,7 +4302,9 @@ class TestPriceDataReloadIntegration: mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config) # _check_and_reload_anthropic_beta_headers now reads through get_generic_data. mock_prisma.get_generic_data = AsyncMock(return_value=mock_config) - mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) + mock_prisma.db.litellm_config.upsert = AsyncMock( + return_value=_reload_schedule_row({}, reload_revision=1) + ) with patch( "litellm.anthropic_beta_headers_manager.reload_beta_headers_config" @@ -4300,7 +4357,9 @@ class TestPriceDataReloadIntegration: mock_prisma.db.litellm_config.find_unique = AsyncMock( return_value=mock_existing ) - mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) + mock_prisma.db.litellm_config.upsert = AsyncMock( + return_value=_reload_schedule_row({}, reload_revision=1) + ) response = client.post("/reload/anthropic_beta_headers") assert response.status_code == 200