diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index b31f5a6d4e5..4c8aeb64e85 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -21,7 +21,7 @@ jobs: pull-requests: read actions: read runs-on: ubuntu-latest - timeout-minutes: 15 + timeout-minutes: 20 services: redis: image: redis@sha256:e2debfb7956fa12c7ddc79d7e645c8cf26b30c99a6e9161ea9bf4171e1668a5f diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index 8b9132cf05a..e7a26c12db8 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -44,7 +44,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `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 fixed budget 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, driven by `.github/workflows/weekly_load_anomaly.yml`), 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 +- `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, driven by `.github/workflows/weekly_load_anomaly.yml`), 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) and `Limit`s (one metric at one percentile under one ceiling). 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 transactions and Redis commands each request cost (`postgres_calls_per_request`, `redis_calls_per_request`, from `pg_stat_database` and Redis `INFO` read by `datastore.py` before and after the run). 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 diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 94e135535ce..9198f6c54d6 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -481,15 +481,40 @@ def get_external[R: BaseModel]( response_type: type[R], headers: BaseModel | None = None, timeout: float = 30.0, + ca_bundle: str | None = None, ) -> Result[R]: """GET an absolute URL outside the proxy (e.g. a public /.well-known document). Unlike the transport wrappers there is no proxy base url and no proxy auth; the - response still gets the same tagged-union classification as every other call.""" + response still gets the same tagged-union classification as every other call. + `ca_bundle` is a PEM path for a server the default trust store does not cover.""" try: resp = requests.get( url, headers={"Accept": "application/json", **(_headers(headers) if headers is not None else {})}, timeout=timeout, + verify=ca_bundle if ca_bundle is not None else True, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return classify(resp, response_type) + + +def post_external[R: BaseModel]( + url: str, + *, + json: BaseModel, + response_type: type[R], + headers: BaseModel | None = None, + timeout: float = 30.0, +) -> Result[R]: + """POST JSON to an absolute URL outside the proxy (e.g. a mock server's control + route), classified the same way as every other call.""" + try: + resp = requests.post( + url, + headers={"Accept": "application/json", **(_headers(headers) if headers is not None else {})}, + json=wire_body(json), + timeout=timeout, ) except requests.RequestException as exc: return NetworkError(message=str(exc)) diff --git a/tests/e2e/gateway/perf_ci_config.yml b/tests/e2e/gateway/perf_ci_config.yml new file mode 100644 index 00000000000..f5ec2e22693 --- /dev/null +++ b/tests/e2e/gateway/perf_ci_config.yml @@ -0,0 +1,77 @@ +# Gateway config for the perf cases (tests/e2e/perf/test_perf_cases.py), the local twin of the +# proxy_config block in project-releaser's .buildkite/e2e-perf-gateway-values.yaml: the same model +# groups over the mock providers, the same two guardrails, so a laptop run and the lane read the same +# config. One deployment per mock provider pod; how long the mocks wait before answering is set by the +# case at setup (mock_latency), not here. Teams, users and keys are created by the case at run time. +general_settings: + master_key: os.environ/LITELLM_MASTER_KEY + database_url: os.environ/DATABASE_URL + store_model_in_db: true + proxy_config_reload_interval_seconds: 7 + enable_jwt_auth: true + litellm_jwtauth: + user_id_jwt_field: sub + user_email_jwt_field: email + team_ids_jwt_field: groups + user_id_upsert: true + proxy_batch_write_at: 5 + database_connection_pool_limit: 10 + +litellm_settings: + cache: true + cache_params: + type: redis + host: os.environ/REDIS_HOST + port: os.environ/REDIS_PORT + supported_call_types: [] + callbacks: + - prometheus + - otel + require_auth_for_metrics_endpoint: false + drop_params: true + request_timeout: 120 + +router_settings: + num_retries: 2 + retry_after: 0 + allowed_fails: 3 + cooldown_time: 5 + routing_strategy: simple-shuffle + +model_list: + - model_name: gpt-6-luna + litellm_params: + model: openai/gpt-6-luna + api_base: os.environ/E2E_PERF_MOCK_OPENAI_URL + api_key: perf-mock-not-checked + - model_name: claude-sonnet-5 + litellm_params: + model: anthropic/claude-sonnet-5 + api_base: os.environ/E2E_PERF_MOCK_ANTHROPIC_URL + api_key: perf-mock-not-checked + - model_name: claude-sonnet-5-bedrock + litellm_params: + model: bedrock/anthropic.claude-sonnet-5 + aws_bedrock_runtime_endpoint: os.environ/E2E_PERF_MOCK_BEDROCK_URL + aws_region_name: us-east-1 + aws_access_key_id: perf-mock-not-checked + aws_secret_access_key: perf-mock-not-checked +guardrails: + - guardrail_name: perf-pre + litellm_params: + guardrail: generic_guardrail_api + mode: pre_call + api_base: os.environ/E2E_PERF_MOCK_GUARDRAIL_URL + default_on: true + streaming_end_of_stream_only: true + additional_provider_specific_params: + latency_s: 0.05 + - guardrail_name: perf-post + litellm_params: + guardrail: generic_guardrail_api + mode: post_call + api_base: os.environ/E2E_PERF_MOCK_GUARDRAIL_URL + default_on: true + streaming_end_of_stream_only: true + additional_provider_specific_params: + latency_s: 0.05 diff --git a/tests/e2e/perf/cases/__init__.py b/tests/e2e/perf/cases/__init__.py new file mode 100644 index 00000000000..076c6f4bd2f --- /dev/null +++ b/tests/e2e/perf/cases/__init__.py @@ -0,0 +1,22 @@ +"""Every module in this package that defines a `CASE` is one perf test.""" + +from __future__ import annotations + +import importlib +import pkgutil +from typing import Final + +from setup.performance_test_case import PerformanceTestCase + + +def _load() -> tuple[PerformanceTestCase, ...]: + found: list[PerformanceTestCase] = [] + for module_info in pkgutil.iter_modules(__path__): + module = importlib.import_module(f"{__name__}.{module_info.name}") + case: object = getattr(module, "CASE", None) + if isinstance(case, PerformanceTestCase): + found.append(case) + return tuple(sorted(found, key=lambda case: case.name)) + + +CASES: Final = _load() diff --git a/tests/e2e/perf/cases/small_bodies_5rps.py b/tests/e2e/perf/cases/small_bodies_5rps.py new file mode 100644 index 00000000000..e118e453396 --- /dev/null +++ b/tests/e2e/perf/cases/small_bodies_5rps.py @@ -0,0 +1,55 @@ +"""Small bodies at a low rate on the small gateway: the shape of github issue #46276 (short chat +requests, master key and virtual keys, no JWT, no guardrails) at 5 rps on one 4-worker pod, far +under what the pod sustains, so nothing queues and gateway_added_seconds is the gateway's own +per-request overhead. A slower middleware chain or an extra Redis round trip per request shows up +directly instead of hiding behind queue wait.""" + +from __future__ import annotations + +from typing import Final + +from setup.performance_test_case import ( + Limit, + MockLatency, + PerformanceTestCase, + RequestShape, + Team, + Traffic, + VirtualKey, +) + +TINY: Final = RequestShape("tiny", input_tokens=(2_500, 4_000)) + +CASE: Final = PerformanceTestCase( + name="small_bodies_5rps", + teams=(Team("team-a", models=("gpt-6-luna", "claude-sonnet-5", "claude-sonnet-5-bedrock")),), + users=(), + virtual_keys=( + VirtualKey("key-a1", team="team-a"), + VirtualKey("key-a2", team="team-a"), + VirtualKey("key-plain"), + ), + model_groups=("gpt-6-luna", "claude-sonnet-5", "claude-sonnet-5-bedrock"), + mock_latency=MockLatency(seconds=1.0, jitter_seconds=0.25), + guardrails=(), + traffic=Traffic( + seed=11, + rps=5.0, + total_requests=3000, + endpoints={"chat_completions": 0.6, "responses": 0.2, "messages": 0.2}, + shapes={TINY: 1.0}, + stream_share=0.5, + intervene_share=0.0, + block_share=0.0, + concurrency=20, + ), + limits=( + Limit("unexpected_status_count", max=3), + Limit("gateway_added_seconds", percentile=50, max=0.38), + Limit("time_to_first_token_seconds", percentile=99, max=3.0), + Limit("memory_growth_megabytes", max=100), + Limit("postgres_calls_per_request", max=13), + Limit("redis_calls_per_request", max=35), + Limit("cpu_seconds_per_request", max=0.19), + ), +) diff --git a/tests/e2e/perf/session_anomaly.py b/tests/e2e/perf/session_anomaly.py index b68482e965f..e105b076943 100644 --- a/tests/e2e/perf/session_anomaly.py +++ b/tests/e2e/perf/session_anomaly.py @@ -70,37 +70,24 @@ class AnomalyReport: @property def warm_cache_read_share(self) -> float: - billed = ( - self.warm_uncached_input_tokens - + self.warm_cache_read_tokens - + self.warm_cache_creation_tokens - ) + billed = self.warm_uncached_input_tokens + self.warm_cache_read_tokens + self.warm_cache_creation_tokens return self.warm_cache_read_tokens / billed if billed else 0.0 def _system_prefix_block(marker: str) -> TextBlock: - text = " ".join( - f"Project context paragraph {index} for session {marker}." for index in range(300) - ) + text = " ".join(f"Project context paragraph {index} for session {marker}." for index in range(300)) return TextBlock(text=text, cache_control=CacheControl()) def _user_turn_text(marker: str, turn_index: int) -> str: - notes = " ".join( - f"Working note {index} of turn {turn_index} in session {marker}." - for index in range(80) - ) + notes = " ".join(f"Working note {index} of turn {turn_index} in session {marker}." for index in range(80)) return f"Reply with one short sentence.\n{notes}" def _reminder_turn() -> RichMessage: return RichMessage( role="system", - content=[ - TextBlock( - text="Keep the answer to one short sentence." - ) - ], + content=[TextBlock(text="Keep the answer to one short sentence.")], ) @@ -127,9 +114,7 @@ def retried( return retried(call, attempts - 1, backoff_seconds, sleep) -def _metric( - result: Result[SessionMessagesResponse], turn_index: int, latency_seconds: float -) -> TurnMetric: +def _metric(result: Result[SessionMessagesResponse], turn_index: int, latency_seconds: float) -> TurnMetric: if isinstance(result, Success): usage = result.data.usage return TurnMetric( @@ -167,11 +152,7 @@ def _drive_turns( return () user_turn = RichMessage( role="user", - content=[ - TextBlock( - text=_user_turn_text(marker, turn_index), cache_control=CacheControl() - ) - ], + content=[TextBlock(text=_user_turn_text(marker, turn_index), cache_control=CacheControl())], ) started = time.monotonic() result = retried( @@ -190,9 +171,7 @@ def _drive_turns( turn = _metric(result, turn_index, time.monotonic() - started) if not isinstance(result, Success): return (turn,) - assistant_turn = RichMessage( - role="assistant", content=[TextBlock(text=result.data.text or "Understood.")] - ) + assistant_turn = RichMessage(role="assistant", content=[TextBlock(text=result.data.text or "Understood.")]) return ( turn, *_drive_turns( @@ -242,9 +221,7 @@ def run_concurrent_sessions( ) -> tuple[TurnMetric, ...]: with ThreadPoolExecutor(max_workers=sessions) as pool: futures = [ - pool.submit( - run_session, transport, key, model, turns_per_session, attempts_per_turn - ) + pool.submit(run_session, transport, key, model, turns_per_session, attempts_per_turn) for _ in range(sessions) ] return tuple(turn for future in futures for turn in future.result()) @@ -296,7 +273,5 @@ def summarize(turns: tuple[TurnMetric, ...], planned_turns: int) -> AnomalyRepor warm_uncached_input_tokens=sum(turn.uncached_input_tokens for turn in warm), warm_cache_read_tokens=sum(turn.cache_read_tokens for turn in warm), warm_cache_creation_tokens=sum(turn.cache_creation_tokens for turn in warm), - p95_turn_seconds=_p95( - tuple(turn.latency_seconds for turn in turns if turn.ok) - ), + p95_turn_seconds=_p95(tuple(turn.latency_seconds for turn in turns if turn.ok)), ) diff --git a/tests/e2e/perf/setup/__init__.py b/tests/e2e/perf/setup/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/e2e/perf/setup/corpus/session.py b/tests/e2e/perf/setup/corpus/session.py new file mode 100644 index 00000000000..9f90048220c --- /dev/null +++ b/tests/e2e/perf/setup/corpus/session.py @@ -0,0 +1,254 @@ +"""A coding-agent session as the gateway sees it. The system prompt and tool set come from this +directory, the turns are the assistant reading files and the files coming back as tool results, and the +files are real Python files from tests/e2e, chosen by the request's marker. The input budget is spent +on turns, so a 300k request is a longer session over bigger files than a 100k one, in the wire format of +the endpoint it goes to.""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from functools import cache +from pathlib import Path +from typing import Final, Literal + +from pydantic import BaseModel, ConfigDict, TypeAdapter + +Endpoint = Literal["chat", "responses", "messages"] + +CORPUS_DIR: Final = Path(__file__).resolve().parent +E2E_ROOT: Final = CORPUS_DIR.parents[2] +CHARS_PER_TOKEN: Final = 4 +_FILE_BYTES: Final = (2_000, 80_000) +_STRIDE: Final = 7919 +_QUESTION: Final = ( + "Which of the files above raise an exception that the caller is not told about? Answer with the paths only." +) +_CACHE_CONTROL: Final = {"type": "ephemeral"} + + +class ToolSpec(BaseModel): + model_config = ConfigDict(frozen=True) + + name: str + description: str + input_schema: dict[str, object] + + +@dataclass(frozen=True, slots=True) +class Step: + """One tool round trip: the assistant asks to read `path`, the user turn carries `result`.""" + + call_id: str + path: str + offset: int + limit: int + result: str + + +@dataclass(frozen=True, slots=True) +class Session: + system: str + tools: tuple[ToolSpec, ...] + steps: tuple[Step, ...] + question: str + + +def estimate_tokens(text: str) -> int: + return max(1, len(text) // CHARS_PER_TOKEN) + + +@cache +def system_prompt() -> str: + return (CORPUS_DIR / "system_prompt.md").read_text() + + +@cache +def tool_specs() -> tuple[ToolSpec, ...]: + return tuple(TypeAdapter(list[ToolSpec]).validate_json((CORPUS_DIR / "tools.json").read_bytes())) + + +@cache +def corpus_files() -> tuple[Path, ...]: + """Python files under tests/e2e between 2 KB and 80 KB, in path order so every run sees the same list. + + The e2e tree is what the runner image ships, so a laptop and the lane read the same files at the same sha. + """ + low, high = _FILE_BYTES + found: Final = tuple(path for path in sorted(E2E_ROOT.rglob("*.py")) if low <= path.stat().st_size <= high) + assert found, f"no corpus files under {E2E_ROOT}" + return found + + +def _digest(seed: str) -> int: + return int.from_bytes(hashlib.sha1(seed.encode()).digest()[:6], "big") + + +def _numbered(lines: tuple[str, ...], offset: int) -> str: + return "\n".join(f"{offset + index:>6}\t{line}" for index, line in enumerate(lines, start=1)) + + +def _step(marker: str, index: int, budget_tokens: int) -> Step: + files: Final = corpus_files() + path: Final = files[(_digest(marker) + index * _STRIDE) % len(files)] + lines: Final = tuple(path.read_text(errors="replace").splitlines()) + whole: Final = _numbered(lines, 1) + if estimate_tokens(whole) <= budget_tokens: + return Step(_call_id(marker, index), f"tests/e2e/{path.relative_to(E2E_ROOT)}", 1, len(lines), whole) + keep: Final = max(1, len(lines) * budget_tokens * CHARS_PER_TOKEN // max(1, len(whole))) + return Step(_call_id(marker, index), f"tests/e2e/{path.relative_to(E2E_ROOT)}", 1, keep, _numbered(lines[:keep], 1)) + + +def _call_id(marker: str, index: int) -> str: + return hashlib.sha1(f"{marker}:call{index}".encode()).hexdigest()[:24] + + +def build_session(marker: str, input_tokens: int, *, tools: bool) -> Session: + """Spend `input_tokens` on file-reading turns after the system prompt, the tool set and the question.""" + specs: Final = tool_specs() if tools else () + fixed: Final = estimate_tokens(system_prompt()) + sum(estimate_tokens(spec.model_dump_json()) for spec in specs) + question: Final = f"{marker} {_QUESTION}" + steps: list[Step] = [] # mutable-ok: each step's size decides whether another fits + spent = fixed + estimate_tokens(question) # rebind-ok: running total over the steps + while spent < input_tokens: + steps.append(step := _step(marker, len(steps), input_tokens - spent)) + spent += estimate_tokens(step.result) + 40 + return Session(system_prompt(), specs, tuple(steps), question) + + +def _read_arguments(step: Step) -> dict[str, object]: + return {"file_path": f"/workspace/litellm/{step.path}", "offset": step.offset, "limit": step.limit} + + +def _chat_turns(session: Session) -> Sequence[Mapping[str, object]]: + if not session.tools: + return [ + turn + for step in session.steps + for turn in ( + {"role": "user", "content": f"Review this file.\n\n{step.path}\n{step.result}"}, + {"role": "assistant", "content": f"Read {step.path}, {step.limit} lines. Continue."}, + ) + ] + return [ + turn + for step in session.steps + for turn in ( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": f"call_{step.call_id}", + "type": "function", + "function": {"name": "Read", "arguments": json.dumps(_read_arguments(step))}, + } + ], + }, + {"role": "tool", "tool_call_id": f"call_{step.call_id}", "content": step.result}, + ) + ] + + +def _responses_input(session: Session) -> Sequence[Mapping[str, object]]: + if not session.tools: + return _chat_turns(session) + return [ + turn + for step in session.steps + for turn in ( + { + "type": "function_call", + "call_id": f"call_{step.call_id}", + "name": "Read", + "arguments": json.dumps(_read_arguments(step)), + }, + {"type": "function_call_output", "call_id": f"call_{step.call_id}", "output": step.result}, + ) + ] + + +def _messages_turns(session: Session) -> Sequence[Mapping[str, object]]: + if not session.tools: + return _chat_turns(session) + last: Final = len(session.steps) - 1 + return [ + turn + for index, step in enumerate(session.steps) + for turn in ( + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": f"toolu_{step.call_id}", "name": "Read", "input": _read_arguments(step)} + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": f"toolu_{step.call_id}", + "content": step.result, + **({"cache_control": _CACHE_CONTROL} if index == last else {}), + } + ], + }, + ) + ] + + +def request_body( + *, endpoint: Endpoint, model: str, marker: str, input_tokens: int, tools: bool, output_tokens: int, stream: bool +) -> dict[str, object]: + """The session as one request in the endpoint's wire format, the marker in the final user message.""" + session: Final = build_session(marker, input_tokens, tools=tools) + question: Final = {"role": "user", "content": session.question} + if endpoint == "responses": + return { + "model": model, + "instructions": session.system, + "input": [*_responses_input(session), question], + "max_output_tokens": output_tokens, + "stream": stream, + **( + { + "tools": [ + {"type": "function", "name": s.name, "description": s.description, "parameters": s.input_schema} + for s in session.tools + ] + } + if session.tools + else {} + ), + } + if endpoint == "messages": + return { + "model": model, + "system": [{"type": "text", "text": session.system, "cache_control": _CACHE_CONTROL}], + "messages": [*_messages_turns(session), question], + "max_tokens": output_tokens, + "stream": stream, + **({"tools": [spec.model_dump() for spec in session.tools]} if session.tools else {}), + } + return { + "model": model, + "messages": [{"role": "system", "content": session.system}, *_chat_turns(session), question], + "max_completion_tokens": output_tokens, + **({"reasoning_effort": "none"} if session.tools else {}), + "stream": stream, + **( + { + "tools": [ + { + "type": "function", + "function": {"name": s.name, "description": s.description, "parameters": s.input_schema}, + } + for s in session.tools + ] + } + if session.tools + else {} + ), + } diff --git a/tests/e2e/perf/setup/corpus/system_prompt.md b/tests/e2e/perf/setup/corpus/system_prompt.md new file mode 100644 index 00000000000..d2e8c76be2d --- /dev/null +++ b/tests/e2e/perf/setup/corpus/system_prompt.md @@ -0,0 +1,132 @@ +You are an interactive CLI tool that helps users with software engineering tasks. Use the instructions below and the tools available to you to assist the user. + +IMPORTANT: Assist with defensive security tasks only. Refuse to create, modify, or improve code that may be used maliciously. Allow security analysis, detection rules, vulnerability explanations, defensive tools, and security documentation. +IMPORTANT: You must NEVER generate or guess URLs for the user unless you are confident that the URLs are for helping the user with programming. You may use URLs provided by the user in their messages or local files. + +If the user asks for help or wants to give feedback inform them of the following: +- /help: Get help with using the tool +- To give feedback, users should report the issue at the project's issue tracker + +When the user directly asks about the tool (eg 'can the tool do...', 'does the tool have...') or asks in second person (eg 'are you able...', 'can you do...'), first use the WebFetch tool to gather information to answer the question from the documentation. + +# Tone and style +You should be concise, direct, and to the point. When you run a non-trivial bash command, you should explain what the command does and why you are running it, to make sure the user understands what you are doing (this is especially important when you are running a command that will make changes to the user's system). +Remember that your output will be displayed on a command line interface. Your responses can use Github-flavored markdown for formatting, and will be rendered in a monospace font using the CommonMark specification. +Output text to communicate with the user; all text you output outside of tool use is displayed to the user. Only use tools to complete tasks. Never use tools like Bash or code comments as means to communicate with the user during the session. +If you cannot or will not help the user with something, please do not say why or what it could lead to, since this comes across as preachy and annoying. Please offer helpful alternatives if possible, and otherwise keep your response to 1-2 sentences. +Only use emojis if the user explicitly requests it. Avoid using emojis in all communication unless asked. +IMPORTANT: You should minimize output tokens as much as possible while maintaining helpfulness, quality, and accuracy. Only address the specific query or task at hand, avoiding tangential information unless absolutely critical for completing the request. If you can answer in 1-3 sentences or a short paragraph, please do. +IMPORTANT: You should NOT answer with unnecessary preamble or postamble (such as explaining your code or summarizing your action), unless the user asks you to. +IMPORTANT: Keep your responses short, since they will be displayed on a command line interface. You MUST answer concisely with fewer than 4 lines (not including tool use or code generation), unless user asks for detail. Answer the user's question directly, without elaboration, explanation, or details. One word answers are best. Avoid introductions, conclusions, and explanations. You MUST avoid text before/after your response, such as "The answer is .", "Here is the content of the file..." or "Based on the information provided, the answer is..." or "Here is what I will do next...". Here are some examples to demonstrate appropriate verbosity: + +user: 2 + 2 +assistant: 4 + + + +user: what is 2+2? +assistant: 4 + + + +user: is 11 a prime number? +assistant: Yes + + + +user: what command should I run to list files in the current directory? +assistant: ls + + + +user: what command should I run to watch files in the current directory? +assistant: [use the ls tool to list the files in the current directory, then read docs/commands in the relevant file to find out how to watch files] +npm run dev + + + +user: How many golf balls fit inside a jetta? +assistant: 150000 + + + +user: what files are in the directory src/? +assistant: [runs ls and sees foo.c, bar.c, baz.c] +user: which file contains the implementation of foo? +assistant: src/foo.c + + + +user: write tests for new feature +assistant: [uses grep and glob search tools to find where similar tests are defined, uses concurrent read file tool use blocks in one single message to read relevant files at the same time, uses edit file tool to write new tests] + + +# Proactiveness +You are allowed to be proactive, but only when the user asks you to do something. You should strive to strike a balance between: +1. Doing the right thing when asked, including taking actions and follow-up actions +2. Not surprising the user with actions you take without asking +For example, if the user asks you how to approach something, you should do your best to answer their question first, and not immediately jump into taking actions. +3. Do not add additional code explanation summary unless requested by the user. After working on a file, just stop, rather than providing an explanation of what you did. + +# Following conventions +When making changes to files, first understand the file's code conventions. Mimic code style, use existing libraries and utilities, and follow existing patterns. +- NEVER assume that a given library is available, even if it is well known. Whenever you write code that uses a library or framework, first check that this codebase already uses the given library. For example, you might look at neighboring files, or check the package.json (or cargo.toml, and so on depending on the language). +- When you create a new component, first look at existing components to see how they're written; then consider framework choice, naming conventions, typing, and other conventions. +- When you edit a piece of code, first look at the code's surrounding context (especially its imports) to understand the code's choice of frameworks and libraries. Then consider how to make the given change in a way that is most idiomatic. +- Always follow security best practices. Never introduce code that exposes or logs secrets and keys. Never commit secrets or keys to the repository. + +# Code style +- IMPORTANT: DO NOT ADD ***ANY*** COMMENTS unless asked + +# Task Management +You have access to the TodoWrite and TodoRead tools to help you manage and plan tasks. Use these tools VERY frequently to ensure that you are tracking your tasks and giving the user visibility into your progress. +These tools are also EXTREMELY helpful for planning tasks, and for breaking down larger complex tasks into smaller steps. If you do not use this tool when planning, you may forget to do important tasks - and that is unacceptable. +It is critical that you mark todos as completed as soon as you are done with a task. Do not batch up multiple tasks before marking them as completed. + +# Doing tasks +The user will primarily request you perform software engineering tasks. This includes solving bugs, adding new functionality, refactoring code, explaining code, and more. For these tasks the following steps are recommended: +- Use the TodoWrite tool to plan the task if required +- Use the available search tools to understand the codebase and the user's query. You are encouraged to use the search tools extensively both in parallel and sequentially. +- Implement the solution using all tools available to you +- Verify the solution if possible with tests. NEVER assume specific test framework or test script. Check the README or search codebase to determine the testing approach. +- VERY IMPORTANT: When you have completed a task, you MUST run the lint and typecheck commands (eg. npm run lint, npm run typecheck, ruff, etc.) with Bash if they were provided to you to ensure your code is correct. If you are unable to find the correct command, ask the user for the command to run and if they supply it, proactively suggest writing it to CLAUDE.md so that you will know to run it next time. +NEVER commit changes unless the user explicitly asks you to. It is VERY IMPORTANT to only commit when explicitly asked, otherwise the user will feel that you are committing without their permission. + +- Tool results and user messages may include tags. tags contain useful information and reminders. They are NOT part of the user's provided input or the tool result. + +# Tool usage policy +- When doing file search, prefer to use the Task tool in order to reduce context usage. +- You should proactively use the Task tool with specialized agents when the task at hand matches the agent's description. +- When WebFetch returns a message about a redirect to a different host, you should immediately make a new WebFetch request with the redirect URL provided in the response. +- You have the capability to call multiple tools in a single response. When multiple independent pieces of information are requested, batch your tool calls together for optimal performance. When making multiple bash tool calls, you MUST send a single message with multiple tools calls to run the calls in parallel. For example, if you need to run "git status" and "git diff", send a single message with two tool calls to run the calls in parallel. + +You MUST answer concisely with fewer than 4 lines of text (not including tool use or code generation), unless user asks for detail. + +Here is useful information about the environment you are running in: + +Working directory: /workspace/litellm +Is directory a git repo: Yes +Platform: linux +OS Version: Linux 6.8.0 +Today's date: 2026-10-14 + +You are powered by the model named the configured model. The exact model ID is the one the gateway routes this request to. + +IMPORTANT: Assist with defensive security tasks only. Refuse to create, modify, or improve code that may be used maliciously. Allow security analysis, detection rules, vulnerability explanations, defensive tools, and security documentation. + +# Code References +When referencing specific functions or pieces of code include the pattern `file_path:line_number` to allow the user to easily navigate to the source code location. + + +user: Where are errors from the client handled? +assistant: Clients are marked as failed in the `connectToServer` function in src/services/process.ts:712. + + +gitStatus: This is the git status at the start of the conversation. Note that this status is a snapshot in time, and will not update during the conversation. +Current branch: main +Main branch (you will usually use this for PRs): main +Status: (clean) +Recent commits: +abcdef0 fix(router): keep the stream retry inside the same deployment group +1234567 docs: router_settings reference for the new fallback flag +89abcde test(e2e): perf cases on an out-of-process rig diff --git a/tests/e2e/perf/setup/corpus/tools.json b/tests/e2e/perf/setup/corpus/tools.json new file mode 100644 index 00000000000..f6ef78ad32e --- /dev/null +++ b/tests/e2e/perf/setup/corpus/tools.json @@ -0,0 +1,95 @@ +[ + { + "name": "Read", + "description": "Reads a file from the local filesystem. You can access any file directly by using this tool. Assume this tool is able to read all files on the machine. If the user provides a path to a file assume that path is valid. It is okay to read a file that does not exist; an error will be returned. The file_path parameter must be an absolute path, not a relative path. By default, it reads up to 2000 lines starting from the beginning of the file. You can optionally specify a line offset and limit (especially handy for long files), but it is recommended to read the whole file by not providing these parameters. Any lines longer than 2000 characters will be truncated. Results are returned using cat -n format, with line numbers starting at 1. This tool can only read files, not directories. To read a directory, use an ls command via the Bash tool. You have the capability to call multiple tools in a single response. It is always better to speculatively read multiple files as a batch that are potentially useful. If you read a file that exists but has empty contents you will receive a system reminder warning in place of file contents.", + "input_schema": { + "type": "object", + "properties": { + "file_path": {"type": "string", "description": "The absolute path to the file to read"}, + "offset": {"type": "number", "description": "The line number to start reading from. Only provide if the file is too large to read at once"}, + "limit": {"type": "number", "description": "The number of lines to read. Only provide if the file is too large to read at once."} + }, + "required": ["file_path"], + "additionalProperties": false + } + }, + { + "name": "Edit", + "description": "Performs exact string replacements in files. You must use your Read tool at least once in the conversation before editing. This tool will error if you attempt an edit without reading the file. When editing text from Read tool output, ensure you preserve the exact indentation (tabs/spaces) as it appears AFTER the line number prefix. The line number prefix format is: spaces + line number + tab. Everything after that tab is the actual file content to match. Never include any part of the line number prefix in the old_string or new_string. ALWAYS prefer editing existing files in the codebase. NEVER write new files unless explicitly required. Only use emojis if the user explicitly requests it. Avoid adding emojis to files unless asked. The edit will FAIL if old_string is not unique in the file. Either provide a larger string with more surrounding context to make it unique or use replace_all to change every instance of old_string. Use replace_all for replacing and renaming strings across the file. This parameter is useful if you want to rename a variable for instance.", + "input_schema": { + "type": "object", + "properties": { + "file_path": {"type": "string", "description": "The absolute path to the file to modify"}, + "old_string": {"type": "string", "description": "The text to replace"}, + "new_string": {"type": "string", "description": "The text to replace it with (must be different from old_string)"}, + "replace_all": {"type": "boolean", "default": false, "description": "Replace all occurences of old_string (default false)"} + }, + "required": ["file_path", "old_string", "new_string"], + "additionalProperties": false + } + }, + { + "name": "Write", + "description": "Writes a file to the local filesystem. This tool will overwrite the existing file if there is one at the provided path. If this is an existing file, you MUST use the Read tool first to read the file's contents. This tool will fail if you did not read the file first. ALWAYS prefer editing existing files in the codebase. NEVER write new files unless explicitly required. NEVER proactively create documentation files (*.md) or README files. Only create documentation files if explicitly requested by the User. Only use emojis if the user explicitly requests it. Avoid writing emojis to files unless asked.", + "input_schema": { + "type": "object", + "properties": { + "file_path": {"type": "string", "description": "The absolute path to the file to write (must be absolute, not relative)"}, + "content": {"type": "string", "description": "The content to write to the file"} + }, + "required": ["file_path", "content"], + "additionalProperties": false + } + }, + { + "name": "Bash", + "description": "Executes a given bash command in a persistent shell session with optional timeout, ensuring proper handling and security measures. Before executing the command, please follow these steps: 1. Directory Verification: If the command will create new directories or files, first use the LS tool to verify the parent directory exists and is the correct location. 2. Command Execution: Always quote file paths that contain spaces with double quotes. After ensuring proper quoting, execute the command. Capture the output of the command. Usage notes: The command argument is required. You can specify an optional timeout in milliseconds (up to 600000ms / 10 minutes). If not specified, commands will timeout after 120000ms (2 minutes). It is very helpful if you write a clear, concise description of what this command does in 5-10 words. If the output exceeds 30000 characters, output will be truncated before being returned to you. VERY IMPORTANT: You MUST avoid using search commands like find and grep. Instead use Grep, Glob, or Task to search. You MUST avoid read tools like cat, head, tail, and ls, and use Read and LS to read files. If you still need to run grep, STOP. ALWAYS USE ripgrep at rg first, which all Claude Code users have pre-installed. When issuing multiple commands, use the ; or && operator to separate them. DO NOT use newlines to separate commands (newlines are ok in quoted strings). Try to maintain your current working directory throughout the session by using absolute paths and avoiding usage of cd. You may use cd if the User explicitly requests it.", + "input_schema": { + "type": "object", + "properties": { + "command": {"type": "string", "description": "The command to execute"}, + "timeout": {"type": "number", "description": "Optional timeout in milliseconds (max 600000)"}, + "description": {"type": "string", "description": "Clear, concise description of what this command does in 5-10 words."}, + "run_in_background": {"type": "boolean", "description": "Set to true to run this command in the background. Use BashOutput to read the output later."} + }, + "required": ["command"], + "additionalProperties": false + } + }, + { + "name": "Grep", + "description": "A powerful search tool built on ripgrep. ALWAYS use Grep for search tasks. NEVER invoke grep or rg as a Bash command. The Grep tool has been optimized for correct permissions and access. Supports full regex syntax (e.g., \"log.*Error\", \"function\\s+\\w+\"). Filter files with glob parameter (e.g., \"*.js\", \"**/*.tsx\") or type parameter (e.g., \"js\", \"py\", \"rust\"). Output modes: \"content\" shows matching lines, \"files_with_matches\" shows only file paths (default), \"count\" shows match counts. Use Task tool for open-ended searches requiring multiple rounds. Pattern syntax: Uses ripgrep (not grep), so literal braces need escaping (use interface\\{\\} to find interface{} in Go code). Multiline matching: By default patterns match within single lines only. For cross-line patterns like struct \\{[\\s\\S]*?field, use multiline: true.", + "input_schema": { + "type": "object", + "properties": { + "pattern": {"type": "string", "description": "The regular expression pattern to search for in file contents"}, + "path": {"type": "string", "description": "File or directory to search in (rg PATH). Defaults to current working directory."}, + "glob": {"type": "string", "description": "Glob pattern to filter files (e.g. \"*.js\", \"*.{ts,tsx}\") - maps to rg --glob"}, + "output_mode": {"type": "string", "enum": ["content", "files_with_matches", "count"], "description": "Output mode: \"content\" shows matching lines, \"files_with_matches\" shows file paths (supports head_limit), \"count\" shows match counts (supports head_limit). Defaults to \"files_with_matches\"."}, + "-B": {"type": "number", "description": "Number of lines to show before each match (rg -B). Requires output_mode: \"content\", ignored otherwise."}, + "-A": {"type": "number", "description": "Number of lines to show after each match (rg -A). Requires output_mode: \"content\", ignored otherwise."}, + "-C": {"type": "number", "description": "Number of lines to show before and after each match (rg -C). Requires output_mode: \"content\", ignored otherwise."}, + "-n": {"type": "boolean", "description": "Show line numbers in output (rg -n). Requires output_mode: \"content\", ignored otherwise. Defaults to true."}, + "-i": {"type": "boolean", "description": "Case insensitive search (rg -i)"}, + "type": {"type": "string", "description": "File type to search (rg --type). Common types: js, py, rust, go, java, etc. More efficient than include for standard file types."}, + "head_limit": {"type": "number", "description": "Limit output to first N lines/entries, equivalent to \"| head -N\". Works across all output modes: content (limits output lines), files_with_matches (limits file paths), count (limits count entries). Defaults to 250 when unspecified."}, + "multiline": {"type": "boolean", "description": "Enable multiline mode where . matches newlines and patterns can span lines (rg -U --multiline-dotall). Default: false."} + }, + "required": ["pattern"], + "additionalProperties": false + } + }, + { + "name": "Glob", + "description": "Fast file pattern matching tool that works with any codebase size. Supports glob patterns like \"**/*.js\" or \"src/**/*.ts\". Returns matching file paths sorted by modification time. Use this tool when you need to find files by name patterns. When you are doing an open ended search that may require multiple rounds of globbing and grepping, use the Task tool instead. You have the capability to call multiple tools in a single response. It is always better to speculatively perform multiple searches as a batch that are potentially useful.", + "input_schema": { + "type": "object", + "properties": { + "pattern": {"type": "string", "description": "The glob pattern to match files against"}, + "path": {"type": "string", "description": "The directory to search in. If not specified, the current working directory will be used. IMPORTANT: Omit this parameter if you do not provide a value - this is the default."} + }, + "required": ["pattern"], + "additionalProperties": false + } + } +] diff --git a/tests/e2e/perf/setup/datastore.py b/tests/e2e/perf/setup/datastore.py new file mode 100644 index 00000000000..e123d129cc5 --- /dev/null +++ b/tests/e2e/perf/setup/datastore.py @@ -0,0 +1,64 @@ +"""Postgres and Redis work counters, read before and after the traffic phase so a case can put a ceiling +on how much database and cache work the gateway does per request. Both come from the servers' own +statistics (`pg_stat_database`, Redis `INFO`), so nothing here depends on how the gateway talks to them. +`DATABASE_URL` and `REDIS_HOST` are the runner's own env, as in the other suites; without them the +counters are None and any Limit on them reports as missing.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Final + +from pydantic import TypeAdapter + +_DB_QUERY: Final = "SELECT xact_commit FROM pg_stat_database WHERE datname = current_database()" +_DB_ROW: Final = TypeAdapter(tuple[int]) +_INFO: Final = TypeAdapter(dict[str, object]) +_INT: Final = TypeAdapter(int) + + +@dataclass(frozen=True, slots=True) +class DbCounters: + commits: int + + +@dataclass(frozen=True, slots=True) +class RedisCounters: + commands: int + connected_clients: int + + +def read_db_counters() -> DbCounters | None: + url: Final = os.environ.get("DATABASE_URL") + if url is None: + return None + import psycopg + + with psycopg.connect(url) as conn: + row: Final = conn.execute(_DB_QUERY).fetchone() + (commits,) = _DB_ROW.validate_python(row) + return DbCounters(commits=commits) + + +def read_redis_counters() -> RedisCounters | None: + """Standalone Redis only, which is what the perf lane's chart-bundled Redis and a local run are.""" + host: Final = os.environ.get("REDIS_HOST") + if host is None: + return None + import redis + + client: Final = redis.Redis( + host=host, + port=int(os.environ.get("REDIS_PORT", "6379")), + ssl=os.environ.get("REDIS_SSL", "false").lower() == "true", + password=os.environ.get("REDIS_PASSWORD") or None, + decode_responses=True, + socket_connect_timeout=2, + ) + stats: Final = _INFO.validate_python(client.info("stats")) + clients: Final = _INFO.validate_python(client.info("clients")) + return RedisCounters( + commands=_INT.validate_python(stats["total_commands_processed"]), + connected_clients=_INT.validate_python(clients["connected_clients"]), + ) diff --git a/tests/e2e/perf/setup/locustfile.py b/tests/e2e/perf/setup/locustfile.py new file mode 100644 index 00000000000..fe8080959e5 --- /dev/null +++ b/tests/e2e/perf/setup/locustfile.py @@ -0,0 +1,266 @@ +"""The Locust side of a perf case: simulated users take the next request off the pre-generated plan, +send it at the planned rate, and write one record per request with what the client saw: status, +time to headers, time to first byte, time to first token (streams), total time, the text read back +and whether it matched. Runs in its own process (`setup/run.py` starts it), never in the gateway's. + +Reads `PERF_PLAN_FILE` (JSON list of wire requests), writes `PERF_RECORDS_FILE` (one JSON object per +line), paces to `PERF_RPS` over `PERF_USERS` users, and quits once the last request is recorded. +""" + +from __future__ import annotations + +import os +import time +from collections.abc import Iterator +from pathlib import Path +from typing import Final, cast + +import gevent +from locust import FastHttpUser, constant_throughput, task +from locust.contrib.fasthttp import FastResponse, ResponseContextManager +from setup.wire import Endpoint, WireRecord, WireRequest, request_body +from pydantic import TypeAdapter + +_PLAN_ADAPTER: Final = TypeAdapter(list[WireRequest]) +_JSON: Final = TypeAdapter[object](object) +_PLAN: Final = _PLAN_ADAPTER.validate_json(Path(os.environ["PERF_PLAN_FILE"]).read_bytes()) +_RECORDS_PATH: Final = Path(os.environ["PERF_RECORDS_FILE"]) +_RPS: Final = float(os.environ["PERF_RPS"]) +_USERS: Final = int(os.environ["PERF_USERS"]) +_NEXT: Final[Iterator[WireRequest]] = iter(_PLAN) +_RECORDS: Final = _RECORDS_PATH.open("a", encoding="utf-8") +_written = 0 +BLOCK_REASON: Final = "perf-test-block" +_TOKEN_MAX_AGE_S: Final = 240.0 +_tokens: dict[str, tuple[str, float]] = {} + + +def _as_dict(value: object) -> dict[str, object]: + return cast(dict[str, object], value) if isinstance(value, dict) else {} + + +def _as_list(value: object) -> list[object]: + return cast(list[object], value) if isinstance(value, list) else [] + + +def _first_choice_field(payload: dict[str, object], holder: str) -> str: + choices: Final = _as_list(payload.get("choices")) + if not choices: + return "" + content: Final = _as_dict(_as_dict(choices[0]).get(holder)).get("content") + return content if isinstance(content, str) else "" + + +def text_of_json(endpoint: Endpoint, payload: dict[str, object]) -> str: + """The assistant text in a non-streaming body of each endpoint.""" + if endpoint == "chat_completions": + return _first_choice_field(payload, "message") + if endpoint == "responses": + parts: list[str] = [] + for item in _as_list(payload.get("output")): + for part in _as_list(_as_dict(item).get("content")): + text = _as_dict(part).get("text") + if _as_dict(part).get("type") == "output_text" and isinstance(text, str): + parts.append(text) + return "".join(parts) + blocks: list[str] = [] + for block in _as_list(payload.get("content")): + text = _as_dict(block).get("text") + if isinstance(text, str): + blocks.append(text) + return "".join(blocks) + + +def delta_of_event(endpoint: Endpoint, payload: dict[str, object]) -> str: + """The text delta one SSE event of each endpoint carries, empty when it carries none.""" + if endpoint == "chat_completions": + return _first_choice_field(payload, "delta") + if endpoint == "responses": + delta: Final = payload.get("delta") + return delta if payload.get("type") == "response.output_text.delta" and isinstance(delta, str) else "" + if payload.get("type") == "content_block_delta": + text: Final = _as_dict(payload.get("delta")).get("text") + return text if isinstance(text, str) else "" + return "" + + +def _events(buffer: bytes) -> tuple[list[dict[str, object]], bytes]: + """Complete SSE events in `buffer` as their JSON data, plus the unfinished tail.""" + events: list[dict[str, object]] = [] + while b"\n\n" in buffer: + raw, buffer = buffer.split(b"\n\n", 1) + for line in raw.split(b"\n"): + if line.startswith(b"data:"): + data = line[5:].strip() + if data and data != b"[DONE]": + events.append(_as_dict(_JSON.validate_json(data))) + return events, buffer + + +def _header(response: FastResponse, name: str) -> str | None: + headers: Final = response.headers + if headers is None: + return None + value: object = headers.get(name) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] + return value if isinstance(value, str) else None + + +def _quit(user: FastHttpUser) -> None: + """End the run once every record is written. Quitting from a separate greenlet: `runner.quit` + kills the user greenlets, and the user calling it would be killed mid-shutdown.""" + runner = user.environment.runner # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] + if runner is not None: + gevent.spawn(runner.quit) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType, reportUnusedCallResult] + + +def _judge(item: WireRequest, status: int, body_text: str) -> tuple[bool, str | None]: + if item.blocked: + return ( + status == 400 and BLOCK_REASON in body_text, + None if status == 400 else f"expected a 400 block, got {status}", + ) + if status != 200: + return False, f"status {status}: {body_text[:200]}" + return True, None + + +def _record(record: WireRecord) -> None: + global _written + _ = _RECORDS.write(record.model_dump_json() + "\n") + _RECORDS.flush() + _written += 1 + + +def _authorization(user: FastHttpUser, item: WireRequest) -> str: + """The bearer for this request: the key as given, or the caller's JWT, re-minted from Keycloak + once it is older than the realm's access token lifespan allows.""" + if item.token_source is None: + assert item.authorization is not None, f"request {item.seq} has neither a key nor a token source" + return item.authorization + source: Final = item.token_source + cached: Final = _tokens.get(source.username) + if cached is not None and time.monotonic() - cached[1] < _TOKEN_MAX_AGE_S: + return f"Bearer {cached[0]}" + minted = "" + with user.client.request( # pyright: ignore[reportUnknownMemberType] + "POST", + source.token_url, + data={ + "grant_type": "password", + "client_id": source.client_id, + "username": source.username, + "password": source.password, + }, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + catch_response=True, + name="keycloak token", + ) as response: + payload: Final = _as_dict(_JSON.validate_json(response.text or "{}")) + token: Final = payload.get("access_token") + if not isinstance(token, str): + response.failure(f"no access_token in {response.status_code}") # pyright: ignore[reportUnknownMemberType] + raise RuntimeError( + f"could not mint a token for {source.username}: {response.status_code} {response.text!r}" + ) + response.success() + minted = token + _tokens[source.username] = (minted, time.monotonic()) + return f"Bearer {minted}" + + +class PlanUser(FastHttpUser): + """One simulated client; all users share the plan and together hold the planned rate.""" + + wait_time = constant_throughput(_RPS / _USERS) + network_timeout = 120.0 + connection_timeout = 30.0 + + @task + def send(self) -> None: + item: Final = next(_NEXT, None) + if item is None: + if _written >= len(_PLAN): + _quit(self) + return + started_at: Final = time.time() + started: Final = time.perf_counter() + first_byte_s: float | None = None + first_token_s: float | None = None + text = "" + body_text = "" + status = 0 + ok = False + error: str | None = None + header_s = 0.0 + response: FastResponse | None = None + try: + headers: Final = {"Authorization": _authorization(self, item), "Content-Type": "application/json"} + context: ResponseContextManager + with self.client.request( # pyright: ignore[reportUnknownMemberType] + "POST", + item.path, + json=request_body(item), + headers=headers, + stream=item.stream, + catch_response=True, + name=f"{item.endpoint}{' stream' if item.stream else ''}", + ) as context: + response = context + header_s = time.perf_counter() - started + status = int(context.status_code) + if item.stream and status == 200: + buffer = b"" + raw = context.stream # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the socket-backed response + while True: + line = cast(bytes, raw.readline(b"\n")) # pyright: ignore[reportUnknownMemberType] # returns as soon as a line is complete + if not line: + break + now = time.perf_counter() - started + if first_byte_s is None: + first_byte_s = now + buffer += line + events, buffer = _events(buffer) + for event in events: + delta = delta_of_event(item.endpoint, event) + if delta: + if first_token_s is None: + first_token_s = now + text += delta + else: + body_text = context.text or "" + first_byte_s = time.perf_counter() - started + if status == 200: + text = text_of_json(item.endpoint, _as_dict(_JSON.validate_json(body_text))) + ok, error = _judge(item, status, body_text) + if ok: + context.success() + else: + context.failure(error or "failed") # pyright: ignore[reportUnknownMemberType] + except Exception as exc: # noqa: BLE001 # anything the client raised is the record's error + ok = False + error = f"{type(exc).__name__}: {exc}" + total_s: Final = time.perf_counter() - started + delay_header: Final = _header(response, "llm_provider-x-mock-delay-s") if response is not None else None + _record( + WireRecord( + seq=item.seq, + marker=item.marker, + caller=item.caller, + endpoint=item.endpoint, + stream=item.stream, + blocked=item.blocked, + started_at=started_at, + status=status, + header_s=header_s, + first_byte_s=first_byte_s, + first_token_s=first_token_s, + total_s=total_s, + text=text, + ok=ok, + mock_delay_s=float(delay_header) if delay_header else None, + call_id=_header(response, "x-litellm-call-id") if response is not None else None, + error=error, + ) + ) + if _written >= len(_PLAN): + _quit(self) diff --git a/tests/e2e/perf/setup/mocks/__init__.py b/tests/e2e/perf/setup/mocks/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/e2e/perf/setup/mocks/mock_anthropic.py b/tests/e2e/perf/setup/mocks/mock_anthropic.py new file mode 100644 index 00000000000..c7a767d023d --- /dev/null +++ b/tests/e2e/perf/setup/mocks/mock_anthropic.py @@ -0,0 +1,40 @@ +"""The Anthropic-shaped perf mock: /v1/messages, JSON and SSE, answering each request with a generated reply in +api.anthropic.com's wire format (thinking blocks, tool_use, cache usage, the message_start to message_stop event +order and its chunk sizes).""" + +from __future__ import annotations + +from typing import Final + +from fastapi import FastAPI, Request +from fastapi.responses import Response + +from mock_common import ( + RequestLog, + bool_field, + json_reply, + mount_control, + parse_json, + str_field, + stream_reply, +) + +SSE: Final = "text/event-stream" + + +def build_app() -> FastAPI: + app: Final = FastAPI(title="perf mock anthropic") + log: Final = RequestLog("anthropic") + mount_control(app, log) + + async def messages(request: Request) -> Response: + body: Final = parse_json(await request.body()) + stream: Final = bool_field(body, "stream") + decision: Final = log.decide( + route="messages", stream=stream, body=body, model=str_field(body, "model", "claude-sonnet-5") + ) + return stream_reply(decision, SSE) if stream else await json_reply(decision) + + app.add_api_route("/v1/messages", messages, methods=["POST"]) + app.add_api_route("/messages", messages, methods=["POST"]) + return app diff --git a/tests/e2e/perf/setup/mocks/mock_bedrock.py b/tests/e2e/perf/setup/mocks/mock_bedrock.py new file mode 100644 index 00000000000..bbf145d90b4 --- /dev/null +++ b/tests/e2e/perf/setup/mocks/mock_bedrock.py @@ -0,0 +1,51 @@ +"""The Bedrock-shaped perf mock: Converse, ConverseStream, InvokeModel and InvokeModelWithResponseStream under +/model/{model_id}/..., answering each request with a generated reply in bedrock-runtime's wire format. Streams are application/vnd.amazon.eventstream +frames (prelude, headers, CRCs), so the gateway's botocore-based decoder runs against the real framing.""" + +from __future__ import annotations + +from typing import Final, Literal + +from fastapi import FastAPI, Request +from fastapi.responses import Response + +from mock_common import ( + RequestLog, + json_reply, + mount_control, + parse_json, + stream_reply, +) + +EVENT_STREAM: Final = "application/vnd.amazon.eventstream" + + +def build_app() -> FastAPI: + app: Final = FastAPI(title="perf mock bedrock") + log: Final = RequestLog("bedrock") + mount_control(app, log) + + async def handle(model_id: str, request: Request, route: Literal["converse", "invoke"], stream: bool) -> Response: + body: Final = parse_json(await request.body()) + decision: Final = log.decide(route=route, stream=stream, body=body, model=model_id) + return stream_reply(decision, EVENT_STREAM) if stream else await json_reply(decision) + + async def converse(model_id: str, request: Request) -> Response: + return await handle(model_id, request, "converse", False) + + async def converse_stream(model_id: str, request: Request) -> Response: + return await handle(model_id, request, "converse", True) + + async def invoke(model_id: str, request: Request) -> Response: + return await handle(model_id, request, "invoke", False) + + async def invoke_with_response_stream(model_id: str, request: Request) -> Response: + return await handle(model_id, request, "invoke", True) + + app.add_api_route("/model/{model_id:path}/converse", converse, methods=["POST"]) + app.add_api_route("/model/{model_id:path}/converse-stream", converse_stream, methods=["POST"]) + app.add_api_route("/model/{model_id:path}/invoke", invoke, methods=["POST"]) + app.add_api_route( + "/model/{model_id:path}/invoke-with-response-stream", invoke_with_response_stream, methods=["POST"] + ) + return app diff --git a/tests/e2e/perf/setup/mocks/mock_common.py b/tests/e2e/perf/setup/mocks/mock_common.py new file mode 100644 index 00000000000..3599702bb05 --- /dev/null +++ b/tests/e2e/perf/setup/mocks/mock_common.py @@ -0,0 +1,227 @@ +"""Shared behaviour of the perf mock providers. + +Every reply is generated for the request just received (see `replies.py`): sized from its prompt and output +cap, a tool call when tools were offered, in the provider's own wire format with its usage fields, ids, +event order and chunk sizes, so the gateway parses, prices and streams it as it would the real thing. The +model id stays a real one (`gpt-6-luna`, `anthropic.claude-sonnet-5`) so every name-based branch in the +gateway runs, and one mock pod per provider serves any number of deployments. How long a mock waits before answering is the one setting a case +controls, `POST /mock/latency {"seconds": 1.0, "jitter_seconds": 0.25}`, set by the test's setup directly +on the pod, never through the gateway. The wait, the reply's size and its ids are pure functions of the request marker (`perfreq-000123-n`, put in the prompt by the plan +generator), so the test can join what each provider received against what each client saw. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import json +import re +from collections import deque +from dataclasses import dataclass +from typing import Final, Literal, cast + +from fastapi import FastAPI +from fastapi.responses import Response, StreamingResponse +from pydantic import BaseModel, ConfigDict, Field + +from replies import Route, render_chunks, render_json, reply_for +from reply_common import Provider, Reply, chunks_after, unit + +MARKER: Final = re.compile(r"perfreq-\d{6}-[nib]") +DELAY_HEADER: Final = "x-mock-delay-s" +Flag = Literal["n", "i", "b"] + + +def marker_in(text: str) -> str | None: + match: Final = MARKER.search(text) + return None if match is None else match.group(0) + + +def flag_of(marker: str) -> Flag: + last: Final = marker[-1] + assert last in ("n", "i", "b"), marker + return "n" if last == "n" else "i" if last == "i" else "b" + + +def collect_text(value: object) -> list[str]: + """Every prompt string, whatever the wire format nests it in: `text` and `content` parts, the + Responses API's bare `input`, and `system` prompts.""" + if isinstance(value, dict): + items: Final = cast(dict[object, object], value).items() + return [ + part + for key, item in items + for part in ( + [item] + if key in ("text", "content", "input", "system") and isinstance(item, str) + else collect_text(item) + ) + ] # comprehension-ok: flattening one nesting level of the request body + if isinstance(value, list): + return [part for item in cast(list[object], value) for part in collect_text(item)] # comprehension-ok: flatten + return [] + + +def token_count(text: str) -> int: + return max(1, len(text) // 4) + + +def str_field(body: object, key: str, default: str) -> str: + if isinstance(body, dict): + value: object = cast(dict[object, object], body).get(key) + if isinstance(value, str): + return value + return default + + +def bool_field(body: object, key: str) -> bool: + if isinstance(body, dict): + value: object = cast(dict[object, object], body).get(key) + return value is True + return False + + +class MockRequest(BaseModel): + model_config = ConfigDict(frozen=True) + + provider: Provider + marker: str | None + route: str + stream: bool + output_tokens: int + text_sha1: str + delay_s: float + + +class MockRequestsResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + requests: list[MockRequest] + + +class ResetResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + cleared: int + + +class Latency(BaseModel): + """Seconds a mock waits before its first byte: `seconds` plus or minus up to `jitter_seconds`, drawn + from the request marker so the same request always waits the same.""" + + model_config = ConfigDict(frozen=True) + + seconds: float = Field(default=0.0, ge=0.0) + jitter_seconds: float = Field(default=0.0, ge=0.0) + + def for_seed(self, seed: str) -> float: + return max(0.0, self.seconds + self.jitter_seconds * (2 * unit(f"{seed}:latency") - 1)) + + +@dataclass(frozen=True, slots=True) +class Decision: + """What this mock does with one request: the reply it generates, on which route, after how long.""" + + marker: str | None + reply: Reply + route: Route + delay_s: float + """Seconds before the first byte, what `gateway_added_seconds` subtracts from the client's time.""" + + @property + def seed(self) -> str: + return self.marker or "unmarked" + + @property + def headers(self) -> dict[str, str]: + return {DELAY_HEADER: f"{self.delay_s:.4f}"} + + +class RequestLog: + """What one mock pod received, in order, readable by the test at GET /mock/requests.""" + + def __init__(self, provider: Provider, capacity: int = 100_000) -> None: + self._provider: Final = provider + self._entries: deque[MockRequest] = deque(maxlen=capacity) + self._latency = Latency() + + @property + def latency(self) -> Latency: + return self._latency + + def set_latency(self, latency: Latency) -> Latency: + self._latency = latency + return self._latency + + def decide(self, *, route: Route, stream: bool, body: object, model: str) -> Decision: + prompt_text: Final = " ".join(collect_text(body)) + marker: Final = marker_in(prompt_text) + seed: Final = marker or "unmarked" + reply: Final = reply_for( + provider=self._provider, body=body, seed=seed, model=model, prompt_tokens=token_count(prompt_text) + ) + decision: Final = Decision(marker=marker, reply=reply, route=route, delay_s=self._latency.for_seed(seed)) + self._entries.append( + MockRequest( + provider=self._provider, + marker=marker, + route=route, + stream=stream, + output_tokens=reply.output_tokens, + text_sha1=hashlib.sha1(reply.text.encode()).hexdigest(), + delay_s=decision.delay_s, + ) + ) + return decision + + def snapshot(self) -> list[MockRequest]: + return list(self._entries) + + def reset(self) -> int: + count: Final = len(self._entries) + self._entries.clear() + return count + + +async def json_reply(decision: Decision) -> Response: + """The generated body after the case's latency.""" + if decision.delay_s > 0: + await asyncio.sleep(decision.delay_s) + return Response( + content=render_json(decision.reply, decision.route), + media_type="application/json", + headers=decision.headers, + ) + + +def stream_reply(decision: Decision, media_type: str) -> Response: + """The generated chunks back to back after the case's latency.""" + return StreamingResponse( + chunks_after(render_chunks(decision.reply, decision.route), decision.delay_s), + media_type=media_type, + headers=decision.headers, + ) + + +def parse_json(raw: bytes) -> object: + parsed: object = json.loads(raw) if raw else {} + return parsed + + +def mount_control(app: FastAPI, log: RequestLog) -> None: + async def health() -> dict[str, str]: + return {"status": "ok"} + + async def requests_seen() -> MockRequestsResponse: + return MockRequestsResponse(requests=log.snapshot()) + + async def reset() -> ResetResponse: + return ResetResponse(cleared=log.reset()) + + async def latency(request: Latency) -> Latency: + return log.set_latency(request) + + app.add_api_route("/health", health, methods=["GET"]) + app.add_api_route("/mock/requests", requests_seen, methods=["GET"]) + app.add_api_route("/mock/reset", reset, methods=["POST"]) + app.add_api_route("/mock/latency", latency, methods=["POST"]) diff --git a/tests/e2e/perf/setup/mocks/mock_guardrail.py b/tests/e2e/perf/setup/mocks/mock_guardrail.py new file mode 100644 index 00000000000..7b6f43285a2 --- /dev/null +++ b/tests/e2e/perf/setup/mocks/mock_guardrail.py @@ -0,0 +1,115 @@ +"""The perf mock guardrail: a server speaking litellm's Generic Guardrail API contract +(POST /beta/litellm_basic_guardrail_api), registered once as pre_call and once as post_call. +What it does with a request is decided by the marker's flag: `-n` passes it through +(`action: NONE`), `-i` rewrites the text (`[pre]` on the request, `[post]` on the response) so the +test can see the rewrite reach the provider and the client, and `-b` blocks the request with the +reason `perf-test-block`. Its own latency comes from `additional_provider_specific_params`.""" + +from __future__ import annotations + +import asyncio +from collections import deque +from collections.abc import Mapping +from typing import Final, Literal + +from fastapi import FastAPI +from fastapi.responses import JSONResponse +from mock_common import flag_of, marker_in +from pydantic import BaseModel, ConfigDict + +BLOCK_REASON: Final = "perf-test-block" +Action = Literal["NONE", "BLOCKED", "GUARDRAIL_INTERVENED"] + + +class GuardrailRequest(BaseModel): + model_config = ConfigDict(extra="allow") + input_type: Literal["request", "response"] + texts: list[str] | None = None + litellm_call_id: str | None = None + additional_provider_specific_params: dict[str, object] | None = None + + def latency_s(self) -> float: + value: Final = (self.additional_provider_specific_params or {}).get("latency_s") + return float(value) if isinstance(value, (int, float)) else 0.0 + + +class GuardrailCall(BaseModel): + marker: str | None + input_type: Literal["request", "response"] + action: Action + + +class GuardrailCallsResponse(BaseModel): + calls: list[GuardrailCall] + + +class GuardrailResetResponse(BaseModel): + cleared: int + + +def marker_of(body: GuardrailRequest, remembered: Mapping[str, str]) -> str | None: + """The request's marker: in the prompt on the request hook, remembered by `litellm_call_id` on the response + hook, since a generated provider reply does not contain it.""" + texts: Final = body.texts or [] + in_text: Final = next((found for text in texts if (found := marker_in(text)) is not None), None) + if in_text is not None or body.litellm_call_id is None: + return in_text + return remembered.get(body.litellm_call_id) + + +def decide(body: GuardrailRequest, marker: str | None) -> tuple[Action, dict[str, object]]: + texts: Final = body.texts or [] + if marker is None: + return "NONE", {"action": "NONE"} + flag: Final = flag_of(marker) + if flag == "b" and body.input_type == "request": + return "BLOCKED", {"action": "BLOCKED", "blocked_reason": BLOCK_REASON} + if flag == "i": + suffix: Final = " [pre]" if body.input_type == "request" else " [post]" + rewritten: Final = ( + [text + suffix if marker_in(text) else text for text in texts] + if body.input_type == "request" + else [text + suffix for text in texts] + ) + return "GUARDRAIL_INTERVENED", {"action": "GUARDRAIL_INTERVENED", "texts": rewritten} + return "NONE", {"action": "NONE"} + + +def build_app() -> FastAPI: + app: Final = FastAPI(title="perf mock guardrail") + calls: deque[GuardrailCall] = deque(maxlen=200_000) + markers: dict[str, str] = {} # mutable-ok: the request hook's marker, looked up by the response hook + + async def health() -> dict[str, str]: + return {"status": "ok"} + + async def guard(body: GuardrailRequest) -> JSONResponse: + latency: Final = body.latency_s() + if latency > 0: + await asyncio.sleep(latency) + marker: Final = marker_of(body, markers) + if marker is not None and body.litellm_call_id is not None and body.input_type == "request": + markers[body.litellm_call_id] = marker + action, content = decide(body, marker) + calls.append( + GuardrailCall( + marker=marker, + input_type=body.input_type, + action=action, + ) + ) + return JSONResponse(content=content) + + async def calls_seen() -> GuardrailCallsResponse: + return GuardrailCallsResponse(calls=list(calls)) + + async def reset() -> GuardrailResetResponse: + count: Final = len(calls) + calls.clear() + return GuardrailResetResponse(cleared=count) + + app.add_api_route("/health", health, methods=["GET"]) + app.add_api_route("/beta/litellm_basic_guardrail_api", guard, methods=["POST"]) + app.add_api_route("/mock/calls", calls_seen, methods=["GET"]) + app.add_api_route("/mock/reset", reset, methods=["POST"]) + return app diff --git a/tests/e2e/perf/setup/mocks/mock_openai.py b/tests/e2e/perf/setup/mocks/mock_openai.py new file mode 100644 index 00000000000..ccbb3d46b3a --- /dev/null +++ b/tests/e2e/perf/setup/mocks/mock_openai.py @@ -0,0 +1,48 @@ +"""The OpenAI-shaped perf mock: /v1/chat/completions and /v1/responses, JSON and SSE, answering each request +with a generated reply in api.openai.com's wire format, so the gateway's real OpenAI client and stream parsers +run against the provider's field set and chunk sizes.""" + +from __future__ import annotations + +from typing import Final, Literal + +from fastapi import FastAPI, Request +from fastapi.responses import Response + +from mock_common import ( + RequestLog, + bool_field, + json_reply, + mount_control, + parse_json, + str_field, + stream_reply, +) + +SSE: Final = "text/event-stream" + + +def build_app() -> FastAPI: + app: Final = FastAPI(title="perf mock openai") + log: Final = RequestLog("openai") + mount_control(app, log) + + async def handle(request: Request, route: Literal["chat", "responses"]) -> Response: + body: Final = parse_json(await request.body()) + stream: Final = bool_field(body, "stream") + decision: Final = log.decide( + route=route, stream=stream, body=body, model=str_field(body, "model", "gpt-6-luna") + ) + return stream_reply(decision, SSE) if stream else await json_reply(decision) + + async def chat_completions(request: Request) -> Response: + return await handle(request, "chat") + + async def responses(request: Request) -> Response: + return await handle(request, "responses") + + app.add_api_route("/v1/chat/completions", chat_completions, methods=["POST"]) + app.add_api_route("/chat/completions", chat_completions, methods=["POST"]) + app.add_api_route("/v1/responses", responses, methods=["POST"]) + app.add_api_route("/responses", responses, methods=["POST"]) + return app diff --git a/tests/e2e/perf/setup/mocks/replies.py b/tests/e2e/perf/setup/mocks/replies.py new file mode 100644 index 00000000000..e472ba5b057 --- /dev/null +++ b/tests/e2e/perf/setup/mocks/replies.py @@ -0,0 +1,80 @@ +"""A generated provider reply for one request: `plan_reply` reads the request (model, prompt size, output +cap, tools offered) and `render_json` / `render_chunks` write it in the wire format of the provider and +route. The words are synthetic; the schema, ids, event order, usage fields and delta and chunk sizes follow +what each provider sends, so the gateway's parsers, cost tracking and stream handling do the same work.""" + +from __future__ import annotations + +from typing import Final, Literal, cast + +import reply_anthropic +import reply_bedrock +import reply_openai +from reply_common import Provider, Reply, plan_reply + +Route = Literal["chat", "responses", "messages", "converse", "invoke"] +_MAX_TOKEN_KEYS: Final = ("max_output_tokens", "max_completion_tokens", "max_tokens") +DEFAULT_MAX_TOKENS: Final = 1024 + + +def _mapping(value: object) -> dict[str, object]: + return cast(dict[str, object], value) if isinstance(value, dict) else {} + + +def _items(value: object) -> tuple[object, ...]: + return tuple(cast(list[object], value)) if isinstance(value, list) else () + + +def max_tokens_of(body: object) -> int: + mapping: Final = _mapping(body) + found: Final = tuple(mapping[key] for key in _MAX_TOKEN_KEYS if isinstance(mapping.get(key), int)) + inference: Final = _mapping(mapping.get("inferenceConfig")).get("maxTokens") + if found: + return cast(int, found[0]) + return inference if isinstance(inference, int) else DEFAULT_MAX_TOKENS + + +def _tool_name(tool: object) -> str | None: + mapping: Final = _mapping(tool) + nested: Final = _mapping(mapping.get("function")) or _mapping(mapping.get("toolSpec")) + name: Final = nested.get("name", mapping.get("name")) + return name if isinstance(name, str) else None + + +def tool_names_of(body: object) -> tuple[str, ...]: + mapping: Final = _mapping(body) + tools: Final = _items(mapping.get("tools")) or _items(_mapping(mapping.get("toolConfig")).get("tools")) + return tuple(name for name in map(_tool_name, tools) if name is not None) + + +def reply_for(*, provider: Provider, body: object, seed: str, model: str, prompt_tokens: int) -> Reply: + return plan_reply( + provider=provider, + seed=seed, + model=model, + prompt_tokens=prompt_tokens, + max_tokens=max_tokens_of(body), + tool_names=tool_names_of(body), + ) + + +def render_json(reply: Reply, route: Route) -> bytes: + if route == "chat": + return reply_openai.chat_json(reply) + if route == "responses": + return reply_openai.responses_json(reply) + if route in ("messages", "invoke"): + return reply_anthropic.messages_json(reply) + return reply_bedrock.converse_json(reply) + + +def render_chunks(reply: Reply, route: Route) -> tuple[bytes, ...]: + if route == "chat": + return reply_openai.chat_chunks(reply) + if route == "responses": + return reply_openai.responses_chunks(reply) + if route == "messages": + return reply_anthropic.messages_chunks(reply) + if route == "invoke": + return reply_bedrock.invoke_chunks(reply) + return reply_bedrock.converse_chunks(reply) diff --git a/tests/e2e/perf/setup/mocks/reply_anthropic.py b/tests/e2e/perf/setup/mocks/reply_anthropic.py new file mode 100644 index 00000000000..b44e9a2c514 --- /dev/null +++ b/tests/e2e/perf/setup/mocks/reply_anthropic.py @@ -0,0 +1,132 @@ +"""Anthropic wire format for a generated `Reply`: `/v1/messages` JSON and the SSE event sequence +(`message_start`, `ping`, content blocks of `thinking`, `text` or `tool_use` with their deltas, +`message_delta`, `message_stop`) with the usage fields, delta and chunk sizes api.anthropic.com uses.""" + +from __future__ import annotations + +import json +from typing import Final + +from reply_common import Reply, grouped, ident, pieces, sse + +DELTA_CHARS: Final = 40 +CHUNK_BYTES: Final = 1000 + + +def _usage(reply: Reply) -> dict[str, object]: + return { + "input_tokens": reply.uncached_tokens, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": reply.cached_tokens, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": reply.output_tokens, + "service_tier": "standard", + } + + +def _signature(reply: Reply) -> str: + return ident(reply.seed, "", 64, hex_digits=True) + + +def _blocks(reply: Reply) -> tuple[dict[str, object], ...]: + thinking: Final[tuple[dict[str, object], ...]] = ( + ({"type": "thinking", "thinking": reply.reasoning, "signature": _signature(reply)},) + if reply.reasoning_tokens + else () + ) + if reply.tool_call is None: + return (*thinking, _block(type="text", text=reply.text)) + tool: Final = _block( + type="tool_use", + id=ident(reply.seed, "toolu_", 24), + name=reply.tool_call.name, + input=reply.tool_call.parsed_arguments, + ) + return (*thinking, tool) + + +def _block(**fields: object) -> dict[str, object]: + return dict(fields) + + +def _stop_reason(reply: Reply) -> str: + return "end_turn" if reply.tool_call is None else "tool_use" + + +def messages_json(reply: Reply) -> bytes: + body: Final = { + "model": reply.model, + "id": ident(reply.seed, "msg_", 24), + "type": "message", + "role": "assistant", + "content": list(_blocks(reply)), + "stop_reason": _stop_reason(reply), + "stop_sequence": None, + "usage": _usage(reply), + } + return json.dumps(body, separators=(",", ":")).encode() + + +def _block_events( + index: int, start: dict[str, object], deltas: tuple[dict[str, object], ...] +) -> tuple[dict[str, object], ...]: + return ( + {"type": "content_block_start", "index": index, "content_block": start}, + *({"type": "content_block_delta", "index": index, "delta": delta} for delta in deltas), + {"type": "content_block_stop", "index": index}, + ) + + +def messages_events(reply: Reply) -> tuple[dict[str, object], ...]: + """The stream as event payloads, `message_start` to `message_stop`, each one's `type` its SSE event name.""" + opening: Final[dict[str, object]] = { + "model": reply.model, + "id": ident(reply.seed, "msg_", 24), + "type": "message", + "role": "assistant", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {**_usage(reply), "output_tokens": 1}, + } + thinking: Final = ( + _block_events( + 0, + {"type": "thinking", "thinking": "", "signature": ""}, + ( + *({"type": "thinking_delta", "thinking": piece} for piece in pieces(reply.reasoning, DELTA_CHARS)), + {"type": "signature_delta", "signature": _signature(reply)}, + ), + ) + if reply.reasoning_tokens + else () + ) + index: Final = 1 if reply.reasoning_tokens else 0 + content: Final = ( + _block_events( + index, + {"type": "text", "text": ""}, + tuple({"type": "text_delta", "text": p} for p in pieces(reply.text, DELTA_CHARS)), + ) + if reply.tool_call is None + else _block_events( + index, + {"type": "tool_use", "id": ident(reply.seed, "toolu_", 24), "name": reply.tool_call.name, "input": {}}, + tuple( + {"type": "input_json_delta", "partial_json": p} for p in pieces(reply.tool_call.arguments, DELTA_CHARS) + ), + ) + ) + closing: Final[tuple[dict[str, object], ...]] = ( + { + "type": "message_delta", + "delta": {"stop_reason": _stop_reason(reply), "stop_sequence": None}, + "usage": {**_usage(reply), "output_tokens": reply.output_tokens}, + }, + {"type": "message_stop"}, + ) + return ({"type": "message_start", "message": opening}, {"type": "ping"}, *thinking, *content, *closing) + + +def messages_chunks(reply: Reply) -> tuple[bytes, ...]: + return grouped(tuple(sse(event, str(event["type"])) for event in messages_events(reply)), CHUNK_BYTES) diff --git a/tests/e2e/perf/setup/mocks/reply_bedrock.py b/tests/e2e/perf/setup/mocks/reply_bedrock.py new file mode 100644 index 00000000000..c9a995ebe51 --- /dev/null +++ b/tests/e2e/perf/setup/mocks/reply_bedrock.py @@ -0,0 +1,174 @@ +"""Bedrock wire formats for a generated `Reply`: the `converse` JSON body and `converse-stream` frames, and +for Anthropic models the `invoke` body (Anthropic's own messages JSON) and `invoke-with-response-stream` frames +(each a `chunk` event carrying one base64 Anthropic stream event). Frames are `application/vnd.amazon.eventstream` +(prelude, string headers, JSON payload, CRC32s), what the gateway's botocore-based decoder reads, one per chunk.""" + +from __future__ import annotations + +import base64 +import json +import struct +import zlib +from collections.abc import Callable +from typing import Final + +from reply_anthropic import messages_events +from reply_common import Reply, ident, pick, pieces + +DELTA_CHARS: Final = 15 +_PADDING: Final = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123" + + +def _usage(reply: Reply) -> dict[str, object]: + return { + "cacheReadInputTokenCount": reply.cached_tokens, + "cacheReadInputTokens": reply.cached_tokens, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": reply.uncached_tokens, + "outputTokens": reply.output_tokens, + "serverToolUsage": {}, + "totalTokens": reply.total_tokens, + } + + +def _tool_use_id(reply: Reply) -> str: + return ident(reply.seed, "tooluse_", 22) + + +def _stop_reason(reply: Reply) -> str: + return "end_turn" if reply.tool_call is None else "tool_use" + + +def _text_block(text: str) -> dict[str, object]: + return {"text": text} + + +def _content(reply: Reply) -> tuple[dict[str, object], ...]: + reasoning: Final[tuple[dict[str, object], ...]] = ( + ( + { + "reasoningContent": { + "reasoningText": {"text": reply.reasoning, "signature": ident(reply.seed, "", 64, hex_digits=True)} + } + }, + ) + if reply.reasoning_tokens + else () + ) + if reply.tool_call is None: + return (*reasoning, _text_block(reply.text)) + tool: Final[dict[str, object]] = { + "toolUse": { + "toolUseId": _tool_use_id(reply), + "name": reply.tool_call.name, + "input": reply.tool_call.parsed_arguments, + } + } + return (*reasoning, tool) + + +def converse_json(reply: Reply) -> bytes: + body: Final = { + "metrics": {"latencyMs": 40 * reply.output_tokens}, + "output": {"message": {"role": "assistant", "content": list(_content(reply))}}, + "stopReason": _stop_reason(reply), + "usage": _usage(reply), + } + return json.dumps(body, separators=(",", ":")).encode() + + +def _header(name: str, value: str) -> bytes: + encoded: Final = value.encode() + return struct.pack("!B", len(name)) + name.encode() + struct.pack("!BH", 7, len(encoded)) + encoded + + +def frame(event_type: str, payload: dict[str, object]) -> bytes: + """One eventstream message: prelude (total length, headers length, CRC), the three string headers + bedrock-runtime sets, the JSON payload, and the CRC of everything before it.""" + headers: Final = ( + _header(":event-type", event_type) + + _header(":content-type", "application/json") + + _header(":message-type", "event") + ) + body: Final = json.dumps(payload, separators=(",", ":")).encode() + total: Final = 4 + 4 + 4 + len(headers) + len(body) + 4 + prelude: Final = struct.pack("!II", total, len(headers)) + head: Final = prelude + struct.pack("!I", zlib.crc32(prelude)) + message: Final = head + headers + body + return message + struct.pack("!I", zlib.crc32(message)) + + +def _pad(seed: str, index: int) -> str: + return _PADDING[: 2 + pick(f"{seed}:pad", index, 50)] + + +def _text_events(reply: Reply, index: int, event: Callable[[int, str, dict[str, object]], bytes]) -> tuple[bytes, ...]: + return ( + *( + event(i, "contentBlockDelta", {"contentBlockIndex": index, "delta": {"text": piece}}) + for i, piece in enumerate(pieces(reply.text, DELTA_CHARS), start=1) + ), + event(0, "contentBlockStop", {"contentBlockIndex": index}), + ) + + +def _tool_events(reply: Reply, index: int, event: Callable[[int, str, dict[str, object]], bytes]) -> tuple[bytes, ...]: + assert reply.tool_call is not None + start: Final = {"toolUse": {"name": reply.tool_call.name, "toolUseId": _tool_use_id(reply)}} + return ( + event(0, "contentBlockStart", {"contentBlockIndex": index, "start": start}), + *( + event(i, "contentBlockDelta", {"contentBlockIndex": index, "delta": {"toolUse": {"input": piece}}}) + for i, piece in enumerate(pieces(reply.tool_call.arguments, DELTA_CHARS), start=1) + ), + event(0, "contentBlockStop", {"contentBlockIndex": index}), + ) + + +def converse_chunks(reply: Reply) -> tuple[bytes, ...]: + def event(index: int, event_type: str, payload: dict[str, object]) -> bytes: + return frame(event_type, {**payload, "p": _pad(reply.seed, index)}) + + reasoning: Final = ( + ( + *( + event(i, "contentBlockDelta", {"contentBlockIndex": 0, "delta": {"reasoningContent": {"text": piece}}}) + for i, piece in enumerate(pieces(reply.reasoning, DELTA_CHARS), start=1) + ), + event( + 0, + "contentBlockDelta", + { + "contentBlockIndex": 0, + "delta": {"reasoningContent": {"signature": ident(reply.seed, "", 64, hex_digits=True)}}, + }, + ), + event(0, "contentBlockStop", {"contentBlockIndex": 0}), + ) + if reply.reasoning_tokens + else () + ) + index: Final = 1 if reply.reasoning_tokens else 0 + content: Final = _text_events(reply, index, event) if reply.tool_call is None else _tool_events(reply, index, event) + return ( + event(0, "messageStart", {"role": "assistant"}), + *reasoning, + *content, + event(0, "messageStop", {"stopReason": _stop_reason(reply)}), + event(0, "metadata", {"metrics": {"latencyMs": 40 * reply.output_tokens}, "usage": _usage(reply)}), + ) + + +def invoke_chunks(reply: Reply) -> tuple[bytes, ...]: + """The Anthropic stream events, each base64 in the `bytes` of one `chunk` frame.""" + return tuple( + frame( + "chunk", + { + "bytes": base64.b64encode(json.dumps(event, separators=(",", ":")).encode()).decode(), + "p": _pad(reply.seed, index), + }, + ) + for index, event in enumerate(messages_events(reply)) + ) diff --git a/tests/e2e/perf/setup/mocks/reply_common.py b/tests/e2e/perf/setup/mocks/reply_common.py new file mode 100644 index 00000000000..fa4daa039cf --- /dev/null +++ b/tests/e2e/perf/setup/mocks/reply_common.py @@ -0,0 +1,182 @@ +"""What every generated provider reply is made of: a `Reply` sized from the request (prompt tokens, the +output cap, whether tools were offered), deterministic ids and text drawn from the request marker, and the +chunking that makes a stream arrive the way the provider's TCP writes do. The provider modules render a +`Reply` in their own wire format; nothing here is provider specific.""" + +from __future__ import annotations + +import asyncio +import hashlib +import json +from collections.abc import AsyncIterator, Sequence +from dataclasses import dataclass +from itertools import accumulate, groupby +from typing import Final, Literal + +from pydantic import TypeAdapter + +Provider = Literal["openai", "anthropic", "bedrock"] +CHARS_PER_TOKEN: Final = 4 +CACHED_PREFIX_TOKENS: Final = 1024 +"""Prompt tokens the provider reports as a cache hit (system prompt and tool set) on the requests that hit.""" +CACHE_HIT_SHARE: Final = 0.7 +REASONING_SHARE: Final = 0.6 +"""Share of the output spent on reasoning when the request offered tools, as the captured replies did.""" +_BASE62: Final = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789" +_DIRS: Final = ( + "management", + "guardrails", + "llm_translation", + "router", + "quota_management", + "mcp", + "logging", + "security", +) +_WORDS: Final = ( + "spend", "budget", "team", "key", "fallback", "cooldown", "retry", "stream", "cache", "tag", + "alias", "quota", "webhook", "secret", "model", "org", "audit", "rerank", "batch", "realtime", +) # fmt: skip +_THOUGHT: Final = "Checking {path}: the except clause around the client call logs and returns None, so the caller never sees the error. " + + +def unit(seed: str) -> float: + """A float in [0, 1) that is the same for the same seed on every mock pod.""" + digest: Final = hashlib.sha1(seed.encode()).digest() + return int.from_bytes(digest[:8], "big") / 2**64 + + +def pick(seed: str, index: int, size: int) -> int: + return int(hashlib.sha1(f"{seed}:{index}".encode()).hexdigest()[:8], 16) % size + + +def ident(seed: str, prefix: str, length: int, *, hex_digits: bool = False) -> str: + """A provider-shaped id (`chatcmpl-...`, `msg_...`, `toolu_...`) that is a pure function of the seed.""" + digest: Final = hashlib.sha256(f"{seed}:{prefix}".encode()) + if hex_digits: + return prefix + digest.hexdigest()[:length] + return prefix + "".join(_BASE62[byte % len(_BASE62)] for byte in digest.digest()[:length]) + + +def file_paths(seed: str, count: int) -> tuple[str, ...]: + return tuple( + f"tests/e2e/{_DIRS[pick(seed, i, len(_DIRS))]}/test_{_WORDS[pick(seed, i + 1000, len(_WORDS))]}" + f"_{_WORDS[pick(seed, i + 2000, len(_WORDS))]}_e2e.py" + for i in range(count) + ) + + +def answer_text(seed: str, tokens: int) -> str: + """The assistant's answer to the corpus question (which files swallow exceptions): paths, one per line.""" + return "\n".join(file_paths(seed, max(1, tokens * CHARS_PER_TOKEN // 56))) + + +def thinking_text(seed: str, tokens: int) -> str: + if tokens == 0: + return "" + paths: Final = file_paths(f"{seed}:thinking", max(1, tokens * CHARS_PER_TOKEN // len(_THOUGHT))) + return "".join(_THOUGHT.format(path=path) for path in paths).rstrip() + + +def grep_arguments(seed: str, tokens: int) -> str: + """A `Grep` call over the files the session read, its glob sized to the output budget.""" + paths: Final = file_paths(f"{seed}:grep", max(1, tokens * CHARS_PER_TOKEN // 56)) + arguments: Final = { + "pattern": r"\braise\s+\w+", + "path": "/workspace/litellm/tests/e2e", + "glob": "{" + ",".join(paths) + "}", + } + return json.dumps(arguments, separators=(",", ":")) + + +_ARGUMENTS: Final = TypeAdapter(dict[str, object]) + + +@dataclass(frozen=True, slots=True) +class ToolCall: + name: str + arguments: str + + @property + def parsed_arguments(self) -> dict[str, object]: + return _ARGUMENTS.validate_json(self.arguments) + + +@dataclass(frozen=True, slots=True) +class Reply: + """One generated reply, every number the gateway will read already decided.""" + + provider: Provider + seed: str + model: str + prompt_tokens: int + cached_tokens: int + reasoning_tokens: int + output_tokens: int + reasoning: str + text: str + tool_call: ToolCall | None + + @property + def visible_tokens(self) -> int: + return self.output_tokens - self.reasoning_tokens + + @property + def uncached_tokens(self) -> int: + return self.prompt_tokens - self.cached_tokens + + @property + def total_tokens(self) -> int: + return self.prompt_tokens + self.output_tokens + + +def plan_reply( + *, provider: Provider, seed: str, model: str, prompt_tokens: int, max_tokens: int, tool_names: Sequence[str] +) -> Reply: + """Size the reply from the request: a quarter to three quarters of the output cap, a tool call (with + reasoning first) when tools were offered, otherwise a text answer, a prefix cache hit on most requests.""" + output_tokens: Final = max(32, int(max_tokens * (0.25 + 0.5 * unit(f"{seed}:length")))) + tool: Final = "Grep" if "Grep" in tool_names else next(iter(tool_names), None) + reasoning_tokens: Final = int(output_tokens * REASONING_SHARE) if tool is not None else 0 + visible: Final = output_tokens - reasoning_tokens + cached: Final = ( + CACHED_PREFIX_TOKENS if unit(f"{seed}:cache") < CACHE_HIT_SHARE and prompt_tokens > CACHED_PREFIX_TOKENS else 0 + ) + return Reply( + provider=provider, + seed=seed, + model=model, + prompt_tokens=prompt_tokens, + cached_tokens=cached, + reasoning_tokens=reasoning_tokens, + output_tokens=output_tokens, + reasoning=thinking_text(seed, reasoning_tokens), + text="" if tool is not None else answer_text(seed, visible), + tool_call=ToolCall(tool, grep_arguments(seed, visible)) if tool is not None else None, + ) + + +def pieces(text: str, size: int) -> tuple[str, ...]: + """`text` cut into the deltas a provider streams it as, `size` characters each.""" + return tuple(text[i : i + size] for i in range(0, len(text), size)) + + +def sse(payload: dict[str, object], event: str | None = None) -> bytes: + data: Final = json.dumps(payload, separators=(",", ":")) + prefix: Final = f"event: {event}\n" if event is not None else "" + return f"{prefix}data: {data}\n\n".encode() + + +def grouped(events: Sequence[bytes], target_bytes: int) -> tuple[bytes, ...]: + """Consecutive events joined until a chunk reaches `target_bytes`, the size the provider's writes arrive in.""" + cumulative: Final = tuple(accumulate(len(event) for event in events)) + keyed: Final = tuple(zip((total // target_bytes for total in cumulative), events, strict=True)) + return tuple(b"".join(event for _, event in group) for _, group in groupby(keyed, key=lambda pair: pair[0])) + + +async def chunks_after(chunks: Sequence[bytes], delay_s: float) -> AsyncIterator[bytes]: + """The chunks back to back once `delay_s` has passed.""" + if delay_s > 0: + await asyncio.sleep(delay_s) + for chunk in chunks: + yield chunk diff --git a/tests/e2e/perf/setup/mocks/reply_openai.py b/tests/e2e/perf/setup/mocks/reply_openai.py new file mode 100644 index 00000000000..31ba7d31fef --- /dev/null +++ b/tests/e2e/perf/setup/mocks/reply_openai.py @@ -0,0 +1,340 @@ +"""OpenAI wire formats for a generated `Reply`: `/v1/chat/completions` (JSON and SSE chunks) and +`/v1/responses` (JSON and typed SSE events), with the field set, ids, event order, delta sizes and chunk +sizes of what api.openai.com sends gpt-6-luna replies as.""" + +from __future__ import annotations + +import json +import time +from collections.abc import Callable +from typing import Final + +from reply_common import Reply, grouped, ident, pieces, sse + +CHAT_DELTA_CHARS: Final = 4 +CHAT_CHUNK_BYTES: Final = 560 +RESPONSES_DELTA_CHARS: Final = 3 +RESPONSES_CHUNK_BYTES: Final = 250 + + +def _chat_usage(reply: Reply) -> dict[str, object]: + return { + "prompt_tokens": reply.prompt_tokens, + "completion_tokens": reply.output_tokens, + "total_tokens": reply.total_tokens, + "prompt_tokens_details": {"cached_tokens": reply.cached_tokens, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": reply.reasoning_tokens, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + } + + +def _chat_message(reply: Reply) -> dict[str, object]: + if reply.tool_call is None: + return {"role": "assistant", "content": reply.text, "refusal": None, "annotations": []} + call: Final = { + "id": ident(reply.seed, "call_", 24), + "type": "function", + "function": {"name": reply.tool_call.name, "arguments": reply.tool_call.arguments}, + } + return {"role": "assistant", "content": None, "refusal": None, "annotations": [], "tool_calls": [call]} + + +def chat_json(reply: Reply) -> bytes: + finish: Final = "stop" if reply.tool_call is None else "tool_calls" + body: Final = { + "id": ident(reply.seed, "chatcmpl-", 29), + "object": "chat.completion", + "created": int(time.time()), + "model": reply.model, + "choices": [{"index": 0, "message": _chat_message(reply), "logprobs": None, "finish_reason": finish}], + "usage": _chat_usage(reply), + "service_tier": "default", + "system_fingerprint": None, + } + return json.dumps(body, separators=(",", ":")).encode() + + +def _chat_deltas(reply: Reply, chunk: Callable[[dict[str, object]], bytes]) -> tuple[bytes, ...]: + if reply.tool_call is None: + return tuple(chunk({"content": piece}) for piece in pieces(reply.text, CHAT_DELTA_CHARS)) + opening: Final = { + "index": 0, + "id": ident(reply.seed, "call_", 24), + "type": "function", + "function": {"name": reply.tool_call.name, "arguments": ""}, + } + argument_pieces: Final = pieces(reply.tool_call.arguments, CHAT_DELTA_CHARS) + return ( + chunk({"tool_calls": [opening]}), + *(chunk({"tool_calls": [{"index": 0, "function": {"arguments": piece}}]}) for piece in argument_pieces), + ) + + +def chat_chunks(reply: Reply) -> tuple[bytes, ...]: + head: Final = { + "id": ident(reply.seed, "chatcmpl-", 29), + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": reply.model, + "service_tier": "default", + "system_fingerprint": None, + } + + def chunk(delta: dict[str, object], finish: str | None = None) -> bytes: + choice: Final = {"index": 0, "delta": delta, "logprobs": None, "finish_reason": finish} + return sse({**head, "choices": [choice], "usage": None}) + + deltas: Final = _chat_deltas(reply, chunk) + finish: Final = "stop" if reply.tool_call is None else "tool_calls" + events: Final = ( + chunk({"role": "assistant", "content": "", "refusal": None}), + *deltas, + chunk({}, finish), + sse({**head, "choices": [], "usage": _chat_usage(reply)}), + b"data: [DONE]\n\n", + ) + return grouped(events, CHAT_CHUNK_BYTES) + + +def _responses_usage(reply: Reply) -> dict[str, object]: + return { + "input_tokens": reply.prompt_tokens, + "input_tokens_details": {"cached_tokens": reply.cached_tokens}, + "output_tokens": reply.output_tokens, + "output_tokens_details": {"reasoning_tokens": reply.reasoning_tokens}, + "total_tokens": reply.total_tokens, + } + + +def _reasoning_item(reply: Reply) -> dict[str, object]: + return {"id": ident(reply.seed, "rs_", 50, hex_digits=True), "type": "reasoning", "summary": [], "content": []} + + +def _message_item(reply: Reply, text: str) -> dict[str, object]: + content: Final[dict[str, object]] = {"type": "output_text", "annotations": [], "logprobs": [], "text": text} + return { + "id": ident(reply.seed, "msg_", 50, hex_digits=True), + "type": "message", + "status": "completed", + "content": [content], + "role": "assistant", + } + + +def _function_item(reply: Reply, arguments: str, status: str) -> dict[str, object]: + assert reply.tool_call is not None + return { + "id": ident(reply.seed, "fc_", 50, hex_digits=True), + "type": "function_call", + "status": status, + "arguments": arguments, + "call_id": ident(reply.seed, "call_", 24), + "name": reply.tool_call.name, + } + + +def _output(reply: Reply) -> tuple[dict[str, object], ...]: + final: Final = ( + _message_item(reply, reply.text) + if reply.tool_call is None + else _function_item(reply, reply.tool_call.arguments, "completed") + ) + return (_reasoning_item(reply), final) if reply.reasoning_tokens else (final,) + + +def _response(reply: Reply, status: str, output: tuple[dict[str, object], ...]) -> dict[str, object]: + created: Final = int(time.time()) + return { + "id": ident(reply.seed, "resp_", 50, hex_digits=True), + "object": "response", + "created_at": created, + "status": status, + "background": False, + "completed_at": created if status == "completed" else None, + "error": None, + "incomplete_details": None, + "instructions": None, + "max_output_tokens": None, + "model": reply.model, + "output": list(output), + "parallel_tool_calls": True, + "previous_response_id": None, + "reasoning": {"effort": "medium", "summary": None}, + "service_tier": "default", + "store": True, + "temperature": 1.0, + "text": {"format": {"type": "text"}, "verbosity": "medium"}, + "tool_choice": "auto", + "tools": [], + "top_p": 1.0, + "truncation": "disabled", + "usage": _responses_usage(reply) if status == "completed" else None, + "user": None, + "metadata": {}, + } + + +def responses_json(reply: Reply) -> bytes: + return json.dumps(_response(reply, "completed", _output(reply)), separators=(",", ":")).encode() + + +def _message_events(reply: Reply, index: int, start: int) -> tuple[bytes, ...]: + item_id: Final = ident(reply.seed, "msg_", 50, hex_digits=True) + at: Final = {"item_id": item_id, "output_index": index, "content_index": 0} + deltas: Final = pieces(reply.text, RESPONSES_DELTA_CHARS) + part: Final[dict[str, object]] = {"type": "output_text", "annotations": [], "logprobs": [], "text": ""} + added: Final[dict[str, object]] = { + "id": item_id, + "type": "message", + "status": "in_progress", + "content": [], + "role": "assistant", + } + middle: Final = tuple( + sse( + {"type": "response.output_text.delta", **at, "delta": delta, "sequence_number": start + 2 + i}, + "response.output_text.delta", + ) + for i, delta in enumerate(deltas) + ) + end: Final = start + 2 + len(deltas) + return ( + sse( + {"type": "response.output_item.added", "output_index": index, "item": added, "sequence_number": start}, + "response.output_item.added", + ), + sse( + {"type": "response.content_part.added", **at, "part": part, "sequence_number": start + 1}, + "response.content_part.added", + ), + *middle, + sse( + {"type": "response.output_text.done", **at, "text": reply.text, "sequence_number": end}, + "response.output_text.done", + ), + sse( + { + "type": "response.content_part.done", + **at, + "part": {**part, "text": reply.text}, + "sequence_number": end + 1, + }, + "response.content_part.done", + ), + sse( + { + "type": "response.output_item.done", + "output_index": index, + "item": _message_item(reply, reply.text), + "sequence_number": end + 2, + }, + "response.output_item.done", + ), + ) + + +def _function_events(reply: Reply, index: int, start: int) -> tuple[bytes, ...]: + assert reply.tool_call is not None + item_id: Final = ident(reply.seed, "fc_", 50, hex_digits=True) + deltas: Final = pieces(reply.tool_call.arguments, RESPONSES_DELTA_CHARS) + middle: Final = tuple( + sse( + { + "type": "response.function_call_arguments.delta", + "item_id": item_id, + "output_index": index, + "delta": delta, + "sequence_number": start + 1 + i, + }, + "response.function_call_arguments.delta", + ) + for i, delta in enumerate(deltas) + ) + end: Final = start + 1 + len(deltas) + return ( + sse( + { + "type": "response.output_item.added", + "output_index": index, + "item": _function_item(reply, "", "in_progress"), + "sequence_number": start, + }, + "response.output_item.added", + ), + *middle, + sse( + { + "type": "response.function_call_arguments.done", + "item_id": item_id, + "output_index": index, + "arguments": reply.tool_call.arguments, + "sequence_number": end, + }, + "response.function_call_arguments.done", + ), + sse( + { + "type": "response.output_item.done", + "output_index": index, + "item": _function_item(reply, reply.tool_call.arguments, "completed"), + "sequence_number": end + 1, + }, + "response.output_item.done", + ), + ) + + +def responses_chunks(reply: Reply) -> tuple[bytes, ...]: + opening: Final = ( + sse( + {"type": "response.created", "response": _response(reply, "in_progress", ()), "sequence_number": 0}, + "response.created", + ), + sse( + {"type": "response.in_progress", "response": _response(reply, "in_progress", ()), "sequence_number": 1}, + "response.in_progress", + ), + ) + reasoning: Final = ( + ( + sse( + { + "type": "response.output_item.added", + "output_index": 0, + "item": _reasoning_item(reply), + "sequence_number": 2, + }, + "response.output_item.added", + ), + sse( + { + "type": "response.output_item.done", + "output_index": 0, + "item": _reasoning_item(reply), + "sequence_number": 3, + }, + "response.output_item.done", + ), + ) + if reply.reasoning_tokens + else () + ) + index: Final = len(reasoning) // 2 + start: Final = 2 + len(reasoning) + body: Final = ( + _message_events(reply, index, start) if reply.tool_call is None else _function_events(reply, index, start) + ) + sequence: Final = start + len(body) + closing: Final = sse( + { + "type": "response.completed", + "response": _response(reply, "completed", _output(reply)), + "sequence_number": sequence, + }, + "response.completed", + ) + return grouped((*opening, *reasoning, *body, closing), RESPONSES_CHUNK_BYTES) diff --git a/tests/e2e/perf/setup/mocks/server.py b/tests/e2e/perf/setup/mocks/server.py new file mode 100644 index 00000000000..6978d818602 --- /dev/null +++ b/tests/e2e/perf/setup/mocks/server.py @@ -0,0 +1,43 @@ +"""Run one perf mock out of process: `python server.py openai --port 8090`. + +Each kind is its own pod (or local process) with its own request log; the gateway talks to it over +HTTP exactly as it would to the real provider, and the test reads `/mock/requests` (providers) or +`/mock/calls` (guardrail) afterwards to join what was received against what each client saw. +""" + +from __future__ import annotations + +import argparse +from collections.abc import Callable +from typing import Final + +import uvicorn +from fastapi import FastAPI + +import mock_anthropic +import mock_bedrock +import mock_guardrail +import mock_openai + +KINDS: Final[dict[str, Callable[[], FastAPI]]] = { + "openai": mock_openai.build_app, + "anthropic": mock_anthropic.build_app, + "bedrock": mock_bedrock.build_app, + "guardrail": mock_guardrail.build_app, +} + + +def main() -> None: + parser: Final = argparse.ArgumentParser(description=__doc__) + _ = parser.add_argument("kind", choices=sorted(KINDS)) + _ = parser.add_argument("--host", default="0.0.0.0") + _ = parser.add_argument("--port", type=int, default=8090) + args: Final = parser.parse_args() + kind: Final = str(args.kind) # pyright: ignore[reportAny] + port: Final = int(args.port) # pyright: ignore[reportAny] + host: Final = str(args.host) # pyright: ignore[reportAny] + uvicorn.run(KINDS[kind](), host=host, port=port, log_level="warning", access_log=False) + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/perf/setup/observe.py b/tests/e2e/perf/setup/observe.py new file mode 100644 index 00000000000..af8295e8ffd --- /dev/null +++ b/tests/e2e/perf/setup/observe.py @@ -0,0 +1,298 @@ +"""Turns what a run left behind (the client records, what each mock received, the guardrail's +calls, the workers' memory) into the case's metrics, and the case's limits into a verdict.""" + +from __future__ import annotations + +import hashlib +import statistics +from collections import Counter +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Final + +from e2e_http import NoBody, get_external, post_external, unwrap +from memory_readings import RssCapture, RssReading, WorkerKey +from pydantic import BaseModel + +from setup.datastore import DbCounters, RedisCounters +from setup.mocks.reply_common import Provider +from setup.performance_test_case import Limit, Metric, MockLatency, PerformanceTestCase +from setup.wire import WireRecord + + +class MockRequestSeen(BaseModel): + provider: Provider + marker: str | None + route: str + stream: bool + output_tokens: int + text_sha1: str + delay_s: float + + +class MockRequestsSeen(BaseModel): + requests: list[MockRequestSeen] + + +class GuardrailCallSeen(BaseModel): + marker: str | None + input_type: str + action: str + + +class GuardrailCallsSeen(BaseModel): + calls: list[GuardrailCallSeen] + + +class Cleared(BaseModel): + cleared: int + + +class LatencyBody(BaseModel): + seconds: float + jitter_seconds: float + + +def read_mock_requests(url: str) -> list[MockRequestSeen]: + return unwrap(get_external(f"{url}/mock/requests", response_type=MockRequestsSeen, timeout=60.0)).requests + + +def read_guardrail_calls(url: str) -> list[GuardrailCallSeen]: + return unwrap(get_external(f"{url}/mock/calls", response_type=GuardrailCallsSeen, timeout=60.0)).calls + + +def reset_mock(url: str) -> int: + return unwrap(post_external(f"{url}/mock/reset", json=NoBody(), response_type=Cleared)).cleared + + +def set_mock_latency(url: str, latency: MockLatency) -> LatencyBody: + body: Final = LatencyBody(seconds=latency.seconds, jitter_seconds=latency.jitter_seconds) + return unwrap(post_external(f"{url}/mock/latency", json=body, response_type=LatencyBody)) + + +def percentile(values: Iterable[float], p: int) -> float | None: + ordered: Final = sorted(values) + if not ordered: + return None + index: Final = min(len(ordered) - 1, max(0, round(p / 100 * (len(ordered) - 1)))) + return ordered[index] + + +def worst_worker_growth_mb(before: RssCapture, after: RssCapture) -> float | None: + """Worst RSS growth across workers. A worker missing from `before` (a pod replaced during the + run) is measured against the median of the old workers.""" + old: Final[dict[WorkerKey, RssReading]] = {reading.worker: reading for reading in before.readings} + if not old or not after.readings: + return None + median_before: Final = statistics.median(reading.ram_usage_mb for reading in old.values()) + growth: list[float] = [] + for reading in after.readings: + baseline = old[reading.worker].ram_usage_mb if reading.worker in old else median_before + growth.append(reading.ram_usage_mb - baseline) + return max(growth) + + +@dataclass(frozen=True, slots=True) +class Observation: + records: tuple[WireRecord, ...] + mock_requests: tuple[MockRequestSeen, ...] + guardrail_calls: tuple[GuardrailCallSeen, ...] + rss_before: RssCapture + rss_after: RssCapture + wall_seconds: float + spend_logged_share: float | None = None + db_before: DbCounters | None = None + db_after: DbCounters | None = None + redis_before: RedisCounters | None = None + redis_after: RedisCounters | None = None + gateway_cpu_seconds: float | None = None + + def per_request(self, before: int | None, after: int | None) -> float | None: + if before is None or after is None or not self.records: + return None + return (after - before) / len(self.records) + + @property + def retried_markers(self) -> frozenset[str]: + counts: Final = Counter(seen.marker for seen in self.mock_requests if seen.marker is not None) + return frozenset(marker for marker, count in counts.items() if count > 1) + + @property + def served(self) -> tuple[WireRecord, ...]: + return tuple(record for record in self.records if record.ok and not record.blocked) + + @property + def _delay_by_marker(self) -> dict[str, float]: + return {seen.marker: seen.delay_s for seen in self.mock_requests if seen.marker is not None} + + def _clean(self, *, stream: bool) -> list[tuple[WireRecord, float]]: + """Served requests of one kind that were answered on the first attempt, with the mock's delay.""" + delays: Final = self._delay_by_marker + retried: Final = self.retried_markers + pairs: list[tuple[WireRecord, float]] = [] + for record in self.served: + delay = record.mock_delay_s if record.mock_delay_s is not None else delays.get(record.marker) + if record.stream == stream and delay is not None and record.marker not in retried: + pairs.append((record, delay)) + return pairs + + def unexpected_status_count(self) -> int: + """Requests the plan expected to succeed that came back with any other status (0 is a client timeout).""" + return sum(1 for record in self.records if not record.blocked and record.status != 200) + + def gateway_added_seconds(self) -> list[float]: + """Seconds the gateway added to each first-attempt request: total time minus the mock's own delay for a + non-streamed request, time to the first token minus the mock's delay to its first byte for a stream.""" + non_stream: Final = [record.total_s - delay for record, delay in self._clean(stream=False)] + stream: Final = [ + record.first_token_s - delay + for record, delay in self._clean(stream=True) + if record.first_token_s is not None + ] + return non_stream + stream + + def time_to_first_token_seconds(self) -> list[float]: + """Seconds from the client sending each first-attempt streamed request to its first token, mock delay + included, comparable across versions because the mocks generate the same replies after the same wait.""" + return [record.first_token_s for record, _ in self._clean(stream=True) if record.first_token_s is not None] + + def bodies_off_reply(self) -> int: + """Served requests whose text is not the reply the mock generated for them (plus the post-call rewrite + for requests in the intervene slice), a translation or streaming fault somewhere in the gateway.""" + seen_by_marker: Final = {seen.marker: seen for seen in self.mock_requests if seen.marker is not None} + return sum( + 1 + for record in self.served + if (seen := seen_by_marker.get(record.marker)) is not None + and _text_sha1(record.text, intervened=record.marker.endswith("-i")) != seen.text_sha1 + ) + + def guardrail_missed(self) -> int: + """Served requests that did not get exactly one pre-call and one post-call guardrail call.""" + pre: Final = Counter(call.marker for call in self.guardrail_calls if call.input_type == "request") + post: Final = Counter(call.marker for call in self.guardrail_calls if call.input_type == "response") + return sum(1 for record in self.served if pre[record.marker] != 1 or post[record.marker] != 1) + + +def _text_sha1(text: str, *, intervened: bool) -> str: + original: Final = text.removesuffix(" [post]") if intervened else text + return hashlib.sha1(original.encode()).hexdigest() + + +def observe(metric: Metric, percentile_rank: int | None, observation: Observation) -> float | None: + """The observed value of one metric, at a percentile when the metric is per request.""" + per_request: dict[str, list[float]] = { + "gateway_added_seconds": observation.gateway_added_seconds(), + "time_to_first_token_seconds": observation.time_to_first_token_seconds(), + } + if metric in per_request: + return percentile(per_request[metric], 50 if percentile_rank is None else percentile_rank) + if metric == "unexpected_status_count": + return float(observation.unexpected_status_count()) + if metric == "postgres_calls_per_request": + db_before, db_after = observation.db_before, observation.db_after + return observation.per_request(db_before and db_before.commits, db_after and db_after.commits) + if metric == "redis_calls_per_request": + redis_before, redis_after = observation.redis_before, observation.redis_after + return observation.per_request(redis_before and redis_before.commands, redis_after and redis_after.commands) + if metric == "cpu_seconds_per_request": + cpu: Final = observation.gateway_cpu_seconds + return None if cpu is None or not observation.records else cpu / len(observation.records) + return worst_worker_growth_mb(observation.rss_before, observation.rss_after) + + +@dataclass(frozen=True, slots=True) +class Row: + label: str + observed: float | None + limit: float | None + + @property + def over(self) -> bool: + return self.limit is not None and self.observed is not None and self.observed > self.limit + + @property + def missing(self) -> bool: + return self.limit is not None and self.observed is None + + +REPORTED: Final[tuple[tuple[Metric, int | None], ...]] = ( + ("unexpected_status_count", None), + ("gateway_added_seconds", 50), + ("gateway_added_seconds", 99), + ("time_to_first_token_seconds", 50), + ("time_to_first_token_seconds", 99), + ("memory_growth_megabytes", None), + ("postgres_calls_per_request", None), + ("redis_calls_per_request", None), + ("cpu_seconds_per_request", None), +) + + +def evaluate(case: PerformanceTestCase, observation: Observation) -> tuple[list[Row], list[str]]: + """Every reported metric with its observed value and the case's limit if it has one, plus one line per + limit that was exceeded or that the run could not observe.""" + limits: Final[dict[tuple[Metric, int | None], Limit]] = { + (limit.metric, limit.percentile): limit for limit in case.limits + } + keys: list[tuple[Metric, int | None]] = list(REPORTED) + keys.extend(key for key in limits if key not in keys) + rows: list[Row] = [] + for metric, rank in keys: + limit = limits.get((metric, rank)) + label = f"{metric} percentile {rank}" if rank is not None else metric + rows.append(Row(label=label, observed=observe(metric, rank, observation), limit=limit.max if limit else None)) + exceeded: Final = [ + f"{row.label}: observed {row.observed:.4f} > limit {row.limit:.4f}" + if row.over and row.observed is not None and row.limit is not None + else f"{row.label}: limit set but nothing observed" + for row in rows + if row.over or row.missing + ] + return rows, exceeded + + +def render(case: PerformanceTestCase, observation: Observation, rows: list[Row]) -> str: + by_endpoint: Final = Counter((record.endpoint, record.stream) for record in observation.records) + failures: Final = [record for record in observation.records if not record.ok][:10] + lines: list[str] = [ + f"perf case {case.name}: {len(observation.records)} requests, {len(observation.served)} served, " + f"{sum(1 for r in observation.records if r.blocked and r.ok)} blocked as planned, " + f"{sum(1 for r in observation.records if not r.ok)} wrong, wall {observation.wall_seconds:.1f}s", + " " + + ", ".join( + f"{endpoint}{' stream' if stream else ''}={count}" + for (endpoint, stream), count in sorted(by_endpoint.items()) + ), + f" {'metric':32} {'observed':>12} {'limit':>12}", + ] + for row in rows: + observed = "n/a" if row.observed is None else f"{row.observed:.4f}" + limit = "" if row.limit is None else f"{row.limit:.4f}" + flag = " OVER" if row.over else " MISSING" if row.missing else "" + lines.append(f" {row.label:32} {observed:>12} {limit:>12}{flag}") + lines.extend(_diagnostics(observation)) + for record in failures: + lines.append( + f" wrong: seq {record.seq} {record.caller} {record.endpoint}{' stream' if record.stream else ''} status {record.status}: {record.error}" + ) + return "\n".join(lines) + + +def _diagnostics(observation: Observation) -> list[str]: + """Functional side-channels printed for whoever reads a red run, never limits of a case.""" + count: Final = max(1, len(observation.records)) + share: Final = observation.spend_logged_share + return [ + " diagnostics (not limits):", + f" bodies off reply {observation.bodies_off_reply()}", + f" guardrail calls per request {len(observation.guardrail_calls) / count:.3f}", + f" guardrail missed {observation.guardrail_missed() if observation.guardrail_calls else 'n/a'}", + f" retried share {len(observation.retried_markers) / count:.4f}", + f" spend logged share {'n/a' if share is None else f'{share:.3f}'}", + f" redis connected clients {_clients(observation.redis_before)} -> {_clients(observation.redis_after)}", + ] + + +def _clients(counters: RedisCounters | None) -> str: + return "n/a" if counters is None else str(counters.connected_clients) diff --git a/tests/e2e/perf/setup/performance_test_case.py b/tests/e2e/perf/setup/performance_test_case.py new file mode 100644 index 00000000000..d90b262329c --- /dev/null +++ b/tests/e2e/perf/setup/performance_test_case.py @@ -0,0 +1,191 @@ +"""The schema of a perf test: a `PerformanceTestCase` names what the gateway is deployed with +(model groups and guardrails, both written out in the gateway's config) and lists what the test +itself creates and sends: teams, Keycloak users, virtual keys, the request shapes, the traffic plan +and the `Limit`s that decide pass or fail. Add a case by adding a module to `cases/`. + +The same shape as `CostTrackingTestCase` in the cost suite: frozen dataclasses with plain defaults, +so a case is a literal a reviewer can read top to bottom. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Final, Literal + +Endpoint = Literal["chat_completions", "responses", "messages"] +MockProvider = Literal["openai", "anthropic", "bedrock"] +GuardrailMode = Literal["pre_call", "post_call"] +Metric = Literal[ + "unexpected_status_count", + "gateway_added_seconds", + "time_to_first_token_seconds", + "memory_growth_megabytes", + "postgres_calls_per_request", + "cpu_seconds_per_request", + "redis_calls_per_request", +] + + +@dataclass(frozen=True, slots=True) +class Team: + """One team (`/team/new`). `models` names the model groups the team may call; empty means all.""" + + name: str + models: tuple[str, ...] = () + max_budget: float | None = None + rpm_limit: int | None = None + tpm_limit: int | None = None + + +@dataclass(frozen=True, slots=True) +class User: + """One Keycloak user who calls the gateway with a JWT. `team` puts them in that team's group.""" + + name: str + team: str | None = None + max_budget: float | None = None + + +@dataclass(frozen=True, slots=True) +class VirtualKey: + """One virtual key (`/key/generate`), scoped to a team and/or a user when named.""" + + name: str + team: str | None = None + user: str | None = None + models: tuple[str, ...] = () + rpm_limit: int | None = None + max_budget: float | None = None + + +@dataclass(frozen=True, slots=True) +class Guardrail: + """A guardrail the gateway's config defines (a `generic_guardrail_api` pointed at the mock + guardrail, `default_on`); the case names it so the plan knows which hooks every request sees.""" + + name: str + mode: GuardrailMode + + +@dataclass(frozen=True, slots=True) +class RequestShape: + """What one request looks like: a coding-agent session (`setup/corpus/session.py`, the system prompt, the + Read/Edit/Write/Bash/Grep/Glob tools when `tools` is set, then file-reading turns carrying real source files) + whose size in input tokens is drawn per request from the `input_tokens` range, and the output cap. The mocks + answer with a reply generated for the request, sized from that cap, a tool call when `tools` is set.""" + + name: str + input_tokens: tuple[int, int] + tools: bool = False + output_tokens: int = 1_024 + + +SHORT: Final = RequestShape("short", input_tokens=(5_000, 100_000), tools=True) +LONG: Final = RequestShape("long", input_tokens=(300_000, 1_000_000), tools=True) + + +@dataclass(frozen=True, slots=True) +class Traffic: + """The plan: `total_requests` drawn from `seed`, sent at `rps` by `concurrency` simulated + users. Shares are fractions of the requests: `endpoints` and `shapes` must each sum to 1, + `intervene_share` picks non-streaming requests the guardrail rewrites, `block_share` picks + requests it blocks.""" + + seed: int + rps: float + total_requests: int + endpoints: Mapping[Endpoint, float] = field(default_factory=lambda: {"chat_completions": 1.0}) + shapes: Mapping[RequestShape, float] = field(default_factory=lambda: {SHORT: 1.0}) + stream_share: float = 0.0 + intervene_share: float = 0.0 + block_share: float = 0.0 + concurrency: int = 20 + + @property + def expected_seconds(self) -> float: + return self.total_requests / self.rps + + +@dataclass(frozen=True, slots=True) +class Limit: + """One pass/fail rule: `metric` (at `percentile`, or the whole-run value) must not exceed `max`. + Metrics without a Limit are still reported.""" + + metric: Metric + max: float + percentile: int | None = None + + @property + def label(self) -> str: + return f"{self.metric} percentile {self.percentile}" if self.percentile is not None else self.metric + + +@dataclass(frozen=True, slots=True) +class MockLatency: + """Seconds a mock provider waits before its first byte, plus or minus up to `jitter_seconds` per request + (drawn from the request marker, so a rerun waits the same). Zero means the mocks answer at once.""" + + seconds: float = 0.0 + jitter_seconds: float = 0.0 + + +@dataclass(frozen=True, slots=True) +class PerformanceTestCase: + """`model_groups` and `guardrails` name what the gateway config deploys (setup checks they are + there); everything else the test creates or sends. `mock_latency` is how long every mock provider waits + before answering, set on the mock pods at setup, so the provider side of each request is a known + constant and the rest of the client's time is the gateway's.""" + + name: str + teams: tuple[Team, ...] + users: tuple[User, ...] + virtual_keys: tuple[VirtualKey, ...] + model_groups: tuple[str, ...] + guardrails: tuple[Guardrail, ...] + traffic: Traffic + limits: tuple[Limit, ...] + mock_latency: MockLatency = field(default_factory=lambda: MockLatency()) + + def team(self, name: str) -> Team: + return next(team for team in self.teams if team.name == name) + + @property + def has_post_guardrail(self) -> bool: + return any(guardrail.mode == "post_call" for guardrail in self.guardrails) + + def validate(self) -> None: + """Fail fast on a case that cannot run: dangling names, shares that do not add up, a model + group no caller may use.""" + team_names: Final = {team.name for team in self.teams} + user_names: Final = {user.name for user in self.users} + group_names: Final = set(self.model_groups) + for names, label in ( + ([team.name for team in self.teams], "team"), + ([user.name for user in self.users], "user"), + ([key.name for key in self.virtual_keys], "key"), + (list(self.model_groups), "model group"), + ([guardrail.name for guardrail in self.guardrails], "guardrail"), + ): + assert len(names) == len(set(names)), f"{self.name}: duplicate {label} names in {names}" + for team in self.teams: + assert set(team.models) <= group_names, f"{self.name}: team {team.name} names unknown models {team.models}" + for user in self.users: + assert user.team is None or user.team in team_names, ( + f"{self.name}: user {user.name} in unknown team {user.team}" + ) + for key in self.virtual_keys: + assert key.team is None or key.team in team_names, f"{self.name}: key {key.name} in unknown team {key.team}" + assert key.user is None or key.user in user_names, ( + f"{self.name}: key {key.name} for unknown user {key.user}" + ) + assert set(key.models) <= group_names, f"{self.name}: key {key.name} names unknown models {key.models}" + assert abs(sum(self.traffic.endpoints.values()) - 1.0) < 1e-9, f"{self.name}: endpoint shares must sum to 1" + assert abs(sum(self.traffic.shapes.values()) - 1.0) < 1e-9, f"{self.name}: shape shares must sum to 1" + assert 0 <= self.traffic.stream_share <= 1, f"{self.name}: stream_share out of range" + assert self.traffic.intervene_share + self.traffic.block_share <= 1, f"{self.name}: guardrail shares exceed 1" + assert self.traffic.total_requests > 0 and self.traffic.rps > 0, f"{self.name}: empty traffic" + assert self.model_groups, f"{self.name}: no model groups" + assert 0 <= self.mock_latency.jitter_seconds <= self.mock_latency.seconds or self.mock_latency.seconds == 0, ( + f"{self.name}: mock latency jitter {self.mock_latency.jitter_seconds} s exceeds {self.mock_latency.seconds} s" + ) diff --git a/tests/e2e/perf/setup/plan.py b/tests/e2e/perf/setup/plan.py new file mode 100644 index 00000000000..6b9a5a19662 --- /dev/null +++ b/tests/e2e/perf/setup/plan.py @@ -0,0 +1,122 @@ +"""The deterministic request plan of a `PerformanceTestCase`: every request the run sends, in order, +drawn once from the case's seed so two runs of the same case send the same traffic. Each request +carries a marker (`perfreq-000123-n`) that the runner puts in the final user message and the mock +providers and the mock guardrail read back, so the test can join what the client saw against what +each mock received, and a `RequestShape` the runner hands the corpus to build the body.""" + +from __future__ import annotations + +import random +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final, Literal, TypeVar + +from setup.performance_test_case import Endpoint, PerformanceTestCase, RequestShape, Team + +Flag = Literal["n", "i", "b"] +CallerKind = Literal["jwt", "key"] +T = TypeVar("T") + + +@dataclass(frozen=True, slots=True) +class Caller: + kind: CallerKind + name: str + + @property + def label(self) -> str: + return f"{self.kind}:{self.name}" + + +@dataclass(frozen=True, slots=True) +class PlannedRequest: + seq: int + marker: str + caller: Caller + endpoint: Endpoint + model_group: str + stream: bool + flag: Flag + shape: RequestShape + input_tokens: int + + @property + def blocked(self) -> bool: + return self.flag == "b" + + +def _allowed_models(case: PerformanceTestCase, team: Team | None, own: tuple[str, ...]) -> frozenset[str]: + if own: + return frozenset(own) + if team is not None and team.models: + return frozenset(team.models) + return frozenset(case.model_groups) + + +def _callers_by_group(case: PerformanceTestCase) -> dict[str, tuple[Caller, ...]]: + """Which callers may use each model group, in case order, so a caller never asks for a model its + team or key does not allow and every denial in the run is a bug, not the plan.""" + reach: list[tuple[Caller, frozenset[str]]] = [] + for user in case.users: + user_team = case.team(user.team) if user.team else None + reach.append((Caller("jwt", user.name), _allowed_models(case, user_team, ()))) + for key in case.virtual_keys: + key_team = case.team(key.team) if key.team else None + reach.append((Caller("key", key.name), _allowed_models(case, key_team, key.models))) + by_group: dict[str, tuple[Caller, ...]] = {} + for group in case.model_groups: + by_group[group] = tuple(caller for caller, models in reach if group in models) + assert by_group[group], f"{case.name}: no caller may use model group {group}" + return by_group + + +def _pick(rng: random.Random, shares: Mapping[T, float]) -> T: + draw: Final = rng.random() + cumulative = 0.0 + chosen: T | None = None + for item, share in shares.items(): + cumulative += share + chosen = item + if draw < cumulative: + break + assert chosen is not None, "empty shares" + return chosen + + +def _pick_flag(rng: random.Random, case: PerformanceTestCase, *, stream: bool) -> Flag: + if not case.guardrails: + return "n" + draw: Final = rng.random() + if draw < case.traffic.block_share: + return "b" + if draw < case.traffic.block_share + case.traffic.intervene_share and not stream: + return "i" + return "n" + + +def build_plan(case: PerformanceTestCase) -> tuple[PlannedRequest, ...]: + case.validate() + rng: Final = random.Random(case.traffic.seed) + callers: Final = _callers_by_group(case) + planned: list[PlannedRequest] = [] + for seq in range(1, case.traffic.total_requests + 1): + endpoint = _pick(rng, case.traffic.endpoints) + group = rng.choice(case.model_groups) + stream = rng.random() < case.traffic.stream_share + flag = _pick_flag(rng, case, stream=stream) + shape = _pick(rng, case.traffic.shapes) + low, high = shape.input_tokens + planned.append( + PlannedRequest( + seq=seq, + marker=f"perfreq-{seq:06d}-{flag}", + caller=rng.choice(callers[group]), + endpoint=endpoint, + model_group=group, + stream=stream, + flag=flag, + shape=shape, + input_tokens=rng.randint(low, high), + ) + ) + return tuple(planned) diff --git a/tests/e2e/perf/setup/podcpu.py b/tests/e2e/perf/setup/podcpu.py new file mode 100644 index 00000000000..c2402b38b71 --- /dev/null +++ b/tests/e2e/perf/setup/podcpu.py @@ -0,0 +1,126 @@ +"""CPU seconds the gateway pods burned during the traffic phase, from metrics-server (`metrics.k8s.io` +PodMetrics), so a case can put a ceiling on the CPU the gateway costs per request: a slower gateway shows +up here at 5 rps, long before it saturates. The runner samples the pods matching +`E2E_PERF_GATEWAY_POD_SELECTOR` every few seconds with its own service account; `usage.cpu` is the +pods' average core use over metrics-server's window, integrated over the sampling interval. Without the +selector or the in-cluster token the reading is None and a Limit on it reports as missing.""" + +from __future__ import annotations + +import os +import threading +import time +from collections.abc import Generator +from contextlib import contextmanager +from pathlib import Path +from typing import Final +from urllib.parse import quote + +from e2e_http import AuthHeaders, Success, get_external +from pydantic import BaseModel + +SELECTOR_ENV: Final = "E2E_PERF_GATEWAY_POD_SELECTOR" +_SERVICE_ACCOUNT: Final = Path("/var/run/secrets/kubernetes.io/serviceaccount") +_INTERVAL_S: Final = 5.0 +_SUFFIX_CORES: Final = {"n": 1e-9, "u": 1e-6, "m": 1e-3} + + +class _Usage(BaseModel): + cpu: str + + +class _Container(BaseModel): + usage: _Usage + + +class _PodMetrics(BaseModel): + containers: list[_Container] + + +class _PodMetricsList(BaseModel): + items: list[_PodMetrics] + + +def cores(quantity: str) -> float: + """A Kubernetes CPU quantity (`1500m`, `250000000n`, `2`) as cores.""" + if quantity and quantity[-1] in _SUFFIX_CORES: + return float(quantity[:-1]) * _SUFFIX_CORES[quantity[-1]] + return float(quantity) + + +class GatewayCpuSampler: + """Integrates the matching pods' core use over the time it is running; `seconds` is CPU seconds.""" + + def __init__(self, selector: str, namespace: str, token: str, ca_bundle: str) -> None: + self._url: Final = ( + f"https://kubernetes.default.svc/apis/metrics.k8s.io/v1beta1/namespaces/{namespace}/pods" + f"?labelSelector={quote(selector)}" + ) + self._headers: Final = AuthHeaders(authorization=f"Bearer {token}") + self._ca_bundle: Final = ca_bundle + self._stop: Final = threading.Event() + self._thread: Final = threading.Thread(target=self._run, name="gateway-cpu-sampler", daemon=True) + self._cpu_seconds = 0.0 + self._samples = 0 + + @property + def seconds(self) -> float | None: + return self._cpu_seconds if self._samples > 0 else None + + def _read_cores(self) -> float | None: + result: Final = get_external( + self._url, response_type=_PodMetricsList, headers=self._headers, timeout=10.0, ca_bundle=self._ca_bundle + ) + if not isinstance(result, Success): + return None + return sum(cores(container.usage.cpu) for pod in result.data.items for container in pod.containers) + + def _run(self) -> None: + last: float = time.monotonic() # rebind-ok: the sampling clock advances every tick + while not self._stop.wait(_INTERVAL_S): + now: float = time.monotonic() + current: float | None = self._read_cores() + if current is not None: + self._cpu_seconds += current * (now - last) + self._samples += 1 + last = now # rebind-ok: the sampling clock advances every tick + + def start(self) -> None: + self._thread.start() + + def stop(self) -> None: + self._stop.set() + self._thread.join(timeout=_INTERVAL_S * 3) + + +class NoSampler: + seconds: Final[None] = None + + def start(self) -> None: + return None + + def stop(self) -> None: + return None + + +def _in_cluster() -> tuple[str, str, str] | None: + token: Final = _SERVICE_ACCOUNT / "token" + ca: Final = _SERVICE_ACCOUNT / "ca.crt" + namespace: Final = _SERVICE_ACCOUNT / "namespace" + if not (token.is_file() and ca.is_file() and namespace.is_file()): + return None + return token.read_text().strip(), str(ca), namespace.read_text().strip() + + +@contextmanager +def gateway_cpu_sampler() -> Generator[GatewayCpuSampler | NoSampler]: + selector: Final = os.environ.get(SELECTOR_ENV, "") + cluster: Final = _in_cluster() + sampler: Final = ( + GatewayCpuSampler(selector, cluster[2], cluster[0], cluster[1]) if selector and cluster else NoSampler() + ) + sampler.start() + try: + yield sampler + finally: + sampler.stop() diff --git a/tests/e2e/perf/setup/provision.py b/tests/e2e/perf/setup/provision.py new file mode 100644 index 00000000000..05bdcf43140 --- /dev/null +++ b/tests/e2e/perf/setup/provision.py @@ -0,0 +1,190 @@ +"""Builds the caller side of a `PerformanceTestCase` on the gateway: the teams, the Keycloak users +(JWT callers) and the virtual keys, each suffixed with the run marker and deferred to the +`ResourceManager`, so two runs never collide and a failed run still tears down. Model groups and +guardrails are not created here: they are written out in the gateway's config (the helm values in +the lane, `gateway/perf_ci_config.yml` locally), and setup only checks the deployed gateway has what +the case names.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Final + +from e2e_config import unique_marker +from e2e_http import NoBody, unwrap +from idp import TESTS_CLIENT_ID, Keycloak +from lifecycle import ResourceManager +from models import KeyGenerateBody, TeamNewBody, UserNewBody, UserNewResponse +from proxy_client import ProxyClient +from pydantic import BaseModel, ConfigDict +from setup.performance_test_case import MockProvider, PerformanceTestCase +from setup.plan import Caller +from setup.wire import TokenSource + +MOCK_PROVIDERS: Final[tuple[MockProvider, ...]] = ("openai", "anthropic", "bedrock") + +MOCK_URL_ENV: Final[dict[MockProvider, str]] = { + "openai": "E2E_PERF_MOCK_OPENAI_URL", + "anthropic": "E2E_PERF_MOCK_ANTHROPIC_URL", + "bedrock": "E2E_PERF_MOCK_BEDROCK_URL", +} +MOCK_URL_DEFAULT: Final[dict[MockProvider, str]] = { + "openai": "http://127.0.0.1:8090/v1", + "anthropic": "http://127.0.0.1:8091", + "bedrock": "http://127.0.0.1:8092", +} +GUARDRAIL_URL_ENV: Final = "E2E_PERF_MOCK_GUARDRAIL_URL" +GUARDRAIL_URL_DEFAULT: Final = "http://127.0.0.1:8093" + + +def mock_url(provider: MockProvider) -> str: + """The mock's root, where its control routes live. The same variable feeds the gateway's + `api_base`, which for OpenAI ends in `/v1`, so that suffix is dropped here.""" + url: Final = os.environ.get(MOCK_URL_ENV[provider], MOCK_URL_DEFAULT[provider]).rstrip("/") + return url.removesuffix("/v1") + + +def guardrail_url() -> str: + return os.environ.get(GUARDRAIL_URL_ENV, GUARDRAIL_URL_DEFAULT).rstrip("/") + + +class PerfTeamNewBody(TeamNewBody): + """`/team/new` with the per-team settings a `Team` carries.""" + + max_budget: float | None = None + rpm_limit: int | None = None + tpm_limit: int | None = None + + +class ListedGuardrailParams(BaseModel): + model_config = ConfigDict(frozen=True) + + mode: str | list[str] + + +class ListedGuardrail(BaseModel): + model_config = ConfigDict(frozen=True) + + guardrail_name: str + litellm_params: ListedGuardrailParams + + +class GuardrailList(BaseModel): + model_config = ConfigDict(frozen=True) + + guardrails: list[ListedGuardrail] + + +@dataclass(frozen=True, slots=True) +class CaseSetup: + marker: str + team_ids: dict[str, str] + key_bearers: dict[str, str] + token_sources: dict[str, TokenSource] + user_ids: dict[str, str] + + def authorization(self, caller: Caller) -> str | None: + return f"Bearer {self.key_bearers[caller.name]}" if caller.kind == "key" else None + + def token_source(self, caller: Caller) -> TokenSource | None: + return self.token_sources[caller.name] if caller.kind == "jwt" else None + + +def check_deployed(proxy: ProxyClient, case: PerformanceTestCase) -> None: + """The gateway was deployed with every model group and guardrail the case names, with the + guardrail on the hook the case expects. A mismatch is a config drift, not a perf result.""" + deployed: Final = sorted({entry.model_name for entry in proxy.model_info()}) + missing: Final = [group for group in case.model_groups if group not in deployed] + assert not missing, f"{case.name}: model groups {missing} are not in the gateway config (deployed: {deployed})" + if not case.guardrails: + return + listed: Final = unwrap( + proxy.transport.get( + "/guardrails/list", headers=proxy.transport.master, params=NoBody(), response_type=GuardrailList + ) + ).guardrails + modes: Final = {guardrail.guardrail_name: guardrail.litellm_params.mode for guardrail in listed} + for guardrail in case.guardrails: + mode = modes.get(guardrail.name) + hooks = mode if isinstance(mode, list) else [mode] + assert guardrail.mode in hooks, ( + f"{case.name}: guardrail {guardrail.name} should be on {guardrail.mode}, gateway has {modes}" + ) + + +def build_setup( + case: PerformanceTestCase, *, proxy: ProxyClient, idp: Keycloak, resources: ResourceManager +) -> CaseSetup: + check_deployed(proxy, case) + marker: Final = unique_marker() + team_ids: dict[str, str] = {} + keycloak_groups: dict[str, str] = {} + for team in case.teams: + team_id = f"perf-{team.name}-{marker}" + team_ids[team.name] = proxy.create_team( + PerfTeamNewBody( + team_alias=team_id, + team_id=team_id, + models=list(team.models), + max_budget=team.max_budget, + rpm_limit=team.rpm_limit, + tpm_limit=team.tpm_limit, + ) + ) + resources.defer(lambda team_id=team_id: proxy.delete_team(team_id)) + group_id = idp.create_group(team_id) + keycloak_groups[team_id] = group_id + resources.defer(lambda group_id=group_id: idp.delete_group(group_id)) + + user_ids: dict[str, str] = {} + token_sources: dict[str, TokenSource] = {} + for user in case.users: + groups = (team_ids[user.team],) if user.team else () + identity = idp.provision_user( + marker=f"{marker}-{user.name}", + groups=groups, + group_ids=tuple(keycloak_groups[group] for group in groups), + defer=resources.defer, + ) + user_ids[user.name] = identity.user_id + token_sources[user.name] = TokenSource( + token_url=idp.token_url(idp.realm), + client_id=TESTS_CLIENT_ID, + username=identity.username, + password=identity.password, + ) + created = unwrap( + proxy.transport.post( + "/user/new", + headers=proxy.transport.master, + json=UserNewBody( + user_email=f"{identity.username}@example.com", + user_role="internal_user", + user_id=identity.user_id, + auto_create_key=False, + max_budget=user.max_budget, + ), + response_type=UserNewResponse, + ) + ) + resources.defer(lambda user_id=created.user_id: proxy.delete_user(user_id)) + + key_bearers: dict[str, str] = {} + for key in case.virtual_keys: + generated = proxy.generate_key( + KeyGenerateBody( + models=list(key.models), + team_id=team_ids[key.team] if key.team else None, + user_id=user_ids[key.user] if key.user else None, + rpm_limit=key.rpm_limit, + max_budget=key.max_budget, + key_alias=f"perf-{key.name}-{marker}", + ) + ) + resources.defer(lambda generated=generated: proxy.delete_key(generated)) + key_bearers[key.name] = generated + + return CaseSetup( + marker=marker, team_ids=team_ids, key_bearers=key_bearers, token_sources=token_sources, user_ids=user_ids + ) diff --git a/tests/e2e/perf/setup/run.py b/tests/e2e/perf/setup/run.py new file mode 100644 index 00000000000..9ba9b45104d --- /dev/null +++ b/tests/e2e/perf/setup/run.py @@ -0,0 +1,129 @@ +"""Sends a case's plan through the gateway from a separate process: the plan is written to disk as +wire requests, Locust (`locustfile.py`) replays it at the planned rate, and the records it +wrote come back as `WireRecord`s.""" + +from __future__ import annotations + +import os +import shutil +import subprocess +import sys +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +from e2e_config import PROXY_BASE_URL +from pydantic import TypeAdapter + +from setup.performance_test_case import PerformanceTestCase +from setup.plan import PlannedRequest +from setup.provision import CaseSetup +from setup.wire import WireRecord, WireRequest, WireShape + +_LOCUSTFILE: Final = Path(__file__).with_name("locustfile.py") +_PERF_DIR: Final = Path(__file__).resolve().parents[1] +_PLAN_ADAPTER: Final = TypeAdapter(list[WireRequest]) +_RECORD_ADAPTER: Final = TypeAdapter(WireRecord) +_PATHS: Final = {"chat_completions": "/v1/chat/completions", "responses": "/v1/responses", "messages": "/v1/messages"} + + +@dataclass(frozen=True, slots=True) +class RunOutcome: + records: tuple[WireRecord, ...] + locust_exit: int + locust_output: str + + @property + def wall_seconds(self) -> float: + if not self.records: + return 0.0 + first: Final = min(record.started_at for record in self.records) + last: Final = max(record.started_at + record.total_s for record in self.records) + return last - first + + +def wire_request(item: PlannedRequest, setup: CaseSetup) -> WireRequest: + return WireRequest( + seq=item.seq, + marker=item.marker, + caller=item.caller.label, + endpoint=item.endpoint, + path=_PATHS[item.endpoint], + model=item.model_group, + shape=WireShape( + name=item.shape.name, + input_tokens=item.input_tokens, + tools=item.shape.tools, + output_tokens=item.shape.output_tokens, + ), + authorization=setup.authorization(item.caller), + token_source=setup.token_source(item.caller), + stream=item.stream, + blocked=item.blocked, + ) + + +def _keep_records_artifact(records_path: Path) -> None: + """Copy the per-request records next to the JUnit file when the runner exposes an artifact + directory (`E2E_ARTIFACT_DIR`), so a red lane run can be broken down after the pod is gone.""" + artifact_dir: Final = os.environ.get("E2E_ARTIFACT_DIR") + if not artifact_dir or not Path(artifact_dir).is_dir(): + return + _ = shutil.copyfile(records_path, Path(artifact_dir) / "perf-records.jsonl") + + +def run_plan(case: PerformanceTestCase, plan: tuple[PlannedRequest, ...], setup: CaseSetup) -> RunOutcome: + """Replay `plan` with Locust in a child process and collect its records. The run-time ceiling + is well past the planned duration so Locust's own clock never cuts a slow run short; the + locustfile quits itself once the last record is written.""" + wire: Final = [wire_request(item, setup) for item in plan] + ceiling_s: Final = int(case.traffic.expected_seconds * 2 + 180) + with tempfile.TemporaryDirectory(prefix="e2e-perf-") as directory: + plan_path: Final = Path(directory) / "plan.json" + records_path: Final = Path(directory) / "records.jsonl" + _ = plan_path.write_bytes(_PLAN_ADAPTER.dump_json(wire)) + records_path.touch() + env: Final = { + **os.environ, + "PYTHONPATH": os.pathsep.join(filter(None, (str(_PERF_DIR), os.environ.get("PYTHONPATH")))), + "PERF_PLAN_FILE": str(plan_path), + "PERF_RECORDS_FILE": str(records_path), + "PERF_RPS": str(case.traffic.rps), + "PERF_USERS": str(case.traffic.concurrency), + } + completed: Final = subprocess.run( + [ + sys.executable, + "-P", + "-m", + "locust", + "--headless", + "--only-summary", + "--locustfile", + str(_LOCUSTFILE), + "--host", + PROXY_BASE_URL, + "--users", + str(case.traffic.concurrency), + "--spawn-rate", + str(case.traffic.concurrency), + "--run-time", + f"{ceiling_s}s", + "--exit-code-on-error", + "0", + ], + env=env, + capture_output=True, + text=True, + timeout=ceiling_s + 120, + check=False, + ) + records: Final = tuple( + _RECORD_ADAPTER.validate_json(line) + for line in records_path.read_text(encoding="utf-8").splitlines() + if line.strip() + ) + _keep_records_artifact(records_path) + output: Final = (completed.stdout + completed.stderr)[-6000:] + return RunOutcome(records=records, locust_exit=completed.returncode, locust_output=output) diff --git a/tests/e2e/perf/setup/wire.py b/tests/e2e/perf/setup/wire.py new file mode 100644 index 00000000000..07dc5d170ae --- /dev/null +++ b/tests/e2e/perf/setup/wire.py @@ -0,0 +1,91 @@ +"""What crosses the process boundary between the test and the Locust runner: the wire plan (one +entry per planned request, with the shape the runner expands into the body right before sending, +so a 300k-token request costs a few hundred bytes in the plan file) and the record each request +leaves behind.""" + +from __future__ import annotations + +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final, Literal + +from pydantic import BaseModel, ConfigDict, Field + +from setup.corpus.session import Endpoint as CorpusEndpoint +from setup.corpus.session import request_body as corpus_body + +Endpoint = Literal["chat_completions", "responses", "messages"] + + +class TokenSource(BaseModel): + """How the runner mints a JWT for a Keycloak user, so a run longer than the realm's access + token lifespan keeps authenticating as that user.""" + + token_url: str + client_id: str + username: str + password: str = Field(repr=False) + + +class WireShape(BaseModel): + """`RequestShape` on the wire: what the runner hands the corpus to build the body.""" + + model_config = ConfigDict(frozen=True) + + name: str + input_tokens: int + tools: bool + output_tokens: int + + +class WireRequest(BaseModel): + seq: int + marker: str + caller: str + endpoint: Endpoint + path: str + model: str + shape: WireShape + authorization: str | None = Field(default=None, repr=False) + token_source: TokenSource | None = None + stream: bool + blocked: bool + + +class WireRecord(BaseModel): + seq: int + marker: str + caller: str + endpoint: Endpoint + stream: bool + blocked: bool + started_at: float + status: int + header_s: float + first_byte_s: float | None + first_token_s: float | None + total_s: float + text: str + ok: bool + mock_delay_s: float | None + call_id: str | None + error: str | None + + +_CORPUS_ENDPOINT: Final[Mapping[Endpoint, CorpusEndpoint]] = MappingProxyType( + {"chat_completions": "chat", "responses": "responses", "messages": "messages"} +) + + +def request_body(item: WireRequest) -> dict[str, object]: + """The request in the endpoint's own wire format: the planned shape's coding-agent session, the marker in + the final user message, built right before sending so the plan file stays small.""" + return corpus_body( + endpoint=_CORPUS_ENDPOINT[item.endpoint], + model=item.model, + marker=item.marker, + input_tokens=item.shape.input_tokens, + tools=item.shape.tools, + output_tokens=item.shape.output_tokens, + stream=item.stream, + ) diff --git a/tests/e2e/perf/test_perf_cases.py b/tests/e2e/perf/test_perf_cases.py new file mode 100644 index 00000000000..a69c6bd63cc --- /dev/null +++ b/tests/e2e/perf/test_perf_cases.py @@ -0,0 +1,106 @@ +"""One perf run per `PerformanceTestCase` in `cases/`: check the gateway was deployed with the case's +model groups and guardrails, build its teams, users and keys, replay its deterministic plan from a separate Locust +process against the out-of-process mock providers, then check every `Limit` against what the +clients saw, what the mocks received and how the workers' memory moved.""" + +from __future__ import annotations + +import time +from typing import Final + +import pytest +from idp import Keycloak +from lifecycle import ResourceManager +from memory_readings import read_rss_everywhere +from models import SpendLogsParams +from cases import CASES +from setup.observe import ( + Observation, + evaluate, + read_guardrail_calls, + read_mock_requests, + render, + reset_mock, + set_mock_latency, +) +from setup.performance_test_case import PerformanceTestCase +from setup.plan import build_plan +from setup.datastore import read_db_counters, read_redis_counters +from setup.podcpu import gateway_cpu_sampler +from setup.provision import MOCK_PROVIDERS, build_setup, guardrail_url, mock_url +from setup.run import run_plan +from proxy_client import ProxyClient + +pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack] +_SPEND_SAMPLE: Final = 40 + + +def _spend_logged_share(proxy: ProxyClient, call_ids: tuple[str, ...]) -> float | None: + """Share of a sample of served requests whose spend row the gateway wrote, read back by the + request id the gateway returned in `x-litellm-call-id`.""" + if not call_ids: + return None + found = 0 + for call_id in call_ids: + rows = proxy.spend_logs(SpendLogsParams(request_id=call_id)) + found += 1 if rows else 0 + return found / len(call_ids) + + +@pytest.mark.parametrize("case", CASES, ids=[case.name for case in CASES]) +def test_perf_case(case: PerformanceTestCase, proxy: ProxyClient, idp: Keycloak, resources: ResourceManager) -> None: + plan: Final = build_plan(case) + mock_urls: Final = {provider: mock_url(provider) for provider in MOCK_PROVIDERS} + for url in (*mock_urls.values(), guardrail_url()): + _ = reset_mock(url) + for url in mock_urls.values(): + _ = set_mock_latency(url, case.mock_latency) + setup_started: Final = time.monotonic() + setup: Final = build_setup(case, proxy=proxy, idp=idp, resources=resources) + rss_before: Final = read_rss_everywhere(proxy) + db_before: Final = read_db_counters() + redis_before: Final = read_redis_counters() + run_started: Final = time.monotonic() + with gateway_cpu_sampler() as cpu: + outcome: Final = run_plan(case, plan, setup) + run_finished: Final = time.monotonic() + rss_after: Final = read_rss_everywhere(proxy) + rss_read: Final = time.monotonic() + assert outcome.records, f"locust wrote no records (exit {outcome.locust_exit}):\n{outcome.locust_output}" + served_call_ids: Final = tuple( + record.call_id for record in outcome.records if record.ok and not record.blocked and record.call_id + ) + spend_share: Final = _spend_logged_share( + proxy, served_call_ids[:: max(1, len(served_call_ids) // _SPEND_SAMPLE)][:_SPEND_SAMPLE] + ) + spend_read: Final = time.monotonic() + db_after: Final = read_db_counters() + redis_after: Final = read_redis_counters() + mock_requests: Final = tuple(seen for url in mock_urls.values() for seen in read_mock_requests(url)) + guardrail_calls: Final = tuple(read_guardrail_calls(guardrail_url())) if case.guardrails else () + mocks_read: Final = time.monotonic() + observation: Final = Observation( + records=outcome.records, + mock_requests=mock_requests, + guardrail_calls=guardrail_calls, + rss_before=rss_before, + rss_after=rss_after, + wall_seconds=outcome.wall_seconds, + spend_logged_share=spend_share, + db_before=db_before, + db_after=db_after, + redis_before=redis_before, + redis_after=redis_after, + gateway_cpu_seconds=cpu.seconds, + ) + rows, exceeded_limits = evaluate(case, observation) + print(render(case, observation, rows)) + print( + f" phases: setup {run_started - setup_started:.0f}s, locust {run_finished - run_started:.0f}s, " + f"observe {time.monotonic() - run_finished:.0f}s (rss {rss_read - run_finished:.0f}s, spend sample " + f"{spend_read - rss_read:.0f}s, mock logs {mocks_read - spend_read:.0f}s); locust exit {outcome.locust_exit}" + ) + assert len(outcome.records) == len(plan), ( + f"{len(plan)} planned, {len(outcome.records)} recorded; locust said:\n{outcome.locust_output}" + ) + assert not exceeded_limits, "limits exceeded (observed > max):\n" + "\n".join(exceeded_limits)