mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
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).
This commit is contained in:
parent
b7ac3e237b
commit
11014897a7
6 changed files with 320 additions and 145 deletions
|
|
@ -10,6 +10,7 @@ export LITELLM_LOCAL_MODEL_COST_MAP=True
|
|||
|
||||
import json
|
||||
import os
|
||||
from datetime import datetime, timezone
|
||||
from importlib.resources import files
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
|
|
@ -169,6 +170,7 @@ class ModelCostMapSourceInfo:
|
|||
url: Optional[str] = None
|
||||
is_env_forced: bool = False
|
||||
fallback_reason: Optional[str] = None
|
||||
loaded_at: "datetime | None" = None
|
||||
|
||||
|
||||
# Module-level singleton tracking the source of the current cost map
|
||||
|
|
@ -193,6 +195,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.
|
||||
|
|
@ -272,6 +279,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":
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
"""
|
||||
Persistence for the admin-configured periodic reload schedules (model cost map,
|
||||
Anthropic beta headers) stored in ``LiteLLM_Config``.
|
||||
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
|
||||
|
|
@ -34,7 +34,6 @@ if TYPE_CHECKING:
|
|||
from prisma.models import LiteLLM_Config
|
||||
|
||||
MODEL_COST_MAP_RELOAD_PARAM_NAME = "model_cost_map_reload_config"
|
||||
ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME = "anthropic_beta_headers_reload_config"
|
||||
|
||||
|
||||
class _ConfigRowWrite(TypedDict, total=False):
|
||||
|
|
|
|||
|
|
@ -327,7 +327,6 @@ from litellm.proxy.common_utils.openai_endpoint_utils import (
|
|||
remove_sensitive_info_from_deployment,
|
||||
)
|
||||
from litellm.proxy.common_utils.periodic_reload_schedule import (
|
||||
ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME,
|
||||
MODEL_COST_MAP_RELOAD_PARAM_NAME,
|
||||
pod_reload_is_due,
|
||||
read_reload_schedule,
|
||||
|
|
@ -2067,6 +2066,9 @@ celery_fn = None # Redis Queue for handling requests
|
|||
|
||||
scheduler = None
|
||||
|
||||
# Global variable for anthropic beta headers reload scheduling
|
||||
last_anthropic_beta_headers_reload = None
|
||||
|
||||
|
||||
### DB WRITER ###
|
||||
db_writer_client: Optional[AsyncHTTPHandler] = None
|
||||
|
|
@ -3856,8 +3858,11 @@ class ProxyConfig:
|
|||
self._last_semantic_filter_config: Optional[Dict[str, Any]] = None
|
||||
self._last_hashicorp_vault_config: Optional[Dict[str, Any]] = None
|
||||
self.worker_registry: List["WorkerRegistryEntry"] = []
|
||||
self.model_cost_map_loaded_at: datetime = utc_now()
|
||||
self.anthropic_beta_headers_loaded_at: datetime = utc_now()
|
||||
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()
|
||||
|
||||
def is_yaml(self, config_file_path: str) -> bool:
|
||||
if not os.path.isfile(config_file_path):
|
||||
|
|
@ -6197,6 +6202,7 @@ class ProxyConfig:
|
|||
"router_settings",
|
||||
"litellm_settings",
|
||||
"environment_variables",
|
||||
"anthropic_beta_headers_reload_config",
|
||||
],
|
||||
)
|
||||
|
||||
|
|
@ -6259,6 +6265,9 @@ 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="anthropic_beta_headers"):
|
||||
await self._check_and_reload_anthropic_beta_headers(prisma_client=prisma_client)
|
||||
|
||||
if self._should_load_db_object(object_type="sso_settings"):
|
||||
await self._init_sso_settings_in_db(prisma_client=prisma_client)
|
||||
if self._should_load_db_object(object_type="cache_settings"):
|
||||
|
|
@ -6419,17 +6428,14 @@ class ProxyConfig:
|
|||
|
||||
async def check_periodic_reloads(self, prisma_client: PrismaClient):
|
||||
"""
|
||||
Run the admin-configured periodic reloads (model cost map, Anthropic beta headers).
|
||||
Run the admin-configured periodic model cost map reload.
|
||||
|
||||
Scheduled on its own job so schedules configured from the Admin UI fire whether
|
||||
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)
|
||||
|
||||
if self._should_load_db_object(object_type="anthropic_beta_headers"):
|
||||
await self._check_and_reload_anthropic_beta_headers(prisma_client=prisma_client)
|
||||
|
||||
async def _check_and_reload_model_cost_map(self, prisma_client: PrismaClient):
|
||||
"""
|
||||
Check if model cost map needs to be reloaded based on database configuration.
|
||||
|
|
@ -6462,35 +6468,96 @@ class ProxyConfig:
|
|||
async def _check_and_reload_anthropic_beta_headers(self, prisma_client: PrismaClient):
|
||||
"""
|
||||
Check if anthropic beta headers config needs to be reloaded based on database configuration.
|
||||
Runs on the periodic reload job, independently of `store_model_in_db`.
|
||||
This function runs every 10 seconds as part of _init_non_llm_objects_in_db.
|
||||
"""
|
||||
try:
|
||||
schedule = await read_reload_schedule(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME)
|
||||
if schedule is None:
|
||||
return
|
||||
# Get anthropic beta headers reload configuration from database
|
||||
config_record = await get_config_param(prisma_client, "anthropic_beta_headers_reload_config")
|
||||
|
||||
current_time = utc_now()
|
||||
if not pod_reload_is_due(
|
||||
schedule=schedule,
|
||||
pod_data_loaded_at=self.anthropic_beta_headers_loaded_at,
|
||||
current_time=current_time,
|
||||
description="Anthropic beta headers",
|
||||
):
|
||||
return
|
||||
if config_record is None or config_record.param_value is None:
|
||||
return # No configuration found, skip reload
|
||||
|
||||
from litellm.anthropic_beta_headers_manager import (
|
||||
reload_beta_headers_config,
|
||||
)
|
||||
config = config_record.param_value
|
||||
interval_hours = config.get("interval_hours")
|
||||
force_reload = config.get("force_reload", False)
|
||||
|
||||
new_config = await asyncio.to_thread(reload_beta_headers_config)
|
||||
if interval_hours is None and force_reload is False:
|
||||
return # No interval configured, skip reload
|
||||
|
||||
self.anthropic_beta_headers_loaded_at = current_time
|
||||
await record_reload_run(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME, current_time)
|
||||
current_time = datetime.utcnow()
|
||||
|
||||
provider_count = sum(1 for k in new_config.keys() if k != "provider_aliases" and k != "description")
|
||||
verbose_proxy_logger.info(
|
||||
f"Anthropic beta headers config reloaded successfully. Providers: {provider_count}"
|
||||
)
|
||||
# Check if we need to reload based on interval or force reload
|
||||
should_reload = False
|
||||
|
||||
if force_reload:
|
||||
should_reload = True
|
||||
verbose_proxy_logger.info("Anthropic beta headers reload triggered by force reload flag")
|
||||
elif interval_hours is not None:
|
||||
# Use pod's in-memory last reload time
|
||||
global last_anthropic_beta_headers_reload
|
||||
if last_anthropic_beta_headers_reload is not None:
|
||||
try:
|
||||
last_reload_time = datetime.fromisoformat(last_anthropic_beta_headers_reload)
|
||||
time_since_last_reload = current_time - last_reload_time
|
||||
hours_since_last_reload = time_since_last_reload.total_seconds() / 3600
|
||||
|
||||
if hours_since_last_reload >= interval_hours:
|
||||
should_reload = True
|
||||
verbose_proxy_logger.info(
|
||||
f"Anthropic beta headers reload triggered by interval. Hours since last reload: {hours_since_last_reload:.2f}, Interval: {interval_hours}"
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(f"Error parsing last reload time: {e}")
|
||||
# If we can't parse the last reload time, reload anyway
|
||||
should_reload = True
|
||||
else:
|
||||
# No last reload time recorded, reload now
|
||||
should_reload = True
|
||||
verbose_proxy_logger.info(
|
||||
"Anthropic beta headers reload triggered - no previous reload time recorded"
|
||||
)
|
||||
|
||||
if should_reload:
|
||||
# Perform the reload
|
||||
from litellm.anthropic_beta_headers_manager import (
|
||||
reload_beta_headers_config,
|
||||
)
|
||||
|
||||
new_config = reload_beta_headers_config()
|
||||
|
||||
# Update pod's in-memory last reload time
|
||||
last_anthropic_beta_headers_reload = current_time.isoformat()
|
||||
|
||||
# Clear force reload flag in database
|
||||
await ConfigRepository(prisma_client).table.upsert(
|
||||
where={"param_name": "anthropic_beta_headers_reload_config"},
|
||||
data={
|
||||
"create": {
|
||||
"param_name": "anthropic_beta_headers_reload_config",
|
||||
"param_value": safe_dumps(
|
||||
{
|
||||
"interval_hours": interval_hours,
|
||||
"force_reload": False,
|
||||
}
|
||||
),
|
||||
},
|
||||
"update": {
|
||||
"param_value": safe_dumps(
|
||||
{
|
||||
"interval_hours": interval_hours,
|
||||
"force_reload": False,
|
||||
}
|
||||
)
|
||||
},
|
||||
},
|
||||
)
|
||||
await invalidate_config_param("anthropic_beta_headers_reload_config")
|
||||
|
||||
# Count providers in config
|
||||
provider_count = sum(1 for k in new_config.keys() if k != "provider_aliases" and k != "description")
|
||||
verbose_proxy_logger.info(
|
||||
f"Anthropic beta headers config reloaded successfully. Providers: {provider_count}"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(f"Error in _check_and_reload_anthropic_beta_headers: {str(e)}")
|
||||
|
|
@ -15797,13 +15864,32 @@ async def reload_anthropic_beta_headers(
|
|||
# Immediately reload the beta headers config in the current pod
|
||||
from litellm.anthropic_beta_headers_manager import reload_beta_headers_config
|
||||
|
||||
new_config = await asyncio.to_thread(reload_beta_headers_config)
|
||||
new_config = reload_beta_headers_config()
|
||||
|
||||
current_time = utc_now()
|
||||
proxy_config.anthropic_beta_headers_loaded_at = current_time
|
||||
# Update pod's in-memory last reload time
|
||||
global last_anthropic_beta_headers_reload
|
||||
current_time = datetime.utcnow()
|
||||
last_anthropic_beta_headers_reload = current_time.isoformat()
|
||||
|
||||
# Stamp the shared last run and the reload request so other pods reload too
|
||||
await record_manual_reload(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME, current_time)
|
||||
# Set force reload flag in database for other pods, preserving existing interval_hours
|
||||
existing_beta_config = await ConfigRepository(prisma_client).table.find_unique(
|
||||
where={"param_name": "anthropic_beta_headers_reload_config"}
|
||||
)
|
||||
existing_beta_interval = None
|
||||
if existing_beta_config and existing_beta_config.param_value:
|
||||
existing_beta_interval = existing_beta_config.param_value.get("interval_hours")
|
||||
|
||||
await ConfigRepository(prisma_client).table.upsert(
|
||||
where={"param_name": "anthropic_beta_headers_reload_config"},
|
||||
data={
|
||||
"create": {
|
||||
"param_name": "anthropic_beta_headers_reload_config",
|
||||
"param_value": safe_dumps({"interval_hours": None, "force_reload": True}),
|
||||
},
|
||||
"update": {"param_value": safe_dumps({"interval_hours": existing_beta_interval, "force_reload": True})},
|
||||
},
|
||||
)
|
||||
await invalidate_config_param("anthropic_beta_headers_reload_config")
|
||||
|
||||
provider_count = sum(1 for k in new_config.keys() if k not in ["provider_aliases", "description"])
|
||||
verbose_proxy_logger.info(
|
||||
|
|
@ -15852,7 +15938,18 @@ async def schedule_anthropic_beta_headers_reload(
|
|||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail="Database connection not available")
|
||||
|
||||
await write_reload_interval(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME, hours)
|
||||
# Update database with new reload configuration
|
||||
await ConfigRepository(prisma_client).table.upsert(
|
||||
where={"param_name": "anthropic_beta_headers_reload_config"},
|
||||
data={
|
||||
"create": {
|
||||
"param_name": "anthropic_beta_headers_reload_config",
|
||||
"param_value": safe_dumps({"interval_hours": hours, "force_reload": False}),
|
||||
},
|
||||
"update": {"param_value": safe_dumps({"interval_hours": hours, "force_reload": False})},
|
||||
},
|
||||
)
|
||||
await invalidate_config_param("anthropic_beta_headers_reload_config")
|
||||
|
||||
verbose_proxy_logger.info(f"Anthropic beta headers reload scheduled for every {hours} hours")
|
||||
|
||||
|
|
@ -15860,7 +15957,7 @@ async def schedule_anthropic_beta_headers_reload(
|
|||
"message": f"Anthropic beta headers reload scheduled for every {hours} hours",
|
||||
"status": "success",
|
||||
"interval_hours": hours,
|
||||
"timestamp": utc_now().isoformat(),
|
||||
"timestamp": datetime.utcnow().isoformat(),
|
||||
}
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(f"Failed to schedule anthropic beta headers reload: {str(e)}")
|
||||
|
|
@ -15897,17 +15994,15 @@ async def cancel_anthropic_beta_headers_reload(
|
|||
raise HTTPException(status_code=500, detail="Database connection not available")
|
||||
|
||||
# Remove reload configuration from database
|
||||
await ConfigRepository(prisma_client).table.delete(
|
||||
where={"param_name": ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME}
|
||||
)
|
||||
await invalidate_config_param(ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME)
|
||||
await ConfigRepository(prisma_client).table.delete(where={"param_name": "anthropic_beta_headers_reload_config"})
|
||||
await invalidate_config_param("anthropic_beta_headers_reload_config")
|
||||
|
||||
verbose_proxy_logger.info("Anthropic beta headers reload schedule cancelled")
|
||||
|
||||
return {
|
||||
"message": "Anthropic beta headers reload schedule cancelled",
|
||||
"status": "success",
|
||||
"timestamp": utc_now().isoformat(),
|
||||
"timestamp": datetime.utcnow().isoformat(),
|
||||
}
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(f"Failed to cancel anthropic beta headers reload: {str(e)}")
|
||||
|
|
@ -15939,15 +16034,68 @@ async def get_anthropic_beta_headers_reload_status(
|
|||
)
|
||||
|
||||
try:
|
||||
global prisma_client
|
||||
global prisma_client, last_anthropic_beta_headers_reload
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
f"Checking anthropic beta headers reload status. Last reload: {last_anthropic_beta_headers_reload}"
|
||||
)
|
||||
|
||||
if prisma_client is None:
|
||||
verbose_proxy_logger.info("No database connection, returning not scheduled")
|
||||
return reload_schedule_status(None)
|
||||
return {
|
||||
"scheduled": False,
|
||||
"interval_hours": None,
|
||||
"last_run": None,
|
||||
"next_run": None,
|
||||
}
|
||||
|
||||
return reload_schedule_status(
|
||||
await read_reload_schedule(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME)
|
||||
# Get reload configuration from database
|
||||
config_record = await ConfigRepository(prisma_client).table.find_unique(
|
||||
where={"param_name": "anthropic_beta_headers_reload_config"}
|
||||
)
|
||||
|
||||
if config_record is None or config_record.param_value is None:
|
||||
verbose_proxy_logger.info("No anthropic beta headers reload configuration found")
|
||||
return {
|
||||
"scheduled": False,
|
||||
"interval_hours": None,
|
||||
"last_run": None,
|
||||
"next_run": None,
|
||||
}
|
||||
|
||||
config = config_record.param_value
|
||||
interval_hours = config.get("interval_hours")
|
||||
|
||||
if interval_hours is None:
|
||||
verbose_proxy_logger.info("No interval configured, returning not scheduled")
|
||||
return {
|
||||
"scheduled": False,
|
||||
"interval_hours": None,
|
||||
"last_run": None,
|
||||
"next_run": None,
|
||||
}
|
||||
|
||||
current_time = datetime.utcnow()
|
||||
next_run = None
|
||||
|
||||
# Use pod's in-memory last reload time
|
||||
if last_anthropic_beta_headers_reload is not None:
|
||||
try:
|
||||
last_reload_time = datetime.fromisoformat(last_anthropic_beta_headers_reload)
|
||||
time_since_last_reload = current_time - last_reload_time
|
||||
hours_since_last_reload = time_since_last_reload.total_seconds() / 3600
|
||||
|
||||
if hours_since_last_reload < interval_hours:
|
||||
next_run = (last_reload_time + timedelta(hours=interval_hours)).isoformat()
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(f"Error parsing last reload time: {e}")
|
||||
|
||||
return {
|
||||
"scheduled": True,
|
||||
"interval_hours": interval_hours,
|
||||
"last_run": last_anthropic_beta_headers_reload,
|
||||
"next_run": next_run,
|
||||
}
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(f"Failed to get anthropic beta headers reload status: {str(e)}")
|
||||
raise HTTPException(
|
||||
|
|
|
|||
|
|
@ -248,3 +248,25 @@ def test_azure_ai_claude_1m_context_entries(cost_map: dict):
|
|||
"azure_ai/claude-haiku-4-5",
|
||||
]:
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -9,9 +9,11 @@ Routes covered:
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from .conftest import VOLATILE_KEYS, normalize
|
||||
|
||||
|
|
@ -97,24 +99,22 @@ def test_reload_anthropic_beta_headers_admin_success(client, auth_as, monkeypatc
|
|||
"providers_count": 2,
|
||||
"timestamp": "<VOLATILE>",
|
||||
}
|
||||
# And the upsert was actually invoked (force_reload write).
|
||||
prisma.db.litellm_config.upsert.assert_awaited_once()
|
||||
|
||||
|
||||
def test_reload_anthropic_beta_headers_preserves_existing_interval(
|
||||
client, auth_as, monkeypatch
|
||||
):
|
||||
"""The manual-reload write only stamps the job-owned timestamp columns; param_value
|
||||
(the admin-owned interval) is not in the payload, so it cannot be clobbered."""
|
||||
from datetime import datetime
|
||||
|
||||
"""When an existing reload config has an interval set, the force-reload
|
||||
write must preserve that interval (the route reads it back then upserts
|
||||
with the same number). This pins the read-then-write behaviour."""
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
_stub_reload_beta_headers(monkeypatch)
|
||||
existing = SimpleNamespace(
|
||||
param_name="anthropic_beta_headers_reload_config",
|
||||
param_value={"interval_hours": 12},
|
||||
reload_requested_at=None,
|
||||
last_run_at=None,
|
||||
param_value={"interval_hours": 12, "force_reload": False},
|
||||
)
|
||||
prisma = _make_prisma_with_config(config_record=existing)
|
||||
_install_prisma(monkeypatch, prisma)
|
||||
|
|
@ -123,12 +123,17 @@ def test_reload_anthropic_beta_headers_preserves_existing_interval(
|
|||
response = client.post("/reload/anthropic_beta_headers")
|
||||
|
||||
assert response.status_code == 200
|
||||
# The update branch's interval_hours was sourced from the existing record.
|
||||
call_kwargs = prisma.db.litellm_config.upsert.await_args.kwargs
|
||||
update_payload = call_kwargs["data"]["update"]
|
||||
assert set(update_payload) == {"last_run_at", "reload_requested_at"}
|
||||
assert isinstance(update_payload["last_run_at"], datetime)
|
||||
assert update_payload["reload_requested_at"] == update_payload["last_run_at"]
|
||||
assert "param_value" not in call_kwargs["data"]["create"]
|
||||
data = call_kwargs["data"]
|
||||
update_payload = data["update"]["param_value"]
|
||||
parsed = (
|
||||
json.loads(update_payload)
|
||||
if isinstance(update_payload, str)
|
||||
else update_payload
|
||||
)
|
||||
assert parsed["interval_hours"] == 12
|
||||
assert parsed["force_reload"] is True
|
||||
|
||||
|
||||
def test_reload_anthropic_beta_headers_not_admin_forbidden(client, auth_as):
|
||||
|
|
@ -286,16 +291,17 @@ def test_get_anthropic_beta_headers_reload_status_scheduled(
|
|||
):
|
||||
"""When a config row with ``interval_hours`` is present, ``scheduled`` is True
|
||||
and ``interval_hours`` echoes the DB value. Pins the full response shape."""
|
||||
from litellm.proxy import proxy_server as ps
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
record = SimpleNamespace(
|
||||
param_name="anthropic_beta_headers_reload_config",
|
||||
param_value={"interval_hours": 6},
|
||||
reload_requested_at=None,
|
||||
last_run_at=None,
|
||||
param_value={"interval_hours": 6, "force_reload": False},
|
||||
)
|
||||
prisma = _make_prisma_with_config(config_record=record)
|
||||
_install_prisma(monkeypatch, prisma)
|
||||
# No prior reload — next_run stays None.
|
||||
monkeypatch.setattr(ps, "last_anthropic_beta_headers_reload", None)
|
||||
|
||||
with auth_as(LitellmUserRoles.PROXY_ADMIN):
|
||||
response = client.get("/schedule/anthropic_beta_headers_reload/status")
|
||||
|
|
@ -309,35 +315,6 @@ def test_get_anthropic_beta_headers_reload_status_scheduled(
|
|||
}
|
||||
|
||||
|
||||
def test_get_anthropic_beta_headers_reload_status_reports_persisted_last_run(
|
||||
client, auth_as, monkeypatch
|
||||
):
|
||||
"""last_run/next_run come from the DB row, so status survives pod restarts."""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
record = SimpleNamespace(
|
||||
param_name="anthropic_beta_headers_reload_config",
|
||||
param_value={"interval_hours": 6},
|
||||
reload_requested_at=None,
|
||||
last_run_at=datetime(2024, 1, 1, 6, 0, 0, tzinfo=timezone.utc),
|
||||
)
|
||||
prisma = _make_prisma_with_config(config_record=record)
|
||||
_install_prisma(monkeypatch, prisma)
|
||||
|
||||
with auth_as(LitellmUserRoles.PROXY_ADMIN):
|
||||
response = client.get("/schedule/anthropic_beta_headers_reload/status")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert normalize(response.json()) == {
|
||||
"scheduled": True,
|
||||
"interval_hours": 6,
|
||||
"last_run": "2024-01-01T06:00:00+00:00",
|
||||
"next_run": "2024-01-01T12:00:00+00:00",
|
||||
}
|
||||
|
||||
|
||||
def test_get_anthropic_beta_headers_reload_status_not_scheduled_no_db(
|
||||
client, auth_as, monkeypatch
|
||||
):
|
||||
|
|
@ -362,15 +339,11 @@ def test_get_anthropic_beta_headers_reload_status_no_interval_unscheduled(
|
|||
client, auth_as, monkeypatch
|
||||
):
|
||||
"""Config row present but ``interval_hours`` is None → unscheduled response."""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
record = SimpleNamespace(
|
||||
param_name="anthropic_beta_headers_reload_config",
|
||||
param_value={"interval_hours": None},
|
||||
reload_requested_at=datetime(2024, 1, 1, 6, 0, 0, tzinfo=timezone.utc),
|
||||
last_run_at=None,
|
||||
param_value={"interval_hours": None, "force_reload": True},
|
||||
)
|
||||
prisma = _make_prisma_with_config(config_record=record)
|
||||
_install_prisma(monkeypatch, prisma)
|
||||
|
|
|
|||
|
|
@ -3972,6 +3972,20 @@ class TestPriceDataReloadIntegration:
|
|||
litellm.model_cost = original_model_cost
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
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
|
||||
|
||||
fetch_time = datetime(2024, 1, 1, 6, 0, tzinfo=timezone.utc)
|
||||
with patch(
|
||||
"litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map_loaded_at",
|
||||
return_value=fetch_time,
|
||||
):
|
||||
assert ProxyConfig().model_cost_map_loaded_at == fetch_time
|
||||
|
||||
def test_distributed_reload_check_function(self):
|
||||
"""
|
||||
A reload requested by another pod after this pod's last reload takes effect here
|
||||
|
|
@ -4215,39 +4229,52 @@ class TestPriceDataReloadIntegration:
|
|||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
def test_anthropic_beta_headers_reload_preserves_interval_hours(self):
|
||||
"""Same column-only write contract as the cost map job, for the beta headers schedule"""
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
"""Test that _check_and_reload_anthropic_beta_headers preserves interval_hours after reload.
|
||||
|
||||
Regression test: the update branch of the upsert was dropping interval_hours,
|
||||
identical to the model cost map bug.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
from litellm.proxy.utils import litellm_config_cache
|
||||
|
||||
litellm_config_cache.flush_cache()
|
||||
proxy_config = ProxyConfig()
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(
|
||||
return_value=_reload_schedule_row({"interval_hours": 12})
|
||||
)
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
|
||||
mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None)
|
||||
frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.anthropic_beta_headers_manager.reload_beta_headers_config"
|
||||
) as mock_reload,
|
||||
patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now),
|
||||
):
|
||||
# Set up config with interval_hours=12 and force_reload=True to trigger reload
|
||||
mock_config = MagicMock()
|
||||
mock_config.param_value = {"interval_hours": 12, "force_reload": True}
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config)
|
||||
# _check_and_reload_anthropic_beta_headers now reads through get_generic_data.
|
||||
mock_prisma.get_generic_data = AsyncMock(return_value=mock_config)
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
|
||||
|
||||
with patch(
|
||||
"litellm.anthropic_beta_headers_manager.reload_beta_headers_config"
|
||||
) as mock_reload:
|
||||
mock_reload.return_value = {"anthropic": {"beta_header": "test-value"}}
|
||||
|
||||
asyncio.run(
|
||||
proxy_config._check_and_reload_anthropic_beta_headers(mock_prisma)
|
||||
)
|
||||
|
||||
assert proxy_config.anthropic_beta_headers_loaded_at == frozen_now
|
||||
assert mock_prisma.db.litellm_config.update_many.call_args[1] == {
|
||||
"data": {"last_run_at": frozen_now},
|
||||
"where": {"param_name": "anthropic_beta_headers_reload_config"},
|
||||
}
|
||||
mock_prisma.db.litellm_config.upsert.assert_not_called()
|
||||
# Verify the upsert update branch preserves interval_hours
|
||||
mock_prisma.db.litellm_config.upsert.assert_called()
|
||||
call_args = mock_prisma.db.litellm_config.upsert.call_args
|
||||
param_value_json = call_args[1]["data"]["update"]["param_value"]
|
||||
param_value_dict = json.loads(param_value_json)
|
||||
assert param_value_dict["force_reload"] == False
|
||||
assert param_value_dict["interval_hours"] == 12, (
|
||||
"interval_hours must be preserved in the update branch; "
|
||||
"dropping it causes the schedule to self-destruct"
|
||||
)
|
||||
|
||||
def test_anthropic_beta_headers_manual_reload_preserves_interval_hours(self):
|
||||
"""Manual beta headers reload stamps the run columns only, leaving the schedule intact"""
|
||||
"""Test that manual reload via /reload/anthropic_beta_headers preserves existing interval_hours.
|
||||
|
||||
Regression test: the manual reload endpoint was overwriting param_value with
|
||||
only force_reload=True, dropping any existing interval_hours schedule.
|
||||
"""
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
from litellm.proxy.proxy_server import cleanup_router_config_variables
|
||||
|
||||
|
|
@ -4260,36 +4287,34 @@ class TestPriceDataReloadIntegration:
|
|||
mock_auth.user_role = LitellmUserRoles.PROXY_ADMIN
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: mock_auth
|
||||
client = TestClient(app)
|
||||
frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.anthropic_beta_headers_manager.reload_beta_headers_config"
|
||||
) as mock_reload,
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now),
|
||||
):
|
||||
with patch(
|
||||
"litellm.anthropic_beta_headers_manager.reload_beta_headers_config"
|
||||
) as mock_reload:
|
||||
mock_reload.return_value = {"anthropic": {"beta_header": "test-value"}}
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
|
||||
|
||||
response = client.post("/reload/anthropic_beta_headers")
|
||||
assert response.status_code == 200
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
|
||||
# Simulate existing config with a schedule
|
||||
mock_existing = MagicMock()
|
||||
mock_existing.param_value = {"interval_hours": 8, "force_reload": False}
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(
|
||||
return_value=mock_existing
|
||||
)
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
|
||||
|
||||
mock_prisma.db.litellm_config.find_unique.assert_not_called()
|
||||
call_args = mock_prisma.db.litellm_config.upsert.call_args
|
||||
assert call_args[1]["data"]["update"] == {
|
||||
"last_run_at": frozen_now,
|
||||
"reload_requested_at": frozen_now,
|
||||
}
|
||||
assert call_args[1]["data"]["create"] == {
|
||||
"param_name": "anthropic_beta_headers_reload_config",
|
||||
"last_run_at": frozen_now,
|
||||
"reload_requested_at": frozen_now,
|
||||
}
|
||||
assert (
|
||||
proxy_server_module.proxy_config.anthropic_beta_headers_loaded_at
|
||||
== frozen_now
|
||||
)
|
||||
response = client.post("/reload/anthropic_beta_headers")
|
||||
assert response.status_code == 200
|
||||
|
||||
# Verify interval_hours was preserved in the upsert
|
||||
mock_prisma.db.litellm_config.upsert.assert_called()
|
||||
call_args = mock_prisma.db.litellm_config.upsert.call_args
|
||||
param_value_json = call_args[1]["data"]["update"]["param_value"]
|
||||
param_value_dict = json.loads(param_value_json)
|
||||
assert param_value_dict["force_reload"] == True
|
||||
assert param_value_dict["interval_hours"] == 8, (
|
||||
"interval_hours must be preserved when manual reload sets force_reload; "
|
||||
"dropping it destroys any existing schedule"
|
||||
)
|
||||
|
||||
def test_config_file_parsing(self):
|
||||
"""Test parsing of config file with reload settings"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue