fix(cuda): MLX CUDA context teardown race turns green test runs into exit 101
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Problem / Background
On the GB10 self-hosted CUDA runner a test binary sometimes aborts after it has already printed `test result: ok`. The tail of the run is:
```
test result: ok. N passed; 0 failed; ...
terminate called after throwing an instance of 'std::runtime_error'
what(): Destroy(handle_) failed: driver shutting down
error: test failed, to rerun pass `--lib`
Caused by: process didn't exit successfully: ... (signal: 6, SIGABRT)
```
Cargo then exits 101 with zero failed tests, so a green run is reported red and the epic gate has to be re-run by hand. Observed frequency:
- 1 of 10 runs of `cargo test --profile test-fast --features cuda --lib embeddings::` at `917815ef` (PR #1410).
- 2 of 17 runs of `cargo test --profile test-fast --features cuda --lib models::modernbert` at `0a407d48` (PR #1412).
- Earlier on unmodified main (2026-08-22) in the `mlxcel-surgery` crate's `replace_integration` test.
It is load dependent: never seen in an isolated single run on a quiet box, seen when a sibling `cargo` or test process shares the GPU. The exception text comes from the CUDA driver refusing a `cuGraphExecDestroy` / `cudaStreamDestroy` style call after `cuDevicePrimaryCtxRelease` has already run, which means an MLX static (the CUDA graph cache, a stream pool, or an allocator) is being destroyed after the driver has torn down the primary context. This is the same abort class as the graph-cache thrashing throw that `MLX_CUDA_GRAPH_CACHE_SIZE=2000` mitigates (issue #818, `docs/upstream/mlx-cuda-graph-cache-lifetime-miss-abort.md`), but it fires at process exit rather than during an eval.
## Current Behavior
- mlxcel does not perform any explicit MLX shutdown. `src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp` exposes `clear_memory_cache()` (line 4410, wraps `mlx::core::clear_cache()`) and `synchronize_default()` (line 4441, wraps `mlx::core::synchronize()`), but nothing calls them at exit; there is no `atexit`, destructor attribute, or `Drop`-based shutdown anywhere in `src/lib/mlxcel-core/cpp/*.cpp` or `src/lib/mlxcel-core/src/*.rs`.
- Test binaries end through libtest's normal return from `main`, so MLX's static destructors run in whatever order the C++ runtime picks relative to the CUDA runtime's own `atexit` handlers. When the driver's teardown wins the race, the MLX destructor throws and, being inside a destructor, terminates the process.
- The MLX commit in effect is pinned through `src/lib/mlxcel-core/build.rs:37-47` (`mlx_pin::read_pinned_commit`), so the root-cause reading must be done against that commit's `mlx/backend/cuda/device.cpp` and `mlx/backend/cuda/lru_cache.h`.
## Proposed Solution
1. Root cause: read the pinned MLX's `mlx/backend/cuda/device.cpp` (the `Device` and `CommandEncoder` destructors, `graph_cache_`, the stream pool) and identify which static owns the handle that `Destroy(handle_)` fails on, and why it outlives the driver. Record the file and line in the PR body, and file the finding under `docs/upstream/` next to the existing lifetime-miss report so it can be sent upstream.
2. Add an explicit teardown hook in `mlxcel-core`: `pub fn shutdown()` in `src/lib/mlxcel-core/src/lib.rs`, backed by a bridge function that runs `mlx::core::synchronize()`, `mlx::core::clear_cache()`, and whatever the root-cause reading shows is needed to drain the CUDA graph cache (for example evaluating nothing further and releasing cached graphs) before `main` returns. The function must be idempotent and a no-op on Metal and CPU-only builds.
3. Call it from every process exit path that matters: `src/main.rs:2059` (`mlxcel`) and `src/bin/mlx_server.rs:1348` (`mlxcel-server`) at the end of `main`, and test binaries through a `#[ctor::dtor]` hook in `mlxcel-core` (`ctor = "1.0.13"` is already a dependency at `src/lib/mlxcel-core/Cargo.toml:32`) so that every crate that links MLX gets it without each test module opting in. `bench_decode` and `speculative_bench` (`src/bin/`) get the same hook through the dtor, so they need no edit.
4. If the root cause is a static destructor in MLX itself that no host-side call can pre-empt, the hook should instead register an `atexit` handler that runs before MLX's statics are destroyed (registration order is reverse of execution, so registering it during `mlxcel_core` initialization, which happens after MLX statics are constructed, gives it priority) and drains the caches there.
Rejected alternative: retrying the cargo command in the Makefile or CI when the exit is 101 with zero failures. That hides the abort and the same teardown race would hit `mlxcel-server` on SIGTERM.
## Scope
**In scope:** `src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp`, `src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h`, `src/lib/mlxcel-core/src/lib.rs`, `src/main.rs`, `src/bin/mlx_server.rs`, a `docs/upstream/` note.
**Out of scope:** the in-eval graph-cache thrashing abort (#818, #821), raising `MLX_CUDA_GRAPH_CACHE_SIZE`, changing the GB10 runner to serialize jobs.
## Implementation Notes
- **Reuse**: `clear_memory_cache` and `synchronize_default` already exist in the bridge; the `no_cuda` stub pattern used by `cuda_is_available` (`src/lib/mlxcel-core/src/lib.rs:2069-2072`) is how to keep the hook backend-agnostic.
- **Constraints**: the hook must not run while worker threads still hold streams (the server's worker threads must be joined first); it must not throw across the FFI boundary, so wrap the C++ body in try/catch and log rather than propagate.
- **Edge cases**: process exit through `std::process::exit` (bypasses `Drop`, so `atexit` is the only reliable path); exit while a `spawn_blocking` eval is still running (synchronize first); a CPU-only build where none of the CUDA calls exist.
- **Error handling**: a failure inside the hook is logged at `warn` and swallowed; the process exit code must never change because of the hook.
## Acceptance Criteria
- [ ] The PR names the MLX static and the destructor ordering that produces `Destroy(handle_) failed: driver shutting down`, with file and line at the pinned MLX commit.
- [ ] `mlxcel_core::shutdown()` exists, is idempotent, and is reached on normal exit of `mlxcel`, `mlxcel-server`, and every test binary that links `mlxcel-core`.
- [ ] 20 consecutive runs of `cargo test --profile test-fast --features cuda --lib embeddings::` and 20 of `cargo test --profile test-fast --features cuda --lib models::modernbert` on GB10, with another cargo build or test process sharing the GPU during the runs, exit 0 with no SIGABRT.
- [ ] Integrated into the real code flow: the hook is wired into the binaries' `main` and the test harness, not left as an unused function.
## Verification
```bash
cargo fmt --all -- --check
cargo clippy --lib --tests --profile test-fast --features cuda -- -D warnings
for i in $(seq 1 20); do cargo test --profile test-fast --features cuda --lib embeddings:: || { echo "abort on run $i"; break; }; done
for i in $(seq 1 20); do cargo test --profile test-fast --features cuda --lib models::modernbert || { echo "abort on run $i"; break; }; done
```
Run the two loops while a second shell executes `cargo test --profile test-fast --features cuda --lib server::` to reproduce the shared-GPU load. Pass: all 40 runs exit 0 and no run prints `driver shutting down`.
## Technical Considerations
Related: #818 and #821 (graph-cache thrashing abort and its upstream report), `docs/CONTINUOUS_BATCHING.md:427-429` (why the server raises `MLX_CUDA_GRAPH_CACHE_SIZE` at startup), the `make verify-test-cuda` gate at `Makefile:688-690`.
Originating PRs: #1410, #1412.
Follow-up from epic #1348.
Contributor guide
Research direction
Start by reading the pinned MLX CUDA implementation referenced by src/lib/mlxcel-core/build.rs:37-47, especially device.cpp and lru_cache.h, to identify the destructor ordering. Then inspect the bridge and shutdown entry points in mlx_cxx_bridge.cpp, src/lib/mlxcel-core/src/lib.rs, src/main.rs:2059, and src/bin/mlx_server.rs:1348. Done means an idempotent backend-safe hook is wired into exits and tests, documented, and the two 20-run shared-GPU loops exit 0 without driver-shutdown aborts.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, rust
- Domain
- backend, performance, testing
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100