mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(proxy): persist periodic reload schedule state so status survives restarts and fires without store_model_in_db (#35165)
* fix(proxy): persist periodic reload schedule state so status survives restarts and fires without store_model_in_db The model cost map and Anthropic beta headers reload schedules kept their last-run time in a per-pod module global, so GET /schedule/*/status reported last_run null after any restart and the Admin UI showed the reload as never having run. The reload check also only ran from the add_deployment job, which is registered only when store_model_in_db is true, so config-file deployments stored a schedule that never fired. Persist last_run_at and reload_requested_at as dedicated columns on LiteLLM_Config, owned by the reload job and manual reload endpoints, while the schedule endpoints own the param_value JSON (interval_hours); no writer can clobber another's fields. Serve status entirely from the row. Register the check as its own periodic_reload_job outside the store_model_in_db gate. Replace the force_reload boolean with a reload_requested_at timestamp each pod compares against its own in-memory last reload, so a manual reload reaches every pod exactly once instead of being cleared by the first poller. Run the blocking fetches via asyncio.to_thread, and stamp last_run_at with update_many so a schedule cancelled mid-poll is not resurrected. * fix(proxy): compare reload requests against pod data age seeded at boot A pod that had never reloaded kept its in-memory clock at None, and with no interval configured nothing ever set it, so every manual reload request was ignored by every pod except the one serving the click (Greptile P1 on the previous commit). Seed the per-pod timestamp at boot as the time its data was loaded and reload whenever a request or the interval is older than that, which also removes both None special cases from the due predicate. A schedule whose row has no last_run_at fires on the next tick so the first run does not wait a full interval. * fix(proxy): scope reload persistence to the model cost map and seed the pod clock from the actual load time Revert the Anthropic beta headers reload path to its previous JSON-flag implementation so this PR only changes the price data reload; the beta headers path keeps working exactly as before and can migrate to the shared module in a follow-up. The unused columns on its config row are inert. Seed model_cost_map_loaded_at from the timestamp get_model_cost_map records at the actual import-time fetch instead of ProxyConfig construction time, closing the startup window where a manual reload request stamped between the fetch and the constructor compared as older than the pod's data and was skipped (Greptile P1 on the previous commit). * refactor(proxy): drop the legacy force_reload backfill from the reload tracking migration The backfill only carried over a manual reload clicked in the seconds before an upgrade, and every upgrade restarts the pods, which re-fetch the cost map at import and so already deliver what that request asked for. Removing it makes the migration schema-only, so prisma db push and prisma migrate deploy leave the database in the same state instead of diverging on a data statement that only one of them runs. * fix(proxy): stamp reload timestamps at the precision they are stored at Postgres stores these columns as TIMESTAMP(3) while Python stamps microseconds, so a pod comparing its in-memory clock against the persisted copy of the same instant read as newer and skipped the reload request it had just recorded. Truncate every stamp to milliseconds at the source, and floor the boot seed the same way, so the in-memory value and its persisted copy compare exactly. * fix(proxy): identify manual reloads by revision instead of comparing timestamps Comparing a request timestamp against each pod's data age made correctness depend on clock resolution: Postgres stores TIMESTAMP(3) while Python stamps microseconds, and two events inside the same millisecond are indistinguishable no matter how the comparison is written. Replace reload_requested_at with a reload_revision counter the manual reload endpoint increments atomically in the database. Each pod records the revision it last applied and reloads whenever the row's differs, so a request reaches every pod exactly once regardless of clock skew or precision, and concurrent requests publish distinct revisions instead of overwriting one another. A pod adopts the current revision on its first poll, since data it loaded at boot already satisfies any earlier request. Interval reloads still key off the pod's own data age, where hour scale comparisons make precision irrelevant. * fix(proxy): seed the applied reload revision at startup A pod adopted whatever revision it found on its first poll, so a manual reload published while the pod was starting was marked applied without ever being served and the pod kept the prices it fetched at import. Read the row once at startup instead, right after that fetch, and treat a missing row as revision 0 * style(tests): revert incidental reformatting of test_proxy_server.py An earlier ruff format run reflowed the whole file from its 88-column formatting, adding ~1150 lines of churn unrelated to this PR. Replay only the real test changes onto the original formatting * fix(proxy): serve an outstanding reload request on a booting pod Seeding the applied revision at startup left a window: a manual reload published after the import-time cost map fetch but before startup read the row was marked applied without ever being fetched, stranding that pod on stale prices when no interval was configured. A pod now starts unapplied and serves any outstanding request on its first poll, which costs one redundant fetch per boot and removes the window along with the seeding step * fix(proxy): accept a reload interval still encoded as JSON text param_value is written with safe_dumps, and a raw row read can return it decoded or as a string depending on the driver. Strict validation rejected the string, so the schedule read as disabled and an admin's configured reloads silently stopped. Mirrors the guard ConfigRepository.get_param already carries for the same column * fix(proxy): cancel a reload schedule without resetting the revision * fix(proxy): null the interval in JSON so cancelling keeps the revision prisma rejects a null literal for a Json? column, so update_many writes an interval-less object instead. The fake config table now rejects the same input the database does, which is what the live run caught and the mock did not. Also records the run before adopting the revision, so a failed status write leaves the request unserved for the next poll rather than reporting a run that never landed. * fix(ui): match the CI-generated user_role union order in schema.d.ts
This commit is contained in:
ryan-crabbe-berri
GitHub
parent
d1ca826ff6
commit
9ea5cfce0e
12 changed files with 1247 additions and 492 deletions
|
|
@ -0,0 +1,3 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_Config" ADD COLUMN IF NOT EXISTS "last_run_at" TIMESTAMP(3),
|
||||
ADD COLUMN IF NOT EXISTS "reload_revision" BIGINT NOT NULL DEFAULT 0;
|
||||
|
|
@ -601,6 +601,8 @@ model LiteLLM_TagTable {
|
|||
model LiteLLM_Config {
|
||||
param_name String @id
|
||||
param_value Json?
|
||||
last_run_at DateTime?
|
||||
reload_revision BigInt @default(0)
|
||||
}
|
||||
|
||||
// View spend, model, api_key per request
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ import os
|
|||
import random
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from importlib.resources import files
|
||||
from typing import Final, Protocol
|
||||
|
||||
|
|
@ -325,6 +326,7 @@ class ModelCostMapSourceInfo:
|
|||
url: str | None = None
|
||||
is_env_forced: bool = False
|
||||
fallback_reason: str | None = None
|
||||
loaded_at: "datetime | None" = None
|
||||
|
||||
|
||||
# Module-level singleton tracking the source of the current cost map
|
||||
|
|
@ -349,6 +351,11 @@ def get_model_cost_map_source_info() -> dict:
|
|||
}
|
||||
|
||||
|
||||
def get_model_cost_map_loaded_at() -> "datetime | None":
|
||||
"""When this process last loaded its cost map, stamped at the start of every load"""
|
||||
return _cost_map_source_info.loaded_at
|
||||
|
||||
|
||||
def _expand_model_aliases(model_cost: dict) -> dict:
|
||||
"""
|
||||
Expand ``aliases`` lists in model cost entries into top-level entries.
|
||||
|
|
@ -428,6 +435,7 @@ def get_model_cost_map(url: str) -> dict:
|
|||
The full backup dict is only parsed when it must be *returned* as a
|
||||
fallback — it is never held in memory long-term.
|
||||
"""
|
||||
_cost_map_source_info.loaded_at = datetime.now(timezone.utc)
|
||||
# Note: can't use get_secret_bool here — this runs during litellm.__init__
|
||||
# before litellm._key_management_settings is set.
|
||||
if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true":
|
||||
|
|
|
|||
231
litellm/proxy/common_utils/periodic_reload_schedule.py
Normal file
231
litellm/proxy/common_utils/periodic_reload_schedule.py
Normal file
|
|
@ -0,0 +1,231 @@
|
|||
"""
|
||||
Persistence for the admin-configured periodic model cost map reload schedule stored in
|
||||
``LiteLLM_Config``.
|
||||
|
||||
Field ownership is split by writer so concurrent writers never overwrite each other:
|
||||
the schedule endpoints own the ``param_value`` JSON (``interval_hours``), while the
|
||||
reload job and the manual reload endpoints own the dedicated ``last_run_at`` /
|
||||
``reload_revision`` columns. ``last_run_at`` lives in the row rather than process memory
|
||||
so the Admin UI still reports the last execution after a restart and across pods.
|
||||
``reload_revision`` is a monotonic counter a manual reload increments; each pod records
|
||||
the revision it last applied and reloads whenever the row's differs, so a request reaches
|
||||
every pod exactly once without any pod clearing it and without comparing clocks. A booting
|
||||
pod starts at revision 0 rather than adopting the published one, because it cannot know
|
||||
whether that request predates the prices it fetched at import. Interval reloads stay
|
||||
per-pod, driven by when that pod's own copy of the data was loaded.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Protocol,
|
||||
TypedDict,
|
||||
cast, # noqa: TID251 # prisma table access is untyped (PrismaWrapper.__getattr__)
|
||||
)
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.proxy.utils import PrismaClient, evict_config_param
|
||||
from litellm.repositories.config_repository import ConfigRepository
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma.models import LiteLLM_Config
|
||||
|
||||
MODEL_COST_MAP_RELOAD_PARAM_NAME = "model_cost_map_reload_config"
|
||||
|
||||
|
||||
class _RevisionIncrement(TypedDict):
|
||||
increment: int
|
||||
|
||||
|
||||
class _ConfigRowWrite(TypedDict, total=False):
|
||||
param_name: str
|
||||
param_value: str
|
||||
last_run_at: datetime
|
||||
reload_revision: int | _RevisionIncrement
|
||||
|
||||
|
||||
class _ConfigUpsertData(TypedDict):
|
||||
create: _ConfigRowWrite
|
||||
update: _ConfigRowWrite
|
||||
|
||||
|
||||
class _ConfigTable(Protocol):
|
||||
async def find_unique(self, where: Mapping[str, str]) -> "LiteLLM_Config | None": ...
|
||||
|
||||
async def upsert(self, where: Mapping[str, str], data: _ConfigUpsertData) -> "LiteLLM_Config": ...
|
||||
|
||||
async def update_many(self, data: _ConfigRowWrite, where: Mapping[str, str]) -> int: ...
|
||||
|
||||
|
||||
def _config_table(prisma_client: PrismaClient) -> _ConfigTable:
|
||||
return cast(_ConfigTable, ConfigRepository(prisma_client).table) # cast-ok: prisma table is untyped (Any)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ReloadSchedule:
|
||||
interval_hours: int | None = None
|
||||
reload_revision: int = 0
|
||||
last_run_at: datetime | None = None
|
||||
|
||||
|
||||
class ReloadScheduleStatus(TypedDict):
|
||||
scheduled: bool
|
||||
interval_hours: int | None
|
||||
last_run: str | None
|
||||
next_run: str | None
|
||||
|
||||
|
||||
class _IntervalConfig(BaseModel):
|
||||
model_config = ConfigDict(strict=True)
|
||||
|
||||
interval_hours: int | None = None
|
||||
|
||||
|
||||
def utc_now() -> datetime:
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
def _parse_interval_hours(param_value: object) -> int | None:
|
||||
"""``param_value`` is written as serialized JSON, and a raw row read can hand it back
|
||||
either decoded or still as a string depending on the driver, so accept both rather than
|
||||
reading a string as no schedule at all. Mirrors ``ConfigRepository.get_param``"""
|
||||
try:
|
||||
if isinstance(param_value, str):
|
||||
return _IntervalConfig.model_validate_json(param_value).interval_hours
|
||||
return _IntervalConfig.model_validate(param_value).interval_hours
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _as_utc(value: datetime | None) -> datetime | None:
|
||||
if value is None:
|
||||
return None
|
||||
return value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc)
|
||||
|
||||
|
||||
def parse_reload_schedule(row: "LiteLLM_Config") -> ReloadSchedule:
|
||||
return ReloadSchedule(
|
||||
interval_hours=_parse_interval_hours(row.param_value),
|
||||
reload_revision=int(row.reload_revision or 0),
|
||||
last_run_at=_as_utc(row.last_run_at),
|
||||
)
|
||||
|
||||
|
||||
def next_run_at(schedule: ReloadSchedule) -> datetime | None:
|
||||
if schedule.interval_hours is None or schedule.last_run_at is None:
|
||||
return None
|
||||
return schedule.last_run_at + timedelta(hours=schedule.interval_hours)
|
||||
|
||||
|
||||
def reload_schedule_status(schedule: ReloadSchedule | None) -> ReloadScheduleStatus:
|
||||
if schedule is None:
|
||||
return {"scheduled": False, "interval_hours": None, "last_run": None, "next_run": None}
|
||||
next_run = next_run_at(schedule)
|
||||
return {
|
||||
"scheduled": schedule.interval_hours is not None,
|
||||
"interval_hours": schedule.interval_hours,
|
||||
"last_run": schedule.last_run_at.isoformat() if schedule.last_run_at is not None else None,
|
||||
"next_run": next_run.isoformat() if next_run is not None else None,
|
||||
}
|
||||
|
||||
|
||||
def pod_reload_is_due(
|
||||
*,
|
||||
schedule: ReloadSchedule,
|
||||
pod_applied_revision: int,
|
||||
pod_data_loaded_at: datetime,
|
||||
current_time: datetime,
|
||||
description: str,
|
||||
) -> bool:
|
||||
"""
|
||||
Whether this pod should reload now. A revision it has not applied means a manual reload
|
||||
it has not served. A pod starts at revision 0, so it serves any request published before
|
||||
it booted; that costs one redundant fetch per boot and is what keeps a request from being
|
||||
marked applied against data fetched before it. Interval reloads compare against this pod's
|
||||
own data, and a schedule that has never run anywhere fires immediately rather than one
|
||||
interval later
|
||||
"""
|
||||
if schedule.reload_revision != pod_applied_revision:
|
||||
verbose_proxy_logger.info("%s reload triggered by manual reload request", description)
|
||||
return True
|
||||
if schedule.interval_hours is None:
|
||||
return False
|
||||
if schedule.last_run_at is None:
|
||||
verbose_proxy_logger.info("%s reload triggered - schedule has never run", description)
|
||||
return True
|
||||
hours_since_data_loaded = (current_time - pod_data_loaded_at).total_seconds() / 3600
|
||||
if hours_since_data_loaded < schedule.interval_hours:
|
||||
return False
|
||||
verbose_proxy_logger.info(
|
||||
"%s reload triggered by interval. Hours since data loaded: %.2f, Interval: %s",
|
||||
description,
|
||||
hours_since_data_loaded,
|
||||
schedule.interval_hours,
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
async def read_reload_schedule(prisma_client: PrismaClient, param_name: str) -> ReloadSchedule | None:
|
||||
row = await _config_table(prisma_client).find_unique(where={"param_name": param_name})
|
||||
if row is None:
|
||||
return None
|
||||
return parse_reload_schedule(row)
|
||||
|
||||
|
||||
async def write_reload_interval(prisma_client: PrismaClient, param_name: str, interval_hours: int) -> None:
|
||||
"""Admin-owned write: replaces ``param_value`` without touching the job-owned columns"""
|
||||
param_value = safe_dumps({"interval_hours": interval_hours})
|
||||
await _config_table(prisma_client).upsert(
|
||||
where={"param_name": param_name},
|
||||
data={
|
||||
"create": {"param_name": param_name, "param_value": param_value},
|
||||
"update": {"param_value": param_value},
|
||||
},
|
||||
)
|
||||
await evict_config_param(param_name)
|
||||
|
||||
|
||||
async def clear_reload_interval(prisma_client: PrismaClient, param_name: str) -> None:
|
||||
"""Admin-owned write: drops the schedule but keeps the row, because the revision counter
|
||||
identifies a request rather than ordering one and so can never reuse a number. Deleting
|
||||
the row restarts it, and a reissued revision matches what pods already applied, so their
|
||||
next manual reload is silently skipped. The interval is nulled inside the JSON rather
|
||||
than by nulling the column, which prisma rejects for a ``Json?`` field"""
|
||||
await _config_table(prisma_client).update_many(
|
||||
data={"param_value": safe_dumps({"interval_hours": None})},
|
||||
where={"param_name": param_name},
|
||||
)
|
||||
await evict_config_param(param_name)
|
||||
|
||||
|
||||
async def record_reload_run(prisma_client: PrismaClient, param_name: str, ran_at: datetime) -> None:
|
||||
"""Job-owned write after this pod reloaded: stamps the shared last run only if the row
|
||||
still exists, so a schedule deleted mid-poll is not resurrected"""
|
||||
await _config_table(prisma_client).update_many(
|
||||
data={"last_run_at": ran_at},
|
||||
where={"param_name": param_name},
|
||||
)
|
||||
await evict_config_param(param_name)
|
||||
|
||||
|
||||
async def record_manual_reload(prisma_client: PrismaClient, param_name: str, ran_at: datetime) -> int:
|
||||
"""
|
||||
After a manual in-pod reload: stamp the shared last run and bump the revision every other
|
||||
pod compares against. The increment is atomic, so concurrent requests each publish a
|
||||
distinct revision instead of overwriting one another. Returns the published revision so
|
||||
the serving pod can adopt it rather than reloading again on its next poll
|
||||
"""
|
||||
row = await _config_table(prisma_client).upsert(
|
||||
where={"param_name": param_name},
|
||||
data={
|
||||
"create": {"param_name": param_name, "last_run_at": ran_at, "reload_revision": 1},
|
||||
"update": {"last_run_at": ran_at, "reload_revision": {"increment": 1}},
|
||||
},
|
||||
)
|
||||
await evict_config_param(param_name)
|
||||
return int(row.reload_revision)
|
||||
|
|
@ -319,6 +319,17 @@ from litellm.proxy.common_utils.model_listing_utils import TeamModelNameTranslat
|
|||
from litellm.proxy.common_utils.openai_endpoint_utils import (
|
||||
remove_sensitive_info_from_deployment,
|
||||
)
|
||||
from litellm.proxy.common_utils.periodic_reload_schedule import (
|
||||
MODEL_COST_MAP_RELOAD_PARAM_NAME,
|
||||
clear_reload_interval,
|
||||
pod_reload_is_due,
|
||||
read_reload_schedule,
|
||||
record_manual_reload,
|
||||
record_reload_run,
|
||||
reload_schedule_status,
|
||||
utc_now,
|
||||
write_reload_interval,
|
||||
)
|
||||
from litellm.proxy.common_utils.proxy_state import ProxyState
|
||||
from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob
|
||||
from litellm.proxy.common_utils.swagger_utils import ERROR_RESPONSES
|
||||
|
|
@ -2062,9 +2073,7 @@ async_result: Final = None
|
|||
celery_app_conn: Final = None
|
||||
celery_fn: Final = None # Redis Queue for handling requests
|
||||
|
||||
# Global variables for model cost map reload scheduling
|
||||
scheduler = None
|
||||
last_model_cost_map_reload = None
|
||||
|
||||
# Global variable for anthropic beta headers reload scheduling
|
||||
last_anthropic_beta_headers_reload = None
|
||||
|
|
@ -3839,6 +3848,17 @@ def resolve_complexity_router_plugins(
|
|||
)
|
||||
|
||||
|
||||
def _swap_in_model_cost_map(new_model_cost_map: dict) -> int:
|
||||
"""Adopt a freshly fetched cost map into this process's litellm state, return the model count"""
|
||||
litellm.model_cost = new_model_cost_map
|
||||
# Invalidate case-insensitive lookup map since model_cost was replaced
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
# Repopulate provider model sets (e.g. litellm.anthropic_models) so that
|
||||
# wildcard patterns like "anthropic/*" include any newly added models.
|
||||
litellm.add_known_models(model_cost_map=new_model_cost_map)
|
||||
return len(new_model_cost_map) if new_model_cost_map else 0
|
||||
|
||||
|
||||
class ProxyConfig:
|
||||
"""
|
||||
Abstraction class on top of config loading/updating logic. Gives us one place to control all config updating logic.
|
||||
|
|
@ -3850,6 +3870,15 @@ class ProxyConfig:
|
|||
self._last_hashicorp_vault_config: dict[str, Any] | None = None
|
||||
self.worker_registry: list[WorkerRegistryEntry] = []
|
||||
self.config_sync_subscriber: ConfigSyncSubscriber | None = None
|
||||
from litellm.litellm_core_utils.get_model_cost_map import (
|
||||
get_model_cost_map_loaded_at,
|
||||
)
|
||||
|
||||
self.model_cost_map_loaded_at: datetime = get_model_cost_map_loaded_at() or utc_now()
|
||||
# Starts unapplied rather than adopting the published revision: this pod cannot tell
|
||||
# whether an existing request predates the prices it just fetched, and re-serving one
|
||||
# costs a single fetch where skipping one leaves it priced wrong indefinitely
|
||||
self.model_cost_map_applied_revision: int = 0
|
||||
|
||||
def is_yaml(self, config_file_path: str) -> bool:
|
||||
if not os.path.isfile(config_file_path):
|
||||
|
|
@ -6242,7 +6271,6 @@ class ProxyConfig:
|
|||
"router_settings",
|
||||
"litellm_settings",
|
||||
"environment_variables",
|
||||
"model_cost_map_reload_config",
|
||||
"anthropic_beta_headers_reload_config",
|
||||
],
|
||||
)
|
||||
|
|
@ -6340,9 +6368,6 @@ class ProxyConfig:
|
|||
if self._should_load_db_object(object_type="tools"):
|
||||
await self._init_tool_policy_in_db(prisma_client=prisma_client)
|
||||
|
||||
if self._should_load_db_object(object_type="model_cost_map"):
|
||||
await self._check_and_reload_model_cost_map(prisma_client=prisma_client)
|
||||
|
||||
if self._should_load_db_object(object_type="anthropic_beta_headers"):
|
||||
await self._check_and_reload_anthropic_beta_headers(prisma_client=prisma_client)
|
||||
|
||||
|
|
@ -6504,111 +6529,60 @@ class ProxyConfig:
|
|||
str(e),
|
||||
)
|
||||
|
||||
async def check_periodic_reloads(self, prisma_client: PrismaClient):
|
||||
"""
|
||||
Run the admin-configured periodic model cost map reload.
|
||||
|
||||
Scheduled on its own job so a schedule configured from the Admin UI fires whether
|
||||
or not `store_model_in_db` is enabled.
|
||||
"""
|
||||
if self._should_load_db_object(object_type="model_cost_map"):
|
||||
await self._check_and_reload_model_cost_map(prisma_client=prisma_client)
|
||||
|
||||
async def _check_and_reload_model_cost_map(self, prisma_client: PrismaClient):
|
||||
"""
|
||||
Check if model cost map needs to be reloaded based on database configuration.
|
||||
This function runs every 10 seconds as part of _init_non_llm_objects_in_db.
|
||||
Runs on the periodic reload job, independently of `store_model_in_db`.
|
||||
"""
|
||||
try:
|
||||
# Get model cost map reload configuration from database
|
||||
config_record: Final = await get_config_param(prisma_client, "model_cost_map_reload_config")
|
||||
schedule = await read_reload_schedule(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME)
|
||||
if schedule is None:
|
||||
return
|
||||
|
||||
if config_record is None or config_record.param_value is None:
|
||||
return # No configuration found, skip reload
|
||||
current_time = utc_now()
|
||||
is_due = pod_reload_is_due(
|
||||
schedule=schedule,
|
||||
pod_applied_revision=self.model_cost_map_applied_revision,
|
||||
pod_data_loaded_at=self.model_cost_map_loaded_at,
|
||||
current_time=current_time,
|
||||
description="Model cost map",
|
||||
)
|
||||
if not is_due:
|
||||
return
|
||||
|
||||
config: Final = config_record.param_value
|
||||
interval_hours: Final = config.get("interval_hours")
|
||||
force_reload: Final = config.get("force_reload", False)
|
||||
from litellm.litellm_core_utils.get_model_cost_map import (
|
||||
ModelCostMapReloadUnavailable,
|
||||
refetch_model_cost_map,
|
||||
)
|
||||
|
||||
if interval_hours is None and force_reload is False:
|
||||
return # No interval configured, skip reload
|
||||
|
||||
current_time: Final = datetime.utcnow()
|
||||
|
||||
# Check if we need to reload based on interval or force reload
|
||||
should_reload = False
|
||||
|
||||
if force_reload:
|
||||
should_reload = True
|
||||
verbose_proxy_logger.info("Model cost map reload triggered by force reload flag")
|
||||
elif interval_hours is not None:
|
||||
# Use pod's in-memory last reload time
|
||||
global last_model_cost_map_reload
|
||||
if last_model_cost_map_reload is not None:
|
||||
try:
|
||||
last_reload_time: Final = datetime.fromisoformat(last_model_cost_map_reload)
|
||||
time_since_last_reload: Final = current_time - last_reload_time
|
||||
hours_since_last_reload: Final = time_since_last_reload.total_seconds() / 3600
|
||||
|
||||
if hours_since_last_reload >= interval_hours:
|
||||
should_reload = True
|
||||
verbose_proxy_logger.info(
|
||||
f"Model cost map reload triggered by interval. Hours since last reload: {hours_since_last_reload:.2f}, Interval: {interval_hours}"
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning("Error parsing last reload time: %s", e)
|
||||
# If we can't parse the last reload time, reload anyway
|
||||
should_reload = True
|
||||
else:
|
||||
# No last reload time recorded, reload now
|
||||
should_reload = True
|
||||
verbose_proxy_logger.info("Model cost map reload triggered - no previous reload time recorded")
|
||||
|
||||
if should_reload:
|
||||
# Perform the reload
|
||||
from litellm.litellm_core_utils.get_model_cost_map import (
|
||||
ModelCostMapReloadUnavailable,
|
||||
refetch_model_cost_map,
|
||||
reload_result = await refetch_model_cost_map(url=litellm.model_cost_map_url)
|
||||
if isinstance(reload_result, ModelCostMapReloadUnavailable):
|
||||
verbose_proxy_logger.warning(
|
||||
"Model cost map reload failed (%s); keeping current pricing data. The revision stays "
|
||||
"unapplied so this pod retries on its next poll",
|
||||
reload_result.reason,
|
||||
)
|
||||
return
|
||||
|
||||
model_cost_map_url: Final = litellm.model_cost_map_url
|
||||
reload_result: Final = await refetch_model_cost_map(url=model_cost_map_url)
|
||||
if isinstance(reload_result, ModelCostMapReloadUnavailable):
|
||||
verbose_proxy_logger.warning(
|
||||
"Model cost map reload failed (%s); keeping current pricing data, will retry on the next config poll",
|
||||
reload_result.reason,
|
||||
)
|
||||
return
|
||||
new_model_cost_map: Final = reload_result.model_cost_map
|
||||
litellm.model_cost = new_model_cost_map
|
||||
# Invalidate case-insensitive lookup map since model_cost was replaced
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
# Repopulate provider model sets (e.g. litellm.anthropic_models) so that
|
||||
# wildcard patterns like "anthropic/*" include any newly added models.
|
||||
litellm.add_known_models(model_cost_map=new_model_cost_map)
|
||||
models_count = _swap_in_model_cost_map(reload_result.model_cost_map)
|
||||
self.model_cost_map_loaded_at = current_time
|
||||
await record_reload_run(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME, current_time)
|
||||
# Adopted last, so neither a failed fetch nor a failed status write is recorded
|
||||
# as served; either way the next poll retries instead of leaving the card
|
||||
# reporting a run that never landed
|
||||
self.model_cost_map_applied_revision = schedule.reload_revision
|
||||
|
||||
# Update pod's in-memory last reload time
|
||||
last_model_cost_map_reload = current_time.isoformat()
|
||||
|
||||
# Clear force reload flag in database
|
||||
await ConfigRepository(prisma_client).table.upsert(
|
||||
where={"param_name": "model_cost_map_reload_config"},
|
||||
data={
|
||||
"create": {
|
||||
"param_name": "model_cost_map_reload_config",
|
||||
"param_value": safe_dumps(
|
||||
{
|
||||
"interval_hours": interval_hours,
|
||||
"force_reload": False,
|
||||
}
|
||||
),
|
||||
},
|
||||
"update": {
|
||||
"param_value": safe_dumps(
|
||||
{
|
||||
"interval_hours": interval_hours,
|
||||
"force_reload": False,
|
||||
}
|
||||
)
|
||||
},
|
||||
},
|
||||
)
|
||||
await evict_config_param("model_cost_map_reload_config")
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
"Model cost map reloaded successfully. Models count: %s",
|
||||
len(new_model_cost_map) if new_model_cost_map else 0,
|
||||
)
|
||||
verbose_proxy_logger.info("Model cost map reloaded successfully. Models count: %s", models_count)
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error in _check_and_reload_model_cost_map: %s", e)
|
||||
|
|
@ -8254,15 +8228,26 @@ class ProxyStartupEvent:
|
|||
except Exception as e:
|
||||
verbose_proxy_logger.debug("Failed to check DB for store_model_in_db: %s", str(e))
|
||||
|
||||
if store_model_in_db is True:
|
||||
config_reload_interval_seconds = proxy_config_reload_interval_seconds
|
||||
if not isinstance(config_reload_interval_seconds, int) or config_reload_interval_seconds <= 0:
|
||||
verbose_proxy_logger.warning(
|
||||
"proxy_config_reload_interval_seconds=%s must be a positive integer; falling back to 30s",
|
||||
config_reload_interval_seconds,
|
||||
)
|
||||
config_reload_interval_seconds = 30
|
||||
config_reload_interval_seconds = proxy_config_reload_interval_seconds
|
||||
if not isinstance(config_reload_interval_seconds, int) or config_reload_interval_seconds <= 0:
|
||||
verbose_proxy_logger.warning(
|
||||
"proxy_config_reload_interval_seconds=%s must be a positive integer; falling back to 30s",
|
||||
config_reload_interval_seconds,
|
||||
)
|
||||
config_reload_interval_seconds = 30
|
||||
|
||||
### PERIODIC RELOADS (model cost map, anthropic beta headers) ###
|
||||
scheduler.add_job(
|
||||
proxy_config.check_periodic_reloads,
|
||||
"interval",
|
||||
seconds=config_reload_interval_seconds,
|
||||
args=[prisma_client],
|
||||
id="periodic_reload_job",
|
||||
replace_existing=True,
|
||||
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
)
|
||||
|
||||
if store_model_in_db is True:
|
||||
# MEMORY LEAK FIX: Increase interval from 10s to 30s minimum
|
||||
# Frequent polling was causing excessive memory allocations
|
||||
scheduler.add_job(
|
||||
|
|
@ -15868,47 +15853,23 @@ async def reload_model_cost_map(
|
|||
refetch_model_cost_map,
|
||||
)
|
||||
|
||||
model_cost_map_url: Final = litellm.model_cost_map_url
|
||||
reload_result: Final = await refetch_model_cost_map(url=model_cost_map_url)
|
||||
reload_result = await refetch_model_cost_map(url=litellm.model_cost_map_url)
|
||||
if isinstance(reload_result, ModelCostMapReloadUnavailable):
|
||||
raise HTTPException(
|
||||
status_code=502,
|
||||
detail=f"Failed to reload model cost map: {reload_result.reason}. Current pricing data was kept.",
|
||||
)
|
||||
new_model_cost_map: Final = reload_result.model_cost_map
|
||||
litellm.model_cost = new_model_cost_map
|
||||
# Invalidate case-insensitive lookup map since model_cost was replaced
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
# Repopulate provider model sets (e.g. litellm.anthropic_models) so that
|
||||
# wildcard patterns like "anthropic/*" include any newly added models.
|
||||
litellm.add_known_models(model_cost_map=new_model_cost_map)
|
||||
|
||||
# Update pod's in-memory last reload time
|
||||
global last_model_cost_map_reload
|
||||
current_time: Final = datetime.utcnow()
|
||||
last_model_cost_map_reload = current_time.isoformat()
|
||||
models_count = _swap_in_model_cost_map(reload_result.model_cost_map)
|
||||
current_time = utc_now()
|
||||
proxy_config.model_cost_map_loaded_at = current_time
|
||||
|
||||
# Set force reload flag in database for other pods, preserving existing interval_hours
|
||||
existing_config: Final = await ConfigRepository(prisma_client).table.find_unique(
|
||||
where={"param_name": "model_cost_map_reload_config"}
|
||||
# Publish a new revision so every other pod reloads on its next poll; this pod has
|
||||
# already served it, so adopt it here rather than reloading again a tick later
|
||||
proxy_config.model_cost_map_applied_revision = await record_manual_reload(
|
||||
prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME, current_time
|
||||
)
|
||||
existing_interval = None
|
||||
if existing_config and existing_config.param_value:
|
||||
existing_interval = existing_config.param_value.get("interval_hours")
|
||||
|
||||
await ConfigRepository(prisma_client).table.upsert(
|
||||
where={"param_name": "model_cost_map_reload_config"},
|
||||
data={
|
||||
"create": {
|
||||
"param_name": "model_cost_map_reload_config",
|
||||
"param_value": safe_dumps({"interval_hours": None, "force_reload": True}),
|
||||
},
|
||||
"update": {"param_value": safe_dumps({"interval_hours": existing_interval, "force_reload": True})},
|
||||
},
|
||||
)
|
||||
await invalidate_config_param("model_cost_map_reload_config")
|
||||
|
||||
models_count: Final = len(new_model_cost_map) if new_model_cost_map else 0
|
||||
verbose_proxy_logger.info("Model cost map reloaded successfully in current pod. Models count: %s", models_count)
|
||||
|
||||
return {
|
||||
|
|
@ -15955,18 +15916,7 @@ async def schedule_model_cost_map_reload(
|
|||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail="Database connection not available")
|
||||
|
||||
# Update database with new reload configuration
|
||||
await ConfigRepository(prisma_client).table.upsert(
|
||||
where={"param_name": "model_cost_map_reload_config"},
|
||||
data={
|
||||
"create": {
|
||||
"param_name": "model_cost_map_reload_config",
|
||||
"param_value": safe_dumps({"interval_hours": hours, "force_reload": False}),
|
||||
},
|
||||
"update": {"param_value": safe_dumps({"interval_hours": hours, "force_reload": False})},
|
||||
},
|
||||
)
|
||||
await invalidate_config_param("model_cost_map_reload_config")
|
||||
await write_reload_interval(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME, hours)
|
||||
|
||||
verbose_proxy_logger.info("Model cost map reload scheduled for every %s hours", hours)
|
||||
|
||||
|
|
@ -15974,7 +15924,7 @@ async def schedule_model_cost_map_reload(
|
|||
"message": f"Model cost map reload scheduled for every {hours} hours",
|
||||
"status": "success",
|
||||
"interval_hours": hours,
|
||||
"timestamp": datetime.utcnow().isoformat(),
|
||||
"timestamp": utc_now().isoformat(),
|
||||
}
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Failed to schedule model cost map reload: %s", e)
|
||||
|
|
@ -16010,16 +15960,14 @@ async def cancel_model_cost_map_reload(
|
|||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail="Database connection not available")
|
||||
|
||||
# Remove reload configuration from database
|
||||
await ConfigRepository(prisma_client).table.delete(where={"param_name": "model_cost_map_reload_config"})
|
||||
await invalidate_config_param("model_cost_map_reload_config")
|
||||
await clear_reload_interval(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME)
|
||||
|
||||
verbose_proxy_logger.info("Model cost map reload schedule cancelled")
|
||||
|
||||
return {
|
||||
"message": "Model cost map reload schedule cancelled",
|
||||
"status": "success",
|
||||
"timestamp": datetime.utcnow().isoformat(),
|
||||
"timestamp": utc_now().isoformat(),
|
||||
}
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Failed to cancel model cost map reload: %s", e)
|
||||
|
|
@ -16048,66 +15996,13 @@ async def get_model_cost_map_reload_status(
|
|||
)
|
||||
|
||||
try:
|
||||
global prisma_client, last_model_cost_map_reload
|
||||
|
||||
verbose_proxy_logger.info("Checking model cost map reload status. Last reload: %s", last_model_cost_map_reload)
|
||||
global prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
verbose_proxy_logger.info("No database connection, returning not scheduled")
|
||||
return {
|
||||
"scheduled": False,
|
||||
"interval_hours": None,
|
||||
"last_run": None,
|
||||
"next_run": None,
|
||||
}
|
||||
return reload_schedule_status(None)
|
||||
|
||||
# Get reload configuration from database
|
||||
config_record: Final = await ConfigRepository(prisma_client).table.find_unique(
|
||||
where={"param_name": "model_cost_map_reload_config"}
|
||||
)
|
||||
|
||||
if config_record is None or config_record.param_value is None:
|
||||
verbose_proxy_logger.info("No model cost map reload configuration found")
|
||||
return {
|
||||
"scheduled": False,
|
||||
"interval_hours": None,
|
||||
"last_run": None,
|
||||
"next_run": None,
|
||||
}
|
||||
|
||||
config: Final = config_record.param_value
|
||||
interval_hours: Final = config.get("interval_hours")
|
||||
|
||||
if interval_hours is None:
|
||||
verbose_proxy_logger.info("No interval configured, returning not scheduled")
|
||||
return {
|
||||
"scheduled": False,
|
||||
"interval_hours": None,
|
||||
"last_run": None,
|
||||
"next_run": None,
|
||||
}
|
||||
|
||||
current_time: Final = datetime.utcnow()
|
||||
next_run = None
|
||||
|
||||
# Use pod's in-memory last reload time
|
||||
if last_model_cost_map_reload is not None:
|
||||
try:
|
||||
last_reload_time: Final = datetime.fromisoformat(last_model_cost_map_reload)
|
||||
time_since_last_reload: Final = current_time - last_reload_time
|
||||
hours_since_last_reload: Final = time_since_last_reload.total_seconds() / 3600
|
||||
|
||||
if hours_since_last_reload < interval_hours:
|
||||
next_run = (last_reload_time + timedelta(hours=interval_hours)).isoformat()
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning("Error parsing last reload time: %s", e)
|
||||
|
||||
return {
|
||||
"scheduled": True,
|
||||
"interval_hours": interval_hours,
|
||||
"last_run": last_model_cost_map_reload,
|
||||
"next_run": next_run,
|
||||
}
|
||||
return reload_schedule_status(await read_reload_schedule(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME))
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Failed to get model cost map reload status: %s", e)
|
||||
raise HTTPException(
|
||||
|
|
|
|||
|
|
@ -601,6 +601,8 @@ model LiteLLM_TagTable {
|
|||
model LiteLLM_Config {
|
||||
param_name String @id
|
||||
param_value Json?
|
||||
last_run_at DateTime?
|
||||
reload_revision BigInt @default(0)
|
||||
}
|
||||
|
||||
// View spend, model, api_key per request
|
||||
|
|
|
|||
|
|
@ -601,6 +601,8 @@ model LiteLLM_TagTable {
|
|||
model LiteLLM_Config {
|
||||
param_name String @id
|
||||
param_value Json?
|
||||
last_run_at DateTime?
|
||||
reload_revision BigInt @default(0)
|
||||
}
|
||||
|
||||
// View spend, model, api_key per request
|
||||
|
|
|
|||
|
|
@ -250,6 +250,27 @@ def test_azure_ai_claude_1m_context_entries(cost_map: dict):
|
|||
assert cost_map[model]["max_input_tokens"] == 200000, model
|
||||
|
||||
|
||||
def test_get_model_cost_map_stamps_loaded_at(monkeypatch):
|
||||
"""The load time feeds each pod's reload-due decision; a load that does not stamp it
|
||||
would make manual reload requests race the proxy's startup"""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from litellm.litellm_core_utils import get_model_cost_map as module
|
||||
|
||||
monkeypatch.setattr(module._cost_map_source_info, "loaded_at", None)
|
||||
monkeypatch.setattr(
|
||||
module.GetModelCostMap,
|
||||
"fetch_remote_model_cost_map",
|
||||
staticmethod(lambda url, timeout=5: _load_root_cost_map()),
|
||||
)
|
||||
|
||||
before = datetime.now(timezone.utc)
|
||||
module.get_model_cost_map(url="https://example.invalid/cost_map.json")
|
||||
loaded_at = module.get_model_cost_map_loaded_at()
|
||||
|
||||
assert loaded_at is not None
|
||||
assert before <= loaded_at <= datetime.now(timezone.utc)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# refetch_model_cost_map: retry/backoff behavior for runtime reloads
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -757,9 +757,13 @@ async def test_evict_config_param_does_not_publish() -> None:
|
|||
def _reload_config_prisma_client() -> MagicMock:
|
||||
config_record = MagicMock()
|
||||
config_record.param_value = {"interval_hours": 6, "force_reload": True}
|
||||
config_record.reload_revision = 0
|
||||
config_record.last_run_at = None
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.get_generic_data = AsyncMock(return_value=config_record)
|
||||
prisma_client.db.litellm_config.upsert = AsyncMock(return_value=None)
|
||||
prisma_client.db.litellm_config.find_unique = AsyncMock(return_value=config_record)
|
||||
prisma_client.db.litellm_config.upsert = AsyncMock(return_value=config_record)
|
||||
prisma_client.db.litellm_config.update_many = AsyncMock(return_value=1)
|
||||
return prisma_client
|
||||
|
||||
|
||||
|
|
@ -790,7 +794,7 @@ async def test_model_cost_map_reload_does_not_publish_config_change() -> None:
|
|||
_invalidate_model_cost_lowercase_map()
|
||||
_set_redis_usage_cache(previous_cache)
|
||||
|
||||
prisma_client.db.litellm_config.upsert.assert_awaited_once()
|
||||
prisma_client.db.litellm_config.update_many.assert_awaited_once()
|
||||
assert client.published == []
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,359 @@
|
|||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.common_utils.periodic_reload_schedule import (
|
||||
ReloadSchedule,
|
||||
clear_reload_interval,
|
||||
next_run_at,
|
||||
parse_reload_schedule,
|
||||
pod_reload_is_due,
|
||||
read_reload_schedule,
|
||||
record_manual_reload,
|
||||
record_reload_run,
|
||||
reload_schedule_status,
|
||||
write_reload_interval,
|
||||
)
|
||||
|
||||
LAST_RUN = datetime(2024, 1, 1, 6, 0, 0, tzinfo=timezone.utc)
|
||||
NOW = datetime(2024, 1, 1, 12, 0, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
def _row(param_value=None, reload_revision=0, last_run_at=None):
|
||||
return SimpleNamespace(
|
||||
param_name="model_cost_map_reload_config",
|
||||
param_value=param_value,
|
||||
reload_revision=reload_revision,
|
||||
last_run_at=last_run_at,
|
||||
)
|
||||
|
||||
|
||||
def _mock_prisma(row=None, upserted_revision=1):
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_config.find_unique = AsyncMock(return_value=row)
|
||||
prisma_client.db.litellm_config.upsert = AsyncMock(return_value=_row(reload_revision=upserted_revision))
|
||||
prisma_client.db.litellm_config.update_many = AsyncMock(return_value=1)
|
||||
return prisma_client
|
||||
|
||||
|
||||
class _FakeConfigTable:
|
||||
"""In-memory stand-in for the prisma LiteLLM_Config actions, faithful on the parts the
|
||||
revision depends on: upsert creates at the column default and applies ``{"increment": 1}``,
|
||||
and delete drops the row along with its counter"""
|
||||
|
||||
def __init__(self):
|
||||
self._rows = {}
|
||||
|
||||
async def find_unique(self, where):
|
||||
return self._rows.get(where["param_name"])
|
||||
|
||||
async def upsert(self, where, data):
|
||||
row = self._rows.get(where["param_name"])
|
||||
if row is None:
|
||||
created = data["create"]
|
||||
row = _row(
|
||||
param_value=created.get("param_value"),
|
||||
reload_revision=created.get("reload_revision", 0),
|
||||
last_run_at=created.get("last_run_at"),
|
||||
)
|
||||
self._rows[where["param_name"]] = row
|
||||
return row
|
||||
self._apply(row, data["update"])
|
||||
return row
|
||||
|
||||
async def update_many(self, data, where):
|
||||
row = self._rows.get(where["param_name"])
|
||||
if row is None:
|
||||
return 0
|
||||
self._apply(row, data)
|
||||
return 1
|
||||
|
||||
async def delete(self, where):
|
||||
return self._rows.pop(where["param_name"], None)
|
||||
|
||||
@staticmethod
|
||||
def _apply(row, data):
|
||||
if "param_value" in data and data["param_value"] is None:
|
||||
raise ValueError("`data.param_value`: A value is required but not set")
|
||||
for field, value in data.items():
|
||||
increment = value.get("increment") if isinstance(value, dict) else None
|
||||
setattr(row, field, getattr(row, field) + increment if increment is not None else value)
|
||||
|
||||
|
||||
def _fake_prisma(table):
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_config = table
|
||||
return prisma_client
|
||||
|
||||
|
||||
def test_parse_reads_interval_from_json_and_state_from_columns():
|
||||
schedule = parse_reload_schedule(_row(param_value={"interval_hours": 6}, reload_revision=7, last_run_at=LAST_RUN))
|
||||
|
||||
assert schedule == ReloadSchedule(interval_hours=6, reload_revision=7, last_run_at=LAST_RUN)
|
||||
|
||||
|
||||
def test_parse_treats_naive_column_timestamps_as_utc():
|
||||
schedule = parse_reload_schedule(_row(last_run_at=LAST_RUN.replace(tzinfo=None)))
|
||||
|
||||
assert schedule.last_run_at == LAST_RUN
|
||||
|
||||
|
||||
def test_parse_defaults_revision_when_the_column_is_null():
|
||||
"""Rows written before the column existed read back as NULL and must not crash the
|
||||
comparison; nobody has applied revision 0, so treating it as 0 is a no-op"""
|
||||
assert parse_reload_schedule(_row(reload_revision=None)).reload_revision == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"param_value",
|
||||
[None, "not-a-dict", '{"interval_hours": "6"}', {}, {"interval_hours": "6"}, {"interval_hours": None}],
|
||||
)
|
||||
def test_parse_tolerates_unusable_param_values(param_value):
|
||||
assert parse_reload_schedule(_row(param_value=param_value)).interval_hours is None
|
||||
|
||||
|
||||
def test_parse_reads_an_interval_still_encoded_as_json_text():
|
||||
"""The interval is written with safe_dumps, so a raw row read can return it either
|
||||
decoded or as a string; reading a string as no schedule would silently stop the
|
||||
reloads an admin configured"""
|
||||
assert parse_reload_schedule(_row(param_value='{"interval_hours": 6}')).interval_hours == 6
|
||||
|
||||
|
||||
def test_parse_ignores_legacy_json_force_reload():
|
||||
"""Rows written by pre-column versions carry force_reload in the JSON; honoring it
|
||||
would re-trigger a reload every poll because nothing clears the JSON copy"""
|
||||
schedule = parse_reload_schedule(_row(param_value={"interval_hours": 6, "force_reload": True}))
|
||||
|
||||
assert schedule == ReloadSchedule(interval_hours=6, reload_revision=0, last_run_at=None)
|
||||
|
||||
|
||||
def test_status_reports_persisted_last_run_and_next_run():
|
||||
status = reload_schedule_status(ReloadSchedule(interval_hours=6, last_run_at=LAST_RUN))
|
||||
|
||||
assert status == {
|
||||
"scheduled": True,
|
||||
"interval_hours": 6,
|
||||
"last_run": "2024-01-01T06:00:00+00:00",
|
||||
"next_run": "2024-01-01T12:00:00+00:00",
|
||||
}
|
||||
|
||||
|
||||
def test_status_without_interval_is_not_scheduled():
|
||||
assert reload_schedule_status(None)["scheduled"] is False
|
||||
assert reload_schedule_status(ReloadSchedule(last_run_at=LAST_RUN)) == {
|
||||
"scheduled": False,
|
||||
"interval_hours": None,
|
||||
"last_run": "2024-01-01T06:00:00+00:00",
|
||||
"next_run": None,
|
||||
}
|
||||
|
||||
|
||||
def test_next_run_needs_both_interval_and_last_run():
|
||||
assert next_run_at(ReloadSchedule(interval_hours=6)) is None
|
||||
assert next_run_at(ReloadSchedule(last_run_at=LAST_RUN)) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"schedule, pod_applied_revision, pod_data_loaded_at, expected",
|
||||
[
|
||||
(ReloadSchedule(reload_revision=4), 3, NOW, True),
|
||||
(ReloadSchedule(interval_hours=6, reload_revision=4, last_run_at=LAST_RUN), 3, NOW, True),
|
||||
(ReloadSchedule(reload_revision=3), 3, LAST_RUN, False),
|
||||
(ReloadSchedule(reload_revision=4), 0, LAST_RUN, True),
|
||||
(ReloadSchedule(), 0, LAST_RUN, False),
|
||||
(ReloadSchedule(interval_hours=6), 0, datetime(2024, 1, 1, 11, 59, tzinfo=timezone.utc), True),
|
||||
(
|
||||
ReloadSchedule(interval_hours=6, last_run_at=datetime(2024, 1, 1, 6, 30, tzinfo=timezone.utc)),
|
||||
0,
|
||||
datetime(2024, 1, 1, 11, 0, tzinfo=timezone.utc),
|
||||
False,
|
||||
),
|
||||
(ReloadSchedule(interval_hours=6, last_run_at=LAST_RUN), 0, LAST_RUN, True),
|
||||
],
|
||||
)
|
||||
def test_pod_reload_is_due(schedule, pod_applied_revision, pod_data_loaded_at, expected):
|
||||
assert (
|
||||
pod_reload_is_due(
|
||||
schedule=schedule,
|
||||
pod_applied_revision=pod_applied_revision,
|
||||
pod_data_loaded_at=pod_data_loaded_at,
|
||||
current_time=NOW,
|
||||
description="test",
|
||||
)
|
||||
is expected
|
||||
)
|
||||
|
||||
|
||||
def test_manual_request_is_identified_not_ordered():
|
||||
"""Comparing revisions for inequality rather than ordering timestamps: a pod applies a
|
||||
request once and is not due again, no matter how the clocks or precisions line up"""
|
||||
unapplied = pod_reload_is_due(
|
||||
schedule=ReloadSchedule(reload_revision=9),
|
||||
pod_applied_revision=8,
|
||||
pod_data_loaded_at=NOW,
|
||||
current_time=NOW,
|
||||
description="test",
|
||||
)
|
||||
applied = pod_reload_is_due(
|
||||
schedule=ReloadSchedule(reload_revision=9),
|
||||
pod_applied_revision=9,
|
||||
pod_data_loaded_at=NOW,
|
||||
current_time=NOW,
|
||||
description="test",
|
||||
)
|
||||
|
||||
assert (unapplied, applied) == (True, False)
|
||||
|
||||
|
||||
def test_booting_pod_serves_a_request_it_cannot_prove_it_already_has():
|
||||
"""A pod that just booted cannot tell whether an outstanding request predates the prices
|
||||
it fetched at import, so it serves it. Adopting instead would strand it on stale prices
|
||||
with no interval configured to rescue it"""
|
||||
assert (
|
||||
pod_reload_is_due(
|
||||
schedule=ReloadSchedule(reload_revision=12),
|
||||
pod_applied_revision=0,
|
||||
pod_data_loaded_at=NOW,
|
||||
current_time=NOW,
|
||||
description="test",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_pod_reload_decision_ignores_persisted_last_run():
|
||||
"""A pod with stale data must refresh even when another pod already stamped
|
||||
last_run_at within the interval"""
|
||||
assert (
|
||||
pod_reload_is_due(
|
||||
schedule=ReloadSchedule(interval_hours=6, last_run_at=datetime(2024, 1, 1, 11, 59, tzinfo=timezone.utc)),
|
||||
pod_applied_revision=0,
|
||||
pod_data_loaded_at=LAST_RUN,
|
||||
current_time=NOW,
|
||||
description="test",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_schedule_that_never_ran_fires_immediately():
|
||||
"""A fresh schedule must not wait a full interval for its first run, even on a pod
|
||||
whose own data is boot-fresh"""
|
||||
assert (
|
||||
pod_reload_is_due(
|
||||
schedule=ReloadSchedule(interval_hours=6, last_run_at=None),
|
||||
pod_applied_revision=0,
|
||||
pod_data_loaded_at=datetime(2024, 1, 1, 11, 59, tzinfo=timezone.utc),
|
||||
current_time=NOW,
|
||||
description="test",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_reload_schedule_returns_none_for_missing_row():
|
||||
assert await read_reload_schedule(_mock_prisma(row=None), "model_cost_map_reload_config") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_reload_schedule_surfaces_revision_on_interval_less_row():
|
||||
"""A manual reload on a proxy with no schedule creates a row with only the columns
|
||||
set; the revision must still reach other pods"""
|
||||
prisma_client = _mock_prisma(row=_row(param_value=None, reload_revision=3))
|
||||
|
||||
schedule = await read_reload_schedule(prisma_client, "model_cost_map_reload_config")
|
||||
|
||||
assert schedule == ReloadSchedule(interval_hours=None, reload_revision=3, last_run_at=None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_reload_interval_touches_only_param_value():
|
||||
prisma_client = _mock_prisma()
|
||||
|
||||
await write_reload_interval(prisma_client, "model_cost_map_reload_config", 12)
|
||||
|
||||
data = prisma_client.db.litellm_config.upsert.await_args.kwargs["data"]
|
||||
assert set(data["update"]) == {"param_value"}
|
||||
assert set(data["create"]) == {"param_name", "param_value"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_reload_run_updates_last_run_without_creating_or_bumping():
|
||||
"""update_many so a schedule deleted mid-poll stays deleted, and the untouched revision
|
||||
keeps fanning the request out to pods that have not applied it"""
|
||||
prisma_client = _mock_prisma()
|
||||
|
||||
await record_reload_run(prisma_client, "model_cost_map_reload_config", LAST_RUN)
|
||||
|
||||
kwargs = prisma_client.db.litellm_config.update_many.await_args.kwargs
|
||||
assert kwargs == {"data": {"last_run_at": LAST_RUN}, "where": {"param_name": "model_cost_map_reload_config"}}
|
||||
prisma_client.db.litellm_config.upsert.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelling_a_schedule_never_reissues_a_revision():
|
||||
"""Cancelling must keep the row. The revision identifies a request rather than ordering
|
||||
one, so a counter restarted by a delete reissues a number pods already applied and their
|
||||
next manual reload is skipped everywhere but the pod that served it"""
|
||||
prisma_client = _fake_prisma(_FakeConfigTable())
|
||||
param_name = "model_cost_map_reload_config"
|
||||
await write_reload_interval(prisma_client, param_name, 6)
|
||||
pod_applied_revision = await record_manual_reload(prisma_client, param_name, LAST_RUN)
|
||||
|
||||
await clear_reload_interval(prisma_client, param_name)
|
||||
republished = await record_manual_reload(prisma_client, param_name, NOW)
|
||||
|
||||
assert (pod_applied_revision, republished) == (1, 2)
|
||||
schedule = await read_reload_schedule(prisma_client, param_name)
|
||||
assert schedule is not None
|
||||
assert (
|
||||
pod_reload_is_due(
|
||||
schedule=schedule,
|
||||
pod_applied_revision=pod_applied_revision,
|
||||
pod_data_loaded_at=NOW,
|
||||
current_time=NOW,
|
||||
description="test",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelling_a_schedule_stops_it_while_keeping_the_recorded_run():
|
||||
"""Dropping only the admin-owned param_value: the card must report no schedule without
|
||||
losing the last run it already showed"""
|
||||
prisma_client = _fake_prisma(_FakeConfigTable())
|
||||
param_name = "model_cost_map_reload_config"
|
||||
await write_reload_interval(prisma_client, param_name, 6)
|
||||
await record_reload_run(prisma_client, param_name, LAST_RUN)
|
||||
|
||||
await clear_reload_interval(prisma_client, param_name)
|
||||
|
||||
status = reload_schedule_status(await read_reload_schedule(prisma_client, param_name))
|
||||
assert status == {
|
||||
"scheduled": False,
|
||||
"interval_hours": None,
|
||||
"last_run": "2024-01-01T06:00:00+00:00",
|
||||
"next_run": None,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_manual_reload_bumps_the_revision_atomically():
|
||||
"""The increment must be delegated to the database: two concurrent requests that both
|
||||
read then wrote a computed value would publish the same revision and one would be lost"""
|
||||
prisma_client = _mock_prisma(upserted_revision=5)
|
||||
|
||||
published = await record_manual_reload(prisma_client, "model_cost_map_reload_config", LAST_RUN)
|
||||
|
||||
data = prisma_client.db.litellm_config.upsert.await_args.kwargs["data"]
|
||||
assert data["update"] == {"last_run_at": LAST_RUN, "reload_revision": {"increment": 1}}
|
||||
assert data["create"] == {
|
||||
"param_name": "model_cost_map_reload_config",
|
||||
"last_run_at": LAST_RUN,
|
||||
"reload_revision": 1,
|
||||
}
|
||||
assert published == 5
|
||||
|
|
@ -10,9 +10,9 @@ Routes covered:
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from .conftest import VOLATILE_KEYS, normalize
|
||||
|
||||
|
|
@ -35,6 +35,7 @@ def _attach_litellm_config(mock_prisma):
|
|||
table.upsert = AsyncMock()
|
||||
table.create = AsyncMock()
|
||||
table.update = AsyncMock()
|
||||
table.update_many = AsyncMock(return_value=1)
|
||||
table.delete = AsyncMock()
|
||||
table.delete_many = AsyncMock()
|
||||
mock_prisma.db.litellm_config = table
|
||||
|
|
@ -84,6 +85,9 @@ def test_reload_model_cost_map_happy(client, auth_as, monkeypatch, mock_prisma):
|
|||
"timestamp": "<VOLATILE>",
|
||||
}
|
||||
assert table.upsert.await_count == 1
|
||||
update_payload = table.upsert.await_args.kwargs["data"]["update"]
|
||||
assert set(update_payload) == {"last_run_at", "reload_revision"}
|
||||
assert update_payload["reload_revision"] == {"increment": 1}
|
||||
|
||||
|
||||
def test_reload_model_cost_map_fetch_failure_502_keeps_map(
|
||||
|
|
@ -179,6 +183,9 @@ def test_schedule_model_cost_map_reload_happy(
|
|||
"timestamp": "<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(
|
||||
|
|
@ -213,18 +220,15 @@ def test_schedule_model_cost_map_reload_not_admin_forbidden(client, auth_as):
|
|||
|
||||
|
||||
def test_cancel_model_cost_map_reload_happy(client, auth_as, monkeypatch, mock_prisma):
|
||||
"""Admin cancellation deletes config row and returns success body."""
|
||||
"""Admin cancellation clears the interval and returns success body. The row itself stays:
|
||||
it also holds the reload revision, and a counter restarted by a delete reissues a number
|
||||
pods already applied, silently skipping their next manual reload."""
|
||||
from litellm.proxy import proxy_server as ps
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
table = _attach_litellm_config(mock_prisma)
|
||||
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
|
||||
|
||||
async def _fake_invalidate(name):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(ps, "invalidate_config_param", _fake_invalidate)
|
||||
|
||||
with auth_as(LitellmUserRoles.PROXY_ADMIN):
|
||||
response = client.delete("/schedule/model_cost_map_reload")
|
||||
assert response.status_code == 200
|
||||
|
|
@ -234,7 +238,8 @@ def test_cancel_model_cost_map_reload_happy(client, auth_as, monkeypatch, mock_p
|
|||
"status": "success",
|
||||
"timestamp": "<VOLATILE>",
|
||||
}
|
||||
assert table.delete.await_count == 1
|
||||
assert json.loads(table.update_many.await_args.kwargs["data"]["param_value"]) == {"interval_hours": None}
|
||||
assert table.delete.await_count == 0
|
||||
|
||||
|
||||
def test_cancel_model_cost_map_reload_not_admin_forbidden(client, auth_as):
|
||||
|
|
@ -290,10 +295,11 @@ def test_get_model_cost_map_reload_status_scheduled(
|
|||
|
||||
table = _attach_litellm_config(mock_prisma)
|
||||
config_row = MagicMock()
|
||||
config_row.param_value = {"interval_hours": 12, "force_reload": False}
|
||||
config_row.param_value = {"interval_hours": 12}
|
||||
config_row.reload_revision = 0
|
||||
config_row.last_run_at = None
|
||||
table.find_unique = AsyncMock(return_value=config_row)
|
||||
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
|
||||
monkeypatch.setattr(ps, "last_model_cost_map_reload", None)
|
||||
|
||||
with auth_as(LitellmUserRoles.PROXY_ADMIN):
|
||||
response = client.get("/schedule/model_cost_map_reload/status")
|
||||
|
|
@ -306,19 +312,50 @@ def test_get_model_cost_map_reload_status_scheduled(
|
|||
}
|
||||
|
||||
|
||||
def test_get_model_cost_map_reload_status_no_config_not_scheduled(
|
||||
def test_get_model_cost_map_reload_status_reports_persisted_last_run(
|
||||
client, auth_as, monkeypatch, mock_prisma
|
||||
):
|
||||
"""Config row exists but interval_hours=None → not scheduled."""
|
||||
"""last_run/next_run come from the DB row, so status survives pod restarts."""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from litellm.proxy import proxy_server as ps
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
table = _attach_litellm_config(mock_prisma)
|
||||
config_row = MagicMock()
|
||||
config_row.param_value = {"interval_hours": None, "force_reload": True}
|
||||
config_row.param_value = {"interval_hours": 6}
|
||||
config_row.reload_revision = 0
|
||||
config_row.last_run_at = datetime(2024, 1, 1, 6, 0, 0, tzinfo=timezone.utc)
|
||||
table.find_unique = AsyncMock(return_value=config_row)
|
||||
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
|
||||
|
||||
with auth_as(LitellmUserRoles.PROXY_ADMIN):
|
||||
response = client.get("/schedule/model_cost_map_reload/status")
|
||||
assert response.status_code == 200
|
||||
assert normalize(response.json()) == {
|
||||
"scheduled": True,
|
||||
"interval_hours": 6,
|
||||
"last_run": "2024-01-01T06:00:00+00:00",
|
||||
"next_run": "2024-01-01T12:00:00+00:00",
|
||||
}
|
||||
|
||||
|
||||
def test_get_model_cost_map_reload_status_no_config_not_scheduled(
|
||||
client, auth_as, monkeypatch, mock_prisma
|
||||
):
|
||||
"""A row left behind by a manual reload (interval_hours=None) → not scheduled."""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from litellm.proxy import proxy_server as ps
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
table = _attach_litellm_config(mock_prisma)
|
||||
config_row = MagicMock()
|
||||
config_row.param_value = {"interval_hours": None}
|
||||
config_row.reload_revision = 3
|
||||
config_row.last_run_at = None
|
||||
table.find_unique = AsyncMock(return_value=config_row)
|
||||
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
|
||||
monkeypatch.setattr(ps, "last_model_cost_map_reload", None)
|
||||
|
||||
with auth_as(LitellmUserRoles.PROXY_ADMIN):
|
||||
response = client.get("/schedule/model_cost_map_reload/status")
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ import litellm
|
|||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
from litellm.caching.caching import RedisCache
|
||||
from litellm.caching.redis_cluster_cache import RedisClusterCache
|
||||
from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
|
@ -755,6 +756,48 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch):
|
|||
assert len(mock_scheduler_calls) > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_periodic_reload_job_scheduled_without_store_model_in_db(monkeypatch):
|
||||
"""
|
||||
Regression (LIT-4882): reload schedules configured from the Admin UI live in the DB and
|
||||
must fire even without store_model_in_db, which used to gate the job that ran them
|
||||
"""
|
||||
monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False)
|
||||
monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False)
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
|
||||
from litellm.proxy.proxy_server import ProxyStartupEvent
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=None)
|
||||
mock_proxy_logging = MagicMock(spec=ProxyLogging)
|
||||
mock_proxy_logging.slack_alerting_instance = MagicMock()
|
||||
mock_proxy_config = AsyncMock()
|
||||
scheduler = AsyncIOScheduler()
|
||||
|
||||
try:
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config),
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", False),
|
||||
patch("litellm.proxy.proxy_server.get_secret_bool", return_value=False),
|
||||
patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler),
|
||||
):
|
||||
await ProxyStartupEvent.initialize_scheduled_background_jobs(
|
||||
general_settings={},
|
||||
prisma_client=mock_prisma_client,
|
||||
proxy_budget_rescheduler_min_time=1,
|
||||
proxy_budget_rescheduler_max_time=2,
|
||||
proxy_batch_write_at=5,
|
||||
proxy_logging_obj=mock_proxy_logging,
|
||||
)
|
||||
|
||||
assert scheduler.get_job("periodic_reload_job") is not None
|
||||
assert scheduler.get_job("add_deployment_job") is None
|
||||
finally:
|
||||
scheduler.shutdown(wait=False)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_scheduled_jobs_uses_configured_config_reload_interval(monkeypatch):
|
||||
"""
|
||||
|
|
@ -3596,6 +3639,20 @@ async def test_chat_completion_result_no_nested_none_values():
|
|||
# ============================================================================
|
||||
|
||||
|
||||
def _reload_schedule_row(
|
||||
param_value: dict,
|
||||
*,
|
||||
reload_revision: int = 0,
|
||||
last_run_at: datetime | None = None,
|
||||
) -> types.SimpleNamespace:
|
||||
"""LiteLLM_Config row shape: admin-owned interval in param_value, run state in dedicated columns"""
|
||||
return types.SimpleNamespace(
|
||||
param_value=param_value,
|
||||
reload_revision=reload_revision,
|
||||
last_run_at=last_run_at,
|
||||
)
|
||||
|
||||
|
||||
class TestPriceDataReloadAPI:
|
||||
"""Test cases for price data reload API endpoints"""
|
||||
|
||||
|
|
@ -3636,10 +3693,9 @@ class TestPriceDataReloadAPI:
|
|||
):
|
||||
# Mock the database connection
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(
|
||||
return_value=None
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(
|
||||
return_value=_reload_schedule_row({}, reload_revision=1)
|
||||
)
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
|
||||
|
||||
response = client_with_auth.post("/reload/model_cost_map")
|
||||
|
||||
|
|
@ -3694,7 +3750,9 @@ class TestPriceDataReloadAPI:
|
|||
# Mock the database connection
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None)
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(
|
||||
return_value=_reload_schedule_row({}, reload_revision=1)
|
||||
)
|
||||
|
||||
response = client_with_auth.post("/reload/model_cost_map")
|
||||
|
||||
|
|
@ -3705,10 +3763,10 @@ class TestPriceDataReloadAPI:
|
|||
assert "Failed to reload model cost map" in data["detail"]
|
||||
|
||||
def test_schedule_model_cost_map_reload_admin_access(self, client_with_auth):
|
||||
"""Test that admin users can schedule periodic reload"""
|
||||
"""Admin schedule write owns param_value only, so it can't clobber the job-owned run columns"""
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
|
||||
# Mock database upsert
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=_reload_schedule_row({}, reload_revision=1))
|
||||
|
||||
response = client_with_auth.post("/schedule/model_cost_map_reload?hours=6")
|
||||
|
||||
|
|
@ -3719,6 +3777,15 @@ class TestPriceDataReloadAPI:
|
|||
assert "message" in data
|
||||
assert "timestamp" in data
|
||||
|
||||
call_args = mock_prisma.db.litellm_config.upsert.call_args
|
||||
assert call_args[1]["where"] == {"param_name": "model_cost_map_reload_config"}
|
||||
update_payload = call_args[1]["data"]["update"]
|
||||
assert set(update_payload.keys()) == {"param_value"}
|
||||
assert json.loads(update_payload["param_value"]) == {"interval_hours": 6}
|
||||
create_payload = call_args[1]["data"]["create"]
|
||||
assert set(create_payload.keys()) == {"param_name", "param_value"}
|
||||
assert json.loads(create_payload["param_value"]) == {"interval_hours": 6}
|
||||
|
||||
def test_schedule_model_cost_map_reload_non_admin_access(self, client_with_auth):
|
||||
"""Test that non-admin users cannot schedule periodic reload"""
|
||||
# Mock non-admin user
|
||||
|
|
@ -3744,7 +3811,7 @@ class TestPriceDataReloadAPI:
|
|||
def test_cancel_model_cost_map_reload_admin_access(self, client_with_auth):
|
||||
"""Test that admin users can cancel periodic reload"""
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
|
||||
# Mock database delete
|
||||
mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=1)
|
||||
mock_prisma.db.litellm_config.delete = AsyncMock(return_value=None)
|
||||
|
||||
response = client_with_auth.delete("/schedule/model_cost_map_reload")
|
||||
|
|
@ -3754,6 +3821,10 @@ class TestPriceDataReloadAPI:
|
|||
assert data["status"] == "success"
|
||||
assert "message" in data
|
||||
assert "timestamp" in data
|
||||
assert json.loads(mock_prisma.db.litellm_config.update_many.await_args.kwargs["data"]["param_value"]) == {
|
||||
"interval_hours": None
|
||||
}
|
||||
mock_prisma.db.litellm_config.delete.assert_not_called()
|
||||
|
||||
def test_cancel_model_cost_map_reload_non_admin_access(self, client_with_auth):
|
||||
"""Test that non-admin users cannot cancel periodic reload"""
|
||||
|
|
@ -3770,35 +3841,28 @@ class TestPriceDataReloadAPI:
|
|||
assert "Admin role required" in data["detail"]
|
||||
|
||||
def test_get_model_cost_map_reload_status_admin_access(self, client_with_auth):
|
||||
"""Test that admin users can get reload status"""
|
||||
"""
|
||||
Regression (LIT-4882): status is served purely from the DB row, so a restarted pod
|
||||
(whose in-memory clock only knows its own boot) still reports the real last/next run
|
||||
"""
|
||||
proxy_server_module.proxy_config.model_cost_map_loaded_at = datetime(2030, 6, 1, tzinfo=timezone.utc)
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
|
||||
# Mock database config record
|
||||
mock_config = MagicMock()
|
||||
mock_config.param_value = {"interval_hours": 6, "force_reload": False}
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(
|
||||
return_value=mock_config
|
||||
return_value=_reload_schedule_row(
|
||||
{"interval_hours": 6},
|
||||
last_run_at=datetime(2024, 1, 1, 6, 0, tzinfo=timezone.utc),
|
||||
)
|
||||
)
|
||||
|
||||
# Mock the last reload time and current time
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.last_model_cost_map_reload",
|
||||
"2024-01-01T06:00:00",
|
||||
):
|
||||
with patch("litellm.proxy.proxy_server.datetime") as mock_datetime:
|
||||
# Mock current time to be 1 hour after last reload
|
||||
mock_datetime.utcnow.return_value = datetime(2024, 1, 1, 7, 0, 0)
|
||||
mock_datetime.fromisoformat = datetime.fromisoformat
|
||||
response = client_with_auth.get("/schedule/model_cost_map_reload/status")
|
||||
|
||||
response = client_with_auth.get(
|
||||
"/schedule/model_cost_map_reload/status"
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["scheduled"] == True
|
||||
assert data["interval_hours"] == 6
|
||||
assert data["last_run"] == "2024-01-01T06:00:00"
|
||||
assert data["next_run"] == "2024-01-01T12:00:00"
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["scheduled"] is True
|
||||
assert data["interval_hours"] == 6
|
||||
assert data["last_run"] == "2024-01-01T06:00:00+00:00"
|
||||
assert data["next_run"] == "2024-01-01T12:00:00+00:00"
|
||||
|
||||
def test_get_model_cost_map_reload_status_non_admin_access(self, client_with_auth):
|
||||
"""Test that non-admin users cannot get reload status"""
|
||||
|
|
@ -3829,36 +3893,44 @@ class TestPriceDataReloadAPI:
|
|||
assert data["next_run"] == None
|
||||
|
||||
def test_get_model_cost_map_reload_status_no_interval(self, client_with_auth):
|
||||
"""Test that status returns not scheduled when no interval is configured"""
|
||||
"""A row left behind by a manual reload (no interval) must not read as scheduled"""
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
|
||||
# Mock config with no interval
|
||||
mock_config = MagicMock()
|
||||
mock_config.param_value = {"interval_hours": None, "force_reload": False}
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(
|
||||
return_value=mock_config
|
||||
return_value=_reload_schedule_row(
|
||||
{"interval_hours": None},
|
||||
reload_revision=3,
|
||||
)
|
||||
)
|
||||
|
||||
response = client_with_auth.get("/schedule/model_cost_map_reload/status")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["scheduled"] == False
|
||||
assert data["interval_hours"] == None
|
||||
assert data["last_run"] == None
|
||||
assert data["next_run"] == None
|
||||
assert data["scheduled"] is False
|
||||
assert data["interval_hours"] is None
|
||||
assert data["last_run"] is None
|
||||
assert data["next_run"] is None
|
||||
|
||||
def test_get_model_cost_map_reload_status_before_first_run(self, client_with_auth):
|
||||
"""Scheduled but never executed: no last_run_at means no next_run can be computed"""
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(
|
||||
return_value=_reload_schedule_row({"interval_hours": 6})
|
||||
)
|
||||
|
||||
response = client_with_auth.get("/schedule/model_cost_map_reload/status")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["scheduled"] is True
|
||||
assert data["interval_hours"] == 6
|
||||
assert data["last_run"] is None
|
||||
assert data["next_run"] is None
|
||||
|
||||
|
||||
class TestPriceDataReloadIntegration:
|
||||
"""Integration tests for the complete price data reload feature"""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _flush_litellm_config_cache(self):
|
||||
from litellm.proxy.utils import litellm_config_cache
|
||||
|
||||
litellm_config_cache.flush_cache()
|
||||
yield
|
||||
litellm_config_cache.flush_cache()
|
||||
|
||||
@pytest.fixture
|
||||
def client_with_auth(self):
|
||||
"""Create a test client with authentication"""
|
||||
|
|
@ -3900,10 +3972,9 @@ class TestPriceDataReloadIntegration:
|
|||
):
|
||||
# Mock the database connection
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(
|
||||
return_value=None
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(
|
||||
return_value=_reload_schedule_row({}, reload_revision=1)
|
||||
)
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
|
||||
|
||||
# Test reload endpoint
|
||||
response = client_with_auth.post("/reload/model_cost_map")
|
||||
|
|
@ -3916,172 +3987,352 @@ class TestPriceDataReloadIntegration:
|
|||
litellm.model_cost = original_model_cost
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
def test_distributed_reload_check_function(self):
|
||||
"""Test the _check_and_reload_model_cost_map function"""
|
||||
def test_pod_data_clock_seeded_from_actual_cost_map_load(self):
|
||||
"""Regression: seeding from ProxyConfig construction time instead of the real
|
||||
import-time fetch let a manual request stamped during startup be skipped"""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
from litellm.proxy.utils import litellm_config_cache
|
||||
|
||||
proxy_config = ProxyConfig()
|
||||
|
||||
# Mock prisma client
|
||||
mock_prisma = MagicMock()
|
||||
|
||||
# Test case 1: No config in database
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None)
|
||||
# _check_and_reload_model_cost_map routes through get_config_param,
|
||||
# which calls prisma.get_generic_data on a cache miss.
|
||||
mock_prisma.get_generic_data = AsyncMock(return_value=None)
|
||||
|
||||
# Should return early without reloading
|
||||
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
|
||||
|
||||
# Test case 2: Config with interval but not time to reload
|
||||
litellm_config_cache.flush_cache()
|
||||
mock_config = MagicMock()
|
||||
mock_config.param_value = {"interval_hours": 6, "force_reload": False}
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config)
|
||||
mock_prisma.get_generic_data = AsyncMock(return_value=mock_config)
|
||||
|
||||
# Mock current time and last reload time
|
||||
fetch_time = datetime(2024, 1, 1, 6, 0, tzinfo=timezone.utc)
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.last_model_cost_map_reload",
|
||||
"2024-01-01T06:00:00",
|
||||
"litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map_loaded_at",
|
||||
return_value=fetch_time,
|
||||
):
|
||||
with patch("litellm.proxy.proxy_server.datetime") as mock_datetime:
|
||||
mock_datetime.utcnow.return_value = datetime(
|
||||
2024, 1, 1, 7, 0, 0
|
||||
) # 1 hour later
|
||||
assert ProxyConfig().model_cost_map_loaded_at == fetch_time
|
||||
|
||||
# Should not reload (only 1 hour passed, need 6)
|
||||
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
|
||||
|
||||
# Test case 3: Config with force reload
|
||||
litellm_config_cache.flush_cache()
|
||||
mock_config.param_value = {"interval_hours": 6, "force_reload": True}
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config)
|
||||
mock_prisma.get_generic_data = AsyncMock(return_value=mock_config)
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
|
||||
|
||||
from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded
|
||||
|
||||
original_model_cost = litellm.model_cost.copy()
|
||||
try:
|
||||
with patch(
|
||||
"litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map",
|
||||
new=AsyncMock(
|
||||
return_value=ModelCostMapReloaded(
|
||||
model_cost_map={"gpt-3.5-turbo": {"input_cost_per_token": 0.001}}
|
||||
)
|
||||
),
|
||||
):
|
||||
# Should reload due to force flag
|
||||
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
|
||||
|
||||
# Verify force_reload was reset to False
|
||||
mock_prisma.db.litellm_config.upsert.assert_called()
|
||||
call_args = mock_prisma.db.litellm_config.upsert.call_args
|
||||
# The param_value is now a JSON string, so we need to parse it
|
||||
param_value_json = call_args[1]["data"]["update"]["param_value"]
|
||||
param_value_dict = json.loads(param_value_json)
|
||||
assert param_value_dict["force_reload"] == False
|
||||
assert param_value_dict.get("interval_hours") == 6
|
||||
finally:
|
||||
litellm.model_cost = original_model_cost
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
def test_distributed_reload_preserves_interval_hours(self):
|
||||
"""Test that _check_and_reload_model_cost_map preserves interval_hours after reload.
|
||||
|
||||
Regression test: the update branch of the upsert was previously dropping
|
||||
interval_hours, causing scheduled reloads to self-destruct after first execution.
|
||||
def test_distributed_reload_check_function(self):
|
||||
"""
|
||||
A revision this pod has not applied takes effect here even one minute into a 6h
|
||||
interval; a missing row is a no-op
|
||||
"""
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
proxy_config = ProxyConfig()
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=_reload_schedule_row({}, reload_revision=1))
|
||||
mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None)
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
# Set up config with interval_hours=24 and force_reload=True to trigger reload
|
||||
mock_config = MagicMock()
|
||||
mock_config.param_value = {"interval_hours": 24, "force_reload": True}
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config)
|
||||
# _check_and_reload_model_cost_map now reads through get_generic_data.
|
||||
mock_prisma.get_generic_data = AsyncMock(return_value=mock_config)
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
|
||||
boot_loaded_at = proxy_config.model_cost_map_loaded_at
|
||||
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
|
||||
|
||||
mock_prisma.db.litellm_config.update_many.assert_not_called()
|
||||
assert proxy_config.model_cost_map_loaded_at == boot_loaded_at
|
||||
|
||||
frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc)
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(
|
||||
return_value=_reload_schedule_row(
|
||||
{"interval_hours": 6},
|
||||
reload_revision=4,
|
||||
last_run_at=datetime(2024, 1, 1, 6, 59, 30, tzinfo=timezone.utc),
|
||||
)
|
||||
)
|
||||
proxy_config.model_cost_map_loaded_at = frozen_now - timedelta(minutes=1)
|
||||
proxy_config.model_cost_map_applied_revision = 3
|
||||
|
||||
from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded
|
||||
|
||||
original_model_cost = litellm.model_cost.copy()
|
||||
try:
|
||||
with patch(
|
||||
"litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map",
|
||||
new=AsyncMock(
|
||||
return_value=ModelCostMapReloaded(
|
||||
model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}}
|
||||
)
|
||||
),
|
||||
with (
|
||||
patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) as mock_get_map,
|
||||
patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now),
|
||||
):
|
||||
mock_get_map.return_value = ModelCostMapReloaded(model_cost_map={"gpt-3.5-turbo": {"input_cost_per_token": 0.001}})
|
||||
|
||||
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
|
||||
|
||||
assert litellm.model_cost["gpt-3.5-turbo"] == {"input_cost_per_token": 0.001}
|
||||
assert proxy_config.model_cost_map_loaded_at == frozen_now
|
||||
assert mock_prisma.db.litellm_config.update_many.call_args[1] == {
|
||||
"data": {"last_run_at": frozen_now},
|
||||
"where": {"param_name": "model_cost_map_reload_config"},
|
||||
}
|
||||
mock_prisma.db.litellm_config.upsert.assert_not_called()
|
||||
assert proxy_config.model_cost_map_applied_revision == 4
|
||||
finally:
|
||||
litellm.model_cost = original_model_cost
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
def test_distributed_reload_ignores_already_applied_request(self):
|
||||
"""
|
||||
A revision this pod already applied must not re-trigger on every job tick for the
|
||||
rest of the interval
|
||||
"""
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
proxy_config = ProxyConfig()
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None)
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(
|
||||
return_value=_reload_schedule_row(
|
||||
{"interval_hours": 6},
|
||||
reload_revision=4,
|
||||
last_run_at=datetime(2024, 1, 1, 6, 0, tzinfo=timezone.utc),
|
||||
)
|
||||
)
|
||||
frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc)
|
||||
pod_data_loaded_at = frozen_now - timedelta(minutes=1)
|
||||
proxy_config.model_cost_map_loaded_at = pod_data_loaded_at
|
||||
proxy_config.model_cost_map_applied_revision = 4
|
||||
|
||||
original_model_cost = litellm.model_cost.copy()
|
||||
try:
|
||||
with (
|
||||
patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) as mock_get_map,
|
||||
patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now),
|
||||
):
|
||||
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
|
||||
|
||||
# Verify the upsert update branch preserves interval_hours
|
||||
mock_prisma.db.litellm_config.upsert.assert_called()
|
||||
call_args = mock_prisma.db.litellm_config.upsert.call_args
|
||||
param_value_json = call_args[1]["data"]["update"]["param_value"]
|
||||
param_value_dict = json.loads(param_value_json)
|
||||
assert param_value_dict["force_reload"] == False
|
||||
assert param_value_dict["interval_hours"] == 24, (
|
||||
"interval_hours must be preserved in the update branch; "
|
||||
"dropping it causes the schedule to self-destruct"
|
||||
)
|
||||
mock_get_map.assert_not_called()
|
||||
mock_prisma.db.litellm_config.update_many.assert_not_called()
|
||||
assert proxy_config.model_cost_map_loaded_at == pod_data_loaded_at
|
||||
finally:
|
||||
litellm.model_cost = original_model_cost
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
def test_periodic_reload_uses_pod_local_data_age(self):
|
||||
"""
|
||||
Each pod decides from the age of its own data, so a pod holding a stale copy
|
||||
refreshes even when the shared row was just stamped by another pod, and stays
|
||||
put while its copy is inside the interval
|
||||
"""
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
proxy_config = ProxyConfig()
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(
|
||||
return_value=_reload_schedule_row(
|
||||
{"interval_hours": 6},
|
||||
last_run_at=datetime(2024, 1, 1, 6, 59, tzinfo=timezone.utc),
|
||||
)
|
||||
)
|
||||
mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None)
|
||||
frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc)
|
||||
proxy_config.model_cost_map_loaded_at = datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc)
|
||||
|
||||
original_model_cost = litellm.model_cost.copy()
|
||||
try:
|
||||
with (
|
||||
patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) as mock_get_map,
|
||||
patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now),
|
||||
):
|
||||
mock_get_map.return_value = ModelCostMapReloaded(model_cost_map={"gpt-4-test": {"input_cost_per_token": 0.5}})
|
||||
|
||||
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
|
||||
|
||||
assert litellm.model_cost["gpt-4-test"] == {"input_cost_per_token": 0.5}
|
||||
assert proxy_config.model_cost_map_loaded_at == frozen_now
|
||||
assert mock_prisma.db.litellm_config.update_many.call_args[1]["data"] == {"last_run_at": frozen_now}
|
||||
|
||||
mock_get_map.reset_mock()
|
||||
mock_prisma.db.litellm_config.update_many.reset_mock()
|
||||
proxy_config.model_cost_map_loaded_at = frozen_now - timedelta(hours=1)
|
||||
|
||||
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
|
||||
|
||||
mock_get_map.assert_not_called()
|
||||
mock_prisma.db.litellm_config.update_many.assert_not_called()
|
||||
assert proxy_config.model_cost_map_loaded_at == frozen_now - timedelta(hours=1)
|
||||
finally:
|
||||
litellm.model_cost = original_model_cost
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
def test_every_pod_applies_a_manual_revision_exactly_once(self):
|
||||
"""The fleet property: no pod clears the revision, so each one reloads on the tick
|
||||
after it is published and then stops, whatever order the pods poll in"""
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
pods = [ProxyConfig(), ProxyConfig(), ProxyConfig()]
|
||||
frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc)
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None)
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=_reload_schedule_row({}, reload_revision=1))
|
||||
for pod in pods:
|
||||
pod.model_cost_map_applied_revision = 0
|
||||
pod.model_cost_map_loaded_at = frozen_now
|
||||
|
||||
original_model_cost = litellm.model_cost.copy()
|
||||
try:
|
||||
with (
|
||||
patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) as mock_get_map,
|
||||
patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now),
|
||||
):
|
||||
mock_get_map.return_value = ModelCostMapReloaded(model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}})
|
||||
|
||||
for _ in range(3):
|
||||
for pod in pods:
|
||||
asyncio.run(pod._check_and_reload_model_cost_map(mock_prisma))
|
||||
|
||||
assert mock_get_map.call_count == len(pods)
|
||||
assert all(p.model_cost_map_applied_revision == 1 for p in pods)
|
||||
finally:
|
||||
litellm.model_cost = original_model_cost
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"published_revision, expect_reload",
|
||||
[(4, True), (0, False)],
|
||||
)
|
||||
def test_booting_pod_serves_an_outstanding_request_once(self, published_revision, expect_reload):
|
||||
"""
|
||||
Regression: a manual reload published while this pod was starting must still be
|
||||
served. The pod cannot prove its import-time fetch already covers that request, so
|
||||
it applies it on the first poll and adopts the revision, leaving later polls quiet.
|
||||
A row nobody has ever reloaded (revision 0) costs the pod nothing
|
||||
"""
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
proxy_config = ProxyConfig()
|
||||
frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc)
|
||||
proxy_config.model_cost_map_loaded_at = frozen_now
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None)
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(
|
||||
return_value=_reload_schedule_row({}, reload_revision=published_revision)
|
||||
)
|
||||
|
||||
original_model_cost = litellm.model_cost.copy()
|
||||
try:
|
||||
with (
|
||||
patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) as mock_get_map,
|
||||
patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now),
|
||||
):
|
||||
mock_get_map.return_value = ModelCostMapReloaded(model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}})
|
||||
|
||||
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
|
||||
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
|
||||
|
||||
assert mock_get_map.call_count == (1 if expect_reload else 0)
|
||||
assert proxy_config.model_cost_map_applied_revision == published_revision
|
||||
finally:
|
||||
litellm.model_cost = original_model_cost
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
def test_distributed_reload_stamps_last_run_without_creating_row(self):
|
||||
"""
|
||||
Regression: the job's write carries neither param_value (which would clobber the
|
||||
admin-configured interval) nor a create branch (which would resurrect a schedule
|
||||
a concurrent cancel just deleted)
|
||||
"""
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
proxy_config = ProxyConfig()
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=_reload_schedule_row({"interval_hours": 24}))
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=_reload_schedule_row({}, reload_revision=1))
|
||||
mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None)
|
||||
frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc)
|
||||
|
||||
original_model_cost = litellm.model_cost.copy()
|
||||
try:
|
||||
with (
|
||||
patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) as mock_get_map,
|
||||
patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now),
|
||||
):
|
||||
mock_get_map.return_value = ModelCostMapReloaded(model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}})
|
||||
|
||||
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
|
||||
|
||||
assert mock_prisma.db.litellm_config.update_many.call_args[1] == {
|
||||
"data": {"last_run_at": frozen_now},
|
||||
"where": {"param_name": "model_cost_map_reload_config"},
|
||||
}
|
||||
mock_prisma.db.litellm_config.upsert.assert_not_called()
|
||||
mock_prisma.db.litellm_config.create.assert_not_called()
|
||||
finally:
|
||||
litellm.model_cost = original_model_cost
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
def test_distributed_reload_leaves_request_unserved_when_status_write_fails(self):
|
||||
"""
|
||||
A run that never reached the row must not be recorded as served. Adopting the
|
||||
revision here would leave the card reporting the previous run until someone clicks
|
||||
again, because a manual request is published once and never republished
|
||||
"""
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
proxy_config = ProxyConfig()
|
||||
frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc)
|
||||
proxy_config.model_cost_map_loaded_at = frozen_now - timedelta(hours=9)
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(
|
||||
return_value=_reload_schedule_row({"interval_hours": 6}, reload_revision=7)
|
||||
)
|
||||
mock_prisma.db.litellm_config.update_many = AsyncMock(side_effect=Exception("connection reset"))
|
||||
|
||||
original_model_cost = litellm.model_cost.copy()
|
||||
try:
|
||||
with (
|
||||
patch(
|
||||
"litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_get_map,
|
||||
patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now),
|
||||
):
|
||||
mock_get_map.return_value = ModelCostMapReloaded(model_cost_map={"gpt-4": {"input_cost_per_token": 0.1}})
|
||||
|
||||
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
|
||||
|
||||
assert proxy_config.model_cost_map_applied_revision == 0
|
||||
finally:
|
||||
litellm.model_cost = original_model_cost
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
def test_distributed_reload_keeps_current_map_when_fetch_fails(self):
|
||||
"""Fetch failure during a periodic/forced reload must not downgrade the pod.
|
||||
"""Fetch failure during a periodic reload must not downgrade the pod or count the
|
||||
request as served.
|
||||
|
||||
Regression: a 429/network failure used to silently replace litellm.model_cost
|
||||
with the stale packaged backup, stamp last_run, and clear force_reload.
|
||||
Regression: a 429/network failure used to silently replace litellm.model_cost with
|
||||
the stale packaged backup and stamp last_run. Adopting the revision here would be
|
||||
the same bug one level up: a manual request is published once and never republished,
|
||||
so a pod that records it as applied without the data stays mispriced until someone
|
||||
clicks again
|
||||
"""
|
||||
from litellm.litellm_core_utils.get_model_cost_map import (
|
||||
ModelCostMapReloadUnavailable,
|
||||
)
|
||||
from litellm.proxy import proxy_server as ps
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
proxy_config = ProxyConfig()
|
||||
frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc)
|
||||
pod_data_loaded_at = frozen_now - timedelta(hours=9)
|
||||
proxy_config.model_cost_map_loaded_at = pod_data_loaded_at
|
||||
mock_prisma = MagicMock()
|
||||
|
||||
mock_config = MagicMock()
|
||||
mock_config.param_value = {"interval_hours": 6, "force_reload": True}
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config)
|
||||
mock_prisma.get_generic_data = AsyncMock(return_value=mock_config)
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(
|
||||
return_value=_reload_schedule_row({"interval_hours": 6}, reload_revision=7)
|
||||
)
|
||||
mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None)
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
|
||||
|
||||
original_model_cost = litellm.model_cost
|
||||
with patch(
|
||||
"litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map",
|
||||
new=AsyncMock(
|
||||
return_value=ModelCostMapReloadUnavailable(reason="HTTP 429 from upstream")
|
||||
with (
|
||||
patch(
|
||||
"litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map",
|
||||
new=AsyncMock(return_value=ModelCostMapReloadUnavailable(reason="HTTP 429 from upstream")),
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now),
|
||||
):
|
||||
with patch("litellm.proxy.proxy_server.last_model_cost_map_reload", None):
|
||||
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
|
||||
assert ps.last_model_cost_map_reload is None, (
|
||||
"a failed reload must not stamp the pod's last reload time, "
|
||||
"otherwise the retry waits a full interval"
|
||||
)
|
||||
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
|
||||
|
||||
assert litellm.model_cost is original_model_cost, (
|
||||
"a failed reload must keep the currently loaded cost map, "
|
||||
"not swap in the packaged backup"
|
||||
)
|
||||
assert proxy_config.model_cost_map_loaded_at == pod_data_loaded_at, (
|
||||
"a failed reload must not stamp the pod's data age, otherwise the retry waits a full interval"
|
||||
)
|
||||
assert proxy_config.model_cost_map_applied_revision == 0, (
|
||||
"a failed reload must leave the revision unapplied so the next poll retries it"
|
||||
)
|
||||
mock_prisma.db.litellm_config.update_many.assert_not_called()
|
||||
mock_prisma.db.litellm_config.upsert.assert_not_called()
|
||||
|
||||
def test_manual_reload_preserves_interval_hours(self):
|
||||
"""Test that manual reload via /reload/model_cost_map preserves existing interval_hours.
|
||||
|
||||
Regression test: the manual reload endpoint was overwriting param_value with
|
||||
only force_reload=True, dropping any existing interval_hours schedule.
|
||||
"""
|
||||
Regression: manual reload owns only the run columns, so it never reads or rewrites
|
||||
param_value and cannot destroy an existing schedule
|
||||
"""
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
from litellm.proxy.proxy_server import cleanup_router_config_variables
|
||||
|
|
@ -4095,44 +4346,40 @@ class TestPriceDataReloadIntegration:
|
|||
mock_auth.user_role = LitellmUserRoles.PROXY_ADMIN
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: mock_auth
|
||||
client = TestClient(app)
|
||||
frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc)
|
||||
|
||||
from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded
|
||||
|
||||
original_model_cost = litellm.model_cost.copy()
|
||||
try:
|
||||
with patch(
|
||||
"litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map",
|
||||
new=AsyncMock(
|
||||
return_value=ModelCostMapReloaded(
|
||||
model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}}
|
||||
)
|
||||
),
|
||||
with (
|
||||
patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) as mock_get_map,
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now),
|
||||
):
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
|
||||
# Simulate existing config with a schedule
|
||||
mock_existing = MagicMock()
|
||||
mock_existing.param_value = {
|
||||
"interval_hours": 12,
|
||||
"force_reload": False,
|
||||
}
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(
|
||||
return_value=mock_existing
|
||||
)
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
|
||||
mock_get_map.return_value = ModelCostMapReloaded(model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}})
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(
|
||||
return_value=_reload_schedule_row({}, reload_revision=9)
|
||||
)
|
||||
|
||||
response = client.post("/reload/model_cost_map")
|
||||
assert response.status_code == 200
|
||||
response = client.post("/reload/model_cost_map")
|
||||
assert response.status_code == 200
|
||||
|
||||
# Verify interval_hours was preserved in the upsert
|
||||
mock_prisma.db.litellm_config.upsert.assert_called()
|
||||
call_args = mock_prisma.db.litellm_config.upsert.call_args
|
||||
param_value_json = call_args[1]["data"]["update"]["param_value"]
|
||||
param_value_dict = json.loads(param_value_json)
|
||||
assert param_value_dict["force_reload"] == True
|
||||
assert param_value_dict["interval_hours"] == 12, (
|
||||
"interval_hours must be preserved when manual reload sets force_reload; "
|
||||
"dropping it destroys any existing schedule"
|
||||
)
|
||||
mock_prisma.db.litellm_config.find_unique.assert_not_called()
|
||||
call_args = mock_prisma.db.litellm_config.upsert.call_args
|
||||
assert call_args[1]["data"]["update"] == {
|
||||
"last_run_at": frozen_now,
|
||||
"reload_revision": {"increment": 1},
|
||||
}
|
||||
assert call_args[1]["data"]["create"] == {
|
||||
"param_name": "model_cost_map_reload_config",
|
||||
"last_run_at": frozen_now,
|
||||
"reload_revision": 1,
|
||||
}
|
||||
assert proxy_server_module.proxy_config.model_cost_map_loaded_at == frozen_now
|
||||
assert proxy_server_module.proxy_config.model_cost_map_applied_revision == 9, (
|
||||
"the serving pod must adopt the revision it published, not reload again"
|
||||
)
|
||||
finally:
|
||||
litellm.model_cost = original_model_cost
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
|
@ -4144,7 +4391,9 @@ class TestPriceDataReloadIntegration:
|
|||
identical to the model cost map bug.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
from litellm.proxy.utils import litellm_config_cache
|
||||
|
||||
litellm_config_cache.flush_cache()
|
||||
proxy_config = ProxyConfig()
|
||||
mock_prisma = MagicMock()
|
||||
|
||||
|
|
@ -4154,7 +4403,7 @@ class TestPriceDataReloadIntegration:
|
|||
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config)
|
||||
# _check_and_reload_anthropic_beta_headers now reads through get_generic_data.
|
||||
mock_prisma.get_generic_data = AsyncMock(return_value=mock_config)
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=_reload_schedule_row({}, reload_revision=1))
|
||||
|
||||
with patch(
|
||||
"litellm.anthropic_beta_headers_manager.reload_beta_headers_config"
|
||||
|
|
@ -4204,10 +4453,10 @@ class TestPriceDataReloadIntegration:
|
|||
# Simulate existing config with a schedule
|
||||
mock_existing = MagicMock()
|
||||
mock_existing.param_value = {"interval_hours": 8, "force_reload": False}
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(
|
||||
return_value=mock_existing
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_existing)
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(
|
||||
return_value=_reload_schedule_row({}, reload_revision=1)
|
||||
)
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
|
||||
|
||||
response = client.post("/reload/anthropic_beta_headers")
|
||||
assert response.status_code == 200
|
||||
|
|
@ -4251,64 +4500,6 @@ model_list:
|
|||
assert "model_list" in config
|
||||
assert len(config["model_list"]) == 2
|
||||
|
||||
def test_database_config_storage(self):
|
||||
"""Test that configuration is properly stored in database"""
|
||||
# Mock prisma client
|
||||
mock_prisma = MagicMock()
|
||||
|
||||
# Test the database upsert call that would be made by the schedule endpoint
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
|
||||
|
||||
# Simulate the database call that the schedule endpoint would make
|
||||
asyncio.run(
|
||||
mock_prisma.db.litellm_config.upsert(
|
||||
where={"param_name": "model_cost_map_reload_config"},
|
||||
data={
|
||||
"create": {
|
||||
"param_name": "model_cost_map_reload_config",
|
||||
"param_value": {"interval_hours": 6, "force_reload": False},
|
||||
},
|
||||
"update": {
|
||||
"param_value": {"interval_hours": 6, "force_reload": False}
|
||||
},
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
# Verify database upsert was called with correct data
|
||||
mock_prisma.db.litellm_config.upsert.assert_called_once()
|
||||
call_args = mock_prisma.db.litellm_config.upsert.call_args
|
||||
assert call_args[1]["where"]["param_name"] == "model_cost_map_reload_config"
|
||||
assert call_args[1]["data"]["create"]["param_value"]["interval_hours"] == 6
|
||||
assert call_args[1]["data"]["create"]["param_value"]["force_reload"] == False
|
||||
|
||||
def test_manual_reload_force_flag(self):
|
||||
"""Test that manual reload sets force flag correctly"""
|
||||
# Mock prisma client
|
||||
mock_prisma = MagicMock()
|
||||
|
||||
# Test the database upsert call that would be made by the manual reload endpoint
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
|
||||
|
||||
# Simulate the database call that the manual reload endpoint would make
|
||||
asyncio.run(
|
||||
mock_prisma.db.litellm_config.upsert(
|
||||
where={"param_name": "model_cost_map_reload_config"},
|
||||
data={
|
||||
"create": {
|
||||
"param_name": "model_cost_map_reload_config",
|
||||
"param_value": {"interval_hours": None, "force_reload": True},
|
||||
},
|
||||
"update": {"param_value": {"force_reload": True}},
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
# Verify force_reload flag was set
|
||||
mock_prisma.db.litellm_config.upsert.assert_called_once()
|
||||
call_args = mock_prisma.db.litellm_config.upsert.call_args
|
||||
assert call_args[1]["data"]["update"]["param_value"]["force_reload"] == True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_router_settings_from_db_config_merge_logic():
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue