scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 267

fix(proxy): persist periodic reload schedule state so status survives restarts and fires without store_model_in_db

Browse source 
The model cost map and Anthropic beta headers reload schedules kept their
last-run time in a per-pod module global, so GET /schedule/*/status reported
last_run null after any restart and the Admin UI showed the reload as never
having run. The reload check also only ran from the add_deployment job, which
is registered only when store_model_in_db is true, so config-file deployments
stored a schedule that never fired.

Persist last_run_at and reload_requested_at as dedicated columns on
LiteLLM_Config, owned by the reload job and manual reload endpoints, while the
schedule endpoints own the param_value JSON (interval_hours); no writer can
clobber another's fields. Serve status entirely from the row. Register the
check as its own periodic_reload_job outside the store_model_in_db gate.
Replace the force_reload boolean with a reload_requested_at timestamp each pod
compares against its own in-memory last reload, so a manual reload reaches
every pod exactly once instead of being cleared by the first poller. Run the
blocking fetches via asyncio.to_thread, and stamp last_run_at with update_many
so a schedule cancelled mid-poll is not resurrected.
This commit is contained in:
ryan-crabbe-berri 2026-07-29 17:37:27 -07:00
parent b930e2fc2b
commit bfa47535b5
11 changed files with 1002 additions and 709 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm-proxy-extras/litellm_proxy_extras/migrations/20260729000000_add_reload_tracking_to_litellm_config/migration.sql Normal file
View file

@ -0,0 +1,7 @@
-- AlterTable
ALTER TABLE "LiteLLM_Config" ADD COLUMN IF NOT EXISTS "last_run_at" TIMESTAMP(3),
ADD COLUMN IF NOT EXISTS "reload_requested_at" TIMESTAMP(3);
-- Carry over a manual reload that was still pending under the older boolean flag, so
-- pods that have not picked it up yet still fan out after this deploy.
UPDATE "LiteLLM_Config" SET "reload_requested_at" = CURRENT_TIMESTAMP WHERE "param_value"->'force_reload' = 'true'::jsonb;

2
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -601,6 +601,8 @@ model LiteLLM_TagTable {
model LiteLLM_Config {
param_name String @id
param_value Json?
last_run_at DateTime?
reload_requested_at DateTime?
}
// View spend, model, api_key per request

204
litellm/proxy/common_utils/periodic_reload_schedule.py Normal file
View file

@ -0,0 +1,204 @@
"""
Persistence for the admin-configured periodic reload schedules (model cost map,
Anthropic beta headers) stored in ``LiteLLM_Config``.
Field ownership is split by writer so concurrent writers never overwrite each other:
the schedule endpoints own the ``param_value`` JSON (``interval_hours``), while the
reload job and the manual reload endpoints own the dedicated ``last_run_at`` /
``reload_requested_at`` columns. ``last_run_at`` lives in the row rather than process
memory so the Admin UI still reports the last execution after a restart and across pods.
``reload_requested_at`` fans a manual reload out to every pod without any pod having to
write the flag back down, so no pod can starve another of the request. Each pod tracks
its own last reload in memory and compares it against those timestamps to decide when to
refresh its in-memory copy of the data.
"""
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import (
TYPE_CHECKING,
Protocol,
TypedDict,
cast, # noqa: TID251 # prisma table access is untyped (PrismaWrapper.__getattr__)
)
from pydantic import BaseModel, ConfigDict, ValidationError
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy.utils import PrismaClient, invalidate_config_param
from litellm.repositories.config_repository import ConfigRepository
if TYPE_CHECKING:
from prisma.models import LiteLLM_Config
MODEL_COST_MAP_RELOAD_PARAM_NAME = "model_cost_map_reload_config"
ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME = "anthropic_beta_headers_reload_config"
class _ConfigRowWrite(TypedDict, total=False):
param_name: str
param_value: str
last_run_at: datetime
reload_requested_at: datetime
class _ConfigUpsertData(TypedDict):
create: _ConfigRowWrite
update: _ConfigRowWrite
class _ConfigTable(Protocol):
async def find_unique(self, where: Mapping[str, str]) -> "LiteLLM_Config | None": ...
async def upsert(self, where: Mapping[str, str], data: _ConfigUpsertData) -> "LiteLLM_Config": ...
async def update_many(self, data: _ConfigRowWrite, where: Mapping[str, str]) -> int: ...
def _config_table(prisma_client: PrismaClient) -> _ConfigTable:
return cast(_ConfigTable, ConfigRepository(prisma_client).table) # cast-ok: prisma table is untyped (Any)
@dataclass(frozen=True, slots=True)
class ReloadSchedule:
interval_hours: int | None = None
reload_requested_at: datetime | None = None
last_run_at: datetime | None = None
class ReloadScheduleStatus(TypedDict):
scheduled: bool
interval_hours: int | None
last_run: str | None
next_run: str | None
class _IntervalConfig(BaseModel):
model_config = ConfigDict(strict=True)
interval_hours: int | None = None
def utc_now() -> datetime:
return datetime.now(timezone.utc)
def _parse_interval_hours(param_value: object) -> int | None:
try:
return _IntervalConfig.model_validate(param_value).interval_hours
except ValidationError:
return None
def _as_utc(value: datetime | None) -> datetime | None:
if value is None:
return None
return value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc)
def parse_reload_schedule(row: "LiteLLM_Config") -> ReloadSchedule:
return ReloadSchedule(
interval_hours=_parse_interval_hours(row.param_value),
reload_requested_at=_as_utc(row.reload_requested_at),
last_run_at=_as_utc(row.last_run_at),
)
def next_run_at(schedule: ReloadSchedule) -> datetime | None:
if schedule.interval_hours is None or schedule.last_run_at is None:
return None
return schedule.last_run_at + timedelta(hours=schedule.interval_hours)
def reload_schedule_status(schedule: ReloadSchedule | None) -> ReloadScheduleStatus:
if schedule is None:
return {"scheduled": False, "interval_hours": None, "last_run": None, "next_run": None}
next_run = next_run_at(schedule)
return {
"scheduled": schedule.interval_hours is not None,
"interval_hours": schedule.interval_hours,
"last_run": schedule.last_run_at.isoformat() if schedule.last_run_at is not None else None,
"next_run": next_run.isoformat() if next_run is not None else None,
}
def pod_reload_is_due(
*,
schedule: ReloadSchedule,
pod_last_reload: datetime | None,
current_time: datetime,
description: str,
) -> bool:
"""
Whether this pod should reload now, based on its own last reload rather than the
persisted one, so every pod refreshes its in-memory data on the configured interval.
A pod that has never reloaded ignores ``reload_requested_at``: whatever it booted with
is at least as fresh as the request
"""
if (
schedule.reload_requested_at is not None
and pod_last_reload is not None
and schedule.reload_requested_at > pod_last_reload
):
verbose_proxy_logger.info("%s reload triggered by manual reload request", description)
return True
if schedule.interval_hours is None:
return False
if pod_last_reload is None:
verbose_proxy_logger.info("%s reload triggered - no previous reload time recorded", description)
return True
hours_since_last_reload = (current_time - pod_last_reload).total_seconds() / 3600
if hours_since_last_reload < schedule.interval_hours:
return False
verbose_proxy_logger.info(
"%s reload triggered by interval. Hours since last reload: %.2f, Interval: %s",
description,
hours_since_last_reload,
schedule.interval_hours,
)
return True
async def read_reload_schedule(prisma_client: PrismaClient, param_name: str) -> ReloadSchedule | None:
row = await _config_table(prisma_client).find_unique(where={"param_name": param_name})
if row is None:
return None
return parse_reload_schedule(row)
async def write_reload_interval(prisma_client: PrismaClient, param_name: str, interval_hours: int) -> None:
"""Admin-owned write: replaces ``param_value`` without touching the job-owned columns"""
param_value = safe_dumps({"interval_hours": interval_hours})
await _config_table(prisma_client).upsert(
where={"param_name": param_name},
data={
"create": {"param_name": param_name, "param_value": param_value},
"update": {"param_value": param_value},
},
)
await invalidate_config_param(param_name)
async def record_reload_run(prisma_client: PrismaClient, param_name: str, ran_at: datetime) -> None:
"""Job-owned write after this pod reloaded: stamps the shared last run only if the row
still exists, so a schedule deleted mid-poll is not resurrected"""
await _config_table(prisma_client).update_many(
data={"last_run_at": ran_at},
where={"param_name": param_name},
)
await invalidate_config_param(param_name)
async def record_manual_reload(prisma_client: PrismaClient, param_name: str, ran_at: datetime) -> None:
"""After a manual in-pod reload: stamps the shared last run and the request every other
pod compares against on its next poll"""
await _config_table(prisma_client).upsert(
where={"param_name": param_name},
data={
"create": {"param_name": param_name, "last_run_at": ran_at, "reload_requested_at": ran_at},
"update": {"last_run_at": ran_at, "reload_requested_at": ran_at},
},
)
await invalidate_config_param(param_name)

520
litellm/proxy/proxy_server.py
View file

@ -326,6 +326,17 @@ from litellm.proxy.common_utils.model_listing_utils import TeamModelNameTranslat
from litellm.proxy.common_utils.openai_endpoint_utils import (
remove_sensitive_info_from_deployment,
)
from litellm.proxy.common_utils.periodic_reload_schedule import (
ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME,
MODEL_COST_MAP_RELOAD_PARAM_NAME,
pod_reload_is_due,
read_reload_schedule,
record_manual_reload,
record_reload_run,
reload_schedule_status,
utc_now,
write_reload_interval,
)
from litellm.proxy.common_utils.proxy_state import ProxyState
from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob
from litellm.proxy.common_utils.swagger_utils import ERROR_RESPONSES
@ -2054,12 +2065,7 @@ async_result = None
celery_app_conn = None
celery_fn = None # Redis Queue for handling requests
# Global variables for model cost map reload scheduling
scheduler = None
last_model_cost_map_reload = None
# Global variable for anthropic beta headers reload scheduling
last_anthropic_beta_headers_reload = None
### DB WRITER ###
@ -3826,6 +3832,20 @@ def resolve_complexity_router_plugins(
)
def _reload_model_cost_map_in_process() -> int:
"""Fetch the latest model cost map, swap it into this process's litellm state, return the model count"""
from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map
new_model_cost_map = get_model_cost_map(url=litellm.model_cost_map_url)
litellm.model_cost = new_model_cost_map
# Invalidate case-insensitive lookup map since model_cost was replaced
_invalidate_model_cost_lowercase_map()
# Repopulate provider model sets (e.g. litellm.anthropic_models) so that
# wildcard patterns like "anthropic/*" include any newly added models.
litellm.add_known_models(model_cost_map=new_model_cost_map)
return len(new_model_cost_map) if new_model_cost_map else 0
class ProxyConfig:
"""
Abstraction class on top of config loading/updating logic. Gives us one place to control all config updating logic.
@ -3836,6 +3856,8 @@ class ProxyConfig:
self._last_semantic_filter_config: Optional[Dict[str, Any]] = None
self._last_hashicorp_vault_config: Optional[Dict[str, Any]] = None
self.worker_registry: List["WorkerRegistryEntry"] = []
self.last_model_cost_map_reload: datetime | None = None
self.last_anthropic_beta_headers_reload: datetime | None = None
def is_yaml(self, config_file_path: str) -> bool:
if not os.path.isfile(config_file_path):
@ -6175,8 +6197,6 @@ class ProxyConfig:
"router_settings",
"litellm_settings",
"environment_variables",
"model_cost_map_reload_config",
"anthropic_beta_headers_reload_config",
],
)
@ -6239,12 +6259,6 @@ class ProxyConfig:
if self._should_load_db_object(object_type="tools"):
await self._init_tool_policy_in_db(prisma_client=prisma_client)
if self._should_load_db_object(object_type="model_cost_map"):
await self._check_and_reload_model_cost_map(prisma_client=prisma_client)
if self._should_load_db_object(object_type="anthropic_beta_headers"):
await self._check_and_reload_anthropic_beta_headers(prisma_client=prisma_client)
if self._should_load_db_object(object_type="sso_settings"):
await self._init_sso_settings_in_db(prisma_client=prisma_client)
if self._should_load_db_object(object_type="cache_settings"):
@ -6403,102 +6417,44 @@ class ProxyConfig:
str(e),
)
async def check_periodic_reloads(self, prisma_client: PrismaClient):
"""
Run the admin-configured periodic reloads (model cost map, Anthropic beta headers).
Scheduled on its own job so schedules configured from the Admin UI fire whether
or not `store_model_in_db` is enabled.
"""
if self._should_load_db_object(object_type="model_cost_map"):
await self._check_and_reload_model_cost_map(prisma_client=prisma_client)
if self._should_load_db_object(object_type="anthropic_beta_headers"):
await self._check_and_reload_anthropic_beta_headers(prisma_client=prisma_client)
async def _check_and_reload_model_cost_map(self, prisma_client: PrismaClient):
"""
Check if model cost map needs to be reloaded based on database configuration.
This function runs every 10 seconds as part of _init_non_llm_objects_in_db.
Runs on the periodic reload job, independently of `store_model_in_db`.
"""
try:
# Get model cost map reload configuration from database
config_record = await get_config_param(prisma_client, "model_cost_map_reload_config")
schedule = await read_reload_schedule(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME)
if schedule is None:
return
if config_record is None or config_record.param_value is None:
return # No configuration found, skip reload
current_time = utc_now()
if not pod_reload_is_due(
schedule=schedule,
pod_last_reload=self.last_model_cost_map_reload,
current_time=current_time,
description="Model cost map",
):
return
config = config_record.param_value
interval_hours = config.get("interval_hours")
force_reload = config.get("force_reload", False)
models_count = await asyncio.to_thread(_reload_model_cost_map_in_process)
if interval_hours is None and force_reload is False:
return # No interval configured, skip reload
self.last_model_cost_map_reload = current_time
await record_reload_run(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME, current_time)
current_time = datetime.utcnow()
# Check if we need to reload based on interval or force reload
should_reload = False
if force_reload:
should_reload = True
verbose_proxy_logger.info("Model cost map reload triggered by force reload flag")
elif interval_hours is not None:
# Use pod's in-memory last reload time
global last_model_cost_map_reload
if last_model_cost_map_reload is not None:
try:
last_reload_time = datetime.fromisoformat(last_model_cost_map_reload)
time_since_last_reload = current_time - last_reload_time
hours_since_last_reload = time_since_last_reload.total_seconds() / 3600
if hours_since_last_reload >= interval_hours:
should_reload = True
verbose_proxy_logger.info(
f"Model cost map reload triggered by interval. Hours since last reload: {hours_since_last_reload:.2f}, Interval: {interval_hours}"
)
except Exception as e:
verbose_proxy_logger.warning(f"Error parsing last reload time: {e}")
# If we can't parse the last reload time, reload anyway
should_reload = True
else:
# No last reload time recorded, reload now
should_reload = True
verbose_proxy_logger.info("Model cost map reload triggered - no previous reload time recorded")
if should_reload:
# Perform the reload
from litellm.litellm_core_utils.get_model_cost_map import (
get_model_cost_map,
)
model_cost_map_url = litellm.model_cost_map_url
new_model_cost_map = get_model_cost_map(url=model_cost_map_url)
litellm.model_cost = new_model_cost_map
# Invalidate case-insensitive lookup map since model_cost was replaced
_invalidate_model_cost_lowercase_map()
# Repopulate provider model sets (e.g. litellm.anthropic_models) so that
# wildcard patterns like "anthropic/*" include any newly added models.
litellm.add_known_models(model_cost_map=new_model_cost_map)
# Update pod's in-memory last reload time
last_model_cost_map_reload = current_time.isoformat()
# Clear force reload flag in database
await ConfigRepository(prisma_client).table.upsert(
where={"param_name": "model_cost_map_reload_config"},
data={
"create": {
"param_name": "model_cost_map_reload_config",
"param_value": safe_dumps(
{
"interval_hours": interval_hours,
"force_reload": False,
}
),
},
"update": {
"param_value": safe_dumps(
{
"interval_hours": interval_hours,
"force_reload": False,
}
)
},
},
)
await invalidate_config_param("model_cost_map_reload_config")
verbose_proxy_logger.info(
f"Model cost map reloaded successfully. Models count: {len(new_model_cost_map) if new_model_cost_map else 0}"
)
verbose_proxy_logger.info(f"Model cost map reloaded successfully. Models count: {models_count}")
except Exception as e:
verbose_proxy_logger.exception(f"Error in _check_and_reload_model_cost_map: {str(e)}")
@ -6506,96 +6462,35 @@ class ProxyConfig:
async def _check_and_reload_anthropic_beta_headers(self, prisma_client: PrismaClient):
"""
Check if anthropic beta headers config needs to be reloaded based on database configuration.
This function runs every 10 seconds as part of _init_non_llm_objects_in_db.
Runs on the periodic reload job, independently of `store_model_in_db`.
"""
try:
# Get anthropic beta headers reload configuration from database
config_record = await get_config_param(prisma_client, "anthropic_beta_headers_reload_config")
schedule = await read_reload_schedule(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME)
if schedule is None:
return
if config_record is None or config_record.param_value is None:
return # No configuration found, skip reload
current_time = utc_now()
if not pod_reload_is_due(
schedule=schedule,
pod_last_reload=self.last_anthropic_beta_headers_reload,
current_time=current_time,
description="Anthropic beta headers",
):
return
config = config_record.param_value
interval_hours = config.get("interval_hours")
force_reload = config.get("force_reload", False)
from litellm.anthropic_beta_headers_manager import (
reload_beta_headers_config,
)
if interval_hours is None and force_reload is False:
return # No interval configured, skip reload
new_config = await asyncio.to_thread(reload_beta_headers_config)
current_time = datetime.utcnow()
self.last_anthropic_beta_headers_reload = current_time
await record_reload_run(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME, current_time)
# Check if we need to reload based on interval or force reload
should_reload = False
if force_reload:
should_reload = True
verbose_proxy_logger.info("Anthropic beta headers reload triggered by force reload flag")
elif interval_hours is not None:
# Use pod's in-memory last reload time
global last_anthropic_beta_headers_reload
if last_anthropic_beta_headers_reload is not None:
try:
last_reload_time = datetime.fromisoformat(last_anthropic_beta_headers_reload)
time_since_last_reload = current_time - last_reload_time
hours_since_last_reload = time_since_last_reload.total_seconds() / 3600
if hours_since_last_reload >= interval_hours:
should_reload = True
verbose_proxy_logger.info(
f"Anthropic beta headers reload triggered by interval. Hours since last reload: {hours_since_last_reload:.2f}, Interval: {interval_hours}"
)
except Exception as e:
verbose_proxy_logger.warning(f"Error parsing last reload time: {e}")
# If we can't parse the last reload time, reload anyway
should_reload = True
else:
# No last reload time recorded, reload now
should_reload = True
verbose_proxy_logger.info(
"Anthropic beta headers reload triggered - no previous reload time recorded"
)
if should_reload:
# Perform the reload
from litellm.anthropic_beta_headers_manager import (
reload_beta_headers_config,
)
new_config = reload_beta_headers_config()
# Update pod's in-memory last reload time
last_anthropic_beta_headers_reload = current_time.isoformat()
# Clear force reload flag in database
await ConfigRepository(prisma_client).table.upsert(
where={"param_name": "anthropic_beta_headers_reload_config"},
data={
"create": {
"param_name": "anthropic_beta_headers_reload_config",
"param_value": safe_dumps(
{
"interval_hours": interval_hours,
"force_reload": False,
}
),
},
"update": {
"param_value": safe_dumps(
{
"interval_hours": interval_hours,
"force_reload": False,
}
)
},
},
)
await invalidate_config_param("anthropic_beta_headers_reload_config")
# Count providers in config
provider_count = sum(1 for k in new_config.keys() if k != "provider_aliases" and k != "description")
verbose_proxy_logger.info(
f"Anthropic beta headers config reloaded successfully. Providers: {provider_count}"
)
provider_count = sum(1 for k in new_config.keys() if k != "provider_aliases" and k != "description")
verbose_proxy_logger.info(
f"Anthropic beta headers config reloaded successfully. Providers: {provider_count}"
)
except Exception as e:
verbose_proxy_logger.exception(f"Error in _check_and_reload_anthropic_beta_headers: {str(e)}")
@ -8059,15 +7954,26 @@ class ProxyStartupEvent:
except Exception as e:
verbose_proxy_logger.debug("Failed to check DB for store_model_in_db: %s", str(e))
if store_model_in_db is True:
config_reload_interval_seconds = proxy_config_reload_interval_seconds
if not isinstance(config_reload_interval_seconds, int) or config_reload_interval_seconds <= 0:
verbose_proxy_logger.warning(
"proxy_config_reload_interval_seconds=%s must be a positive integer; falling back to 30s",
config_reload_interval_seconds,
)
config_reload_interval_seconds = 30
config_reload_interval_seconds = proxy_config_reload_interval_seconds
if not isinstance(config_reload_interval_seconds, int) or config_reload_interval_seconds <= 0:
verbose_proxy_logger.warning(
"proxy_config_reload_interval_seconds=%s must be a positive integer; falling back to 30s",
config_reload_interval_seconds,
)
config_reload_interval_seconds = 30
### PERIODIC RELOADS (model cost map, anthropic beta headers) ###
scheduler.add_job(
proxy_config.check_periodic_reloads,
"interval",
seconds=config_reload_interval_seconds,
args=[prisma_client],
id="periodic_reload_job",
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
if store_model_in_db is True:
# MEMORY LEAK FIX: Increase interval from 10s to 30s minimum
# Frequent polling was causing excessive memory allocations
scheduler.add_job(
@ -15661,43 +15567,14 @@ async def reload_model_cost_map(
raise HTTPException(status_code=500, detail="Database connection not available")
# Immediately reload the model cost map in the current pod
from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map
models_count = await asyncio.to_thread(_reload_model_cost_map_in_process)
model_cost_map_url = litellm.model_cost_map_url
new_model_cost_map = get_model_cost_map(url=model_cost_map_url)
litellm.model_cost = new_model_cost_map
# Invalidate case-insensitive lookup map since model_cost was replaced
_invalidate_model_cost_lowercase_map()
# Repopulate provider model sets (e.g. litellm.anthropic_models) so that
# wildcard patterns like "anthropic/*" include any newly added models.
litellm.add_known_models(model_cost_map=new_model_cost_map)
current_time = utc_now()
proxy_config.last_model_cost_map_reload = current_time
# Update pod's in-memory last reload time
global last_model_cost_map_reload
current_time = datetime.utcnow()
last_model_cost_map_reload = current_time.isoformat()
# Stamp the shared last run and the reload request so other pods reload too
await record_manual_reload(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME, current_time)
# Set force reload flag in database for other pods, preserving existing interval_hours
existing_config = await ConfigRepository(prisma_client).table.find_unique(
where={"param_name": "model_cost_map_reload_config"}
)
existing_interval = None
if existing_config and existing_config.param_value:
existing_interval = existing_config.param_value.get("interval_hours")
await ConfigRepository(prisma_client).table.upsert(
where={"param_name": "model_cost_map_reload_config"},
data={
"create": {
"param_name": "model_cost_map_reload_config",
"param_value": safe_dumps({"interval_hours": None, "force_reload": True}),
},
"update": {"param_value": safe_dumps({"interval_hours": existing_interval, "force_reload": True})},
},
)
await invalidate_config_param("model_cost_map_reload_config")
models_count = len(new_model_cost_map) if new_model_cost_map else 0
verbose_proxy_logger.info(f"Model cost map reloaded successfully in current pod. Models count: {models_count}")
return {
@ -15742,18 +15619,7 @@ async def schedule_model_cost_map_reload(
if prisma_client is None:
raise HTTPException(status_code=500, detail="Database connection not available")
# Update database with new reload configuration
await ConfigRepository(prisma_client).table.upsert(
where={"param_name": "model_cost_map_reload_config"},
data={
"create": {
"param_name": "model_cost_map_reload_config",
"param_value": safe_dumps({"interval_hours": hours, "force_reload": False}),
},
"update": {"param_value": safe_dumps({"interval_hours": hours, "force_reload": False})},
},
)
await invalidate_config_param("model_cost_map_reload_config")
await write_reload_interval(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME, hours)
verbose_proxy_logger.info(f"Model cost map reload scheduled for every {hours} hours")
@ -15761,7 +15627,7 @@ async def schedule_model_cost_map_reload(
"message": f"Model cost map reload scheduled for every {hours} hours",
"status": "success",
"interval_hours": hours,
"timestamp": datetime.utcnow().isoformat(),
"timestamp": utc_now().isoformat(),
}
except Exception as e:
verbose_proxy_logger.exception(f"Failed to schedule model cost map reload: {str(e)}")
@ -15798,15 +15664,15 @@ async def cancel_model_cost_map_reload(
raise HTTPException(status_code=500, detail="Database connection not available")
# Remove reload configuration from database
await ConfigRepository(prisma_client).table.delete(where={"param_name": "model_cost_map_reload_config"})
await invalidate_config_param("model_cost_map_reload_config")
await ConfigRepository(prisma_client).table.delete(where={"param_name": MODEL_COST_MAP_RELOAD_PARAM_NAME})
await invalidate_config_param(MODEL_COST_MAP_RELOAD_PARAM_NAME)
verbose_proxy_logger.info("Model cost map reload schedule cancelled")
return {
"message": "Model cost map reload schedule cancelled",
"status": "success",
"timestamp": datetime.utcnow().isoformat(),
"timestamp": utc_now().isoformat(),
}
except Exception as e:
verbose_proxy_logger.exception(f"Failed to cancel model cost map reload: {str(e)}")
@ -15835,66 +15701,13 @@ async def get_model_cost_map_reload_status(
)
try:
global prisma_client, last_model_cost_map_reload
verbose_proxy_logger.info(f"Checking model cost map reload status. Last reload: {last_model_cost_map_reload}")
global prisma_client
if prisma_client is None:
verbose_proxy_logger.info("No database connection, returning not scheduled")
return {
"scheduled": False,
"interval_hours": None,
"last_run": None,
"next_run": None,
}
return reload_schedule_status(None)
# Get reload configuration from database
config_record = await ConfigRepository(prisma_client).table.find_unique(
where={"param_name": "model_cost_map_reload_config"}
)
if config_record is None or config_record.param_value is None:
verbose_proxy_logger.info("No model cost map reload configuration found")
return {
"scheduled": False,
"interval_hours": None,
"last_run": None,
"next_run": None,
}
config = config_record.param_value
interval_hours = config.get("interval_hours")
if interval_hours is None:
verbose_proxy_logger.info("No interval configured, returning not scheduled")
return {
"scheduled": False,
"interval_hours": None,
"last_run": None,
"next_run": None,
}
current_time = datetime.utcnow()
next_run = None
# Use pod's in-memory last reload time
if last_model_cost_map_reload is not None:
try:
last_reload_time = datetime.fromisoformat(last_model_cost_map_reload)
time_since_last_reload = current_time - last_reload_time
hours_since_last_reload = time_since_last_reload.total_seconds() / 3600
if hours_since_last_reload < interval_hours:
next_run = (last_reload_time + timedelta(hours=interval_hours)).isoformat()
except Exception as e:
verbose_proxy_logger.warning(f"Error parsing last reload time: {e}")
return {
"scheduled": True,
"interval_hours": interval_hours,
"last_run": last_model_cost_map_reload,
"next_run": next_run,
}
return reload_schedule_status(await read_reload_schedule(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME))
except Exception as e:
verbose_proxy_logger.exception(f"Failed to get model cost map reload status: {str(e)}")
raise HTTPException(
@ -15984,32 +15797,13 @@ async def reload_anthropic_beta_headers(
# Immediately reload the beta headers config in the current pod
from litellm.anthropic_beta_headers_manager import reload_beta_headers_config
new_config = reload_beta_headers_config()
new_config = await asyncio.to_thread(reload_beta_headers_config)
# Update pod's in-memory last reload time
global last_anthropic_beta_headers_reload
current_time = datetime.utcnow()
last_anthropic_beta_headers_reload = current_time.isoformat()
current_time = utc_now()
proxy_config.last_anthropic_beta_headers_reload = current_time
# Set force reload flag in database for other pods, preserving existing interval_hours
existing_beta_config = await ConfigRepository(prisma_client).table.find_unique(
where={"param_name": "anthropic_beta_headers_reload_config"}
)
existing_beta_interval = None
if existing_beta_config and existing_beta_config.param_value:
existing_beta_interval = existing_beta_config.param_value.get("interval_hours")
await ConfigRepository(prisma_client).table.upsert(
where={"param_name": "anthropic_beta_headers_reload_config"},
data={
"create": {
"param_name": "anthropic_beta_headers_reload_config",
"param_value": safe_dumps({"interval_hours": None, "force_reload": True}),
},
"update": {"param_value": safe_dumps({"interval_hours": existing_beta_interval, "force_reload": True})},
},
)
await invalidate_config_param("anthropic_beta_headers_reload_config")
# Stamp the shared last run and the reload request so other pods reload too
await record_manual_reload(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME, current_time)
provider_count = sum(1 for k in new_config.keys() if k not in ["provider_aliases", "description"])
verbose_proxy_logger.info(
@ -16058,18 +15852,7 @@ async def schedule_anthropic_beta_headers_reload(
if prisma_client is None:
raise HTTPException(status_code=500, detail="Database connection not available")
# Update database with new reload configuration
await ConfigRepository(prisma_client).table.upsert(
where={"param_name": "anthropic_beta_headers_reload_config"},
data={
"create": {
"param_name": "anthropic_beta_headers_reload_config",
"param_value": safe_dumps({"interval_hours": hours, "force_reload": False}),
},
"update": {"param_value": safe_dumps({"interval_hours": hours, "force_reload": False})},
},
)
await invalidate_config_param("anthropic_beta_headers_reload_config")
await write_reload_interval(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME, hours)
verbose_proxy_logger.info(f"Anthropic beta headers reload scheduled for every {hours} hours")
@ -16077,7 +15860,7 @@ async def schedule_anthropic_beta_headers_reload(
"message": f"Anthropic beta headers reload scheduled for every {hours} hours",
"status": "success",
"interval_hours": hours,
"timestamp": datetime.utcnow().isoformat(),
"timestamp": utc_now().isoformat(),
}
except Exception as e:
verbose_proxy_logger.exception(f"Failed to schedule anthropic beta headers reload: {str(e)}")
@ -16114,15 +15897,17 @@ async def cancel_anthropic_beta_headers_reload(
raise HTTPException(status_code=500, detail="Database connection not available")
# Remove reload configuration from database
await ConfigRepository(prisma_client).table.delete(where={"param_name": "anthropic_beta_headers_reload_config"})
await invalidate_config_param("anthropic_beta_headers_reload_config")
await ConfigRepository(prisma_client).table.delete(
where={"param_name": ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME}
)
await invalidate_config_param(ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME)
verbose_proxy_logger.info("Anthropic beta headers reload schedule cancelled")
return {
"message": "Anthropic beta headers reload schedule cancelled",
"status": "success",
"timestamp": datetime.utcnow().isoformat(),
"timestamp": utc_now().isoformat(),
}
except Exception as e:
verbose_proxy_logger.exception(f"Failed to cancel anthropic beta headers reload: {str(e)}")
@ -16154,68 +15939,15 @@ async def get_anthropic_beta_headers_reload_status(
)
try:
global prisma_client, last_anthropic_beta_headers_reload
verbose_proxy_logger.info(
f"Checking anthropic beta headers reload status. Last reload: {last_anthropic_beta_headers_reload}"
)
global prisma_client
if prisma_client is None:
verbose_proxy_logger.info("No database connection, returning not scheduled")
return {
"scheduled": False,
"interval_hours": None,
"last_run": None,
"next_run": None,
}
return reload_schedule_status(None)
# Get reload configuration from database
config_record = await ConfigRepository(prisma_client).table.find_unique(
where={"param_name": "anthropic_beta_headers_reload_config"}
return reload_schedule_status(
await read_reload_schedule(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME)
)
if config_record is None or config_record.param_value is None:
verbose_proxy_logger.info("No anthropic beta headers reload configuration found")
return {
"scheduled": False,
"interval_hours": None,
"last_run": None,
"next_run": None,
}
config = config_record.param_value
interval_hours = config.get("interval_hours")
if interval_hours is None:
verbose_proxy_logger.info("No interval configured, returning not scheduled")
return {
"scheduled": False,
"interval_hours": None,
"last_run": None,
"next_run": None,
}
current_time = datetime.utcnow()
next_run = None
# Use pod's in-memory last reload time
if last_anthropic_beta_headers_reload is not None:
try:
last_reload_time = datetime.fromisoformat(last_anthropic_beta_headers_reload)
time_since_last_reload = current_time - last_reload_time
hours_since_last_reload = time_since_last_reload.total_seconds() / 3600
if hours_since_last_reload < interval_hours:
next_run = (last_reload_time + timedelta(hours=interval_hours)).isoformat()
except Exception as e:
verbose_proxy_logger.warning(f"Error parsing last reload time: {e}")
return {
"scheduled": True,
"interval_hours": interval_hours,
"last_run": last_anthropic_beta_headers_reload,
"next_run": next_run,
}
except Exception as e:
verbose_proxy_logger.exception(f"Failed to get anthropic beta headers reload status: {str(e)}")
raise HTTPException(

2
litellm/proxy/schema.prisma
View file

@ -601,6 +601,8 @@ model LiteLLM_TagTable {
model LiteLLM_Config {
param_name String @id
param_value Json?
last_run_at DateTime?
reload_requested_at DateTime?
}
// View spend, model, api_key per request

2
ruff.toml
View file

@ -6,7 +6,7 @@ lint.extend-select = ["T20", "PGH004", "RUF008", "RUF009", "RUF100"]
# litellm's own ruff config both rely on suppressions this config can't see.
lint.external = [
# Enforced by the strict-rule gate (scripts/ruff_strict_gate.py + ruff-strict.toml)
"C901",
"C901", "TID251",
# Enforced by upstream litellm's ruff config, but not run in this repo's CI
"PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405",
]

2
schema.prisma
View file

@ -601,6 +601,8 @@ model LiteLLM_TagTable {
model LiteLLM_Config {
param_name String @id
param_value Json?
last_run_at DateTime?
reload_requested_at DateTime?
}
// View spend, model, api_key per request

213
tests/test_litellm/proxy/common_utils/test_periodic_reload_schedule.py Normal file
View file

@ -0,0 +1,213 @@
from datetime import datetime, timezone
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy.common_utils.periodic_reload_schedule import (
ReloadSchedule,
next_run_at,
parse_reload_schedule,
pod_reload_is_due,
read_reload_schedule,
record_manual_reload,
record_reload_run,
reload_schedule_status,
write_reload_interval,
)
LAST_RUN = datetime(2024, 1, 1, 6, 0, 0, tzinfo=timezone.utc)
NOW = datetime(2024, 1, 1, 12, 0, tzinfo=timezone.utc)
def _row(param_value=None, reload_requested_at=None, last_run_at=None):
return SimpleNamespace(
param_name="model_cost_map_reload_config",
param_value=param_value,
reload_requested_at=reload_requested_at,
last_run_at=last_run_at,
)
def _mock_prisma(row=None):
prisma_client = MagicMock()
prisma_client.db.litellm_config.find_unique = AsyncMock(return_value=row)
prisma_client.db.litellm_config.upsert = AsyncMock(return_value=None)
prisma_client.db.litellm_config.update_many = AsyncMock(return_value=1)
return prisma_client
def test_parse_reads_interval_from_json_and_state_from_columns():
schedule = parse_reload_schedule(
_row(param_value={"interval_hours": 6}, reload_requested_at=NOW, last_run_at=LAST_RUN)
)
assert schedule == ReloadSchedule(interval_hours=6, reload_requested_at=NOW, last_run_at=LAST_RUN)
def test_parse_treats_naive_column_timestamps_as_utc():
schedule = parse_reload_schedule(
_row(reload_requested_at=NOW.replace(tzinfo=None), last_run_at=LAST_RUN.replace(tzinfo=None))
)
assert (schedule.reload_requested_at, schedule.last_run_at) == (NOW, LAST_RUN)
@pytest.mark.parametrize(
"param_value",
[None, "not-a-dict", {}, {"interval_hours": "6"}, {"interval_hours": None}],
)
def test_parse_tolerates_unusable_param_values(param_value):
assert parse_reload_schedule(_row(param_value=param_value)).interval_hours is None
def test_parse_ignores_legacy_json_force_reload():
"""Rows written by pre-column versions carry force_reload in the JSON; honoring it
would re-trigger a reload every poll because nothing clears the JSON copy"""
schedule = parse_reload_schedule(_row(param_value={"interval_hours": 6, "force_reload": True}))
assert schedule == ReloadSchedule(interval_hours=6, reload_requested_at=None, last_run_at=None)
def test_status_reports_persisted_last_run_and_next_run():
status = reload_schedule_status(ReloadSchedule(interval_hours=6, last_run_at=LAST_RUN))
assert status == {
"scheduled": True,
"interval_hours": 6,
"last_run": "2024-01-01T06:00:00+00:00",
"next_run": "2024-01-01T12:00:00+00:00",
}
def test_status_without_interval_is_not_scheduled():
assert reload_schedule_status(None)["scheduled"] is False
assert reload_schedule_status(ReloadSchedule(last_run_at=LAST_RUN)) == {
"scheduled": False,
"interval_hours": None,
"last_run": "2024-01-01T06:00:00+00:00",
"next_run": None,
}
def test_next_run_needs_both_interval_and_last_run():
assert next_run_at(ReloadSchedule(interval_hours=6)) is None
assert next_run_at(ReloadSchedule(last_run_at=LAST_RUN)) is None
@pytest.mark.parametrize(
"schedule, pod_last_reload, expected",
[
(ReloadSchedule(reload_requested_at=datetime(2024, 1, 1, 11, 0, tzinfo=timezone.utc)), LAST_RUN, True),
(
ReloadSchedule(interval_hours=6, reload_requested_at=datetime(2024, 1, 1, 11, 30, tzinfo=timezone.utc)),
datetime(2024, 1, 1, 11, 0, tzinfo=timezone.utc),
True,
),
(ReloadSchedule(reload_requested_at=LAST_RUN), LAST_RUN, False),
(ReloadSchedule(reload_requested_at=datetime(2024, 1, 1, 5, 0, tzinfo=timezone.utc)), LAST_RUN, False),
(ReloadSchedule(interval_hours=6, reload_requested_at=LAST_RUN), LAST_RUN, True),
(ReloadSchedule(reload_requested_at=NOW), None, False),
(ReloadSchedule(interval_hours=6, reload_requested_at=NOW), None, True),
(ReloadSchedule(), None, False),
(ReloadSchedule(interval_hours=6), None, True),
(ReloadSchedule(interval_hours=6), datetime(2024, 1, 1, 11, 0, tzinfo=timezone.utc), False),
(ReloadSchedule(interval_hours=6), LAST_RUN, True),
],
)
def test_pod_reload_is_due(schedule, pod_last_reload, expected):
assert (
pod_reload_is_due(
schedule=schedule,
pod_last_reload=pod_last_reload,
current_time=NOW,
description="test",
)
is expected
)
def test_manual_request_reaches_every_pod_that_has_not_seen_it():
"""Nothing clears reload_requested_at, so each pod reloads exactly once per request:
due while its own last reload predates the request, not due afterwards"""
schedule = ReloadSchedule(reload_requested_at=datetime(2024, 1, 1, 11, 0, tzinfo=timezone.utc))
assert pod_reload_is_due(schedule=schedule, pod_last_reload=LAST_RUN, current_time=NOW, description="test") is True
assert (
pod_reload_is_due(
schedule=schedule,
pod_last_reload=datetime(2024, 1, 1, 11, 0, 1, tzinfo=timezone.utc),
current_time=NOW,
description="test",
)
is False
)
def test_pod_reload_decision_ignores_persisted_last_run():
"""A pod that has never reloaded must refresh its own copy even when another pod
already stamped last_run_at within the interval"""
assert (
pod_reload_is_due(
schedule=ReloadSchedule(interval_hours=6, last_run_at=datetime(2024, 1, 1, 11, 59, tzinfo=timezone.utc)),
pod_last_reload=None,
current_time=NOW,
description="test",
)
is True
)
@pytest.mark.asyncio
async def test_read_reload_schedule_returns_none_for_missing_row():
assert await read_reload_schedule(_mock_prisma(row=None), "model_cost_map_reload_config") is None
@pytest.mark.asyncio
async def test_read_reload_schedule_surfaces_request_on_interval_less_row():
"""A manual reload on a proxy with no schedule creates a row with only the columns
set; the request must still reach other pods"""
prisma_client = _mock_prisma(row=_row(param_value=None, reload_requested_at=NOW))
schedule = await read_reload_schedule(prisma_client, "model_cost_map_reload_config")
assert schedule == ReloadSchedule(interval_hours=None, reload_requested_at=NOW, last_run_at=None)
@pytest.mark.asyncio
async def test_write_reload_interval_touches_only_param_value():
prisma_client = _mock_prisma()
await write_reload_interval(prisma_client, "model_cost_map_reload_config", 12)
data = prisma_client.db.litellm_config.upsert.await_args.kwargs["data"]
assert set(data["update"]) == {"param_value"}
assert set(data["create"]) == {"param_name", "param_value"}
@pytest.mark.asyncio
async def test_record_reload_run_updates_last_run_without_creating_or_clearing():
"""update_many so a schedule deleted mid-poll stays deleted, and the untouched
reload_requested_at keeps fanning the request out to pods that have not seen it"""
prisma_client = _mock_prisma()
await record_reload_run(prisma_client, "model_cost_map_reload_config", LAST_RUN)
kwargs = prisma_client.db.litellm_config.update_many.await_args.kwargs
assert kwargs == {"data": {"last_run_at": LAST_RUN}, "where": {"param_name": "model_cost_map_reload_config"}}
prisma_client.db.litellm_config.upsert.assert_not_called()
@pytest.mark.asyncio
async def test_record_manual_reload_stamps_request_and_last_run():
prisma_client = _mock_prisma()
await record_manual_reload(prisma_client, "model_cost_map_reload_config", LAST_RUN)
data = prisma_client.db.litellm_config.upsert.await_args.kwargs["data"]
assert data["update"] == {"last_run_at": LAST_RUN, "reload_requested_at": LAST_RUN}
assert data["create"] == {
"param_name": "model_cost_map_reload_config",
"last_run_at": LAST_RUN,
"reload_requested_at": LAST_RUN,
}

71
tests/test_litellm/proxy/proxy_server/test_routes_anthropic_beta.py
View file

@ -9,11 +9,9 @@ Routes covered:
from __future__ import annotations
import json
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
import pytest
from .conftest import VOLATILE_KEYS, normalize
@ -99,22 +97,24 @@ def test_reload_anthropic_beta_headers_admin_success(client, auth_as, monkeypatc
"providers_count": 2,
"timestamp": "<VOLATILE>",
}
# And the upsert was actually invoked (force_reload write).
prisma.db.litellm_config.upsert.assert_awaited_once()
def test_reload_anthropic_beta_headers_preserves_existing_interval(
client, auth_as, monkeypatch
):
"""When an existing reload config has an interval set, the force-reload
write must preserve that interval (the route reads it back then upserts
with the same number). This pins the read-then-write behaviour."""
"""The manual-reload write only stamps the job-owned timestamp columns; param_value
(the admin-owned interval) is not in the payload, so it cannot be clobbered."""
from datetime import datetime
from litellm.proxy._types import LitellmUserRoles
_stub_reload_beta_headers(monkeypatch)
existing = SimpleNamespace(
param_name="anthropic_beta_headers_reload_config",
param_value={"interval_hours": 12, "force_reload": False},
param_value={"interval_hours": 12},
reload_requested_at=None,
last_run_at=None,
)
prisma = _make_prisma_with_config(config_record=existing)
_install_prisma(monkeypatch, prisma)
@ -123,17 +123,12 @@ def test_reload_anthropic_beta_headers_preserves_existing_interval(
response = client.post("/reload/anthropic_beta_headers")
assert response.status_code == 200
# The update branch's interval_hours was sourced from the existing record.
call_kwargs = prisma.db.litellm_config.upsert.await_args.kwargs
data = call_kwargs["data"]
update_payload = data["update"]["param_value"]
parsed = (
json.loads(update_payload)
if isinstance(update_payload, str)
else update_payload
)
assert parsed["interval_hours"] == 12
assert parsed["force_reload"] is True
update_payload = call_kwargs["data"]["update"]
assert set(update_payload) == {"last_run_at", "reload_requested_at"}
assert isinstance(update_payload["last_run_at"], datetime)
assert update_payload["reload_requested_at"] == update_payload["last_run_at"]
assert "param_value" not in call_kwargs["data"]["create"]
def test_reload_anthropic_beta_headers_not_admin_forbidden(client, auth_as):
@ -291,17 +286,16 @@ def test_get_anthropic_beta_headers_reload_status_scheduled(
):
"""When a config row with ``interval_hours`` is present, ``scheduled`` is True
and ``interval_hours`` echoes the DB value. Pins the full response shape."""
from litellm.proxy import proxy_server as ps
from litellm.proxy._types import LitellmUserRoles
record = SimpleNamespace(
param_name="anthropic_beta_headers_reload_config",
param_value={"interval_hours": 6, "force_reload": False},
param_value={"interval_hours": 6},
reload_requested_at=None,
last_run_at=None,
)
prisma = _make_prisma_with_config(config_record=record)
_install_prisma(monkeypatch, prisma)
# No prior reload — next_run stays None.
monkeypatch.setattr(ps, "last_anthropic_beta_headers_reload", None)
with auth_as(LitellmUserRoles.PROXY_ADMIN):
response = client.get("/schedule/anthropic_beta_headers_reload/status")
@ -315,6 +309,35 @@ def test_get_anthropic_beta_headers_reload_status_scheduled(
}
def test_get_anthropic_beta_headers_reload_status_reports_persisted_last_run(
client, auth_as, monkeypatch
):
"""last_run/next_run come from the DB row, so status survives pod restarts."""
from datetime import datetime, timezone
from litellm.proxy._types import LitellmUserRoles
record = SimpleNamespace(
param_name="anthropic_beta_headers_reload_config",
param_value={"interval_hours": 6},
reload_requested_at=None,
last_run_at=datetime(2024, 1, 1, 6, 0, 0, tzinfo=timezone.utc),
)
prisma = _make_prisma_with_config(config_record=record)
_install_prisma(monkeypatch, prisma)
with auth_as(LitellmUserRoles.PROXY_ADMIN):
response = client.get("/schedule/anthropic_beta_headers_reload/status")
assert response.status_code == 200
assert normalize(response.json()) == {
"scheduled": True,
"interval_hours": 6,
"last_run": "2024-01-01T06:00:00+00:00",
"next_run": "2024-01-01T12:00:00+00:00",
}
def test_get_anthropic_beta_headers_reload_status_not_scheduled_no_db(
client, auth_as, monkeypatch
):
@ -339,11 +362,15 @@ def test_get_anthropic_beta_headers_reload_status_no_interval_unscheduled(
client, auth_as, monkeypatch
):
"""Config row present but ``interval_hours`` is None → unscheduled response."""
from datetime import datetime, timezone
from litellm.proxy._types import LitellmUserRoles
record = SimpleNamespace(
param_name="anthropic_beta_headers_reload_config",
param_value={"interval_hours": None, "force_reload": True},
param_value={"interval_hours": None},
reload_requested_at=datetime(2024, 1, 1, 6, 0, 0, tzinfo=timezone.utc),
last_run_at=None,
)
prisma = _make_prisma_with_config(config_record=record)
_install_prisma(monkeypatch, prisma)

51
tests/test_litellm/proxy/proxy_server/test_routes_model_cost_map.py
View file

@ -12,7 +12,6 @@ from __future__ import annotations
from unittest.mock import AsyncMock, MagicMock
import pytest
from .conftest import VOLATILE_KEYS, normalize
@ -83,6 +82,9 @@ def test_reload_model_cost_map_happy(client, auth_as, monkeypatch, mock_prisma):
"timestamp": "<VOLATILE>",
}
assert table.upsert.await_count == 1
update_payload = table.upsert.await_args.kwargs["data"]["update"]
assert set(update_payload) == {"last_run_at", "reload_requested_at"}
assert update_payload["reload_requested_at"] == update_payload["last_run_at"]
def test_reload_model_cost_map_not_admin_forbidden(client, auth_as):
@ -138,6 +140,9 @@ def test_schedule_model_cost_map_reload_happy(
"timestamp": "<VOLATILE>",
}
assert table.upsert.await_count == 1
upsert_data = table.upsert.await_args.kwargs["data"]
assert set(upsert_data["update"]) == {"param_value"}
assert set(upsert_data["create"]) == {"param_name", "param_value"}
def test_schedule_model_cost_map_reload_invalid_hours(
@ -249,10 +254,11 @@ def test_get_model_cost_map_reload_status_scheduled(
table = _attach_litellm_config(mock_prisma)
config_row = MagicMock()
config_row.param_value = {"interval_hours": 12, "force_reload": False}
config_row.param_value = {"interval_hours": 12}
config_row.reload_requested_at = None
config_row.last_run_at = None
table.find_unique = AsyncMock(return_value=config_row)
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
monkeypatch.setattr(ps, "last_model_cost_map_reload", None)
with auth_as(LitellmUserRoles.PROXY_ADMIN):
response = client.get("/schedule/model_cost_map_reload/status")
@ -265,19 +271,50 @@ def test_get_model_cost_map_reload_status_scheduled(
}
def test_get_model_cost_map_reload_status_no_config_not_scheduled(
def test_get_model_cost_map_reload_status_reports_persisted_last_run(
client, auth_as, monkeypatch, mock_prisma
):
"""Config row exists but interval_hours=None → not scheduled."""
"""last_run/next_run come from the DB row, so status survives pod restarts."""
from datetime import datetime, timezone
from litellm.proxy import proxy_server as ps
from litellm.proxy._types import LitellmUserRoles
table = _attach_litellm_config(mock_prisma)
config_row = MagicMock()
config_row.param_value = {"interval_hours": None, "force_reload": True}
config_row.param_value = {"interval_hours": 6}
config_row.reload_requested_at = None
config_row.last_run_at = datetime(2024, 1, 1, 6, 0, 0, tzinfo=timezone.utc)
table.find_unique = AsyncMock(return_value=config_row)
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
with auth_as(LitellmUserRoles.PROXY_ADMIN):
response = client.get("/schedule/model_cost_map_reload/status")
assert response.status_code == 200
assert normalize(response.json()) == {
"scheduled": True,
"interval_hours": 6,
"last_run": "2024-01-01T06:00:00+00:00",
"next_run": "2024-01-01T12:00:00+00:00",
}
def test_get_model_cost_map_reload_status_no_config_not_scheduled(
client, auth_as, monkeypatch, mock_prisma
):
"""A row left behind by a manual reload (interval_hours=None) → not scheduled."""
from datetime import datetime, timezone
from litellm.proxy import proxy_server as ps
from litellm.proxy._types import LitellmUserRoles
table = _attach_litellm_config(mock_prisma)
config_row = MagicMock()
config_row.param_value = {"interval_hours": None}
config_row.reload_requested_at = datetime(2024, 1, 1, 6, 0, 0, tzinfo=timezone.utc)
config_row.last_run_at = None
table.find_unique = AsyncMock(return_value=config_row)
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
monkeypatch.setattr(ps, "last_model_cost_map_reload", None)
with auth_as(LitellmUserRoles.PROXY_ADMIN):
response = client.get("/schedule/model_cost_map_reload/status")

637
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -755,6 +755,50 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch):
assert len(mock_scheduler_calls) > 0
@pytest.mark.asyncio
async def test_periodic_reload_job_scheduled_without_store_model_in_db(monkeypatch):
"""
Regression (LIT-4882): reload schedules configured from the Admin UI live in the DB and
must fire even without store_model_in_db, which used to gate the job that ran them
"""
monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False)
monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False)
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from litellm.proxy.proxy_server import ProxyStartupEvent
from litellm.proxy.utils import ProxyLogging
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=None)
mock_proxy_logging = MagicMock(spec=ProxyLogging)
mock_proxy_logging.slack_alerting_instance = MagicMock()
mock_proxy_config = AsyncMock()
scheduler = AsyncIOScheduler()
try:
with (
patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config),
patch("litellm.proxy.proxy_server.store_model_in_db", False),
patch("litellm.proxy.proxy_server.get_secret_bool", return_value=False),
patch(
"litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler
),
):
await ProxyStartupEvent.initialize_scheduled_background_jobs(
general_settings={},
prisma_client=mock_prisma_client,
proxy_budget_rescheduler_min_time=1,
proxy_budget_rescheduler_max_time=2,
proxy_batch_write_at=5,
proxy_logging_obj=mock_proxy_logging,
)
assert scheduler.get_job("periodic_reload_job") is not None
assert scheduler.get_job("add_deployment_job") is None
finally:
scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_initialize_scheduled_jobs_uses_configured_config_reload_interval(monkeypatch):
"""
@ -3592,6 +3636,20 @@ async def test_chat_completion_result_no_nested_none_values():
# ============================================================================
def _reload_schedule_row(
param_value: dict,
*,
reload_requested_at: datetime | None = None,
last_run_at: datetime | None = None,
) -> types.SimpleNamespace:
"""LiteLLM_Config row shape: admin-owned interval in param_value, run state in dedicated columns"""
return types.SimpleNamespace(
param_value=param_value,
reload_requested_at=reload_requested_at,
last_run_at=last_run_at,
)
class TestPriceDataReloadAPI:
"""Test cases for price data reload API endpoints"""
@ -3628,9 +3686,6 @@ class TestPriceDataReloadAPI:
}
# Mock the database connection
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
mock_prisma.db.litellm_config.find_unique = AsyncMock(
return_value=None
)
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
response = client_with_auth.post("/reload/model_cost_map")
@ -3698,7 +3753,7 @@ class TestPriceDataReloadAPI:
assert "Failed to reload model cost map" in data["detail"]
def test_schedule_model_cost_map_reload_admin_access(self, client_with_auth):
"""Test that admin users can schedule periodic reload"""
"""Admin schedule write owns param_value only, so it can't clobber the job-owned run columns"""
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
# Mock database upsert
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
@ -3712,6 +3767,17 @@ class TestPriceDataReloadAPI:
assert "message" in data
assert "timestamp" in data
call_args = mock_prisma.db.litellm_config.upsert.call_args
assert call_args[1]["where"] == {
"param_name": "model_cost_map_reload_config"
}
update_payload = call_args[1]["data"]["update"]
assert set(update_payload.keys()) == {"param_value"}
assert json.loads(update_payload["param_value"]) == {"interval_hours": 6}
create_payload = call_args[1]["data"]["create"]
assert set(create_payload.keys()) == {"param_name", "param_value"}
assert json.loads(create_payload["param_value"]) == {"interval_hours": 6}
def test_schedule_model_cost_map_reload_non_admin_access(self, client_with_auth):
"""Test that non-admin users cannot schedule periodic reload"""
# Mock non-admin user
@ -3763,35 +3829,28 @@ class TestPriceDataReloadAPI:
assert "Admin role required" in data["detail"]
def test_get_model_cost_map_reload_status_admin_access(self, client_with_auth):
"""Test that admin users can get reload status"""
"""
Regression (LIT-4882): status is served purely from the DB row, so a restarted pod
with no in-memory last-reload time still reports the real last/next run
"""
proxy_server_module.proxy_config.last_model_cost_map_reload = None
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
# Mock database config record
mock_config = MagicMock()
mock_config.param_value = {"interval_hours": 6, "force_reload": False}
mock_prisma.db.litellm_config.find_unique = AsyncMock(
return_value=mock_config
return_value=_reload_schedule_row(
{"interval_hours": 6},
last_run_at=datetime(2024, 1, 1, 6, 0, tzinfo=timezone.utc),
)
)
# Mock the last reload time and current time
with patch(
"litellm.proxy.proxy_server.last_model_cost_map_reload",
"2024-01-01T06:00:00",
):
with patch("litellm.proxy.proxy_server.datetime") as mock_datetime:
# Mock current time to be 1 hour after last reload
mock_datetime.utcnow.return_value = datetime(2024, 1, 1, 7, 0, 0)
mock_datetime.fromisoformat = datetime.fromisoformat
response = client_with_auth.get("/schedule/model_cost_map_reload/status")
response = client_with_auth.get(
"/schedule/model_cost_map_reload/status"
)
assert response.status_code == 200
data = response.json()
assert data["scheduled"] == True
assert data["interval_hours"] == 6
assert data["last_run"] == "2024-01-01T06:00:00"
assert data["next_run"] == "2024-01-01T12:00:00"
assert response.status_code == 200
data = response.json()
assert data["scheduled"] is True
assert data["interval_hours"] == 6
assert data["last_run"] == "2024-01-01T06:00:00+00:00"
assert data["next_run"] == "2024-01-01T12:00:00+00:00"
def test_get_model_cost_map_reload_status_non_admin_access(self, client_with_auth):
"""Test that non-admin users cannot get reload status"""
@ -3822,36 +3881,44 @@ class TestPriceDataReloadAPI:
assert data["next_run"] == None
def test_get_model_cost_map_reload_status_no_interval(self, client_with_auth):
"""Test that status returns not scheduled when no interval is configured"""
"""A row left behind by a manual reload (no interval) must not read as scheduled"""
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
# Mock config with no interval
mock_config = MagicMock()
mock_config.param_value = {"interval_hours": None, "force_reload": False}
mock_prisma.db.litellm_config.find_unique = AsyncMock(
return_value=mock_config
return_value=_reload_schedule_row(
{"interval_hours": None},
reload_requested_at=datetime(2024, 1, 1, 6, 0, tzinfo=timezone.utc),
)
)
response = client_with_auth.get("/schedule/model_cost_map_reload/status")
assert response.status_code == 200
data = response.json()
assert data["scheduled"] == False
assert data["interval_hours"] == None
assert data["last_run"] == None
assert data["next_run"] == None
assert data["scheduled"] is False
assert data["interval_hours"] is None
assert data["last_run"] is None
assert data["next_run"] is None
def test_get_model_cost_map_reload_status_before_first_run(self, client_with_auth):
"""Scheduled but never executed: no last_run_at means no next_run can be computed"""
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
mock_prisma.db.litellm_config.find_unique = AsyncMock(
return_value=_reload_schedule_row({"interval_hours": 6})
)
response = client_with_auth.get("/schedule/model_cost_map_reload/status")
assert response.status_code == 200
data = response.json()
assert data["scheduled"] is True
assert data["interval_hours"] == 6
assert data["last_run"] is None
assert data["next_run"] is None
class TestPriceDataReloadIntegration:
"""Integration tests for the complete price data reload feature"""
@pytest.fixture(autouse=True)
def _flush_litellm_config_cache(self):
from litellm.proxy.utils import litellm_config_cache
litellm_config_cache.flush_cache()
yield
litellm_config_cache.flush_cache()
@pytest.fixture
def client_with_auth(self):
"""Create a test client with authentication"""
@ -3890,9 +3957,6 @@ class TestPriceDataReloadIntegration:
# Mock the database connection
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
mock_prisma.db.litellm_config.find_unique = AsyncMock(
return_value=None
)
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
# Test reload endpoint
@ -3907,122 +3971,193 @@ class TestPriceDataReloadIntegration:
_invalidate_model_cost_lowercase_map()
def test_distributed_reload_check_function(self):
"""Test the _check_and_reload_model_cost_map function"""
from litellm.proxy.proxy_server import ProxyConfig
from litellm.proxy.utils import litellm_config_cache
proxy_config = ProxyConfig()
# Mock prisma client
mock_prisma = MagicMock()
# Test case 1: No config in database
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None)
# _check_and_reload_model_cost_map routes through get_config_param,
# which calls prisma.get_generic_data on a cache miss.
mock_prisma.get_generic_data = AsyncMock(return_value=None)
# Should return early without reloading
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
# Test case 2: Config with interval but not time to reload
litellm_config_cache.flush_cache()
mock_config = MagicMock()
mock_config.param_value = {"interval_hours": 6, "force_reload": False}
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config)
mock_prisma.get_generic_data = AsyncMock(return_value=mock_config)
# Mock current time and last reload time
with patch(
"litellm.proxy.proxy_server.last_model_cost_map_reload",
"2024-01-01T06:00:00",
):
with patch("litellm.proxy.proxy_server.datetime") as mock_datetime:
mock_datetime.utcnow.return_value = datetime(
2024, 1, 1, 7, 0, 0
) # 1 hour later
# Should not reload (only 1 hour passed, need 6)
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
# Test case 3: Config with force reload
litellm_config_cache.flush_cache()
mock_config.param_value = {"interval_hours": 6, "force_reload": True}
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config)
mock_prisma.get_generic_data = AsyncMock(return_value=mock_config)
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
original_model_cost = litellm.model_cost.copy()
try:
with patch(
"litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map"
) as mock_get_map:
mock_get_map.return_value = {
"gpt-3.5-turbo": {"input_cost_per_token": 0.001}
}
# Should reload due to force flag
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
# Verify force_reload was reset to False
mock_prisma.db.litellm_config.upsert.assert_called()
call_args = mock_prisma.db.litellm_config.upsert.call_args
# The param_value is now a JSON string, so we need to parse it
param_value_json = call_args[1]["data"]["update"]["param_value"]
param_value_dict = json.loads(param_value_json)
assert param_value_dict["force_reload"] == False
assert param_value_dict.get("interval_hours") == 6
finally:
litellm.model_cost = original_model_cost
_invalidate_model_cost_lowercase_map()
def test_distributed_reload_preserves_interval_hours(self):
"""Test that _check_and_reload_model_cost_map preserves interval_hours after reload.
Regression test: the update branch of the upsert was previously dropping
interval_hours, causing scheduled reloads to self-destruct after first execution.
"""
A reload requested by another pod after this pod's last reload takes effect here
even one minute into a 6h interval; a missing row is a no-op
"""
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
mock_prisma = MagicMock()
# Set up config with interval_hours=24 and force_reload=True to trigger reload
mock_config = MagicMock()
mock_config.param_value = {"interval_hours": 24, "force_reload": True}
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config)
# _check_and_reload_model_cost_map now reads through get_generic_data.
mock_prisma.get_generic_data = AsyncMock(return_value=mock_config)
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None)
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None)
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
mock_prisma.db.litellm_config.update_many.assert_not_called()
assert proxy_config.last_model_cost_map_reload is None
frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc)
requested_at = datetime(2024, 1, 1, 6, 59, 30, tzinfo=timezone.utc)
mock_prisma.db.litellm_config.find_unique = AsyncMock(
return_value=_reload_schedule_row(
{"interval_hours": 6},
reload_requested_at=requested_at,
last_run_at=requested_at,
)
)
proxy_config.last_model_cost_map_reload = frozen_now - timedelta(minutes=1)
original_model_cost = litellm.model_cost.copy()
try:
with patch(
"litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map"
) as mock_get_map:
mock_get_map.return_value = {"gpt-4": {"input_cost_per_token": 0.001}}
with (
patch(
"litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map"
) as mock_get_map,
patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now),
):
mock_get_map.return_value = {
"gpt-3.5-turbo": {"input_cost_per_token": 0.001}
}
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
# Verify the upsert update branch preserves interval_hours
mock_prisma.db.litellm_config.upsert.assert_called()
call_args = mock_prisma.db.litellm_config.upsert.call_args
param_value_json = call_args[1]["data"]["update"]["param_value"]
param_value_dict = json.loads(param_value_json)
assert param_value_dict["force_reload"] == False
assert param_value_dict["interval_hours"] == 24, (
"interval_hours must be preserved in the update branch; "
"dropping it causes the schedule to self-destruct"
assert litellm.model_cost["gpt-3.5-turbo"] == {
"input_cost_per_token": 0.001
}
assert proxy_config.last_model_cost_map_reload == frozen_now
assert mock_prisma.db.litellm_config.update_many.call_args[1] == {
"data": {"last_run_at": frozen_now},
"where": {"param_name": "model_cost_map_reload_config"},
}
mock_prisma.db.litellm_config.upsert.assert_not_called()
finally:
litellm.model_cost = original_model_cost
_invalidate_model_cost_lowercase_map()
def test_distributed_reload_ignores_already_applied_request(self):
"""
A reload request older than this pod's own last reload was already applied here, so
it must not re-trigger on every job tick for the rest of the interval
"""
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
mock_prisma = MagicMock()
mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None)
mock_prisma.db.litellm_config.find_unique = AsyncMock(
return_value=_reload_schedule_row(
{"interval_hours": 6},
reload_requested_at=datetime(2024, 1, 1, 6, 0, tzinfo=timezone.utc),
last_run_at=datetime(2024, 1, 1, 6, 0, tzinfo=timezone.utc),
)
)
frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc)
pod_last_reload = frozen_now - timedelta(minutes=1)
proxy_config.last_model_cost_map_reload = pod_last_reload
original_model_cost = litellm.model_cost.copy()
try:
with (
patch(
"litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map"
) as mock_get_map,
patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now),
):
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
mock_get_map.assert_not_called()
mock_prisma.db.litellm_config.update_many.assert_not_called()
assert proxy_config.last_model_cost_map_reload == pod_last_reload
finally:
litellm.model_cost = original_model_cost
_invalidate_model_cost_lowercase_map()
def test_periodic_reload_uses_pod_local_last_reload(self):
"""
Each pod decides from its own last reload, so a pod that never reloaded refreshes even
when the shared row was just stamped, and stays put while inside the interval
"""
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
mock_prisma = MagicMock()
mock_prisma.db.litellm_config.find_unique = AsyncMock(
return_value=_reload_schedule_row(
{"interval_hours": 6},
last_run_at=datetime(2024, 1, 1, 6, 59, tzinfo=timezone.utc),
)
)
mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None)
frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc)
original_model_cost = litellm.model_cost.copy()
try:
with (
patch(
"litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map"
) as mock_get_map,
patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now),
):
mock_get_map.return_value = {"gpt-4-test": {"input_cost_per_token": 0.5}}
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
assert litellm.model_cost["gpt-4-test"] == {"input_cost_per_token": 0.5}
assert proxy_config.last_model_cost_map_reload == frozen_now
assert mock_prisma.db.litellm_config.update_many.call_args[1][
"data"
] == {"last_run_at": frozen_now}
mock_get_map.reset_mock()
mock_prisma.db.litellm_config.update_many.reset_mock()
proxy_config.last_model_cost_map_reload = frozen_now - timedelta(hours=1)
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
mock_get_map.assert_not_called()
mock_prisma.db.litellm_config.update_many.assert_not_called()
assert proxy_config.last_model_cost_map_reload == frozen_now - timedelta(
hours=1
)
finally:
litellm.model_cost = original_model_cost
_invalidate_model_cost_lowercase_map()
def test_manual_reload_preserves_interval_hours(self):
"""Test that manual reload via /reload/model_cost_map preserves existing interval_hours.
def test_distributed_reload_stamps_last_run_without_creating_row(self):
"""
Regression: the job's write carries neither param_value (which would clobber the
admin-configured interval) nor a create branch (which would resurrect a schedule
a concurrent cancel just deleted)
"""
from litellm.proxy.proxy_server import ProxyConfig
Regression test: the manual reload endpoint was overwriting param_value with
only force_reload=True, dropping any existing interval_hours schedule.
proxy_config = ProxyConfig()
mock_prisma = MagicMock()
mock_prisma.db.litellm_config.find_unique = AsyncMock(
return_value=_reload_schedule_row({"interval_hours": 24})
)
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None)
frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc)
original_model_cost = litellm.model_cost.copy()
try:
with (
patch(
"litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map"
) as mock_get_map,
patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now),
):
mock_get_map.return_value = {"gpt-4": {"input_cost_per_token": 0.001}}
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
assert mock_prisma.db.litellm_config.update_many.call_args[1] == {
"data": {"last_run_at": frozen_now},
"where": {"param_name": "model_cost_map_reload_config"},
}
mock_prisma.db.litellm_config.upsert.assert_not_called()
mock_prisma.db.litellm_config.create.assert_not_called()
finally:
litellm.model_cost = original_model_cost
_invalidate_model_cost_lowercase_map()
def test_manual_reload_preserves_interval_hours(self):
"""
Regression: manual reload owns only the run columns, so it never reads or rewrites
param_value and cannot destroy an existing schedule
"""
from litellm.proxy._types import LitellmUserRoles
from litellm.proxy.proxy_server import cleanup_router_config_variables
@ -4036,88 +4171,76 @@ class TestPriceDataReloadIntegration:
mock_auth.user_role = LitellmUserRoles.PROXY_ADMIN
app.dependency_overrides[user_api_key_auth] = lambda: mock_auth
client = TestClient(app)
frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc)
original_model_cost = litellm.model_cost.copy()
try:
with patch(
"litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map"
) as mock_get_map:
with (
patch(
"litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map"
) as mock_get_map,
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now),
):
mock_get_map.return_value = {"gpt-4": {"input_cost_per_token": 0.001}}
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
# Simulate existing config with a schedule
mock_existing = MagicMock()
mock_existing.param_value = {
"interval_hours": 12,
"force_reload": False,
}
mock_prisma.db.litellm_config.find_unique = AsyncMock(
return_value=mock_existing
)
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
response = client.post("/reload/model_cost_map")
assert response.status_code == 200
response = client.post("/reload/model_cost_map")
assert response.status_code == 200
# Verify interval_hours was preserved in the upsert
mock_prisma.db.litellm_config.upsert.assert_called()
call_args = mock_prisma.db.litellm_config.upsert.call_args
param_value_json = call_args[1]["data"]["update"]["param_value"]
param_value_dict = json.loads(param_value_json)
assert param_value_dict["force_reload"] == True
assert param_value_dict["interval_hours"] == 12, (
"interval_hours must be preserved when manual reload sets force_reload; "
"dropping it destroys any existing schedule"
)
mock_prisma.db.litellm_config.find_unique.assert_not_called()
call_args = mock_prisma.db.litellm_config.upsert.call_args
assert call_args[1]["data"]["update"] == {
"last_run_at": frozen_now,
"reload_requested_at": frozen_now,
}
assert call_args[1]["data"]["create"] == {
"param_name": "model_cost_map_reload_config",
"last_run_at": frozen_now,
"reload_requested_at": frozen_now,
}
assert (
proxy_server_module.proxy_config.last_model_cost_map_reload
== frozen_now
)
finally:
litellm.model_cost = original_model_cost
_invalidate_model_cost_lowercase_map()
def test_anthropic_beta_headers_reload_preserves_interval_hours(self):
"""Test that _check_and_reload_anthropic_beta_headers preserves interval_hours after reload.
Regression test: the update branch of the upsert was dropping interval_hours,
identical to the model cost map bug.
"""
"""Same column-only write contract as the cost map job, for the beta headers schedule"""
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
mock_prisma = MagicMock()
# Set up config with interval_hours=12 and force_reload=True to trigger reload
mock_config = MagicMock()
mock_config.param_value = {"interval_hours": 12, "force_reload": True}
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config)
# _check_and_reload_anthropic_beta_headers now reads through get_generic_data.
mock_prisma.get_generic_data = AsyncMock(return_value=mock_config)
mock_prisma.db.litellm_config.find_unique = AsyncMock(
return_value=_reload_schedule_row({"interval_hours": 12})
)
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None)
frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc)
with patch(
"litellm.anthropic_beta_headers_manager.reload_beta_headers_config"
) as mock_reload:
with (
patch(
"litellm.anthropic_beta_headers_manager.reload_beta_headers_config"
) as mock_reload,
patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now),
):
mock_reload.return_value = {"anthropic": {"beta_header": "test-value"}}
asyncio.run(
proxy_config._check_and_reload_anthropic_beta_headers(mock_prisma)
)
# Verify the upsert update branch preserves interval_hours
mock_prisma.db.litellm_config.upsert.assert_called()
call_args = mock_prisma.db.litellm_config.upsert.call_args
param_value_json = call_args[1]["data"]["update"]["param_value"]
param_value_dict = json.loads(param_value_json)
assert param_value_dict["force_reload"] == False
assert param_value_dict["interval_hours"] == 12, (
"interval_hours must be preserved in the update branch; "
"dropping it causes the schedule to self-destruct"
)
assert proxy_config.last_anthropic_beta_headers_reload == frozen_now
assert mock_prisma.db.litellm_config.update_many.call_args[1] == {
"data": {"last_run_at": frozen_now},
"where": {"param_name": "anthropic_beta_headers_reload_config"},
}
mock_prisma.db.litellm_config.upsert.assert_not_called()
def test_anthropic_beta_headers_manual_reload_preserves_interval_hours(self):
"""Test that manual reload via /reload/anthropic_beta_headers preserves existing interval_hours.
Regression test: the manual reload endpoint was overwriting param_value with
only force_reload=True, dropping any existing interval_hours schedule.
"""
"""Manual beta headers reload stamps the run columns only, leaving the schedule intact"""
from litellm.proxy._types import LitellmUserRoles
from litellm.proxy.proxy_server import cleanup_router_config_variables
@ -4130,34 +4253,36 @@ class TestPriceDataReloadIntegration:
mock_auth.user_role = LitellmUserRoles.PROXY_ADMIN
app.dependency_overrides[user_api_key_auth] = lambda: mock_auth
client = TestClient(app)
frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc)
with patch(
"litellm.anthropic_beta_headers_manager.reload_beta_headers_config"
) as mock_reload:
with (
patch(
"litellm.anthropic_beta_headers_manager.reload_beta_headers_config"
) as mock_reload,
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now),
):
mock_reload.return_value = {"anthropic": {"beta_header": "test-value"}}
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
# Simulate existing config with a schedule
mock_existing = MagicMock()
mock_existing.param_value = {"interval_hours": 8, "force_reload": False}
mock_prisma.db.litellm_config.find_unique = AsyncMock(
return_value=mock_existing
)
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
response = client.post("/reload/anthropic_beta_headers")
assert response.status_code == 200
response = client.post("/reload/anthropic_beta_headers")
assert response.status_code == 200
# Verify interval_hours was preserved in the upsert
mock_prisma.db.litellm_config.upsert.assert_called()
call_args = mock_prisma.db.litellm_config.upsert.call_args
param_value_json = call_args[1]["data"]["update"]["param_value"]
param_value_dict = json.loads(param_value_json)
assert param_value_dict["force_reload"] == True
assert param_value_dict["interval_hours"] == 8, (
"interval_hours must be preserved when manual reload sets force_reload; "
"dropping it destroys any existing schedule"
)
mock_prisma.db.litellm_config.find_unique.assert_not_called()
call_args = mock_prisma.db.litellm_config.upsert.call_args
assert call_args[1]["data"]["update"] == {
"last_run_at": frozen_now,
"reload_requested_at": frozen_now,
}
assert call_args[1]["data"]["create"] == {
"param_name": "anthropic_beta_headers_reload_config",
"last_run_at": frozen_now,
"reload_requested_at": frozen_now,
}
assert (
proxy_server_module.proxy_config.last_anthropic_beta_headers_reload
== frozen_now
)
def test_config_file_parsing(self):
"""Test parsing of config file with reload settings"""
@ -4187,64 +4312,6 @@ model_list:
assert "model_list" in config
assert len(config["model_list"]) == 2
def test_database_config_storage(self):
"""Test that configuration is properly stored in database"""
# Mock prisma client
mock_prisma = MagicMock()
# Test the database upsert call that would be made by the schedule endpoint
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
# Simulate the database call that the schedule endpoint would make
asyncio.run(
mock_prisma.db.litellm_config.upsert(
where={"param_name": "model_cost_map_reload_config"},
data={
"create": {
"param_name": "model_cost_map_reload_config",
"param_value": {"interval_hours": 6, "force_reload": False},
},
"update": {
"param_value": {"interval_hours": 6, "force_reload": False}
},
},
)
)
# Verify database upsert was called with correct data
mock_prisma.db.litellm_config.upsert.assert_called_once()
call_args = mock_prisma.db.litellm_config.upsert.call_args
assert call_args[1]["where"]["param_name"] == "model_cost_map_reload_config"
assert call_args[1]["data"]["create"]["param_value"]["interval_hours"] == 6
assert call_args[1]["data"]["create"]["param_value"]["force_reload"] == False
def test_manual_reload_force_flag(self):
"""Test that manual reload sets force flag correctly"""
# Mock prisma client
mock_prisma = MagicMock()
# Test the database upsert call that would be made by the manual reload endpoint
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
# Simulate the database call that the manual reload endpoint would make
asyncio.run(
mock_prisma.db.litellm_config.upsert(
where={"param_name": "model_cost_map_reload_config"},
data={
"create": {
"param_name": "model_cost_map_reload_config",
"param_value": {"interval_hours": None, "force_reload": True},
},
"update": {"param_value": {"force_reload": True}},
},
)
)
# Verify force_reload flag was set
mock_prisma.db.litellm_config.upsert.assert_called_once()
call_args = mock_prisma.db.litellm_config.upsert.call_args
assert call_args[1]["data"]["update"]["param_value"]["force_reload"] == True
@pytest.mark.asyncio
async def test_add_router_settings_from_db_config_merge_logic():