Skip to content

Pull-based error boundary: typed MLXError, native throws, poison propagation (#270) - #439

Draft
kapellirohith wants to merge 2 commits into
ml-explore:mainfrom
kapellirohith:exception-boundary-fix
Draft

Pull-based error boundary: typed MLXError, native throws, poison propagation (#270)#439
kapellirohith wants to merge 2 commits into
ml-explore:mainfrom
kapellirohith:exception-boundary-fix

Conversation

@kapellirohith

@kapellirohith kapellirohith commented Jul 7, 2026

Copy link
Copy Markdown

Motivation

Addresses #270 - full design rationale in this comment. Companion to ml-explore/mlx-c#120, which this builds on.

The current handler-based design has three gaps this PR closes:

  1. Task-local hole - errors raised on threads without task-local state (async eval, Metal completion, GCD workers) fall through to fatalError, even inside withError.
  2. Stringly-typed errors - MLXError.caught(String) cannot distinguish recoverable OOM from a programmer shape bug.
  3. Zombie cascade - a failed op returns a valid-looking MLXArray over an empty handle; subsequent use produces misleading secondary errors.

Design (pull model, opt-in)

  • MLXError struct with typed code (invalidArgument / outOfRange / outOfMemory / io / runtime / unknown), mirroring mlx_error_code
  • checkStatus(_:) consumes the calling thread's mlx-c error slot at the call site - where a real Swift frame exists - and throws natively. Because the slot is per-thread and read via the returned status, errors surface to whoever synchronizes, regardless of task context
  • Poison propagation - non-throwing operator paths attach the error to the result (checkStatus(_:poisoning:), identity-keyed weak side-table); the first error rides inside the value and rethrows at the next sync point
  • installPullErrorBarrier() - opt-in no-op global handler so the push path never exits the process; without it, historical behavior (handler → fatalError) is fully preserved
  • loadArrays classifies through the slot (corrupt file → .io), withError/ErrorBox now surface the structured type

Verification

xcodebuild test -only-testing:MLXTests/ErrorBoundaryTests - 6/6 passed:

Test Proves
testBroadcastMismatchThrows eager shape error → .invalidArgument in do/catch
testOutOfMemoryClassified deferred [metal::malloc] over max buffer size → .outOfMemory
testCorruptLoadThrowsIO corrupt safetensors → .io, distinguishable from shape bugs
testErrorOnBackgroundThreadReachesCaller the case that fatalErrors today, caught in caller's do/catch
testNoCrossThreadBleed 64-way concurrent errors, no cross-thread slot contamination
testPoisonCarriesFirstError original error propagates through derived values

Draft caveats (deliberately left open for maintainer input)

  • Submodule: Source/Cmlx/mlx-c points at my fork's v0.6.0-based branch (the mlx-c changes rebased onto the current pin, since mlx-c main targets MLX 0.31.2 vs the pinned 0.31.1). Needs Add structured thread-local error state with typed exception classification mlx-c#120 to land + a submodule bump before this can leave draft.
  • enum MLXErrorstruct MLXError is source-breaking; happy to reshape (new name + deprecation, or land in a minor bump) per your preference.
  • Throwing sync points (try eval / item / asArray) are implemented as a test-target shim pending the API-shape decision in error handling in mlx-swift #270 - the mechanism is entirely in the library; only the public throwing overloads await consensus.
  • Cross-thread testing avoids concurrent GPU eval (Metal encoder crash under 64-way eval - an existing core constraint, independent of this change); isolation is proven at graph-construction time.

Pull-based Swift side of the exception boundary (issue ml-explore#270):

- MLXError struct with a typed `code` (invalidArgument / outOfRange /
  outOfMemory / io / runtime / unknown) mirroring mlx_error_code
- checkStatus(_:) consumes the calling thread's mlx-c error slot and
  throws natively at the call site, closing the task-local thread hole
- checkStatus(_:poisoning:) + MLXArray poison side-table carry the
  first error inside the value for non-throwing operator paths, ending
  the zombie-value cascade
- ErrorBoundaryTests: broadcast mismatch, deterministic OOM via
  max-buffer-size allocation, corrupt safetensors -> .io,
  background-thread error reaching the caller's do/catch, no
  cross-thread bleed, first-error poison attribution

Draft: depends on the mlx-c exception-boundary-fix branch (submodule
bump) and the throwing sync-point changes (eval/item/asArray); the
existing MLXError enum in ErrorHandler.swift is superseded pending the
enum-vs-struct API decision on ml-explore#270.
…y tests pass

Verified locally via `xcodebuild test -only-testing:MLXTests/ErrorBoundaryTests`
(** TEST SUCCEEDED **, 6/6):

- broadcast mismatch -> MLXError(.invalidArgument) via typed catch ladder
- max-buffer OOM     -> MLXError(.outOfMemory) via [metal::malloc] refinement
- corrupt safetensors-> MLXError(.io) via [read]/[load prefix refinement
- error on GCD thread reaches caller's do/catch (fatalError case fixed)
- no cross-thread slot bleed (64-way concurrent graph-construction errors)
- poison carries the first error through derived values to `try eval`

Changes:
- submodule mlx-c -> exception-boundary-fix-v0.6.0 (thread-local slot +
  typed catch ladder + IO refinement, rebased on the v0.6.0 pin)
- vendored include/mlx/c/error.h synced with the submodule header
- MLXError struct (typed .code) replaces the message-only enum;
  withError/ErrorBox now surface the structured type
- checkStatus() pull bridge + installPullErrorBarrier() opt-in;
  historical push->fatalError behaviour preserved when not installed
- `+ (MLXArray, MLXArray)` and add() capture status -> poison; poison
  propagates through derived values (first error wins)
- loadArrays classifies through the slot instead of withError
- ErrorBoundaryTests with a test-local throwing eval shim pending the
  ml-explore#270 throwing-sync-point API decision

Note: concurrent GPU eval is not exercised — mlx core does not
guarantee thread safety for concurrent evaluation (Metal encoder
crash); isolation is proven at graph-construction time instead.
@kapellirohith
kapellirohith marked this pull request as ready for review July 7, 2026 17:33
@kapellirohith
kapellirohith marked this pull request as draft July 8, 2026 06:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant