Repository navigation
fix custom io reader/writer API - #130
Merged
Merged
Conversation
davidkoski
approved these changes
Sep 10, 2026
davidkoski
left a comment
davidkoski
left a comment
Member
There was a problem hiding this comment.
Looks good, see ml-explore/mlx-swift#476
aleroot
added a commit
to aleroot/mlx-swift
that referenced
this pull request
Sep 15, 2026
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.
3 of 4 tasks
aleroot
added a commit
to aleroot/mlx-swift
that referenced
this pull request
Oct 13, 2026
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.
davidkoski
pushed a commit
to ml-explore/mlx-swift
that referenced
this pull request
Oct 21, 2026
* 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.
* 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.
* 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.
* Harden progress reads and cover lazy handler lifetime
8 of 9 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.