Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/common_utils/periodic_reload_schedule.py: 48%
98 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""
2Persistence for the admin-configured periodic model cost map reload schedule stored in
3``LiteLLM_Config``.
5Field ownership is split by writer so concurrent writers never overwrite each other:
6the schedule endpoints own the ``param_value`` JSON (``interval_hours``), while the
7reload job and the manual reload endpoints own the dedicated ``last_run_at`` /
8``reload_revision`` columns. ``last_run_at`` lives in the row rather than process memory
9so the Admin UI still reports the last execution after a restart and across pods.
10``reload_revision`` is a monotonic counter a manual reload increments; each pod records
11the revision it last applied and reloads whenever the row's differs, so a request reaches
12every pod exactly once without any pod clearing it and without comparing clocks. A booting
13pod starts at revision 0 rather than adopting the published one, because it cannot know
14whether that request predates the prices it fetched at import. Interval reloads stay
15per-pod, driven by when that pod's own copy of the data was loaded.
16"""
18from collections.abc import Mapping
19from dataclasses import dataclass
20from datetime import datetime, timedelta, timezone
21from typing import (
22 TYPE_CHECKING,
23 Protocol,
24 TypedDict,
25 cast, # noqa: TID251 # prisma table access is untyped (PrismaWrapper.__getattr__)
26)
28from pydantic import BaseModel, ConfigDict, ValidationError
30from litellm._logging import verbose_proxy_logger
31from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
32from litellm.proxy.utils import PrismaClient, evict_config_param
33from litellm.repositories.config_repository import ConfigRepository
35if TYPE_CHECKING: 35 ↛ 36line 35 didn't jump to line 36 because the condition on line 35 was never true
36 from prisma.models import LiteLLM_Config
38MODEL_COST_MAP_RELOAD_PARAM_NAME = "model_cost_map_reload_config"
41class _RevisionIncrement(TypedDict):
42 increment: int
45class _ConfigRowWrite(TypedDict, total=False):
46 param_name: str
47 param_value: str
48 last_run_at: datetime
49 reload_revision: int | _RevisionIncrement
52class _ConfigUpsertData(TypedDict):
53 create: _ConfigRowWrite
54 update: _ConfigRowWrite
57class _ConfigTable(Protocol):
58 async def find_unique(self, where: Mapping[str, str]) -> "LiteLLM_Config | None": ... 58 ↛ exitline 58 didn't return from function 'find_unique' because
60 async def upsert(self, where: Mapping[str, str], data: _ConfigUpsertData) -> "LiteLLM_Config": ... 60 ↛ exitline 60 didn't return from function 'upsert' because
62 async def update_many(self, data: _ConfigRowWrite, where: Mapping[str, str]) -> int: ... 62 ↛ exitline 62 didn't return from function 'update_many' because
65def _config_table(prisma_client: PrismaClient) -> _ConfigTable:
66 return cast(_ConfigTable, ConfigRepository(prisma_client).table) # cast-ok: prisma table is untyped (Any)
69@dataclass(frozen=True, slots=True)
70class ReloadSchedule:
71 interval_hours: int | None = None
72 reload_revision: int = 0
73 last_run_at: datetime | None = None
76class ReloadScheduleStatus(TypedDict):
77 scheduled: bool
78 interval_hours: int | None
79 last_run: str | None
80 next_run: str | None
83class _IntervalConfig(BaseModel):
84 model_config = ConfigDict(strict=True)
86 interval_hours: int | None = None
89def utc_now() -> datetime:
90 return datetime.now(timezone.utc)
93def _parse_interval_hours(param_value: object) -> int | None:
94 """``param_value`` is written as serialized JSON, and a raw row read can hand it back
95 either decoded or still as a string depending on the driver, so accept both rather than
96 reading a string as no schedule at all. Mirrors ``ConfigRepository.get_param``"""
97 try:
98 if isinstance(param_value, str):
99 return _IntervalConfig.model_validate_json(param_value).interval_hours
100 return _IntervalConfig.model_validate(param_value).interval_hours
101 except ValidationError:
102 return None
105def _as_utc(value: datetime | None) -> datetime | None:
106 if value is None:
107 return None
108 return value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc)
111def parse_reload_schedule(row: "LiteLLM_Config") -> ReloadSchedule:
112 return ReloadSchedule(
113 interval_hours=_parse_interval_hours(row.param_value),
114 reload_revision=int(row.reload_revision or 0),
115 last_run_at=_as_utc(row.last_run_at),
116 )
119def next_run_at(schedule: ReloadSchedule) -> datetime | None:
120 if schedule.interval_hours is None or schedule.last_run_at is None:
121 return None
122 return schedule.last_run_at + timedelta(hours=schedule.interval_hours)
125def reload_schedule_status(schedule: ReloadSchedule | None) -> ReloadScheduleStatus:
126 if schedule is None:
127 return {"scheduled": False, "interval_hours": None, "last_run": None, "next_run": None}
128 next_run = next_run_at(schedule)
129 return {
130 "scheduled": schedule.interval_hours is not None,
131 "interval_hours": schedule.interval_hours,
132 "last_run": schedule.last_run_at.isoformat() if schedule.last_run_at is not None else None,
133 "next_run": next_run.isoformat() if next_run is not None else None,
134 }
137def pod_reload_is_due(
138 *,
139 schedule: ReloadSchedule,
140 pod_applied_revision: int,
141 pod_data_loaded_at: datetime,
142 current_time: datetime,
143 description: str,
144) -> bool:
145 """
146 Whether this pod should reload now. A revision it has not applied means a manual reload
147 it has not served. A pod starts at revision 0, so it serves any request published before
148 it booted; that costs one redundant fetch per boot and is what keeps a request from being
149 marked applied against data fetched before it. Interval reloads compare against this pod's
150 own data, and a schedule that has never run anywhere fires immediately rather than one
151 interval later
152 """
153 if schedule.reload_revision != pod_applied_revision:
154 verbose_proxy_logger.info("%s reload triggered by manual reload request", description)
155 return True
156 if schedule.interval_hours is None:
157 return False
158 if schedule.last_run_at is None:
159 verbose_proxy_logger.info("%s reload triggered - schedule has never run", description)
160 return True
161 hours_since_data_loaded = (current_time - pod_data_loaded_at).total_seconds() / 3600
162 if hours_since_data_loaded < schedule.interval_hours:
163 return False
164 verbose_proxy_logger.info(
165 "%s reload triggered by interval. Hours since data loaded: %.2f, Interval: %s",
166 description,
167 hours_since_data_loaded,
168 schedule.interval_hours,
169 )
170 return True
173async def read_reload_schedule(prisma_client: PrismaClient, param_name: str) -> ReloadSchedule | None:
174 row = await _config_table(prisma_client).find_unique(where={"param_name": param_name})
175 if row is None: 175 ↛ 177line 175 didn't jump to line 177 because the condition on line 175 was always true
176 return None
177 return parse_reload_schedule(row)
180async def write_reload_interval(prisma_client: PrismaClient, param_name: str, interval_hours: int) -> None:
181 """Admin-owned write: replaces ``param_value`` without touching the job-owned columns"""
182 param_value = safe_dumps({"interval_hours": interval_hours})
183 await _config_table(prisma_client).upsert(
184 where={"param_name": param_name},
185 data={
186 "create": {"param_name": param_name, "param_value": param_value},
187 "update": {"param_value": param_value},
188 },
189 )
190 await evict_config_param(param_name)
193async def clear_reload_interval(prisma_client: PrismaClient, param_name: str) -> None:
194 """Admin-owned write: drops the schedule but keeps the row, because the revision counter
195 identifies a request rather than ordering one and so can never reuse a number. Deleting
196 the row restarts it, and a reissued revision matches what pods already applied, so their
197 next manual reload is silently skipped. The interval is nulled inside the JSON rather
198 than by nulling the column, which prisma rejects for a ``Json?`` field"""
199 await _config_table(prisma_client).update_many(
200 data={"param_value": safe_dumps({"interval_hours": None})},
201 where={"param_name": param_name},
202 )
203 await evict_config_param(param_name)
206async def record_reload_run(prisma_client: PrismaClient, param_name: str, ran_at: datetime) -> None:
207 """Job-owned write after this pod reloaded: stamps the shared last run only if the row
208 still exists, so a schedule deleted mid-poll is not resurrected"""
209 await _config_table(prisma_client).update_many(
210 data={"last_run_at": ran_at},
211 where={"param_name": param_name},
212 )
213 await evict_config_param(param_name)
216async def record_manual_reload(prisma_client: PrismaClient, param_name: str, ran_at: datetime) -> int:
217 """
218 After a manual in-pod reload: stamp the shared last run and bump the revision every other
219 pod compares against. The increment is atomic, so concurrent requests each publish a
220 distinct revision instead of overwriting one another. Returns the published revision so
221 the serving pod can adopt it rather than reloading again on its next poll
222 """
223 row = await _config_table(prisma_client).upsert(
224 where={"param_name": param_name},
225 data={
226 "create": {"param_name": param_name, "last_run_at": ran_at, "reload_revision": 1},
227 "update": {"last_run_at": ran_at, "reload_revision": {"increment": 1}},
228 },
229 )
230 await evict_config_param(param_name)
231 return int(row.reload_revision)