Skip to content

hold evalLock during mlx_save_safetensors etc#379

Merged
davidkoski merged 1 commit into
mainfrom
save-unsafe
Apr 1, 2026
Merged

hold evalLock during mlx_save_safetensors etc#379
davidkoski merged 1 commit into
mainfrom
save-unsafe

Conversation

@davidkoski

@davidkoski davidkoski commented Mar 27, 2026

Copy link
Copy Markdown
Collaborator

Proposed changes

Verified via concurrent tests in mlx-swift-lm.

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

- fix #378
- save functions call eval internally and must be run under evalLock
@davidkoski davidkoski requested a review from angeloskath March 27, 2026 22:52
Comment thread Source/MLX/IO.swift
_ = try withError {
mlx_save_safetensors(path.cString(using: .utf8), mlx_arrays, mlx_metadata)
_ = evalLock.withLock {
mlx_save_safetensors(path.cString(using: .utf8), mlx_arrays, mlx_metadata)

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

internally these functions prepare contiguous arrays and use eval to realize them. eval is not thread safe so the entire thing has to be run under the evalLock.

See #378 for the use case.

davidkoski added a commit to ml-explore/mlx-swift-lm that referenced this pull request Mar 27, 2026
- some concurrency issues snuck in for test support
- fix and prevent more from coming in by switching Package to swift 6
- depends on ml-explore/mlx-swift#379

@angeloskath angeloskath left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks great!

@davidkoski davidkoski merged commit 61b9e01 into main Apr 1, 2026
7 checks passed
@davidkoski davidkoski deleted the save-unsafe branch April 1, 2026 16:45
davidkoski added a commit to ml-explore/mlx-swift-lm that referenced this pull request Apr 1, 2026
…es (#165)

* switch to swift 6 -- prevent concurrency issues, fix concurrency issues

- some concurrency issues snuck in for test support
- fix and prevent more from coming in by switching Package to swift 6
- depends on ml-explore/mlx-swift#379

* pick up mlx-swift with fix for save_safetensors concurrency issue
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] mlx_save_safetensors is not thread safe

2 participants