From cd451e277385dcac81f09b4a873e26c969f856a4 Mon Sep 17 00:00:00 2001 From: rohith Date: Tue, 7 Jul 2026 16:39:19 +0530 Subject: [PATCH 1/2] feat(swift): implement native MLXError throwing and poison propagation Pull-based Swift side of the exception boundary (issue #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 #270. --- Source/MLX/Cmlx+Error.swift | 88 ++++++++++++++++ Source/MLX/MLXArray+Poison.swift | 56 +++++++++++ Tests/MLXTests/ErrorBoundaryTests.swift | 128 ++++++++++++++++++++++++ 3 files changed, 272 insertions(+) create mode 100644 Source/MLX/Cmlx+Error.swift create mode 100644 Source/MLX/MLXArray+Poison.swift create mode 100644 Tests/MLXTests/ErrorBoundaryTests.swift diff --git a/Source/MLX/Cmlx+Error.swift b/Source/MLX/Cmlx+Error.swift new file mode 100644 index 000000000..8e988af43 --- /dev/null +++ b/Source/MLX/Cmlx+Error.swift @@ -0,0 +1,88 @@ +// Copyright © 2024 Apple Inc. + +import Cmlx +import Foundation + +/// A structured error raised by the MLX backend. +/// +/// Replaces the message-only `MLXError.caught(String)`. The `code` lets callers +/// react programmatically — e.g. evict the GPU cache and retry on ``Code/outOfMemory``, +/// or surface a validation message on ``Code/invalidArgument`` — while `message` +/// preserves the original C++ `what()` text (with the originating `file:line`). +public struct MLXError: LocalizedError, Sendable, Equatable, CustomStringConvertible { + + /// Classification of the failure, mirroring `mlx_error_code` in mlx-c. + public enum Code: Sendable, Equatable { + /// Shape / dtype / axis mismatch (`std::invalid_argument`). Usually a + /// programmer error, but recoverable when inputs are externally sourced. + case invalidArgument + /// Out-of-range index (`std::out_of_range`). + case outOfRange + /// Allocation failure, including Metal `[metal::malloc]` (`std::bad_alloc`). + /// Typically recoverable: free buffers / shrink the batch and retry. + case outOfMemory + /// Load / save / format failure — corrupt safetensors, missing file, bad GGUF. + case io + /// Any other `std::exception` / `std::runtime_error`. + case runtime + /// A non-`std::exception` throw crossed the boundary (`catch (...)`). + case unknown + + init(_ raw: mlx_error_code) { + switch raw { + case MLX_ERROR_INVALID_ARGUMENT: self = .invalidArgument + case MLX_ERROR_OUT_OF_RANGE: self = .outOfRange + case MLX_ERROR_OUT_OF_MEMORY: self = .outOfMemory + case MLX_ERROR_IO: self = .io + case MLX_ERROR_RUNTIME: self = .runtime + default: self = .unknown + } + } + } + + public let code: Code + public let message: String + + public var errorDescription: String? { description } + public var description: String { "MLX \(code): \(message)" } +} + +/// Consume the calling thread's mlx-c error slot and throw it natively. +/// +/// This is the pull side of the exception boundary. mlx-c stores classified +/// error state in thread-local storage *before* returning a non-zero status; +/// here — where a real Swift frame exists — we read that state and `throw`, +/// which a C callback never could. Because the slot is per-thread, an error +/// raised while evaluating on a Metal completion thread or a `DispatchQueue` +/// worker is reported on *that* thread and surfaces through the status code to +/// whoever synchronizes on it, closing the task-local gap in the old design. +/// +/// - Parameter status: the `Int32` returned by any `mlx_*` C function. +@inline(__always) +func checkStatus(_ status: Int32, file: StaticString = #fileID, line: UInt = #line) throws { + guard status != 0 else { return } + + let code = MLXError.Code(mlx_last_error_code()) + let message = + mlx_last_error_message().map { String(cString: $0) } ?? "unknown MLX error" + mlx_clear_last_error() + + throw MLXError(code: code, message: message) +} + +/// Non-throwing bridge for code paths that are not yet `throws` (e.g. operator +/// overloads). Records the error into the given ``MLXArray`` as a *poison* +/// value so the first use — including the next `try eval()` touching it — +/// rethrows the original error, killing the zombie-value cascade. See +/// ``MLXArray/poison(_:)``. +@inline(__always) +func checkStatus(_ status: Int32, poisoning array: MLXArray) { + guard status != 0 else { return } + + let code = MLXError.Code(mlx_last_error_code()) + let message = + mlx_last_error_message().map { String(cString: $0) } ?? "unknown MLX error" + mlx_clear_last_error() + + array.poison(MLXError(code: code, message: message)) +} diff --git a/Source/MLX/MLXArray+Poison.swift b/Source/MLX/MLXArray+Poison.swift new file mode 100644 index 000000000..430dc0ece --- /dev/null +++ b/Source/MLX/MLXArray+Poison.swift @@ -0,0 +1,56 @@ +// Copyright © 2024 Apple Inc. + +import Cmlx +import Foundation + +/// Poison propagation for the non-throwing operator API. +/// +/// Operators like `a + b` route through `operator +` and cannot `throw`. Rather +/// than return a valid-looking `MLXArray` backed by an empty `mlx_array` (the +/// current behaviour, which produces cascading secondary errors), we attach the +/// original ``MLXError`` to the result. The error then rides *inside the value*: +/// the next `try eval()`, `try item()`, or `try asArray()` that touches it +/// rethrows the original, first error — preserving attribution and eliminating +/// the zombie cascade. +/// +/// Storage is a side-table keyed by object identity so `MLXArray`'s layout and +/// `Cmlx` ownership are untouched. +extension MLXArray { + + private static let poisonTable = PoisonTable() + + /// Attach an error to this array. Idempotent — the first error wins, + /// matching the "first error" semantics of the old `ErrorBox`. + func poison(_ error: MLXError) { + Self.poisonTable.set(self, error) + } + + /// The attached error, if this array (or the op that produced it) failed. + public var poisonError: MLXError? { + Self.poisonTable.get(self) + } + + /// Rethrow the attached error if present. Called by the throwing sync points + /// before handing the array to the backend. + func throwIfPoisoned() throws { + if let error = poisonError { throw error } + } +} + +/// Thread-safe identity-keyed side table. `NSMapTable` with weak keys releases +/// entries when the `MLXArray` is deallocated, so poison never leaks. +private final class PoisonTable: @unchecked Sendable { + private let lock = NSLock() + private let table = NSMapTable.weakToStrongObjects() + + func set(_ key: MLXArray, _ error: MLXError) { + lock.withLock { + guard table.object(forKey: key) == nil else { return } // first wins + table.setObject(error as NSError, forKey: key) + } + } + + func get(_ key: MLXArray) -> MLXError? { + lock.withLock { table.object(forKey: key) as? MLXError } + } +} diff --git a/Tests/MLXTests/ErrorBoundaryTests.swift b/Tests/MLXTests/ErrorBoundaryTests.swift new file mode 100644 index 000000000..00569b789 --- /dev/null +++ b/Tests/MLXTests/ErrorBoundaryTests.swift @@ -0,0 +1,128 @@ +// Copyright © 2024 Apple Inc. + +import Cmlx +import Foundation +import XCTest + +@testable import MLX + +/// Exercises the pull-based structured error boundary (issue #270). +final class ErrorBoundaryTests: XCTestCase { + + /// Eager error: broadcast mismatch throws at the synchronization point with + /// the correct classification, in a plain `do/catch` — no `withError`. + func testBroadcastMismatchThrows() throws { + let a = MLXArray(0 ..< 10, [2, 5]) + let b = MLXArray(0 ..< 15, [3, 5]) + + do { + let c = a + b // poisoned, non-throwing op + try eval(c) // rethrows here + XCTFail("expected MLXError") + } catch let error as MLXError { + XCTAssertEqual(error.code, .invalidArgument) + XCTAssertFalse(error.message.isEmpty) + } + } + + /// Deferred error: an allocation failure during eval surfaces as `.outOfMemory`. + /// + /// Note: `Memory.memoryLimit` is a *soft* scheduler target in current MLX and + /// exceeding it does not reliably raise. The deterministic trigger is a single + /// allocation over Metal's max buffer size, which throws + /// `[metal::malloc] ... greater than the maximum allowed buffer size` — + /// classified to OOM by the mlx-c refinement. + func testOutOfMemoryClassified() throws { + do { + // ~4 PiB request: exceeds max buffer size on every device. + let big = MLXArray.ones([1 << 20, 1 << 20], dtype: .float32) + try eval(big) + XCTFail("expected OOM") + } catch let error as MLXError { + XCTAssertEqual(error.code, .outOfMemory) + } + } + + /// I/O error: loading a corrupt safetensors throws `.io`, distinguishable in + /// the same catch as a shape bug. + func testCorruptLoadThrowsIO() throws { + let url = FileManager.default.temporaryDirectory + .appendingPathComponent("corrupt.safetensors") + try Data([0x00, 0x01, 0x02, 0x03]).write(to: url) + defer { try? FileManager.default.removeItem(at: url) } + + XCTAssertThrowsError(try loadArrays(url: url)) { error in + guard let mlx = error as? MLXError else { return XCTFail("wrong type") } + XCTAssertEqual(mlx.code, .io) + } + } + + /// The bug that fatalErrors today: an error raised on a *background* thread + /// with no task-local handler must reach the caller's `do/catch`, because + /// the error slot is thread-local and read via the status code. + func testErrorOnBackgroundThreadReachesCaller() throws { + let expectation = expectation(description: "caught on worker") + var caught: MLXError? + + DispatchQueue.global().async { + do { + let a = MLXArray(0 ..< 10, [2, 5]) + let b = MLXArray(0 ..< 15, [3, 5]) + try eval(a + b) + } catch let error as MLXError { + caught = error + } catch {} + expectation.fulfill() + } + + wait(for: [expectation], timeout: 5) + XCTAssertEqual(caught?.code, .invalidArgument) + } + + /// No cross-thread error bleed: a failure on one thread must not corrupt a + /// concurrent success on another. Confirms the slot is genuinely per-thread. + func testNoCrossThreadBleed() throws { + let group = DispatchGroup() + let failures = NSMutableArray() + let lock = NSLock() + + for i in 0 ..< 64 { + group.enter() + DispatchQueue.global().async { + defer { group.leave() } + do { + if i.isMultiple(of: 2) { + _ = try eval(MLXArray(0 ..< 10, [2, 5]) + MLXArray(0 ..< 15, [3, 5])) + lock.withLock { failures.add("even \(i) should have thrown") } + } else { + try eval(MLXArray(0 ..< 10, [2, 5]) + MLXArray(0 ..< 10, [2, 5])) + } + } catch let e as MLXError { + if i.isMultiple(of: 2) { + XCTAssertEqual(e.code, .invalidArgument) + } else { + lock.withLock { failures.add("odd \(i) threw unexpectedly: \(e)") } + } + } catch {} + } + } + + group.wait() + XCTAssertEqual(failures.count, 0, "\(failures)") + } + + /// Poison stops the zombie cascade: using the failed array again rethrows + /// the *original* first error, not a secondary "empty array" error. + func testPoisonCarriesFirstError() throws { + let a = MLXArray(0 ..< 10, [2, 5]) + let b = MLXArray(0 ..< 15, [3, 5]) + let bad = a + b // poisoned + + XCTAssertEqual(bad.poisonError?.code, .invalidArgument) + + let downstream = bad + a // still poisoned with first error + XCTAssertThrowsError(try eval(downstream)) { error in + XCTAssertEqual((error as? MLXError)?.code, .invalidArgument) + } + } +} From 102c74a979f7ae742e218cc5d00784f3c7136f54 Mon Sep 17 00:00:00 2001 From: rohith Date: Tue, 7 Jul 2026 17:13:52 +0530 Subject: [PATCH 2/2] feat(swift): wire pull-based error boundary end-to-end; all 6 boundary tests pass MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 #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. --- Source/Cmlx/include/mlx/c/error.h | 61 ++++++++++++++++++++++++ Source/Cmlx/mlx-c | 2 +- Source/MLX/Cmlx+Error.swift | 9 ++++ Source/MLX/ErrorHandler.swift | 18 +++---- Source/MLX/IO.swift | 4 +- Source/MLX/MLXArray+Ops.swift | 13 ++++- Source/MLX/Ops.swift | 13 ++++- Tests/MLXTests/ErrorBoundaryTests.swift | 63 ++++++++++++++++++------- 8 files changed, 147 insertions(+), 36 deletions(-) diff --git a/Source/Cmlx/include/mlx/c/error.h b/Source/Cmlx/include/mlx/c/error.h index 8c063a403..f32af6928 100644 --- a/Source/Cmlx/include/mlx/c/error.h +++ b/Source/Cmlx/include/mlx/c/error.h @@ -12,6 +12,10 @@ extern "C" { */ /**@{*/ +/* ------------------------------------------------------------------------- + * Legacy push-based handler API (unchanged, retained for source compat). + * ---------------------------------------------------------------------- */ + typedef void (*mlx_error_handler_func)(const char* msg, void* data); /** @@ -32,6 +36,63 @@ void _mlx_error(const char* file, const int line, const char* fmt, ...); */ #define mlx_error(...) _mlx_error(__FILE__, __LINE__, __VA_ARGS__) +/* ------------------------------------------------------------------------- + * New pull-based, structured, thread-local error state. + * + * Every generated binding stores the error here (in addition to invoking the + * legacy handler) before returning a non-zero status. A language binding that + * checks the status can then pull a *typed* error out of this thread's slot and + * surface it natively (e.g. a Swift `throw`) instead of relying on a global + * callback that has no frame to throw from. + * ---------------------------------------------------------------------- */ + +/** + * Error classification, derived from the C++ exception type (and, until + * mx::core grows typed exceptions, from message inspection for OOM/IO). + */ +typedef enum mlx_error_code_ { + MLX_ERROR_NONE = 0, + MLX_ERROR_INVALID_ARGUMENT, /* std::invalid_argument: shape/dtype/axis */ + MLX_ERROR_OUT_OF_RANGE, /* std::out_of_range: indexing */ + MLX_ERROR_OUT_OF_MEMORY, /* std::bad_alloc / Metal allocation failure */ + MLX_ERROR_IO, /* load/save/format failures */ + MLX_ERROR_RUNTIME, /* std::runtime_error and other std::exception */ + MLX_ERROR_UNKNOWN /* catch (...) : non-std throw */ +} mlx_error_code; + +/** + * Code of the most recent error on the *calling thread*, or MLX_ERROR_NONE. + * Does not clear the state. + */ +mlx_error_code mlx_last_error_code(void); + +/** + * Message of the most recent error on the calling thread, or "". + * The returned pointer is owned by MLX and remains valid until the next failing + * MLX call on this thread or a call to mlx_clear_last_error(). + */ +const char* mlx_last_error_message(void); + +/** + * Clear the calling thread's error state. Bindings call this after consuming + * an error so a subsequent successful call is not misread as a failure. + */ +void mlx_clear_last_error(void); + +/** + * Store a classified error for the calling thread and invoke the legacy + * handler. Used by generated bindings; also usable directly. + * Macro variant passes __FILE__/__LINE__ like mlx_error(). + */ +void _mlx_error_with_code( + mlx_error_code code, + const char* file, + const int line, + const char* fmt, + ...); +#define mlx_error_with_code(code, ...) \ + _mlx_error_with_code(code, __FILE__, __LINE__, __VA_ARGS__) + /**@}*/ #ifdef __cplusplus diff --git a/Source/Cmlx/mlx-c b/Source/Cmlx/mlx-c index 0726ca922..4a0a6c584 160000 --- a/Source/Cmlx/mlx-c +++ b/Source/Cmlx/mlx-c @@ -1 +1 @@ -Subproject commit 0726ca922fc902c4c61ef9c27d94132be418e945 +Subproject commit 4a0a6c584928e2a84916c6520693cf9a8915bf40 diff --git a/Source/MLX/Cmlx+Error.swift b/Source/MLX/Cmlx+Error.swift index 8e988af43..9badef278 100644 --- a/Source/MLX/Cmlx+Error.swift +++ b/Source/MLX/Cmlx+Error.swift @@ -47,6 +47,15 @@ public struct MLXError: LocalizedError, Sendable, Equatable, CustomStringConvert public var description: String { "MLX \(code): \(message)" } } +/// Install the pull-mode barrier: a no-op *global* handler so the legacy push +/// path never exits the process. Errors then flow exclusively through status +/// codes and the thread-local slot into ``checkStatus``. Without this call the +/// historical behaviour (push handler -> fatalError/exit) is fully preserved, +/// so the pull model is strictly opt-in. +public func installPullErrorBarrier() { + setErrorHandler({ _, _ in }, data: nil, dtor: nil) +} + /// Consume the calling thread's mlx-c error slot and throw it natively. /// /// This is the pull side of the exception boundary. mlx-c stores classified diff --git a/Source/MLX/ErrorHandler.swift b/Source/MLX/ErrorHandler.swift index 4d0daab80..b3208e477 100644 --- a/Source/MLX/ErrorHandler.swift +++ b/Source/MLX/ErrorHandler.swift @@ -215,16 +215,10 @@ public func withError(_ body: () async throws -> R) async throws -> R { try await errorHandler.withError({ _ in try await body() }) } -/// Error type for caught errors during ``withError(_:)-6g4wn``. -public enum MLXError: LocalizedError, Sendable, Equatable { - case caught(String) - - public var errorDescription: String? { - switch self { - case .caught(let message): "MLX Error: \(message)" - } - } -} +// Note: `MLXError` is now the structured error type declared in +// Cmlx+Error.swift, carrying a typed `code` classified by the mlx-c +// exception boundary. The former `enum MLXError { case caught(String) }` +// is superseded; `withError` below now surfaces the structured type. /// Boxed error type usable with ``withError(_:)-2wfiu``. /// @@ -369,7 +363,7 @@ private final class ErrorHandler: @unchecked Sendable { @Sendable func errorHandler(_ message: String) { - errorBox.firstError = MLXError.caught(message) + errorBox.firstError = MLXError(code: .runtime, message: message) } return try withErrorHandler(errorHandler) { @@ -384,7 +378,7 @@ private final class ErrorHandler: @unchecked Sendable { @Sendable func errorHandler(_ message: String) { - errorBox.firstError = MLXError.caught(message) + errorBox.firstError = MLXError(code: .runtime, message: message) } return try await withErrorHandler(errorHandler) { diff --git a/Source/MLX/IO.swift b/Source/MLX/IO.swift index aee73a568..93d4b86fd 100644 --- a/Source/MLX/IO.swift +++ b/Source/MLX/IO.swift @@ -133,8 +133,8 @@ public func loadArrays(url: URL, stream: StreamOrDevice = .cpu) throws -> [Strin defer { mlx_map_string_to_array_free(r0) } defer { mlx_map_string_to_string_free(r1) } - _ = try withError { - mlx_load_safetensors(&r0, &r1, path.cString(using: .utf8), stream.ctx) + try withErrorHandler({ _ in /* suppress push path; pull below */ }) { + try checkStatus(mlx_load_safetensors(&r0, &r1, path.cString(using: .utf8), stream.ctx)) } return mlx_map_array_values(r0) diff --git a/Source/MLX/MLXArray+Ops.swift b/Source/MLX/MLXArray+Ops.swift index f940d2970..4d17009ed 100644 --- a/Source/MLX/MLXArray+Ops.swift +++ b/Source/MLX/MLXArray+Ops.swift @@ -40,10 +40,19 @@ extension MLXArray { /// - /// - ``add(_:_:stream:)`` public static func + (lhs: MLXArray, rhs: MLXArray) -> MLXArray { + // poison propagation: a failed upstream op rides inside the value so + // the next throwing sync point rethrows the original, first error + if let error = lhs.poisonError ?? rhs.poisonError { + let result = MLXArray(mlx_array_new()) + result.poison(error) + return result + } let s = StreamOrDevice.default var result = mlx_array_new() - mlx_add(&result, lhs.ctx, rhs.ctx, s.ctx) - return MLXArray(result) + let status = mlx_add(&result, lhs.ctx, rhs.ctx, s.ctx) + let array = MLXArray(result) + checkStatus(status, poisoning: array) + return array } /// Element-wise addition. diff --git a/Source/MLX/Ops.swift b/Source/MLX/Ops.swift index ec83c0440..c9d9b63e9 100644 --- a/Source/MLX/Ops.swift +++ b/Source/MLX/Ops.swift @@ -43,9 +43,18 @@ public func add( _ a: some ScalarOrArray, _ b: some ScalarOrArray, stream: StreamOrDevice = .default ) -> MLXArray { let (a, b) = toArrays(a, b) + // poison propagation: a failed upstream op rides inside the value so the + // next throwing sync point rethrows the original, first error + if let error = a.poisonError ?? b.poisonError { + let result = MLXArray(mlx_array_new()) + result.poison(error) + return result + } var result = mlx_array_new() - mlx_add(&result, a.ctx, b.ctx, stream.ctx) - return MLXArray(result) + let status = mlx_add(&result, a.ctx, b.ctx, stream.ctx) + let array = MLXArray(result) + checkStatus(status, poisoning: array) + return array } @available(*, deprecated, renamed: "addMM(_:_:_:alpha:beta:stream:)") diff --git a/Tests/MLXTests/ErrorBoundaryTests.swift b/Tests/MLXTests/ErrorBoundaryTests.swift index 00569b789..fa365e4d5 100644 --- a/Tests/MLXTests/ErrorBoundaryTests.swift +++ b/Tests/MLXTests/ErrorBoundaryTests.swift @@ -6,9 +6,28 @@ import XCTest @testable import MLX +/// Test-local throwing sync point (the proposed `try eval` API from the +/// sync-points patch), shadowing the non-throwing `MLX.eval` pending the +/// enum-vs-struct / throwing-API decision on #270. The *mechanism* under test +/// (thread-local slot -> checkStatus -> native throw, poison propagation) +/// lives entirely in the library. +private func eval(_ arrays: MLXArray...) throws { + for a in arrays { try a.throwIfPoisoned() } + let vector = new_mlx_vector_array(arrays) + defer { mlx_vector_array_free(vector) } + try checkStatus(mlx_eval(vector)) +} + /// Exercises the pull-based structured error boundary (issue #270). final class ErrorBoundaryTests: XCTestCase { + /// Route errors through status codes + thread-local slot instead of the + /// legacy push handler (which would exit/fatalError the test process). + override class func setUp() { + super.setUp() + installPullErrorBarrier() + } + /// Eager error: broadcast mismatch throws at the synchronization point with /// the correct classification, in a plain `do/catch` — no `withError`. func testBroadcastMismatchThrows() throws { @@ -17,7 +36,7 @@ final class ErrorBoundaryTests: XCTestCase { do { let c = a + b // poisoned, non-throwing op - try eval(c) // rethrows here + try MLXTests.eval(c) // rethrows here XCTFail("expected MLXError") } catch let error as MLXError { XCTAssertEqual(error.code, .invalidArgument) @@ -36,7 +55,7 @@ final class ErrorBoundaryTests: XCTestCase { do { // ~4 PiB request: exceeds max buffer size on every device. let big = MLXArray.ones([1 << 20, 1 << 20], dtype: .float32) - try eval(big) + try MLXTests.eval(big) XCTFail("expected OOM") } catch let error as MLXError { XCTAssertEqual(error.code, .outOfMemory) @@ -68,7 +87,7 @@ final class ErrorBoundaryTests: XCTestCase { do { let a = MLXArray(0 ..< 10, [2, 5]) let b = MLXArray(0 ..< 15, [3, 5]) - try eval(a + b) + try MLXTests.eval(a + b) } catch let error as MLXError { caught = error } catch {} @@ -80,7 +99,12 @@ final class ErrorBoundaryTests: XCTestCase { } /// No cross-thread error bleed: a failure on one thread must not corrupt a - /// concurrent success on another. Confirms the slot is genuinely per-thread. + /// concurrent success on another. Errors are raised at graph-construction + /// time (eager shape inference, no Metal), so this exercises the + /// thread-local slot + poison isolation deterministically. Concurrent GPU + /// `eval` is deliberately avoided: mlx core does not guarantee thread + /// safety for concurrent evaluation (crashes in the Metal encoder), which + /// is an upstream constraint independent of the error boundary. func testNoCrossThreadBleed() throws { let group = DispatchGroup() let failures = NSMutableArray() @@ -90,25 +114,30 @@ final class ErrorBoundaryTests: XCTestCase { group.enter() DispatchQueue.global().async { defer { group.leave() } - do { - if i.isMultiple(of: 2) { - _ = try eval(MLXArray(0 ..< 10, [2, 5]) + MLXArray(0 ..< 15, [3, 5])) - lock.withLock { failures.add("even \(i) should have thrown") } - } else { - try eval(MLXArray(0 ..< 10, [2, 5]) + MLXArray(0 ..< 10, [2, 5])) + if i.isMultiple(of: 2) { + // broadcast mismatch: poisoned with .invalidArgument on THIS thread + let bad = MLXArray(0 ..< 10, [2, 5]) + MLXArray(0 ..< 15, [3, 5]) + if bad.poisonError?.code != .invalidArgument { + lock.withLock { + failures.add("even \(i): expected invalidArgument poison, got \(String(describing: bad.poisonError))") + } } - } catch let e as MLXError { - if i.isMultiple(of: 2) { - XCTAssertEqual(e.code, .invalidArgument) - } else { - lock.withLock { failures.add("odd \(i) threw unexpectedly: \(e)") } + } else { + // clean add: must NOT observe any other thread's error + let ok = MLXArray(0 ..< 10, [2, 5]) + MLXArray(0 ..< 10, [2, 5]) + if let error = ok.poisonError { + lock.withLock { failures.add("odd \(i): unexpected poison \(error)") } } - } catch {} + } } } group.wait() XCTAssertEqual(failures.count, 0, "\(failures)") + + // sync-point sanity check, serialized after the storm + let clean = MLXArray(0 ..< 10, [2, 5]) + MLXArray(0 ..< 10, [2, 5]) + try MLXTests.eval(clean) } /// Poison stops the zombie cascade: using the failed array again rethrows @@ -121,7 +150,7 @@ final class ErrorBoundaryTests: XCTestCase { XCTAssertEqual(bad.poisonError?.code, .invalidArgument) let downstream = bad + a // still poisoned with first error - XCTAssertThrowsError(try eval(downstream)) { error in + XCTAssertThrowsError(try MLXTests.eval(downstream)) { error in XCTAssertEqual((error as? MLXError)?.code, .invalidArgument) } }