litellm/tests/e2e/perf/setup/locustfile.py

"""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 random
import time
from collections.abc import Iterator
from pathlib import Path
from typing import Final, cast

import gevent
from locust import FastHttpUser, task
from locust.contrib.fasthttp import FastResponse, ResponseContextManager
from pydantic import TypeAdapter

from setup.wire import Endpoint, PlanProgress, WireRecord, WireRequest, request_body, response_id

_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)
_INTERVAL_S: Final = _USERS / _RPS
_JITTER: Final = 0.25
_RNG: Final = random.Random(_PLAN[0].marker if _PLAN else "empty-plan")
_RECORDS: Final = _RECORDS_PATH.open("a", encoding="utf-8")
_PROGRESS: Final = PlanProgress(len(_PLAN))
_STOP_FILE: Final = os.environ.get("PERF_STOP_FILE")
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 ""


_RESPONSES_TOKEN_EVENTS: Final = frozenset(
    {
        "response.output_text.delta",
        "response.refusal.delta",
        "response.reasoning_text.delta",
        "response.reasoning_summary_text.delta",
        "response.function_call_arguments.delta",
    }
)


def token_in_event(endpoint: Endpoint, payload: dict[str, object]) -> bool:
    """Whether one SSE event carries a generated token of any kind: visible text, reasoning or tool-call arguments.
    The first such event is the stream's first token. `delta_of_event` stays visible text only because the body
    check compares it with the text the mock generated, and a tool-call reply has none."""
    if delta_of_event(endpoint, payload):
        return True
    if endpoint == "chat_completions":
        choices: Final = _as_list(payload.get("choices"))
        delta: Final = _as_dict(_as_dict(choices[0]).get("delta")) if choices else {}
        return any(delta.get(key) for key in ("reasoning_content", "reasoning", "tool_calls", "function_call"))
    if endpoint == "responses":
        return payload.get("type") in _RESPONSES_TOKEN_EVENTS
    return payload.get("type") == "content_block_delta"


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:
    _ = _RECORDS.write(record.model_dump_json() + "\n")
    _RECORDS.flush()
    _PROGRESS.completed += 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. Each user starts at
    a random point of its first interval and jitters every later interval by up to `_JITTER` either way, so the
    users never fire in lockstep; the mean interval is unchanged and `rps` stays the total rate."""

    network_timeout = float(os.environ.get("PERF_REQUEST_TIMEOUT", "120"))
    connection_timeout = 30.0

    def on_start(self) -> None:
        self._last_run = 0.0
        self._last_wait = 0.0
        gevent.sleep(_RNG.uniform(0.0, _INTERVAL_S))

    def wait_time(self) -> float:
        target: Final = _INTERVAL_S * _RNG.uniform(1.0 - _JITTER, 1.0 + _JITTER)
        now: Final = time.monotonic()
        run_s: Final = now - self._last_run - self._last_wait if self._last_run else 0.0
        self._last_wait = max(0.0, target - run_s)
        self._last_run = now
        return self._last_wait

    @task
    def send(self) -> None:
        take: Final = _PROGRESS.take(stop=_STOP_FILE is not None and Path(_STOP_FILE).exists())
        item: Final = next(_NEXT, None) if take else None
        if item is None:
            if _PROGRESS.drained:
                _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
        identifier: str | None = None  # rebind-ok: the ID arrives in a JSON body or the first SSE event
        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:
                            identifier = identifier or response_id(event)
                            if first_token_s is None and token_in_event(item.endpoint, event):
                                first_token_s = now
                            text += delta_of_event(item.endpoint, event)
                else:
                    body_text = context.text or ""
                    first_byte_s = time.perf_counter() - started
                    if status == 200:
                        payload: Final = _as_dict(_JSON.validate_json(body_text))
                        identifier = response_id(payload)
                        text = text_of_json(item.endpoint, payload)
                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,
                response_id=identifier,
            )
        )
        if _PROGRESS.drained:
            _quit(self)