blackjax-devs / blackjax-devs/tuningfork

Roadmap: per-cell crash isolation for nightly benchmark suite

Open
#138 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.