From 00b12ae0674b1122b092d3bf5cd64e4c1c00aade Mon Sep 17 00:00:00 2001 From: milan Date: Tue, 28 Jul 2026 14:52:05 +0000 Subject: [PATCH] fix(proxy): persist periodic price data reload schedule across restarts Store last_run in the LiteLLM_Config row and schedule the reload check outside the store_model_in_db gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../common_utils/periodic_reload_schedule.py | 151 ++++++ litellm/proxy/proxy_server.py | 507 ++++++------------ .../test_periodic_reload_schedule.py | 132 +++++ tests/test_litellm/proxy/test_proxy_server.py | 186 ++++++- .../src/components/price_data_reload.tsx | 30 +- 5 files changed, 618 insertions(+), 388 deletions(-) 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/common_utils/periodic_reload_schedule.py b/litellm/proxy/common_utils/periodic_reload_schedule.py new file mode 100644 index 00000000000..14d91f47325 --- /dev/null +++ b/litellm/proxy/common_utils/periodic_reload_schedule.py @@ -0,0 +1,151 @@ +""" +Persistence for the admin-configured periodic reload schedules (model cost map, +Anthropic beta headers) stored in the ``LiteLLM_Config`` table. + +``last_run`` is kept in the row rather than in process memory so the Admin UI still +reports the schedule and its last execution after a restart and across pods. Each pod +tracks its own last reload separately to decide when to refresh its in-memory copy of +the data. +""" + +from dataclasses import dataclass, replace +from datetime import datetime, timedelta, timezone +from typing import TypedDict + +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 + +MODEL_COST_MAP_RELOAD_PARAM_NAME = "model_cost_map_reload_config" +ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME = "anthropic_beta_headers_reload_config" + + +@dataclass(frozen=True, slots=True) +class ReloadSchedule: + interval_hours: int | None = None + force_reload: bool = False + last_run: datetime | None = None + + +class ReloadScheduleStatus(TypedDict): + scheduled: bool + interval_hours: int | None + last_run: str | None + next_run: str | None + + +def utc_now() -> datetime: + return datetime.now(timezone.utc) + + +def _parse_last_run(raw: object) -> datetime | None: + if not isinstance(raw, str): + return None + try: + parsed = datetime.fromisoformat(raw) + except ValueError: + verbose_proxy_logger.warning("Ignoring unparseable reload last_run: %s", raw) + return None + return parsed if parsed.tzinfo is not None else parsed.replace(tzinfo=timezone.utc) + + +def parse_reload_schedule(param_value: object) -> ReloadSchedule: + if not isinstance(param_value, dict): + return ReloadSchedule() + interval_hours = param_value.get("interval_hours") + return ReloadSchedule( + interval_hours=interval_hours if isinstance(interval_hours, int) else None, + force_reload=param_value.get("force_reload") is True, + last_run=_parse_last_run(param_value.get("last_run")), + ) + + +def serialize_reload_schedule(schedule: ReloadSchedule) -> str: + return safe_dumps( + { + "interval_hours": schedule.interval_hours, + "force_reload": schedule.force_reload, + "last_run": schedule.last_run.isoformat() if schedule.last_run is not None else None, + } + ) + + +def next_run_at(schedule: ReloadSchedule) -> datetime | None: + if schedule.interval_hours is None or schedule.last_run is None: + return None + return schedule.last_run + timedelta(hours=schedule.interval_hours) + + +def reload_schedule_status(schedule: ReloadSchedule | None) -> ReloadScheduleStatus: + last_run = schedule.last_run if schedule is not None else None + next_run = next_run_at(schedule) if schedule is not None else None + return { + "scheduled": schedule is not None and schedule.interval_hours is not None, + "interval_hours": schedule.interval_hours if schedule is not None else None, + "last_run": last_run.isoformat() if last_run 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 that every pod refreshes its in-memory data on the configured interval + """ + if schedule.force_reload: + verbose_proxy_logger.info("%s reload triggered by force reload flag", 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 ConfigRepository(prisma_client).table.find_unique(where={"param_name": param_name}) + if row is None or row.param_value is None: + return None + return parse_reload_schedule(row.param_value) + + +async def write_reload_schedule(prisma_client: PrismaClient, param_name: str, schedule: ReloadSchedule) -> None: + param_value = serialize_reload_schedule(schedule) + await ConfigRepository(prisma_client).table.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: + """ + Persist a completed reload: stamp ``last_run`` and clear ``force_reload`` while keeping + the configured interval + """ + existing = await read_reload_schedule(prisma_client, param_name) or ReloadSchedule() + await write_reload_schedule( + prisma_client, + param_name, + replace(existing, force_reload=False, last_run=ran_at), + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4486cd7de59..ddede2fe08c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -326,6 +326,18 @@ 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, + ReloadSchedule, + parse_reload_schedule, + pod_reload_is_due, + read_reload_schedule, + record_reload_run, + reload_schedule_status, + utc_now, + write_reload_schedule, +) 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 +2066,10 @@ 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 +last_model_cost_map_reload: datetime | None = None +last_anthropic_beta_headers_reload: datetime | None = None ### DB WRITER ### @@ -6175,8 +6185,6 @@ class ProxyConfig: "router_settings", "litellm_settings", "environment_variables", - "model_cost_map_reload_config", - "anthropic_beta_headers_reload_config", ], ) @@ -6239,12 +6247,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 +6405,65 @@ 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 the 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") + config_record = await get_config_param(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME) if config_record is None or config_record.param_value is None: return # No configuration found, skip reload - config = config_record.param_value - interval_hours = config.get("interval_hours") - force_reload = config.get("force_reload", False) + schedule = parse_reload_schedule(config_record.param_value) - if interval_hours is None and force_reload is False: + if schedule.interval_hours is None and schedule.force_reload is False: return # No interval configured, skip reload - current_time = datetime.utcnow() + current_time = utc_now() - # Check if we need to reload based on interval or force reload - should_reload = False + global last_model_cost_map_reload + if not pod_reload_is_due( + schedule=schedule, + pod_last_reload=last_model_cost_map_reload, + current_time=current_time, + description="Model cost map", + ): + return - 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 + from litellm.litellm_core_utils.get_model_cost_map import ( + get_model_cost_map, + ) - 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") + 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) - if should_reload: - # Perform the reload - from litellm.litellm_core_utils.get_model_cost_map import ( - get_model_cost_map, - ) + last_model_cost_map_reload = current_time + await record_reload_run(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME, current_time) - 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: {len(new_model_cost_map) if new_model_cost_map else 0}" + ) except Exception as e: verbose_proxy_logger.exception(f"Error in _check_and_reload_model_cost_map: {str(e)}") @@ -6506,96 +6471,43 @@ 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") + config_record = await get_config_param(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME) if config_record is None or config_record.param_value is None: return # No configuration found, skip reload - config = config_record.param_value - interval_hours = config.get("interval_hours") - force_reload = config.get("force_reload", False) + schedule = parse_reload_schedule(config_record.param_value) - if interval_hours is None and force_reload is False: + if schedule.interval_hours is None and schedule.force_reload is False: return # No interval configured, skip reload - current_time = datetime.utcnow() + current_time = utc_now() - # Check if we need to reload based on interval or force reload - should_reload = False + global last_anthropic_beta_headers_reload + if not pod_reload_is_due( + schedule=schedule, + pod_last_reload=last_anthropic_beta_headers_reload, + current_time=current_time, + description="Anthropic beta headers", + ): + return - 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 + from litellm.anthropic_beta_headers_manager import ( + reload_beta_headers_config, + ) - 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" - ) + new_config = reload_beta_headers_config() - if should_reload: - # Perform the reload - from litellm.anthropic_beta_headers_manager import ( - reload_beta_headers_config, - ) + last_anthropic_beta_headers_reload = current_time + await record_reload_run(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME, current_time) - 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 +7971,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( @@ -15674,28 +15597,20 @@ async def reload_model_cost_map( # 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() + current_time = utc_now() + last_model_cost_map_reload = 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_schedule = await read_reload_schedule(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME) + await write_reload_schedule( + prisma_client, + MODEL_COST_MAP_RELOAD_PARAM_NAME, + ReloadSchedule( + interval_hours=existing_schedule.interval_hours if existing_schedule is not None else None, + force_reload=True, + last_run=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 = 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}") @@ -15742,18 +15657,16 @@ 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})}, - }, + existing_schedule = await read_reload_schedule(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME) + await write_reload_schedule( + prisma_client, + MODEL_COST_MAP_RELOAD_PARAM_NAME, + ReloadSchedule( + interval_hours=hours, + force_reload=False, + last_run=existing_schedule.last_run if existing_schedule is not None else None, + ), ) - await invalidate_config_param("model_cost_map_reload_config") verbose_proxy_logger.info(f"Model cost map reload scheduled for every {hours} hours") @@ -15761,7 +15674,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,8 +15711,8 @@ 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") @@ -15835,66 +15748,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( @@ -15988,28 +15848,20 @@ async def reload_anthropic_beta_headers( # 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() + 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_schedule = await read_reload_schedule(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME) + await write_reload_schedule( + prisma_client, + ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME, + ReloadSchedule( + interval_hours=existing_schedule.interval_hours if existing_schedule is not None else None, + force_reload=True, + last_run=current_time, + ), ) - 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") provider_count = sum(1 for k in new_config.keys() if k not in ["provider_aliases", "description"]) verbose_proxy_logger.info( @@ -16058,18 +15910,16 @@ 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})}, - }, + existing_schedule = await read_reload_schedule(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME) + await write_reload_schedule( + prisma_client, + ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME, + ReloadSchedule( + interval_hours=hours, + force_reload=False, + last_run=existing_schedule.last_run if existing_schedule is not None else None, + ), ) - await invalidate_config_param("anthropic_beta_headers_reload_config") verbose_proxy_logger.info(f"Anthropic beta headers reload scheduled for every {hours} hours") @@ -16077,7 +15927,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,8 +15964,10 @@ 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") @@ -16154,68 +16006,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/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..448a4e68e5f --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_periodic_reload_schedule.py @@ -0,0 +1,132 @@ +import json +from datetime import datetime, timezone +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_reload_run, + reload_schedule_status, + serialize_reload_schedule, +) + +LAST_RUN = datetime(2024, 1, 1, 6, 0, 0, tzinfo=timezone.utc) + + +def test_parse_round_trips_through_serialization(): + schedule = ReloadSchedule(interval_hours=6, force_reload=True, last_run=LAST_RUN) + + assert parse_reload_schedule(json.loads(serialize_reload_schedule(schedule))) == schedule + + +def test_parse_treats_naive_last_run_as_utc(): + """Rows written before last_run was persisted as tz-aware must still parse""" + schedule = parse_reload_schedule({"interval_hours": 6, "last_run": "2024-01-01T06:00:00"}) + + assert schedule.last_run == LAST_RUN + + +@pytest.mark.parametrize( + "param_value", + [None, "not-a-dict", {}, {"interval_hours": "6"}, {"last_run": "garbage"}], +) +def test_parse_tolerates_unusable_values(param_value): + schedule = parse_reload_schedule(param_value) + + assert schedule.interval_hours is None + assert schedule.last_run is None + + +def test_status_reports_persisted_last_run_and_next_run(): + status = reload_schedule_status(ReloadSchedule(interval_hours=6, last_run=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=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=LAST_RUN)) is None + + +@pytest.mark.parametrize( + "schedule, pod_last_reload, expected", + [ + (ReloadSchedule(force_reload=True), LAST_RUN, 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), datetime(2024, 1, 1, 6, 0, tzinfo=timezone.utc), 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=datetime(2024, 1, 1, 12, 0, tzinfo=timezone.utc), + description="test", + ) + is expected + ) + + +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 within the interval + """ + assert ( + pod_reload_is_due( + schedule=ReloadSchedule(interval_hours=6, last_run=datetime(2024, 1, 1, 11, 59, tzinfo=timezone.utc)), + pod_last_reload=None, + current_time=datetime(2024, 1, 1, 12, 0, tzinfo=timezone.utc), + description="test", + ) + is True + ) + + +@pytest.mark.asyncio +async def test_record_reload_run_stamps_last_run_and_keeps_interval(): + existing = MagicMock() + existing.param_value = {"interval_hours": 6, "force_reload": True} + prisma_client = MagicMock() + prisma_client.db.litellm_config.find_unique = AsyncMock(return_value=existing) + prisma_client.db.litellm_config.upsert = AsyncMock(return_value=None) + + await record_reload_run(prisma_client, "model_cost_map_reload_config", LAST_RUN) + + written = json.loads(prisma_client.db.litellm_config.upsert.call_args[1]["data"]["update"]["param_value"]) + assert written == { + "interval_hours": 6, + "force_reload": False, + "last_run": "2024-01-01T06:00:00+00:00", + } + + +@pytest.mark.asyncio +async def test_read_reload_schedule_returns_none_for_missing_row(): + prisma_client = MagicMock() + prisma_client.db.litellm_config.find_unique = AsyncMock(return_value=None) + + assert await read_reload_schedule(prisma_client, "model_cost_map_reload_config") is None diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 536d24d4b4e..580afef43a2 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -755,6 +755,46 @@ 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): + """ + The periodic reload job (model cost map / anthropic beta headers) must be scheduled + even when store_model_in_db is False. + + Regression test: the reload checks only ran from the store_model_in_db-gated + add_deployment job, so a schedule configured in the Admin UI silently never fired + """ + monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False) + monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) + 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() + + with ( + patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", False), + ): + 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, + ) + + from litellm.proxy.proxy_server import scheduler + + assert scheduler is not None + assert scheduler.get_job("periodic_reload_job") is not None + assert scheduler.get_job("add_deployment_job") is None + + @pytest.mark.asyncio async def test_initialize_scheduled_jobs_uses_configured_config_reload_interval(monkeypatch): """ @@ -3701,6 +3741,7 @@ class TestPriceDataReloadAPI: """Test that admin users can schedule periodic reload""" with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: # Mock database upsert + 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("/schedule/model_cost_map_reload?hours=6") @@ -3763,35 +3804,27 @@ 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""" + """Test that admin users can get reload status, sourced from the persisted last_run""" 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_config.param_value = { + "interval_hours": 6, + "force_reload": False, + "last_run": "2024-01-01T06:00:00+00:00", + } mock_prisma.db.litellm_config.find_unique = AsyncMock( return_value=mock_config ) - # 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"] == 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""" @@ -3935,15 +3968,15 @@ class TestPriceDataReloadIntegration: # Mock current time and last reload time with patch( "litellm.proxy.proxy_server.last_model_cost_map_reload", - "2024-01-01T06:00:00", + datetime(2024, 1, 1, 6, 0, 0, tzinfo=timezone.utc), ): - 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 - + with patch( + "litellm.proxy.proxy_server.utc_now", + return_value=datetime(2024, 1, 1, 7, 0, 0, tzinfo=timezone.utc), + ): # Should not reload (only 1 hour passed, need 6) asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) + mock_prisma.db.litellm_config.upsert.assert_not_called() # Test case 3: Config with force reload litellm_config_cache.flush_cache() @@ -3972,10 +4005,109 @@ class TestPriceDataReloadIntegration: param_value_dict = json.loads(param_value_json) assert param_value_dict["force_reload"] == False assert param_value_dict.get("interval_hours") == 6 + assert param_value_dict.get("last_run") is not None finally: litellm.model_cost = original_model_cost _invalidate_model_cost_lowercase_map() + def test_distributed_reload_persists_last_run(self): + """A completed reload must stamp last_run in the DB row. + + Regression test: last_run used to live only in the pod's memory, so the Admin UI + reported "Never" for a schedule that had been running for days + """ + 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() + + 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_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) + + reload_time = datetime(2024, 1, 1, 7, 0, 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", + return_value={"gpt-4": {"input_cost_per_token": 0.001}}, + ): + with patch( + "litellm.proxy.proxy_server.last_model_cost_map_reload", None + ): + with patch( + "litellm.proxy.proxy_server.utc_now", return_value=reload_time + ): + asyncio.run( + proxy_config._check_and_reload_model_cost_map(mock_prisma) + ) + + call_args = mock_prisma.db.litellm_config.upsert.call_args + param_value_dict = json.loads( + call_args[1]["data"]["update"]["param_value"] + ) + assert param_value_dict["last_run"] == reload_time.isoformat() + assert param_value_dict["interval_hours"] == 6 + assert param_value_dict["force_reload"] is False + finally: + litellm.model_cost = original_model_cost + _invalidate_model_cost_lowercase_map() + litellm_config_cache.flush_cache() + + def test_schedule_model_cost_map_reload_preserves_last_run(self): + """Re-scheduling must not wipe the recorded last_run""" + from litellm.proxy.proxy_server import schedule_model_cost_map_reload + + mock_prisma = MagicMock() + existing = MagicMock() + existing.param_value = { + "interval_hours": 6, + "force_reload": False, + "last_run": "2024-01-01T06:00:00+00:00", + } + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=existing) + mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) + + admin = MagicMock() + admin.user_role = LitellmUserRoles.PROXY_ADMIN + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + asyncio.run( + schedule_model_cost_map_reload(hours=12, user_api_key_dict=admin) + ) + + param_value_dict = json.loads( + mock_prisma.db.litellm_config.upsert.call_args[1]["data"]["update"][ + "param_value" + ] + ) + assert param_value_dict["interval_hours"] == 12 + assert param_value_dict["last_run"] == "2024-01-01T06:00:00+00:00" + + def test_check_periodic_reloads_runs_both_checks(self): + """check_periodic_reloads drives both admin-configurable reloads""" + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + mock_prisma = MagicMock() + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + mock_prisma.get_generic_data = AsyncMock(return_value=None) + + with patch.object( + proxy_config, "_check_and_reload_model_cost_map", new=AsyncMock() + ) as mock_cost_map_check: + with patch.object( + proxy_config, "_check_and_reload_anthropic_beta_headers", new=AsyncMock() + ) as mock_beta_headers_check: + asyncio.run(proxy_config.check_periodic_reloads(mock_prisma)) + + mock_cost_map_check.assert_awaited_once() + mock_beta_headers_check.assert_awaited_once() + def test_distributed_reload_preserves_interval_hours(self): """Test that _check_and_reload_model_cost_map preserves interval_hours after reload. diff --git a/ui/litellm-dashboard/src/components/price_data_reload.tsx b/ui/litellm-dashboard/src/components/price_data_reload.tsx index e0ec0fb794d..4139f6170c2 100644 --- a/ui/litellm-dashboard/src/components/price_data_reload.tsx +++ b/ui/litellm-dashboard/src/components/price_data_reload.tsx @@ -60,6 +60,7 @@ const PriceDataReload: React.FC = ({ const [showScheduleModal, setShowScheduleModal] = useState(false); const [hours, setHours] = useState(6); const [reloadStatus, setReloadStatus] = useState(null); + const [statusError, setStatusError] = useState(false); const [loadingStatus, setLoadingStatus] = useState(false); const [sourceInfo, setSourceInfo] = useState(null); const [loadingSource, setLoadingSource] = useState(false); @@ -85,15 +86,10 @@ const PriceDataReload: React.FC = ({ try { const status = await getModelCostMapReloadStatus(accessToken); setReloadStatus(status); + setStatusError(false); } catch (error) { console.error("Failed to fetch reload status:", error); - // Set a default status to prevent UI issues - setReloadStatus({ - scheduled: false, - interval_hours: null, - last_run: null, - next_run: null, - }); + setStatusError(true); } finally { setLoadingStatus(false); } @@ -420,6 +416,26 @@ const PriceDataReload: React.FC = ({ )} {/* Reload Schedule Status Card */} + {statusError && ( + + + + + Could not load the reload schedule. {reloadStatus ? "Showing the last known state; refresh" : "Refresh"}{" "} + the page or sign in again. + + + + )} + {reloadStatus && (