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