From 72d3629886b562be54eee3b81a583091e6f1ad3f Mon Sep 17 00:00:00 2001 From: Aegis-AI Date: Mon, 27 Apr 2026 18:30:50 -0700 Subject: [PATCH 1/3] fix: replace fatalError with error latch in SSD streaming Convert ThreadSafeError.check() from fatalError (which crashes the entire app) to a global SSDStreamingErrorLatch pattern. When a pread I/O error occurs on truncated/corrupted safetensors files, the error is now posted to the latch instead of killing the process. The generation loop in Evaluate.swift checks the latch: - After model.prepare() during prefill (catches errors during prompt processing) - After each token in the generation loop (catches errors during decoding) This allows the consuming code (InferenceEngine) to surface the error gracefully in the UI and prompt the user to re-download the model. Also adds SSDStreamingError and SSDStreamingErrorLatch as public types for downstream consumers. --- Libraries/MLXLMCommon/ConcurrentError.swift | 53 ++++++++++++++++++++- Libraries/MLXLMCommon/Evaluate.swift | 21 ++++++++ 2 files changed, 73 insertions(+), 1 deletion(-) diff --git a/Libraries/MLXLMCommon/ConcurrentError.swift b/Libraries/MLXLMCommon/ConcurrentError.swift index fc14ebb5e..263c2f4d2 100644 --- a/Libraries/MLXLMCommon/ConcurrentError.swift +++ b/Libraries/MLXLMCommon/ConcurrentError.swift @@ -1,6 +1,49 @@ import Foundation import MLX +/// Error thrown when SSD expert streaming encounters a corrupted, truncated, +/// or incomplete safetensors file during pread I/O. +public struct SSDStreamingError: Error, LocalizedError { + public let underlyingError: Error + + public var errorDescription: String? { + "MLX SSD Streaming Error: \(underlyingError.localizedDescription). The model safetensors file may be corrupted, truncated, or incomplete. Try re-downloading the model." + } +} + +/// Global error latch for SSD streaming errors that occur inside non-throwing +/// `callAsFunction` paths. Set by `ThreadSafeError.check()`, cleared and +/// inspected by the generation loop after each token. +public final class SSDStreamingErrorLatch: @unchecked Sendable { + public static let shared = SSDStreamingErrorLatch() + private let lock = NSLock() + private var _error: Error? + + /// Record an error (first-wins semantics). + public func set(_ error: Error) { + lock.withLock { + if _error == nil { _error = error } + } + } + + /// Consume and return the recorded error, resetting the latch. + /// Returns nil if no error was recorded. + public func consume() -> Error? { + lock.withLock { + let e = _error + _error = nil + return e + } + } + + /// Throw the recorded error if one exists, then clear it. + public func throwIfSet() throws { + if let error = consume() { + throw error + } + } +} + package final class ThreadSafeError: @unchecked Sendable { package let lock = NSLock() package var error: Swift.Error? @@ -19,9 +62,17 @@ package final class ThreadSafeError: @unchecked Sendable { } } + /// Check if any error was recorded during concurrent I/O. + /// + /// Instead of calling `fatalError` (which crashes the entire app), this + /// posts the error to the global `SSDStreamingErrorLatch` so the generation + /// loop can detect it after the current token and surface it gracefully + /// in the UI (e.g., prompting a re-download). package func check() { if let error = error { - fatalError("MLX SSD Streaming Error: \(error.localizedDescription). (The model safetensors file may be corrupted, truncated, or incomplete).") + SSDStreamingErrorLatch.shared.set( + SSDStreamingError(underlyingError: error) + ) } } } diff --git a/Libraries/MLXLMCommon/Evaluate.swift b/Libraries/MLXLMCommon/Evaluate.swift index 36b9f2af7..1f42b6453 100644 --- a/Libraries/MLXLMCommon/Evaluate.swift +++ b/Libraries/MLXLMCommon/Evaluate.swift @@ -650,12 +650,24 @@ public struct TokenIterator: TokenIteratorProtocol { case .tokens(let tokens): y = tokens + // Check for SSD streaming errors that occurred during prefill. + // The MoE expert pread path uses a non-throwing callAsFunction, + // so errors are posted to the global latch instead. + try SSDStreamingErrorLatch.shared.throwIfSet() + // evaluate the remainder of the prompt -- this primes the pump let token = step(previous: y) + + // Check again after step() which also runs through MoE layers + try SSDStreamingErrorLatch.shared.throwIfSet() + y = .init(tokens: token) asyncEval(y.tokens) case .logits(let result): + // Check for SSD streaming errors during logits computation + try SSDStreamingErrorLatch.shared.throwIfSet() + y = .init(tokens: convertToToken(logits: result.logits)) asyncEval(y.tokens) @@ -1705,6 +1717,15 @@ private func generateLoopTask( break } + // Check for SSD streaming errors (truncated/corrupted safetensors). + // These are set by ThreadSafeError.check() inside SwitchGLU's non-throwing + // callAsFunction path via the global error latch. + if let ssdError = SSDStreamingErrorLatch.shared.consume() { + print("[MLXLMCommon] SSD streaming error detected: \(ssdError.localizedDescription)") + stopReason = .cancelled + break + } + if promptTime == 0 { let now = Date.timeIntervalSinceReferenceDate promptTime = now - start From 88faf358f28256d2147f9e18b29fda08bec955d6 Mon Sep 17 00:00:00 2001 From: Aegis-AI Date: Tue, 28 Apr 2026 08:16:24 -0700 Subject: [PATCH 2/3] Fix SSD streaming review issues --- Libraries/MLXLMCommon/ConcurrentError.swift | 52 +++++++++++--- Libraries/MLXLMCommon/Evaluate.swift | 72 +++++++++++-------- .../MLXLMTests/CorruptSafetensorsTests.swift | 26 +++++++ 3 files changed, 113 insertions(+), 37 deletions(-) diff --git a/Libraries/MLXLMCommon/ConcurrentError.swift b/Libraries/MLXLMCommon/ConcurrentError.swift index 263c2f4d2..cf2c2a837 100644 --- a/Libraries/MLXLMCommon/ConcurrentError.swift +++ b/Libraries/MLXLMCommon/ConcurrentError.swift @@ -6,19 +6,48 @@ import MLX public struct SSDStreamingError: Error, LocalizedError { public let underlyingError: Error + public init(underlyingError: Error) { + self.underlyingError = underlyingError + } + public var errorDescription: String? { "MLX SSD Streaming Error: \(underlyingError.localizedDescription). The model safetensors file may be corrupted, truncated, or incomplete. Try re-downloading the model." } } -/// Global error latch for SSD streaming errors that occur inside non-throwing -/// `callAsFunction` paths. Set by `ThreadSafeError.check()`, cleared and -/// inspected by the generation loop after each token. +private enum SSDStreamingErrorLatchContext { + static let threadDictionaryKey = "MLXLMCommon.SSDStreamingErrorLatch.active" +} + +/// Error latch for SSD streaming errors that occur inside non-throwing +/// `callAsFunction` paths. A generation installs its own active latch around +/// model execution so concurrent sessions do not cross-contaminate each other. public final class SSDStreamingErrorLatch: @unchecked Sendable { - public static let shared = SSDStreamingErrorLatch() private let lock = NSLock() private var _error: Error? + public init() {} + + package static func withActive(_ latch: SSDStreamingErrorLatch, _ body: () throws -> T) rethrows -> T { + let key = SSDStreamingErrorLatchContext.threadDictionaryKey as NSString + let threadDictionary = Thread.current.threadDictionary + let previous = threadDictionary[key] + threadDictionary[key] = latch + defer { + if let previous { + threadDictionary[key] = previous + } else { + threadDictionary.removeObject(forKey: key) + } + } + return try body() + } + + package static var active: SSDStreamingErrorLatch? { + let key = SSDStreamingErrorLatchContext.threadDictionaryKey as NSString + return Thread.current.threadDictionary[key] as? SSDStreamingErrorLatch + } + /// Record an error (first-wins semantics). public func set(_ error: Error) { lock.withLock { @@ -47,8 +76,11 @@ public final class SSDStreamingErrorLatch: @unchecked Sendable { package final class ThreadSafeError: @unchecked Sendable { package let lock = NSLock() package var error: Swift.Error? + private let latch: SSDStreamingErrorLatch? - package init() {} + package init(latch: SSDStreamingErrorLatch? = SSDStreamingErrorLatch.active) { + self.latch = latch + } package func catchError(_ block: () throws -> Void) { do { @@ -68,11 +100,13 @@ package final class ThreadSafeError: @unchecked Sendable { /// posts the error to the global `SSDStreamingErrorLatch` so the generation /// loop can detect it after the current token and surface it gracefully /// in the UI (e.g., prompting a re-download). - package func check() { + @discardableResult + package func check() -> SSDStreamingError? { if let error = error { - SSDStreamingErrorLatch.shared.set( - SSDStreamingError(underlyingError: error) - ) + let streamingError = SSDStreamingError(underlyingError: error) + latch?.set(streamingError) + return streamingError } + return nil } } diff --git a/Libraries/MLXLMCommon/Evaluate.swift b/Libraries/MLXLMCommon/Evaluate.swift index 1f42b6453..d9ef4b3fa 100644 --- a/Libraries/MLXLMCommon/Evaluate.swift +++ b/Libraries/MLXLMCommon/Evaluate.swift @@ -502,6 +502,7 @@ protocol TokenIteratorProtocol: Sequence, IteratorProtocol where Element == Int var maxTokens: Int? { get } var tokenCount: Int { get } var promptPrefillTime: TimeInterval { get } + var streamingError: SSDStreamingError? { get } } /// Generator of tokens. @@ -546,6 +547,8 @@ public struct TokenIterator: TokenIteratorProtocol { // Internal metrics var promptPrefillTime: TimeInterval = 0.0 + var streamingError: SSDStreamingError? + let ssdErrorLatch = SSDStreamingErrorLatch() /// Initialize a `TokenIterator` with the given tokens. Note: this has been /// replaced with ``init(input:model:cache:parameters:)``. @@ -646,27 +649,24 @@ public struct TokenIterator: TokenIteratorProtocol { mutating func prepare(input: LMInput, windowSize: Int? = nil) throws { processor?.prompt(input.text.tokens) - switch try model.prepare(input, cache: cache, windowSize: windowSize) { + let preparation = try SSDStreamingErrorLatch.withActive(ssdErrorLatch) { + try model.prepare(input, cache: cache, windowSize: windowSize) + } + + switch preparation { case .tokens(let tokens): y = tokens - // Check for SSD streaming errors that occurred during prefill. - // The MoE expert pread path uses a non-throwing callAsFunction, - // so errors are posted to the global latch instead. - try SSDStreamingErrorLatch.shared.throwIfSet() + try ssdErrorLatch.throwIfSet() // evaluate the remainder of the prompt -- this primes the pump - let token = step(previous: y) - - // Check again after step() which also runs through MoE layers - try SSDStreamingErrorLatch.shared.throwIfSet() + let token = try step(previous: y) y = .init(tokens: token) asyncEval(y.tokens) case .logits(let result): - // Check for SSD streaming errors during logits computation - try SSDStreamingErrorLatch.shared.throwIfSet() + try ssdErrorLatch.throwIfSet() y = .init(tokens: convertToToken(logits: result.logits)) asyncEval(y.tokens) @@ -689,11 +689,14 @@ public struct TokenIterator: TokenIteratorProtocol { } /// Evaluate the next token and return the new token (y), updating cache state - mutating func step(previous: LMInput.Text) -> MLXArray { - let result = model( - previous[text: .newAxis], cache: cache.isEmpty ? nil : cache, state: state) + mutating func step(previous: LMInput.Text) throws -> MLXArray { + let result = SSDStreamingErrorLatch.withActive(ssdErrorLatch) { + model(previous[text: .newAxis], cache: cache.isEmpty ? nil : cache, state: state) + } self.state = result.state + try ssdErrorLatch.throwIfSet() + // Apply dynamic cache quantization after each step maybeQuantizeKVCache( cache: &cache, @@ -706,6 +709,10 @@ public struct TokenIterator: TokenIteratorProtocol { } mutating public func next() -> Int? { + if streamingError != nil { + return nil + } + if let maxTokens, tokenCount >= maxTokens { return nil } @@ -714,7 +721,17 @@ public struct TokenIterator: TokenIteratorProtocol { let previousY = y // compute the next state and async eval the next token - let token = step(previous: previousY) + let token: MLXArray + do { + token = try step(previous: previousY) + } catch let error as SSDStreamingError { + streamingError = error + return nil + } catch { + streamingError = SSDStreamingError(underlyingError: error) + return nil + } + y = .init(tokens: token) asyncEval(token) @@ -758,6 +775,7 @@ public struct SpeculativeTokenIterator: TokenIteratorProtocol { let draftModel: any LanguageModel var mainState: LMOutput.State? + public let streamingError: SSDStreamingError? = nil var mainCache: [KVCache] var draftCache: [KVCache] let quantizeKVCache: (inout [KVCache]) -> Void @@ -1141,7 +1159,9 @@ private func runSynchronousGenerationLoop( // If the iterator ends naturally, the max-token limit was reached. if stopReason == nil { - if let maxTokens = iterator.maxTokens, iterator.tokenCount >= maxTokens { + if iterator.streamingError != nil { + stopReason = .error + } else if let maxTokens = iterator.maxTokens, iterator.tokenCount >= maxTokens { stopReason = .length } else { stopReason = .cancelled @@ -1697,7 +1717,7 @@ private func generateLoopTask( // Launch a Task to perform iteration asynchronously. let task = Task { let performIteration = { - let iterator = iterator.consume() + var iterator = iterator.consume() var handler = handler.consume() var start = Date.timeIntervalSinceReferenceDate @@ -1710,22 +1730,13 @@ private func generateLoopTask( tokenizer: tokenizer ) - for token in iterator { + while let token = iterator.next() { // Check for cancellation on every loop iteration. if Task.isCancelled { stopReason = .cancelled break } - // Check for SSD streaming errors (truncated/corrupted safetensors). - // These are set by ThreadSafeError.check() inside SwitchGLU's non-throwing - // callAsFunction path via the global error latch. - if let ssdError = SSDStreamingErrorLatch.shared.consume() { - print("[MLXLMCommon] SSD streaming error detected: \(ssdError.localizedDescription)") - stopReason = .cancelled - break - } - if promptTime == 0 { let now = Date.timeIntervalSinceReferenceDate promptTime = now - start @@ -1753,7 +1764,9 @@ private func generateLoopTask( } if stopReason == nil { - if Task.isCancelled { + if iterator.streamingError != nil { + stopReason = .error + } else if Task.isCancelled { stopReason = .cancelled } else if let maxTokens = iterator.maxTokens, tokenCount >= maxTokens { stopReason = .length @@ -1821,6 +1834,9 @@ public enum GenerateStopReason: Sendable { /// Generation stopped due to explicit task cancellation or early stream termination. case cancelled + + /// Generation stopped because model evaluation failed. + case error } /// Represents metadata and statistics related to token generation. diff --git a/Tests/MLXLMTests/CorruptSafetensorsTests.swift b/Tests/MLXLMTests/CorruptSafetensorsTests.swift index 446501bca..cf78e8597 100644 --- a/Tests/MLXLMTests/CorruptSafetensorsTests.swift +++ b/Tests/MLXLMTests/CorruptSafetensorsTests.swift @@ -5,6 +5,32 @@ import Testing @Suite struct CorruptSafetensorsTests { + @Test + func testThreadSafeErrorCheckPublishesToActiveLatch() throws { + let latch = SSDStreamingErrorLatch() + + SSDStreamingErrorLatch.withActive(latch) { + let errState = ThreadSafeError() + errState.catchError { + throw NSError(domain: "CorruptSafetensorsTests", code: 13, userInfo: [ + NSLocalizedDescriptionKey: "truncated shard" + ]) + } + + let latched = errState.check() + #expect(latched != nil) + } + + do { + try latch.throwIfSet() + Issue.record("Expected latch.throwIfSet() to surface an SSDStreamingError") + } catch let error as SSDStreamingError { + #expect(error.localizedDescription.contains("truncated shard")) + } catch { + Issue.record("Unexpected error type: \(error)") + } + } + @Test func testDeadlock() throws { let tempDir = FileManager.default.temporaryDirectory From 38d7ff2840ab6b91a84b8f168c3cc2539f9356e1 Mon Sep 17 00:00:00 2001 From: Aegis-AI Date: Tue, 28 Apr 2026 08:22:48 -0700 Subject: [PATCH 3/3] Restore downstream compatibility --- Libraries/MLXLMCommon/ConcurrentError.swift | 2 ++ Libraries/MLXLMCommon/Evaluate.swift | 11 ++--------- 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/Libraries/MLXLMCommon/ConcurrentError.swift b/Libraries/MLXLMCommon/ConcurrentError.swift index cf2c2a837..528cc815c 100644 --- a/Libraries/MLXLMCommon/ConcurrentError.swift +++ b/Libraries/MLXLMCommon/ConcurrentError.swift @@ -23,6 +23,7 @@ private enum SSDStreamingErrorLatchContext { /// `callAsFunction` paths. A generation installs its own active latch around /// model execution so concurrent sessions do not cross-contaminate each other. public final class SSDStreamingErrorLatch: @unchecked Sendable { + public static let shared = SSDStreamingErrorLatch() private let lock = NSLock() private var _error: Error? @@ -105,6 +106,7 @@ package final class ThreadSafeError: @unchecked Sendable { if let error = error { let streamingError = SSDStreamingError(underlyingError: error) latch?.set(streamingError) + SSDStreamingErrorLatch.shared.set(streamingError) return streamingError } return nil diff --git a/Libraries/MLXLMCommon/Evaluate.swift b/Libraries/MLXLMCommon/Evaluate.swift index d9ef4b3fa..41460cb86 100644 --- a/Libraries/MLXLMCommon/Evaluate.swift +++ b/Libraries/MLXLMCommon/Evaluate.swift @@ -1159,9 +1159,7 @@ private func runSynchronousGenerationLoop( // If the iterator ends naturally, the max-token limit was reached. if stopReason == nil { - if iterator.streamingError != nil { - stopReason = .error - } else if let maxTokens = iterator.maxTokens, iterator.tokenCount >= maxTokens { + if let maxTokens = iterator.maxTokens, iterator.tokenCount >= maxTokens { stopReason = .length } else { stopReason = .cancelled @@ -1764,9 +1762,7 @@ private func generateLoopTask( } if stopReason == nil { - if iterator.streamingError != nil { - stopReason = .error - } else if Task.isCancelled { + if Task.isCancelled || iterator.streamingError != nil { stopReason = .cancelled } else if let maxTokens = iterator.maxTokens, tokenCount >= maxTokens { stopReason = .length @@ -1834,9 +1830,6 @@ public enum GenerateStopReason: Sendable { /// Generation stopped due to explicit task cancellation or early stream termination. case cancelled - - /// Generation stopped because model evaluation failed. - case error } /// Represents metadata and statistics related to token generation.