scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 136
litellm/tests/e2e/proxy_client.py
Yuneng Jiang 09a98f5505
test(e2e): settle control-plane writes across every replica, not just one 
The suite already waits for a new model or agent to become servable before
handing it back, but that wait returns on the first successful read. Every
request opens a fresh connection (e2e_http calls requests.* with no Session), so
a load-balanced Service routes each one independently: one successful read proves
one replica converged, and the caller's next request re-rolls and can land on a
replica that has not reloaded yet.

At replicaCount: 2 this surfaced as 30 failures on a SHA that is green at 1
replica -- 400 "Invalid model name passed", 404 "Guardrail not found", "no
healthy deployments for this model", and a /model/info listing that contained
one of two models created moments apart.

Add PROPAGATION_TIMEOUT (default 15s, override E2E_PROPAGATION_TIMEOUT) and
settle_propagation(), sized off the proxy's proxy_config_reload_interval_seconds
(30s by default, 7s on the e2e stack) plus margin, and settle after every
control-plane create whose object the suite then uses:

- ProxyClient.create_model and A2AClient.register_agent, after their existing
  polls -- the poll still fails loudly if the object never appears at all
- GuardrailsClient.register, which had no barrier; create_content_filter_guardrail
  and create_bedrock_guardrail now route through it instead of POSTing directly
- the guardrail creates in mcp_client and logging_client
- the vertex passthrough model, whose body cannot go through create_model

Left alone: the /model/new calls that assert a 403 or read back a status code,
since they never use the model.
2026-08-07 19:36:30 -07:00

433 lines
15 KiB
Python
Raw Blame History

"""ProxyClient: the shared proxy operations, DI'd into every client (composition).
A frozen-slots dataclass holding a Transport plus poll config. Clients hold a
ProxyClient and add their own route methods; the lifecycle ResourceManager uses the
ProxyClient's key/customer methods for cleanup. Read-backs are eventually consistent
(proxy_batch_write_at ~60s) so they poll to a deadline.
"""
from __future__ import annotations
import time
import warnings
from collections.abc import Callable
from dataclasses import dataclass
from datetime import datetime
from e2e_http import (
AnthropicHeaders,
NoBody,
ProbeResult,
Result,
StreamingResponse,
Success,
is_ok,
unwrap,
)
from models import (
AnthropicMessagesBody,
AnthropicMessagesResponse,
ChatBody,
ChatResponse,
CountTokensBody,
CountTokensResponse,
CredentialCreateBody,
CredentialCreateResponse,
CustomerDeleteBody,
EmbedBody,
EmbedResponse,
FileListResponse,
FineTuningJobsParams,
FineTuningJobsResponse,
KeyDeleteBody,
KeyGenerateBody,
KeyGenerateResponse,
KeyInfo,
KeyInfoParams,
KeyInfoResponse,
LiteLLMParamsBody,
ModelDeleteBody,
ModelInfoBody,
ModelInfoEntry,
ModelInfoResponse,
ModelMode,
ModelNewBody,
ModelNewResponse,
ModelsListResponse,
ModelUpdateBody,
OcrBody,
OcrResponse,
SpendLogRow,
SpendLogs,
SpendLogsPage,
SpendLogsPageParams,
SpendLogsParams,
)
from e2e_config import (
CONTROL_PLANE_BASE_URL,
MASTER_KEY,
POLL_INTERVAL,
POLL_TIMEOUT,
PROXY_BASE_URL,
REQUEST_TIMEOUT,
settle_propagation,
)
from transport import HttpTransport, SplitTransport, Transport
RowsPredicate = Callable[[list[SpendLogRow]], bool]
@dataclass(frozen=True, slots=True)
class ProxyClient:
transport: Transport
poll_timeout: float = 120.0
poll_interval: float = 5.0
# ---- keys / customers (satisfies lifecycle.ResourceClient) ----------
def generate_key(self, body: KeyGenerateBody) -> str:
return unwrap(
self.transport.post(
"/key/generate",
headers=self.transport.master,
json=body,
response_type=KeyGenerateResponse,
)
).key
def delete_key(self, key: str) -> None:
_ = self.transport.post(
"/key/delete",
headers=self.transport.master,
json=KeyDeleteBody(keys=[key]),
response_type=NoBody,
)
def delete_customers(self, user_ids: list[str]) -> None:
if not user_ids:
return
_ = self.transport.post(
"/customer/delete",
headers=self.transport.master,
json=CustomerDeleteBody(user_ids=user_ids),
response_type=NoBody,
)
def key_info(self, key: str) -> KeyInfo:
return unwrap(
self.transport.get(
"/key/info",
headers=self.transport.master,
params=KeyInfoParams(key=key),
response_type=KeyInfoResponse,
)
).info
def model_info(self) -> list[ModelInfoEntry]:
"""Every configured deployment with the price the proxy resolved for it
(config override merged over cost-map defaults)."""
return unwrap(
self.transport.get(
"/model/info",
headers=self.transport.master,
params=NoBody(),
response_type=ModelInfoResponse,
)
).data
def list_files(self, key: str) -> Result[FileListResponse]:
return self.transport.get(
"/v1/files",
headers=self.transport.bearer(key),
params=NoBody(),
response_type=FileListResponse,
)
def list_fine_tuning_jobs(
self, key: str, params: FineTuningJobsParams
) -> Result[FineTuningJobsResponse]:
return self.transport.get(
"/v1/fine_tuning/jobs",
headers=self.transport.bearer(key),
params=params,
response_type=FineTuningJobsResponse,
)
def create_model(
self,
model_name: str,
litellm_params: LiteLLMParamsBody,
mode: ModelMode | None = None,
) -> str:
"""Register a deployment under `model_name` and return its proxy-assigned
model_id, once the model is actually servable on the data plane.
/model/new is a control-plane route; the data plane (which serves /chat,
/ocr, ...) only picks the new model up on its next DB reload, so a call
issued the instant this returns can race the reload and 400 with "Invalid
model name passed". We poll the data-plane /v1/models until the model
appears, then settle for the remainder of the propagation budget.
Both steps are needed, and the second is the one that matters at >1 replica.
The poll proves *a* replica is serving the model; it cannot prove they all
are, because every request opens a fresh connection and a load-balanced
Service routes each one independently -- so the caller's next request
re-rolls and can land on a replica that has not reloaded yet. Waiting out
PROPAGATION_TIMEOUT is what makes the model safe to use anywhere."""
model_id = unwrap(
self.transport.post(
"/model/new",
headers=self.transport.master,
json=ModelNewBody(
model_name=model_name,
litellm_params=litellm_params,
model_info=ModelInfoBody(mode=mode),
),
response_type=ModelNewResponse,
)
).model_id
written_at = time.monotonic()
self._await_model_servable(model_name)
settle_propagation(written_at)
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(
"/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
):
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}"
)
def update_model(self, model_id: str, litellm_params: LiteLLMParamsBody) -> None:
"""Merge `litellm_params` over the deployment `model_id`'s stored params via
POST /model/update. The proxy overlays only the non-null fields and clears
its model cache, so a later /model/info read reflects the change (eventually,
after the reload)."""
unwrap(
self.transport.post(
"/model/update",
headers=self.transport.master,
json=ModelUpdateBody(
litellm_params=litellm_params,
model_info=ModelInfoBody(id=model_id),
),
response_type=NoBody,
)
)
def delete_model(self, model_id: str) -> None:
result = self.transport.post(
"/model/delete",
headers=self.transport.master,
json=ModelDeleteBody(id=model_id),
response_type=NoBody,
)
if not is_ok(result):
warnings.warn(f"delete_model({model_id!r}) failed: {result}", stacklevel=2)
def create_credential(self, body: CredentialCreateBody) -> None:
unwrap(
self.transport.post(
"/credentials",
headers=self.transport.master,
json=body,
response_type=CredentialCreateResponse,
)
)
def delete_credential(self, credential_name: str) -> None:
result = self.transport.delete(
f"/credentials/{credential_name}",
headers=self.transport.master,
json=NoBody(),
response_type=NoBody,
)
if not is_ok(result):
warnings.warn(f"delete_credential({credential_name!r}) failed: {result}", stacklevel=2)
# ---- LLM calls ------------------------------------------------------
def chat(self, key: str, body: ChatBody) -> Result[ChatResponse]:
return self.transport.post(
"/chat/completions",
headers=self.transport.bearer(key),
json=body,
response_type=ChatResponse,
)
def chat_stream(self, key: str, body: ChatBody) -> StreamingResponse:
return self.transport.stream("/chat/completions", headers=self.transport.bearer(key), json=body)
def messages_stream(self, key: str, body: AnthropicMessagesBody) -> StreamingResponse:
return self.transport.stream("/v1/messages", headers=self.transport.bearer(key), json=body)
def embed(self, key: str, body: EmbedBody) -> Result[EmbedResponse]:
return self.transport.post(
"/embeddings",
headers=self.transport.bearer(key),
json=body,
response_type=EmbedResponse,
)
def ocr(self, key: str, body: OcrBody) -> Result[OcrResponse]:
return self.transport.post(
"/v1/ocr",
headers=self.transport.bearer(key),
json=body,
response_type=OcrResponse,
)
def count_tokens(self, key: str, body: CountTokensBody) -> Result[CountTokensResponse]:
"""POST /v1/messages/count_tokens (Anthropic-native). Sends the
anthropic-version header so the native path accepts it; harmless on the
other providers the proxy fronts."""
return self.transport.post(
"/v1/messages/count_tokens",
headers=self._anthropic_headers(key),
json=body,
response_type=CountTokensResponse,
)
def messages(self, key: str, body: AnthropicMessagesBody) -> Result[AnthropicMessagesResponse]:
"""POST /v1/messages (Anthropic-native). The response is either the
Anthropic-shape passthrough (`content`) or the OpenAI-normalized shape
(`choices`); AnthropicMessagesResponse models both."""
return self.transport.post(
"/v1/messages",
headers=self._anthropic_headers(key),
json=body,
response_type=AnthropicMessagesResponse,
)
def _anthropic_headers(self, key: str) -> AnthropicHeaders:
return AnthropicHeaders(authorization=self.transport.bearer(key).authorization)
# ---- spend read-back ------------------------------------------------
def spend_logs(self, params: SpendLogsParams) -> list[SpendLogRow]:
result = self.transport.get(
"/spend/logs",
headers=self.transport.master,
params=params,
response_type=SpendLogs,
)
match result:
case Success(data=logs):
return logs.root
case _:
return []
def spend_logs_window(self, *, start: datetime, end: datetime) -> list[SpendLogRow]:
def fetch(page: int) -> SpendLogsPage:
return unwrap(
self.transport.get(
"/spend/logs/v2",
headers=self.transport.master,
params=SpendLogsPageParams(
start_date=start.strftime("%Y-%m-%d %H:%M:%S"),
end_date=end.strftime("%Y-%m-%d %H:%M:%S"),
page=page,
page_size=100,
),
response_type=SpendLogsPage,
)
)
first = fetch(1)
return [
*first.data,
*(row for page in range(2, first.total_pages + 1) for row in fetch(page).data),
]
def poll_logs_for_key(
self, key: str, *, min_rows: int = 1, predicate: RowsPredicate | None = None
) -> list[SpendLogRow]:
return self._poll(lambda: self.spend_logs(SpendLogsParams(api_key=key)), min_rows, predicate)
def poll_logs_for_request_id(
self,
request_id: str,
*,
min_rows: int = 1,
predicate: RowsPredicate | None = None,
) -> list[SpendLogRow]:
return self._poll(
lambda: self.spend_logs(SpendLogsParams(request_id=request_id)),
min_rows,
predicate,
)
def _poll(
self,
fetch: Callable[[], list[SpendLogRow]],
min_rows: int,
predicate: RowsPredicate | None,
) -> list[SpendLogRow]:
deadline = time.monotonic() + self.poll_timeout
rows: list[SpendLogRow] = []
while time.monotonic() < deadline:
rows = fetch()
if len(rows) >= min_rows and (predicate is None or predicate(rows)):
return rows
time.sleep(self.poll_interval)
return rows
# ---- route probe ----------------------------------------------------
def probe(self, path: str, *, params: NoBody) -> ProbeResult:
return self.transport.probe(path, params=params)
def build_proxy_client(
*,
base_url: str = PROXY_BASE_URL,
master_key: str = MASTER_KEY,
control_plane_base_url: str = CONTROL_PLANE_BASE_URL,
) -> ProxyClient:
"""The ProxyClient every suite's client is built from: a SplitTransport that routes
LLM calls to the data plane (PROXY_BASE_URL) and management/admin calls to the
control plane (CONTROL_PLANE_BASE_URL), with the shared poll budget. The two
base URLs are the same for a monolithic proxy, so routing is then a no-op.
The endpoints are injectable for callers that resolve the proxy some other
way than ``e2e_config``'s env names (see ``claude_code/_env.py``); they must
pass all three together, since a caller that overrides only the data plane
would leave management calls pointed at the env default."""
return ProxyClient(
transport=SplitTransport(
data=HttpTransport(
base_url=base_url,
master_key=master_key,
request_timeout=REQUEST_TIMEOUT,
),
control=HttpTransport(
base_url=control_plane_base_url,
master_key=master_key,
request_timeout=REQUEST_TIMEOUT,
),
),
poll_timeout=POLL_TIMEOUT,
poll_interval=POLL_INTERVAL,
)
Reference in a new issue View git blame Copy permalink