[OMNIML-6178] [1/11] Run the unit tests in parallel with pytest-xdist

.github/workflows/unit_tests.yml

@@ -69,6 +69,10 @@ jobs:
           COVERAGE_FILE: ${{ github.workspace }}/.coverage
           # Everything but PRs also runs the slow tests marked ``manual`` (e.g. AutoNAS).
           PYTEST_EXTRA: ${{ github.event_name != 'pull_request' && '--run-manual' || '' }}
+          # One single-threaded xdist worker per vCPU (`-n auto` counts physical cores); torch's
+          # default threads per worker oversubscribe the runner and run slower than serial.
+          PYTEST_ADDOPTS: -n 4 --dist loadfile
+          OMP_NUM_THREADS: "1"
         run: pip install nox uv && nox -s "unit-3.12(torch_214, tf_latest)" -- $PYTEST_EXTRA
       - name: Upload coverage reports to Codecov
         uses: codecov/codecov-action@v7
@@ -146,6 +150,9 @@ jobs:
         with:
           python-version: ${{ matrix.python_version }}
       - name: Run unit tests
+        env:
+          PYTEST_ADDOPTS: -n 4 --dist loadfile
+          OMP_NUM_THREADS: "1"
         run: pip install nox uv && nox -s "${{ matrix.nox_session }}"
   partial-install:
     if: needs.check-file-changes.outputs.any_changed == 'true'

noxfile.py

@@ -79,11 +79,12 @@ def unit(session, torch_ver, tf_ver):
     if run_manual:
         # A separate process: some manual tests (e.g. multi-threaded from_pretrained) leak global
         # state that breaks later tests. No --cov here, so this run writes no coverage data.
+        # Serial too (CI's PYTEST_ADDOPTS asks for xdist): a worker would run them in some order.
         session.run(
             "python",
             # Caller args first: the last ``-m`` wins, and it must be ours.
             *("-m", "pytest", "tests/unit", *pytest_args, "--run-manual", "-m", "manual"),
-            env={**_CPU_ONLY_ENV, "COVERAGE_PROCESS_START": ""},
+            env={**_CPU_ONLY_ENV, "COVERAGE_PROCESS_START": "", "PYTEST_ADDOPTS": ""},
         )
 
 

tests/_test_utils/torch/quantization/onnx_export.py

@@ -23,7 +23,7 @@
 
 from .models import SimpleConv, SimpleConvLinear, SimpleLinear
 
-TEST_MODELS = {SimpleLinear, SimpleConv, SimpleConvLinear}
+TEST_MODELS = (SimpleLinear, SimpleConv, SimpleConvLinear)
 
 
 def onnx_export_tester(model, device, num_bits, per_channel_quantization, constant_folding, dtype):

tests/conftest.py

@@ -58,7 +58,10 @@ def pytest_addoption(parser):
     "gpu_trtllm": 60,
     "gpu_vllm": 60,
     "regression": 180,
-    "unit": 120 if platform.system() == "Windows" else 60,
+    # Tests that spawn processes take 2-3x longer under pytest-xdist: 4 workers and their children
+    # share the 4 vCPUs of a CI runner (test_dist.py::test_data_parallel: 8-15 s serial, 25-35 s with
+    # xdist, over 60 s on a slow runner).
+    "unit": 120,
 }
 
 

tests/unit/torch/nas/test_full_algorithms_vision.py

@@ -117,7 +117,7 @@ def run_search():
         pytest.param(False, False, marks=pytest.mark.manual),
     ],
 )
-def test_search_constraints(flops_only: bool, bounded_latency: bool):
+def test_search_constraints(flops_only: bool, bounded_latency: bool, monkeypatch):
     def _fake_latency(self, model, precomputed=None):
         coeff_params = 0.00 if flops_only else 0.75
         if not hasattr(_fake_latency, "call_counter"):
@@ -130,8 +130,9 @@ def _fake_latency(self, model, precomputed=None):
     dummy_input = (*args, kwargs)
     model = mtn.convert(model, mode="autonas")
 
-    # get fake latency and spoof ConstraintsFunc class
-    ConstraintsFunc._get_true_latency = _fake_latency
+    # get fake latency and spoof ConstraintsFunc class (undone after the test: a class-level patch
+    # that stays behind changes the latency of every later test in the same process)
+    monkeypatch.setattr(ConstraintsFunc, "_get_true_latency", _fake_latency)
 
     # now setup constraints functor
     constraints_func = ConstraintsFunc(model, {}, dummy_input, {"device": None})

tests/unit/torch/speculative/plugins/test_hf_dflash2.py

@@ -355,7 +355,8 @@ def test_overfits_a_single_batch(self):
         Guards the target/predecessor alignment: a misaligned selector objective still
         produces a finite decreasing loss, but its accuracy does not reach 1.
         """
-        model = get_tiny_llama(num_hidden_layers=4)
+        # float32: CPU bf16 matmuls are emulated (~8x slower); this tiny model does not need bf16
+        model = get_tiny_llama(num_hidden_layers=4, dtype=torch.float32)
         mtsp.convert(model, [("dflash", _get_dflash2_config())])
         model.train()