"""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 import time from collections.abc import Callable from dataclasses import dataclass from pathlib import Path from typing import BinaryIO, 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, send_at_s=item.send_at_s, ) 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") class RecordStream: def __init__(self, source: BinaryIO) -> None: self.source: Final = source self.partial = b"" def read(self) -> tuple[WireRecord, ...]: chunks: Final = (self.partial + self.source.read()).split(b"\n") self.partial = chunks[-1] return tuple(_RECORD_ADAPTER.validate_json(line) for line in chunks[:-1] if line) def run_until( command: list[str], env: dict[str, str], records_path: Path, stop_path: Path, stop_when: Callable[[tuple[WireRecord, ...]], bool], timeout: int, ) -> subprocess.CompletedProcess[str]: with records_path.open("rb") as records, records_path.with_suffix(".log").open("w+") as output: with subprocess.Popen(command, env=env, stdout=output, stderr=subprocess.STDOUT, text=True) as process: deadline: Final = time.monotonic() + timeout reader: Final = RecordStream(records) try: while process.poll() is None: if stop_when(reader.read()): stop_path.touch() if time.monotonic() >= deadline: raise subprocess.TimeoutExpired(command, timeout) try: _ = process.wait(timeout=1) except subprocess.TimeoutExpired: pass finally: stop_path.touch() try: _ = process.wait(timeout=35) except subprocess.TimeoutExpired: process.kill() _ = process.wait(timeout=10) _keep_records_artifact(records_path) _ = output.seek(0) return subprocess.CompletedProcess(command, process.returncode, output.read(), "") def run_plan( case: PerformanceTestCase, plan: tuple[PlannedRequest, ...], setup: CaseSetup, *, stop_when: Callable[[tuple[WireRecord, ...]], bool] | None = None, request_timeout: float = 120.0, ) -> 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" stop_path: Final = Path(directory) / "stop" _ = 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), "PERF_REQUEST_TIMEOUT": str(request_timeout), **({"PERF_STOP_FILE": str(stop_path)} if stop_when is not None else {}), } command: Final = [ 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", ] completed: Final = ( run_until(command, env, records_path, stop_path, stop_when, ceiling_s + 120) if stop_when is not None else subprocess.run(command, 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)