From b18188dee2044de608d6b4389bbff7cf13366ff4 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 4 Nov 2026 09:55:53 -0800 Subject: [PATCH] test(e2e): retire the quiet_stack lock and batch the claude_code model registration (#49984) * 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 --- tests/e2e/AGENTS.md | 8 +- tests/e2e/claude_code/conftest.py | 35 ++-- tests/e2e/conftest.py | 21 +-- .../coverage_registry/quota_management.yaml | 2 - tests/e2e/e2e_config.py | 2 +- .../mcp/test_mcp_semantic_tool_filter_e2e.py | 5 +- tests/e2e/models.py | 21 --- tests/e2e/perf/test_database_upgrade.py | 2 +- tests/e2e/perf/test_perf_cases.py | 2 +- .../test_reliability_memory_e2e.py | 111 ++++++------- tests/e2e/perf/test_ui_journeys.py | 1 - tests/e2e/proxy_client.py | 96 +++++++---- tests/e2e/pytest.ini | 2 +- ...2e.py => test_guardrail_cost_spend_e2e.py} | 152 +----------------- tests/e2e/stack_lock.py | 45 ------ tests/e2e_harness/AGENTS.md | 2 +- tests/e2e_harness/test_proxy_client.py | 90 ++++++++++- tests/e2e_harness/test_stack_lock.py | 117 -------------- .../pricing/test_cost_margin_discount.py | 32 +++- .../test_failing_request_stored_snapshot.py | 84 ++++++++++ 20 files changed, 359 insertions(+), 471 deletions(-) rename tests/e2e/{router => perf}/test_reliability_memory_e2e.py (77%) rename tests/e2e/quota_management/spend_tracking/{test_pricing_config_spend_e2e.py => test_guardrail_cost_spend_e2e.py} (52%) delete mode 100644 tests/e2e/stack_lock.py delete mode 100644 tests/e2e_harness/test_stack_lock.py create mode 100644 tests/integration/spend/test_failing_request_stored_snapshot.py diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index 781d8fd051b..b4ca8bd3d50 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -40,11 +40,11 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`) - `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials (API surface, not Playwright) - `a2a/` - the A2A (agent-to-agent) surface: admin registration via `/v1/agents`, proxy-fronted card discovery at `/.well-known/agent-card.json`, and JSON-RPC `message/send` invocation, driving agents backed by the litellm completion bridge (a real provider) and asserting protocol-version normalization (0.3 vs 1.0) -- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below) +- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below). The semantic tool filter test turns on a proxy-wide setting, so it is marked `proxy_wide_config` and runs on the redeployed managed-files stack after the main run, never next to it - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection -- `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory tests (`test_reliability_memory_e2e.py`: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the release-gate check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a loose ceiling over the measured batches nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak) -- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, run by hand since its weekly workflow never had working Bedrock credentials), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares, or a `Ramp` that steps the rate up through the run, built with `ramped()`, whose plan carries each request's send second and whose run reports `sustained_rps`, the last step the pod held, scored in `perf/setup/ramp.py`, with every step's latency, completed rate, non-200s, CPU per request, Redis latency and Postgres connections cut from the samplers' time-stamped `Series` (`perf/setup/series.py`) to the step's window and emitted as `perf_ramp_step__` JUnit properties for the dashboard) and `Limit`s (one metric at one percentile under one ceiling, or over one floor for `sustained_rps`). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; the same overhead as the gateway itself reports it, read from the `x-litellm-overhead-duration-ms` header per request (`litellm_overhead_seconds`, gated next to `gateway_added_seconds`, the outside reading of the same thing, so the two can be charted together) and the gateway's own end-to-end time from `x-litellm-response-duration-ms` (`litellm_response_seconds`), the run mean of the gateway's `litellm_redis_latency` histogram from a `/metrics` scrape before and after the run (`telemetry.py`, `prometheus_redis_latency_mean_seconds`), and the database's own view of the run from the same `pg_stat_statements` rows and a once-a-second `pg_stat_activity` sample: execution seconds per wall second (`postgres_active_sessions`, the average-active-sessions number RDS reports as DB load), the slowest statement's mean (`postgres_slowest_statement_seconds`, named in the detail) and the peak client connection count (`postgres_peak_connections`), each under a ceiling set from the first local runs and due a retune once the lane has its ten-run baseline; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` +- `router/` - routing and reliability behavior (fallbacks, cooldowns) +- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, run by hand since its weekly workflow never had working Bedrock credentials), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). The memory tests (`test_reliability_memory_e2e.py`, run by the perf lane on its own fresh stack: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a loose ceiling over the measured batches, the check for the v1.100.0 retry-breadcrumb leak). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares, or a `Ramp` that steps the rate up through the run, built with `ramped()`, whose plan carries each request's send second and whose run reports `sustained_rps`, the last step the pod held, scored in `perf/setup/ramp.py`, with every step's latency, completed rate, non-200s, CPU per request, Redis latency and Postgres connections cut from the samplers' time-stamped `Series` (`perf/setup/series.py`) to the step's window and emitted as `perf_ramp_step__` JUnit properties for the dashboard) and `Limit`s (one metric at one percentile under one ceiling, or over one floor for `sustained_rps`). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; the same overhead as the gateway itself reports it, read from the `x-litellm-overhead-duration-ms` header per request (`litellm_overhead_seconds`, gated next to `gateway_added_seconds`, the outside reading of the same thing, so the two can be charted together) and the gateway's own end-to-end time from `x-litellm-response-duration-ms` (`litellm_response_seconds`), the run mean of the gateway's `litellm_redis_latency` histogram from a `/metrics` scrape before and after the run (`telemetry.py`, `prometheus_redis_latency_mean_seconds`), and the database's own view of the run from the same `pg_stat_statements` rows and a once-a-second `pg_stat_activity` sample: execution seconds per wall second (`postgres_active_sessions`, the average-active-sessions number RDS reports as DB load), the slowest statement's mean (`postgres_slowest_statement_seconds`, named in the detail) and the peak client connection count (`postgres_peak_connections`), each under a ceiling set from the first local runs and due a retune once the lane has its ten-run baseline; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at, see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `secret_manager/` - the gateway's `key_management_system` against a real secret manager: deployment keys resolved from it (`os.environ/` where the name exists only in the manager) and virtual keys written to and deleted from it. The tests are backend-agnostic and each backend is its own lane, because the setting is global to the proxy: `E2E_SECRET_MANAGER=` opts in and picks the backend from `secret_backends.BACKENDS`, the proxy is booted from `gateway/secret_manager__ci_config.yml` against the live manager, and the tests reach that manager through the backend's `SecretStore` (`secret_store_.py`). A test needing something not every backend does carries `requires_capability(...)` and is deselected on lanes that lack it. `secret_manager/backend.sh up ` runs a backend in Docker and writes the proxy's and the tests' env. Marked `secret_manager`, deselected unless `E2E_SECRET_MANAGER` is set, and kept out of the per-PR selector. Backends today: `hashicorp_vault` and `cyberark` (CyberArk Conjur, which cannot delete, so the delete test is Vault-only) - `gateway/` - proxy configuration only (`litellm-config.yml`). No tests @@ -100,7 +100,7 @@ Each suite provides its own `client` fixture (see `llm_translation/passthrough_c Request and response bodies are typed pydantic models in `models.py`. Only the fields a test reads are modelled, and nothing passes raw dicts. Outcomes come back as a `Result[R]` tagged union (`Success`, `NetworkError`, `UnauthorizedError`, `RateLimitedError`, `ValidationError`, `UnknownApiError`). Handle them with `match`, or call `unwrap(...)` when a non-success should fail the test. The harness hard-fails and never skips: a test marked `e2e` fails when no proxy answers its liveness probe, and once a request reaches the proxy any wrong behavior is likewise a hard failure, so a missing proxy turns the run red instead of being mistaken for a pass -Mark live tests with `@pytest.mark.e2e` (on the class or the module). Coverage of the harness itself lives outside the suite in `tests/e2e_harness/` (see its `AGENTS.md`) and runs without a proxy, so nothing under `tests/e2e/` is a markerless test. Add `@pytest.mark.quiet_stack` to a test that measures the proxy itself (RSS, latency): the shared stack lock in `stack_lock.py` then runs it while no other test on the host is hitting the stack, marked or not, so the reading depends only on the test's own traffic. Use `scoped_key` for a fresh all-models key that auto-deletes, `resources` when you need to create and tear down more than a key, and `unique_marker()` from `e2e_config` to keep prompts, tags, and customer ids from colliding across concurrent runs and the shared response cache +Mark live tests with `@pytest.mark.e2e` (on the class or the module). Coverage of the harness itself lives outside the suite in `tests/e2e_harness/` (see its `AGENTS.md`) and runs without a proxy, so nothing under `tests/e2e/` is a markerless test. A test that measures the proxy itself (RSS, latency) cannot share the stack with the rest of the suite, so it goes in `perf/`, marked `perf`. Use `scoped_key` for a fresh all-models key that auto-deletes, `resources` when you need to create and tear down more than a key, and `unique_marker()` from `e2e_config` to keep prompts, tags, and customer ids from colliding across concurrent runs and the shared response cache ## Record and replay fixtures diff --git a/tests/e2e/claude_code/conftest.py b/tests/e2e/claude_code/conftest.py index b2442d492cf..14201ef6846 100644 --- a/tests/e2e/claude_code/conftest.py +++ b/tests/e2e/claude_code/conftest.py @@ -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 diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 4c7c57ca18f..8fb96bef2f9 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -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. diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index cbaf46130d6..49997aef162 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -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"} diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 724aa2362e7..6316d22d3fa 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -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: diff --git a/tests/e2e/mcp/test_mcp_semantic_tool_filter_e2e.py b/tests/e2e/mcp/test_mcp_semantic_tool_filter_e2e.py index 6e564f4625a..1ec2a6822e3 100644 --- a/tests/e2e/mcp/test_mcp_semantic_tool_filter_e2e.py +++ b/tests/e2e/mcp/test_mcp_semantic_tool_filter_e2e.py @@ -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 diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 9c12e2146c9..25683d64313 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -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" diff --git a/tests/e2e/perf/test_database_upgrade.py b/tests/e2e/perf/test_database_upgrade.py index fba6af30014..a9f5cc43fe6 100644 --- a/tests/e2e/perf/test_database_upgrade.py +++ b/tests/e2e/perf/test_database_upgrade.py @@ -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") diff --git a/tests/e2e/perf/test_perf_cases.py b/tests/e2e/perf/test_perf_cases.py index db657a04b00..0a137e84572 100644 --- a/tests/e2e/perf/test_perf_cases.py +++ b/tests/e2e/perf/test_perf_cases.py @@ -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 diff --git a/tests/e2e/router/test_reliability_memory_e2e.py b/tests/e2e/perf/test_reliability_memory_e2e.py similarity index 77% rename from tests/e2e/router/test_reliability_memory_e2e.py rename to tests/e2e/perf/test_reliability_memory_e2e.py index fb114092303..2066137bb72 100644 --- a/tests/e2e/router/test_reliability_memory_e2e.py +++ b/tests/e2e/perf/test_reliability_memory_e2e.py @@ -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, diff --git a/tests/e2e/perf/test_ui_journeys.py b/tests/e2e/perf/test_ui_journeys.py index 63eefedec68..11b2cd1051d 100644 --- a/tests/e2e/perf/test_ui_journeys.py +++ b/tests/e2e/perf/test_ui_journeys.py @@ -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"), ] diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 8cc7bf0c2b4..9b842a723d8 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -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: diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index f15cfd4e98b..891c5fc282c 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -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__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 diff --git a/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py b/tests/e2e/quota_management/spend_tracking/test_guardrail_cost_spend_e2e.py similarity index 52% rename from tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py rename to tests/e2e/quota_management/spend_tracking/test_guardrail_cost_spend_e2e.py index 8f3bed3f89b..a77ca91ad3a 100644 --- a/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_guardrail_cost_spend_e2e.py @@ -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"] diff --git a/tests/e2e/stack_lock.py b/tests/e2e/stack_lock.py deleted file mode 100644 index 06df7a20b6a..00000000000 --- a/tests/e2e/stack_lock.py +++ /dev/null @@ -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 diff --git a/tests/e2e_harness/AGENTS.md b/tests/e2e_harness/AGENTS.md index 83e573ae249..4f507c2fe7b 100644 --- a/tests/e2e_harness/AGENTS.md +++ b/tests/e2e_harness/AGENTS.md @@ -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 diff --git a/tests/e2e_harness/test_proxy_client.py b/tests/e2e_harness/test_proxy_client.py index e1615ec5f65..d352e5c8d36 100644 --- a/tests/e2e_harness/test_proxy_client.py +++ b/tests/e2e_harness/test_proxy_client.py @@ -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() diff --git a/tests/e2e_harness/test_stack_lock.py b/tests/e2e_harness/test_stack_lock.py deleted file mode 100644 index b275ca8c755..00000000000 --- a/tests/e2e_harness/test_stack_lock.py +++ /dev/null @@ -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") diff --git a/tests/integration/pricing/test_cost_margin_discount.py b/tests/integration/pricing/test_cost_margin_discount.py index b6cda8f125f..613fe72ed70 100644 --- a/tests/integration/pricing/test_cost_margin_discount.py +++ b/tests/integration/pricing/test_cost_margin_discount.py @@ -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 = ( ( diff --git a/tests/integration/spend/test_failing_request_stored_snapshot.py b/tests/integration/spend/test_failing_request_stored_snapshot.py new file mode 100644 index 00000000000..4f016c33f37 --- /dev/null +++ b/tests/integration/spend/test_failing_request_stored_snapshot.py @@ -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" + )