Skip to content

Make Scheduler::enqueue thread safe - #3423

Merged
zcbenz merged 1 commit into
ml-explore:mainfrom
zcbenz:scheduler-lock
Apr 20, 2026
Merged

zcbenz merged 1 commit into
ml-explore:mainfrom
zcbenz:scheduler-lock

Conversation

@zcbenz

@zcbenz zcbenz commented Apr 18, 2026

Copy link
Copy Markdown
Member

Fix a race condition caused by Scheduler::enqueue not using locks when reading threads_ which might be updated concurrently.

===============================================================================
/home/runner/work/mlx/mlx/tests/scheduler_tests.cpp:107:
TEST CASE:  test thread local stream

/home/runner/work/mlx/mlx/tests/scheduler_tests.cpp:107: FATAL ERROR: test case CRASHED: SIGSEGV - Segmentation violation signal

===============================================================================
[doctest] test cases:  212 |  211 passed | 1 failed | 20 skipped
[doctest] assertions: 3024 | 3024 passed | 0 failed |
[doctest] Status: FAILURE!
    #0 0xff050d0d3410 in __pthread_mutex_lock (/lib/aarch64-linux-gnu/libc.so.6+0x83410)
    #1 0xab086196f06c in __gthread_mutex_lock /usr/include/aarch64-linux-gnu/c++/11/bits/gthr-default.h:749
    #2 0xab086196f7fc in std::mutex::lock() /usr/include/c++/11/bits/std_mutex.h:100
    #3 0xab086197a204 in std::lock_guard<std::mutex>::lock_guard(std::mutex&) /usr/include/c++/11/bits/std_mutex.h:229
    #4 0xab08639fcdf8 in enqueue<mlx::core::Fence::wait(mlx::core::Stream, const mlx::core::array&)::<lambda()> > /home/runner/work/mlx/mlx/mlx/scheduler.h:56
    #5 0xab08639fcae0 in enqueue<mlx::core::Fence::wait(mlx::core::Stream, const mlx::core::array&)::<lambda()> > /home/runner/work/mlx/mlx/mlx/scheduler.h:124
    #6 0xab08639fc838 in enqueue<mlx::core::Fence::wait(mlx::core::Stream, const mlx::core::array&)::<lambda()> > /home/runner/work/mlx/mlx/mlx/scheduler.h:131
    #7 0xab08639fbf4c in mlx::core::Fence::wait(mlx::core::Stream, mlx::core::array const&) /home/runner/work/mlx/mlx/mlx/backend/no_gpu/fence.cpp:26
    #8 0xab08626e0610 in mlx::core::eval_impl(std::vector<mlx::core::array, std::allocator<mlx::core::array> >, bool) /home/runner/work/mlx/mlx/mlx/transforms.cpp:247
    #9 0xab08626e1fb0 in mlx::core::eval(std::vector<mlx::core::array, std::allocator<mlx::core::array> >) /home/runner/work/mlx/mlx/mlx/transforms.cpp:344
    #10 0xab08621ca684 in mlx::core::array::eval() /home/runner/work/mlx/mlx/mlx/array.cpp:157

@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.

👍

@zcbenz
zcbenz merged commit 1f5a413 into ml-explore:main Apr 20, 2026
16 checks passed
@zcbenz
zcbenz deleted the scheduler-lock branch April 20, 2026 05:30
@BrewTestBot BrewTestBot mentioned this pull request Apr 22, 2026
1 task done
ekryski added a commit to ekryski/mlx that referenced this pull request May 6, 2026
…dLocalStream, linalg.det)

Highlights:
- ThreadLocalStream in C++ (ml-explore#3405) — joins StreamOrDevice variant
- Segmented mm NAX kernel (ml-explore#3419) — new steel/gemm/kernels/steel_gemm_segmented{,_nax}
- mlx.core.linalg additions: determinant, sign-log-determinant (ml-explore#3416)
- NAX split-K tuning + addmm fix (ml-explore#3422)
- gather_mm CUDA + qmm CUDA fixes (ml-explore#3414, ml-explore#3417, ml-explore#3445)
- Scheduler::enqueue thread safety (ml-explore#3423)
- Random kernel bytes_per_key truncation fix (ml-explore#3432)
- Many smaller bug + CI fixes; nanobind 2.12.0 pin

Conflict resolution (5 files; all auto-merged cleanly, no manual edits):
- mlx/utils.h — `ThreadLocalStream` joins StreamOrDevice; our
  metal_retain_bound_buffers() helper preserved.
- mlx/backend/metal/device.cpp — upstream's "no Metal device available"
  error message coexists with our signpost/MLX_METAL_PROFILE perf
  diagnostic + hazard-mode retain logic.
- mlx/backend/metal/CMakeLists.txt — upstream's steel_gemm_segmented{,_nax}
  joins our kernel registrations.
- mlx/backend/metal/kernels/CMakeLists.txt — same: upstream segmented +
  our fused_gate_activation kernel.
- tests/scheduler_tests.cpp — upstream's "test thread local stream"
  joins our scheduler-test alignment for lazy command-encoder init.

No disputed conflicts; the auto-merge produced a coherent result with both
sides' work intact.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
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.

2 participants