mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
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>
This commit is contained in:
milan
parent
daf22ec871
commit
00b12ae067
5 changed files with 618 additions and 388 deletions
151
litellm/proxy/common_utils/periodic_reload_schedule.py
Normal file
151
litellm/proxy/common_utils/periodic_reload_schedule.py
Normal file
|
|
@ -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),
|
||||
)
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -60,6 +60,7 @@ const PriceDataReload: React.FC<PriceDataReloadProps> = ({
|
|||
const [showScheduleModal, setShowScheduleModal] = useState(false);
|
||||
const [hours, setHours] = useState<number>(6);
|
||||
const [reloadStatus, setReloadStatus] = useState<ReloadStatus | null>(null);
|
||||
const [statusError, setStatusError] = useState(false);
|
||||
const [loadingStatus, setLoadingStatus] = useState(false);
|
||||
const [sourceInfo, setSourceInfo] = useState<CostMapSourceInfo | null>(null);
|
||||
const [loadingSource, setLoadingSource] = useState(false);
|
||||
|
|
@ -85,15 +86,10 @@ const PriceDataReload: React.FC<PriceDataReloadProps> = ({
|
|||
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<PriceDataReloadProps> = ({
|
|||
)}
|
||||
|
||||
{/* Reload Schedule Status Card */}
|
||||
{statusError && (
|
||||
<Card
|
||||
size="small"
|
||||
style={{
|
||||
backgroundColor: "#fff7e6",
|
||||
border: "1px solid #ffd591",
|
||||
borderRadius: 8,
|
||||
marginBottom: reloadStatus ? 12 : 0,
|
||||
}}
|
||||
>
|
||||
<Space size="small">
|
||||
<WarningOutlined style={{ color: "#fa8c16", fontSize: 12 }} />
|
||||
<Text style={{ fontSize: "12px", color: "#614700" }}>
|
||||
Could not load the reload schedule. {reloadStatus ? "Showing the last known state; refresh" : "Refresh"}{" "}
|
||||
the page or sign in again.
|
||||
</Text>
|
||||
</Space>
|
||||
</Card>
|
||||
)}
|
||||
|
||||
{reloadStatus && (
|
||||
<Card
|
||||
size="small"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue