Running the benchmarks¶
Every published table comes from a harness that ships inside the package: bonsai.bench. There are two ways to stand it up, and they differ only in where the package comes from.
From the installed wheel (the normal path):
pip install "bonsai-gbt[bench]"
python -m bonsai.bench.grinsztajn out.jsonl
The [bench] extra pulls the reference libraries the suites compare against (XGBoost, LightGBM, CatBoost, pandas, openml). Plain import bonsai.bench needs none of them; heavy libraries load lazily per suite.
From a source tree (when you are hacking on bonsai itself):
make python
PYTHONPATH=build/python python -m bonsai.bench.grinsztajn out.jsonl
make python builds the extension into build/python/bonsai; PYTHONPATH=build/python makes that tree the package. Everything below works identically under either setup.
The suites¶
| suite | division | command | notes |
|---|---|---|---|
grinsztajn |
quality | python -m bonsai.bench.grinsztajn out.jsonl then --report |
The external standings suite (55 third-party tasks). Datasets fetch from OpenML on first run; rows already in out.jsonl are skipped, so an interrupted run resumes by re-running the same command. --report renders the standings from the jsonl. |
scaling |
perf | python -m bonsai.bench.scaling --smoke |
Fit seconds vs rows/cols/bins/threads against the reference libraries. --smoke is the laptop mode (small cells, minutes); the full grid (--axis all) takes hours and wants a CUDA build (make python-cuda) for the GPU variants. --dry-run prints the grid without running it. |
datasets |
(fetcher) | python -m bonsai.bench.datasets --list |
Lists the pinned datasets and their cache state; python -m bonsai.bench.datasets <name> fetches one ahead of time. |
Custom ladders: the spec-driven CLI¶
Cells that are not on a suite's built-in grid run through the unified CLI: a JSON spec names the cells (or a generator), the variants, threads, and repeats, and the driver handles child processes, resume, and row emission.
python -m bonsai.bench specs
python -m bonsai.bench plan --spec gpu-tall
python -m bonsai.bench run --spec gpu-tall --out gpu-tall.jsonl
python -m bonsai.bench variants
plan prints the expansion (every cell, variant, timeout, repeat count) without fitting anything. run resumes by default when the output file already exists: finished rows are skipped, failures re-attempt, so an interrupted sweep continues by re-running the same command. Campaign specs ship inside the package (bench/specs/), so a bare name works from any install and --spec also takes a path to your own JSON; the iso_volume generator holds rows x cols constant while sweeping the aspect ratio, and GPU rows record measured peak device memory (dev_mem) sampled while the child runs.
Every measurement runs in its own child process, so an out-of-memory, a segfault, or a timeout lands as a jsonl row instead of killing the sweep. Cells the driver judges infeasible from host RAM and VRAM are recorded as status="skipped" with the estimate that ruled them out, because the feasibility frontier is data rather than an error; a spec can switch the gates off (gates.mem_gate: off) when the measured failure is the point.
GPU and pod runs¶
The GPU variants need the extension built against the interpreter that will run it, since the nanobind module is ABI-tied to one Python: make python-cuda PYTHON=$(uv python find 3.12). The bench-scaling target picks up build-cuda/python on its own when that build exists.
For a campaign on a rented pod, copy the committed driver scripts/pod_bench_driver.sh up and launch it with HOST_TAG, SPEC, OUT, and RUN_LABEL set: it clones, builds, and runs the spec, and re-invoking resumes. Bring the rows home with ssh pod cat <out>.jsonl >> benchmarks/results/<file>.jsonl, since jsonl appends compose, then commit and delete the pod. The RunPod runbook covers the image, pod acceptance, and the rest of the loop.
Reading what comes out¶
Suites append one JSON row per measurement to the output file, self-describing enough to reproduce: the command, the knob set (hashed for grouping), the git sha, and the host down to library versions (bonsai.bench.runlog). The published tables in the results ledger are rendered from committed rows of exactly this shape.
One honest caveat before comparing numbers across machines: identical-model GPUs across rental fleets measure up to ~25% apart, so only same-host comparisons mean anything; the benchmark protocol is the full set of rules the published numbers follow.
The building blocks are importable directly when you want a custom harness:
from bonsai.bench import metrics, synth
X, y, X_test, y_test = synth.gen_data(rows=1000, cols=8, seed=0, n_test=200, informative=4)
print("train shape:", X.shape, "| r2 of predicting zero:", round(metrics.r2(y_test, y_test * 0), 3))