inspect_ai #4839: transcript notification benchmark (script + raw results)

README.md

Benchmark: transcript notification cost (inspect_ai PR #4839)

Provenance for the performance table in UKGovernmentBEIS/inspect_ai#4839. One script, one machine, one interpreter and dependency set — only the inspect_ai source tree differs between the two columns, selected by PYTHONPATH.

• origin/main at 0161c68f51
• perf/bounded-transcript-linear at 008a1b4cda

What it measures

Drives one sample's growing conversation through the realistic four-notification-per-turn shape (pending → call registration → completion → timestamp stamping) with a real SampleBufferDatabase subscribed the way _eval/task/run.py subscribes one.

counter	meaning	load-dependent?
call_hashes	calls to inspect_ai.event._pool._call_hash	no — deterministic
msg_hashes	calls to inspect_ai.event._pool._msg_hash	no — deterministic
bytes_shipped	total string-parameter bytes handed to sqlite3 across every write path (events, attachments, message_pool, call_pool). INSERT OR IGNORE discards duplicates only after the content has crossed, so this is bytes shipped, not stored.	no — deterministic
bytes_unique	distinct >100-char values among those — content that had to cross at least once	no — deterministic
wall_s	median wall-clock over --repeat runs	yes

Deterministic counters (the primary evidence)

These are exact and reproduce anywhere, regardless of machine load:

shape	call hashes	bytes shipped	bytes unique
400 turns × 1.5KB	160,400 → 800	125.8MB → 1.79MB	1.35MB (both)
200 turns × 1.5KB	40,200 → 400	32.0MB → 0.90MB	0.68MB (both)
200 turns × 32KB	40,200 → 400	666.0MB → 7.14MB	6.92MB (both)

bytes_unique is identical on both sides — the correctness check on the measurement. The branch ships less because it stops re-shipping content, not because less content exists; on the branch bytes_shipped lands close to bytes_unique, i.e. each distinct payload crosses roughly once.

msg_hashes is identical on both sides at every shape (400/400, 200/200): message pooling is not what changed here, call-payload hashing is.

Wall clock

Absolute timings come from the minimum of 4 interleaved runs — main, branch, main, branch within one session — on a shared machine with unrelated load. Minimum is the right statistic here: it approximates the uncontended time, whereas a mean or median absorbs whatever contention happened to land.

shape	main	branch	ratio
400 turns × 1.5KB	2.041s	1.261s	1.62×
200 turns × 1.5KB	0.562s	0.366s	1.54×
200 turns × 32KB	0.755s	0.370s	2.04×

Three independent sessions on this machine agree to within a few percent (2.07/2.11/2.04 main and 1.26/1.30/1.26 branch for the 400-turn shape), across different origin/main bases — which also shows the intervening upstream commits do not touch these paths.

Raw interleaved iterations: interleaved-final2.txt (this run), interleaved-quiet.txt (earlier session).

Reproducing

git worktree add --detach /tmp/main-bench origin/main
PYTHONPATH=/tmp/main-bench/src   python bench_transcript_notify.py --repeat 3 --json > main.json
PYTHONPATH=<branch-worktree>/src python bench_transcript_notify.py --repeat 3 --json > branch.json

Both columns must use the same interpreter so dependency versions are held constant. On a quiet machine the counters will match this gist exactly; wall times will differ.

Raw output: main3.json, branch3.json.

bench_transcript_notify.py

"""Benchmark the per-notification cost of transcript event notification.

Drives a single sample's growing conversation through the realistic
four-notification-per-turn shape (pending, call registration, completion,
timestamp stamping) with a real ``SampleBufferDatabase`` subscribed the way
``_eval/task/run.py`` subscribes one, and reports primitive counters so the
derived figures can be re-checked:

  wall_s        median wall-clock seconds to drive the sample (with
                wall_min/wall_max over --repeat runs, so a reader can tell a
                real difference from noise)
  msg_hashes    calls to ``_pool._msg_hash``   (message pooling work)
  call_hashes   calls to ``_pool._call_hash``  (call pooling work)
  msgs_offered  messages presented to the pool across all notifications
  bytes_shipped total string-parameter bytes handed to sqlite3 across every
                write path (events, attachments, message_pool, call_pool).
                ``INSERT OR IGNORE`` discards duplicates only after the
                content has crossed, so this is bytes *shipped*, not stored.
  bytes_unique  distinct >100-char values among those, i.e. the content that
                had to cross at least once

Counters are deterministic for a given shape; only wall time varies.

Runs unmodified against this branch and against origin/main; select the
source tree with PYTHONPATH so both runs share one interpreter and one set
of dependency versions:

    PYTHONPATH=<tree>/src python scripts/bench_transcript_notify.py --repeat 3

Usage:
    python scripts/bench_transcript_notify.py [--json] [--repeat N]
                                              [--turns N --msg-bytes N]
"""

from __future__ import annotations

import argparse
import json
import sqlite3
import sys
import tempfile
import time
from pathlib import Path
from typing import Any

import inspect_ai.event._pool as pool
from inspect_ai.event._model import ModelEvent
from inspect_ai.log._log import EvalSampleSummary
from inspect_ai.log._recorders.buffer import SampleBufferDatabase
from inspect_ai.log._recorders.types import SampleEvent
from inspect_ai.log._transcript import Transcript
from inspect_ai.model._chat_message import ChatMessageUser
from inspect_ai.model._generate_config import GenerateConfig
from inspect_ai.model._model_call import ModelCall
from inspect_ai.model._model_output import ModelOutput


class Counters:
    def __init__(self) -> None:
        self.msg_hashes = 0
        self.call_hashes = 0
        self.msgs_offered = 0
        self.bytes_shipped = 0
        self.unique: set[str] = set()

    @property
    def bytes_unique(self) -> int:
        return sum(len(h) for h in self.unique)


def _instrument(counters: Counters) -> list[tuple[Any, str, Any]]:
    """Wrap the hash and attachment-insert paths; returns undo records."""
    undo: list[tuple[Any, str, Any]] = []

    real_msg_hash = pool._msg_hash
    real_call_hash = pool._call_hash

    def msg_hash(m: Any) -> Any:
        counters.msg_hashes += 1
        return real_msg_hash(m)

    def call_hash(m: Any) -> Any:
        counters.call_hashes += 1
        return real_call_hash(m)

    pool._msg_hash = msg_hash  # type: ignore[assignment]
    pool._call_hash = call_hash  # type: ignore[assignment]
    undo.append((pool, "_msg_hash", real_msg_hash))
    undo.append((pool, "_call_hash", real_call_hash))

    # Count at the SQLite boundary so every write path is covered (attachments,
    # message_pool, call_pool, events) rather than whichever one a given
    # workload happens to use. INSERT OR IGNORE discards duplicates only after
    # the content has crossed, so parameters handed to sqlite3 are the bytes
    # actually shipped.
    real_connect = sqlite3.connect

    class CountingConnection(sqlite3.Connection):
        def execute(self, sql: str, parameters: Any = (), /) -> Any:  # type: ignore[override]
            _tally(counters, parameters)
            return super().execute(sql, parameters)

        def executemany(self, sql: str, parameters: Any, /) -> Any:  # type: ignore[override]
            rows = list(parameters)
            for row in rows:
                _tally(counters, row)
            return super().executemany(sql, rows)

    def connect(*args: Any, **kwargs: Any) -> Any:
        kwargs.setdefault("factory", CountingConnection)
        return real_connect(*args, **kwargs)

    sqlite3.connect = connect  # type: ignore[assignment]
    undo.append((sqlite3, "connect", real_connect))

    return undo


def _tally(counters: Counters, parameters: Any) -> None:
    if isinstance(parameters, dict):
        values: Any = parameters.values()
    elif isinstance(parameters, (list, tuple)):
        values = parameters
    else:
        return
    for value in values:
        if isinstance(value, str):
            counters.bytes_shipped += len(value)
            if len(value) > 100:  # pooled/attachment-scale content only
                counters.unique.add(value)


def _restore(undo: list[tuple[Any, str, Any]]) -> None:
    for target, name, original in undo:
        setattr(target, name, original)


def run_shape(turns: int, msg_bytes: int) -> dict[str, float | int]:
    """Drive one sample of `turns` turns whose messages are ~`msg_bytes` each."""
    counters = Counters()
    # Distinct per turn so nothing dedups by accident, and large enough to be
    # pooled as attachment content rather than left inline.
    filler = "x" * max(msg_bytes - 32, 1)
    undo = _instrument(counters)
    try:
        with tempfile.TemporaryDirectory() as d:
            db = SampleBufferDatabase(location="bench", db_dir=Path(d))
            db.start_sample(EvalSampleSummary(id="s1", epoch=1, input="in", target="t"))
            transcript = Transcript(bounded=True, resident_tail=100, log_model_api=True)
            transcript._subscribe(
                lambda e: db.log_events([SampleEvent(id="s1", epoch=1, event=e)])
            )

            history: list[ChatMessageUser] = []
            wire: list[dict[str, Any]] = []
            start = time.perf_counter()
            for turn in range(turns):
                body = f"turn {turn} {filler}"
                history.append(ChatMessageUser(content=body))
                wire.append({"role": "user", "content": body})
                counters.msgs_offered += len(history)

                event = ModelEvent(
                    model="bench",
                    input=list(history),
                    tools=[],
                    tool_choice="auto",
                    config=GenerateConfig(),
                    output=ModelOutput(),
                    pending=True,
                )
                transcript._event(event)  # 1: pending
                event.call = ModelCall(
                    request={"model": "bench", "messages": [dict(m) for m in wire]}
                )
                transcript._event_updated(event)  # 2: call registration
                event.output = ModelOutput.from_content("bench", "response")
                event.call = ModelCall(
                    request={"model": "bench", "messages": [dict(m) for m in wire]},
                    response={"id": f"r{turn}"},
                )
                event.pending = None
                transcript._event_updated(event)  # 3: completion
                transcript._event_updated(event)  # 4: timestamp stamping
            wall = time.perf_counter() - start
            db.cleanup()
    finally:
        _restore(undo)

    offered = counters.msgs_offered or 1
    return {
        "turns": turns,
        "msg_bytes": msg_bytes,
        "wall_s": round(wall, 3),
        "msg_hashes": counters.msg_hashes,
        "call_hashes": counters.call_hashes,
        "msgs_offered": counters.msgs_offered,
        "hit_rate": round(1.0 - counters.msg_hashes / offered, 4),
        "bytes_shipped": counters.bytes_shipped,
        "bytes_unique": counters.bytes_unique,
    }


SHAPES = [(400, 1536), (200, 1536), (200, 32768)]


def run_repeated(turns: int, msg_bytes: int, repeat: int) -> dict[str, float | int]:
    """Run one shape `repeat` times; report median wall and its spread.

    Counters are deterministic for a given shape, so only the timing needs
    repeating; the median is reported with min/max so a reader can see how
    much of a difference is noise.
    """
    runs = [run_shape(turns, msg_bytes) for _ in range(repeat)]
    walls = sorted(float(r["wall_s"]) for r in runs)
    result = dict(runs[-1])
    result["wall_s"] = walls[len(walls) // 2]
    result["wall_min"] = walls[0]
    result["wall_max"] = walls[-1]
    result["repeat"] = repeat
    return result


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--turns", type=int)
    parser.add_argument("--msg-bytes", type=int)
    parser.add_argument("--repeat", type=int, default=3)
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()

    shapes = [(args.turns, args.msg_bytes)] if args.turns and args.msg_bytes else SHAPES
    rows = [run_repeated(t, b, args.repeat) for t, b in shapes]

    if args.json:
        print(json.dumps({"source": str(Path(pool.__file__).resolve()), "rows": rows}))
    else:
        print(f"source: {Path(pool.__file__).resolve()}")
        for r in rows:
            print(
                f"turns={r['turns']:4d} msg={r['msg_bytes']:6d}B  "
                f"wall={r['wall_s']:7.3f}s"
                f"[{r.get('wall_min', 0):.3f}-{r.get('wall_max', 0):.3f}]  "
                f"msg_hashes={r['msg_hashes']:7d}  "
                f"call_hashes={r['call_hashes']:8d}  offered={r['msgs_offered']:8d}  "
                f"shipped={r['bytes_shipped'] / 1e6:9.1f}MB  "
                f"unique={r['bytes_unique'] / 1e6:7.1f}MB"
            )
    return 0


if __name__ == "__main__":
    sys.exit(main())

branch3.json

{"source": "/home/faber/src/aisi/inspect_ai/.claude/worktrees/bounded-transcript-linear/src/inspect_ai/event/_pool.py", "rows": [{"turns": 400, "msg_bytes": 1536, "wall_s": 1.458, "msg_hashes": 400, "call_hashes": 800, "msgs_offered": 80200, "hit_rate": 0.995, "bytes_shipped": 1790449, "bytes_unique": 1350026, "wall_min": 1.437, "wall_max": 1.57, "repeat": 3}, {"turns": 200, "msg_bytes": 1536, "wall_s": 0.469, "msg_hashes": 200, "call_hashes": 400, "msgs_offered": 20100, "hit_rate": 0.99, "bytes_shipped": 894773, "bytes_unique": 674719, "wall_min": 0.435, "wall_max": 0.482, "repeat": 3}, {"turns": 200, "msg_bytes": 32768, "wall_s": 0.443, "msg_hashes": 200, "call_hashes": 400, "msgs_offered": 20100, "hit_rate": 0.99, "bytes_shipped": 7141201, "bytes_unique": 6921140, "wall_min": 0.436, "wall_max": 0.475, "repeat": 3}]}

interleaved-final.txt

400x1536|2.114 160400 125822504 1351271|1.317 800 1792089 1351256
200x1536|0.57 40200 31973288 675334|0.367 400 895553 675304
200x32768|0.789 40200 665982888 6921734|0.378 400 7141937 6921692
400x1536|2.14 160400 125822496 1351265|1.3 800 1792145 1351298
200x1536|0.562 40200 31973268 675319|0.358 400 895565 675313
200x32768|0.772 40200 665982860 6921713|0.384 400 7141969 6921716
400x1536|2.255 160400 125822464 1351241|1.316 800 1792057 1351232
200x1536|0.574 40200 31973260 675313|0.366 400 895557 675307
200x32768|0.77 40200 665982888 6921734|0.395 400 7141977 6921722
400x1536|2.157 160400 125822468 1351244|1.337 800 1792089 1351256
200x1536|0.597 40200 31973288 675334|0.397 400 895577 675322
200x32768|0.802 40200 665982852 6921707|0.388 400 7141989 6921731

interleaved-final2.txt

400x1536|2.104 160400 125822480 1351253|1.279 800 1792089 1351256
200x1536|0.567 40200 31973256 675310|0.367 400 895569 675316
200x32768|0.758 40200 665982860 6921713|0.373 400 7141965 6921713
400x1536|2.08 160400 125822500 1351268|1.308 800 1792081 1351250
200x1536|0.578 40200 31973260 675313|0.366 400 895569 675316
200x32768|0.758 40200 665982888 6921734|0.37 400 7141977 6921722
400x1536|2.041 160400 125822496 1351265|1.261 800 1792125 1351283
200x1536|0.562 40200 31973252 675307|0.379 400 895597 675337
200x32768|0.755 40200 665982868 6921719|0.376 400 7141957 6921707
400x1536|2.439 160400 125822492 1351262|1.271 800 1792033 1351214
200x1536|0.568 40200 31973244 675301|0.368 400 895581 675325
200x32768|0.758 40200 665982876 6921725|0.372 400 7141989 6921731

interleaved-quiet.txt

400x1536 iter1 main=2.071 branch=1.281
200x1536 iter1 main=0.563 branch=0.365
200x32768 iter1 main=0.765 branch=1.124
400x1536 iter2 main=2.863 branch=1.255
200x1536 iter2 main=1.722 branch=0.808
200x32768 iter2 main=1.748 branch=1.446
400x1536 iter3 main=10.604 branch=2.979
200x1536 iter3 main=0.631 branch=0.382
200x32768 iter3 main=0.764 branch=0.373
400x1536 iter4 main=2.961 branch=1.258
200x1536 iter4 main=0.571 branch=0.368
200x32768 iter4 main=0.758 branch=0.368

main3.json

{"source": "/mnt/data/scratch/main-bench/src/inspect_ai/event/_pool.py", "rows": [{"turns": 400, "msg_bytes": 1536, "wall_s": 2.08, "msg_hashes": 400, "call_hashes": 160400, "msgs_offered": 80200, "hit_rate": 0.995, "bytes_shipped": 125820856, "bytes_unique": 1350035, "wall_min": 2.079, "wall_max": 2.157, "repeat": 3}, {"turns": 200, "msg_bytes": 1536, "wall_s": 0.581, "msg_hashes": 200, "call_hashes": 40200, "msgs_offered": 20100, "hit_rate": 0.99, "bytes_shipped": 31972480, "bytes_unique": 674728, "wall_min": 0.575, "wall_max": 0.663, "repeat": 3}, {"turns": 200, "msg_bytes": 32768, "wall_s": 0.772, "msg_hashes": 200, "call_hashes": 40200, "msgs_offered": 20100, "hit_rate": 0.99, "bytes_shipped": 665982072, "bytes_unique": 6921122, "wall_min": 0.759, "wall_max": 0.775, "repeat": 3}]}