Merge pull request #36288 from BerriAI/litellm_sync_main_into_internal_staging

chore(ci): sync main into internal staging
This commit is contained in:
yuneng-jiang 2026-08-08 11:25:43 -07:00 committed by GitHub
commit b0fd3e1e30
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 150 additions and 24 deletions

View file

@ -76,12 +76,127 @@ from transport import HttpTransport, SplitTransport, Transport
RowsPredicate = Callable[[list[SpendLogRow]], bool]
# After /model/new, poll data-plane /v1/models until the model is listed (or fail).
# Bound by MODEL_SERVABLE_TIMEOUT so a stuck reload does not burn the spend
# poll_timeout (120s). Return on first listing; settle_propagation owns the separate
# wait that lets every worker and replica reload before the caller uses the model.
MODEL_SERVABLE_TIMEOUT = 40.0
MODEL_SERVABLE_DB_SYNC_SECONDS = 0.0
MODEL_SERVABLE_INTERVAL = 2.0
# Cap each /v1/models poll so one slow request cannot outlast the remaining budget.
MODEL_SERVABLE_REQUEST_TIMEOUT = 5.0
@dataclass(frozen=True, slots=True)
class Servable:
"""The data plane listed the model within the deadline."""
@dataclass(frozen=True, slots=True)
class NotServable:
"""The deadline passed without the data plane listing the model.
`last_result` is the final /v1/models read, so the caller can tell "the proxy
answered but omitted the model" (propagation) from "the read itself failed"
(network/auth) when reporting."""
last_result: Result[ModelsListResponse] | None
ServableOutcome = Servable | NotServable
def await_servable(
list_models: Callable[[float], Result[ModelsListResponse]],
*,
model_name: str,
timeout: float,
interval: float,
request_timeout: float,
db_sync_seconds: float,
now: Callable[[], float],
sleep: Callable[[float], None],
) -> ServableOutcome:
"""Poll until `model_name` is listed long enough for every worker to DB-sync.
First listing must happen within `timeout`. After that, the model must stay
listed continuously for `db_sync_seconds` (any miss resets the continuous
window). `db_sync_seconds=0` returns on the first listing. Each poll's request
timeout is clamped to the remaining budget. Sleeps only min(interval, time left)
so a final deadline-clamped poll is never skipped just because a full interval
does not fit. Clock and sleep are injected."""
started = now()
first_seen_at: float | None = None
last_result: Result[ModelsListResponse] | None = None
while True:
t = now()
phase_deadline = (
started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds
)
remaining = phase_deadline - t
if remaining <= 0:
if (
last_result is not None
and first_seen_at is not None
and (db_sync_seconds <= 0 or t - first_seen_at >= db_sync_seconds)
):
return Servable()
return NotServable(last_result=last_result)
poll_timeout = min(request_timeout, remaining)
last_result = list_models(poll_timeout)
listed = isinstance(last_result, Success) and any(
entry.id == model_name for entry in last_result.data.data
)
t = now()
if not listed:
first_seen_at = None
elif first_seen_at is None:
if t > started + timeout:
return NotServable(last_result=last_result)
first_seen_at = t
if db_sync_seconds <= 0:
return Servable()
elif t - first_seen_at >= db_sync_seconds:
return Servable()
phase_deadline = (
started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds
)
wait = min(interval, phase_deadline - now())
if wait > 0:
sleep(wait)
def servable_timeout_message(
*,
model_name: str,
timeout: float,
db_sync_seconds: float,
last_result: Result[ModelsListResponse] | None,
) -> str:
last_error = (
f"; last /v1/models poll did not succeed: {last_result}"
if last_result is not None and not isinstance(last_result, Success)
else ""
)
return (
f"model {model_name!r} was created but never became servable on the data "
f"plane within {timeout}s of first listing (plus {db_sync_seconds}s continuous "
f"DB sync) after /model/new (control/data-plane propagation or "
f"STORE_MODEL_IN_DB reload issue){last_error}"
)
@dataclass(frozen=True, slots=True)
class ProxyClient:
transport: Transport
poll_timeout: float = 120.0
poll_interval: float = 5.0
model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT
model_servable_db_sync_seconds: float = MODEL_SERVABLE_DB_SYNC_SECONDS
model_servable_interval: float = MODEL_SERVABLE_INTERVAL
model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT
# ---- keys / customers (satisfies lifecycle.ResourceClient) ----------
@ -192,33 +307,35 @@ class ProxyClient:
return model_id
def _await_model_servable(self, model_name: str) -> None:
"""Block until the data plane lists `model_name`, or fail loudly if it does
not within poll_timeout (a real propagation/config problem, surfaced here
instead of as a downstream "Invalid model name passed")."""
deadline = time.monotonic() + self.poll_timeout
last_result: Result[ModelsListResponse] | None = None
while time.monotonic() < deadline:
last_result = self.transport.get(
"""Block until the data plane lists `model_name`, or fail at model_servable_timeout."""
outcome = await_servable(
lambda poll_timeout: self.transport.get(
"/v1/models",
headers=self.transport.master,
params=NoBody(),
response_type=ModelsListResponse,
)
if isinstance(last_result, Success) and any(
entry.id == model_name for entry in last_result.data.data
):
timeout=poll_timeout,
),
model_name=model_name,
timeout=self.model_servable_timeout,
interval=self.model_servable_interval,
request_timeout=self.model_servable_request_timeout,
db_sync_seconds=self.model_servable_db_sync_seconds,
now=time.monotonic,
sleep=time.sleep,
)
match outcome:
case Servable():
return
time.sleep(self.poll_interval)
last_error = (
f"; last /v1/models poll did not succeed: {last_result}"
if last_result is not None and not isinstance(last_result, Success)
else ""
)
raise AssertionError(
f"model {model_name!r} was created but never became servable on the data "
f"plane within {self.poll_timeout}s of /model/new (control/data-plane "
f"propagation or STORE_MODEL_IN_DB reload issue){last_error}"
)
case NotServable(last_result=last_result):
raise AssertionError(
servable_timeout_message(
model_name=model_name,
timeout=self.model_servable_timeout,
db_sync_seconds=self.model_servable_db_sync_seconds,
last_result=last_result,
)
)
def update_model(self, model_id: str, litellm_params: LiteLLMParamsBody) -> None:
"""Merge `litellm_params` over the deployment `model_id`'s stored params via

View file

@ -58,6 +58,7 @@ class Transport(Protocol):
headers: BaseModel,
params: BaseModel,
response_type: type[R],
timeout: float | None = None,
) -> Result[R]: ...
def delete[R: BaseModel](
@ -136,13 +137,16 @@ class HttpTransport:
headers: BaseModel,
params: BaseModel,
response_type: type[R],
timeout: float | None = None,
) -> Result[R]:
"""`timeout` overrides the transport-wide request_timeout for this call, for
pollers whose own deadline is shorter than it."""
return e2e_http.get(
self._url(path),
headers=headers,
params=params,
response_type=response_type,
timeout=self.request_timeout,
timeout=self.request_timeout if timeout is None else timeout,
)
def delete[R: BaseModel](
@ -336,9 +340,14 @@ class SplitTransport:
headers: BaseModel,
params: BaseModel,
response_type: type[R],
timeout: float | None = None,
) -> Result[R]:
return self._route(path).get(
path, headers=headers, params=params, response_type=response_type
path,
headers=headers,
params=params,
response_type=response_type,
timeout=timeout,
)
def delete[R: BaseModel](