Make api_test.py work when test cases are run using multiple threads.

* keep track of all known config.State objects so we can find them by name.
* change `@jtu.with_config` to default to setting thread-local configurations.
* add a `@jtu.with_global_config` for those things that truly need to be set globally.
* add a `@jtu.thread_local_config_context` that overrides thread-local configuration options, just as `jtu.global_config_context` overrides global configuration options.
* change the pretty printer color option to be a State so it can be set locally.
* tag a number of tests as thread-hostile, in particular tests that check counters for numbers of compilations, rely on garbage collection having particular semantics, or look at log output.

PiperOrigin-RevId: 713411171

Author: hawkinsp

jax/_src/config.py

@@ -233,6 +233,8 @@ def trace_context():
 class NoDefault: pass
 no_default = NoDefault()
 
+config_states = {}
+
 class State(config_ext.Config[_T]):
 
   __slots__ = (
@@ -265,6 +267,7 @@ def __init__(
       self._validator(default)
     if self._update_global_hook:
       self._update_global_hook(default)
+    config_states[name] = self
 
   def __bool__(self) -> NoReturn:
     raise TypeError(

jax/_src/pretty_printer.py

@@ -42,9 +42,9 @@
   colorama = None
 
 
-_PPRINT_USE_COLOR = config.bool_flag(
+_PPRINT_USE_COLOR = config.bool_state(
     'jax_pprint_use_color',
-    config.bool_env('JAX_PPRINT_USE_COLOR', True),
+    True,
     help='Enable jaxpr pretty-printing with colorful syntax highlighting.'
 )
 

jax/_src/test_util.py

@@ -1023,6 +1023,10 @@ def _run_one_test(test: unittest.TestCase, result: ThreadSafeTestResult):
   @contextmanager
   def thread_hostile_test():
     "Decorator for tests that are not thread-safe."
+    if TEST_NUM_THREADS.value <= 0:
+      yield
+      return
+
     _test_rwlock.assert_reader_held()
     _test_rwlock.reader_unlock()
     _test_rwlock.writer_lock()
@@ -1157,10 +1161,21 @@ def with_config(**kwds):
   """Test case decorator for subclasses of JaxTestCase"""
   def decorator(cls):
     assert inspect.isclass(cls) and issubclass(cls, JaxTestCase), "@with_config can only wrap JaxTestCase class definitions."
-    cls._default_config = {}
+    cls._default_thread_local_config = {}
     for b in cls.__bases__:
-      cls._default_config.update(b._default_config)
-    cls._default_config.update(kwds)
+      cls._default_thread_local_config.update(b._default_thread_local_config)
+    cls._default_thread_local_config.update(kwds)
+    return cls
+  return decorator
+
+def with_global_config(**kwds):
+  """Test case decorator for subclasses of JaxTestCase"""
+  def decorator(cls):
+    assert inspect.isclass(cls) and issubclass(cls, JaxTestCase), "@with_config can only wrap JaxTestCase class definitions."
+    cls._default_global_config = {}
+    for b in cls.__bases__:
+      cls._default_global_config.update(b._default_global_config)
+    cls._default_global_config.update(kwds)
     return cls
   return decorator
 
@@ -1191,6 +1206,15 @@ def global_config_context(**kwds):
     for key, value in original_config.items():
       config.update(key, value)
 
+@contextmanager
+def thread_local_config_context(**kwds):
+  stack = ExitStack()
+  for config_name, value in kwds.items():
+    stack.enter_context(config.config_states[config_name](value))
+  try:
+    yield
+  finally:
+    stack.close()
 
 class NotPresent:
   def __repr__(self):
@@ -1214,7 +1238,8 @@ def assert_global_configs_unchanged():
 
 class JaxTestCase(parameterized.TestCase):
   """Base class for JAX tests including numerical checks and boilerplate."""
-  _default_config = {
+  _default_global_config: dict[str, Any] = {}
+  _default_thread_local_config = {
     'jax_enable_checks': True,
     'jax_numpy_dtype_promotion': 'strict',
     'jax_numpy_rank_promotion': 'raise',
@@ -1239,7 +1264,9 @@ def setUp(self):
     self._context_stack = ExitStack()
     self.addCleanup(self._context_stack.close)
     stack = self._context_stack
-    stack.enter_context(global_config_context(**self._default_config))
+    stack.enter_context(global_config_context(**self._default_global_config))
+    for config_name, value in self._default_thread_local_config.items():
+      stack.enter_context(jax._src.config.config_states[config_name](value))
 
     if TEST_WITH_PERSISTENT_COMPILATION_CACHE.value:
       assert TEST_NUM_THREADS.value <= 1, "Persistent compilation cache is not thread-safe."
@@ -1365,7 +1392,7 @@ def wrapped_fun(*args):
 
     cache_misses = dispatch.xla_primitive_callable.cache_info().misses
     python_ans = fun(*args)
-    if check_cache_misses:
+    if check_cache_misses and TEST_NUM_THREADS.value <= 1:
       self.assertEqual(
           cache_misses, dispatch.xla_primitive_callable.cache_info().misses,
           "Compilation detected during second call of {} in op-by-op "

jax/experimental/jax2tf/tests/shape_poly_test.py

@@ -2670,7 +2670,7 @@ def test_harness(self, harness: PolyHarness):
     if harness.group_name == "eig" and not jtu.test_device_matches(["cpu"]):
       raise unittest.SkipTest("JAX implements eig only on CPU.")
 
-    with jtu.global_config_context(**harness.override_jax_config_flags):
+    with jtu.thread_local_config_context(**harness.override_jax_config_flags):
       harness.run_test(self)
 
 

tests/api_test.py

@@ -632,6 +632,7 @@ def f(x):
     python_should_be_executing = False
     jit(f)(3)
 
+  @jtu.thread_hostile_test()
   def test_jit_cache_clear(self):
     @jit
     def f(x, y):
@@ -2591,6 +2592,7 @@ def test_block_until_ready_mixed(self):
     self.assertAllClose(pytree[2], np.ones(3), check_dtypes=False)
     self.assertEqual(pytree[3], 4)
 
+  @jtu.thread_hostile_test()
   def test_devicearray_weakref_friendly(self):
     x = device_put(1.)
     y = weakref.ref(x)
@@ -2739,6 +2741,7 @@ def f(x):
 
     self.assertEqual(count(), 1)
 
+  @jtu.thread_hostile_test()
   def test_jit_infer_params_cache(self):
     def f(x):
       return x
@@ -3329,6 +3332,7 @@ def test_grad_object_array_error(self):
     with self.assertRaisesRegex(TypeError, ".*is not a valid JAX type"):
       jax.grad(lambda x: x)(x)
 
+  @jtu.thread_hostile_test()
   def test_jit_compilation_time_logging(self):
     @api.jit
     def f(x):
@@ -3417,6 +3421,7 @@ def test_trivial_computations(self):
     self.assertNotEqual(z3.unsafe_buffer_pointer(), x1.unsafe_buffer_pointer())
     self.assertEqual(z2, 1)
 
+  @jtu.thread_hostile_test()
   def test_nested_jit_hoisting(self):
     @api.jit
     def f(x, y):
@@ -3454,6 +3459,7 @@ def mlir_jaxpr_subcomp_and_collect(c, jaxpr, *args, **kwargs):
     self.assertEqual(inner_jaxpr.eqns[-2].primitive.name, 'mul')
     self.assertEqual(inner_jaxpr.eqns[-1].primitive.name, 'add')
 
+  @jtu.thread_hostile_test()
   def test_primitive_compilation_cache(self):
     with jtu.count_primitive_compiles() as count:
       lax.add(1, 2)
@@ -4013,13 +4019,17 @@ def __jax_array__(self):
     a2 = jnp.array(((x, x), [x, x]))
     self.assertAllClose(np.array(((1, 1), (1, 1))), a2)
 
+  @jtu.thread_hostile_test()
   def test_eval_shape_weak_type(self):
     # https://github.com/jax-ml/jax/issues/23302
     arr = jax.numpy.array(1)
 
+    def f(x):
+      return jax.numpy.array(x)
+
     with jtu.count_jit_tracing_cache_miss() as count:
-      jax.eval_shape(jax.numpy.array, 1)
-      out = jax.eval_shape(jax.numpy.array, 1)
+      jax.eval_shape(f, 1)
+      out = jax.eval_shape(f, 1)
 
     self.assertEqual(count(), 1)
     self.assertTrue(out.weak_type)
@@ -4138,6 +4148,7 @@ def test_dot_precision_flag(self):
       jaxpr = jax.make_jaxpr(jnp.dot)(x, x)
     self.assertIn('Precision.HIGH', str(jaxpr))
 
+  @jtu.thread_hostile_test()
   def test_dot_precision_forces_retrace(self):
     num_traces = 0
 
@@ -4310,6 +4321,7 @@ def test_jnp_array_doesnt_device_put(self):
       api.make_jaxpr(lambda: jnp.array(3))()
     self.assertEqual(count(), 0)
 
+  @jtu.thread_hostile_test()
   def test_rank_promotion_forces_retrace(self):
     num_traces = 0
 
@@ -4328,7 +4340,7 @@ def f_jit(x):
 
     for f in [f_jit, f_cond]:
       # Use _read() to read the flag value rather than threadlocal value.
-      allow_promotion = config._read("jax_numpy_rank_promotion")
+      allow_promotion = jax.numpy_rank_promotion.get_global()
       try:
         config.update("jax_numpy_rank_promotion", "allow")
         num_traces = 0
@@ -4350,9 +4362,9 @@ def f(x):
           self.assertGreaterEqual(num_traces, 2)
         nt = num_traces
         f(x)
-        self.assertEqual(num_traces, nt + 1)
+        self.assertEqual(num_traces, nt)
         f(x)
-        self.assertEqual(num_traces, nt + 1)
+        self.assertEqual(num_traces, nt)
       finally:
         config.update("jax_numpy_rank_promotion", allow_promotion)
 
@@ -4450,6 +4462,7 @@ def foo(x, y, z):
     self.assertEqual(jfoo.__qualname__, f"make_jaxpr({foo.__qualname__})")
     self.assertEqual(jfoo.__module__, "jax")
 
+  @jtu.thread_hostile_test()
   def test_inner_jit_function_retracing(self):
     # https://github.com/jax-ml/jax/issues/7155
     inner_count = outer_count = 0
@@ -4691,6 +4704,7 @@ def test_mesh_creation_error_message(self):
     with self.assertRaisesRegex(ValueError, "ndim of its first argument"):
       jax.sharding.Mesh(jax.devices(), ("x", "y"))
 
+  @jtu.thread_hostile_test()
   def test_jit_boundmethod_reference_cycle(self):
     class A:
       def __init__(self):
@@ -4829,6 +4843,7 @@ class RematTest(jtu.JaxTestCase):
           ('_policy', partial(jax.remat, policy=lambda *_, **__: False)),
           ('_new', partial(new_checkpoint, policy=lambda *_, **__: False)),
       ])
+  @jtu.thread_hostile_test()
   def test_remat_basic(self, remat):
     @remat
     def g(x):
@@ -5166,6 +5181,7 @@ def f_yesremat(x):
           ('_policy', partial(jax.remat, policy=lambda *_, **__: False)),
           ('_new', partial(new_checkpoint, policy=lambda *_, **__: False)),
       ])
+  @jtu.thread_hostile_test()
   def test_remat_no_redundant_flops(self, remat):
     # see https://github.com/jax-ml/jax/pull/1749#issuecomment-558267584
 
@@ -6409,6 +6425,7 @@ def f(x):
     self.assertIn(' sin ', str(jaxpr))
     self.assertIn(' cos ', str(jaxpr))
 
+  @jtu.thread_hostile_test()
   def test_remat_residual_logging(self):
     def f(x):
       x = jnp.sin(x)
@@ -9626,11 +9643,8 @@ def foo_bwd(_, g):
 
     foo.defvjp(foo_fwd, foo_bwd)
 
-    try:
-      jax.config.update('jax_custom_vjp_disable_shape_check', True)
+    with config.custom_vjp_disable_shape_check(True):
       jax.grad(lambda x, y: foo(x, y).sum(), 1)(jnp.ones(3), jnp.ones(4))
-    finally:
-      jax.config.update('jax_custom_vjp_disable_shape_check', False)
 
   def test_bwd_rule_can_produce_list_or_tuple(self):
     @jax.custom_vjp
@@ -11114,6 +11128,8 @@ def test_autodidax_smoketest(self):
     spec.loader.exec_module(autodidax_module)
 
 class GarbageCollectionTest(jtu.JaxTestCase):
+
+  @jtu.thread_hostile_test()
   def test_xla_gc_callback(self):
     # https://github.com/jax-ml/jax/issues/14882
     x_np = np.arange(10, dtype='int32')

tests/mock_gpu_test.py

@@ -28,7 +28,7 @@
 NUM_SHARDS = 4
 
 
-@jtu.with_config(mock_num_gpu_processes=NUM_SHARDS)
+@jtu.with_global_config(mock_num_gpu_processes=NUM_SHARDS)
 class MockGPUTest(jtu.JaxTestCase):
 
   def setUp(self):

tests/mock_gpu_topology_test.py

@@ -25,7 +25,7 @@
 NUM_HOSTS_PER_SLICE = 4
 
 
-@jtu.with_config(
+@jtu.with_global_config(
   jax_mock_gpu_topology=f"{NUM_SLICES}x{NUM_HOSTS_PER_SLICE}x1",
   jax_cuda_visible_devices="0")
 class MockGPUTopologyTest(jtu.JaxTestCase):

tests/pmap_test.py

@@ -2185,21 +2185,22 @@ def test_cache_hits_across_threads(self):
     f = lambda x: x+1
     inputs = np.zeros([jax.device_count()], dtype=np.float32)
     pmaped_f = self.pmap(f)
-    pmaped_f(inputs)
-    self.assertEqual(pmaped_f._cache_size, 1)
+    self.assertEqual(pmaped_f._cache_size, 0)
 
-    # Note: We do not call jax.pmap in the other thread but we reuse the same
-    # object.
+    # We only call pmaped_f in the thread pool to make sure that any
+    # thread-local config settings are identical.
     futures = []
-    with ThreadPoolExecutor(max_workers=1) as executor:
-      futures.append(executor.submit(lambda: pmaped_f(inputs)))
+    with ThreadPoolExecutor(max_workers=2) as executor:
+      for _ in range(8):
+        futures.append(executor.submit(lambda: pmaped_f(inputs)))
       outputs = [f.result() for f in futures]
 
-    np.testing.assert_array_equal(pmaped_f(inputs), outputs[0])
     if pmaped_f._cache_size != 1:
       print(pmaped_f._debug_cache_keys())
     self.assertEqual(pmaped_f._cache_size, 1)
 
+    np.testing.assert_array_equal(pmaped_f(inputs), outputs[0])
+
   def test_cache_uses_jax_key(self):
     f = lambda x: x+1
     inputs = np.zeros([jax.device_count()], dtype=np.float32)

tests/shape_poly_test.py

@@ -3683,7 +3683,7 @@ def test_harness(self, harness: PolyHarness):
     if "cholesky" in harness.group_name and jtu.test_device_matches(["tpu"]):
       harness.tol = 5e-5
 
-    with jtu.global_config_context(**config_flags):
+    with jtu.thread_local_config_context(**config_flags):
       harness.run_test(self)