[davidtai]

Summary

mx.compile(fun) in the Python binding keys its traces by the address of the Python callable (PyCompiledFun::fun_id). Compile caches are thread-local, and each call records its trace in the cache of the calling thread. call_impl stored the cache of the calling thread in PyCompiledFun::cache, so cache was the cache of the thread that called the function last. ~PyCompiledFun called compile_erase(cache, fun_id) and erased only that cache. Traces that other threads recorded stayed after the free:

• They keep their arrays alive (inputs, outputs, tape, captured constants) until that thread exits.
• The key is the address of a Python object, which Python reuses. A later mx.compile of a callable at the same address, called on that thread with the same input structure, shapes, dtypes and default stream, replays the old trace and does not trace the new function.

Fix: the same as mlx #43 for compile(std::function). The destructor erases the key from every live thread cache through the fork's cache registry (detail::compile_erase(fun_id), PR #10), before it releases fun. PyCompiledFun no longer keeps a cache, and call_impl no longer reads one on each call. A moved-from object (fun_id == 0) erases nothing.

This changes only the Python binding (python/src). mlx-c and mlx-swift do not compile it, so no pin bump is needed in mlx-c, mlx-swift or mlx-swift-lm.

Changes

• python/src/transforms.cpp: ~PyCompiledFun calls mx::detail::compile_erase(fun_id); the cache member and its assignment in call_impl are removed.
• python/tests/test_compile.py: test_compile_free_erases_every_thread_cache. A worker thread makes its default stream, then calls the compiled function two times. The test checks that the second call used the trace (one trace), which shows that the worker finds the entry. The main thread then calls the function (so the main thread called it last), frees it, and the worker compiles the same callable again and calls it. The test expects a new trace there (3 traces in all). It uses only the trace count, the same observable as upstream's test_compile_release_on_another_thread. No test-only production API.
• fork.yaml: section "Freed Python compiled functions leave every thread cache".

Test first

• f40a6f60 (first test version): passed without the fix (Fork Tests https://github.com/Layr-Labs/mlx/actions/runs/40117260246). The reason: a trace is cached under the default stream of the calling thread (peek_default_stream in CompileCache::find), and the first op of a new thread makes that stream. The worker's first trace was stored before its stream existed, so no later call on the worker could match it. The compile after the free traced again with or without the erase.
• ee8eae07 (test fixed, no fix): FAILED on the new test only, AssertionError: 2 != 3, with the GPU and with the CPU as the default device: Fork Tests https://github.com/Layr-Labs/mlx/actions/runs/40118257732 (python-gpu (1) job 120944438596, python-cpu (1) job 120944438617).
• bef32f93 (fix): ALL GREEN. Fork Tests https://github.com/Layr-Labs/mlx/actions/runs/40119450075 (the new test passed with the GPU, job 120948067665, and with the CPU, job 120948067783), Build and Test https://github.com/Layr-Labs/mlx/actions/runs/40119449948, Fork Diff Check https://github.com/Layr-Labs/mlx/actions/runs/40119450076.

Note: upstream's test_compile_release_on_another_thread has the same property as the first version of this test. Its worker traces before its default stream exists, so the test passes whether or not the release erases the entry. It is unchanged here.

Draft upstream issue (not opened)

Python mx.compile: freeing a compiled function leaves its traces in other threads' compile caches

PyCompiledFun (python/src/transforms.cpp) keys traces by fun_id, the address of the Python callable. Since compile caches became thread-local, each call records its trace in the cache of the calling thread. call_impl sets cache = detail::compile_cache() on every call, and ~PyCompiledFun calls detail::compile_erase(cache, fun_id). Only the cache of the thread that called the function last is erased.

Effects: (1) traces on other threads keep their arrays alive until those threads exit, which leaks in a long-lived thread pool; (2) Python reuses object addresses, so a new callable can get the same fun_id and replay an old trace on a thread that still holds it (same input structure, shapes, dtypes and default stream).

Repro (Python): on a worker thread, call mx.default_stream(mx.default_device()), then call f = mx.compile(fun) once (one trace). Call f on the main thread, then del f; gc.collect() there. On the still-running worker, call mx.compile(fun)(x): it does not trace again (expected: a new trace).

Note on testing: test_compile_release_on_another_thread does not detect this class of bug. Its worker calls the compiled function before the worker's default stream exists, so the trace is stored with no stream and no later call can match it.

Possible fix: keep a process-wide registry of live thread caches and erase fun_id from every one of them in ~PyCompiledFun, before the callable is released. Layr-Labs/mlx carries this registry and fix.

🤖 Generated with Claude Code