From bfa47535b51ffac7c729cc87563adf87f169f521 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 29 Jul 2026 17:37:27 -0700 Subject: [PATCH] fix(proxy): persist periodic reload schedule state so status survives restarts and fires without store_model_in_db The model cost map and Anthropic beta headers reload schedules kept their last-run time in a per-pod module global, so GET /schedule/*/status reported last_run null after any restart and the Admin UI showed the reload as never having run. The reload check also only ran from the add_deployment job, which is registered only when store_model_in_db is true, so config-file deployments stored a schedule that never fired. Persist last_run_at and reload_requested_at as dedicated columns on LiteLLM_Config, owned by the reload job and manual reload endpoints, while the schedule endpoints own the param_value JSON (interval_hours); no writer can clobber another's fields. Serve status entirely from the row. Register the check as its own periodic_reload_job outside the store_model_in_db gate. Replace the force_reload boolean with a reload_requested_at timestamp each pod compares against its own in-memory last reload, so a manual reload reaches every pod exactly once instead of being cleared by the first poller. Run the blocking fetches via asyncio.to_thread, and stamp last_run_at with update_many so a schedule cancelled mid-poll is not resurrected. --- .../migration.sql | 7 + .../litellm_proxy_extras/schema.prisma | 2 + .../common_utils/periodic_reload_schedule.py | 204 ++++++ litellm/proxy/proxy_server.py | 520 ++++---------- litellm/proxy/schema.prisma | 2 + ruff.toml | 2 +- schema.prisma | 2 + .../test_periodic_reload_schedule.py | 213 ++++++ .../test_routes_anthropic_beta.py | 71 +- .../test_routes_model_cost_map.py | 51 +- tests/test_litellm/proxy/test_proxy_server.py | 637 ++++++++++-------- 11 files changed, 1002 insertions(+), 709 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260729000000_add_reload_tracking_to_litellm_config/migration.sql create mode 100644 litellm/proxy/common_utils/periodic_reload_schedule.py create mode 100644 tests/test_litellm/proxy/common_utils/test_periodic_reload_schedule.py 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 new file mode 100644 index 00000000000..c288aef76ff --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260729000000_add_reload_tracking_to_litellm_config/migration.sql @@ -0,0 +1,7 @@ +-- 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); + +-- Carry over a manual reload that was still pending under the older boolean flag, so +-- pods that have not picked it up yet still fan out after this deploy. +UPDATE "LiteLLM_Config" SET "reload_requested_at" = CURRENT_TIMESTAMP WHERE "param_value"->'force_reload' = 'true'::jsonb; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 37ea55f8c13..cca69741357 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -601,6 +601,8 @@ model LiteLLM_TagTable { model LiteLLM_Config { param_name String @id param_value Json? + last_run_at DateTime? + reload_requested_at DateTime? } // 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 new file mode 100644 index 00000000000..ede10efc432 --- /dev/null +++ b/litellm/proxy/common_utils/periodic_reload_schedule.py @@ -0,0 +1,204 @@ +""" +Persistence for the admin-configured periodic reload schedules (model cost map, +Anthropic beta headers) stored in ``LiteLLM_Config``. + +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 +its own last reload in memory and compares it against those timestamps to decide when to +refresh its in-memory copy of the data. +""" + +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import ( + TYPE_CHECKING, + Protocol, + TypedDict, + cast, # noqa: TID251 # prisma table access is untyped (PrismaWrapper.__getattr__) +) + +from pydantic import BaseModel, ConfigDict, ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.proxy.utils import PrismaClient, invalidate_config_param +from litellm.repositories.config_repository import ConfigRepository + +if TYPE_CHECKING: + from prisma.models import LiteLLM_Config + +MODEL_COST_MAP_RELOAD_PARAM_NAME = "model_cost_map_reload_config" +ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME = "anthropic_beta_headers_reload_config" + + +class _ConfigRowWrite(TypedDict, total=False): + param_name: str + param_value: str + last_run_at: datetime + reload_requested_at: datetime + + +class _ConfigUpsertData(TypedDict): + create: _ConfigRowWrite + update: _ConfigRowWrite + + +class _ConfigTable(Protocol): + async def find_unique(self, where: Mapping[str, str]) -> "LiteLLM_Config | None": ... + + async def upsert(self, where: Mapping[str, str], data: _ConfigUpsertData) -> "LiteLLM_Config": ... + + async def update_many(self, data: _ConfigRowWrite, where: Mapping[str, str]) -> int: ... + + +def _config_table(prisma_client: PrismaClient) -> _ConfigTable: + return cast(_ConfigTable, ConfigRepository(prisma_client).table) # cast-ok: prisma table is untyped (Any) + + +@dataclass(frozen=True, slots=True) +class ReloadSchedule: + interval_hours: int | None = None + reload_requested_at: datetime | None = None + last_run_at: datetime | None = None + + +class ReloadScheduleStatus(TypedDict): + scheduled: bool + interval_hours: int | None + last_run: str | None + next_run: str | None + + +class _IntervalConfig(BaseModel): + model_config = ConfigDict(strict=True) + + interval_hours: int | None = None + + +def utc_now() -> datetime: + return datetime.now(timezone.utc) + + +def _parse_interval_hours(param_value: object) -> int | None: + try: + return _IntervalConfig.model_validate(param_value).interval_hours + except ValidationError: + return None + + +def _as_utc(value: datetime | None) -> datetime | None: + if value is None: + return None + return value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc) + + +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), + last_run_at=_as_utc(row.last_run_at), + ) + + +def next_run_at(schedule: ReloadSchedule) -> datetime | None: + if schedule.interval_hours is None or schedule.last_run_at is None: + return None + return schedule.last_run_at + timedelta(hours=schedule.interval_hours) + + +def reload_schedule_status(schedule: ReloadSchedule | None) -> ReloadScheduleStatus: + if schedule is None: + return {"scheduled": False, "interval_hours": None, "last_run": None, "next_run": None} + next_run = next_run_at(schedule) + return { + "scheduled": schedule.interval_hours is not None, + "interval_hours": schedule.interval_hours, + "last_run": schedule.last_run_at.isoformat() if schedule.last_run_at is not None else None, + "next_run": next_run.isoformat() if next_run is not None else None, + } + + +def pod_reload_is_due( + *, + schedule: ReloadSchedule, + pod_last_reload: datetime | None, + current_time: datetime, + description: str, +) -> bool: + """ + Whether this pod should reload now, based on its own last reload rather than the + persisted one, so every pod refreshes its in-memory data on the configured interval. + A pod that has never reloaded ignores ``reload_requested_at``: whatever it booted with + is at least as fresh as the request + """ + if ( + schedule.reload_requested_at is not None + and pod_last_reload is not None + and schedule.reload_requested_at > pod_last_reload + ): + verbose_proxy_logger.info("%s reload triggered by manual reload request", description) + return True + if schedule.interval_hours is None: + return False + if pod_last_reload is None: + verbose_proxy_logger.info("%s reload triggered - no previous reload time recorded", description) + return True + hours_since_last_reload = (current_time - pod_last_reload).total_seconds() / 3600 + if hours_since_last_reload < schedule.interval_hours: + return False + verbose_proxy_logger.info( + "%s reload triggered by interval. Hours since last reload: %.2f, Interval: %s", + description, + hours_since_last_reload, + schedule.interval_hours, + ) + return True + + +async def read_reload_schedule(prisma_client: PrismaClient, param_name: str) -> ReloadSchedule | None: + row = await _config_table(prisma_client).find_unique(where={"param_name": param_name}) + if row is None: + return None + return parse_reload_schedule(row) + + +async def write_reload_interval(prisma_client: PrismaClient, param_name: str, interval_hours: int) -> None: + """Admin-owned write: replaces ``param_value`` without touching the job-owned columns""" + param_value = safe_dumps({"interval_hours": interval_hours}) + await _config_table(prisma_client).upsert( + where={"param_name": param_name}, + data={ + "create": {"param_name": param_name, "param_value": param_value}, + "update": {"param_value": param_value}, + }, + ) + await invalidate_config_param(param_name) + + +async def record_reload_run(prisma_client: PrismaClient, param_name: str, ran_at: datetime) -> None: + """Job-owned write after this pod reloaded: stamps the shared last run only if the row + still exists, so a schedule deleted mid-poll is not resurrected""" + await _config_table(prisma_client).update_many( + data={"last_run_at": ran_at}, + where={"param_name": param_name}, + ) + 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( + 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}, + }, + ) + await invalidate_config_param(param_name) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4486cd7de59..7e0f496f525 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -326,6 +326,17 @@ from litellm.proxy.common_utils.model_listing_utils import TeamModelNameTranslat from litellm.proxy.common_utils.openai_endpoint_utils import ( remove_sensitive_info_from_deployment, ) +from litellm.proxy.common_utils.periodic_reload_schedule import ( + ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME, + MODEL_COST_MAP_RELOAD_PARAM_NAME, + pod_reload_is_due, + read_reload_schedule, + record_manual_reload, + record_reload_run, + reload_schedule_status, + utc_now, + write_reload_interval, +) from litellm.proxy.common_utils.proxy_state import ProxyState from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob from litellm.proxy.common_utils.swagger_utils import ERROR_RESPONSES @@ -2054,12 +2065,7 @@ async_result = None celery_app_conn = None celery_fn = None # Redis Queue for handling requests -# Global variables for model cost map reload scheduling scheduler = None -last_model_cost_map_reload = None - -# Global variable for anthropic beta headers reload scheduling -last_anthropic_beta_headers_reload = None ### DB WRITER ### @@ -3826,6 +3832,20 @@ def resolve_complexity_router_plugins( ) +def _reload_model_cost_map_in_process() -> int: + """Fetch the latest model cost map, swap it into this process's litellm state, return the model count""" + from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map + + new_model_cost_map = get_model_cost_map(url=litellm.model_cost_map_url) + litellm.model_cost = new_model_cost_map + # Invalidate case-insensitive lookup map since model_cost was replaced + _invalidate_model_cost_lowercase_map() + # Repopulate provider model sets (e.g. litellm.anthropic_models) so that + # wildcard patterns like "anthropic/*" include any newly added models. + litellm.add_known_models(model_cost_map=new_model_cost_map) + return len(new_model_cost_map) if new_model_cost_map else 0 + + class ProxyConfig: """ Abstraction class on top of config loading/updating logic. Gives us one place to control all config updating logic. @@ -3836,6 +3856,8 @@ class ProxyConfig: self._last_semantic_filter_config: Optional[Dict[str, Any]] = None self._last_hashicorp_vault_config: Optional[Dict[str, Any]] = None self.worker_registry: List["WorkerRegistryEntry"] = [] + self.last_model_cost_map_reload: datetime | None = None + self.last_anthropic_beta_headers_reload: datetime | None = None def is_yaml(self, config_file_path: str) -> bool: if not os.path.isfile(config_file_path): @@ -6175,8 +6197,6 @@ class ProxyConfig: "router_settings", "litellm_settings", "environment_variables", - "model_cost_map_reload_config", - "anthropic_beta_headers_reload_config", ], ) @@ -6239,12 +6259,6 @@ class ProxyConfig: if self._should_load_db_object(object_type="tools"): await self._init_tool_policy_in_db(prisma_client=prisma_client) - if self._should_load_db_object(object_type="model_cost_map"): - await self._check_and_reload_model_cost_map(prisma_client=prisma_client) - - if self._should_load_db_object(object_type="anthropic_beta_headers"): - await self._check_and_reload_anthropic_beta_headers(prisma_client=prisma_client) - if self._should_load_db_object(object_type="sso_settings"): await self._init_sso_settings_in_db(prisma_client=prisma_client) if self._should_load_db_object(object_type="cache_settings"): @@ -6403,102 +6417,44 @@ class ProxyConfig: str(e), ) + async def check_periodic_reloads(self, prisma_client: PrismaClient): + """ + Run the admin-configured periodic reloads (model cost map, Anthropic beta headers). + + Scheduled on its own job so schedules configured from the Admin UI fire whether + or not `store_model_in_db` is enabled. + """ + if self._should_load_db_object(object_type="model_cost_map"): + await self._check_and_reload_model_cost_map(prisma_client=prisma_client) + + if self._should_load_db_object(object_type="anthropic_beta_headers"): + await self._check_and_reload_anthropic_beta_headers(prisma_client=prisma_client) + async def _check_and_reload_model_cost_map(self, prisma_client: PrismaClient): """ Check if model cost map needs to be reloaded based on database configuration. - This function runs every 10 seconds as part of _init_non_llm_objects_in_db. + Runs on the periodic reload job, independently of `store_model_in_db`. """ try: - # Get model cost map reload configuration from database - config_record = await get_config_param(prisma_client, "model_cost_map_reload_config") + schedule = await read_reload_schedule(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME) + if schedule is None: + return - if config_record is None or config_record.param_value is None: - return # No configuration found, skip reload + current_time = utc_now() + if not pod_reload_is_due( + schedule=schedule, + pod_last_reload=self.last_model_cost_map_reload, + current_time=current_time, + description="Model cost map", + ): + return - config = config_record.param_value - interval_hours = config.get("interval_hours") - force_reload = config.get("force_reload", False) + models_count = await asyncio.to_thread(_reload_model_cost_map_in_process) - if interval_hours is None and force_reload is False: - return # No interval configured, skip reload + self.last_model_cost_map_reload = current_time + await record_reload_run(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME, current_time) - current_time = datetime.utcnow() - - # Check if we need to reload based on interval or force reload - should_reload = False - - if force_reload: - should_reload = True - verbose_proxy_logger.info("Model cost map reload triggered by force reload flag") - elif interval_hours is not None: - # Use pod's in-memory last reload time - global last_model_cost_map_reload - if last_model_cost_map_reload is not None: - try: - last_reload_time = datetime.fromisoformat(last_model_cost_map_reload) - time_since_last_reload = current_time - last_reload_time - hours_since_last_reload = time_since_last_reload.total_seconds() / 3600 - - if hours_since_last_reload >= interval_hours: - should_reload = True - verbose_proxy_logger.info( - f"Model cost map reload triggered by interval. Hours since last reload: {hours_since_last_reload:.2f}, Interval: {interval_hours}" - ) - except Exception as e: - verbose_proxy_logger.warning(f"Error parsing last reload time: {e}") - # If we can't parse the last reload time, reload anyway - should_reload = True - else: - # No last reload time recorded, reload now - should_reload = True - verbose_proxy_logger.info("Model cost map reload triggered - no previous reload time recorded") - - if should_reload: - # Perform the reload - from litellm.litellm_core_utils.get_model_cost_map import ( - get_model_cost_map, - ) - - model_cost_map_url = litellm.model_cost_map_url - new_model_cost_map = get_model_cost_map(url=model_cost_map_url) - litellm.model_cost = new_model_cost_map - # Invalidate case-insensitive lookup map since model_cost was replaced - _invalidate_model_cost_lowercase_map() - # Repopulate provider model sets (e.g. litellm.anthropic_models) so that - # wildcard patterns like "anthropic/*" include any newly added models. - litellm.add_known_models(model_cost_map=new_model_cost_map) - - # Update pod's in-memory last reload time - last_model_cost_map_reload = current_time.isoformat() - - # Clear force reload flag in database - await ConfigRepository(prisma_client).table.upsert( - where={"param_name": "model_cost_map_reload_config"}, - data={ - "create": { - "param_name": "model_cost_map_reload_config", - "param_value": safe_dumps( - { - "interval_hours": interval_hours, - "force_reload": False, - } - ), - }, - "update": { - "param_value": safe_dumps( - { - "interval_hours": interval_hours, - "force_reload": False, - } - ) - }, - }, - ) - await invalidate_config_param("model_cost_map_reload_config") - - verbose_proxy_logger.info( - f"Model cost map reloaded successfully. Models count: {len(new_model_cost_map) if new_model_cost_map else 0}" - ) + verbose_proxy_logger.info(f"Model cost map reloaded successfully. Models count: {models_count}") except Exception as e: verbose_proxy_logger.exception(f"Error in _check_and_reload_model_cost_map: {str(e)}") @@ -6506,96 +6462,35 @@ class ProxyConfig: async def _check_and_reload_anthropic_beta_headers(self, prisma_client: PrismaClient): """ Check if anthropic beta headers config needs to be reloaded based on database configuration. - This function runs every 10 seconds as part of _init_non_llm_objects_in_db. + Runs on the periodic reload job, independently of `store_model_in_db`. """ try: - # Get anthropic beta headers reload configuration from database - config_record = await get_config_param(prisma_client, "anthropic_beta_headers_reload_config") + schedule = await read_reload_schedule(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME) + if schedule is None: + return - if config_record is None or config_record.param_value is None: - return # No configuration found, skip reload + current_time = utc_now() + if not pod_reload_is_due( + schedule=schedule, + pod_last_reload=self.last_anthropic_beta_headers_reload, + current_time=current_time, + description="Anthropic beta headers", + ): + return - config = config_record.param_value - interval_hours = config.get("interval_hours") - force_reload = config.get("force_reload", False) + from litellm.anthropic_beta_headers_manager import ( + reload_beta_headers_config, + ) - if interval_hours is None and force_reload is False: - return # No interval configured, skip reload + new_config = await asyncio.to_thread(reload_beta_headers_config) - current_time = datetime.utcnow() + self.last_anthropic_beta_headers_reload = current_time + await record_reload_run(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME, current_time) - # Check if we need to reload based on interval or force reload - should_reload = False - - if force_reload: - should_reload = True - verbose_proxy_logger.info("Anthropic beta headers reload triggered by force reload flag") - elif interval_hours is not None: - # Use pod's in-memory last reload time - global last_anthropic_beta_headers_reload - if last_anthropic_beta_headers_reload is not None: - try: - last_reload_time = datetime.fromisoformat(last_anthropic_beta_headers_reload) - time_since_last_reload = current_time - last_reload_time - hours_since_last_reload = time_since_last_reload.total_seconds() / 3600 - - if hours_since_last_reload >= interval_hours: - should_reload = True - verbose_proxy_logger.info( - f"Anthropic beta headers reload triggered by interval. Hours since last reload: {hours_since_last_reload:.2f}, Interval: {interval_hours}" - ) - except Exception as e: - verbose_proxy_logger.warning(f"Error parsing last reload time: {e}") - # If we can't parse the last reload time, reload anyway - should_reload = True - else: - # No last reload time recorded, reload now - should_reload = True - verbose_proxy_logger.info( - "Anthropic beta headers reload triggered - no previous reload time recorded" - ) - - if should_reload: - # Perform the reload - from litellm.anthropic_beta_headers_manager import ( - reload_beta_headers_config, - ) - - new_config = reload_beta_headers_config() - - # Update pod's in-memory last reload time - last_anthropic_beta_headers_reload = current_time.isoformat() - - # Clear force reload flag in database - await ConfigRepository(prisma_client).table.upsert( - where={"param_name": "anthropic_beta_headers_reload_config"}, - data={ - "create": { - "param_name": "anthropic_beta_headers_reload_config", - "param_value": safe_dumps( - { - "interval_hours": interval_hours, - "force_reload": False, - } - ), - }, - "update": { - "param_value": safe_dumps( - { - "interval_hours": interval_hours, - "force_reload": False, - } - ) - }, - }, - ) - await invalidate_config_param("anthropic_beta_headers_reload_config") - - # Count providers in config - provider_count = sum(1 for k in new_config.keys() if k != "provider_aliases" and k != "description") - verbose_proxy_logger.info( - f"Anthropic beta headers config reloaded successfully. Providers: {provider_count}" - ) + provider_count = sum(1 for k in new_config.keys() if k != "provider_aliases" and k != "description") + verbose_proxy_logger.info( + f"Anthropic beta headers config reloaded successfully. Providers: {provider_count}" + ) except Exception as e: verbose_proxy_logger.exception(f"Error in _check_and_reload_anthropic_beta_headers: {str(e)}") @@ -8059,15 +7954,26 @@ class ProxyStartupEvent: except Exception as e: verbose_proxy_logger.debug("Failed to check DB for store_model_in_db: %s", str(e)) - if store_model_in_db is True: - config_reload_interval_seconds = proxy_config_reload_interval_seconds - if not isinstance(config_reload_interval_seconds, int) or config_reload_interval_seconds <= 0: - verbose_proxy_logger.warning( - "proxy_config_reload_interval_seconds=%s must be a positive integer; falling back to 30s", - config_reload_interval_seconds, - ) - config_reload_interval_seconds = 30 + config_reload_interval_seconds = proxy_config_reload_interval_seconds + if not isinstance(config_reload_interval_seconds, int) or config_reload_interval_seconds <= 0: + verbose_proxy_logger.warning( + "proxy_config_reload_interval_seconds=%s must be a positive integer; falling back to 30s", + config_reload_interval_seconds, + ) + config_reload_interval_seconds = 30 + ### PERIODIC RELOADS (model cost map, anthropic beta headers) ### + scheduler.add_job( + proxy_config.check_periodic_reloads, + "interval", + seconds=config_reload_interval_seconds, + args=[prisma_client], + id="periodic_reload_job", + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + + if store_model_in_db is True: # MEMORY LEAK FIX: Increase interval from 10s to 30s minimum # Frequent polling was causing excessive memory allocations scheduler.add_job( @@ -15661,43 +15567,14 @@ async def reload_model_cost_map( raise HTTPException(status_code=500, detail="Database connection not available") # Immediately reload the model cost map in the current pod - from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map + models_count = await asyncio.to_thread(_reload_model_cost_map_in_process) - model_cost_map_url = litellm.model_cost_map_url - new_model_cost_map = get_model_cost_map(url=model_cost_map_url) - litellm.model_cost = new_model_cost_map - # Invalidate case-insensitive lookup map since model_cost was replaced - _invalidate_model_cost_lowercase_map() - # Repopulate provider model sets (e.g. litellm.anthropic_models) so that - # wildcard patterns like "anthropic/*" include any newly added models. - litellm.add_known_models(model_cost_map=new_model_cost_map) + current_time = utc_now() + proxy_config.last_model_cost_map_reload = current_time - # Update pod's in-memory last reload time - global last_model_cost_map_reload - current_time = datetime.utcnow() - last_model_cost_map_reload = current_time.isoformat() + # 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) - # Set force reload flag in database for other pods, preserving existing interval_hours - existing_config = await ConfigRepository(prisma_client).table.find_unique( - where={"param_name": "model_cost_map_reload_config"} - ) - existing_interval = None - if existing_config and existing_config.param_value: - existing_interval = existing_config.param_value.get("interval_hours") - - await ConfigRepository(prisma_client).table.upsert( - where={"param_name": "model_cost_map_reload_config"}, - data={ - "create": { - "param_name": "model_cost_map_reload_config", - "param_value": safe_dumps({"interval_hours": None, "force_reload": True}), - }, - "update": {"param_value": safe_dumps({"interval_hours": existing_interval, "force_reload": True})}, - }, - ) - await invalidate_config_param("model_cost_map_reload_config") - - models_count = len(new_model_cost_map) if new_model_cost_map else 0 verbose_proxy_logger.info(f"Model cost map reloaded successfully in current pod. Models count: {models_count}") return { @@ -15742,18 +15619,7 @@ async def schedule_model_cost_map_reload( if prisma_client is None: raise HTTPException(status_code=500, detail="Database connection not available") - # Update database with new reload configuration - await ConfigRepository(prisma_client).table.upsert( - where={"param_name": "model_cost_map_reload_config"}, - data={ - "create": { - "param_name": "model_cost_map_reload_config", - "param_value": safe_dumps({"interval_hours": hours, "force_reload": False}), - }, - "update": {"param_value": safe_dumps({"interval_hours": hours, "force_reload": False})}, - }, - ) - await invalidate_config_param("model_cost_map_reload_config") + await write_reload_interval(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME, hours) verbose_proxy_logger.info(f"Model cost map reload scheduled for every {hours} hours") @@ -15761,7 +15627,7 @@ async def schedule_model_cost_map_reload( "message": f"Model cost map reload scheduled for every {hours} hours", "status": "success", "interval_hours": hours, - "timestamp": datetime.utcnow().isoformat(), + "timestamp": utc_now().isoformat(), } except Exception as e: verbose_proxy_logger.exception(f"Failed to schedule model cost map reload: {str(e)}") @@ -15798,15 +15664,15 @@ async def cancel_model_cost_map_reload( raise HTTPException(status_code=500, detail="Database connection not available") # Remove reload configuration from database - await ConfigRepository(prisma_client).table.delete(where={"param_name": "model_cost_map_reload_config"}) - await invalidate_config_param("model_cost_map_reload_config") + await ConfigRepository(prisma_client).table.delete(where={"param_name": MODEL_COST_MAP_RELOAD_PARAM_NAME}) + await invalidate_config_param(MODEL_COST_MAP_RELOAD_PARAM_NAME) verbose_proxy_logger.info("Model cost map reload schedule cancelled") return { "message": "Model cost map reload schedule cancelled", "status": "success", - "timestamp": datetime.utcnow().isoformat(), + "timestamp": utc_now().isoformat(), } except Exception as e: verbose_proxy_logger.exception(f"Failed to cancel model cost map reload: {str(e)}") @@ -15835,66 +15701,13 @@ async def get_model_cost_map_reload_status( ) try: - global prisma_client, last_model_cost_map_reload - - verbose_proxy_logger.info(f"Checking model cost map reload status. Last reload: {last_model_cost_map_reload}") + global prisma_client if prisma_client is None: verbose_proxy_logger.info("No database connection, returning not scheduled") - return { - "scheduled": False, - "interval_hours": None, - "last_run": None, - "next_run": None, - } + return reload_schedule_status(None) - # Get reload configuration from database - config_record = await ConfigRepository(prisma_client).table.find_unique( - where={"param_name": "model_cost_map_reload_config"} - ) - - if config_record is None or config_record.param_value is None: - verbose_proxy_logger.info("No model cost map reload configuration found") - return { - "scheduled": False, - "interval_hours": None, - "last_run": None, - "next_run": None, - } - - config = config_record.param_value - interval_hours = config.get("interval_hours") - - if interval_hours is None: - verbose_proxy_logger.info("No interval configured, returning not scheduled") - return { - "scheduled": False, - "interval_hours": None, - "last_run": None, - "next_run": None, - } - - current_time = datetime.utcnow() - next_run = None - - # Use pod's in-memory last reload time - if last_model_cost_map_reload is not None: - try: - last_reload_time = datetime.fromisoformat(last_model_cost_map_reload) - time_since_last_reload = current_time - last_reload_time - hours_since_last_reload = time_since_last_reload.total_seconds() / 3600 - - if hours_since_last_reload < interval_hours: - next_run = (last_reload_time + timedelta(hours=interval_hours)).isoformat() - except Exception as e: - verbose_proxy_logger.warning(f"Error parsing last reload time: {e}") - - return { - "scheduled": True, - "interval_hours": interval_hours, - "last_run": last_model_cost_map_reload, - "next_run": next_run, - } + return reload_schedule_status(await read_reload_schedule(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME)) except Exception as e: verbose_proxy_logger.exception(f"Failed to get model cost map reload status: {str(e)}") raise HTTPException( @@ -15984,32 +15797,13 @@ async def reload_anthropic_beta_headers( # Immediately reload the beta headers config in the current pod from litellm.anthropic_beta_headers_manager import reload_beta_headers_config - new_config = reload_beta_headers_config() + new_config = await asyncio.to_thread(reload_beta_headers_config) - # Update pod's in-memory last reload time - global last_anthropic_beta_headers_reload - current_time = datetime.utcnow() - last_anthropic_beta_headers_reload = current_time.isoformat() + current_time = utc_now() + proxy_config.last_anthropic_beta_headers_reload = current_time - # Set force reload flag in database for other pods, preserving existing interval_hours - existing_beta_config = await ConfigRepository(prisma_client).table.find_unique( - where={"param_name": "anthropic_beta_headers_reload_config"} - ) - existing_beta_interval = None - if existing_beta_config and existing_beta_config.param_value: - existing_beta_interval = existing_beta_config.param_value.get("interval_hours") - - await ConfigRepository(prisma_client).table.upsert( - where={"param_name": "anthropic_beta_headers_reload_config"}, - data={ - "create": { - "param_name": "anthropic_beta_headers_reload_config", - "param_value": safe_dumps({"interval_hours": None, "force_reload": True}), - }, - "update": {"param_value": safe_dumps({"interval_hours": existing_beta_interval, "force_reload": True})}, - }, - ) - await invalidate_config_param("anthropic_beta_headers_reload_config") + # Stamp the shared last run and the reload request so other pods reload too + await record_manual_reload(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME, current_time) provider_count = sum(1 for k in new_config.keys() if k not in ["provider_aliases", "description"]) verbose_proxy_logger.info( @@ -16058,18 +15852,7 @@ async def schedule_anthropic_beta_headers_reload( if prisma_client is None: raise HTTPException(status_code=500, detail="Database connection not available") - # Update database with new reload configuration - await ConfigRepository(prisma_client).table.upsert( - where={"param_name": "anthropic_beta_headers_reload_config"}, - data={ - "create": { - "param_name": "anthropic_beta_headers_reload_config", - "param_value": safe_dumps({"interval_hours": hours, "force_reload": False}), - }, - "update": {"param_value": safe_dumps({"interval_hours": hours, "force_reload": False})}, - }, - ) - await invalidate_config_param("anthropic_beta_headers_reload_config") + await write_reload_interval(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME, hours) verbose_proxy_logger.info(f"Anthropic beta headers reload scheduled for every {hours} hours") @@ -16077,7 +15860,7 @@ async def schedule_anthropic_beta_headers_reload( "message": f"Anthropic beta headers reload scheduled for every {hours} hours", "status": "success", "interval_hours": hours, - "timestamp": datetime.utcnow().isoformat(), + "timestamp": utc_now().isoformat(), } except Exception as e: verbose_proxy_logger.exception(f"Failed to schedule anthropic beta headers reload: {str(e)}") @@ -16114,15 +15897,17 @@ async def cancel_anthropic_beta_headers_reload( raise HTTPException(status_code=500, detail="Database connection not available") # Remove reload configuration from database - await ConfigRepository(prisma_client).table.delete(where={"param_name": "anthropic_beta_headers_reload_config"}) - await invalidate_config_param("anthropic_beta_headers_reload_config") + await ConfigRepository(prisma_client).table.delete( + where={"param_name": ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME} + ) + await invalidate_config_param(ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME) verbose_proxy_logger.info("Anthropic beta headers reload schedule cancelled") return { "message": "Anthropic beta headers reload schedule cancelled", "status": "success", - "timestamp": datetime.utcnow().isoformat(), + "timestamp": utc_now().isoformat(), } except Exception as e: verbose_proxy_logger.exception(f"Failed to cancel anthropic beta headers reload: {str(e)}") @@ -16154,68 +15939,15 @@ async def get_anthropic_beta_headers_reload_status( ) try: - global prisma_client, last_anthropic_beta_headers_reload - - verbose_proxy_logger.info( - f"Checking anthropic beta headers reload status. Last reload: {last_anthropic_beta_headers_reload}" - ) + global prisma_client if prisma_client is None: verbose_proxy_logger.info("No database connection, returning not scheduled") - return { - "scheduled": False, - "interval_hours": None, - "last_run": None, - "next_run": None, - } + return reload_schedule_status(None) - # Get reload configuration from database - config_record = await ConfigRepository(prisma_client).table.find_unique( - where={"param_name": "anthropic_beta_headers_reload_config"} + return reload_schedule_status( + await read_reload_schedule(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME) ) - - if config_record is None or config_record.param_value is None: - verbose_proxy_logger.info("No anthropic beta headers reload configuration found") - return { - "scheduled": False, - "interval_hours": None, - "last_run": None, - "next_run": None, - } - - config = config_record.param_value - interval_hours = config.get("interval_hours") - - if interval_hours is None: - verbose_proxy_logger.info("No interval configured, returning not scheduled") - return { - "scheduled": False, - "interval_hours": None, - "last_run": None, - "next_run": None, - } - - current_time = datetime.utcnow() - next_run = None - - # Use pod's in-memory last reload time - if last_anthropic_beta_headers_reload is not None: - try: - last_reload_time = datetime.fromisoformat(last_anthropic_beta_headers_reload) - time_since_last_reload = current_time - last_reload_time - hours_since_last_reload = time_since_last_reload.total_seconds() / 3600 - - if hours_since_last_reload < interval_hours: - next_run = (last_reload_time + timedelta(hours=interval_hours)).isoformat() - except Exception as e: - verbose_proxy_logger.warning(f"Error parsing last reload time: {e}") - - return { - "scheduled": True, - "interval_hours": interval_hours, - "last_run": last_anthropic_beta_headers_reload, - "next_run": next_run, - } except Exception as e: verbose_proxy_logger.exception(f"Failed to get anthropic beta headers reload status: {str(e)}") raise HTTPException( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 37ea55f8c13..cca69741357 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -601,6 +601,8 @@ model LiteLLM_TagTable { model LiteLLM_Config { param_name String @id param_value Json? + last_run_at DateTime? + reload_requested_at DateTime? } // View spend, model, api_key per request diff --git a/ruff.toml b/ruff.toml index 2ea9d7260fb..b652e206f41 100644 --- a/ruff.toml +++ b/ruff.toml @@ -6,7 +6,7 @@ lint.extend-select = ["T20", "PGH004", "RUF008", "RUF009", "RUF100"] # litellm's own ruff config both rely on suppressions this config can't see. lint.external = [ # Enforced by the strict-rule gate (scripts/ruff_strict_gate.py + ruff-strict.toml) - "C901", + "C901", "TID251", # Enforced by upstream litellm's ruff config, but not run in this repo's CI "PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405", ] diff --git a/schema.prisma b/schema.prisma index 37ea55f8c13..cca69741357 100644 --- a/schema.prisma +++ b/schema.prisma @@ -601,6 +601,8 @@ model LiteLLM_TagTable { model LiteLLM_Config { param_name String @id param_value Json? + last_run_at DateTime? + reload_requested_at DateTime? } // 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 new file mode 100644 index 00000000000..d05a2f53a89 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_periodic_reload_schedule.py @@ -0,0 +1,213 @@ +from datetime import datetime, timezone +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy.common_utils.periodic_reload_schedule import ( + ReloadSchedule, + next_run_at, + parse_reload_schedule, + pod_reload_is_due, + read_reload_schedule, + record_manual_reload, + record_reload_run, + reload_schedule_status, + write_reload_interval, +) + +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): + return SimpleNamespace( + param_name="model_cost_map_reload_config", + param_value=param_value, + reload_requested_at=reload_requested_at, + last_run_at=last_run_at, + ) + + +def _mock_prisma(row=None): + 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.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) + ) + + assert schedule == ReloadSchedule(interval_hours=6, reload_requested_at=NOW, 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)) + ) + + assert (schedule.reload_requested_at, schedule.last_run_at) == (NOW, LAST_RUN) + + +@pytest.mark.parametrize( + "param_value", + [None, "not-a-dict", {}, {"interval_hours": "6"}, {"interval_hours": None}], +) +def test_parse_tolerates_unusable_param_values(param_value): + assert parse_reload_schedule(_row(param_value=param_value)).interval_hours is None + + +def test_parse_ignores_legacy_json_force_reload(): + """Rows written by pre-column versions carry force_reload in the JSON; honoring it + 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) + + +def test_status_reports_persisted_last_run_and_next_run(): + status = reload_schedule_status(ReloadSchedule(interval_hours=6, last_run_at=LAST_RUN)) + + assert status == { + "scheduled": True, + "interval_hours": 6, + "last_run": "2024-01-01T06:00:00+00:00", + "next_run": "2024-01-01T12:00:00+00:00", + } + + +def test_status_without_interval_is_not_scheduled(): + assert reload_schedule_status(None)["scheduled"] is False + assert reload_schedule_status(ReloadSchedule(last_run_at=LAST_RUN)) == { + "scheduled": False, + "interval_hours": None, + "last_run": "2024-01-01T06:00:00+00:00", + "next_run": None, + } + + +def test_next_run_needs_both_interval_and_last_run(): + assert next_run_at(ReloadSchedule(interval_hours=6)) is None + assert next_run_at(ReloadSchedule(last_run_at=LAST_RUN)) is None + + +@pytest.mark.parametrize( + "schedule, pod_last_reload, 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)), + 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(interval_hours=6, reload_requested_at=LAST_RUN), LAST_RUN, True), + (ReloadSchedule(reload_requested_at=NOW), None, False), + (ReloadSchedule(interval_hours=6, reload_requested_at=NOW), None, True), + (ReloadSchedule(), None, False), + (ReloadSchedule(interval_hours=6), None, True), + (ReloadSchedule(interval_hours=6), datetime(2024, 1, 1, 11, 0, tzinfo=timezone.utc), False), + (ReloadSchedule(interval_hours=6), LAST_RUN, True), + ], +) +def test_pod_reload_is_due(schedule, pod_last_reload, expected): + assert ( + pod_reload_is_due( + schedule=schedule, + pod_last_reload=pod_last_reload, + current_time=NOW, + description="test", + ) + is expected + ) + + +def test_manual_request_reaches_every_pod_that_has_not_seen_it(): + """Nothing clears reload_requested_at, so each pod reloads exactly once per request: + due while its own last reload predates the request, not due afterwards""" + schedule = ReloadSchedule(reload_requested_at=datetime(2024, 1, 1, 11, 0, tzinfo=timezone.utc)) + + assert pod_reload_is_due(schedule=schedule, pod_last_reload=LAST_RUN, current_time=NOW, description="test") is True + assert ( + pod_reload_is_due( + schedule=schedule, + pod_last_reload=datetime(2024, 1, 1, 11, 0, 1, tzinfo=timezone.utc), + current_time=NOW, + description="test", + ) + is False + ) + + +def test_pod_reload_decision_ignores_persisted_last_run(): + """A pod that has never reloaded must refresh its own copy even when another pod + already stamped last_run_at within the interval""" + assert ( + pod_reload_is_due( + schedule=ReloadSchedule(interval_hours=6, last_run_at=datetime(2024, 1, 1, 11, 59, tzinfo=timezone.utc)), + pod_last_reload=None, + current_time=NOW, + description="test", + ) + is True + ) + + +@pytest.mark.asyncio +async def test_read_reload_schedule_returns_none_for_missing_row(): + assert await read_reload_schedule(_mock_prisma(row=None), "model_cost_map_reload_config") is None + + +@pytest.mark.asyncio +async def test_read_reload_schedule_surfaces_request_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)) + + 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) + + +@pytest.mark.asyncio +async def test_write_reload_interval_touches_only_param_value(): + prisma_client = _mock_prisma() + + await write_reload_interval(prisma_client, "model_cost_map_reload_config", 12) + + data = prisma_client.db.litellm_config.upsert.await_args.kwargs["data"] + assert set(data["update"]) == {"param_value"} + assert set(data["create"]) == {"param_name", "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""" + prisma_client = _mock_prisma() + + await record_reload_run(prisma_client, "model_cost_map_reload_config", LAST_RUN) + + kwargs = prisma_client.db.litellm_config.update_many.await_args.kwargs + assert kwargs == {"data": {"last_run_at": LAST_RUN}, "where": {"param_name": "model_cost_map_reload_config"}} + prisma_client.db.litellm_config.upsert.assert_not_called() + + +@pytest.mark.asyncio +async def test_record_manual_reload_stamps_request_and_last_run(): + prisma_client = _mock_prisma() + + 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["create"] == { + "param_name": "model_cost_map_reload_config", + "last_run_at": LAST_RUN, + "reload_requested_at": LAST_RUN, + } diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_anthropic_beta.py b/tests/test_litellm/proxy/proxy_server/test_routes_anthropic_beta.py index 7ef29b71bf0..2e729751d66 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_anthropic_beta.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_anthropic_beta.py @@ -9,11 +9,9 @@ Routes covered: from __future__ import annotations -import json from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock -import pytest from .conftest import VOLATILE_KEYS, normalize @@ -99,22 +97,24 @@ def test_reload_anthropic_beta_headers_admin_success(client, auth_as, monkeypatc "providers_count": 2, "timestamp": "", } - # And the upsert was actually invoked (force_reload write). prisma.db.litellm_config.upsert.assert_awaited_once() def test_reload_anthropic_beta_headers_preserves_existing_interval( client, auth_as, monkeypatch ): - """When an existing reload config has an interval set, the force-reload - write must preserve that interval (the route reads it back then upserts - with the same number). This pins the read-then-write behaviour.""" + """The manual-reload write only stamps the job-owned timestamp columns; param_value + (the admin-owned interval) is not in the payload, so it cannot be clobbered.""" + from datetime import datetime + from litellm.proxy._types import LitellmUserRoles _stub_reload_beta_headers(monkeypatch) existing = SimpleNamespace( param_name="anthropic_beta_headers_reload_config", - param_value={"interval_hours": 12, "force_reload": False}, + param_value={"interval_hours": 12}, + reload_requested_at=None, + last_run_at=None, ) prisma = _make_prisma_with_config(config_record=existing) _install_prisma(monkeypatch, prisma) @@ -123,17 +123,12 @@ def test_reload_anthropic_beta_headers_preserves_existing_interval( response = client.post("/reload/anthropic_beta_headers") assert response.status_code == 200 - # The update branch's interval_hours was sourced from the existing record. call_kwargs = prisma.db.litellm_config.upsert.await_args.kwargs - data = call_kwargs["data"] - update_payload = data["update"]["param_value"] - parsed = ( - json.loads(update_payload) - if isinstance(update_payload, str) - else update_payload - ) - assert parsed["interval_hours"] == 12 - assert parsed["force_reload"] is True + update_payload = call_kwargs["data"]["update"] + assert set(update_payload) == {"last_run_at", "reload_requested_at"} + assert isinstance(update_payload["last_run_at"], datetime) + assert update_payload["reload_requested_at"] == update_payload["last_run_at"] + assert "param_value" not in call_kwargs["data"]["create"] def test_reload_anthropic_beta_headers_not_admin_forbidden(client, auth_as): @@ -291,17 +286,16 @@ def test_get_anthropic_beta_headers_reload_status_scheduled( ): """When a config row with ``interval_hours`` is present, ``scheduled`` is True and ``interval_hours`` echoes the DB value. Pins the full response shape.""" - from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles record = SimpleNamespace( param_name="anthropic_beta_headers_reload_config", - param_value={"interval_hours": 6, "force_reload": False}, + param_value={"interval_hours": 6}, + reload_requested_at=None, + last_run_at=None, ) prisma = _make_prisma_with_config(config_record=record) _install_prisma(monkeypatch, prisma) - # No prior reload — next_run stays None. - monkeypatch.setattr(ps, "last_anthropic_beta_headers_reload", None) with auth_as(LitellmUserRoles.PROXY_ADMIN): response = client.get("/schedule/anthropic_beta_headers_reload/status") @@ -315,6 +309,35 @@ def test_get_anthropic_beta_headers_reload_status_scheduled( } +def test_get_anthropic_beta_headers_reload_status_reports_persisted_last_run( + client, auth_as, monkeypatch +): + """last_run/next_run come from the DB row, so status survives pod restarts.""" + from datetime import datetime, timezone + + from litellm.proxy._types import LitellmUserRoles + + record = SimpleNamespace( + param_name="anthropic_beta_headers_reload_config", + param_value={"interval_hours": 6}, + reload_requested_at=None, + last_run_at=datetime(2024, 1, 1, 6, 0, 0, tzinfo=timezone.utc), + ) + prisma = _make_prisma_with_config(config_record=record) + _install_prisma(monkeypatch, prisma) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/schedule/anthropic_beta_headers_reload/status") + + assert response.status_code == 200 + assert normalize(response.json()) == { + "scheduled": True, + "interval_hours": 6, + "last_run": "2024-01-01T06:00:00+00:00", + "next_run": "2024-01-01T12:00:00+00:00", + } + + def test_get_anthropic_beta_headers_reload_status_not_scheduled_no_db( client, auth_as, monkeypatch ): @@ -339,11 +362,15 @@ def test_get_anthropic_beta_headers_reload_status_no_interval_unscheduled( client, auth_as, monkeypatch ): """Config row present but ``interval_hours`` is None → unscheduled response.""" + from datetime import datetime, timezone + from litellm.proxy._types import LitellmUserRoles record = SimpleNamespace( param_name="anthropic_beta_headers_reload_config", - param_value={"interval_hours": None, "force_reload": True}, + param_value={"interval_hours": None}, + reload_requested_at=datetime(2024, 1, 1, 6, 0, 0, tzinfo=timezone.utc), + last_run_at=None, ) prisma = _make_prisma_with_config(config_record=record) _install_prisma(monkeypatch, prisma) 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 16e410f1b1e..680ae1b5663 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 @@ -12,7 +12,6 @@ from __future__ import annotations from unittest.mock import AsyncMock, MagicMock -import pytest from .conftest import VOLATILE_KEYS, normalize @@ -83,6 +82,9 @@ def test_reload_model_cost_map_happy(client, auth_as, monkeypatch, mock_prisma): "timestamp": "", } 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"] def test_reload_model_cost_map_not_admin_forbidden(client, auth_as): @@ -138,6 +140,9 @@ def test_schedule_model_cost_map_reload_happy( "timestamp": "", } assert table.upsert.await_count == 1 + upsert_data = table.upsert.await_args.kwargs["data"] + assert set(upsert_data["update"]) == {"param_value"} + assert set(upsert_data["create"]) == {"param_name", "param_value"} def test_schedule_model_cost_map_reload_invalid_hours( @@ -249,10 +254,11 @@ 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, "force_reload": False} + config_row.param_value = {"interval_hours": 12} + config_row.reload_requested_at = None + config_row.last_run_at = None table.find_unique = AsyncMock(return_value=config_row) monkeypatch.setattr(ps, "prisma_client", mock_prisma) - monkeypatch.setattr(ps, "last_model_cost_map_reload", None) with auth_as(LitellmUserRoles.PROXY_ADMIN): response = client.get("/schedule/model_cost_map_reload/status") @@ -265,19 +271,50 @@ def test_get_model_cost_map_reload_status_scheduled( } -def test_get_model_cost_map_reload_status_no_config_not_scheduled( +def test_get_model_cost_map_reload_status_reports_persisted_last_run( client, auth_as, monkeypatch, mock_prisma ): - """Config row exists but interval_hours=None → not scheduled.""" + """last_run/next_run come from the DB row, so status survives pod restarts.""" + from datetime import datetime, timezone + from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles table = _attach_litellm_config(mock_prisma) config_row = MagicMock() - config_row.param_value = {"interval_hours": None, "force_reload": True} + config_row.param_value = {"interval_hours": 6} + config_row.reload_requested_at = None + 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) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/schedule/model_cost_map_reload/status") + assert response.status_code == 200 + assert normalize(response.json()) == { + "scheduled": True, + "interval_hours": 6, + "last_run": "2024-01-01T06:00:00+00:00", + "next_run": "2024-01-01T12:00:00+00:00", + } + + +def test_get_model_cost_map_reload_status_no_config_not_scheduled( + client, auth_as, monkeypatch, mock_prisma +): + """A row left behind by a manual reload (interval_hours=None) → not scheduled.""" + from datetime import datetime, timezone + + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + 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.last_run_at = None table.find_unique = AsyncMock(return_value=config_row) monkeypatch.setattr(ps, "prisma_client", mock_prisma) - monkeypatch.setattr(ps, "last_model_cost_map_reload", None) with auth_as(LitellmUserRoles.PROXY_ADMIN): response = client.get("/schedule/model_cost_map_reload/status") diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 536d24d4b4e..ea332a2e960 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -755,6 +755,50 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch): assert len(mock_scheduler_calls) > 0 +@pytest.mark.asyncio +async def test_periodic_reload_job_scheduled_without_store_model_in_db(monkeypatch): + """ + Regression (LIT-4882): reload schedules configured from the Admin UI live in the DB and + must fire even without store_model_in_db, which used to gate the job that ran them + """ + monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False) + monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) + from apscheduler.schedulers.asyncio import AsyncIOScheduler + + from litellm.proxy.proxy_server import ProxyStartupEvent + from litellm.proxy.utils import ProxyLogging + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=None) + mock_proxy_logging = MagicMock(spec=ProxyLogging) + mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_config = AsyncMock() + scheduler = AsyncIOScheduler() + + try: + with ( + patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", False), + patch("litellm.proxy.proxy_server.get_secret_bool", return_value=False), + patch( + "litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler + ), + ): + await ProxyStartupEvent.initialize_scheduled_background_jobs( + general_settings={}, + prisma_client=mock_prisma_client, + proxy_budget_rescheduler_min_time=1, + proxy_budget_rescheduler_max_time=2, + proxy_batch_write_at=5, + proxy_logging_obj=mock_proxy_logging, + ) + + assert scheduler.get_job("periodic_reload_job") is not None + assert scheduler.get_job("add_deployment_job") is None + finally: + scheduler.shutdown(wait=False) + + @pytest.mark.asyncio async def test_initialize_scheduled_jobs_uses_configured_config_reload_interval(monkeypatch): """ @@ -3592,6 +3636,20 @@ async def test_chat_completion_result_no_nested_none_values(): # ============================================================================ +def _reload_schedule_row( + param_value: dict, + *, + reload_requested_at: datetime | None = None, + 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, + last_run_at=last_run_at, + ) + + class TestPriceDataReloadAPI: """Test cases for price data reload API endpoints""" @@ -3628,9 +3686,6 @@ 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) response = client_with_auth.post("/reload/model_cost_map") @@ -3698,7 +3753,7 @@ class TestPriceDataReloadAPI: assert "Failed to reload model cost map" in data["detail"] def test_schedule_model_cost_map_reload_admin_access(self, client_with_auth): - """Test that admin users can schedule periodic reload""" + """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) @@ -3712,6 +3767,17 @@ class TestPriceDataReloadAPI: assert "message" in data assert "timestamp" in data + call_args = mock_prisma.db.litellm_config.upsert.call_args + assert call_args[1]["where"] == { + "param_name": "model_cost_map_reload_config" + } + update_payload = call_args[1]["data"]["update"] + assert set(update_payload.keys()) == {"param_value"} + assert json.loads(update_payload["param_value"]) == {"interval_hours": 6} + create_payload = call_args[1]["data"]["create"] + assert set(create_payload.keys()) == {"param_name", "param_value"} + assert json.loads(create_payload["param_value"]) == {"interval_hours": 6} + def test_schedule_model_cost_map_reload_non_admin_access(self, client_with_auth): """Test that non-admin users cannot schedule periodic reload""" # Mock non-admin user @@ -3763,35 +3829,28 @@ class TestPriceDataReloadAPI: assert "Admin role required" in data["detail"] def test_get_model_cost_map_reload_status_admin_access(self, client_with_auth): - """Test that admin users can get reload status""" + """ + Regression (LIT-4882): status is served purely from the DB row, so a restarted pod + with no in-memory last-reload time still reports the real last/next run + """ + proxy_server_module.proxy_config.last_model_cost_map_reload = None + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: - # Mock database config record - mock_config = MagicMock() - mock_config.param_value = {"interval_hours": 6, "force_reload": False} mock_prisma.db.litellm_config.find_unique = AsyncMock( - return_value=mock_config + return_value=_reload_schedule_row( + {"interval_hours": 6}, + last_run_at=datetime(2024, 1, 1, 6, 0, tzinfo=timezone.utc), + ) ) - # Mock the last reload time and current time - with patch( - "litellm.proxy.proxy_server.last_model_cost_map_reload", - "2024-01-01T06:00:00", - ): - with patch("litellm.proxy.proxy_server.datetime") as mock_datetime: - # Mock current time to be 1 hour after last reload - mock_datetime.utcnow.return_value = datetime(2024, 1, 1, 7, 0, 0) - mock_datetime.fromisoformat = datetime.fromisoformat + response = client_with_auth.get("/schedule/model_cost_map_reload/status") - response = client_with_auth.get( - "/schedule/model_cost_map_reload/status" - ) - - assert response.status_code == 200 - data = response.json() - assert data["scheduled"] == True - assert data["interval_hours"] == 6 - assert data["last_run"] == "2024-01-01T06:00:00" - assert data["next_run"] == "2024-01-01T12:00:00" + assert response.status_code == 200 + data = response.json() + assert data["scheduled"] is True + assert data["interval_hours"] == 6 + assert data["last_run"] == "2024-01-01T06:00:00+00:00" + assert data["next_run"] == "2024-01-01T12:00:00+00:00" def test_get_model_cost_map_reload_status_non_admin_access(self, client_with_auth): """Test that non-admin users cannot get reload status""" @@ -3822,36 +3881,44 @@ class TestPriceDataReloadAPI: assert data["next_run"] == None def test_get_model_cost_map_reload_status_no_interval(self, client_with_auth): - """Test that status returns not scheduled when no interval is configured""" + """A row left behind by a manual reload (no interval) must not read as scheduled""" with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: - # Mock config with no interval - mock_config = MagicMock() - mock_config.param_value = {"interval_hours": None, "force_reload": False} mock_prisma.db.litellm_config.find_unique = AsyncMock( - return_value=mock_config + return_value=_reload_schedule_row( + {"interval_hours": None}, + reload_requested_at=datetime(2024, 1, 1, 6, 0, tzinfo=timezone.utc), + ) ) response = client_with_auth.get("/schedule/model_cost_map_reload/status") assert response.status_code == 200 data = response.json() - assert data["scheduled"] == False - assert data["interval_hours"] == None - assert data["last_run"] == None - assert data["next_run"] == None + assert data["scheduled"] is False + assert data["interval_hours"] is None + assert data["last_run"] is None + assert data["next_run"] is None + + def test_get_model_cost_map_reload_status_before_first_run(self, client_with_auth): + """Scheduled but never executed: no last_run_at means no next_run can be computed""" + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_prisma.db.litellm_config.find_unique = AsyncMock( + return_value=_reload_schedule_row({"interval_hours": 6}) + ) + + response = client_with_auth.get("/schedule/model_cost_map_reload/status") + + assert response.status_code == 200 + data = response.json() + assert data["scheduled"] is True + assert data["interval_hours"] == 6 + assert data["last_run"] is None + assert data["next_run"] is None class TestPriceDataReloadIntegration: """Integration tests for the complete price data reload feature""" - @pytest.fixture(autouse=True) - def _flush_litellm_config_cache(self): - from litellm.proxy.utils import litellm_config_cache - - litellm_config_cache.flush_cache() - yield - litellm_config_cache.flush_cache() - @pytest.fixture def client_with_auth(self): """Create a test client with authentication""" @@ -3890,9 +3957,6 @@ class TestPriceDataReloadIntegration: # 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) # Test reload endpoint @@ -3907,122 +3971,193 @@ class TestPriceDataReloadIntegration: _invalidate_model_cost_lowercase_map() def test_distributed_reload_check_function(self): - """Test the _check_and_reload_model_cost_map function""" - from litellm.proxy.proxy_server import ProxyConfig - from litellm.proxy.utils import litellm_config_cache - - proxy_config = ProxyConfig() - - # Mock prisma client - mock_prisma = MagicMock() - - # Test case 1: No config in database - mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) - # _check_and_reload_model_cost_map routes through get_config_param, - # which calls prisma.get_generic_data on a cache miss. - mock_prisma.get_generic_data = AsyncMock(return_value=None) - - # Should return early without reloading - asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) - - # Test case 2: Config with interval but not time to reload - litellm_config_cache.flush_cache() - mock_config = MagicMock() - mock_config.param_value = {"interval_hours": 6, "force_reload": False} - mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config) - mock_prisma.get_generic_data = AsyncMock(return_value=mock_config) - - # Mock current time and last reload time - with patch( - "litellm.proxy.proxy_server.last_model_cost_map_reload", - "2024-01-01T06:00:00", - ): - with patch("litellm.proxy.proxy_server.datetime") as mock_datetime: - mock_datetime.utcnow.return_value = datetime( - 2024, 1, 1, 7, 0, 0 - ) # 1 hour later - - # Should not reload (only 1 hour passed, need 6) - asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) - - # Test case 3: Config with force reload - litellm_config_cache.flush_cache() - mock_config.param_value = {"interval_hours": 6, "force_reload": True} - mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config) - mock_prisma.get_generic_data = AsyncMock(return_value=mock_config) - mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) - - 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: - mock_get_map.return_value = { - "gpt-3.5-turbo": {"input_cost_per_token": 0.001} - } - - # Should reload due to force flag - asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) - - # Verify force_reload was reset to False - mock_prisma.db.litellm_config.upsert.assert_called() - call_args = mock_prisma.db.litellm_config.upsert.call_args - # The param_value is now a JSON string, so we need to parse it - param_value_json = call_args[1]["data"]["update"]["param_value"] - param_value_dict = json.loads(param_value_json) - assert param_value_dict["force_reload"] == False - assert param_value_dict.get("interval_hours") == 6 - finally: - litellm.model_cost = original_model_cost - _invalidate_model_cost_lowercase_map() - - def test_distributed_reload_preserves_interval_hours(self): - """Test that _check_and_reload_model_cost_map preserves interval_hours after reload. - - Regression test: the update branch of the upsert was previously dropping - interval_hours, causing scheduled reloads to self-destruct after first execution. + """ + 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 """ from litellm.proxy.proxy_server import ProxyConfig proxy_config = ProxyConfig() mock_prisma = MagicMock() - - # Set up config with interval_hours=24 and force_reload=True to trigger reload - mock_config = MagicMock() - mock_config.param_value = {"interval_hours": 24, "force_reload": True} - mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config) - # _check_and_reload_model_cost_map 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.update_many = AsyncMock(return_value=None) + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + + asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) + + mock_prisma.db.litellm_config.update_many.assert_not_called() + assert proxy_config.last_model_cost_map_reload is None + + 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, + ) + ) + proxy_config.last_model_cost_map_reload = frozen_now - timedelta(minutes=1) 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: - mock_get_map.return_value = {"gpt-4": {"input_cost_per_token": 0.001}} + 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-3.5-turbo": {"input_cost_per_token": 0.001} + } asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) - # Verify the upsert update branch preserves interval_hours - mock_prisma.db.litellm_config.upsert.assert_called() - call_args = mock_prisma.db.litellm_config.upsert.call_args - param_value_json = call_args[1]["data"]["update"]["param_value"] - param_value_dict = json.loads(param_value_json) - assert param_value_dict["force_reload"] == False - assert param_value_dict["interval_hours"] == 24, ( - "interval_hours must be preserved in the update branch; " - "dropping it causes the schedule to self-destruct" + assert litellm.model_cost["gpt-3.5-turbo"] == { + "input_cost_per_token": 0.001 + } + assert proxy_config.last_model_cost_map_reload == frozen_now + assert mock_prisma.db.litellm_config.update_many.call_args[1] == { + "data": {"last_run_at": frozen_now}, + "where": {"param_name": "model_cost_map_reload_config"}, + } + mock_prisma.db.litellm_config.upsert.assert_not_called() + 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 + """ + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + 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( + {"interval_hours": 6}, + reload_requested_at=datetime(2024, 1, 1, 6, 0, tzinfo=timezone.utc), + last_run_at=datetime(2024, 1, 1, 6, 0, tzinfo=timezone.utc), + ) + ) + frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc) + pod_last_reload = frozen_now - timedelta(minutes=1) + proxy_config.last_model_cost_map_reload = pod_last_reload + + 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), + ): + asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) + + mock_get_map.assert_not_called() + mock_prisma.db.litellm_config.update_many.assert_not_called() + assert proxy_config.last_model_cost_map_reload == pod_last_reload + finally: + litellm.model_cost = original_model_cost + _invalidate_model_cost_lowercase_map() + + def test_periodic_reload_uses_pod_local_last_reload(self): + """ + Each pod decides from its own last reload, so a pod that never reloaded refreshes even + when the shared row was just stamped, and stays put while inside the interval + """ + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + mock_prisma = MagicMock() + mock_prisma.db.litellm_config.find_unique = AsyncMock( + return_value=_reload_schedule_row( + {"interval_hours": 6}, + last_run_at=datetime(2024, 1, 1, 6, 59, tzinfo=timezone.utc), + ) + ) + mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None) + frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc) + + 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-test": {"input_cost_per_token": 0.5}} + + asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) + + assert litellm.model_cost["gpt-4-test"] == {"input_cost_per_token": 0.5} + assert proxy_config.last_model_cost_map_reload == frozen_now + assert mock_prisma.db.litellm_config.update_many.call_args[1][ + "data" + ] == {"last_run_at": frozen_now} + + mock_get_map.reset_mock() + mock_prisma.db.litellm_config.update_many.reset_mock() + proxy_config.last_model_cost_map_reload = frozen_now - timedelta(hours=1) + + asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) + + mock_get_map.assert_not_called() + mock_prisma.db.litellm_config.update_many.assert_not_called() + assert proxy_config.last_model_cost_map_reload == frozen_now - timedelta( + hours=1 ) finally: litellm.model_cost = original_model_cost _invalidate_model_cost_lowercase_map() - def test_manual_reload_preserves_interval_hours(self): - """Test that manual reload via /reload/model_cost_map preserves existing interval_hours. + def test_distributed_reload_stamps_last_run_without_creating_row(self): + """ + Regression: the job's write carries neither param_value (which would clobber the + admin-configured interval) nor a create branch (which would resurrect a schedule + a concurrent cancel just deleted) + """ + from litellm.proxy.proxy_server import ProxyConfig - Regression test: the manual reload endpoint was overwriting param_value with - only force_reload=True, dropping any existing interval_hours schedule. + proxy_config = ProxyConfig() + mock_prisma = MagicMock() + 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.update_many = AsyncMock(return_value=None) + frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc) + + 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}} + + asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) + + assert mock_prisma.db.litellm_config.update_many.call_args[1] == { + "data": {"last_run_at": frozen_now}, + "where": {"param_name": "model_cost_map_reload_config"}, + } + mock_prisma.db.litellm_config.upsert.assert_not_called() + mock_prisma.db.litellm_config.create.assert_not_called() + finally: + litellm.model_cost = original_model_cost + _invalidate_model_cost_lowercase_map() + + def test_manual_reload_preserves_interval_hours(self): + """ + Regression: manual reload owns only the run columns, so it never reads or rewrites + param_value and cannot destroy an existing schedule """ from litellm.proxy._types import LitellmUserRoles from litellm.proxy.proxy_server import cleanup_router_config_variables @@ -4036,88 +4171,76 @@ class TestPriceDataReloadIntegration: mock_auth.user_role = LitellmUserRoles.PROXY_ADMIN app.dependency_overrides[user_api_key_auth] = lambda: mock_auth client = TestClient(app) + frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc) 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: + with ( + patch( + "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map" + ) as mock_get_map, + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + 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) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: - # Simulate existing config with a schedule - mock_existing = MagicMock() - mock_existing.param_value = { - "interval_hours": 12, - "force_reload": False, - } - mock_prisma.db.litellm_config.find_unique = AsyncMock( - return_value=mock_existing - ) - mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) + response = client.post("/reload/model_cost_map") + assert response.status_code == 200 - response = client.post("/reload/model_cost_map") - assert response.status_code == 200 - - # Verify interval_hours was preserved in the upsert - mock_prisma.db.litellm_config.upsert.assert_called() - call_args = mock_prisma.db.litellm_config.upsert.call_args - param_value_json = call_args[1]["data"]["update"]["param_value"] - param_value_dict = json.loads(param_value_json) - assert param_value_dict["force_reload"] == True - assert param_value_dict["interval_hours"] == 12, ( - "interval_hours must be preserved when manual reload sets force_reload; " - "dropping it destroys any existing schedule" - ) + mock_prisma.db.litellm_config.find_unique.assert_not_called() + 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, + } + assert call_args[1]["data"]["create"] == { + "param_name": "model_cost_map_reload_config", + "last_run_at": frozen_now, + "reload_requested_at": frozen_now, + } + assert ( + proxy_server_module.proxy_config.last_model_cost_map_reload + == frozen_now + ) finally: litellm.model_cost = original_model_cost _invalidate_model_cost_lowercase_map() def test_anthropic_beta_headers_reload_preserves_interval_hours(self): - """Test that _check_and_reload_anthropic_beta_headers preserves interval_hours after reload. - - Regression test: the update branch of the upsert was dropping interval_hours, - identical to the model cost map bug. - """ + """Same column-only write contract as the cost map job, for the beta headers schedule""" from litellm.proxy.proxy_server import ProxyConfig proxy_config = ProxyConfig() mock_prisma = MagicMock() - - # Set up config with interval_hours=12 and force_reload=True to trigger reload - mock_config = MagicMock() - mock_config.param_value = {"interval_hours": 12, "force_reload": True} - 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.find_unique = AsyncMock( + return_value=_reload_schedule_row({"interval_hours": 12}) + ) mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) + mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None) + frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc) - with patch( - "litellm.anthropic_beta_headers_manager.reload_beta_headers_config" - ) as mock_reload: + with ( + patch( + "litellm.anthropic_beta_headers_manager.reload_beta_headers_config" + ) as mock_reload, + patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), + ): mock_reload.return_value = {"anthropic": {"beta_header": "test-value"}} asyncio.run( proxy_config._check_and_reload_anthropic_beta_headers(mock_prisma) ) - # Verify the upsert update branch preserves interval_hours - mock_prisma.db.litellm_config.upsert.assert_called() - call_args = mock_prisma.db.litellm_config.upsert.call_args - param_value_json = call_args[1]["data"]["update"]["param_value"] - param_value_dict = json.loads(param_value_json) - assert param_value_dict["force_reload"] == False - assert param_value_dict["interval_hours"] == 12, ( - "interval_hours must be preserved in the update branch; " - "dropping it causes the schedule to self-destruct" - ) + assert proxy_config.last_anthropic_beta_headers_reload == frozen_now + assert mock_prisma.db.litellm_config.update_many.call_args[1] == { + "data": {"last_run_at": frozen_now}, + "where": {"param_name": "anthropic_beta_headers_reload_config"}, + } + mock_prisma.db.litellm_config.upsert.assert_not_called() def test_anthropic_beta_headers_manual_reload_preserves_interval_hours(self): - """Test that manual reload via /reload/anthropic_beta_headers preserves existing interval_hours. - - Regression test: the manual reload endpoint was overwriting param_value with - only force_reload=True, dropping any existing interval_hours schedule. - """ + """Manual beta headers reload stamps the run columns only, leaving the schedule intact""" from litellm.proxy._types import LitellmUserRoles from litellm.proxy.proxy_server import cleanup_router_config_variables @@ -4130,34 +4253,36 @@ class TestPriceDataReloadIntegration: mock_auth.user_role = LitellmUserRoles.PROXY_ADMIN app.dependency_overrides[user_api_key_auth] = lambda: mock_auth client = TestClient(app) + frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc) - with patch( - "litellm.anthropic_beta_headers_manager.reload_beta_headers_config" - ) as mock_reload: + with ( + patch( + "litellm.anthropic_beta_headers_manager.reload_beta_headers_config" + ) as mock_reload, + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), + ): mock_reload.return_value = {"anthropic": {"beta_header": "test-value"}} + mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: - # Simulate existing config with a schedule - mock_existing = MagicMock() - mock_existing.param_value = {"interval_hours": 8, "force_reload": False} - mock_prisma.db.litellm_config.find_unique = AsyncMock( - return_value=mock_existing - ) - mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) + response = client.post("/reload/anthropic_beta_headers") + assert response.status_code == 200 - response = client.post("/reload/anthropic_beta_headers") - assert response.status_code == 200 - - # Verify interval_hours was preserved in the upsert - mock_prisma.db.litellm_config.upsert.assert_called() - call_args = mock_prisma.db.litellm_config.upsert.call_args - param_value_json = call_args[1]["data"]["update"]["param_value"] - param_value_dict = json.loads(param_value_json) - assert param_value_dict["force_reload"] == True - assert param_value_dict["interval_hours"] == 8, ( - "interval_hours must be preserved when manual reload sets force_reload; " - "dropping it destroys any existing schedule" - ) + mock_prisma.db.litellm_config.find_unique.assert_not_called() + 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, + } + assert call_args[1]["data"]["create"] == { + "param_name": "anthropic_beta_headers_reload_config", + "last_run_at": frozen_now, + "reload_requested_at": frozen_now, + } + assert ( + proxy_server_module.proxy_config.last_anthropic_beta_headers_reload + == frozen_now + ) def test_config_file_parsing(self): """Test parsing of config file with reload settings""" @@ -4187,64 +4312,6 @@ model_list: assert "model_list" in config assert len(config["model_list"]) == 2 - def test_database_config_storage(self): - """Test that configuration is properly stored in database""" - # Mock prisma client - mock_prisma = MagicMock() - - # Test the database upsert call that would be made by the schedule endpoint - mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) - - # Simulate the database call that the schedule endpoint would make - asyncio.run( - mock_prisma.db.litellm_config.upsert( - where={"param_name": "model_cost_map_reload_config"}, - data={ - "create": { - "param_name": "model_cost_map_reload_config", - "param_value": {"interval_hours": 6, "force_reload": False}, - }, - "update": { - "param_value": {"interval_hours": 6, "force_reload": False} - }, - }, - ) - ) - - # Verify database upsert was called with correct data - mock_prisma.db.litellm_config.upsert.assert_called_once() - call_args = mock_prisma.db.litellm_config.upsert.call_args - assert call_args[1]["where"]["param_name"] == "model_cost_map_reload_config" - assert call_args[1]["data"]["create"]["param_value"]["interval_hours"] == 6 - assert call_args[1]["data"]["create"]["param_value"]["force_reload"] == False - - def test_manual_reload_force_flag(self): - """Test that manual reload sets force flag correctly""" - # Mock prisma client - mock_prisma = MagicMock() - - # Test the database upsert call that would be made by the manual reload endpoint - mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) - - # Simulate the database call that the manual reload endpoint would make - asyncio.run( - mock_prisma.db.litellm_config.upsert( - where={"param_name": "model_cost_map_reload_config"}, - data={ - "create": { - "param_name": "model_cost_map_reload_config", - "param_value": {"interval_hours": None, "force_reload": True}, - }, - "update": {"param_value": {"force_reload": True}}, - }, - ) - ) - - # Verify force_reload flag was set - mock_prisma.db.litellm_config.upsert.assert_called_once() - call_args = mock_prisma.db.litellm_config.upsert.call_args - assert call_args[1]["data"]["update"]["param_value"]["force_reload"] == True - @pytest.mark.asyncio async def test_add_router_settings_from_db_config_merge_logic():