Add a unittest test extension that runs test cases in parallel using threads.

This change does not yet do the work necessary to make any tests pass with threading enabled, which will come in future changes.

This approach is broadly inspired by https://github.com/testing-cabal/testtools/blob/a6d205dd4cac51f3cf9267978d39fc877103aacb/testtools/testsuite.py#L113 and by unittest-ft.

We add a custom TestResult class that batches up any test result actions and applies them under a lock. We also add a custom TestSuite class that runs individual test cases in parallel using a thread-pool.

We need a reader-writer lock to implement a `@jtu.thread_hostile_test` decorator, which we do by adding bindings around absl::Mutex to jaxlib.

PiperOrigin-RevId: 713312937

Author: hawkinsp

jax/_src/test_util.py

@@ -17,6 +17,7 @@
 
 import collections
 from collections.abc import Callable, Generator, Iterable, Sequence
+from concurrent.futures import ThreadPoolExecutor
 from contextlib import ExitStack, contextmanager
 import datetime
 import functools
@@ -31,6 +32,7 @@
 import tempfile
 import textwrap
 import threading
+import time
 from typing import Any, TextIO
 import unittest
 import warnings
@@ -115,6 +117,12 @@
           'deterministic, interactive'),
 )
 
+TEST_NUM_THREADS = config.int_flag(
+    'jax_test_num_threads', 0,
+    help='Number of threads to use for running tests. 0 means run everything '
+    'in the main thread. Using > 1 thread is experimental.'
+)
+
 # We sanitize test names to ensure they work with "unitttest -k" and
 # "pytest -k" test filtering. pytest accepts '[' and ']' but unittest -k
 # does not. We replace sequences of problematic characters with a single '_'.
@@ -998,8 +1006,140 @@ def sample_product(*args, **kw):
   """
   return parameterized.parameters(*sample_product_testcases(*args, **kw))
 
+# We use a reader-writer lock to protect test execution. Tests that may run in
+# parallel acquire a read lock; tests that are not thread-safe acquire a write
+# lock.
+if hasattr(util, 'Mutex'):
+  _test_rwlock = util.Mutex()
+
+  def _run_one_test(test: unittest.TestCase, result: ThreadSafeTestResult):
+    _test_rwlock.reader_lock()
+    try:
+      test(result)  # type: ignore
+    finally:
+      _test_rwlock.reader_unlock()
+
+
+  @contextmanager
+  def thread_hostile_test():
+    "Decorator for tests that are not thread-safe."
+    _test_rwlock.assert_reader_held()
+    _test_rwlock.reader_unlock()
+    _test_rwlock.writer_lock()
+    try:
+      yield
+    finally:
+      _test_rwlock.writer_unlock()
+      _test_rwlock.reader_lock()
+else:
+  # TODO(phawkins): remove this branch when jaxlib 0.5.0 is the minimum.
+  _test_rwlock = threading.Lock()
+
+  def _run_one_test(test: unittest.TestCase, result: ThreadSafeTestResult):
+    _test_rwlock.acquire()
+    try:
+      test(result)  # type: ignore
+    finally:
+      _test_rwlock.release()
+
+
+  @contextmanager
+  def thread_hostile_test():
+    yield  # No reader-writer lock, so we get no parallelism.
+
+class ThreadSafeTestResult:
+  """
+  Wraps a TestResult to make it thread safe.
+
+  We do this by accumulating API calls and applying them in a batch under a
+  lock at the conclusion of each test case.
+
+  We duck type instead of inheriting from TestResult because we aren't actually
+  a perfect implementation of TestResult, and would rather get a loud error
+  for things we haven't implemented.
+  """
+  def __init__(self, lock: threading.Lock, result: unittest.TestResult):
+    self.lock = lock
+    self.test_result = result
+    self.actions: list[Callable] = []
+
+  def startTest(self, test: unittest.TestCase):
+    del test
+    self.start_time = time.time()
+
+  def stopTest(self, test: unittest.TestCase):
+    stop_time = time.time()
+    with self.lock:
+      # We assume test_result is an ABSL _TextAndXMLTestResult, so we can
+      # override how it gets the time.
+      time_getter = self.test_result.time_getter
+      try:
+        self.test_result.time_getter = lambda: self.start_time
+        self.test_result.startTest(test)
+        for callback in self.actions:
+          callback()
+        self.test_result.time_getter = lambda: stop_time
+        self.test_result.stopTest(test)
+      finally:
+        self.test_result.time_getter = time_getter
+
+  def addSuccess(self, test: unittest.TestCase):
+    self.actions.append(lambda: self.test_result.addSuccess(test))
+
+  def addSkip(self, test: unittest.TestCase, reason: str):
+    self.actions.append(lambda: self.test_result.addSkip(test, reason))
+
+  def addError(self, test: unittest.TestCase, err):
+    self.actions.append(lambda: self.test_result.addError(test, err))
+
+  def addFailure(self, test: unittest.TestCase, err):
+    self.actions.append(lambda: self.test_result.addFailure(test, err))
+
+  def addExpectedFailure(self, test: unittest.TestCase, err):
+    self.actions.append(lambda: self.test_result.addExpectedFailure(test, err))
+
+  def addDuration(self, test: unittest.TestCase, elapsed):
+    self.actions.append(lambda: self.test_result.addDuration(test, elapsed))
+
+
+class JaxTestSuite(unittest.TestSuite):
+  """Runs tests in parallel using threads if TEST_NUM_THREADS is > 1.
+
+  Caution: this test suite does not run setUpClass or setUpModule methods if
+  thread parallelism is enabled.
+  """
+
+  def __init__(self, suite: unittest.TestSuite):
+    super().__init__(list(suite))
+
+  def run(self, result: unittest.TestResult, debug: bool = False) -> unittest.TestResult:
+    if TEST_NUM_THREADS.value <= 0:
+      return super().run(result)
+
+    executor = ThreadPoolExecutor(TEST_NUM_THREADS.value)
+    lock = threading.Lock()
+    futures = []
+
+    def run_test(test):
+      "Recursively runs tests in a test suite or test case."
+      if isinstance(test, unittest.TestSuite):
+        for subtest in test:
+          run_test(subtest)
+      else:
+        test_result = ThreadSafeTestResult(lock, result)
+        futures.append(executor.submit(_run_one_test, test, test_result))
+
+    with executor:
+      run_test(self)
+      for future in futures:
+        future.result()
+
+    return result
+
 
 class JaxTestLoader(absltest.TestLoader):
+  suiteClass = JaxTestSuite
+
   def getTestCaseNames(self, testCaseClass):
     names = super().getTestCaseNames(testCaseClass)
     if _TEST_TARGETS.value:
@@ -1102,11 +1242,11 @@ def setUp(self):
     stack.enter_context(global_config_context(**self._default_config))
 
     if TEST_WITH_PERSISTENT_COMPILATION_CACHE.value:
+      assert TEST_NUM_THREADS.value <= 1, "Persistent compilation cache is not thread-safe."
       stack.enter_context(config.enable_compilation_cache(True))
       stack.enter_context(config.raise_persistent_cache_errors(True))
       stack.enter_context(config.persistent_cache_min_compile_time_secs(0))
       stack.enter_context(config.persistent_cache_min_entry_size_bytes(0))
-
       tmp_dir = stack.enter_context(tempfile.TemporaryDirectory())
       stack.enter_context(config.compilation_cache_dir(tmp_dir))
       stack.callback(compilation_cache.reset_cache)

jax/_src/util.py

@@ -685,3 +685,6 @@ def test_event(name: str, *args) -> None:
   if not test_event_listener:
     return
   test_event_listener(name, *args)
+
+if hasattr(jaxlib_utils, "Mutex"):
+  Mutex = jaxlib_utils.Mutex

jaxlib/BUILD

@@ -214,6 +214,7 @@ nanobind_extension(
     deps = [
         "@com_google_absl//absl/cleanup",
         "@com_google_absl//absl/container:inlined_vector",
+        "@com_google_absl//absl/synchronization",
         "@nanobind",
         "@xla//third_party/python_runtime:headers",
     ],

jaxlib/utils.cc

@@ -18,6 +18,7 @@ limitations under the License.
 #include "nanobind/nanobind.h"
 #include "absl/cleanup/cleanup.h"
 #include "absl/container/inlined_vector.h"
+#include "absl/synchronization/mutex.h"
 
 namespace nb = nanobind;
 
@@ -225,4 +226,19 @@ NB_MODULE(utils, m) {
       PyCFunction_NewEx(&safe_map_def, /*self=*/nullptr, module_name.ptr()));
   m.attr("safe_zip") = nb::steal<nb::object>(
       PyCFunction_NewEx(&safe_zip_def, /*self=*/nullptr, module_name.ptr()));
+
+  // Python has no reader-writer lock in its standard library, so we expose
+  // bindings around absl::Mutex.
+  nb::class_<absl::Mutex>(m, "Mutex")
+      .def(nb::init<>())
+      .def("lock", &absl::Mutex::Lock, nb::call_guard<nb::gil_scoped_release>())
+      .def("unlock", &absl::Mutex::Unlock)
+      .def("assert_held", &absl::Mutex::AssertHeld)
+      .def("reader_lock", &absl::Mutex::ReaderLock,
+           nb::call_guard<nb::gil_scoped_release>())
+      .def("reader_unlock", &absl::Mutex::ReaderUnlock)
+      .def("assert_reader_held", &absl::Mutex::AssertReaderHeld)
+      .def("writer_lock", &absl::Mutex::WriterLock,
+           nb::call_guard<nb::gil_scoped_release>())
+      .def("writer_unlock", &absl::Mutex::WriterUnlock);
 }