benchmaker 0.1.0__py3-none-any.whl

benchmaker/trace.py

@@ -0,0 +1,217 @@
+"""Record and replay of request traces for deterministic reproduction.
+
+A *trace* is a JSONL file where each line captures one fired request:
+
+    {"t_rel": 0.123, "method": "POST", "url": "...", "headers": {...},
+     "params": {...}, "json": {...}, "timeout_s": null, "meta": {...}}
+
+`t_rel` is the time (seconds) between bench start and the moment the request
+went out — captured *after* pre-hooks, *before* transport. Binary bodies are
+stored under `body_b64` (base64). The recorder is a runner-level feature; the
+replay side composes through the standard `workload_type + workload + load`
+triple:
+
+  * `ReplayWorkloadType` — `make_request` rebuilds a `Request` from a row.
+  * `TraceWorkload`      — yields rows in order.
+  * `TracePacedLoad`     — emits one ticket at each row's `t_rel`
+    (optionally sped up / slowed down by `speed`).
+
+`build_config` wires all three from a single `replay:` block, so reproducing a
+run is just `replay: {path: ..., speed: 1.0}` plus the same `correctness:` (if
+grading is desired).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import json
+import os
+import time
+from typing import Any, Optional
+
+from benchmaker.load import LoadModel
+from benchmaker.types import Request
+from benchmaker.workloads.base import WorkloadType
+from benchmaker.workloads.datasets import Workload
+
+
+# --------------------------------------------------------------------------- #
+# Wire format
+# --------------------------------------------------------------------------- #
+
+
+def _safe_meta(meta: dict) -> dict:
+    out: dict[str, Any] = {}
+    for k, v in meta.items():
+        try:
+            json.dumps(v)
+            out[k] = v
+        except (TypeError, ValueError):
+            out[k] = repr(v)
+    return out
+
+
+def request_to_row(t_rel: float, req: Request) -> dict:
+    """Encode a `Request` as a JSON-serializable trace row."""
+    row: dict[str, Any] = {
+        "t_rel": float(t_rel),
+        "method": req.method,
+        "url": req.url,
+        "headers": dict(req.headers or {}),
+        "params": dict(req.params or {}),
+        "timeout_s": req.timeout_s,
+        "meta": _safe_meta(req.meta or {}),
+    }
+    if req.json is not None:
+        row["json"] = req.json
+    elif req.body is not None:
+        row["body_b64"] = base64.b64encode(req.body).decode("ascii")
+    return row
+
+
+def row_to_request(row: dict) -> Request:
+    """Rebuild a `Request` from a trace row."""
+    body: Optional[bytes] = None
+    if "body_b64" in row:
+        body = base64.b64decode(row["body_b64"])
+    return Request(
+        method=row.get("method", "GET"),
+        url=row.get("url", ""),
+        headers=dict(row.get("headers") or {}),
+        params=dict(row.get("params") or {}),
+        body=body,
+        json=row.get("json"),
+        timeout_s=row.get("timeout_s"),
+        meta=dict(row.get("meta") or {}),
+    )
+
+
+def load_trace(path: str) -> list[dict]:
+    """Load a trace JSONL into a list of rows, sorted by `t_rel`."""
+    rows: list[dict] = []
+    with open(path) as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            rows.append(json.loads(line))
+    rows.sort(key=lambda r: float(r.get("t_rel", 0.0)))
+    return rows
+
+
+# --------------------------------------------------------------------------- #
+# Recorder
+# --------------------------------------------------------------------------- #
+
+
+class TraceRecorder:
+    """Append `(t_rel, request)` rows to a JSONL file as the bench runs.
+
+    The runner calls `open()` once at the start, `record(req, fire_mono)` per
+    fired request, and `close()` once at the end. `t_rel` is computed as
+    `fire_mono - start_mono`; `start_mono` is captured by `open()` (so it
+    matches whatever clock the runner's progress logger uses).
+    """
+
+    def __init__(self, path: str):
+        self._path = path
+        self._f = None
+        self._lock = asyncio.Lock()
+        self._start_mono: Optional[float] = None
+
+    def open(self, start_mono: Optional[float] = None) -> None:
+        if self._f is not None:
+            return
+        d = os.path.dirname(os.path.abspath(self._path))
+        if d:
+            os.makedirs(d, exist_ok=True)
+        self._f = open(self._path, "w")
+        self._start_mono = start_mono if start_mono is not None else time.monotonic()
+
+    async def record(self, req: Request, fire_mono: float) -> None:
+        if self._f is None:
+            self.open()
+        assert self._start_mono is not None
+        row = request_to_row(fire_mono - self._start_mono, req)
+        line = json.dumps(row, default=str)
+        async with self._lock:
+            self._f.write(line + "\n")
+            self._f.flush()
+
+    def close(self) -> None:
+        if self._f is not None:
+            try:
+                self._f.close()
+            finally:
+                self._f = None
+
+
+# --------------------------------------------------------------------------- #
+# Replay: workload + workload-type + load model
+# --------------------------------------------------------------------------- #
+
+
+class TraceWorkload(Workload):
+    """Yield trace rows (dicts) in order. Halts when the trace is exhausted."""
+
+    def __init__(self, trace: list[dict], name: str = "trace"):
+        self.name = name
+        self._trace = list(trace)
+        self._i = 0
+
+    async def next_item(self) -> Any:
+        if self._i >= len(self._trace):
+            raise StopAsyncIteration
+        row = self._trace[self._i]
+        self._i += 1
+        return row
+
+
+class ReplayWorkloadType(WorkloadType):
+    """Rebuild a `Request` directly from a trace row.
+
+    `streaming` is set per-instance — match the workload-type used during
+    recording so the runner reads chunks the same way (otherwise stream-only
+    metrics like TTFT won't be reproduced).
+    """
+
+    # Recorded Requests already carry their `reference` (and any other meta the
+    # original workload put there) under `meta`. So apply_correctness should
+    # install only the post-hook — no EvalWorkloadType item-splitting wrapper.
+    handles_reference = True
+
+    def __init__(self, streaming: bool = False, name: str = "replay"):
+        self.name = name
+        self.streaming = streaming
+
+    async def make_request(self, item: Any) -> Request:
+        if not isinstance(item, dict):
+            raise TypeError(
+                f"ReplayWorkloadType expects trace rows (dict), got {type(item).__name__}"
+            )
+        return row_to_request(item)
+
+
+class TracePacedLoad(LoadModel):
+    """Open-loop load model that fires one ticket per row at the recorded time.
+
+    `speed` rescales the trace: `2.0` halves all inter-arrival gaps; `0.5`
+    doubles them. Use `1.0` for an as-recorded replay.
+    """
+
+    def __init__(self, trace: list[dict], speed: float = 1.0):
+        if speed <= 0:
+            raise ValueError("speed must be > 0")
+        self._times = [float(r.get("t_rel", 0.0)) for r in trace]
+        self._times.sort()
+        self._speed = float(speed)
+
+    async def tickets(self):
+        start = time.monotonic()
+        for t_rel in self._times:
+            target = start + t_rel / self._speed
+            now = time.monotonic()
+            if target > now:
+                await asyncio.sleep(target - now)
+            yield None

benchmaker/types.py

@@ -0,0 +1,98 @@
+"""Core data types and hook protocols."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Awaitable, Callable, Mapping, Optional, Protocol
+
+
+@dataclass
+class Request:
+    """An outgoing HTTP request.
+
+    `meta` is opaque user-controlled state that travels with the request and is
+    available to post-response hooks (useful for correlating prompts to results,
+    tagging variants, etc.).
+    """
+
+    method: str = "GET"
+    url: str = ""
+    headers: dict[str, str] = field(default_factory=dict)
+    params: dict[str, Any] = field(default_factory=dict)
+    body: Optional[bytes] = None
+    json: Any = None
+    timeout_s: Optional[float] = None
+    meta: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class Response:
+    status: int
+    headers: Mapping[str, str]
+    body: bytes
+    elapsed_s: float
+    ok: bool
+    error: Optional[str] = None
+    # For streaming responses, the chunks captured (bytes) and per-chunk arrival times.
+    stream_chunks: Optional[list[bytes]] = None
+    stream_chunk_times: Optional[list[float]] = None  # seconds since request start
+
+
+@dataclass
+class Sample:
+    """One observation. Workloads may attach extra numeric metrics via `extra`."""
+
+    start_ts: float                # wall-clock start (time.monotonic)
+    latency_s: float               # total request latency
+    status: int                    # HTTP status (0 if connection failure)
+    ok: bool                       # whether the workload considers it a success
+    # Whether the request itself was delivered successfully at the transport/HTTP
+    # layer. Stays True for responses that came back 2xx/3xx even if a post-hook
+    # later flips `ok` to False (e.g. eval scorer judged the output wrong).
+    # Use this to distinguish "request failed" from "request OK but output wrong".
+    request_ok: bool = True
+    bytes_sent: int = 0
+    bytes_recv: int = 0
+    error: Optional[str] = None
+    workload: Optional[str] = None
+    meta: dict[str, Any] = field(default_factory=dict)
+    extra: dict[str, float] = field(default_factory=dict)
+
+
+# Hook signatures. Hooks may be sync or async; the runner awaits coroutines and
+# calls plain callables directly.
+PreRequestHook = Callable[[Request], "Request | Awaitable[Request]"]
+PostResponseHook = Callable[
+    [Request, Response, Sample],
+    "Sample | Awaitable[Sample]",
+]
+
+# A `fire(req)` callable closes over the aiohttp session, applies pre-hooks,
+# times the call, and returns a Response. A WorkloadType.run_ticket implementation
+# may call this 0..N times per ticket.
+FireRequest = Callable[[Request], Awaitable[Response]]
+
+
+@dataclass
+class TicketContext:
+    """Everything a WorkloadType needs to handle a single ticket.
+
+    The runner builds one of these per ticket and hands it to
+    `WorkloadType.run_ticket`. Multi-step workload-types (e.g. sandbox lifecycle)
+    can call `fire(req)` repeatedly; the default base-class implementation calls
+    it exactly once.
+    """
+
+    item: Any
+    start_mono: float
+    fire: FireRequest
+    pre_hooks: tuple = ()
+    post_hooks: tuple = ()
+    workload_name: str = "workload-type"
+
+
+async def maybe_await(value: Any) -> Any:
+    """Await `value` if awaitable, else return it. Lets hooks be sync OR async."""
+    if hasattr(value, "__await__"):
+        return await value
+    return value

benchmaker/workloads/__init__.py

@@ -0,0 +1,53 @@
+from benchmaker.workloads.base import WorkloadType
+from benchmaker.workloads.datasets import (
+    Workload,
+    StaticWorkload,
+    JsonlWorkload,
+    CallableWorkload,
+    IterableWorkload,
+)
+from benchmaker.workloads.http import HttpWorkloadType
+from benchmaker.workloads.llm import OpenAIChatWorkloadType
+from benchmaker.workloads.sandbox import SandboxWorkloadType
+from benchmaker.workloads.hf import HFDatasetWorkload
+from benchmaker.workloads.eval import (
+    EvalWorkloadType,
+    Scorer,
+    correctness_hook,
+    extract_openai_text,
+    extract_raw_text,
+    extract_text,
+    exact_match,
+    contains,
+    regex_match,
+    json_valid,
+    multiple_choice,
+    judge_llm,
+    openai_chat_judge,
+)
+
+__all__ = [
+    "WorkloadType",
+    "Workload",
+    "StaticWorkload",
+    "JsonlWorkload",
+    "CallableWorkload",
+    "IterableWorkload",
+    "HttpWorkloadType",
+    "OpenAIChatWorkloadType",
+    "SandboxWorkloadType",
+    "HFDatasetWorkload",
+    "EvalWorkloadType",
+    "Scorer",
+    "correctness_hook",
+    "extract_openai_text",
+    "extract_raw_text",
+    "extract_text",
+    "exact_match",
+    "contains",
+    "regex_match",
+    "json_valid",
+    "multiple_choice",
+    "judge_llm",
+    "openai_chat_judge",
+]

benchmaker/workloads/agent.py

@@ -0,0 +1,308 @@
+"""User-defined Agent workload-type.
+
+Drives an arbitrary Python class — typically a multi-turn agent that makes its
+own HTTP calls (to a model API, tools, a sandbox, …) — and grades the final
+output via the standard `correctness_hook`.
+
+Compared to the LLM workload-types:
+
+  * No request shape is dictated. The agent owns its own client(s).
+  * One "request" = one full agent run, which may be many internal calls.
+  * Per-trajectory metrics (steps, tool calls, tokens, retries, …) ride on
+    `Sample.extra` via `AgentResult.metrics`.
+
+The expected item shape mirrors the LLM workloads: a dict with at least a
+``prompt`` field (passed to the agent as the task), an optional ``reference``
+(consumed by `correctness_hook`), and any extra columns the dataset carries.
+Items are passed through to the agent verbatim — the workload-type only peels
+``reference`` (and any ``extra_meta_keys``) into ``Sample.meta`` so the grader
+can find them.
+
+Example (subclassing):
+
+    from benchmaker import Agent, AgentContext, AgentResult
+
+    class MyAgent(Agent):
+        def __init__(self, model: str):
+            self.model = model
+
+        async def run(self, ctx: AgentContext) -> AgentResult:
+            task = ctx.item["prompt"]
+            output, steps = await my_pipeline(task, model=self.model)
+            return AgentResult(
+                output=output,
+                ok=True,
+                metrics={"steps": float(steps)},
+            )
+
+YAML:
+
+    workload_type:
+      type: agent
+      agent: 'mypkg.myagent:MyAgent'
+      agent_kwargs:
+        model: 'gpt-4o-mini'
+"""
+
+from __future__ import annotations
+
+import inspect
+import time
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Callable, Optional, Union
+
+from benchmaker.types import (
+    FireRequest,
+    Request,
+    Response,
+    Sample,
+    TicketContext,
+    maybe_await,
+)
+from benchmaker.workloads.base import WorkloadType
+
+
+@dataclass
+class AgentContext:
+    """Per-task state handed to ``Agent.run``.
+
+    ``fire`` is the runner's request-firing callable. When non-None, agents
+    can route their own HTTP calls through bench-maker's session + hook
+    pipeline instead of using a separate client — handy if you want the
+    runner's pre/post hooks (auth headers, request tracing) on every internal
+    call. Most agents will simply ignore it.
+    """
+
+    item: Any
+    workload_name: str = "agent"
+    fire: Optional[FireRequest] = None
+    start_mono: float = 0.0
+
+
+@dataclass
+class AgentResult:
+    """What ``Agent.run`` returns for one task.
+
+    - ``output`` is the text the grader sees (``extract_text`` returns it as-is).
+    - ``ok`` is whether the agent considers this trajectory a success
+      (typically: did it submit?). May be flipped to False by a correctness
+      post-hook when the submission is wrong.
+    - ``request_ok`` is whether the agent ran to completion without an
+      *infrastructure* failure (model endpoint reachable, no Python exceptions,
+      …). Default True. When False the sample is bucketed as "fail" by the
+      progress logger and summary; when True but ``ok`` ends up False it's
+      bucketed as "wrong" — i.e. "agent delivered an answer, but it isn't
+      the right one." Crashed agents are auto-set to ``request_ok=False`` by
+      the workload-type (exceptions caught in ``run_ticket``).
+    - ``metrics`` lands in ``Sample.extra`` (numeric values only).
+    - ``meta`` lands in ``Sample.meta`` (arbitrary JSON-serializable values).
+    """
+
+    output: str = ""
+    ok: bool = True
+    error: Optional[str] = None
+    request_ok: bool = True
+    metrics: dict[str, float] = field(default_factory=dict)
+    meta: dict[str, Any] = field(default_factory=dict)
+    bytes_sent: int = 0
+    bytes_recv: int = 0
+
+
+class Agent(ABC):
+    """Base class for user-provided agents.
+
+    Subclasses can take whatever ``__init__`` kwargs they need; the
+    workload-type instantiates the class once and reuses the same instance
+    across tickets, so any expensive setup (loading tools, opening sessions)
+    happens once. Override ``aclose`` to release resources.
+    """
+
+    @abstractmethod
+    async def run(self, ctx: AgentContext) -> AgentResult:
+        """Handle one task; return its final result."""
+
+    async def aclose(self) -> None:
+        return None
+
+
+class CallableAgent(Agent):
+    """Adapter wrapping a sync/async ``(AgentContext) -> AgentResult|dict|str``."""
+
+    def __init__(self, fn: Callable[[AgentContext], Any], name: str = "callable"):
+        self._fn = fn
+        self._name = name
+
+    async def run(self, ctx: AgentContext) -> AgentResult:
+        result = self._fn(ctx)
+        result = await maybe_await(result)
+        if isinstance(result, AgentResult):
+            return result
+        if isinstance(result, dict):
+            return AgentResult(**result)
+        if isinstance(result, str):
+            return AgentResult(output=result)
+        raise TypeError(
+            f"CallableAgent fn must return AgentResult|dict|str, "
+            f"got {type(result).__name__}"
+        )
+
+
+class AgentWorkloadType(WorkloadType):
+    """Run a user-provided ``Agent`` per ticket.
+
+    Args:
+        agent: an ``Agent`` instance, an ``Agent`` subclass (instantiated with
+            ``agent_kwargs``), or a plain callable wrapped in ``CallableAgent``.
+        agent_kwargs: forwarded to the class constructor when ``agent`` is a class.
+        reference_key: item key carrying the gold reference. Copied into
+            ``Sample.meta`` (and the synthetic ``Request.meta``) so the standard
+            ``correctness_hook`` finds it. Mirrors ``EvalWorkloadType``.
+        extra_meta_keys: additional item keys to copy verbatim into
+            ``Sample.meta`` (e.g. ``task_id``).
+        name: workload-type name (default ``"agent"``).
+    """
+
+    streaming = False
+    # We split `reference` out of items ourselves, so apply_correctness should
+    # NOT wrap us in EvalWorkloadType — install the post-hook directly.
+    handles_reference = True
+
+    def __init__(
+        self,
+        agent: Union[Agent, type, Callable[..., Any]],
+        *,
+        agent_kwargs: Optional[dict[str, Any]] = None,
+        reference_key: str = "reference",
+        extra_meta_keys: tuple[str, ...] = (),
+        name: str = "agent",
+    ):
+        self.name = name
+        self._reference_key = reference_key
+        self._extra_meta_keys = tuple(extra_meta_keys)
+        self._agent = _coerce_agent(agent, agent_kwargs or {})
+
+    async def make_request(self, item: Any) -> Request:
+        return Request(
+            method="AGENT",
+            url=f"agent://{self.name}",
+            meta=self._split_meta(item),
+        )
+
+    async def run_ticket(self, ctx: TicketContext) -> Sample:
+        req = await self.make_request(ctx.item)
+        agent_ctx = AgentContext(
+            item=ctx.item,
+            workload_name=self.name,
+            fire=ctx.fire,
+            start_mono=ctx.start_mono,
+        )
+        start = ctx.start_mono
+
+        try:
+            result = await self._agent.run(agent_ctx)
+        except Exception as e:
+            latency = time.monotonic() - start
+            error = f"{type(e).__name__}: {e}"
+            sample = Sample(
+                start_ts=start,
+                latency_s=latency,
+                status=0,
+                ok=False,
+                request_ok=False,
+                error=error,
+                workload=self.name,
+                meta=dict(req.meta),
+            )
+            resp = Response(
+                status=0, headers={}, body=b"",
+                elapsed_s=latency, ok=False, error=error,
+            )
+            return await self._run_post_hooks(ctx, req, resp, sample)
+
+        latency = time.monotonic() - start
+        body = (result.output or "").encode("utf-8", errors="replace")
+        # `resp.ok` gates the correctness post-hook. We want grading to run
+        # whenever the agent finished its trajectory (even an unsuccessful one
+        # — empty output correctly grades as "wrong"). Only skip grading on
+        # actual infra failure (request_ok=False).
+        resp = Response(
+            status=200 if result.request_ok else 0,
+            headers={},
+            body=body,
+            elapsed_s=latency,
+            ok=result.request_ok,
+            error=result.error,
+        )
+
+        merged_meta = dict(req.meta)
+        merged_meta.update(result.meta or {})
+
+        extra: dict[str, float] = {}
+        for k, v in (result.metrics or {}).items():
+            if isinstance(v, (int, float)):
+                extra[k] = float(v)
+
+        sample = Sample(
+            start_ts=start,
+            latency_s=latency,
+            # When the agent ran to completion (request_ok=True) we record a
+            # synthetic 200 even if `ok` is False — the trajectory was
+            # *delivered*, it just didn't submit a successful answer. That keeps
+            # the ok/wrong/failed bucketing meaningful.
+            status=200 if result.request_ok else 0,
+            ok=result.ok,
+            request_ok=result.request_ok,
+            bytes_sent=result.bytes_sent,
+            bytes_recv=result.bytes_recv or len(body),
+            error=result.error,
+            workload=self.name,
+            meta=merged_meta,
+            extra=extra,
+        )
+        return await self._run_post_hooks(ctx, req, resp, sample)
+
+    @staticmethod
+    async def _run_post_hooks(ctx: TicketContext, req: Request, resp: Response,
+                              sample: Sample) -> Sample:
+        for hook in ctx.post_hooks:
+            sample = await maybe_await(hook(req, resp, sample))
+        return sample
+
+    def _split_meta(self, item: Any) -> dict[str, Any]:
+        if not isinstance(item, dict):
+            return {}
+        keys = {self._reference_key, *self._extra_meta_keys}
+        return {k: item[k] for k in keys if k in item}
+
+    async def aclose(self) -> None:
+        await self._agent.aclose()
+
+
+def _coerce_agent(spec: Any, agent_kwargs: dict[str, Any]) -> Agent:
+    if isinstance(spec, Agent):
+        if agent_kwargs:
+            raise ValueError(
+                "agent_kwargs is only meaningful when 'agent' is a class or "
+                "factory — got an Agent instance"
+            )
+        return spec
+    if inspect.isclass(spec) and issubclass(spec, Agent):
+        return spec(**agent_kwargs)
+    if callable(spec):
+        # Factory path: try calling it with agent_kwargs; if the result is an
+        # Agent, use it. Otherwise treat the original callable as a per-task fn.
+        if agent_kwargs:
+            obj = spec(**agent_kwargs)
+            if isinstance(obj, Agent):
+                return obj
+            if callable(obj):
+                return CallableAgent(obj)
+            raise TypeError(
+                f"agent factory returned {type(obj).__name__}; "
+                "expected Agent or callable"
+            )
+        return CallableAgent(spec)
+    raise TypeError(
+        f"agent must be Agent|Agent-subclass|callable, got {type(spec).__name__}"
+    )