Fix a concurrency problem in jax_explain_cache_misses.

hawkinsp · hawkinsp · commit 0e352deb4839 · 2025-09-16T21:02:25.000Z
Fixes #30163
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
@@ -1293,8 +1293,16 @@ def explain_tracing_cache_miss(
 
   p(f"  for {func_name}{src_info}")
 
+  # Do *not* remove the list() around the call to keys(). The cache may be
+  # updated concurrently by other threads, and we need to perform the iteration
+  # over the dictionary keys in a way that is concurrency safe. Here we are
+  # relying on an implementation behavior of CPython wherein the particular list
+  # constructor used here acts atomically.
+  # See https://github.com/jax-ml/jax/issues/30163
+  cache_keys = list(cache.keys())
+
   diffs = [diff_tracing_cache_keys(key, ok, debug_info)
-           for ok in cache.keys() if key != ok]
+           for ok in cache_keys if key != ok]
   assert diffs, "we must find some diffs if key differs from all cache keys"
   min_diff = min(diffs, key=lambda v: v[1])
   smallest_diffs: Sequence[Sequence[str]]  # the diffs for the closest keys
diff --git a/tests/api_test.py b/tests/api_test.py
@@ -33,6 +33,7 @@
 import re
 import subprocess
 import sys
+import threading
 import traceback
 import types
 from typing import NamedTuple
@@ -4925,6 +4926,32 @@ def test_cache_miss_explanations_no_source_info(self):
     with config.explain_cache_misses(True):
       jax.jit(operator.add)(42, 24)
 
+  def test_cache_miss_explanations_are_thread_safe(self):
+    @jax.jit
+    def f(i):
+      return jnp.sum(i)
+
+    saw_exception = False
+
+    def thread(i0):
+      nonlocal saw_exception
+      try:
+        for i in range(i0, 100, 10):
+          if saw_exception:
+            break
+          with config.explain_cache_misses(True):
+            f(jnp.zeros(i))
+      except Exception:
+        saw_exception = True
+        raise
+
+    t = [threading.Thread(target=thread, args=(i,)) for i in range(10)]
+    for i in t:
+      i.start()
+    for i in t:
+      i.join()
+    self.assertFalse(saw_exception)
+
   @parameterized.named_parameters([
       {"testcase_name": f"{np.dtype(dtype)}", "dtype": dtype}
       for dtype in jtu.dtypes.custom_floats])
