Temporarily disable free-threading build & multithread tests

Why? We need an 0.5.0 release for Python 3.13 support, but the new multi-thread support still results in flaky tests due to some race conditions in both ml-dtypes and NumPy. After the 0.5.0 release we will plan to re-enable these and work on a longer-term fix.
jax-ml · Sep 13, 2024 · 0b08020 · 0b08020
1 parent b39b73c
commit 0b08020
Show file tree

Hide file tree

Showing 6 changed files with 36 additions and 28 deletions.
diff --git a/.github/workflows/wheels.yml b/.github/workflows/wheels.yml
@@ -45,10 +45,11 @@ jobs:
       - name: Build wheels
         run: python -m cibuildwheel --output-dir wheelhouse
         env:
+          # TODO(jakevdp): re-add 313t & free-threading support
           CIBW_ARCHS_LINUX: auto aarch64
           CIBW_ARCHS_MACOS: universal2
-          CIBW_BUILD: cp39-* cp310-* cp311-* cp312-* cp313*
-          CIBW_FREE_THREADED_SUPPORT: True
+          CIBW_BUILD: cp39-* cp310-* cp311-* cp312-* cp313-* # cp313t-*
+          # CIBW_FREE_THREADED_SUPPORT: True
           CIBW_PRERELEASE_PYTHONS: True
           CIBW_SKIP: "*musllinux* *i686* *win32* *t-win*"
           CIBW_TEST_REQUIRES: absl-py pytest pytest-xdist

diff --git a/ml_dtypes/tests/custom_float_test.py b/ml_dtypes/tests/custom_float_test.py
@@ -27,7 +27,7 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import ml_dtypes
-from multi_thread_utils import multi_threaded
+# from multi_thread_utils import multi_threaded
 import numpy as np
 
 bfloat16 = ml_dtypes.bfloat16
@@ -221,11 +221,12 @@ def dtype_is_signed(dtype):
 }
 
 
+# TODO(jakevdp): re-enable multi-threaded tests after 0.5.0 release.
 # pylint: disable=g-complex-comprehension
-@multi_threaded(
-    num_workers=3,
-    skip_tests=["testDiv", "testRoundTripNumpyTypes", "testRoundTripToNumpy"],
-)
+# @multi_threaded(
+#     num_workers=3,
+#     skip_tests=["testDiv", "testRoundTripNumpyTypes", "testRoundTripToNumpy"],
+# )
 @parameterized.named_parameters(
     (
         {"testcase_name": "_" + dtype.__name__, "float_type": dtype}
@@ -660,20 +661,21 @@ def testDtypeFromString(self, float_type):
 ]
 
 
+# TODO(jakevdp): re-enable multi-threaded tests after 0.5.0 release.
 # pylint: disable=g-complex-comprehension
-@multi_threaded(
-    num_workers=3,
-    skip_tests=[
-        "testBinaryUfunc",
-        "testConformNumpyComplex",
-        "testFloordivCornerCases",
-        "testDivmodCornerCases",
-        "testSpacing",
-        "testUnaryUfunc",
-        "testCasts",
-        "testLdexp",
-    ],
-)
+# @multi_threaded(
+#     num_workers=3,
+#     skip_tests=[
+#         "testBinaryUfunc",
+#         "testConformNumpyComplex",
+#         "testFloordivCornerCases",
+#         "testDivmodCornerCases",
+#         "testSpacing",
+#         "testUnaryUfunc",
+#         "testCasts",
+#         "testLdexp",
+#     ],
+# )
 @parameterized.named_parameters(
     (
         {"testcase_name": "_" + dtype.__name__, "float_type": dtype}

diff --git a/ml_dtypes/tests/finfo_test.py b/ml_dtypes/tests/finfo_test.py
@@ -15,7 +15,7 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import ml_dtypes
-from multi_thread_utils import multi_threaded
+# from multi_thread_utils import multi_threaded
 import numpy as np
 
 ALL_DTYPES = [
@@ -55,7 +55,8 @@
 }
 
 
-@multi_threaded(num_workers=3)
+# TODO(jakevdp): re-enable multi-threaded tests after 0.5.0 release.
+# @multi_threaded(num_workers=3)
 class FinfoTest(parameterized.TestCase):
 
   def assertNanEqual(self, x, y):

diff --git a/ml_dtypes/tests/iinfo_test.py b/ml_dtypes/tests/iinfo_test.py
@@ -15,11 +15,12 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import ml_dtypes
-from multi_thread_utils import multi_threaded
+# from multi_thread_utils import multi_threaded
 import numpy as np
 
 
-@multi_threaded(num_workers=3)
+# TODO(jakevdp): re-enable multi-threaded tests after 0.5.0 release.
+# @multi_threaded(num_workers=3)
 class IinfoTest(parameterized.TestCase):
 
   def testIinfoInt2(self):

diff --git a/ml_dtypes/tests/intn_test.py b/ml_dtypes/tests/intn_test.py
@@ -48,8 +48,9 @@ def ignore_warning(**kw):
     yield
 
 
+# TODO(jakevdp): re-enable multi-threaded tests after 0.5.0 release.
 # Tests for the Python scalar type
-@multi_threaded(num_workers=3)
+# @multi_threaded(num_workers=3)
 class ScalarTest(parameterized.TestCase):
 
   @parameterized.product(scalar_type=INTN_TYPES)
@@ -246,8 +247,9 @@ def testCanCast(self, a, b):
     )
 
 
+# TODO(jakevdp): re-enable multi-threaded tests after 0.5.0 release.
 # Tests for the Python scalar type
-@multi_threaded(num_workers=3, skip_tests=["testBinaryUfuncs"])
+# @multi_threaded(num_workers=3, skip_tests=["testBinaryUfuncs"])
 class ArrayTest(parameterized.TestCase):
 
   @parameterized.product(scalar_type=INTN_TYPES)

diff --git a/ml_dtypes/tests/metadata_test.py b/ml_dtypes/tests/metadata_test.py
@@ -16,10 +16,11 @@
 
 from absl.testing import absltest
 import ml_dtypes
-from multi_thread_utils import multi_threaded
+# from multi_thread_utils import multi_threaded
 
 
-@multi_threaded(num_workers=3)
+# TODO(jakevdp): re-enable multi-threaded tests after 0.5.0 release.
+# @multi_threaded(num_workers=3)
 class CustomFloatTest(absltest.TestCase):
 
   def test_version_matches_package_metadata(self):