Regression in 8e00a2d9d (#4248): fatal PyThreadState_Get abort at interpreter exit after mx.compile is used — compile-cache cleanup no longer runs before finalization

[jrp2014]

Summary

Since 8e00a2d9d "Make mx.compile cache erasing thread safe (#4248)", any process that populates the compile cache (e.g. one mlx_vlm.generate() call) and then exits normally aborts with SIGABRT (exit 134) during interpreter finalization, after producing correct output:

Fatal Python error: PyThreadState_Get: the function must be called with the GIL held, after Python initialization and before Python finalization, but the GIL is released (the current Python thread state is NULL)
Python runtime state: finalizing (tstate=0x000000010324ba90)

Bisected: d9e2b0d40 (parent) → clean, exit 0; 8e00a2d9d → abort, exit 134; still present at 06f154bcf (current main). mlx-vlm at 830c7f7e held constant throughout.

Cause (from the diff)

#4248 removed both mechanisms that previously guaranteed the compile cache was cleared before the interpreter finalized:

1. ensure_compile_cache_cleanup() in python/src/transforms.cpp — a thread_local guard whose destructor did nb::gil_scoped_acquire gil; mx::detail::compile_clear_cache();, commented "Make sure each thread using mx.compile would clear its compile cache before python interpreter exits."
2. The atexit.register(compile_clear_cache) hook, commented "Ensure the main thread cleanup will happen before the interpreter goes away… clean tear-down."
   With those gone, the (now shared_ptr-held) CompileCache — whose CacheEntry objects hold Python callables (fun.ptr()) — is destroyed by ordinary C++ static destruction, after finalization. Dropping those nb::object refs then calls into the interpreter with a NULL thread state → the abort. It only reproduces after something has populated the cache, which is why simple mx.array scripts are unaffected and model inference is not.

Reproduction

from mlx_vlm import load, generate          # any mx.compile user will do
m, p = load("mlx-community/MiniCPM-V-4.6-8bit")
print(generate(m, p, prompt="Hi", max_tokens=5, verbose=False).text)
# correct output printed, then: Fatal Python error ... exit 134

python -c "import mlx.core as mx; print((mx.array([1.])*2).tolist())" → clean (cache never populated). Environment: macOS 26.6.1, M5 Max, Python 3.13.14, mlx built from source (0.32.1.dev20260817+8e00a2d9d), nanobind 2.14.0.

Suggested fix

Re-establish the pre-finalization clear — the atexit registration (or an equivalent Py_AtExit/nanobind leak-safe hook) that clears the compile cache while the GIL can still be acquired — on top of the new thread-safe cache. Alternatively, ensure CacheEntry releases its Python references under the GIL rather than in a static destructor.

[zcbenz]

Currently it is required to call mx.clear_streams() at the end of thread for manual cleanup, including the main thread, it is impossible to make automatic cleanup perfectly because it always run too late, after python interpreter and cuda runtime shuts down. An example can be found:

mlx/tests/scheduler_tests.cpp

Lines 89 to 101 in d5841be

	TEST_CASE("test new stream in threads") {
	std::vector<std::thread> threads;
	for (int i = 0; i < 1; ++i) {
	threads.emplace_back([]() {
	auto s = new_stream(default_device());
	eval(arange(10, s));
	clear_streams();
	});
	}
	for (auto& t : threads) {
	t.join();
	}
	}

[zcbenz]

I added #4373 so it would not be required to call mx.clear_streams() in the main thread as I see may apps are crashing because of it, though it is still required to call mx.clear_streams() in each worker thread.