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 new file mode 100644 index 000000000..9badef278 --- /dev/null +++ b/Source/MLX/Cmlx+Error.swift @@ -0,0 +1,97 @@ +// 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)" } +} + +/// 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 +/// 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/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/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/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 new file mode 100644 index 000000000..fa365e4d5 --- /dev/null +++ b/Tests/MLXTests/ErrorBoundaryTests.swift @@ -0,0 +1,157 @@ +// Copyright © 2024 Apple Inc. + +import Cmlx +import Foundation +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 { + let a = MLXArray(0 ..< 10, [2, 5]) + let b = MLXArray(0 ..< 15, [3, 5]) + + do { + let c = a + b // poisoned, non-throwing op + try MLXTests.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 MLXTests.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 MLXTests.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. 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() + let lock = NSLock() + + for i in 0 ..< 64 { + group.enter() + DispatchQueue.global().async { + defer { group.leave() } + 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))") + } + } + } 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)") } + } + } + } + } + + 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 + /// 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 MLXTests.eval(downstream)) { error in + XCTAssertEqual((error as? MLXError)?.code, .invalidArgument) + } + } +}