fix(moe): gracefully catch safetensors I/O errors to prevent deadlock crashes#36
Conversation
There was a problem hiding this comment.
Pull request overview
This PR addresses deadlock/freeze scenarios when safetensors SSD streaming reads fail (e.g., truncated/corrupted .safetensors) during MoE expert loading, by making concurrent pread error handling single-fault and deterministic within SwitchGLU.
Changes:
- Added a shared
ThreadSafeErrorhelper inMLXLMCommonto wrap MLX C++ exceptions viaMLX.withErrorand defer failure to a single post-join crash. - Wrapped multiple
DispatchQueue.concurrentPerformexpert-loadingpreadblocks inSwitchGLUwithThreadSafeError.catchError(...)and addederrState.check()after the parallel region. - Added a new test intended to reproduce the corrupted-safetensors concurrent read scenario.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
Libraries/MLXLMCommon/SwitchLayers.swift |
Wrapes concurrent expert pread blocks with thread-safe MLX error interception and a single post-parallel failure point. |
Libraries/MLXLMCommon/ConcurrentError.swift |
Introduces ThreadSafeError (lock + first-error capture + check() fatalError). |
Tests/MLXLMTests/CorruptSafetensorsTests.swift |
Adds a regression test that corrupts a safetensors file and runs parallel preadIntoOffset calls to confirm an error is captured instead of deadlocking. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| final class ThreadSafeError: @unchecked Sendable { | ||
| let lock = NSLock() | ||
| var error: Error? | ||
| func catchError(_ block: () throws -> Void) { | ||
| do { | ||
| try MLX.withError { | ||
| try block() | ||
| } | ||
| } catch { | ||
| lock.withLock { | ||
| if self.error == nil { self.error = error } | ||
| } | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
ThreadSafeError is redefined inside this test, duplicating the implementation that was added to MLXLMCommon (Libraries/MLXLMCommon/ConcurrentError.swift). This makes the test easier to drift from production behavior and misses check() semantics. Prefer using MLXLMCommon.ThreadSafeError directly rather than defining a local copy.
| package func check() { | ||
| if let error = error { | ||
| fatalError("MLX SSD Streaming Error: \(error.localizedDescription). (The model safetensors file may be corrupted, truncated, or incomplete).") | ||
| } |
There was a problem hiding this comment.
PR description says the new wrapper will trigger a single clean crash "on the main thread", but check() calls fatalError directly on whichever thread invokes check(). If the main-thread behavior is required, check() (or its callers) should explicitly dispatch/ensure main-thread execution; otherwise the PR description should be updated to match the actual behavior.
| let path = "dummy_corrupt.safetensors" | ||
| let dict: [String: MLXArray] = ["weight": MLXArray.zeros([256, 1024])] | ||
| try MLX.save(arrays: dict, url: URL(fileURLWithPath: path)) | ||
|
|
||
| let fd = open(path, O_RDWR) | ||
| ftruncate(fd, 100) | ||
| close(fd) |
There was a problem hiding this comment.
This test writes dummy_corrupt.safetensors into the current working directory and never removes it. This can leave artifacts in the repo/build folder and can cause conflicts when tests are re-run or executed concurrently. Consider writing into FileManager.default.temporaryDirectory with a UUID-based filename and cleaning up the file in a defer block.
| ftruncate(fd, 100) | ||
| close(fd) | ||
|
|
There was a problem hiding this comment.
The return values from open, ftruncate, and close are ignored. If open fails (fd == -1) or ftruncate fails, the test may behave unpredictably or pass/fail for the wrong reason. Please assert/expect on these return codes (or use FileHandle APIs) so the test fails with a clear message when the file can't be truncated.
| ftruncate(fd, 100) | |
| close(fd) | |
| #expect(fd != -1) | |
| guard fd != -1 else { return } | |
| defer { | |
| let closeResult = close(fd) | |
| #expect(closeResult == 0) | |
| } | |
| let truncateResult = ftruncate(fd, 100) | |
| #expect(truncateResult == 0) | |
| guard truncateResult == 0 else { return } |
Resolves SharpAI/SwiftLM#87
Description
If a model's
.safetensorsfile is truncated or corrupted,MLXFast.preadIntoOffsetwill throw anstd::runtime_errorfrom C++. The MLX C API catches this and invokes the SwiftErrorHandler.dispatch.Since
SwitchGLUusesDispatchQueue.concurrentPerformto parallelize expert loading, multiple GCD threads were triggeringfatalErrorsimultaneously. This caused a global lock contention (deadlock) during the crash, freezing the SwiftBuddy process.Fix
Added a
ThreadSafeErrorwrapper that usesMLX.withErrorinside the GCD blocks to safely intercept the MLX exceptions. This ensures that a corrupted file will trigger a single, clean crash with an explicitMLX SSD Streaming Errormessage on the main thread, rather than deadlocking.