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..2a8460a9b60 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260729000000_add_reload_tracking_to_litellm_config/migration.sql @@ -0,0 +1,3 @@ +-- AlterTable +ALTER TABLE "LiteLLM_Config" ADD COLUMN IF NOT EXISTS "last_run_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 0d7fa8692c8..17339541fd9 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_revision BigInt @default(0) } // View spend, model, api_key per request diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index acb3fddfbad..2043a9e2f89 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -14,6 +14,7 @@ import os import random from collections.abc import Awaitable, Callable from dataclasses import dataclass +from datetime import datetime, timezone from importlib.resources import files from typing import Final, Protocol @@ -325,6 +326,7 @@ class ModelCostMapSourceInfo: url: str | None = None is_env_forced: bool = False fallback_reason: str | None = None + loaded_at: "datetime | None" = None # Module-level singleton tracking the source of the current cost map @@ -349,6 +351,11 @@ def get_model_cost_map_source_info() -> dict: } +def get_model_cost_map_loaded_at() -> "datetime | None": + """When this process last loaded its cost map, stamped at the start of every load""" + return _cost_map_source_info.loaded_at + + def _expand_model_aliases(model_cost: dict) -> dict: """ Expand ``aliases`` lists in model cost entries into top-level entries. @@ -428,6 +435,7 @@ def get_model_cost_map(url: str) -> dict: The full backup dict is only parsed when it must be *returned* as a fallback — it is never held in memory long-term. """ + _cost_map_source_info.loaded_at = datetime.now(timezone.utc) # Note: can't use get_secret_bool here — this runs during litellm.__init__ # before litellm._key_management_settings is set. if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true": 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..c3228ee8ac0 --- /dev/null +++ b/litellm/proxy/common_utils/periodic_reload_schedule.py @@ -0,0 +1,231 @@ +""" +Persistence for the admin-configured periodic model cost map reload schedule 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_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. A booting +pod starts at revision 0 rather than adopting the published one, because it cannot know +whether that request predates the prices it fetched at import. Interval reloads stay +per-pod, driven by when that pod's own copy of the data was loaded. +""" + +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, evict_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" + + +class _RevisionIncrement(TypedDict): + increment: int + + +class _ConfigRowWrite(TypedDict, total=False): + param_name: str + param_value: str + last_run_at: datetime + reload_revision: int | _RevisionIncrement + + +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_revision: int = 0 + 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: + """``param_value`` is written as serialized JSON, and a raw row read can hand it back + either decoded or still as a string depending on the driver, so accept both rather than + reading a string as no schedule at all. Mirrors ``ConfigRepository.get_param``""" + try: + if isinstance(param_value, str): + return _IntervalConfig.model_validate_json(param_value).interval_hours + 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_revision=int(row.reload_revision or 0), + 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_applied_revision: int, + pod_data_loaded_at: datetime, + current_time: datetime, + description: str, +) -> bool: + """ + Whether this pod should reload now. A revision it has not applied means a manual reload + it has not served. A pod starts at revision 0, so it serves any request published before + it booted; that costs one redundant fetch per boot and is what keeps a request from being + marked applied against data fetched before it. 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_revision != pod_applied_revision: + verbose_proxy_logger.info("%s reload triggered by manual reload request", description) + return True + if schedule.interval_hours is None: + return False + if schedule.last_run_at is None: + verbose_proxy_logger.info("%s reload triggered - schedule has never run", description) + return True + hours_since_data_loaded = (current_time - pod_data_loaded_at).total_seconds() / 3600 + if hours_since_data_loaded < schedule.interval_hours: + return False + verbose_proxy_logger.info( + "%s reload triggered by interval. Hours since data loaded: %.2f, Interval: %s", + description, + hours_since_data_loaded, + 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 evict_config_param(param_name) + + +async def clear_reload_interval(prisma_client: PrismaClient, param_name: str) -> None: + """Admin-owned write: drops the schedule but keeps the row, because the revision counter + identifies a request rather than ordering one and so can never reuse a number. Deleting + the row restarts it, and a reissued revision matches what pods already applied, so their + next manual reload is silently skipped. The interval is nulled inside the JSON rather + than by nulling the column, which prisma rejects for a ``Json?`` field""" + await _config_table(prisma_client).update_many( + data={"param_value": safe_dumps({"interval_hours": None})}, + where={"param_name": param_name}, + ) + await evict_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 evict_config_param(param_name) + + +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_revision": 1}, + "update": {"last_run_at": ran_at, "reload_revision": {"increment": 1}}, + }, + ) + await evict_config_param(param_name) + return int(row.reload_revision) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3380decabf7..fb9c4e67aad 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -319,6 +319,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 ( + MODEL_COST_MAP_RELOAD_PARAM_NAME, + clear_reload_interval, + 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 @@ -2062,9 +2073,7 @@ async_result: Final = None celery_app_conn: Final = None celery_fn: Final = 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 @@ -3839,6 +3848,17 @@ def resolve_complexity_router_plugins( ) +def _swap_in_model_cost_map(new_model_cost_map: dict) -> int: + """Adopt a freshly fetched cost map into this process's litellm state, return the model count""" + 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. @@ -3850,6 +3870,15 @@ class ProxyConfig: self._last_hashicorp_vault_config: dict[str, Any] | None = None self.worker_registry: list[WorkerRegistryEntry] = [] self.config_sync_subscriber: ConfigSyncSubscriber | None = None + from litellm.litellm_core_utils.get_model_cost_map import ( + get_model_cost_map_loaded_at, + ) + + self.model_cost_map_loaded_at: datetime = get_model_cost_map_loaded_at() or utc_now() + # Starts unapplied rather than adopting the published revision: this pod cannot tell + # whether an existing request predates the prices it just fetched, and re-serving one + # costs a single fetch where skipping one leaves it priced wrong indefinitely + self.model_cost_map_applied_revision: int = 0 def is_yaml(self, config_file_path: str) -> bool: if not os.path.isfile(config_file_path): @@ -6242,7 +6271,6 @@ class ProxyConfig: "router_settings", "litellm_settings", "environment_variables", - "model_cost_map_reload_config", "anthropic_beta_headers_reload_config", ], ) @@ -6340,9 +6368,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) @@ -6504,111 +6529,60 @@ class ProxyConfig: str(e), ) + async def check_periodic_reloads(self, prisma_client: PrismaClient): + """ + Run the admin-configured periodic model cost map reload. + + Scheduled on its own job so a schedule configured from the Admin UI fires 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) + 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: Final = 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() + 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", + ) + if not is_due: + return - config: Final = config_record.param_value - interval_hours: Final = config.get("interval_hours") - force_reload: Final = config.get("force_reload", False) + from litellm.litellm_core_utils.get_model_cost_map import ( + ModelCostMapReloadUnavailable, + refetch_model_cost_map, + ) - if interval_hours is None and force_reload is False: - return # No interval configured, skip reload - - current_time: Final = 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: Final = datetime.fromisoformat(last_model_cost_map_reload) - time_since_last_reload: Final = current_time - last_reload_time - hours_since_last_reload: Final = 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("Error parsing last reload time: %s", 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 ( - ModelCostMapReloadUnavailable, - refetch_model_cost_map, + reload_result = await refetch_model_cost_map(url=litellm.model_cost_map_url) + if isinstance(reload_result, ModelCostMapReloadUnavailable): + verbose_proxy_logger.warning( + "Model cost map reload failed (%s); keeping current pricing data. The revision stays " + "unapplied so this pod retries on its next poll", + reload_result.reason, ) + return - model_cost_map_url: Final = litellm.model_cost_map_url - reload_result: Final = await refetch_model_cost_map(url=model_cost_map_url) - if isinstance(reload_result, ModelCostMapReloadUnavailable): - verbose_proxy_logger.warning( - "Model cost map reload failed (%s); keeping current pricing data, will retry on the next config poll", - reload_result.reason, - ) - return - new_model_cost_map: Final = reload_result.model_cost_map - 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) + models_count = _swap_in_model_cost_map(reload_result.model_cost_map) + self.model_cost_map_loaded_at = current_time + await record_reload_run(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME, current_time) + # Adopted last, so neither a failed fetch nor a failed status write is recorded + # as served; either way the next poll retries instead of leaving the card + # reporting a run that never landed + self.model_cost_map_applied_revision = schedule.reload_revision - # 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 evict_config_param("model_cost_map_reload_config") - - verbose_proxy_logger.info( - "Model cost map reloaded successfully. Models count: %s", - len(new_model_cost_map) if new_model_cost_map else 0, - ) + verbose_proxy_logger.info("Model cost map reloaded successfully. Models count: %s", models_count) except Exception as e: verbose_proxy_logger.exception("Error in _check_and_reload_model_cost_map: %s", e) @@ -8254,15 +8228,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( @@ -15868,47 +15853,23 @@ async def reload_model_cost_map( refetch_model_cost_map, ) - model_cost_map_url: Final = litellm.model_cost_map_url - reload_result: Final = await refetch_model_cost_map(url=model_cost_map_url) + reload_result = await refetch_model_cost_map(url=litellm.model_cost_map_url) if isinstance(reload_result, ModelCostMapReloadUnavailable): raise HTTPException( status_code=502, detail=f"Failed to reload model cost map: {reload_result.reason}. Current pricing data was kept.", ) - new_model_cost_map: Final = reload_result.model_cost_map - 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 - global last_model_cost_map_reload - current_time: Final = datetime.utcnow() - last_model_cost_map_reload = current_time.isoformat() + models_count = _swap_in_model_cost_map(reload_result.model_cost_map) + current_time = utc_now() + proxy_config.model_cost_map_loaded_at = current_time - # Set force reload flag in database for other pods, preserving existing interval_hours - existing_config: Final = await ConfigRepository(prisma_client).table.find_unique( - where={"param_name": "model_cost_map_reload_config"} + # 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 ) - 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: Final = len(new_model_cost_map) if new_model_cost_map else 0 verbose_proxy_logger.info("Model cost map reloaded successfully in current pod. Models count: %s", models_count) return { @@ -15955,18 +15916,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("Model cost map reload scheduled for every %s hours", hours) @@ -15974,7 +15924,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("Failed to schedule model cost map reload: %s", e) @@ -16010,16 +15960,14 @@ async def cancel_model_cost_map_reload( if prisma_client is None: 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 clear_reload_interval(prisma_client, 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("Failed to cancel model cost map reload: %s", e) @@ -16048,66 +15996,13 @@ async def get_model_cost_map_reload_status( ) try: - global prisma_client, last_model_cost_map_reload - - verbose_proxy_logger.info("Checking model cost map reload status. Last reload: %s", 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: Final = 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: Final = config_record.param_value - interval_hours: Final = 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: Final = 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: Final = datetime.fromisoformat(last_model_cost_map_reload) - time_since_last_reload: Final = current_time - last_reload_time - hours_since_last_reload: Final = 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("Error parsing last reload time: %s", 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("Failed to get model cost map reload status: %s", e) raise HTTPException( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 0d7fa8692c8..17339541fd9 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_revision BigInt @default(0) } // View spend, model, api_key per request diff --git a/schema.prisma b/schema.prisma index 0d7fa8692c8..17339541fd9 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_revision BigInt @default(0) } // View spend, model, api_key per request diff --git a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py index 463bdc6161f..94798d77348 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py @@ -250,6 +250,27 @@ def test_azure_ai_claude_1m_context_entries(cost_map: dict): assert cost_map[model]["max_input_tokens"] == 200000, model +def test_get_model_cost_map_stamps_loaded_at(monkeypatch): + """The load time feeds each pod's reload-due decision; a load that does not stamp it + would make manual reload requests race the proxy's startup""" + from datetime import datetime, timezone + + from litellm.litellm_core_utils import get_model_cost_map as module + + monkeypatch.setattr(module._cost_map_source_info, "loaded_at", None) + monkeypatch.setattr( + module.GetModelCostMap, + "fetch_remote_model_cost_map", + staticmethod(lambda url, timeout=5: _load_root_cost_map()), + ) + + before = datetime.now(timezone.utc) + module.get_model_cost_map(url="https://example.invalid/cost_map.json") + loaded_at = module.get_model_cost_map_loaded_at() + + assert loaded_at is not None + assert before <= loaded_at <= datetime.now(timezone.utc) + # --------------------------------------------------------------------------- # refetch_model_cost_map: retry/backoff behavior for runtime reloads # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py index e357f25123b..3bb3f75b9b8 100644 --- a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py +++ b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py @@ -757,9 +757,13 @@ async def test_evict_config_param_does_not_publish() -> None: def _reload_config_prisma_client() -> MagicMock: config_record = MagicMock() config_record.param_value = {"interval_hours": 6, "force_reload": True} + config_record.reload_revision = 0 + config_record.last_run_at = None prisma_client = MagicMock() prisma_client.get_generic_data = AsyncMock(return_value=config_record) - prisma_client.db.litellm_config.upsert = AsyncMock(return_value=None) + prisma_client.db.litellm_config.find_unique = AsyncMock(return_value=config_record) + prisma_client.db.litellm_config.upsert = AsyncMock(return_value=config_record) + prisma_client.db.litellm_config.update_many = AsyncMock(return_value=1) return prisma_client @@ -790,7 +794,7 @@ async def test_model_cost_map_reload_does_not_publish_config_change() -> None: _invalidate_model_cost_lowercase_map() _set_redis_usage_cache(previous_cache) - prisma_client.db.litellm_config.upsert.assert_awaited_once() + prisma_client.db.litellm_config.update_many.assert_awaited_once() assert client.published == [] 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..cabb7452d1a --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_periodic_reload_schedule.py @@ -0,0 +1,359 @@ +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, + clear_reload_interval, + 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_revision=0, last_run_at=None): + return SimpleNamespace( + param_name="model_cost_map_reload_config", + param_value=param_value, + reload_revision=reload_revision, + last_run_at=last_run_at, + ) + + +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=_row(reload_revision=upserted_revision)) + prisma_client.db.litellm_config.update_many = AsyncMock(return_value=1) + return prisma_client + + +class _FakeConfigTable: + """In-memory stand-in for the prisma LiteLLM_Config actions, faithful on the parts the + revision depends on: upsert creates at the column default and applies ``{"increment": 1}``, + and delete drops the row along with its counter""" + + def __init__(self): + self._rows = {} + + async def find_unique(self, where): + return self._rows.get(where["param_name"]) + + async def upsert(self, where, data): + row = self._rows.get(where["param_name"]) + if row is None: + created = data["create"] + row = _row( + param_value=created.get("param_value"), + reload_revision=created.get("reload_revision", 0), + last_run_at=created.get("last_run_at"), + ) + self._rows[where["param_name"]] = row + return row + self._apply(row, data["update"]) + return row + + async def update_many(self, data, where): + row = self._rows.get(where["param_name"]) + if row is None: + return 0 + self._apply(row, data) + return 1 + + async def delete(self, where): + return self._rows.pop(where["param_name"], None) + + @staticmethod + def _apply(row, data): + if "param_value" in data and data["param_value"] is None: + raise ValueError("`data.param_value`: A value is required but not set") + for field, value in data.items(): + increment = value.get("increment") if isinstance(value, dict) else None + setattr(row, field, getattr(row, field) + increment if increment is not None else value) + + +def _fake_prisma(table): + prisma_client = MagicMock() + prisma_client.db.litellm_config = table + 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_revision=7, 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(last_run_at=LAST_RUN.replace(tzinfo=None))) + + 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( + "param_value", + [None, "not-a-dict", '{"interval_hours": "6"}', {}, {"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_reads_an_interval_still_encoded_as_json_text(): + """The interval is written with safe_dumps, so a raw row read can return it either + decoded or as a string; reading a string as no schedule would silently stop the + reloads an admin configured""" + assert parse_reload_schedule(_row(param_value='{"interval_hours": 6}')).interval_hours == 6 + + +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_revision=0, 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_applied_revision, pod_data_loaded_at, expected", + [ + (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), 0, LAST_RUN, 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), 0, LAST_RUN, True), + ], +) +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", + ) + is expected + ) + + +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_booting_pod_serves_a_request_it_cannot_prove_it_already_has(): + """A pod that just booted cannot tell whether an outstanding request predates the prices + it fetched at import, so it serves it. Adopting instead would strand it on stale prices + with no interval configured to rescue it""" + assert ( + pod_reload_is_due( + schedule=ReloadSchedule(reload_revision=12), + pod_applied_revision=0, + pod_data_loaded_at=NOW, + current_time=NOW, + description="test", + ) + is True + ) + + +def test_pod_reload_decision_ignores_persisted_last_run(): + """A pod with stale data must refresh 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_applied_revision=0, + pod_data_loaded_at=LAST_RUN, + current_time=NOW, + description="test", + ) + is True + ) + + +def test_schedule_that_never_ran_fires_immediately(): + """A fresh schedule must not wait a full interval for its first run, even on a pod + whose own data is boot-fresh""" + 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", + ) + 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_revision_on_interval_less_row(): + """A manual reload on a proxy with no schedule creates a row with only the columns + 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_revision=3, 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_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) + + 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_cancelling_a_schedule_never_reissues_a_revision(): + """Cancelling must keep the row. The revision identifies a request rather than ordering + one, so a counter restarted by a delete reissues a number pods already applied and their + next manual reload is skipped everywhere but the pod that served it""" + prisma_client = _fake_prisma(_FakeConfigTable()) + param_name = "model_cost_map_reload_config" + await write_reload_interval(prisma_client, param_name, 6) + pod_applied_revision = await record_manual_reload(prisma_client, param_name, LAST_RUN) + + await clear_reload_interval(prisma_client, param_name) + republished = await record_manual_reload(prisma_client, param_name, NOW) + + assert (pod_applied_revision, republished) == (1, 2) + schedule = await read_reload_schedule(prisma_client, param_name) + assert schedule is not None + assert ( + pod_reload_is_due( + schedule=schedule, + pod_applied_revision=pod_applied_revision, + pod_data_loaded_at=NOW, + current_time=NOW, + description="test", + ) + is True + ) + + +@pytest.mark.asyncio +async def test_cancelling_a_schedule_stops_it_while_keeping_the_recorded_run(): + """Dropping only the admin-owned param_value: the card must report no schedule without + losing the last run it already showed""" + prisma_client = _fake_prisma(_FakeConfigTable()) + param_name = "model_cost_map_reload_config" + await write_reload_interval(prisma_client, param_name, 6) + await record_reload_run(prisma_client, param_name, LAST_RUN) + + await clear_reload_interval(prisma_client, param_name) + + status = reload_schedule_status(await read_reload_schedule(prisma_client, param_name)) + assert status == { + "scheduled": False, + "interval_hours": None, + "last_run": "2024-01-01T06:00:00+00:00", + "next_run": None, + } + + +@pytest.mark.asyncio +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) + + 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_revision": {"increment": 1}} + assert data["create"] == { + "param_name": "model_cost_map_reload_config", + "last_run_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 df1d096b3e2..b75ee1caccf 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 @@ -10,9 +10,9 @@ Routes covered: from __future__ import annotations +import json from unittest.mock import AsyncMock, MagicMock -import pytest from .conftest import VOLATILE_KEYS, normalize @@ -35,6 +35,7 @@ def _attach_litellm_config(mock_prisma): table.upsert = AsyncMock() table.create = AsyncMock() table.update = AsyncMock() + table.update_many = AsyncMock(return_value=1) table.delete = AsyncMock() table.delete_many = AsyncMock() mock_prisma.db.litellm_config = table @@ -84,6 +85,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_revision"} + assert update_payload["reload_revision"] == {"increment": 1} def test_reload_model_cost_map_fetch_failure_502_keeps_map( @@ -179,6 +183,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( @@ -213,18 +220,15 @@ def test_schedule_model_cost_map_reload_not_admin_forbidden(client, auth_as): def test_cancel_model_cost_map_reload_happy(client, auth_as, monkeypatch, mock_prisma): - """Admin cancellation deletes config row and returns success body.""" + """Admin cancellation clears the interval and returns success body. The row itself stays: + it also holds the reload revision, and a counter restarted by a delete reissues a number + pods already applied, silently skipping their next manual reload.""" from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles table = _attach_litellm_config(mock_prisma) monkeypatch.setattr(ps, "prisma_client", mock_prisma) - async def _fake_invalidate(name): - return None - - monkeypatch.setattr(ps, "invalidate_config_param", _fake_invalidate) - with auth_as(LitellmUserRoles.PROXY_ADMIN): response = client.delete("/schedule/model_cost_map_reload") assert response.status_code == 200 @@ -234,7 +238,8 @@ def test_cancel_model_cost_map_reload_happy(client, auth_as, monkeypatch, mock_p "status": "success", "timestamp": "", } - assert table.delete.await_count == 1 + assert json.loads(table.update_many.await_args.kwargs["data"]["param_value"]) == {"interval_hours": None} + assert table.delete.await_count == 0 def test_cancel_model_cost_map_reload_not_admin_forbidden(client, auth_as): @@ -290,10 +295,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_revision = 0 + 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") @@ -306,19 +312,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_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) + + 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_revision = 3 + 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 67e76585485..c7aa376a2f7 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -27,6 +27,7 @@ import litellm import litellm.proxy.proxy_server as proxy_server_module from litellm.caching.caching import RedisCache from litellm.caching.redis_cluster_cache import RedisClusterCache +from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded from litellm.caching.dual_cache import DualCache from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -755,6 +756,48 @@ 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): """ @@ -3596,6 +3639,20 @@ async def test_chat_completion_result_no_nested_none_values(): # ============================================================================ +def _reload_schedule_row( + param_value: dict, + *, + 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_revision=reload_revision, + last_run_at=last_run_at, + ) + + class TestPriceDataReloadAPI: """Test cases for price data reload API endpoints""" @@ -3636,10 +3693,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=_reload_schedule_row({}, reload_revision=1) ) - mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) response = client_with_auth.post("/reload/model_cost_map") @@ -3694,7 +3750,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") @@ -3705,10 +3763,10 @@ 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) + 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") @@ -3719,6 +3777,15 @@ 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 @@ -3744,7 +3811,7 @@ class TestPriceDataReloadAPI: def test_cancel_model_cost_map_reload_admin_access(self, client_with_auth): """Test that admin users can cancel periodic reload""" with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: - # Mock database delete + mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=1) mock_prisma.db.litellm_config.delete = AsyncMock(return_value=None) response = client_with_auth.delete("/schedule/model_cost_map_reload") @@ -3754,6 +3821,10 @@ class TestPriceDataReloadAPI: assert data["status"] == "success" assert "message" in data assert "timestamp" in data + assert json.loads(mock_prisma.db.litellm_config.update_many.await_args.kwargs["data"]["param_value"]) == { + "interval_hours": None + } + mock_prisma.db.litellm_config.delete.assert_not_called() def test_cancel_model_cost_map_reload_non_admin_access(self, client_with_auth): """Test that non-admin users cannot cancel periodic reload""" @@ -3770,35 +3841,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 + (whose in-memory clock only knows its own boot) still reports the real last/next run + """ + proxy_server_module.proxy_config.model_cost_map_loaded_at = datetime(2030, 6, 1, tzinfo=timezone.utc) + 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""" @@ -3829,36 +3893,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_revision=3, + ) ) 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""" @@ -3900,10 +3972,9 @@ 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=_reload_schedule_row({}, reload_revision=1) ) - mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) # Test reload endpoint response = client_with_auth.post("/reload/model_cost_map") @@ -3916,172 +3987,352 @@ class TestPriceDataReloadIntegration: litellm.model_cost = original_model_cost _invalidate_model_cost_lowercase_map() - def test_distributed_reload_check_function(self): - """Test the _check_and_reload_model_cost_map function""" + def test_pod_data_clock_seeded_from_actual_cost_map_load(self): + """Regression: seeding from ProxyConfig construction time instead of the real + import-time fetch let a manual request stamped during startup be skipped""" + from datetime import datetime, timezone + 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 + fetch_time = datetime(2024, 1, 1, 6, 0, tzinfo=timezone.utc) with patch( - "litellm.proxy.proxy_server.last_model_cost_map_reload", - "2024-01-01T06:00:00", + "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map_loaded_at", + return_value=fetch_time, ): - 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 + assert ProxyConfig().model_cost_map_loaded_at == fetch_time - # 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) - - from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded - - original_model_cost = litellm.model_cost.copy() - try: - with patch( - "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", - new=AsyncMock( - return_value=ModelCostMapReloaded( - model_cost_map={"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. + def test_distributed_reload_check_function(self): + """ + 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=_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) - # 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) + boot_loaded_at = proxy_config.model_cost_map_loaded_at + 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.model_cost_map_loaded_at == boot_loaded_at + + frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc) + mock_prisma.db.litellm_config.find_unique = AsyncMock( + return_value=_reload_schedule_row( + {"interval_hours": 6}, + 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 from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded original_model_cost = litellm.model_cost.copy() try: - with patch( - "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", - new=AsyncMock( - return_value=ModelCostMapReloaded( - model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}} - ) - ), + with ( + patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) as mock_get_map, + patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), + ): + mock_get_map.return_value = ModelCostMapReloaded(model_cost_map={"gpt-3.5-turbo": {"input_cost_per_token": 0.001}}) + + asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) + + assert litellm.model_cost["gpt-3.5-turbo"] == {"input_cost_per_token": 0.001} + assert proxy_config.model_cost_map_loaded_at == 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() + 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 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 + + 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_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: + with ( + patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) 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)) - # 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" - ) + mock_get_map.assert_not_called() + mock_prisma.db.litellm_config.update_many.assert_not_called() + assert proxy_config.model_cost_map_loaded_at == pod_data_loaded_at + finally: + litellm.model_cost = original_model_cost + _invalidate_model_cost_lowercase_map() + + def test_periodic_reload_uses_pod_local_data_age(self): + """ + Each pod decides from the age of its own data, so a pod holding a stale copy + refreshes even when the shared row was just stamped by another pod, and stays + put while its copy is 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) + proxy_config.model_cost_map_loaded_at = datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc) + + original_model_cost = litellm.model_cost.copy() + try: + with ( + patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) as mock_get_map, + patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), + ): + mock_get_map.return_value = ModelCostMapReloaded(model_cost_map={"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.model_cost_map_loaded_at == 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.model_cost_map_loaded_at = 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.model_cost_map_loaded_at == frozen_now - timedelta(hours=1) + finally: + 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.refetch_model_cost_map", new_callable=AsyncMock) as mock_get_map, + patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), + ): + mock_get_map.return_value = ModelCostMapReloaded(model_cost_map={"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() + + @pytest.mark.parametrize( + "published_revision, expect_reload", + [(4, True), (0, False)], + ) + def test_booting_pod_serves_an_outstanding_request_once(self, published_revision, expect_reload): + """ + Regression: a manual reload published while this pod was starting must still be + served. The pod cannot prove its import-time fetch already covers that request, so + it applies it on the first poll and adopts the revision, leaving later polls quiet. + A row nobody has ever reloaded (revision 0) costs the pod nothing + """ + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc) + proxy_config.model_cost_map_loaded_at = frozen_now + 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=published_revision) + ) + + original_model_cost = litellm.model_cost.copy() + try: + with ( + patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) as mock_get_map, + patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), + ): + mock_get_map.return_value = ModelCostMapReloaded(model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}}) + + asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) + asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) + + assert mock_get_map.call_count == (1 if expect_reload else 0) + assert proxy_config.model_cost_map_applied_revision == published_revision + 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 + admin-configured interval) nor a create branch (which would resurrect a schedule + a concurrent cancel just deleted) + """ + 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": 24})) + 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) + + original_model_cost = litellm.model_cost.copy() + try: + with ( + patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) as mock_get_map, + patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), + ): + mock_get_map.return_value = ModelCostMapReloaded(model_cost_map={"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_distributed_reload_leaves_request_unserved_when_status_write_fails(self): + """ + A run that never reached the row must not be recorded as served. Adopting the + revision here would leave the card reporting the previous run until someone clicks + again, because a manual request is published once and never republished + """ + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc) + proxy_config.model_cost_map_loaded_at = frozen_now - timedelta(hours=9) + mock_prisma = MagicMock() + mock_prisma.db.litellm_config.find_unique = AsyncMock( + return_value=_reload_schedule_row({"interval_hours": 6}, reload_revision=7) + ) + mock_prisma.db.litellm_config.update_many = AsyncMock(side_effect=Exception("connection reset")) + + original_model_cost = litellm.model_cost.copy() + try: + with ( + patch( + "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", + new_callable=AsyncMock, + ) as mock_get_map, + patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), + ): + mock_get_map.return_value = ModelCostMapReloaded(model_cost_map={"gpt-4": {"input_cost_per_token": 0.1}}) + + asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) + + assert proxy_config.model_cost_map_applied_revision == 0 finally: litellm.model_cost = original_model_cost _invalidate_model_cost_lowercase_map() def test_distributed_reload_keeps_current_map_when_fetch_fails(self): - """Fetch failure during a periodic/forced reload must not downgrade the pod. + """Fetch failure during a periodic reload must not downgrade the pod or count the + request as served. - Regression: a 429/network failure used to silently replace litellm.model_cost - with the stale packaged backup, stamp last_run, and clear force_reload. + Regression: a 429/network failure used to silently replace litellm.model_cost with + the stale packaged backup and stamp last_run. Adopting the revision here would be + the same bug one level up: a manual request is published once and never republished, + so a pod that records it as applied without the data stays mispriced until someone + clicks again """ from litellm.litellm_core_utils.get_model_cost_map import ( ModelCostMapReloadUnavailable, ) - from litellm.proxy import proxy_server as ps from litellm.proxy.proxy_server import ProxyConfig proxy_config = ProxyConfig() + frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc) + pod_data_loaded_at = frozen_now - timedelta(hours=9) + proxy_config.model_cost_map_loaded_at = pod_data_loaded_at mock_prisma = MagicMock() - - mock_config = MagicMock() - 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.find_unique = AsyncMock( + return_value=_reload_schedule_row({"interval_hours": 6}, reload_revision=7) + ) + mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None) mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) original_model_cost = litellm.model_cost - with patch( - "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", - new=AsyncMock( - return_value=ModelCostMapReloadUnavailable(reason="HTTP 429 from upstream") + with ( + patch( + "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", + new=AsyncMock(return_value=ModelCostMapReloadUnavailable(reason="HTTP 429 from upstream")), ), + patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): - with patch("litellm.proxy.proxy_server.last_model_cost_map_reload", None): - asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) - assert ps.last_model_cost_map_reload is None, ( - "a failed reload must not stamp the pod's last reload time, " - "otherwise the retry waits a full interval" - ) + asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) assert litellm.model_cost is original_model_cost, ( "a failed reload must keep the currently loaded cost map, " "not swap in the packaged backup" ) + assert proxy_config.model_cost_map_loaded_at == pod_data_loaded_at, ( + "a failed reload must not stamp the pod's data age, otherwise the retry waits a full interval" + ) + assert proxy_config.model_cost_map_applied_revision == 0, ( + "a failed reload must leave the revision unapplied so the next poll retries it" + ) + mock_prisma.db.litellm_config.update_many.assert_not_called() mock_prisma.db.litellm_config.upsert.assert_not_called() def test_manual_reload_preserves_interval_hours(self): - """Test that manual reload via /reload/model_cost_map 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. + """ + 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 @@ -4095,44 +4346,40 @@ 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) from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded original_model_cost = litellm.model_cost.copy() try: - with patch( - "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", - new=AsyncMock( - return_value=ModelCostMapReloaded( - model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}} - ) - ), + with ( + patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) 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), ): - 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) + mock_get_map.return_value = ModelCostMapReloaded(model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}}) + 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 + 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_revision": {"increment": 1}, + } + assert call_args[1]["data"]["create"] == { + "param_name": "model_cost_map_reload_config", + "last_run_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() @@ -4144,7 +4391,9 @@ class TestPriceDataReloadIntegration: identical to the model cost map bug. """ from litellm.proxy.proxy_server import ProxyConfig + from litellm.proxy.utils import litellm_config_cache + litellm_config_cache.flush_cache() proxy_config = ProxyConfig() mock_prisma = MagicMock() @@ -4154,7 +4403,7 @@ 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" @@ -4204,10 +4453,10 @@ class TestPriceDataReloadIntegration: # 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.find_unique = AsyncMock(return_value=mock_existing) + mock_prisma.db.litellm_config.upsert = AsyncMock( + return_value=_reload_schedule_row({}, reload_revision=1) ) - mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) response = client.post("/reload/anthropic_beta_headers") assert response.status_code == 200 @@ -4251,64 +4500,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():