blackjax-devs / blackjax-devs/tuningfork
Roadmap: per-cell crash isolation for nightly benchmark suite
- Dominant language
- Python
- Stars
- 3
- Forks
- 0
- Avg merge
- 7h 24m
- Merged PRs (30d)
- 4
Description
## Problem
When a benchmark cell hard-crashes the host process (SIGABRT, SIGSEGV, OOM-kill), the current implementation loses all subsequent cells in the same pytest session. The current workaround is a manual `_KNOWN_CRASHERS` allowlist in `benchmarks/conftest.py` — but this requires human intervention after every new crasher (caught by an aborted nightly), and one undiscovered crasher silently kills the rest of the suite.
**This is a stopgap, not a design.** A trustworthy nightly should survive any individual cell crashing.
## Goal
Per-cell crash isolation: each cell runs in a subprocess. If a cell SIGABRTs/OOMs, the harness:
1. Records that cell as CRASHED (with exit code and stderr)
2. Continues to the next cell
3. Emits a `::error::CELL_CRASHED` annotation for that cell in CI
4. Still produces a `bench_results.json` with all non-crashed cells
The nightly goes RED if any cell crashes (exit code non-zero from harness), but the remaining cells all run and contribute data.
## Constraints
- **JAX fork-deadlock caveat**: JAX cannot be safely `fork()`'d after import (GPU state, thread-pool, CUDA context). Each subprocess must import JAX from scratch (`spawn`, not `fork`). This is slow (adds ~5-10s per cell for JAX cold-start) but correct.
- The session-level JIT warmup pass becomes per-subprocess (each cell's subprocess does its own JIT warm before the timed runs).
- pytest-benchmark's `--benchmark-json` output format needs aggregation across subprocesses.
## Implementation sketch
Option A: **subprocess-per-cell in `run_benchmark_cell`** (spawn + JSON IPC) — runs each 7-run cell in a `multiprocessing.get_context("spawn").Process`, communicates results via a temp JSON file. On non-zero exit, records CRASHED.
Option B: **pytest-subprocess plugin or `--forked`** — use `pytest-forked` or a custom `@pytest.mark.subprocess` decorator. Cleaner pytest integration but less control over the result-capture path.
Option C: **wrapper script** — a thin shell wrapper that runs each cell individually (`pytest -k `) and collects results, killing and recording on SIGABRT. Avoids the JAX-fork issue entirely.
## Related
- #137 — root-cause diagnosis for laplace_hmc SIGABRT (the cell that triggered this)
- `benchmarks/conftest.py` `_KNOWN_CRASHERS` — current manual stopgap
## Priority
Roadmap. Does not block the current nightly (the `_KNOWN_CRASHERS` stopgap handles known cases); promote when the stopgap becomes a maintenance burden (>3 entries) or a new crasher is missed in production.
Contributor guide
Research direction
Start in benchmarks/conftest.py, especially the _KNOWN_CRASHERS stopgap and run_benchmark_cell entry point; review the related crash in issue #137. Compare the spawn/JSON IPC, pytest-subprocess, and wrapper-script options, then verify that crashed cells are recorded with exit code and stderr, later cells still run, bench_results.json is aggregated, CI emits ::error::CELL_CRASHED, and the harness exits non-zero.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- performance, testing-qa
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100