From 615f9a85fc47392c267a92041c4ac5197bbd516d Mon Sep 17 00:00:00 2001 From: Alessio Pollero Date: Sat, 20 Jun 2026 18:08:22 +0400 Subject: [PATCH 1/3] Add byte-progress reporting for safetensors file loading Introduces LoadProgress and new loadArrays(url:stream:progressHandler:) and loadArraysAndMetadata(url:stream:progressHandler:) overloads that report bytes read as lazy arrays are evaluated. Uses a custom mlx_io_reader vtable backed by pread() so progress callbacks can be invoked from MLX worker threads. Includes a unit test verifying monotonic progress from 0 to 1. --- .../MLX/Documentation.docc/free-functions.md | 2 + Source/MLX/IO.swift | 271 ++++++++++++++++++ Tests/MLXTests/SaveTests.swift | 180 ++++++++---- 3 files changed, 403 insertions(+), 50 deletions(-) diff --git a/Source/MLX/Documentation.docc/free-functions.md b/Source/MLX/Documentation.docc/free-functions.md index e0592dbec..f4e4a4503 100644 --- a/Source/MLX/Documentation.docc/free-functions.md +++ b/Source/MLX/Documentation.docc/free-functions.md @@ -120,7 +120,9 @@ operations as methods for convenience. - ``loadArray(url:stream:)`` - ``loadArrays(url:stream:)`` +- ``loadArrays(url:stream:progressHandler:)`` - ``loadArraysAndMetadata(url:stream:)`` +- ``loadArraysAndMetadata(url:stream:progressHandler:)`` - ``save(array:url:stream:)`` - ``save(arrays:metadata:url:stream:)`` diff --git a/Source/MLX/IO.swift b/Source/MLX/IO.swift index 83c897d3c..44355f62d 100644 --- a/Source/MLX/IO.swift +++ b/Source/MLX/IO.swift @@ -3,6 +3,25 @@ import Cmlx import Foundation +/// Byte-level progress for loading arrays from disk. +/// +/// `completedUnitCount` and `totalUnitCount` are bytes. Progress callbacks for a +/// single load are delivered in monotonically increasing order. +public struct LoadProgress: Sendable, Equatable { + public let completedUnitCount: Int64 + public let totalUnitCount: Int64 + + public var fractionCompleted: Double { + guard totalUnitCount > 0 else { return 0 } + return min(1, max(0, Double(completedUnitCount) / Double(totalUnitCount))) + } + + public init(completedUnitCount: Int64, totalUnitCount: Int64) { + self.completedUnitCount = completedUnitCount + self.totalUnitCount = totalUnitCount + } +} + public enum LoadSaveError: Error { case unableToOpen(URL, String) case unknownExtension(String) @@ -143,6 +162,27 @@ public func loadArrays(url: URL, stream: StreamOrDevice = .cpu) throws -> [Strin } } +/// Load dictionary of ``MLXArray`` from a `safetensors` file, reporting byte progress as +/// lazy arrays are evaluated. +/// +/// - Parameters: +/// - url: URL of file to load +/// - stream: stream or device to evaluate on +/// - progressHandler: progress callback. This may be called from MLX worker threads. +/// Progress is reported in byte chunks while the returned lazy arrays are evaluated. +/// +/// ### See Also +/// - ``loadArrays(url:stream:)`` +/// - ``loadArraysAndMetadata(url:stream:progressHandler:)`` +public func loadArrays( + url: URL, stream: StreamOrDevice = .cpu, + progressHandler: @Sendable @escaping (LoadProgress) -> Void +) throws -> [String: MLXArray] { + let (arrays, _) = try loadArraysAndMetadata( + url: url, stream: stream, progressHandler: progressHandler) + return arrays +} + /// Load dictionary of ``MLXArray`` and metadata `[String:String]` from a `safetensors` file. /// /// - Parameters: @@ -175,6 +215,44 @@ public func loadArraysAndMetadata(url: URL, stream: StreamOrDevice = .cpu) throw } } +/// Load dictionary of ``MLXArray`` and metadata from a `safetensors` file, reporting byte +/// progress as lazy arrays are evaluated. +/// +/// - Parameters: +/// - url: URL of file to load +/// - stream: stream or device to evaluate on +/// - progressHandler: progress callback. This may be called from MLX worker threads. +/// Progress is reported in byte chunks while the returned lazy arrays are evaluated. +/// +/// ### See Also +/// - ``loadArraysAndMetadata(url:stream:)`` +/// - ``loadArrays(url:stream:progressHandler:)`` +public func loadArraysAndMetadata( + url: URL, stream: StreamOrDevice = .cpu, + progressHandler: @Sendable @escaping (LoadProgress) -> Void +) throws -> ([String: MLXArray], [String: String]) { + precondition(url.isFileURL) + + switch url.pathExtension { + case "safetensors": + var r0 = mlx_map_string_to_array_new() + var r1 = mlx_map_string_to_string_new() + defer { mlx_map_string_to_array_free(r0) } + defer { mlx_map_string_to_string_free(r1) } + + let reader = try new_mlx_io_reader_fileIO(url, progressHandler: progressHandler) + defer { mlx_io_reader_free(reader) } + + _ = try withError { + mlx_load_safetensors_reader(&r0, &r1, reader, stream.ctx) + } + + return (mlx_map_array_values(r0), mlx_map_string_values(r1)) + default: + throw LoadSaveError.unknownExtension(url.pathExtension) + } +} + // MARK: - Memory I/O private class IOState { @@ -187,6 +265,156 @@ private class IOState { } } +private final class FileIOState { + private static let maximumReadChunkSize = 4 * 1024 * 1024 + + private let descriptor: CInt + private let lock = NSLock() + private let progressLock = NSLock() + private var offset: Int64 = 0 + private var completedUnitCount: Int64 = 0 + private var readError: String? + private let progressHandler: @Sendable (LoadProgress) -> Void + private let labelPointer: UnsafeMutablePointer + + let totalUnitCount: Int64 + + init(url: URL, progressHandler: @Sendable @escaping (LoadProgress) -> Void) throws { + let path = url.path(percentEncoded: false) + let descriptor = path.withCString { open($0, O_RDONLY) } + guard descriptor >= 0 else { + throw LoadSaveError.unableToOpen(url, String(cString: strerror(errno))) + } + + var statBuffer = stat() + guard fstat(descriptor, &statBuffer) == 0 else { + let message = String(cString: strerror(errno)) + close(descriptor) + throw LoadSaveError.unableToOpen(url, message) + } + + guard let labelPointer = strdup("file \(path)") else { + close(descriptor) + throw LoadSaveError.unableToOpen(url, String(cString: strerror(errno))) + } + + self.descriptor = descriptor + self.totalUnitCount = max(0, Int64(statBuffer.st_size)) + self.progressHandler = progressHandler + self.labelPointer = labelPointer + + progressHandler(.init(completedUnitCount: 0, totalUnitCount: totalUnitCount)) + } + + deinit { + close(descriptor) + free(labelPointer) + } + + var isOpen: Bool { + descriptor >= 0 + } + + var good: Bool { + lock.withLock { + readError == nil + } + } + + var label: UnsafePointer { + UnsafePointer(labelPointer) + } + + func tell() -> Int { + lock.withLock { + Int(offset) + } + } + + func seek(offset newOffset: Int64, whence: Int32) { + lock.withLock { + switch whence { + case SEEK_SET: + offset = newOffset + case SEEK_CUR: + offset += newOffset + case SEEK_END: + offset = totalUnitCount + newOffset + default: + break + } + } + } + + func read(to data: UnsafeMutablePointer?, count: Int) { + guard let data else { return } + + let readOffset = lock.withLock { + offset + } + let bytesRead = read(to: data, count: count, offset: readOffset) + + lock.withLock { + offset += Int64(bytesRead) + } + } + + func read(to data: UnsafeMutablePointer?, count: Int, offset readOffset: Int64) { + guard let data else { return } + + _ = read(to: data, count: count, offset: readOffset) + } + + @discardableResult + private func read(to data: UnsafeMutablePointer, count: Int, offset readOffset: Int64) + -> Int + { + var totalRead = 0 + while totalRead < count { + let chunkSize = min(count - totalRead, Self.maximumReadChunkSize) + let bytesRead = pread( + descriptor, + UnsafeMutableRawPointer(data.advanced(by: totalRead)), + chunkSize, + off_t(readOffset + Int64(totalRead))) + guard bytesRead > 0 else { + recordReadError(bytesRead: bytesRead, requestedCount: count - totalRead) + break + } + totalRead += bytesRead + reportProgress(bytesRead: bytesRead) + } + + return totalRead + } + + private func recordReadError(bytesRead: Int, requestedCount: Int) { + let message: String + if bytesRead < 0 { + message = String(cString: strerror(errno)) + } else { + message = "unexpected end of file while reading \(requestedCount) bytes" + } + lock.withLock { + if readError == nil { + readError = message + } + } + } + + private func reportProgress(bytesRead: Int) { + guard bytesRead > 0 else { return } + + progressLock.withLock { + completedUnitCount = min(totalUnitCount, completedUnitCount + Int64(bytesRead)) + let progress = LoadProgress( + completedUnitCount: completedUnitCount, + totalUnitCount: totalUnitCount) + progressHandler(progress) + } + } +} + private let label: StaticString = "\0" private func getData(_ writer: mlx_io_writer) -> Data { @@ -262,6 +490,49 @@ private func new_mlx_io_reader_dataIO(_ data: Data) -> mlx_io_reader { return mlx_io_reader_new(ptr, new_mlx_io_vtable_dataIO()) } +private func new_mlx_io_vtable_fileIO() -> mlx_io_vtable { + mlx_io_vtable { ptr in + guard let ptr else { return false } + return Unmanaged.fromOpaque(ptr).takeUnretainedValue().isOpen + } good: { ptr in + guard let ptr else { return false } + let state = Unmanaged.fromOpaque(ptr).takeUnretainedValue() + return state.isOpen && state.good + } tell: { ptr in + let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() + return state.tell() + + } seek: { ptr, offset, whence in + let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() + state.seek(offset: Int64(offset), whence: whence) + + } read: { ptr, data, n in + let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() + state.read(to: data, count: n) + + } read_at_offset: { ptr, data, n, offset in + let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() + state.read(to: data, count: n, offset: Int64(offset)) + + } write: { _, _, _ in + + } label: { ptr in + let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() + return state.label + + } free: { ptr in + Unmanaged.fromOpaque(ptr!).release() + } +} + +private func new_mlx_io_reader_fileIO( + _ url: URL, progressHandler: @Sendable @escaping (LoadProgress) -> Void +) throws -> mlx_io_reader { + let ptr = Unmanaged.passRetained(try FileIOState(url: url, progressHandler: progressHandler)) + .toOpaque() + return mlx_io_reader_new(ptr, new_mlx_io_vtable_fileIO()) +} + private func new_mlx_io_writer_dataIO() -> mlx_io_writer { let ptr = Unmanaged.passRetained(IOState()).toOpaque() return mlx_io_writer_new(ptr, new_mlx_io_vtable_dataIO()) diff --git a/Tests/MLXTests/SaveTests.swift b/Tests/MLXTests/SaveTests.swift index 373ff8ca3..2676a7336 100644 --- a/Tests/MLXTests/SaveTests.swift +++ b/Tests/MLXTests/SaveTests.swift @@ -7,6 +7,23 @@ import MLX import XCTest +import os + +private final class ProgressRecorder: Sendable { + private let fractions = OSAllocatedUnfairLock(initialState: [Double]()) + + func record(_ progress: LoadProgress) { + fractions.withLock { values in + values.append(progress.fractionCompleted) + } + } + + var values: [Double] { + fractions.withLock { values in + values + } + } +} final class SaveTests: XCTestCase { @@ -16,7 +33,6 @@ final class SaveTests: XCTestCase { ) override func setUpWithError() throws { - setDefaultDevice() try FileManager.default.createDirectory( at: temporaryPath, withIntermediateDirectories: false @@ -28,72 +44,136 @@ final class SaveTests: XCTestCase { } public func testSaveArrays() throws { - let safetensorsPath = temporaryPath.appending( - path: "arrays.safetensors", - directoryHint: .notDirectory - ) + try MLX.Device.withDefaultDevice(.cpu) { + let safetensorsPath = temporaryPath.appending( + path: "arrays.safetensors", + directoryHint: .notDirectory + ) - let arrays: [String: MLXArray] = [ - "foo": MLX.ones([1, 2]), - "bar": MLX.zeros([2, 1]), - ] + let arrays: [String: MLXArray] = [ + "foo": MLX.ones([1, 2]), + "bar": MLX.zeros([2, 1]), + ] - try MLX.save(arrays: arrays, url: safetensorsPath) + try MLX.save(arrays: arrays, url: safetensorsPath) + + let loadedArrays = try MLX.loadArrays(url: safetensorsPath) + XCTAssertEqual(loadedArrays.keys.sorted(), arrays.keys.sorted()) + + assertEqual(try XCTUnwrap(loadedArrays["foo"]), try XCTUnwrap(arrays["foo"])) + assertEqual(try XCTUnwrap(loadedArrays["bar"]), try XCTUnwrap(arrays["bar"])) + } + } - let loadedArrays = try MLX.loadArrays(url: safetensorsPath) - XCTAssertEqual(loadedArrays.keys.sorted(), arrays.keys.sorted()) + public func testLoadArraysProgressReportsThroughEvaluation() throws { + try MLX.Device.withDefaultDevice(.cpu) { + let safetensorsPath = temporaryPath.appending( + path: "arrays.safetensors", + directoryHint: .notDirectory + ) + + let arrays: [String: MLXArray] = [ + "foo": MLX.ones([128, 128]), + "bar": MLX.zeros([64, 256]), + ] + try MLX.save(arrays: arrays, url: safetensorsPath) + + let recorder = ProgressRecorder() + let loadedArrays = try MLX.loadArrays( + url: safetensorsPath + ) { @Sendable progress in + recorder.record(progress) + } + + assertEqual(try XCTUnwrap(loadedArrays["foo"]), try XCTUnwrap(arrays["foo"])) + assertEqual(try XCTUnwrap(loadedArrays["bar"]), try XCTUnwrap(arrays["bar"])) + + let fractions = recorder.values + XCTAssertGreaterThan(fractions.count, 1) + XCTAssertEqual(fractions.first, 0) + XCTAssertEqual(fractions.last, 1) + XCTAssertEqual(fractions, fractions.sorted()) + } + } - assertEqual(try XCTUnwrap(loadedArrays["foo"]), try XCTUnwrap(arrays["foo"])) - assertEqual(try XCTUnwrap(loadedArrays["bar"]), try XCTUnwrap(arrays["bar"])) + public func testLoadArraysProgressFailsOnTruncatedTensorData() throws { + try MLX.Device.withDefaultDevice(.cpu) { + let safetensorsPath = temporaryPath.appending( + path: "truncated.safetensors", + directoryHint: .notDirectory + ) + + let arrays: [String: MLXArray] = [ + "foo": MLX.ones([128, 128]), + "bar": MLX.zeros([64, 256]), + ] + try MLX.save(arrays: arrays, url: safetensorsPath) + + var data = try Data(contentsOf: safetensorsPath) + data.removeLast(32) + try data.write(to: safetensorsPath) + + let loadedArrays = try MLX.loadArrays(url: safetensorsPath) { _ in } + + XCTAssertThrowsError( + try checkedEval(Array(loadedArrays.values) as [Any]) + ) + } } public func testSaveArray() throws { - // single array npy file - let path = temporaryPath.appending( - path: "array.npy", - directoryHint: .notDirectory - ) + try MLX.Device.withDefaultDevice(.cpu) { + // single array npy file + let path = temporaryPath.appending( + path: "array.npy", + directoryHint: .notDirectory + ) - let array = MLX.ones([2, 4]) + let array = MLX.ones([2, 4]) - try MLX.save(array: array, url: path) + try MLX.save(array: array, url: path) - let loaded = try MLX.loadArray(url: path) + let loaded = try MLX.loadArray(url: path) - assertEqual(array, loaded) + assertEqual(array, loaded) + } } public func testSaveArraysData() throws { - let arrays: [String: MLXArray] = [ - "foo": MLX.ones([1, 2]), - "bar": MLX.zeros([2, 1]), - ] - - let data = try saveToData(arrays: arrays) - let loadedArrays = try loadArrays(data: data) - XCTAssertEqual(loadedArrays.keys.sorted(), arrays.keys.sorted()) - - assertEqual(try XCTUnwrap(loadedArrays["foo"]), try XCTUnwrap(arrays["foo"])) - assertEqual(try XCTUnwrap(loadedArrays["bar"]), try XCTUnwrap(arrays["bar"])) + try MLX.Device.withDefaultDevice(.cpu) { + let arrays: [String: MLXArray] = [ + "foo": MLX.ones([1, 2]), + "bar": MLX.zeros([2, 1]), + ] + + let data = try saveToData(arrays: arrays) + let loadedArrays = try loadArrays(data: data) + XCTAssertEqual(loadedArrays.keys.sorted(), arrays.keys.sorted()) + + assertEqual(try XCTUnwrap(loadedArrays["foo"]), try XCTUnwrap(arrays["foo"])) + assertEqual(try XCTUnwrap(loadedArrays["bar"]), try XCTUnwrap(arrays["bar"])) + } } public func testSaveArraysMetadataData() throws { - let arrays: [String: MLXArray] = [ - "foo": MLX.ones([1, 2]), - "bar": MLX.zeros([2, 1]), - ] - let metadata = [ - "key": "value", - "key2": "value2", - ] - - let data = try saveToData(arrays: arrays, metadata: metadata) - let (loadedArrays, loadedMetadata) = try loadArraysAndMetadata(data: data) - XCTAssertEqual(loadedArrays.keys.sorted(), arrays.keys.sorted()) - - assertEqual(try XCTUnwrap(loadedArrays["foo"]), try XCTUnwrap(arrays["foo"])) - assertEqual(try XCTUnwrap(loadedArrays["bar"]), try XCTUnwrap(arrays["bar"])) - XCTAssertEqual(loadedMetadata, metadata) + try MLX.Device.withDefaultDevice(.cpu) { + let arrays: [String: MLXArray] = [ + "foo": MLX.ones([1, 2]), + "bar": MLX.zeros([2, 1]), + ] + let metadata = [ + "key": "value", + "key2": "value2", + ] + + let data = try saveToData(arrays: arrays, metadata: metadata) + let (loadedArrays, loadedMetadata) = try loadArraysAndMetadata(data: data) + XCTAssertEqual(loadedArrays.keys.sorted(), arrays.keys.sorted()) + + assertEqual(try XCTUnwrap(loadedArrays["foo"]), try XCTUnwrap(arrays["foo"])) + assertEqual(try XCTUnwrap(loadedArrays["bar"]), try XCTUnwrap(arrays["bar"])) + XCTAssertEqual(loadedMetadata, metadata) + } } /// `loadArrays(data:)` seeks to the end to size its input, so the in-memory IO From 87c8a7ba2bf2ad7b592a70f96fab28cb008ae39f Mon Sep 17 00:00:00 2001 From: Alessio Pollero Date: Tue, 18 Aug 2026 18:50:24 +0400 Subject: [PATCH 2/3] Add a scoped load-progress handler and fix memory reader SEEK_END `loadWeights()` style helpers -- including the one in mlx-swift-lm -- call the plain `loadArrays(url:)` / `loadArraysAndMetadata(url:)`, so a per-call `progressHandler:` argument can never reach them without changing every caller along the way. Add `withLoadProgressHandler(_:_:)` (sync and async), a task local scoped handler in the style of `withErrorHandler(_:_:)`. The plain file loading functions report byte progress to it when one is installed, so an application can drive a precise model loading progress bar around code it does not own: let container = try await withLoadProgressHandler({ tracker.update($0) }) { try await factory.loadContainer(from: directory, using: tokenizerLoader) } `LoadProgress` gains the `url` of the file being read so progress can be aggregated across the shards of a sharded model. Also fix the SEEK_END case of the in-memory reader, which moved the offset relative to the current position instead of the end of the data. mlx 0.32.1 seeks to the end of the stream to validate the tensor data offsets against the size of the file, so `loadArrays(data:)` would fail there ("The JSON header is N bytes long but the file is only 8 bytes"). Finally, restructure the truncated file test: a truncated file may now be reported either eagerly, while the header is parsed, or lazily, when the arrays are evaluated, and neither happens with the currently vendored mlx/mlx-c. --- .../MLX/Documentation.docc/free-functions.md | 2 + Source/MLX/IO.swift | 96 ++++++++++++- Tests/MLXTests/SaveTests.swift | 136 ++++++++++++++++-- 3 files changed, 222 insertions(+), 12 deletions(-) diff --git a/Source/MLX/Documentation.docc/free-functions.md b/Source/MLX/Documentation.docc/free-functions.md index f4e4a4503..9199e5458 100644 --- a/Source/MLX/Documentation.docc/free-functions.md +++ b/Source/MLX/Documentation.docc/free-functions.md @@ -125,6 +125,8 @@ operations as methods for convenience. - ``loadArraysAndMetadata(url:stream:progressHandler:)`` - ``save(array:url:stream:)`` - ``save(arrays:metadata:url:stream:)`` +- ``withLoadProgressHandler(_:_:)-(_,()throws->R)`` +- ``withLoadProgressHandler(_:_:)-(_,()async throws->R)`` ### Logical diff --git a/Source/MLX/IO.swift b/Source/MLX/IO.swift index 44355f62d..876c950b7 100644 --- a/Source/MLX/IO.swift +++ b/Source/MLX/IO.swift @@ -5,9 +5,14 @@ import Foundation /// Byte-level progress for loading arrays from disk. /// -/// `completedUnitCount` and `totalUnitCount` are bytes. Progress callbacks for a -/// single load are delivered in monotonically increasing order. +/// `completedUnitCount` and `totalUnitCount` are bytes and describe a single file. +/// Progress callbacks for a given file are delivered in monotonically increasing +/// order, but callbacks for _different_ files may interleave -- use ``url`` to +/// aggregate progress when loading a model made of several `safetensors` shards. public struct LoadProgress: Sendable, Equatable { + /// The file being read. + public let url: URL + public let completedUnitCount: Int64 public let totalUnitCount: Int64 @@ -16,12 +21,76 @@ public struct LoadProgress: Sendable, Equatable { return min(1, max(0, Double(completedUnitCount) / Double(totalUnitCount))) } - public init(completedUnitCount: Int64, totalUnitCount: Int64) { + public init(url: URL, completedUnitCount: Int64, totalUnitCount: Int64) { + self.url = url self.completedUnitCount = completedUnitCount self.totalUnitCount = totalUnitCount } } +/// Holder for the scoped ``withLoadProgressHandler(_:_:)-(_,()throws->R)`` handler. +enum LoadProgressHandler { + + /// The stack of installed handlers -- the innermost scope wins. + @TaskLocal + static var handlers: [@Sendable (LoadProgress) -> Void] = [] + + static var current: (@Sendable (LoadProgress) -> Void)? { + handlers.last + } +} + +/// Evaluate the block with a scoped byte-progress handler for file loads. +/// +/// Any ``loadArrays(url:stream:)`` or ``loadArraysAndMetadata(url:stream:)`` performed +/// inside `body` reports byte progress to `handler`, without the call site having to pass +/// a progress handler explicitly. This makes it possible to drive a precise loading +/// progress bar for code -- such as a model loading library -- that you do not control: +/// +/// ```swift +/// let tracker = LoadProgressTracker(totalBytes: totalBytesOfSafetensors(in: directory)) +/// let model = try withLoadProgressHandler({ tracker.update($0) }) { +/// try loadModel(from: directory) +/// } +/// ``` +/// +/// Loading is lazy: progress is reported as the returned arrays are evaluated, so `body` +/// should include the evaluation of the loaded arrays. Arrays that are never evaluated +/// are never read, so the reported progress may legitimately stop short of the file size. +/// +/// - Note: `handler` is called from MLX worker threads, potentially concurrently for +/// different files, and is on the critical path of the read. It should be cheap and +/// must not call back into MLX loading. +/// +/// - Parameters: +/// - handler: the scoped progress handler +/// - body: the code where the handler is to be active +/// +/// ### See Also +/// - ``loadArrays(url:stream:progressHandler:)`` +public func withLoadProgressHandler( + _ handler: @escaping @Sendable (LoadProgress) -> Void, _ body: () throws -> R +) rethrows -> R { + try LoadProgressHandler.$handlers.withValue(LoadProgressHandler.handlers + [handler]) { + try body() + } +} + +/// Evaluate the block with a scoped byte-progress handler for file loads (async). +/// +/// See ``withLoadProgressHandler(_:_:)-(_,()throws->R)`` for details. +/// +/// - Parameters: +/// - handler: the scoped progress handler +/// - body: the code where the handler is to be active +public func withLoadProgressHandler( + _ handler: @escaping @Sendable (LoadProgress) -> Void, _ body: () async throws -> R +) async rethrows -> R { + try await LoadProgressHandler.$handlers.withValue(LoadProgressHandler.handlers + [handler]) { + try await body() + } +} + public enum LoadSaveError: Error { case unableToOpen(URL, String) case unknownExtension(String) @@ -136,6 +205,9 @@ public func loadArray(url: URL, stream: StreamOrDevice = .cpu) throws -> MLXArra /// - url: URL of file to load /// - stream: stream or device to evaluate on /// +/// - Note: when a scoped progress handler is installed with +/// ``withLoadProgressHandler(_:_:)-(_,()throws->R)`` this reports byte progress to it. +/// /// ### See Also /// - ``loadArray(url:stream:)`` /// - ``loadArraysAndMetadata(url:stream:)`` @@ -147,6 +219,10 @@ public func loadArrays(url: URL, stream: StreamOrDevice = .cpu) throws -> [Strin switch url.pathExtension { case "safetensors": + if let progressHandler = LoadProgressHandler.current { + return try loadArrays(url: url, stream: stream, progressHandler: progressHandler) + } + var r0 = mlx_map_string_to_array_new() var r1 = mlx_map_string_to_string_new() defer { mlx_map_string_to_array_free(r0) } @@ -189,6 +265,9 @@ public func loadArrays( /// - url: URL of file to load /// - stream: stream or device to evaluate on /// +/// - Note: when a scoped progress handler is installed with +/// ``withLoadProgressHandler(_:_:)-(_,()throws->R)`` this reports byte progress to it. +/// /// ### See Also /// - ``loadArrays(url:stream:)`` /// - ``loadArray(url:stream:)`` @@ -200,6 +279,11 @@ public func loadArraysAndMetadata(url: URL, stream: StreamOrDevice = .cpu) throw switch url.pathExtension { case "safetensors": + if let progressHandler = LoadProgressHandler.current { + return try loadArraysAndMetadata( + url: url, stream: stream, progressHandler: progressHandler) + } + var r0 = mlx_map_string_to_array_new() var r1 = mlx_map_string_to_string_new() defer { mlx_map_string_to_array_free(r0) } @@ -276,6 +360,7 @@ private final class FileIOState { private var readError: String? private let progressHandler: @Sendable (LoadProgress) -> Void private let labelPointer: UnsafeMutablePointer + private let url: URL let totalUnitCount: Int64 @@ -302,8 +387,10 @@ private final class FileIOState { self.totalUnitCount = max(0, Int64(statBuffer.st_size)) self.progressHandler = progressHandler self.labelPointer = labelPointer + self.url = url - progressHandler(.init(completedUnitCount: 0, totalUnitCount: totalUnitCount)) + progressHandler( + .init(url: url, completedUnitCount: 0, totalUnitCount: totalUnitCount)) } deinit { @@ -408,6 +495,7 @@ private final class FileIOState { progressLock.withLock { completedUnitCount = min(totalUnitCount, completedUnitCount + Int64(bytesRead)) let progress = LoadProgress( + url: url, completedUnitCount: completedUnitCount, totalUnitCount: totalUnitCount) progressHandler(progress) diff --git a/Tests/MLXTests/SaveTests.swift b/Tests/MLXTests/SaveTests.swift index 2676a7336..db3a44fcd 100644 --- a/Tests/MLXTests/SaveTests.swift +++ b/Tests/MLXTests/SaveTests.swift @@ -10,19 +10,41 @@ import XCTest import os private final class ProgressRecorder: Sendable { - private let fractions = OSAllocatedUnfairLock(initialState: [Double]()) + private let progress = OSAllocatedUnfairLock(initialState: [LoadProgress]()) func record(_ progress: LoadProgress) { - fractions.withLock { values in - values.append(progress.fractionCompleted) + self.progress.withLock { values in + values.append(progress) } } - var values: [Double] { - fractions.withLock { values in + var reported: [LoadProgress] { + progress.withLock { values in values } } + + var values: [Double] { + reported.map { $0.fractionCompleted } + } + + /// Fractions reported for a single file, in order. + func values(for url: URL) -> [Double] { + reported.filter { $0.url == url }.map { $0.fractionCompleted } + } + + /// Aggregate fraction across every file seen, by bytes. + var aggregateFraction: Double { + var completed = [URL: Int64]() + var total = [URL: Int64]() + for progress in reported { + completed[progress.url] = progress.completedUnitCount + total[progress.url] = progress.totalUnitCount + } + let totalBytes = total.values.reduce(0, +) + guard totalBytes > 0 else { return 0 } + return Double(completed.values.reduce(0, +)) / Double(totalBytes) + } } final class SaveTests: XCTestCase { @@ -113,11 +135,109 @@ final class SaveTests: XCTestCase { data.removeLast(32) try data.write(to: safetensorsPath) - let loadedArrays = try MLX.loadArrays(url: safetensorsPath) { _ in } - - XCTAssertThrowsError( + // A truncated file has to be reported either eagerly, while the header is + // parsed (mlx >= 0.32.1 validates the tensor data offsets against the size of + // the file), or lazily, when the arrays are evaluated and the read fails + // (ml-explore/mlx#3742 + ml-explore/mlx-c#126). + var thrownError: Error? + do { + let loadedArrays = try MLX.loadArrays(url: safetensorsPath) { _ in } try checkedEval(Array(loadedArrays.values) as [Any]) + } catch { + thrownError = error + } + + if thrownError == nil { + throw XCTSkip( + """ + the vendored mlx/mlx-c silently ignores a failed read from a custom \ + io reader -- requires mlx >= 0.32.1 (ml-explore/mlx#3742) and \ + ml-explore/mlx-c#126 + """) + } + } + } + + public func testScopedLoadProgressHandler() throws { + try MLX.Device.withDefaultDevice(.cpu) { + let safetensorsPath = temporaryPath.appending( + path: "scoped.safetensors", + directoryHint: .notDirectory + ) + + let arrays: [String: MLXArray] = [ + "foo": MLX.ones([128, 128]), + "bar": MLX.zeros([64, 256]), + ] + try MLX.save(arrays: arrays, url: safetensorsPath) + + let recorder = ProgressRecorder() + + // note: the plain loadArrays(url:) -- no progress handler passed at the call site + let loadedArrays = try withLoadProgressHandler({ @Sendable in recorder.record($0) }) { + let loadedArrays = try MLX.loadArrays(url: safetensorsPath) + MLX.eval(Array(loadedArrays.values)) + return loadedArrays + } + + assertEqual(try XCTUnwrap(loadedArrays["foo"]), try XCTUnwrap(arrays["foo"])) + + let fractions = recorder.values + XCTAssertGreaterThan(fractions.count, 1) + XCTAssertEqual(fractions.first, 0) + XCTAssertEqual(fractions.last, 1) + XCTAssertEqual(fractions, fractions.sorted()) + XCTAssertEqual(Set(recorder.reported.map(\.url)), [safetensorsPath]) + } + } + + public func testScopedLoadProgressHandlerIsScoped() throws { + try MLX.Device.withDefaultDevice(.cpu) { + let safetensorsPath = temporaryPath.appending( + path: "unscoped.safetensors", + directoryHint: .notDirectory ) + try MLX.save(arrays: ["foo": MLX.ones([128, 128])], url: safetensorsPath) + + let recorder = ProgressRecorder() + withLoadProgressHandler({ @Sendable in recorder.record($0) }) { + } + + // outside the scope nothing is reported + let loadedArrays = try MLX.loadArrays(url: safetensorsPath) + MLX.eval(Array(loadedArrays.values)) + + XCTAssertTrue(recorder.reported.isEmpty) + } + } + + public func testScopedLoadProgressAggregatesAcrossFiles() throws { + try MLX.Device.withDefaultDevice(.cpu) { + let shards = try (0 ..< 3).map { index -> URL in + let url = temporaryPath.appending( + path: "shard-\(index).safetensors", + directoryHint: .notDirectory + ) + try MLX.save(arrays: ["w\(index)": MLX.ones([64, 128])], url: url) + return url + } + + let recorder = ProgressRecorder() + try withLoadProgressHandler({ @Sendable in recorder.record($0) }) { + // this mimics a model loader: several shards loaded lazily, then evaluated + var weights = [String: MLXArray]() + for url in shards { + let (w, _) = try MLX.loadArraysAndMetadata(url: url) + weights.merge(w) { _, new in new } + } + MLX.eval(Array(weights.values)) + } + + XCTAssertEqual(Set(recorder.reported.map(\.url)), Set(shards)) + for url in shards { + XCTAssertEqual(recorder.values(for: url).last, 1) + } + XCTAssertEqual(recorder.aggregateFraction, 1, accuracy: 1e-9) } } From 5800740086532cd0057d6e4a431caa3e6c45b105 Mon Sep 17 00:00:00 2001 From: Alessio Pollero Date: Tue, 15 Sep 2026 14:32:10 +0400 Subject: [PATCH 3/3] Update for the mlx-c custom IO reader/writer API ml-explore/mlx-c#130 changed `mlx_io_vtable` so that the callbacks report whether they succeeded: int (*seek)(void*, int64_t off, int whence); size_t (*read)(void*, char* data, size_t n); size_t (*read_at_offset)(void*, char* data, size_t n, size_t off); size_t (*write)(void*, const char* data, size_t n); and `CReader`/`CWriter` now turn a negative seek or a short read/write into a thrown `std::runtime_error`. Bump the mlx-c submodule to that commit (plus the checked-in copies of `io_types.h` and the CMake `GIT_TAG`) and adapt the Swift side: - `FileIOState.seek` returns `0`/`-1` instead of silently ignoring a bad `whence` or a negative resulting offset. - `FileIOState.read` returns the number of bytes actually read, so a file that is truncated after its header was parsed now fails instead of leaving the destination buffer uninitialized. The private implementation is renamed `readBytes` so it cannot be confused with the two public overloads. - The in-memory reader reports an out-of-bounds read as `0` bytes rather than doing nothing, and its `seek` reports failure for an unknown `whence` or a negative offset. - The file reader's `write` reports `0` bytes, so using it as a writer errors out instead of silently discarding the data. This is the error-propagation path the load progress work was waiting on (ml-explore/mlx#3742 is already in the vendored mlx v0.32.2), so the truncated file test asserts a failure rather than skipping, and a new test truncates the file *after* the header is parsed to cover the lazy read path: [mlx_io_reader] unable to read 65536 bytes (read 65504 instead) in file ... Also fix the `withLoadProgressHandler` documentation links -- the `-(_,()throws->R)` disambiguation does not resolve and `verify-docs.sh` builds with `--warnings-as-errors` -- document why the reported progress is approximate (loading is lazy, so it may stop short of the file size and it counts bytes read rather than bytes covered), and add `LoadProgress` to the MLX topics. --- CMakeLists.txt | 2 +- .../Cmlx/include-framework/mlx-c-io_types.h | 8 +- Source/Cmlx/include/mlx/c/io_types.h | 8 +- Source/Cmlx/mlx-c | 2 +- Source/MLX/Documentation.docc/MLX.md | 4 + .../MLX/Documentation.docc/free-functions.md | 4 +- Source/MLX/IO.swift | 131 ++++++++++++------ Tests/MLXTests/SaveTests.swift | 83 +++++++++-- 8 files changed, 174 insertions(+), 68 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 41404a01d..5771d8525 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -30,7 +30,7 @@ endif() FetchContent_Declare( mlx-c GIT_REPOSITORY "https://github.com/ml-explore/mlx-c.git" - GIT_TAG "c74db5307cc8ce122f48d97ef951b30578674e7f") + GIT_TAG "2f8700a5797a4321597507ea57b3c54ffe4373af") FetchContent_MakeAvailable(mlx-c) # swift-numerics diff --git a/Source/Cmlx/include-framework/mlx-c-io_types.h b/Source/Cmlx/include-framework/mlx-c-io_types.h index 4eccc95a9..7b6ae483a 100644 --- a/Source/Cmlx/include-framework/mlx-c-io_types.h +++ b/Source/Cmlx/include-framework/mlx-c-io_types.h @@ -37,10 +37,10 @@ typedef struct mlx_io_vtable_ { bool (*is_open)(void*); bool (*good)(void*); size_t (*tell)(void*); - void (*seek)(void*, int64_t off, int whence); - void (*read)(void*, char* data, size_t n); - void (*read_at_offset)(void*, char* data, size_t n, size_t off); - void (*write)(void*, const char* data, size_t n); + int (*seek)(void*, int64_t off, int whence); + size_t (*read)(void*, char* data, size_t n); + size_t (*read_at_offset)(void*, char* data, size_t n, size_t off); + size_t (*write)(void*, const char* data, size_t n); const char* (*label)(void*); void (*free)(void*); } mlx_io_vtable; diff --git a/Source/Cmlx/include/mlx/c/io_types.h b/Source/Cmlx/include/mlx/c/io_types.h index 5382e7342..3f15873e3 100644 --- a/Source/Cmlx/include/mlx/c/io_types.h +++ b/Source/Cmlx/include/mlx/c/io_types.h @@ -37,10 +37,10 @@ typedef struct mlx_io_vtable_ { bool (*is_open)(void*); bool (*good)(void*); size_t (*tell)(void*); - void (*seek)(void*, int64_t off, int whence); - void (*read)(void*, char* data, size_t n); - void (*read_at_offset)(void*, char* data, size_t n, size_t off); - void (*write)(void*, const char* data, size_t n); + int (*seek)(void*, int64_t off, int whence); + size_t (*read)(void*, char* data, size_t n); + size_t (*read_at_offset)(void*, char* data, size_t n, size_t off); + size_t (*write)(void*, const char* data, size_t n); const char* (*label)(void*); void (*free)(void*); } mlx_io_vtable; diff --git a/Source/Cmlx/mlx-c b/Source/Cmlx/mlx-c index c74db5307..2f8700a57 160000 --- a/Source/Cmlx/mlx-c +++ b/Source/Cmlx/mlx-c @@ -1 +1 @@ -Subproject commit c74db5307cc8ce122f48d97ef951b30578674e7f +Subproject commit 2f8700a5797a4321597507ea57b3c54ffe4373af diff --git a/Source/MLX/Documentation.docc/MLX.md b/Source/MLX/Documentation.docc/MLX.md index 5e2f98c5c..7736e63cc 100644 --- a/Source/MLX/Documentation.docc/MLX.md +++ b/Source/MLX/Documentation.docc/MLX.md @@ -130,3 +130,7 @@ See for configuration, best practices, and policy guidance. - ``Device`` - ``DeviceType`` - ``Stream`` + +### Load Progress + +- ``LoadProgress`` diff --git a/Source/MLX/Documentation.docc/free-functions.md b/Source/MLX/Documentation.docc/free-functions.md index 9199e5458..6d3b011d6 100644 --- a/Source/MLX/Documentation.docc/free-functions.md +++ b/Source/MLX/Documentation.docc/free-functions.md @@ -125,8 +125,8 @@ operations as methods for convenience. - ``loadArraysAndMetadata(url:stream:progressHandler:)`` - ``save(array:url:stream:)`` - ``save(arrays:metadata:url:stream:)`` -- ``withLoadProgressHandler(_:_:)-(_,()throws->R)`` -- ``withLoadProgressHandler(_:_:)-(_,()async throws->R)`` +- ``withLoadProgressHandler(_:_:)-3ghip`` +- ``withLoadProgressHandler(_:_:)-8cm31`` ### Logical diff --git a/Source/MLX/IO.swift b/Source/MLX/IO.swift index 876c950b7..7814c0bd0 100644 --- a/Source/MLX/IO.swift +++ b/Source/MLX/IO.swift @@ -9,11 +9,25 @@ import Foundation /// Progress callbacks for a given file are delivered in monotonically increasing /// order, but callbacks for _different_ files may interleave -- use ``url`` to /// aggregate progress when loading a model made of several `safetensors` shards. +/// +/// Loading `safetensors` is lazy, which makes the reported progress approximate: +/// +/// - It may stop short of `totalUnitCount`. Only arrays that are actually evaluated are +/// read, so weights that are dropped before evaluation -- the ones a model's +/// `sanitize(weights:metadata:)` discards, for example -- are never read at all. Treat +/// the load returning as completion rather than waiting for ``fractionCompleted`` to +/// reach `1`. +/// - It counts bytes read rather than bytes of the file covered, so a region that happens +/// to be read more than once is counted more than once. ``completedUnitCount`` is +/// clamped to `totalUnitCount`, so ``fractionCompleted`` never exceeds `1`. public struct LoadProgress: Sendable, Equatable { /// The file being read. public let url: URL + /// Bytes read so far, clamped to ``totalUnitCount``. public let completedUnitCount: Int64 + + /// Size of the file in bytes. public let totalUnitCount: Int64 public var fractionCompleted: Double { @@ -28,7 +42,7 @@ public struct LoadProgress: Sendable, Equatable { } } -/// Holder for the scoped ``withLoadProgressHandler(_:_:)-(_,()throws->R)`` handler. +/// Holder for the scoped ``withLoadProgressHandler(_:_:)-3ghip`` handler. enum LoadProgressHandler { /// The stack of installed handlers -- the innermost scope wins. @@ -78,7 +92,7 @@ public func withLoadProgressHandler( /// Evaluate the block with a scoped byte-progress handler for file loads (async). /// -/// See ``withLoadProgressHandler(_:_:)-(_,()throws->R)`` for details. +/// See ``withLoadProgressHandler(_:_:)-3ghip`` for details. /// /// - Parameters: /// - handler: the scoped progress handler @@ -206,7 +220,7 @@ public func loadArray(url: URL, stream: StreamOrDevice = .cpu) throws -> MLXArra /// - stream: stream or device to evaluate on /// /// - Note: when a scoped progress handler is installed with -/// ``withLoadProgressHandler(_:_:)-(_,()throws->R)`` this reports byte progress to it. +/// ``withLoadProgressHandler(_:_:)-3ghip`` this reports byte progress to it. /// /// ### See Also /// - ``loadArray(url:stream:)`` @@ -245,11 +259,8 @@ public func loadArrays(url: URL, stream: StreamOrDevice = .cpu) throws -> [Strin /// - url: URL of file to load /// - stream: stream or device to evaluate on /// - progressHandler: progress callback. This may be called from MLX worker threads. -/// Progress is reported in byte chunks while the returned lazy arrays are evaluated. -/// -/// ### See Also -/// - ``loadArrays(url:stream:)`` -/// - ``loadArraysAndMetadata(url:stream:progressHandler:)`` +/// Progress is reported in byte chunks while the returned lazy arrays are evaluated, +/// so it may stop short of the size of the file -- see ``LoadProgress``. public func loadArrays( url: URL, stream: StreamOrDevice = .cpu, progressHandler: @Sendable @escaping (LoadProgress) -> Void @@ -266,7 +277,7 @@ public func loadArrays( /// - stream: stream or device to evaluate on /// /// - Note: when a scoped progress handler is installed with -/// ``withLoadProgressHandler(_:_:)-(_,()throws->R)`` this reports byte progress to it. +/// ``withLoadProgressHandler(_:_:)-3ghip`` this reports byte progress to it. /// /// ### See Also /// - ``loadArrays(url:stream:)`` @@ -306,7 +317,8 @@ public func loadArraysAndMetadata(url: URL, stream: StreamOrDevice = .cpu) throw /// - url: URL of file to load /// - stream: stream or device to evaluate on /// - progressHandler: progress callback. This may be called from MLX worker threads. -/// Progress is reported in byte chunks while the returned lazy arrays are evaluated. +/// Progress is reported in byte chunks while the returned lazy arrays are evaluated, +/// so it may stop short of the size of the file -- see ``LoadProgress``. /// /// ### See Also /// - ``loadArraysAndMetadata(url:stream:)`` @@ -418,44 +430,65 @@ private final class FileIOState { } } - func seek(offset newOffset: Int64, whence: Int32) { - lock.withLock { + /// Move the read position. + /// + /// Returns `0` on success and `-1` on failure; `CReader` turns a negative result into + /// a thrown error (ml-explore/mlx-c#130). + @discardableResult + func seek(offset newOffset: Int64, whence: Int32) -> Int32 { + lock.withLock { () -> Int32 in + let updated: Int64 switch whence { case SEEK_SET: - offset = newOffset + updated = newOffset case SEEK_CUR: - offset += newOffset + updated = offset + newOffset case SEEK_END: - offset = totalUnitCount + newOffset + // offset is relative to the end of the file, not the current position. + updated = totalUnitCount + newOffset default: - break + return -1 } + guard updated >= 0 else { return -1 } + offset = updated + return 0 } } - func read(to data: UnsafeMutablePointer?, count: Int) { - guard let data else { return } + /// Read `count` bytes at the current position and advance it. + /// + /// Returns the number of bytes actually read -- a short read makes `CReader` throw. + @discardableResult + func read(to data: UnsafeMutablePointer?, count: Int) -> Int { + guard let data else { return 0 } let readOffset = lock.withLock { offset } - let bytesRead = read(to: data, count: count, offset: readOffset) + let bytesRead = readBytes(to: data, count: count, offset: readOffset) lock.withLock { offset += Int64(bytesRead) } + + return bytesRead } - func read(to data: UnsafeMutablePointer?, count: Int, offset readOffset: Int64) { - guard let data else { return } + /// Read `count` bytes at an absolute `offset`, leaving the current position alone. + /// + /// Returns the number of bytes actually read -- a short read makes `CReader` throw. + /// This is the path used to materialize lazily loaded arrays, so it is called from + /// MLX worker threads. + @discardableResult + func read(to data: UnsafeMutablePointer?, count: Int, offset readOffset: Int64) -> Int { + guard let data else { return 0 } - _ = read(to: data, count: count, offset: readOffset) + return readBytes(to: data, count: count, offset: readOffset) } - @discardableResult - private func read(to data: UnsafeMutablePointer, count: Int, offset readOffset: Int64) - -> Int - { + private func readBytes( + to data: UnsafeMutablePointer, count: Int, offset readOffset: Int64 + ) -> Int { var totalRead = 0 while totalRead < count { let chunkSize = min(count - totalRead, Self.maximumReadChunkSize) @@ -524,39 +557,44 @@ private func new_mlx_io_vtable_dataIO() -> mlx_io_vtable { } seek: { ptr, offset, whence in let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() + let updated: Int switch whence { case SEEK_SET: - state.offset = Int(offset) + updated = Int(offset) case SEEK_CUR: - state.offset += Int(offset) + updated = state.offset + Int(offset) case SEEK_END: // offset is relative to the end of the data, not the current position. // mlx's load_safetensors uses seek(0, end) + tell() to size the input. - state.offset = state.data.count + Int(offset) + updated = state.data.count + Int(offset) default: - break + return -1 } + guard updated >= 0 else { return -1 } + state.offset = updated + return 0 + } read: { ptr, data, n in let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() - if n + state.offset <= state.data.count { - guard let data = data else { return } - _ = state.data.withUnsafeBytes { buffer in - memcpy(data, buffer.baseAddress!.advanced(by: state.offset), n) - } - state.offset += n + // report a short read as 0 bytes so that CReader throws rather than + // silently leaving the destination buffer uninitialized + guard let data, n > 0, n + state.offset <= state.data.count else { return 0 } + _ = state.data.withUnsafeBytes { buffer in + memcpy(data, buffer.baseAddress!.advanced(by: state.offset), n) } + state.offset += n + return n } read_at_offset: { ptr, data, n, offset in let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() - if n + offset <= state.data.count { - guard let data = data else { return } - _ = state.data.withUnsafeBytes { buffer in - memcpy(data, buffer.baseAddress!.advanced(by: offset), n) - } - state.offset = offset + guard let data, n > 0, n + offset <= state.data.count else { return 0 } + _ = state.data.withUnsafeBytes { buffer in + memcpy(data, buffer.baseAddress!.advanced(by: offset), n) } + state.offset = offset + return n } write: { ptr, data, n in let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() @@ -564,6 +602,7 @@ private func new_mlx_io_vtable_dataIO() -> mlx_io_vtable { let buffer = UnsafeBufferPointer(start: data, count: n) state.data.append(buffer) state.offset += n + return n } label: { ptr in UnsafeRawPointer(label.utf8Start).assumingMemoryBound(to: Int8.self) @@ -592,17 +631,19 @@ private func new_mlx_io_vtable_fileIO() -> mlx_io_vtable { } seek: { ptr, offset, whence in let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() - state.seek(offset: Int64(offset), whence: whence) + return state.seek(offset: Int64(offset), whence: whence) } read: { ptr, data, n in let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() - state.read(to: data, count: n) + return state.read(to: data, count: n) } read_at_offset: { ptr, data, n, offset in let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() - state.read(to: data, count: n, offset: Int64(offset)) + return state.read(to: data, count: n, offset: Int64(offset)) } write: { _, _, _ in + // read only -- reporting 0 bytes written makes CWriter throw + return 0 } label: { ptr in let state = Unmanaged.fromOpaque(ptr!).takeUnretainedValue() diff --git a/Tests/MLXTests/SaveTests.swift b/Tests/MLXTests/SaveTests.swift index db3a44fcd..26b47c5ed 100644 --- a/Tests/MLXTests/SaveTests.swift +++ b/Tests/MLXTests/SaveTests.swift @@ -137,24 +137,62 @@ final class SaveTests: XCTestCase { // A truncated file has to be reported either eagerly, while the header is // parsed (mlx >= 0.32.1 validates the tensor data offsets against the size of - // the file), or lazily, when the arrays are evaluated and the read fails - // (ml-explore/mlx#3742 + ml-explore/mlx-c#126). - var thrownError: Error? + // the file), or lazily, when the arrays are evaluated and the short read is + // turned into an error (ml-explore/mlx-c#130). do { let loadedArrays = try MLX.loadArrays(url: safetensorsPath) { _ in } try checkedEval(Array(loadedArrays.values) as [Any]) + XCTFail("a truncated safetensors file must not load successfully") } catch { - thrownError = error + // expected } + } + } + + /// mlx validates the tensor data offsets against the size of the file while it parses + /// the header, so a file that is _already_ truncated fails eagerly. A file that is + /// truncated after the header is parsed can only be caught when the lazy arrays are + /// evaluated and the read comes up short -- `mlx_io_reader` turns that into an error + /// rather than leaving the destination buffer uninitialized (ml-explore/mlx-c#130). + public func testLoadFailsWhenFileIsTruncatedBeforeEvaluation() throws { + try MLX.Device.withDefaultDevice(.cpu) { + let safetensorsPath = temporaryPath.appending( + path: "truncated-late.safetensors", + directoryHint: .notDirectory + ) + + let arrays: [String: MLXArray] = [ + "foo": MLX.ones([128, 128]), + "bar": MLX.zeros([64, 256]), + ] + try MLX.save(arrays: arrays, url: safetensorsPath) - if thrownError == nil { - throw XCTSkip( - """ - the vendored mlx/mlx-c silently ignores a failed read from a custom \ - io reader -- requires mlx >= 0.32.1 (ml-explore/mlx#3742) and \ - ml-explore/mlx-c#126 - """) + // the header is parsed and validated against the size of the file here, but + // the tensor data is only read once the lazy arrays are evaluated + let recorder = ProgressRecorder() + let loadedArrays = try MLX.loadArrays(url: safetensorsPath) { + @Sendable in recorder.record($0) + } + let size = try XCTUnwrap(recorder.reported.first?.totalUnitCount) + + // ... so truncate the file out from under them. `truncate()` shortens the + // inode the reader already has open, unlike a rewrite which may replace it. + XCTAssertEqual( + truncate(safetensorsPath.path(percentEncoded: false), off_t(size - 32)), 0) + + do { + try checkedEval(Array(loadedArrays.values) as [Any]) + XCTFail("evaluating arrays read from a truncated file must fail") + } catch { + // expected } + + // the bytes that were read are still accounted for, monotonically, and the + // aggregate stops short of the original size of the file + let fractions = recorder.values + XCTAssertEqual(fractions, fractions.sorted()) + XCTAssertEqual(fractions.first, 0) + XCTAssertLessThan(try XCTUnwrap(fractions.last), 1) } } @@ -312,4 +350,27 @@ final class SaveTests: XCTestCase { assertEqual(try XCTUnwrap(loaded["big"]), try XCTUnwrap(arrays["big"])) } + /// A truncated in-memory buffer must be reported as an error. The in-memory reader + /// reports a read that runs past the end of the buffer as zero bytes so that + /// `mlx_io_reader` throws instead of leaving the destination uninitialized + /// (ml-explore/mlx-c#130). + public func testLoadFromTruncatedDataFails() throws { + try MLX.Device.withDefaultDevice(.cpu) { + let arrays: [String: MLXArray] = [ + "big": MLX.ones([64, 64]) + ] + + var data = try saveToData(arrays: arrays) + data.removeLast(32) + + do { + let loaded = try loadArrays(data: data) + try checkedEval(Array(loaded.values) as [Any]) + XCTFail("a truncated safetensors buffer must not load successfully") + } catch { + // expected + } + } + } + }