lablup / lablup/mlxcel

fix(core): make array_evaluated_bytes fallible so GPU faults don't abort

Open
#1,791 0 comments 0 reactions 0 assignees View on GitHub
area:core area:inference priority:medium status:ready type:bug
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

## Problem

`array_evaluated_bytes` is declared without `Result` in the cxx bridge, so an MLX throw inside it is an uncatchable `std::terminate`. At the current MLX pin it does throw on a failed command buffer: since ml-explore/mlx#3742 every event a failed buffer signals carries one shared encoder error, and `array::eval()` rethrows it through `Event::check_error()`. The production caller is the lookahead decode path, so in `mlxcel-server` one failed command buffer kills the process and every in-flight request with it.

The fault does not have to be mlxcel's. A command buffer discarded because another process's fault forced a GPU recovery (`kIOGPUCommandBufferCallbackErrorInnocentVictim`) now throws at this pin where it used to be swallowed, and it aborted the `mlxcel-core` test binary with SIGABRT during the PR #1777 re-gate. Found during review of PR #1772 (MLX pin bump to 81ba1c6a).

## Evidence

- Declaration `src/lib/mlxcel-core/src/lib.rs:194` (`fn array_evaluated_bytes(arr: &MlxArray) -> Vec;`) and `src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h:210`; body `src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:549-576` (`a.eval()` at `:555`; its non-contiguous fallback to the equally infallible `array_to_raw_bytes` is at `:564`).
- Production caller `lookahead_tokens_to_host`, `src/server/batch/scheduler/mod.rs:771`, used by `pipelined_steady_decode` at `src/server/batch/scheduler/decode_tick.rs:633`.
- Test callers: `src/models/falcon_ocr_rope_tests.rs` (four sites), `src/loading/vlm_falcon_ocr_tests.rs:20`, `tests/falcon_ocr_parity.rs:144`.
- PR #1772 fixed the same class for stashed sampler launches with `stashed_launch_state` (`mlx_cxx_bridge.cpp:5585`), which reads status, event validity, `is_signaled()` and `load_error()` without throwing.

## Proposed fix

Declare it `-> Result>`, so the cxx try/catch also covers the fallback copy, and make `lookahead_tokens_to_host` return a `Result`. On `Err`, `pipelined_steady_decode` takes the #822 path the other decode evals already use (`decode_tick.rs:884-900`): `record_eval_outcome` (`src/server/batch/scheduler/handoff.rs:48`), the same lookahead teardown as its `finishing` branch (`decode_tick.rs:654-671`), `abort_sequence_with_error` (`src/server/batch/scheduler/run_loop.rs:102`) for each row in `la.ids`, then `eval_failures_exhausted()`. Update the test callers.

## Acceptance criteria

- [ ] No non-`Result` bridge call remains on the lookahead read path.
- [ ] A test builds an array whose event is signalled and carries an error (a test-only hook modelled on `sampling_dispatch_stash_failed_launch_for_test`, `mlx_cxx_bridge.cpp:5679`), passes it to `array_evaluated_bytes`, and asserts `Err`.
- [ ] A lookahead read failure fails only the rows of that step; the scheduler keeps serving other requests.

## Verification

The workspace gate: `cargo test --workspace --profile test-fast --features metal,accelerate`, `cargo clippy --workspace --all-targets -- -D warnings`, `cargo fmt --all -- --check`. With the bridge signature reverted to `-> Vec`, the new test must abort the test binary (SIGABRT) instead of passing.

Contributor guide

Open the contributing guide

Research direction

Start with the bridge declaration and implementation in src/lib/mlxcel-core/src/lib.rs, cpp/mlx_cxx_bridge.h, and cpp/mlx_cxx_bridge.cpp, then trace lookahead_tokens_to_host into decode_tick.rs. Compare the error path with the existing decode failure handling and test callers listed in the issue. Done means failed array evaluation returns Err, only affected rows are aborted, and the workspace tests, clippy, and format checks pass.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, rust
Domain
ai-infra-agents, backend, testing
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
65/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.