scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-05 07:37:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 233

test(e2e): retire the quiet_stack lock and batch the claude_code model registration (#49984)

Browse source 
* test(e2e): retire the quiet_stack lock

quiet_stack held an exclusive host-wide lock for six tests, draining every
other xdist worker while they ran. On the full e2e lane that left 35-38% of
worker time idle. Each of those tests now lives where it can run without
pausing the suite

- the two memory tests move to perf/ behind the perf marker, so the perf lane
  runs them on its own fresh stack. Their refusing deployments use a literal
  key because the perf stack carries no provider keys
- the stored-request size check from the memory test becomes an integration
  test on an owned proxy that stores prompts
- the discount and margin spend-row tests are dropped from e2e. Their
  breakdown assertions join the existing integration pricing test, which
  already owns the discount and margin repricing flow
- the Bedrock guardrail cost test stays in e2e without the pricing reset
  fixture, since nothing else on the stack writes pricing config now
- the MCP semantic tool filter test turns on a proxy-wide setting, so it gets
  a new proxy_wide_config opt-in marker and runs after the main run

* test(e2e): register the claude_code deployments as one batch

The claude_code session fixture registered its 24 deployments one at a time on every xdist worker, and each one waited out the 15 s propagation window on its own, about 6 minutes per worker. ProxyClient.register_models posts the whole batch, polls each model until every replica serves it, then waits out propagation once from the last write. A model that fails to post or never becomes servable comes back as ModelNotRegistered, and an unservable one is deleted, so the fixture keeps reporting per-deployment failures the way it did
This commit is contained in:
yuneng-jiang 2026-11-04 09:55:53 -08:00 • committed by GitHub
parent df78b16941
commit b18188dee2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 359 additions and 471 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
tests/e2e/AGENTS.md
View file

File diff suppressed because one or more lines are too long

35
tests/e2e/claude_code/conftest.py
View file

@ -597,14 +597,21 @@ def _build_control_plane_client(proxy_config: ProxyConfig):
)
def _register_deployment(proxy, deployment: CompatDeployment) -> str:
"""Register one deployment and return its proxy-assigned model_id
once it is servable on the data plane. The aliases are shared by every
cell and, under xdist, by every worker, so no call to them belongs to
one test and none is cached: the matrix exists to reach real providers."""
return proxy.create_model(
deployment.model_name,
deployment.litellm_params,
def _register_deployments(proxy, deployments: tuple[CompatDeployment, ...]):
"""The aliases are shared by every cell and, under xdist, by every worker,
so no call to them belongs to one test and none is cached: the matrix
exists to reach real providers."""
from models import ModelInfoBody, ModelNewBody
return proxy.register_models(
tuple(
ModelNewBody(
model_name=deployment.model_name,
litellm_params=deployment.litellm_params,
model_info=ModelInfoBody(),
)
for deployment in deployments
),
provider_live=True,
)
@ -632,18 +639,14 @@ def _compat_models_registered() -> Any:
yield
return
from proxy_client import ModelNotRegistered, ModelRegistered
from requests import RequestException
proxy = _build_control_plane_client(proxy_config)
registered_ids: list[str] = []
failures: list[tuple[str, str]] = []
registrations = _register_deployments(proxy, load_all_deployments())
registered_ids = tuple(r.model_id for r in registrations if isinstance(r, ModelRegistered))
failures = tuple((r.model_name, r.reason) for r in registrations if isinstance(r, ModelNotRegistered))
try:
for deployment in load_all_deployments():
try:
model_id = _register_deployment(proxy, deployment)
registered_ids.append(model_id)
except (AssertionError, RequestException) as exc:
failures.append((deployment.model_name, str(exc)))
if failures:
summary = "\n".join(
f" - {name}: {reason}" for name, reason in failures

21
tests/e2e/conftest.py
View file

@ -35,6 +35,7 @@ from e2e_config import (
PROMPT_CACHING_OPT_IN_ENV,
PROVIDER_EDGE_HOST_OPT_IN_ENV,
PROXY_BASE_URL,
PROXY_WIDE_CONFIG_OPT_IN_ENV,
SAGEMAKER_NOVA2_LITE_OPT_IN_ENV,
SAGEMAKER_NOVA_OPT_IN_ENV,
SECRET_MANAGER_OPT_IN_ENV,
@ -54,7 +55,6 @@ from models import TeamNewBody, UserNewBody, UserNewResponse
from provider_cache_routing import LIVE_PROVIDER_REQUIRED
from provider_edge import replay_leftover_error
from proxy_client import ProxyClient, build_proxy_client
from stack_lock import stack_lock
_E2E_TEST_RAN = pytest.StashKey[bool]()
_CALL_PASSED = pytest.StashKey[bool]()
@ -77,6 +77,7 @@ OPT_IN_MARKERS: Final = MappingProxyType(
"secret_manager": SECRET_MANAGER_OPT_IN_ENV,
"sagemaker_nova": SAGEMAKER_NOVA_OPT_IN_ENV,
"sagemaker_nova2_lite": SAGEMAKER_NOVA2_LITE_OPT_IN_ENV,
"proxy_wide_config": PROXY_WIDE_CONFIG_OPT_IN_ENV,
}
)
@ -162,11 +163,6 @@ def pytest_configure(config: pytest.Config) -> None:
"redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from "
"gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_PERF_TESTS is set",
)
config.addinivalue_line(
"markers",
"quiet_stack: measures the proxy itself, so it runs while no other test on this host is hitting the stack; "
"every other test waits for it to finish",
)
config.addinivalue_line(
"markers",
"provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the "
@ -200,6 +196,11 @@ def pytest_configure(config: pytest.Config) -> None:
"sagemaker_nova2_lite: needs a deployed SageMaker Nova 2 Lite endpoint the proxy's AWS credentials can invoke; "
"deselected unless SAGEMAKER_NOVA2_LITE_ENDPOINT names it",
)
config.addinivalue_line(
"markers",
"proxy_wide_config: turns on a proxy-wide setting through /config/update, so it runs on a stack no other suite "
"is driving; deselected unless E2E_PROXY_WIDE_CONFIG_STACK is set",
)
def pytest_sessionstart(session: pytest.Session) -> None:
@ -269,7 +270,7 @@ def pytest_collection_finish(session: pytest.Session) -> None:
ahead of xdist's own collection-finish report, and the controller schedules no
test until every worker has reported, so this is the idle footprint of a stack
that just passed its readiness gate. The fixture hands the capture to the
idle-budget test in router/test_reliability_memory_e2e.py."""
idle-budget test in perf/test_reliability_memory_e2e.py."""
if session.config.getoption("collectonly") or not any(_uses_idle_rss(item) for item in session.items):
return
session.config.stash[_IDLE_RSS] = read_rss_everywhere(build_proxy_client(), timeout=IDLE_RSS_READ_TIMEOUT_SECONDS)
@ -299,12 +300,6 @@ def _proxy_fail_reason() -> str | None:
return None
@pytest.hookimpl(wrapper=True)
def pytest_runtest_protocol(item: pytest.Item, nextitem: pytest.Item | None) -> Generator[None, object, object]:
with stack_lock(exclusive=item.get_closest_marker("quiet_stack") is not None):
return (yield)
@pytest.hookimpl(tryfirst=True)
def pytest_runtest_setup(item: pytest.Item) -> None:
"""Hard-fail `e2e`-marked tests unless a proxy answers its liveness probe.

2
tests/e2e/coverage_registry/quota_management.yaml
View file

@ -44,8 +44,6 @@
- {id: quota_management.spend_tracking.messages_bridge.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: messages_bridge, assertions: [logs_cost], exercised_on: [messages], source: "llms/anthropic/messages/responses_adapters/handler.py", rationale: "A streaming /v1/messages request served by an openai-provider model is bridged through the anthropic-messages -> Responses adapter and must aggregate the consumed SSE stream into one spend row with nonzero cost and token counts, attributed to custom_llm_provider openai under call_type anthropic_messages"}
- {id: quota_management.spend_tracking.embeddings.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: embeddings, assertions: [logs_cost], exercised_on: [embeddings], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Embedding calls write nonzero spend rows"}
- {id: quota_management.spend_tracking.cache_hit.zero_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: cache_hit, assertions: [zero_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "A response-cache hit logs at zero cost with the cache-hit marker"}
- {id: quota_management.spend_tracking.discount_config.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: discount_config, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/management_endpoints/cost_tracking_settings.py + cost_calculator.py", rationale: "A persisted provider discount changes the real request's spend row"}
- {id: quota_management.spend_tracking.margin_config.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: margin_config, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/management_endpoints/cost_tracking_settings.py + cost_calculator.py", rationale: "A persisted provider margin is added to the real request's spend row"}
- {id: quota_management.spend_tracking.guardrail_cost.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: guardrail_cost, assertions: [logs_cost], exercised_on: [chat_completions], source: "litellm_core_utils/llm_cost_calc/guardrail_cost.py + proxy/spend_tracking/spend_tracking_utils.py", rationale: "A priced Bedrock guardrail charge reaches the persisted spend row"}
- {id: quota_management.spend_tracking.key_rollup.matches_sum_of_logs, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_rollup, assertions: [matches_sum_of_logs], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "A key's rolled-up spend equals the sum of its log rows"}
- {id: quota_management.spend_tracking.concurrent_burst.loses_no_spend, module: quota_management, tier: P1, behavior: spend_tracking, variant: concurrent_burst, assertions: [loses_no_spend], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "Concurrent calls all land as spend; no row lost to write contention"}

2
tests/e2e/e2e_config.py
View file

@ -210,6 +210,7 @@ OTEL_TLS_OPT_IN_ENV: Final = "E2E_OTEL_EXPORTER_ENDPOINT"
SECRET_MANAGER_OPT_IN_ENV: Final = "E2E_SECRET_MANAGER"
SAGEMAKER_NOVA_OPT_IN_ENV: Final = "SAGEMAKER_NOVA_ENDPOINT"
SAGEMAKER_NOVA2_LITE_OPT_IN_ENV: Final = "SAGEMAKER_NOVA2_LITE_ENDPOINT"
PROXY_WIDE_CONFIG_OPT_IN_ENV: Final = "E2E_PROXY_WIDE_CONFIG_STACK"
ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6"))
ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6"))
ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3"))
@ -234,7 +235,6 @@ MEMORY_RSS_SETTLE_SAMPLES = int(os.environ.get("E2E_MEMORY_RSS_SETTLE_SAMPLES",
MEMORY_RSS_SAMPLE_INTERVAL_SECONDS = float(os.environ.get("E2E_MEMORY_RSS_SAMPLE_INTERVAL_SECONDS", "1"))
MEMORY_RSS_CEILING_MB = float(os.environ.get("E2E_MEMORY_RSS_CEILING_MB", "200"))
MEMORY_IDLE_RSS_BUDGET_MB = float(os.environ.get("E2E_MEMORY_IDLE_RSS_BUDGET_MB", "768"))
MEMORY_STORED_REQUEST_BUDGET_KB = float(os.environ.get("E2E_MEMORY_STORED_REQUEST_BUDGET_KB", "64"))
def ws_base_url() -> str:

5
tests/e2e/mcp/test_mcp_semantic_tool_filter_e2e.py
View file

@ -5,7 +5,8 @@ embedding deployment, and turns on `litellm_settings.mcp_semantic_tool_filter` t
/config/update. A chat completion that attaches the server and asks to search logs must report,
in the x-litellm-semantic-filter headers, that the catalog was cut to at most top_k tools, all
of them tools the key lists. Requires DD_API_KEY + DD_APP_KEY and the proxy's
OPENAI_API_KEY for the embedding deployment.
OPENAI_API_KEY for the embedding deployment. The filter is proxy-wide, so the test is
marked `proxy_wide_config` and runs on a stack the main suite is not driving.
"""
from __future__ import annotations
@ -23,7 +24,7 @@ from mcp_client import McpClient
from models import ChatBody, ChatMessage, LiteLLMParamsBody, McpChatTool
from pydantic import BaseModel
pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack]
pytestmark = [pytest.mark.e2e, pytest.mark.proxy_wide_config]
TOP_K: Final = 3

21
tests/e2e/models.py
View file

@ -1262,27 +1262,6 @@ class CostMap(RootModel[dict[str, CostMapEntry]]):
pass
class ConfigPatchResponse(BaseModel):
status: str
values: dict[str, float | dict[str, float]]
class CostDiscountConfig(RootModel[dict[str, float]]):
pass
class CostMarginConfig(RootModel[dict[str, float | dict[str, float]]]):
pass
class CostDiscountConfigResponse(BaseModel):
values: dict[str, float]
class CostMarginConfigResponse(BaseModel):
values: dict[str, float | dict[str, float]]
class BedrockGuardrailParams(BaseModel):
guardrail: Literal["bedrock"] = "bedrock"
mode: Literal["pre_call"] = "pre_call"

2
tests/e2e/perf/test_database_upgrade.py
View file

@ -28,7 +28,7 @@ from setup.run import run_plan
from setup.upgrade import assert_finished, monitor_for, upgrade_case, upgrade_properties
from setup.upgrade_worker import connect, read_state
pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack]
pytestmark = pytest.mark.perf
@pytest.mark.skipif(os.environ.get("E2E_DATABASE_UPGRADE") != "1", reason="Dedicated disposable snapshot lane only")

2
tests/e2e/perf/test_perf_cases.py
View file

@ -48,7 +48,7 @@ from setup.ramp import step_properties
from setup.run import run_plan
from setup.telemetry import read_gateway_telemetry, telemetry_sampler
pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack]
pytestmark = pytest.mark.perf
_SPEND_SAMPLE: Final = 40
_WARMUP_SECONDS: Final = 60

111
tests/e2e/router/test_reliability_memory_e2e.py → tests/e2e/perf/test_reliability_memory_e2e.py
View file

@ -1,17 +1,19 @@
"""Live e2e: a few hundred requests that fail before any provider answers must not
grow the proxy's resident memory past a loose ceiling once the proxy is warm.
grow the proxy's resident memory past a loose ceiling once the proxy is warm. It reads
the proxy's own RSS, so it runs in the perf lane, where nothing else is hitting the stack.
The regression this guards shipped in v1.100.0: every retry breadcrumb copied the
whole request and the copies nested into one router-global list, so a proxy under
retry-heavy failing traffic grew until it was OOM-killed. The traffic here has that
shape: a model group whose deployments refuse at the socket (an unreachable base
URL) with cooldown_time 0 so the router keeps retrying them, per-request retries,
and a fallback group that refuses the same way, each request carrying a long chat
URL, with a literal key since the perf stack carries no provider keys) with
cooldown_time 0 so the router keeps retrying them, per-request retries, and a
fallback group that refuses the same way, each request carrying a long chat
transcript so every whole-request copy costs hundreds of containers instead of a
handful. Under the stack's cooldown policy a
deployment that fails a handful of times in a row is benched (a bad-credential 401
included), the router answers "No deployments available" without retrying, and the
retry loop that leaks stops running; cooldown_time 0 keeps it running.
handful. Under the stack's cooldown policy a deployment that fails a handful of
times in a row is benched (a bad-credential 401 included), the router answers "No
deployments available" without retrying, and the retry loop that leaks stops running;
cooldown_time 0 keeps it running.
Four identical batches run back to back. The first two are the warmup that grows
the proxy's caches and allocator arenas, the last two are the measured window, and
@ -35,19 +37,11 @@ single load-balanced address reaches the workers behind it one answer at a time,
a worker that answered only one checkpoint would otherwise drop out of the
comparison, which is where a leaking worker could hide.
RSS alone is a coarse gauge: on the release stack (spend logs storing prompts,
json logs, prometheus and otel callbacks) the same v1.100.0 breadcrumbs grew RSS
by only about 15 MB per 300 failing requests, while every failing request's stored
request snapshot carried a copy of the request per failed attempt, over 100 KB on
the first call and a couple of MB once the copies nested, against tens of KB with
the fix. So the first check sends one failing request before the phases, reads its
spend log back through /spend/logs, and holds the stored request body to a fixed
size budget: the deterministic catch for a breadcrumb that copies the whole
request. It runs before the phases because the leaking writer drops its own rows
under the phases' traffic (a queue budget hit, a recursion limit on the nested
copies), which would turn the size check into a missing-row check.
The stored-request half of that regression (every failing request's spend log carrying a
copy of the request per failed attempt) is a deterministic size check, so it lives in
tests/integration/spend/test_failing_request_stored_snapshot.py.
A second, cheaper check holds the idle footprint: every worker's RSS as the harness
A cheaper check holds the idle footprint: every worker's RSS as the harness
read it at collection time, before this pytest process sent any traffic (see
conftest.pytest_collection_finish), must sit under a fixed budget. On the
release gate that is a fresh stack right after its readiness gate, one worker per
@ -74,7 +68,6 @@ from typing import Final
import pytest
from complexity_router_client import ComplexityRouterClient
from e2e_config import (
MEMORY_CONCURRENCY,
MEMORY_IDLE_RSS_BUDGET_MB,
@ -83,7 +76,6 @@ from e2e_config import (
MEMORY_RSS_CEILING_MB,
MEMORY_RSS_SAMPLE_INTERVAL_SECONDS,
MEMORY_RSS_SETTLE_SAMPLES,
MEMORY_STORED_REQUEST_BUDGET_KB,
MEMORY_TRANSCRIPT_TURNS,
PROXY_REPLICA_URLS,
unique_marker,
@ -91,11 +83,12 @@ from e2e_config import (
from e2e_metadata import Domain, Mode, Provider, Subject, meta
from lifecycle import ResourceManager
from memory_readings import RssCapture, RssReading, WorkerKey, read_rss_everywhere
from models import ChatMessage, RouterSettingsOverride, SpendLogRow
from models import ChatMessage, LiteLLMParamsBody, ReliabilityChatBody, RouterSettingsOverride
from proxy_client import ProxyClient
from reliability_support import REAL_MODEL, chat_override, create_never_benched_refusing_deployment
pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack]
pytestmark = [pytest.mark.e2e, pytest.mark.perf]
REFUSING_MODEL: Final = "openai/gpt-5.5"
DEPLOYMENTS_PER_GROUP: Final = 2
RSS_SAMPLE_CAP: Final = 4 * MEMORY_RSS_SETTLE_SAMPLES
@ -110,7 +103,6 @@ class FailedCall:
status_code: int
seconds: float
body_head: str
call_id: str | None
@dataclass(frozen=True, slots=True)
@ -123,8 +115,17 @@ class WorkerGrowth:
return self.after.ram_usage_mb - self.warm.ram_usage_mb
def _create_never_benched_refusing_deployment(proxy: ProxyClient, name: str) -> str:
return proxy.create_model(
name,
LiteLLMParamsBody(
model=REFUSING_MODEL, api_key="sk-unreachable", api_base="http://127.0.0.1:9/v1", cooldown_time=0
),
)
def _register_refusing_group(proxy: ProxyClient, resources: ResourceManager, name: str) -> None:
for model_id in tuple(create_never_benched_refusing_deployment(proxy, name) for _ in range(DEPLOYMENTS_PER_GROUP)):
for model_id in tuple(_create_never_benched_refusing_deployment(proxy, name) for _ in range(DEPLOYMENTS_PER_GROUP)):
resources.defer(lambda model_id=model_id: proxy.delete_model(model_id))
@ -141,10 +142,18 @@ TRANSCRIPT: Final = _transcript(MEMORY_TRANSCRIPT_TURNS)
def _fail_once(proxy: ProxyClient, key: str, model: str, override: RouterSettingsOverride) -> FailedCall:
started: Final = time.perf_counter()
resp: Final = chat_override(
proxy, key, model, f"memory regression {unique_marker()}", override=override, history=TRANSCRIPT
resp: Final = proxy.transport.send(
"/chat/completions",
headers=proxy.transport.bearer(key),
json=ReliabilityChatBody(
model=model,
messages=[*TRANSCRIPT, ChatMessage(role="user", content=f"memory regression {unique_marker()}")],
max_tokens=512,
router_settings_override=override,
cache={"no-cache": True},
),
)
return FailedCall(resp.status_code, time.perf_counter() - started, resp.body[:300], resp.call_id)
return FailedCall(resp.status_code, time.perf_counter() - started, resp.body[:300])
def _fail_many(proxy: ProxyClient, key: str, model: str, override: RouterSettingsOverride) -> tuple[FailedCall, ...]:
@ -259,24 +268,6 @@ def _assert_every_call_failed_through_fallback(calls: Sequence[FailedCall], fall
)
def _stored_request_kb(proxy: ProxyClient, call: FailedCall) -> float:
assert call.call_id, (
f"the failing call carried no x-litellm-call-id header, so its spend log cannot be read back: {call.body_head}"
)
rows: Final[Sequence[SpendLogRow]] = proxy.poll_logs_for_request_id(call.call_id)
assert rows, (
f"no spend log row appeared for failing call {call.call_id} within the poll window: either the stack "
"writes no spend logs or its writer dropped the row, which the v1.100.0 one did once the stored "
"request outgrew the writer's queue budget"
)
snapshot: Final = rows[0].proxy_server_request
assert snapshot, (
f"spend log {call.call_id} stored no request body, so the stack is not running with "
"general_settings.store_prompts_in_spend_logs and the stored-request check would pass vacuously"
)
return len(json.dumps(snapshot).encode()) / 1024
class TestReliabilityMemory:
@pytest.mark.covers("reliability.perf.idle_memory.under_slo")
@meta(
@ -309,13 +300,13 @@ class TestReliabilityMemory:
Subject(
domain=Domain.ROUTING,
providers=(Provider.OPENAI,),
models=(REAL_MODEL,),
models=(REFUSING_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_failing_requests_do_not_grow_rss_or_stored_request(
def test_failing_requests_do_not_grow_rss(
self,
client: ComplexityRouterClient,
proxy: ProxyClient,
resources: ResourceManager,
scoped_key: str,
record_property: Callable[[str, object], None],
@ -323,35 +314,25 @@ class TestReliabilityMemory:
marker: Final = unique_marker()
primary: Final = f"reliability-memory-{marker}"
fallback: Final = f"reliability-memory-fb-{marker}"
_register_refusing_group(client.proxy, resources, primary)
_register_refusing_group(client.proxy, resources, fallback)
_register_refusing_group(proxy, resources, primary)
_register_refusing_group(proxy, resources, fallback)
override: Final = RouterSettingsOverride(
num_retries=MEMORY_RETRIES_PER_REQUEST, fallbacks=[{primary: [fallback]}]
)
probe: Final = _fail_once(client.proxy, scoped_key, primary, override)
_assert_every_call_failed_through_fallback((probe,), fallback)
stored_kb: Final = _stored_request_kb(client.proxy, probe)
assert stored_kb <= MEMORY_STORED_REQUEST_BUDGET_KB, (
f"the spend log of one failing request stored a {stored_kb:.0f} KB request body, past the "
f"{MEMORY_STORED_REQUEST_BUDGET_KB:.0f} KB budget for a {len(TRANSCRIPT)}-message transcript with "
f"{MEMORY_RETRIES_PER_REQUEST} retries and a fallback; the retry breadcrumbs are copying the whole "
f"request into the stored snapshot the way the v1.100.0 ones did"
)
warm: Final = _settled_rss_after_batches(
client.proxy,
proxy,
scoped_key,
primary,
override,
fallback,
record_property,
_settled_rss_per_worker(client.proxy),
_settled_rss_per_worker(proxy),
1,
WARMUP_BATCHES,
)
after: Final = _settled_rss_after_batches(
client.proxy,
proxy,
scoped_key,
primary,
override,

1
tests/e2e/perf/test_ui_journeys.py
View file

@ -33,7 +33,6 @@ from transport import HttpTransport
pytestmark = [
pytest.mark.perf,
pytest.mark.quiet_stack,
pytest.mark.skipif(os.environ.get("E2E_UI_PERF_DISPOSABLE") != "1", reason="Dedicated disposable database only"),
]

96
tests/e2e/proxy_client.py
View file

@ -11,7 +11,7 @@ from __future__ import annotations
import os
import time
import warnings
from collections.abc import Callable, Mapping
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass, field, replace
from datetime import datetime
from functools import reduce
@ -155,6 +155,21 @@ class NotServableOn:
last_result: Result[ModelsListResponse] | None
@dataclass(frozen=True, slots=True)
class ModelRegistered:
model_name: str
model_id: str
@dataclass(frozen=True, slots=True)
class ModelNotRegistered:
model_name: str
reason: str
ModelRegistration = ModelRegistered | ModelNotRegistered
def await_servable(
list_models: Callable[[float], Result[ModelsListResponse]],
*,
@ -461,6 +476,7 @@ class ProxyClient:
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
settle: Callable[[float], None] = settle_propagation
def with_caller(self, caller: Caller) -> ProxyClient:
return replace(self, caller=caller)
@ -694,31 +710,59 @@ class ProxyClient:
balancer address is configured (every request opens a fresh connection, so
the caller's next request re-rolls), so waiting out PROPAGATION_TIMEOUT is
what makes the model safe to use anywhere."""
model_id = unwrap(
self.transport.post(
"/model/new",
headers=self.management_headers(),
json=body.model_copy(update={"litellm_params": route_cache_model(
body.litellm_params, provider_edge_base,
enabled=os.environ.get("E2E_PROVIDER_CACHE", "0") == "1" and not provider_live,
mode=body.model_info.mode,
)}),
response_type=ModelNewResponse,
)
).model_id
model_id = unwrap(self._post_model(body, provider_live=provider_live)).model_id
written_at = time.monotonic()
try:
self._await_model_servable(body.model_name, listed_for)
except BaseException:
self.delete_model(model_id)
raise
settle_propagation(written_at)
self.settle(written_at)
return model_id
@step("Add a batch of deployments, then wait once for all of them to propagate")
def register_models(
self, bodies: Sequence[ModelNewBody], *, provider_live: bool = False
) -> tuple[ModelRegistration, ...]:
posted: Final = tuple((body.model_name, self._post_model(body, provider_live=provider_live)) for body in bodies)
written_at: Final = time.monotonic()
registrations: Final = tuple(self._servable_registration(name, result) for name, result in posted)
self.settle(written_at)
return registrations
def _post_model(self, body: ModelNewBody, *, provider_live: bool) -> Result[ModelNewResponse]:
return self.transport.post(
"/model/new",
headers=self.management_headers(),
json=body.model_copy(update={"litellm_params": route_cache_model(
body.litellm_params, provider_edge_base,
enabled=os.environ.get("E2E_PROVIDER_CACHE", "0") == "1" and not provider_live,
mode=body.model_info.mode,
)}),
response_type=ModelNewResponse,
)
def _servable_registration(self, model_name: str, posted: Result[ModelNewResponse]) -> ModelRegistration:
if not isinstance(posted, Success):
return ModelNotRegistered(model_name=model_name, reason=f"/model/new failed: {posted!r}")
match self._servable_outcome(model_name, None):
case Servable():
return ModelRegistered(model_name=model_name, model_id=posted.data.model_id)
case NotServableOn() as missing:
self.delete_model(posted.data.model_id)
return ModelNotRegistered(model_name=model_name, reason=self._not_servable_message(model_name, missing))
def _await_model_servable(self, model_name: str, listed_for: str | None = None) -> None:
"""Block until every replica lists `model_name`, or fail at model_servable_timeout."""
match self._servable_outcome(model_name, listed_for):
case Servable():
return
case NotServableOn() as missing:
raise AssertionError(self._not_servable_message(model_name, missing))
def _servable_outcome(self, model_name: str, listed_for: str | None) -> Servable | NotServableOn:
headers: Final = self.management_headers(listed_for)
outcome: Final = await_servable_everywhere(
return await_servable_everywhere(
{url: self._models_poller(transport, headers) for url, transport in self.replicas.items()},
model_name=model_name,
timeout=self.model_servable_timeout,
@ -728,19 +772,15 @@ class ProxyClient:
now=time.monotonic,
sleep=time.sleep,
)
match outcome:
case Servable():
return
case NotServableOn(replica=replica, last_result=last_result):
raise AssertionError(
servable_timeout_message(
model_name=model_name,
replica=replica,
timeout=self.model_servable_timeout,
db_sync_seconds=self.model_servable_db_sync_seconds,
last_result=last_result,
)
)
def _not_servable_message(self, model_name: str, missing: NotServableOn) -> str:
return servable_timeout_message(
model_name=model_name,
replica=missing.replica,
timeout=self.model_servable_timeout,
db_sync_seconds=self.model_servable_db_sync_seconds,
last_result=missing.last_result,
)
@staticmethod
def _models_poller(transport: Transport, headers: AuthHeaders) -> ModelsPoller:

2
tests/e2e/pytest.ini
View file

@ -13,7 +13,6 @@ markers =
prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set
cli_determinism: drives the real claude CLI for several seconds; deselected unless E2E_CLI_DETERMINISM is set
redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_PERF_TESTS is set
quiet_stack: measures the proxy itself, so it runs while no other test on this host is hitting the stack; every other test waits for it to finish
provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the gateway must reach the pytest host; deselected unless E2E_PROVIDER_EDGE_HOST_REACHABLE is set
owned_gateway: boots its own proxy from source against the stack's Postgres, so it needs DATABASE_URL on the pytest host; deselected unless E2E_OWNED_GATEWAY is set
otel_v2: needs a proxy running with LITELLM_OTEL_V2=true; deselected unless E2E_OTEL_V2 is set
@ -21,3 +20,4 @@ markers =
secret_manager: needs a proxy booted from gateway/secret_manager_<system>_ci_config.yml against that live secret manager; deselected unless E2E_SECRET_MANAGER names the backend (see secret_manager/secret_backends.py)
sagemaker_nova: needs a deployed SageMaker Nova endpoint the proxy's AWS credentials can invoke; deselected unless SAGEMAKER_NOVA_ENDPOINT names it
sagemaker_nova2_lite: needs a deployed SageMaker Nova 2 Lite endpoint the proxy's AWS credentials can invoke; deselected unless SAGEMAKER_NOVA2_LITE_ENDPOINT names it
proxy_wide_config: turns on a proxy-wide setting through /config/update, so it runs on a stack no other suite is driving; deselected unless E2E_PROXY_WIDE_CONFIG_STACK is set

152
tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py → tests/e2e/quota_management/spend_tracking/test_guardrail_cost_spend_e2e.py
View file

@ -17,11 +17,6 @@ from models import (
ChatBody,
ChatMessage,
ChatResponse,
ConfigPatchResponse,
CostDiscountConfig,
CostDiscountConfigResponse,
CostMarginConfig,
CostMarginConfigResponse,
GuardrailCreateBody,
GuardrailCreateResponse,
KeyDeleteBody,
@ -31,14 +26,11 @@ from models import (
)
from spend_e2e_client import SpendClient
pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack]
pytestmark = pytest.mark.e2e
BACKEND_MODEL: Final = "openai/gpt-6.1-sol"
INPUT_RATE: Final = 0.00005
OUTPUT_RATE: Final = 0.0001
DISCOUNT: Final = 0.25
MARGIN_PERCENT: Final = 0.1
MARGIN_FIXED: Final = 0.0005
def _register_bedrock_guardrail(
@ -92,52 +84,6 @@ def _guarded_chat(client: SpendClient, key: str, model: str, name: str) -> Strea
)
def _set_discount(client: SpendClient, values: dict[str, float]) -> float:
unwrap(
client.proxy.transport.patch(
"/config/cost_discount_config",
headers=client.proxy.transport.master,
json=CostDiscountConfig(values),
response_type=ConfigPatchResponse,
)
)
return time.monotonic()
def _set_margin(client: SpendClient, values: dict[str, float | dict[str, float]]) -> float:
unwrap(
client.proxy.transport.patch(
"/config/cost_margin_config",
headers=client.proxy.transport.master,
json=CostMarginConfig(values),
response_type=ConfigPatchResponse,
)
)
return time.monotonic()
def _get_discount(client: SpendClient) -> dict[str, float]:
return unwrap(
client.proxy.transport.get(
"/config/cost_discount_config",
headers=client.proxy.transport.master,
params=NoBody(),
response_type=CostDiscountConfigResponse,
)
).values
def _get_margin(client: SpendClient) -> dict[str, float | dict[str, float]]:
return unwrap(
client.proxy.transport.get(
"/config/cost_margin_config",
headers=client.proxy.transport.master,
params=NoBody(),
response_type=CostMarginConfigResponse,
)
).values
def _register_model(client: SpendClient, resources: ResourceManager, prefix: str) -> str:
model: Final = f"{prefix}-{unique_marker()}"
model_id: Final = client.proxy.create_model(
@ -199,100 +145,7 @@ def scoped_key(client: SpendClient, strict_resources: ResourceManager) -> str:
return key
@pytest.fixture
def restored_pricing_config(client: SpendClient) -> Iterator[None]:
discount: Final = _get_discount(client)
margin: Final = _get_margin(client)
try:
_set_discount(client, {})
reset_at: Final = _set_margin(client, {})
settle_propagation(reset_at)
yield
finally:
try:
_set_discount(client, discount)
finally:
restored_at: Final = _set_margin(client, margin)
settle_propagation(restored_at)
class TestPricingConfigSpend:
@pytest.mark.covers(
"quota_management.spend_tracking.discount_config.logs_cost",
exercised_on=["chat_completions"],
)
@meta(
Subject(
domain=Domain.SPEND_BUDGETS,
route=Route.SPEND_REPORTING,
providers=(Provider.OPENAI,),
models=(BACKEND_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_configured_discount_reaches_persisted_spend_row(
self,
client: SpendClient,
strict_resources: ResourceManager,
scoped_key: str,
restored_pricing_config: None,
) -> None:
discount_set_at: Final = _set_discount(client, {"openai": DISCOUNT})
settle_propagation(discount_set_at)
model: Final = _register_model(client, strict_resources, "discount-priced")
chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16))
assert chat.id and chat.usage
assert chat.usage.prompt_tokens is not None and chat.usage.prompt_tokens > 0
assert chat.usage.completion_tokens is not None and chat.usage.completion_tokens >= 0
row: Final = poll_cost_row(client.proxy, chat.id)
assert row is not None
base_cost: Final = _base_cost(row, chat.usage.prompt_tokens, chat.usage.completion_tokens)
breakdown: Final = row.breakdown
assert breakdown.discount_percent is not None and approx_equal(breakdown.discount_percent, DISCOUNT)
assert breakdown.discount_amount is not None and approx_equal(breakdown.discount_amount, base_cost * DISCOUNT)
assert row.spend is not None and approx_equal(row.spend, base_cost * (1 - DISCOUNT))
@pytest.mark.covers(
"quota_management.spend_tracking.margin_config.logs_cost",
exercised_on=["chat_completions"],
)
@meta(
Subject(
domain=Domain.SPEND_BUDGETS,
route=Route.SPEND_REPORTING,
providers=(Provider.OPENAI,),
models=(BACKEND_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_configured_margin_reaches_persisted_spend_row(
self,
client: SpendClient,
strict_resources: ResourceManager,
scoped_key: str,
restored_pricing_config: None,
) -> None:
margin_set_at: Final = _set_margin(
client, {"openai": {"percentage": MARGIN_PERCENT, "fixed_amount": MARGIN_FIXED}}
)
settle_propagation(margin_set_at)
model: Final = _register_model(client, strict_resources, "margin-priced")
chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16))
assert chat.id and chat.usage
assert chat.usage.prompt_tokens is not None and chat.usage.prompt_tokens > 0
assert chat.usage.completion_tokens is not None and chat.usage.completion_tokens >= 0
row: Final = poll_cost_row(client.proxy, chat.id)
assert row is not None
base_cost: Final = _base_cost(row, chat.usage.prompt_tokens, chat.usage.completion_tokens)
breakdown: Final = row.breakdown
expected_margin: Final = base_cost * MARGIN_PERCENT + MARGIN_FIXED
assert breakdown.margin_percent is not None and approx_equal(breakdown.margin_percent, MARGIN_PERCENT)
assert breakdown.margin_fixed_amount is not None and approx_equal(breakdown.margin_fixed_amount, MARGIN_FIXED)
assert breakdown.margin_total_amount is not None and approx_equal(breakdown.margin_total_amount, expected_margin)
assert row.spend is not None and approx_equal(row.spend, base_cost + expected_margin)
class TestGuardrailCostSpend:
@pytest.mark.covers(
"quota_management.spend_tracking.guardrail_cost.logs_cost",
exercised_on=["chat_completions"],
@ -311,7 +164,6 @@ class TestPricingConfigSpend:
client: SpendClient,
strict_resources: ResourceManager,
scoped_key: str,
restored_pricing_config: None,
) -> None:
identifier: Final = os.environ["BEDROCK_GUARDRAIL_IDENTIFIER"]
version: Final = os.environ["BEDROCK_GUARDRAIL_VERSION"]

45
tests/e2e/stack_lock.py
View file

@ -1,45 +0,0 @@
"""Cross-process reader/writer lock over the proxy stack every xdist worker shares.
Every collected test holds it shared, marker or not, since the Claude Code cells and
other unmarked suites drive the same stack; a `quiet_stack` test holds it exclusive,
and the `gate` file makes a waiting exclusive holder win over readers that arrive
after it."""
from __future__ import annotations
import fcntl
import hashlib
import tempfile
from collections.abc import Generator
from contextlib import ExitStack, contextmanager
from pathlib import Path
from typing import Final
from e2e_config import PROXY_BASE_URL
STACK_DIGEST: Final = hashlib.sha256(PROXY_BASE_URL.encode()).hexdigest()[:12]
LOCK_DIR: Final = Path(tempfile.gettempdir()) / f"litellm-e2e-stack-{STACK_DIGEST}"
GATE_FILE: Final = LOCK_DIR / "gate"
STACK_FILE: Final = LOCK_DIR / "stack"
@contextmanager
def _flock(path: Path, operation: int) -> Generator[None]:
with path.open("a") as handle:
fcntl.flock(handle, operation)
try:
yield
finally:
fcntl.flock(handle, fcntl.LOCK_UN)
@contextmanager
def stack_lock(exclusive: bool) -> Generator[None]:
LOCK_DIR.mkdir(parents=True, exist_ok=True)
if exclusive:
with _flock(GATE_FILE, fcntl.LOCK_EX), _flock(STACK_FILE, fcntl.LOCK_EX):
yield
return
with ExitStack() as held:
with _flock(GATE_FILE, fcntl.LOCK_SH):
held.enter_context(_flock(STACK_FILE, fcntl.LOCK_SH))
yield

2
tests/e2e_harness/AGENTS.md
View file

@ -1,6 +1,6 @@
# e2e harness tests
Tests of the harness under `tests/e2e/` (the transport, the clients, the fixture bundle and replay edge, the stack lock, the IdP launcher, the coverage collector, the JUnit properties and the Claude Code driver), not of the product. They live outside `tests/e2e/` because the Buildkite e2e run copies that folder into the runner image and runs every test in it, so a harness test in there counts as a product test in the nightly numbers. Nothing here needs a proxy, provider keys or the network
Tests of the harness under `tests/e2e/` (the transport, the clients, the fixture bundle and replay edge, the IdP launcher, the coverage collector, the JUnit properties and the Claude Code driver), not of the product. They live outside `tests/e2e/` because the Buildkite e2e run copies that folder into the runner image and runs every test in it, so a harness test in there counts as a product test in the nightly numbers. Nothing here needs a proxy, provider keys or the network
The layout mirrors `tests/e2e/`: `test_e2e_http.py` covers `tests/e2e/e2e_http.py`, `logging/test_datadog_reader.py` covers `tests/e2e/logging/datadog_reader.py`, `claude_code/` covers the driver, builder and probe, `cli_harness/` covers the version resolver, the CLI env lockdown and the rate limiter, and `codex/` covers how the Codex driver reads the CLI's JSONL events and the compactions in a thread's rollout file. Put a new harness test under the folder that mirrors the suite folder whose module it covers

90
tests/e2e_harness/test_proxy_client.py
View file

@ -12,10 +12,11 @@ injected clock, so nothing here monkeypatches anything.
from __future__ import annotations
import json
import time
from builtins import ExceptionGroup
from collections.abc import Callable, Generator, Iterable, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from dataclasses import dataclass, replace
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from itertools import chain, repeat
from queue import SimpleQueue
@ -40,7 +41,10 @@ from models import (
LiteLLMParamsBody,
McpServerCreateBody,
McpServerUpdateBody,
ModelDeleteBody,
ModelInfoBody,
ModelListEntry,
ModelNewBody,
ModelsListResponse,
OrgNewBody,
OrgUpdateBody,
@ -59,6 +63,8 @@ from proxy_client import (
ConvergeOutcome,
CredentialKind,
EverywhereConverged,
ModelNotRegistered,
ModelRegistered,
ModelsPoller,
NeverConvergedOn,
NotConverged,
@ -691,3 +697,85 @@ def test_partial_updates_preserve_explicit_null_at_the_http_boundary(operation:
)
assert json.loads(bodies.get_nowait()) == expected
assert bodies.empty()
@contextmanager
def model_registry(
listed: frozenset[str], refused: frozenset[str]
) -> Generator[tuple[ProxyClient, SimpleQueue[float], SimpleQueue[str]]]:
writes: Final[SimpleQueue[float]] = SimpleQueue()
deleted: Final[SimpleQueue[str]] = SimpleQueue()
class Handler(BaseHTTPRequestHandler):
def log_message(self, format: str, *args: object) -> None:
pass
def reply(self, status: int, body: object) -> None:
self.send_response(status)
self.end_headers()
self.wfile.write(json.dumps(body).encode())
def do_GET(self) -> None:
self.reply(200, {"data": [{"id": name} for name in sorted(listed)]})
def do_POST(self) -> None:
raw: Final = self.rfile.read(int(self.headers.get("Content-Length", "0")))
if self.path == "/model/delete":
deleted.put(ModelDeleteBody.model_validate_json(raw).id)
self.reply(200, {})
return
writes.put(time.monotonic())
name: Final = ModelNewBody.model_validate_json(raw).model_name
if name in refused:
self.reply(500, {"error": "refused"})
return
self.reply(200, {"model_id": f"id-{name}"})
server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
thread: Final = Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True)
thread.start()
url: Final = f"http://127.0.0.1:{server.server_port}"
proxy: Final = build_proxy_client(
base_url=url, control_plane_base_url=url, replica_urls=(url,), control_replica_urls=(url,), master_key="bootstrap"
)
try:
yield replace(proxy, model_servable_timeout=0.3, model_servable_interval=0.01), writes, deleted
finally:
server.shutdown()
server.server_close()
thread.join(timeout=5)
class TestRegisterModels:
def test_a_batch_waits_once_after_its_last_write_and_reports_each_model(self) -> None:
names: Final = ("first", "refused", "ghost", "last")
settled: Final[SimpleQueue[float]] = SimpleQueue()
with (
model_registry(listed=frozenset({"first", "last"}), refused=frozenset({"refused"})) as (
proxy,
writes,
deleted,
),
without_retries(),
):
registrations: Final = replace(proxy, settle=settled.put).register_models(
tuple(
ModelNewBody(
model_name=name, litellm_params=LiteLLMParamsBody(model="synthetic"), model_info=ModelInfoBody()
)
for name in names
)
)
last_write: Final = max(writes.get_nowait() for _ in names)
assert writes.empty()
assert settled.get_nowait() >= last_write
assert settled.empty(), "a batch must wait out propagation once, not once per model"
assert tuple((type(r), r.model_name) for r in registrations) == (
(ModelRegistered, "first"),
(ModelNotRegistered, "refused"),
(ModelNotRegistered, "ghost"),
(ModelRegistered, "last"),
)
assert tuple(r.model_id for r in registrations if isinstance(r, ModelRegistered)) == ("id-first", "id-last")
assert deleted.get_nowait() == "id-ghost"
assert deleted.empty()

117
tests/e2e_harness/test_stack_lock.py
View file

@ -1,117 +0,0 @@
"""Cross-process behavior of the stack lock: readers share it, an exclusive holder waits for
every reader and keeps them out, and a reader arriving behind a waiting exclusive holder
queues behind it instead of starving it."""
from __future__ import annotations
import fcntl
import inspect
import os
import subprocess
import sys
import time
from contextlib import ExitStack
from pathlib import Path
from typing import Final
import pytest
import stack_lock
HARNESS_DIR: Final = Path(inspect.getfile(stack_lock)).resolve().parent
DEADLINE_SECONDS: Final = 30.0
SETTLE_SECONDS: Final = 0.5
HOLDER_SCRIPT: Final = """
import sys, time
from pathlib import Path
from stack_lock import stack_lock
name, mode, release_path, log_path = sys.argv[1:]
def record(event):
with Path(log_path).open("a") as log:
log.write(f"{name} {event}\\n")
record("waiting")
with stack_lock(exclusive=mode == "exclusive"):
record("enter")
while not Path(release_path).exists():
time.sleep(0.02)
record("exit")
"""
def _events(log_path: Path) -> tuple[str, ...]:
return tuple(log_path.read_text().splitlines()) if log_path.exists() else ()
def _wait_for_event(log_path: Path, event: str) -> None:
deadline: Final = time.monotonic() + DEADLINE_SECONDS
while event not in _events(log_path):
if time.monotonic() > deadline:
pytest.fail(f"{event!r} never appeared; events so far: {_events(log_path)}")
time.sleep(0.02)
def _wait_until_gate_is_held_exclusively(gate_path: Path) -> None:
deadline: Final = time.monotonic() + DEADLINE_SECONDS
with gate_path.open("a") as handle:
while True:
try:
fcntl.flock(handle, fcntl.LOCK_SH | fcntl.LOCK_NB)
except BlockingIOError:
return
fcntl.flock(handle, fcntl.LOCK_UN)
if time.monotonic() > deadline:
pytest.fail("no exclusive holder ever took the gate")
time.sleep(0.02)
def _start_holder(held: ExitStack, tmp_path: Path, name: str, mode: str) -> subprocess.Popen[bytes]:
holder: Final = held.enter_context(
subprocess.Popen(
(
sys.executable,
"-P",
"-c",
HOLDER_SCRIPT,
name,
mode,
str(tmp_path / f"release-{name}"),
str(tmp_path / "events"),
),
cwd=HARNESS_DIR,
env={**os.environ, "TMPDIR": str(tmp_path), "PYTHONPATH": str(HARNESS_DIR)},
)
)
held.callback(holder.kill)
return holder
def test_readers_share_exclusive_waits_and_a_waiting_exclusive_beats_later_readers(tmp_path: Path) -> None:
lock_dir: Final = tmp_path / f"litellm-e2e-stack-{stack_lock.STACK_DIGEST}"
lock_dir.mkdir()
log_path: Final = tmp_path / "events"
with ExitStack() as held:
first_reader: Final = _start_holder(held, tmp_path, "A", "shared")
_wait_for_event(log_path, "A enter")
second_reader: Final = _start_holder(held, tmp_path, "R", "shared")
_wait_for_event(log_path, "R enter")
(tmp_path / "release-R").touch()
_wait_for_event(log_path, "R exit")
writer: Final = _start_holder(held, tmp_path, "W", "exclusive")
_wait_until_gate_is_held_exclusively(lock_dir / "gate")
late_reader: Final = _start_holder(held, tmp_path, "B", "shared")
_wait_for_event(log_path, "B waiting")
time.sleep(SETTLE_SECONDS)
(tmp_path / "release-A").touch()
_wait_for_event(log_path, "W enter")
(tmp_path / "release-W").touch()
_wait_for_event(log_path, "B enter")
(tmp_path / "release-B").touch()
for holder in (first_reader, second_reader, writer, late_reader):
assert holder.wait(timeout=DEADLINE_SECONDS) == 0
events: Final = _events(log_path)
assert events.index("R enter") < events.index("A exit")
assert events.index("W enter") > events.index("A exit")
assert events.index("B enter") > events.index("W exit")

32
tests/integration/pricing/test_cost_margin_discount.py
View file

@ -7,7 +7,7 @@ import httpx
import pytest
from integration._support.client import Gateway, eventually, object_value, string_value
from integration._support.database import read_rows
from pydantic import BaseModel, JsonValue
from pydantic import BaseModel, ConfigDict, JsonValue
class CostEstimate(BaseModel):
@ -110,6 +110,25 @@ def _assert_logged_spends(expectations: tuple[tuple[str, float], ...]) -> None:
assert actual == pytest.approx(dict(expectations)), rows
class LoggedBreakdown(BaseModel):
model_config = ConfigDict(frozen=True)
original_cost: float
discount_percent: float | None = None
discount_amount: float | None = None
margin_percent: float | None = None
margin_fixed_amount: float | None = None
margin_total_amount: float | None = None
def _logged_breakdown(request_id: str) -> LoggedBreakdown:
rows: Final = read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,))
assert len(rows) == 1, rows
metadata: Final = rows[0]["metadata"]
parsed: Final = object_value(json.loads(metadata) if isinstance(metadata, str) else metadata)
return LoggedBreakdown.model_validate(parsed["cost_breakdown"])
def _estimated(gateway: Gateway, model: str) -> CostEstimate:
response: Final = gateway.request(
"POST", "/cost/estimate", {"model": model, "input_tokens": 20, "output_tokens": 20}
@ -236,6 +255,17 @@ def test_discount_and_margin_reprice_a_request_and_survive_rejection(gateway: Ga
(streamed_global_request_id, 0.12),
)
)
discounted: Final = _logged_breakdown(discounted_request_id)
assert discounted.original_cost == pytest.approx(0.06), discounted
assert discounted.discount_percent == pytest.approx(0.05), discounted
assert discounted.discount_amount == pytest.approx(0.06 * 0.05), discounted
margined: Final = _logged_breakdown(provider_margin_request_id)
assert margined.original_cost == pytest.approx(0.06), margined
assert margined.discount_percent == pytest.approx(0.05), margined
assert margined.discount_amount == pytest.approx(0.06 * 0.05), margined
assert margined.margin_percent == pytest.approx(0.1), margined
assert margined.margin_fixed_amount == pytest.approx(0.001), margined
assert margined.margin_total_amount == pytest.approx(0.057 * 0.1 + 0.001), margined
rejections: Final = (
(

84
tests/integration/spend/test_failing_request_stored_snapshot.py Normal file
View file

@ -0,0 +1,84 @@
import json
import uuid
from collections.abc import Iterator
from pathlib import Path
from typing import Final
import pytest
import yaml
from integration._support.client import Gateway, eventually, gateway_from_environment, object_value
from integration._support.process import owned_proxy
from pydantic import JsonValue
UNREACHABLE_BASE: Final = "http://127.0.0.1:9/v1"
TRANSCRIPT_TURNS: Final = 40
RETRIES: Final = 2
STORED_REQUEST_BUDGET_KB: Final = 64
def _config_storing_prompts(directory: Path) -> Path:
config: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
general_settings: Final = object_value(config["general_settings"])
path: Final = directory / "store-prompts.yaml"
path.write_text(
yaml.safe_dump({**config, "general_settings": {**general_settings, "store_prompts_in_spend_logs": True}})
)
return path
@pytest.fixture(scope="module")
def gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
with gateway_from_environment() as parent:
directory: Final = tmp_path_factory.mktemp("failing_request_stored_snapshot")
with owned_proxy(parent, directory, {}, config=_config_storing_prompts(directory)) as owned:
yield owned
def _transcript(marker: str) -> list[JsonValue]:
turns: Final[list[JsonValue]] = [
{"role": role, "content": f"turn {turn // 2} {role}"}
for turn, role in enumerate(("user", "assistant") * TRANSCRIPT_TURNS)
]
return [*turns, {"role": "user", "content": marker}]
def _stored_request(gateway: Gateway, call_id: str) -> JsonValue:
rows: Final = eventually(
lambda: gateway.rows('SELECT proxy_server_request FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (call_id,)),
lambda values: len(values) == 1,
seconds=40,
)
stored: Final = rows[0]["proxy_server_request"]
return json.loads(stored) if isinstance(stored, str) else stored
def test_failing_request_with_retries_and_fallback_stores_one_copy_of_the_request(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
primary: Final = f"stored-snapshot-{uuid.uuid4().hex}"
fallback: Final = f"stored-snapshot-fb-{uuid.uuid4().hex}"
for name in (primary, primary, fallback, fallback):
scenario.model(name=name, api_base=UNREACHABLE_BASE, cooldown_time=0)
key: Final = scenario.key(models=[primary, fallback])
marker: Final = f"stored snapshot {uuid.uuid4().hex}"
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": primary,
"max_tokens": 16,
"messages": _transcript(marker),
"router_settings_override": {"num_retries": RETRIES, "fallbacks": [{primary: [fallback]}]},
},
key=key,
)
assert response.status_code != 200, response.text
assert fallback in response.text, response.text
serialized: Final = json.dumps(_stored_request(gateway, response.headers["x-litellm-call-id"]))
stored_kb: Final = len(serialized.encode()) / 1024
assert marker in serialized, serialized[:400]
assert stored_kb <= STORED_REQUEST_BUDGET_KB, (
f"one failing request stored a {stored_kb:.0f} KB request snapshot for a "
f"{2 * TRANSCRIPT_TURNS + 1}-message transcript with {RETRIES} retries and a fallback"
)