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

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.