llm-simulator 0.2.1__py3-none-any.whl

llm_simulator/bench/__init__.py

@@ -0,0 +1,6 @@
+"""llmsim-bench: a load generator for OpenAI-compatible chat completion servers.
+
+It measures latency as the client sees it (the third measurement layer, after
+the engine's simulated timeline and the server's release times) and works
+against any OpenAI-compatible endpoint, not only the simulator.
+"""

llm_simulator/bench/cli.py

@@ -0,0 +1,304 @@
+"""The llmsim-bench command: arguments in, a run, a report out.
+
+The only module that reads arguments and the environment and wires the
+others together, as runtime.py does for the server.
+"""
+
+import argparse
+import asyncio
+import os
+import sys
+from collections import Counter
+from pathlib import Path
+from typing import Callable, List, NamedTuple, Optional, Sequence
+
+from llm_simulator.bench import workload as workloads
+from llm_simulator.bench.client import open_session, stream_chat
+from llm_simulator.bench.config import BenchConfig
+from llm_simulator.bench.records import RunRecords
+from llm_simulator.bench.health import LoopMonitor
+from llm_simulator.bench.report import (
+    format_progress,
+    format_result,
+    format_sweep,
+    write_json,
+    write_sweep_json,
+)
+from llm_simulator.bench.results import BenchResult, failure, summarize
+from llm_simulator.bench.runner import WarmupFailed, run_level
+from llm_simulator.bench.server import ServerMetrics, summary_url
+from llm_simulator.bench.sweep import levels, sweep
+from llm_simulator.bench.workload import Workload
+
+
+class Arguments(NamedTuple):
+    # The run's config; for a sweep, its first level.
+    config: BenchConfig
+    # Where to write the results file, if anywhere.
+    output: Optional[Path]
+    # The swept parameter and its values, when one was given several values.
+    sweep: Optional[str] = None
+    values: Sequence[float] = ()
+
+
+def _levels(kind: Callable[[str], float]) -> Callable[[str], List[float]]:
+    """An argparse type for one value or a comma-separated list of them."""
+
+    def parse(text: str) -> List[float]:
+        try:
+            return [kind(part) for part in text.split(",") if part.strip()]
+        except ValueError:
+            raise argparse.ArgumentTypeError(f"invalid list of numbers: {text!r}")
+
+    return parse
+
+
+def parse_args(argv: Optional[Sequence[str]] = None) -> Arguments:
+    defaults = BenchConfig(model="-")
+    parser = argparse.ArgumentParser(
+        prog="llmsim-bench",
+        description="Load-test an OpenAI-compatible chat completions server.",
+    )
+    parser.add_argument("--model", required=True)
+    parser.add_argument("--base-url", default=defaults.base_url)
+
+    load = parser.add_argument_group(
+        "load", "a list of values (1,8,32) runs a sweep, one level per value"
+    )
+    shape = load.add_mutually_exclusive_group()
+    shape.add_argument(
+        "--concurrency",
+        type=_levels(int),
+        help="closed loop: requests in flight at once (default: 1)",
+    )
+    shape.add_argument(
+        "--rate",
+        type=_levels(float),
+        help="open loop: requests per second, arriving on a Poisson schedule",
+    )
+    load.add_argument(
+        "--max-concurrency",
+        type=int,
+        help="open loop: at most this many in flight; later arrivals wait",
+    )
+    length = load.add_mutually_exclusive_group()
+    length.add_argument(
+        "--requests",
+        type=int,
+        help="requests to measure (default: 100)",
+    )
+    length.add_argument(
+        "--duration",
+        type=float,
+        help="seconds to measure for, instead of a number of requests",
+    )
+    load.add_argument(
+        "--warmup",
+        type=int,
+        help="unmeasured requests sent first (default: one per stream, or 8)",
+    )
+
+    workload = parser.add_argument_group(
+        "workload", "prompts come from --input-tokens, --dataset or --prompt"
+    )
+    workload.add_argument(
+        "--max-tokens",
+        type=int,
+        default=defaults.max_tokens,
+        help="max_completion_tokens of every request (a dataset line may override)",
+    )
+    source = workload.add_mutually_exclusive_group()
+    source.add_argument(
+        "--input-tokens",
+        type=int,
+        help="synthetic prompts of exactly this many prompt tokens on the simulator",
+    )
+    source.add_argument(
+        "--dataset",
+        type=Path,
+        help='JSONL file, one {"prompt": ..., "max_tokens": ...} per line',
+    )
+    workload.add_argument(
+        "--prompt", default=defaults.prompt, help="fixed prompt, used otherwise"
+    )
+    workload.add_argument(
+        "--prompt-prefix", default="", help='prepended to every prompt, e.g. "[short] "'
+    )
+    workload.add_argument(
+        "--seed", type=int, default=defaults.seed, help="seed for synthetic prompts"
+    )
+
+    connection = parser.add_argument_group("connection")
+    connection.add_argument(
+        "--api-key",
+        default=os.environ.get("OPENAI_API_KEY"),
+        help="sent as a Bearer token (default: $OPENAI_API_KEY)",
+    )
+    connection.add_argument(
+        "--timeout",
+        type=float,
+        default=defaults.timeout,
+        help="seconds to wait for the next chunk",
+    )
+
+    server = parser.add_argument_group(
+        "server comparison",
+        "compares with the target's /metrics/summary, when it has one",
+    )
+    server.add_argument(
+        "--no-server-metrics",
+        action="store_true",
+        help="do not reset or read the server's summary",
+    )
+    server.add_argument(
+        "--server-metrics-url",
+        help="the summary endpoint (default: the base URL with /v1 → /metrics/summary)",
+    )
+
+    parser.add_argument("--output", type=Path, help="write the results as JSON here")
+    args = parser.parse_args(argv)
+    parameter, values = None, ()
+    for name in ("concurrency", "rate"):
+        given = getattr(args, name)
+        if given is not None:
+            if not given:
+                parser.error(f"--{name} needs at least one value")
+            parameter, values = name, given
+    try:
+        config = BenchConfig(
+            model=args.model,
+            base_url=args.base_url,
+            concurrency=values[0] if parameter == "concurrency" else None,
+            rate=values[0] if parameter == "rate" else None,
+            max_concurrency=args.max_concurrency,
+            requests=args.requests,
+            duration=args.duration,
+            warmup=args.warmup,
+            max_tokens=args.max_tokens,
+            prompt=args.prompt,
+            input_tokens=args.input_tokens,
+            dataset=args.dataset,
+            prompt_prefix=args.prompt_prefix,
+            seed=args.seed,
+            api_key=args.api_key,
+            timeout=args.timeout,
+            server_metrics=not args.no_server_metrics,
+            server_metrics_url=args.server_metrics_url,
+        )
+        # Every level must be valid, not just the first.
+        for value in values[1:]:
+            levels(config, parameter, [value])[0]
+    except ValueError as exc:
+        parser.error(str(exc))
+    if len(values) > 1:
+        return Arguments(config, args.output, parameter, values)
+    return Arguments(config, args.output)
+
+
+def build_workload(config: BenchConfig) -> Workload:
+    """The workload the config asks for; raises ValueError or OSError if it
+    cannot be built (a too-small input_tokens, an unreadable dataset)."""
+    prefix = config.prompt_prefix
+    if config.input_tokens is not None:
+        return workloads.synthetic(
+            config.input_tokens, config.max_tokens, prefix=prefix, seed=config.seed
+        )
+    if config.dataset is not None:
+        return workloads.from_jsonl(
+            config.dataset, config.max_tokens, prefix, seed=config.seed
+        )
+    return workloads.fixed(config.prompt, config.max_tokens, prefix)
+
+
+async def run(config: BenchConfig, workload: Workload) -> RunRecords:
+    """Warm up, measure and return the measured records with their window.
+
+    Raises WarmupFailed if no warm-up request succeeded.
+    """
+    async with open_session(config) as session:
+        server = None
+        if config.server_metrics:
+            url = config.server_metrics_url or summary_url(config.base_url)
+            server = ServerMetrics(session, url, config.model)
+
+        def send(index: int):
+            return stream_chat(session, config, index, workload.request(index))
+
+        closed = config.mode == "closed"
+        return await run_level(
+            send,
+            concurrency=config.closed_concurrency if closed else None,
+            rate=config.rate,
+            max_concurrency=config.max_concurrency,
+            seed=config.seed,
+            requests=config.measured_requests,
+            duration=config.duration,
+            warmup=config.warmup_requests,
+            warmup_concurrency=config.warmup_concurrency,
+            monitor=LoopMonitor(),
+            server=server,
+        )
+
+
+async def run_result(config: BenchConfig, workload: Workload) -> BenchResult:
+    return summarize(await run(config, workload), config, workload.description)
+
+
+def main(argv: Optional[Sequence[str]] = None) -> int:
+    args = parse_args(argv)
+    config = args.config
+    try:
+        workload = build_workload(config)
+    except (ValueError, OSError) as exc:
+        print(f"llmsim-bench: error: {exc}", file=sys.stderr)
+        return 2
+    if args.sweep is not None:
+        return _main_sweep(args, workload)
+    try:
+        result = asyncio.run(run_result(config, workload))
+    except KeyboardInterrupt:
+        print("interrupted", file=sys.stderr)
+        return 130
+    except WarmupFailed as exc:
+        print(f"llmsim-bench: {exc}; nothing was measured", file=sys.stderr)
+        for message, count in Counter(map(failure, exc.records)).most_common(5):
+            print(f"  {count} × {message}", file=sys.stderr)
+        return 1
+
+    print(format_result(result))
+    if args.output is not None:
+        write_json(result, args.output)
+        print(f"\nresults written to {args.output}")
+    # Failing outright (for example, no server) is an error; some failed
+    # requests are a result.
+    return 0 if result.ok else 1
+
+
+def _main_sweep(args: Arguments, workload: Workload) -> int:
+    configs = levels(args.config, args.sweep, args.values)
+
+    def progress(number, level):
+        print(format_progress(number, len(configs), level), flush=True)
+
+    try:
+        done = asyncio.run(
+            sweep(
+                args.sweep,
+                configs,
+                lambda config: run_result(config, workload),
+                progress,
+            )
+        )
+    except KeyboardInterrupt:
+        print("interrupted", file=sys.stderr)
+        return 130
+    print()
+    print(format_sweep(done, workload.description))
+    if args.output is not None:
+        write_sweep_json(done, args.output)
+        print(f"\nresults written to {args.output}")
+    return 0 if any(level.result and level.result.ok for level in done) else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

llm_simulator/bench/client.py

@@ -0,0 +1,168 @@
+"""One streaming chat completion request, timed as the client sees it.
+
+The only module, with server.py later, that touches the network. It speaks
+the OpenAI streaming protocol directly instead of using the OpenAI SDK, which
+builds a typed object for every chunk: the client's own overhead must stay
+far below the server's, or the benchmark measures the client.
+"""
+
+import asyncio
+import json
+import time
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, Optional
+
+import aiohttp
+
+from llm_simulator.bench.config import BenchConfig
+from llm_simulator.bench.records import RequestRecord
+from llm_simulator.bench.workload import RequestSpec
+
+# Waiting for a connection is not the same as waiting for tokens.
+CONNECT_TIMEOUT = 10.0
+
+
+def build_body(config: BenchConfig, spec: RequestSpec) -> bytes:
+    """The request body, serialised before the clock starts."""
+    return json.dumps(
+        {
+            "model": config.model,
+            "messages": [{"role": "user", "content": spec.prompt}],
+            "stream": True,
+            "stream_options": {"include_usage": True},
+            "max_completion_tokens": spec.max_tokens,
+        }
+    ).encode()
+
+
+@dataclass
+class Event:
+    """What one SSE data line carried that the client cares about."""
+
+    # Non-empty delta.content only: the opening role chunk carries "".
+    content: Optional[str] = None
+    finish_reason: Optional[str] = None
+    usage: Optional[Dict[str, Any]] = None
+    error: Optional[str] = None
+    done: bool = False
+
+
+def parse_line(line: bytes) -> Optional[Event]:
+    """Parse one line of the stream; None for lines without data.
+
+    Assumes one JSON object per "data:" line, as OpenAI, vLLM and the
+    simulator all send. Raises ValueError for data that is not valid JSON.
+    """
+    if not line.startswith(b"data:"):
+        # Blank separators, ":" comments and other SSE fields.
+        return None
+    data = line[5:].strip()
+    if data == b"[DONE]":
+        return Event(done=True)
+    chunk = json.loads(data)
+
+    error = chunk.get("error")
+    if error is not None:
+        message = error.get("message") if isinstance(error, dict) else None
+        return Event(error=message or str(error))
+    if chunk.get("object") == "error":  # vLLM's in-stream error shape
+        return Event(error=chunk.get("message") or str(chunk))
+
+    event = Event(usage=chunk.get("usage"))
+    choices = chunk.get("choices")
+    if choices:
+        choice = choices[0]
+        event.content = (choice.get("delta") or {}).get("content") or None
+        event.finish_reason = choice.get("finish_reason")
+    return event
+
+
+def open_session(config: BenchConfig) -> aiohttp.ClientSession:
+    """A session sized for the run; call it inside the event loop."""
+    headers = {"Content-Type": "application/json"}
+    if config.api_key:
+        headers["Authorization"] = f"Bearer {config.api_key}"
+    return aiohttp.ClientSession(
+        # aiohttp's default limit of 100 connections would silently cap the
+        # run at 100 concurrent streams; 0 is no limit, for an uncapped open
+        # loop.
+        connector=aiohttp.TCPConnector(limit=config.streams),
+        # No total limit, since long generations are expected; a stream that
+        # stalls between chunks fails.
+        timeout=aiohttp.ClientTimeout(
+            total=None, sock_connect=CONNECT_TIMEOUT, sock_read=config.timeout
+        ),
+        headers=headers,
+    )
+
+
+async def stream_chat(
+    session: aiohttp.ClientSession,
+    config: BenchConfig,
+    index: int,
+    spec: RequestSpec,
+    clock: Callable[[], float] = time.perf_counter,
+) -> RequestRecord:
+    """Send one streaming request and record when each part arrived.
+
+    Never raises for a failed request: HTTP errors, broken streams, timeouts
+    and connection errors become a record with status and error set, so one
+    failure does not stop a run. Cancellation still propagates.
+    """
+    body = build_body(config, spec)
+    sent = clock()
+    record = RequestRecord(
+        index=index, sent=sent, end=sent, requested_tokens=spec.max_tokens
+    )
+    try:
+        async with session.post(config.chat_url, data=body) as response:
+            record.status = response.status
+            if response.status != 200:
+                record.error = await _error_message(response)
+                record.end = clock()
+                return record
+
+            async for line in response.content:
+                now = clock()
+                event = parse_line(line)
+                if event is None:
+                    continue
+                if event.done:
+                    record.end = now
+                    return record
+                if event.error is not None:
+                    record.error = event.error
+                if event.content is not None:
+                    if record.last_token is None:
+                        record.first_token = now
+                    else:
+                        record.itls.append(now - record.last_token)
+                    record.last_token = now
+                    record.chunks += 1
+                if event.finish_reason is not None:
+                    record.finish_reason = event.finish_reason
+                if event.usage is not None:
+                    record.prompt_tokens = event.usage.get("prompt_tokens")
+                    record.completion_tokens = event.usage.get("completion_tokens")
+            record.error = record.error or "stream ended before [DONE]"
+    except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
+        record.error = _describe(exc)
+    except ValueError as exc:
+        record.error = f"invalid stream data: {exc}"
+    record.end = clock()
+    return record
+
+
+async def _error_message(response: aiohttp.ClientResponse) -> str:
+    """The message of an OpenAI-style error body, or the start of the body."""
+    text = await response.text()
+    try:
+        return json.loads(text)["error"]["message"]
+    except (ValueError, KeyError, TypeError):
+        return f"HTTP {response.status}: {text[:200]}"
+
+
+def _describe(exc: BaseException) -> str:
+    # Timeouts often have an empty message.
+    message = str(exc)
+    return f"{type(exc).__name__}: {message}" if message else type(exc).__name__

llm_simulator/bench/config.py

@@ -0,0 +1,143 @@
+"""What a benchmark run sends, and how much of it."""
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+# Open-loop warm-up without a concurrency cap to size it by.
+DEFAULT_OPEN_WARMUP = 8
+
+
+@dataclass(frozen=True)
+class BenchConfig:
+    """Settings for one benchmark run.
+
+    Plain data: the CLI reads arguments and the environment, everything else
+    takes a BenchConfig.
+    """
+
+    model: str
+    base_url: str = "http://localhost:8080/v1"
+    # How load is generated. Closed loop: `concurrency` requests in flight,
+    # each worker sending its next as soon as its previous one finishes
+    # (the default, with 1). Open loop: requests arrive at `rate` per second
+    # on a Poisson schedule, whether or not earlier ones have finished;
+    # `max_concurrency` optionally caps how many are in flight, and an
+    # arrival that finds it full waits for a slot.
+    concurrency: Optional[int] = None
+    rate: Optional[float] = None
+    max_concurrency: Optional[int] = None
+    # How long the measure phase runs: a number of requests, or a number of
+    # seconds (requests sent before the deadline are counted and waited for).
+    # With neither, 100 requests.
+    requests: Optional[int] = None
+    duration: Optional[float] = None
+    # Requests sent first and not measured, so connections are open and code
+    # paths warm; None means one per concurrent stream.
+    warmup: Optional[int] = None
+    # max_completion_tokens of every request; a dataset line may set its own.
+    max_tokens: int = 256
+    # Where prompts come from: synthetic prompts of input_tokens prompt
+    # tokens, the lines of a JSONL dataset, or else this fixed prompt.
+    prompt: str = "Benchmark request."
+    input_tokens: Optional[int] = None
+    dataset: Optional[Path] = None
+    # Prepended to every prompt, e.g. "[short] " for bench-responses.yml.
+    prompt_prefix: str = ""
+    # Seeds the synthetic prompts and, separately, the arrival schedule, so
+    # runs send the same prompts at the same times.
+    seed: int = 0
+    # Sent as "Authorization: Bearer <api_key>" when set.
+    api_key: Optional[str] = None
+    # Seconds to wait for the next chunk, not for the whole request: long
+    # generations are expected, a stalled stream is not.
+    timeout: float = 60.0
+    # Compare with the target's /metrics/summary, when it has one. The URL
+    # defaults to the base URL with its /v1 replaced by /metrics/summary.
+    server_metrics: bool = True
+    server_metrics_url: Optional[str] = None
+
+    def __post_init__(self):
+        if not self.model:
+            raise ValueError("model must not be empty")
+        if not self.base_url.startswith(("http://", "https://")):
+            raise ValueError(
+                f"base_url must start with http:// or https://, got {self.base_url!r}"
+            )
+        if self.max_tokens < 1:
+            raise ValueError("max_tokens must be at least 1")
+        if self.concurrency is not None and self.rate is not None:
+            raise ValueError(
+                "use concurrency (closed loop) or rate (open loop), not both"
+            )
+        if self.concurrency is not None and self.concurrency < 1:
+            raise ValueError("concurrency must be at least 1")
+        if self.rate is not None and self.rate <= 0:
+            raise ValueError("rate must be positive")
+        if self.max_concurrency is not None:
+            if self.rate is None:
+                raise ValueError("max_concurrency applies to the open loop: set rate")
+            if self.max_concurrency < 1:
+                raise ValueError("max_concurrency must be at least 1")
+        if self.requests is not None and self.duration is not None:
+            raise ValueError("use requests or duration, not both")
+        if self.requests is not None and self.requests < 1:
+            raise ValueError("requests must be at least 1")
+        if self.duration is not None and self.duration <= 0:
+            raise ValueError("duration must be positive")
+        if self.warmup is not None and self.warmup < 0:
+            raise ValueError("warmup must not be negative")
+        if (
+            self.server_metrics_url is not None
+            and not self.server_metrics_url.startswith(("http://", "https://"))
+        ):
+            raise ValueError(
+                "server_metrics_url must start with http:// or https://, got"
+                f" {self.server_metrics_url!r}"
+            )
+        if self.timeout <= 0:
+            raise ValueError("timeout must be positive")
+        if self.input_tokens is not None and self.dataset is not None:
+            raise ValueError("use input_tokens or dataset, not both")
+        if self.input_tokens is not None and self.input_tokens < 1:
+            raise ValueError("input_tokens must be at least 1")
+
+    @property
+    def measured_requests(self) -> Optional[int]:
+        """Requests to measure; None when the run is bounded by duration."""
+        if self.duration is not None:
+            return None
+        return 100 if self.requests is None else self.requests
+
+    @property
+    def mode(self) -> str:
+        return "open" if self.rate is not None else "closed"
+
+    @property
+    def closed_concurrency(self) -> int:
+        return 1 if self.concurrency is None else self.concurrency
+
+    @property
+    def streams(self) -> int:
+        """Connections the client may open; 0 means no limit (open loop, no cap)."""
+        if self.mode == "closed":
+            return self.closed_concurrency
+        return self.max_concurrency or 0
+
+    @property
+    def warmup_requests(self) -> int:
+        if self.warmup is not None:
+            return self.warmup
+        if self.mode == "closed":
+            return self.closed_concurrency
+        return self.max_concurrency or DEFAULT_OPEN_WARMUP
+
+    @property
+    def warmup_concurrency(self) -> int:
+        """Warm-up is always a closed-loop burst, at most one per stream."""
+        streams = self.streams or self.warmup_requests
+        return max(1, min(self.warmup_requests, streams))
+
+    @property
+    def chat_url(self) -> str:
+        return self.base_url.rstrip("/") + "/chat/completions"

llm_simulator/bench/health.py

@@ -0,0 +1,59 @@
+"""Watching the client itself, so its own limits are not blamed on the server.
+
+The client takes every timestamp on its event loop. If the loop is busy, a
+chunk that arrived on time is read late, and the latency it records is too
+high. The monitor asks to wake at a fixed interval and records how late each
+wake-up was: a late wake-up means timestamps taken then were late by about
+as much.
+"""
+
+import asyncio
+import time
+from array import array
+from contextlib import suppress
+from typing import Awaitable, Callable, Optional
+
+from llm_simulator.bench.records import LoopSamples
+
+Clock = Callable[[], float]
+
+
+class LoopMonitor:
+    """Samples event-loop lateness and process CPU between start() and stop()."""
+
+    def __init__(
+        self,
+        interval: float = 0.010,
+        *,
+        clock: Clock = time.perf_counter,
+        cpu_clock: Clock = time.process_time,
+        sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
+    ):
+        self._interval = interval
+        self._clock = clock
+        self._cpu_clock = cpu_clock
+        self._sleep = sleep
+        self._task: Optional[asyncio.Task] = None
+
+    def start(self) -> None:
+        """Begin sampling; call from inside the running event loop."""
+        self._lags = array("d")
+        self._cpu_started = self._cpu_clock()
+        self._wall_started = self._clock()
+        self._task = asyncio.ensure_future(self._watch())
+
+    async def stop(self) -> LoopSamples:
+        self._task.cancel()
+        with suppress(asyncio.CancelledError):
+            await self._task
+        return LoopSamples(
+            lags=self._lags,
+            cpu_seconds=self._cpu_clock() - self._cpu_started,
+            wall_seconds=self._clock() - self._wall_started,
+        )
+
+    async def _watch(self) -> None:
+        while True:
+            before = self._clock()
+            await self._sleep(self._interval)
+            self._lags.append(max(0.0, self._clock() - before - self._interval))