From 8278ad94cee8290e610aab9df2b4c4a747bd99cb Mon Sep 17 00:00:00 2001 From: Marcos Ashton Iglesias Date: Tue, 6 Oct 2026 12:37:11 +0100 Subject: [PATCH 1/2] Fix cross-process race in CPU compile cache --- mlx/backend/cpu/compiled.cpp | 27 +++++++++++++++++++++++---- python/tests/test_compile.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 4 deletions(-) diff --git a/mlx/backend/cpu/compiled.cpp b/mlx/backend/cpu/compiled.cpp index f5b516e44b..d6a53a5287 100644 --- a/mlx/backend/cpu/compiled.cpp +++ b/mlx/backend/cpu/compiled.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include @@ -116,9 +117,14 @@ void* compile( } if (!lib_exists) { - // Open source file and write source code to it - std::string source_file_name = kernel_file_name + ".cpp"; - auto source_file_path = (output_dir / source_file_name).string(); + // Build under a unique name, then move the files into place. Other + // processes then never load a partly written library. + auto tmp_name = fmt::format("tmp_{:08x}", std::random_device{}()); + std::string source_file_name = tmp_name + ".cpp"; + std::string tmp_lib_name = + tmp_name + std::filesystem::path(shared_lib_name).extension().string(); + auto source_file_path = output_dir / source_file_name; + auto tmp_lib_path = output_dir / tmp_lib_name; std::ofstream source_file(source_file_path); source_file << source_code; @@ -127,7 +133,7 @@ void* compile( try { JitCompiler::exec( JitCompiler::build_command( - output_dir, source_file_name, shared_lib_name)); + output_dir, source_file_name, tmp_lib_name)); } catch (const std::exception& error) { throw std::runtime_error( fmt::format( @@ -135,6 +141,19 @@ void* compile( kernel_name, error.what())); } + + std::error_code error; + std::filesystem::rename( + source_file_path, output_dir / (kernel_file_name + ".cpp"), error); + std::filesystem::rename(tmp_lib_path, shared_lib_path, error); + if (error) { + // On Windows the rename can fail while another process uses the file. + if (std::filesystem::exists(shared_lib_path, error)) { + std::filesystem::remove(tmp_lib_path, error); + } else { + shared_lib_path = tmp_lib_path.string(); + } + } } // load library diff --git a/python/tests/test_compile.py b/python/tests/test_compile.py index 9f1ced1867..d6b2b251ff 100644 --- a/python/tests/test_compile.py +++ b/python/tests/test_compile.py @@ -4,7 +4,12 @@ import inspect import io import math +import os +import subprocess +import sys +import tempfile import threading +import unittest from functools import partial, wraps from io import StringIO from itertools import product @@ -1710,6 +1715,32 @@ def fun(x): x = mx.array([1.0, float("nan"), 3.0]) self.assertTrue(mx.array_equal(mx.compile(fun)(x), mx.array([False, True]))) + @unittest.skipIf(sys.platform == "win32", "Needs a POSIX shell") + def test_compile_cpu_partial_library(self): + # A build that stops after it creates the library file must not break + # the build in later processes. + code = ( + "import mlx.core as mx\n" + "mx.set_default_device(mx.cpu)\n" + "mx.eval(mx.compile(lambda x: mx.exp(x) * 2.0)(mx.ones((4,))))\n" + ) + with tempfile.TemporaryDirectory() as tmp_dir: + gxx = os.path.join(tmp_dir, "g++") + with open(gxx, "w") as f: + f.write( + "#!/bin/sh\n" + 'for a; do [ "$p" = -o ] && : > "$a" && exit 1; p=$a; done\n' + ) + os.chmod(gxx, 0o755) + env = dict(os.environ, TMPDIR=tmp_dir) + env.pop("MLX_DISABLE_COMPILE", None) + cmd = [sys.executable, "-c", code] + path = tmp_dir + os.pathsep + env["PATH"] + bad = subprocess.run(cmd, env=dict(env, PATH=path), capture_output=True) + self.assertNotEqual(bad.returncode, 0) + good = subprocess.run(cmd, env=env, capture_output=True) + self.assertEqual(good.returncode, 0, good.stderr.decode()) + if __name__ == "__main__": mlx_tests.MLXTestRunner() From 84c43de31a5428fc311d158ebb9795e1f718c1b0 Mon Sep 17 00:00:00 2001 From: Cheng Date: Wed, 7 Oct 2026 17:18:25 +0900 Subject: [PATCH 2/2] test could be flaky --- python/tests/test_compile.py | 31 ------------------------------- 1 file changed, 31 deletions(-) diff --git a/python/tests/test_compile.py b/python/tests/test_compile.py index d6b2b251ff..9f1ced1867 100644 --- a/python/tests/test_compile.py +++ b/python/tests/test_compile.py @@ -4,12 +4,7 @@ import inspect import io import math -import os -import subprocess -import sys -import tempfile import threading -import unittest from functools import partial, wraps from io import StringIO from itertools import product @@ -1715,32 +1710,6 @@ def fun(x): x = mx.array([1.0, float("nan"), 3.0]) self.assertTrue(mx.array_equal(mx.compile(fun)(x), mx.array([False, True]))) - @unittest.skipIf(sys.platform == "win32", "Needs a POSIX shell") - def test_compile_cpu_partial_library(self): - # A build that stops after it creates the library file must not break - # the build in later processes. - code = ( - "import mlx.core as mx\n" - "mx.set_default_device(mx.cpu)\n" - "mx.eval(mx.compile(lambda x: mx.exp(x) * 2.0)(mx.ones((4,))))\n" - ) - with tempfile.TemporaryDirectory() as tmp_dir: - gxx = os.path.join(tmp_dir, "g++") - with open(gxx, "w") as f: - f.write( - "#!/bin/sh\n" - 'for a; do [ "$p" = -o ] && : > "$a" && exit 1; p=$a; done\n' - ) - os.chmod(gxx, 0o755) - env = dict(os.environ, TMPDIR=tmp_dir) - env.pop("MLX_DISABLE_COMPILE", None) - cmd = [sys.executable, "-c", code] - path = tmp_dir + os.pathsep + env["PATH"] - bad = subprocess.run(cmd, env=dict(env, PATH=path), capture_output=True) - self.assertNotEqual(bad.returncode, 0) - good = subprocess.run(cmd, env=env, capture_output=True) - self.assertEqual(good.returncode, 0, good.stderr.decode()) - if __name__ == "__main__": mlx_tests.MLXTestRunner()