From 8a318d0de7f08e08e8ffe5133006489ec596ee94 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 20 Oct 2026 18:02:10 +0000 Subject: [PATCH 1/4] Run the unit tests in parallel with pytest-xdist The unit workflow's `linux` job gates every GPU and example workflow: their pr-gate jobs wait for it (poseidon/wait-for-status-checks, `^linux$`). That wait is a median 13.8 min, 23% of a PR's GPU run, and the job itself takes a median 14.8 min. The tests are CPU-bound and run serially, so use the runner's four vCPUs. - unit_tests.yml: `PYTEST_ADDOPTS=-n 4 --dist loadfile` on the `linux` job and the multi-version matrix. One worker per vCPU (`-n auto` counts physical cores), whole files per worker, and `OMP_NUM_THREADS=1`: with torch's default threads per worker the runner is oversubscribed and the run is slower than serial (1,481 s and 41 failures in my local runs). - noxfile.py: the second pytest process that runs the `manual` tests clears PYTEST_ADDOPTS, so it stays serial: those tests leak global state into each other, and a worker would run them in some order. - tests/_test_utils/torch/quantization/onnx_export.py: TEST_MODELS was a set of classes, whose iteration order differs between processes, so xdist workers collected different test ids. Use a tuple. pytest-xdist is already in the dev-test extra on main (#2771); uv.lock is refreshed by the weekly bump_uv_lock workflow. Testing: the full tests/unit suite, on 4 CPUs with coverage as in CI, on the base before the unit-test cleanup of #2823. Serial 825 s; this configuration 332-344 s in three repeated runs (4,724 passed, 0 failed). The other matrix entries (Python 3.10, 3.13, 3.14, torch 2.8, transformers min) take 4.5-5.3 min. Signed-off-by: Shengliang Xu --- .github/workflows/unit_tests.yml | 7 +++++++ noxfile.py | 3 ++- tests/_test_utils/torch/quantization/onnx_export.py | 2 +- 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml index 6fbca185d9f..c3bae2e41be 100644 --- a/.github/workflows/unit_tests.yml +++ b/.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' diff --git a/noxfile.py b/noxfile.py index 321a8924210..32ff854a0f1 100644 --- a/noxfile.py +++ b/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": ""}, ) diff --git a/tests/_test_utils/torch/quantization/onnx_export.py b/tests/_test_utils/torch/quantization/onnx_export.py index e65e6f23500..b3041f780d1 100644 --- a/tests/_test_utils/torch/quantization/onnx_export.py +++ b/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): From 95b7eba6ff47c8e95cc5bbc14852881955b0b6ff Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 17 Oct 2026 00:52:17 +0000 Subject: [PATCH 2/4] Restore ConstraintsFunc._get_true_latency after test_search_constraints test_search_constraints replaced ConstraintsFunc._get_true_latency with a fake by plain assignment and never put the original back, so the fake stayed in place for the rest of the process. test_evaluate_constraints.py::test_evaluate_latency_constraints and ::test_percent_limits use the real method through a fake runtime client (latency 100 and 200) and fail when they run after it: the leaked fake returns 280.5 for the tiny resnet, so the 20% latency limit comes out as 56.1 instead of 40, and a 200 limit is not met. Serial runs never saw it, because test_evaluate_constraints.py sorts before test_full_algorithms_vision.py. With pytest-xdist (-n 4 --dist loadfile), added by the previous commit, a worker can get the files in the other order: in the Python 3.10 and 3.14 legs of a CI run, the same worker ran test_search_constraints and then the two victims, which failed; in the legs where they landed on different workers, or in the other order, they passed. Use monkeypatch, so the patch is undone after the test. Testing: running test_search_constraints[False-True] and then test_evaluate_constraints.py in one process failed both tests with the CI values before the change (56.10 against 40.0, and assert False is True) and passes after it, and ConstraintsFunc._get_true_latency is the original function again afterwards. tests/unit (without puzzletron, which needs hydra) in reversed and in shuffled file order shows no other leak of this kind. Signed-off-by: Shengliang Xu --- tests/unit/torch/nas/test_full_algorithms_vision.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/unit/torch/nas/test_full_algorithms_vision.py b/tests/unit/torch/nas/test_full_algorithms_vision.py index f6c30c1da52..b3d81881696 100644 --- a/tests/unit/torch/nas/test_full_algorithms_vision.py +++ b/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}) From 7226f9179295703b8bf0d0837e5118caa67fd1b4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 17 Oct 2026 00:52:30 +0000 Subject: [PATCH 3/4] Give the unit tests 120 s per test on every platform tests/unit/torch/quantization/test_dist.py::test_data_parallel hit the 60 s default cap in the Python 3.12 / torch 2.13 leg of a CI run with pytest-xdist. The test spawns two processes and takes 8-15 s in serial runs. With pytest-xdist it takes 25-35 s in every other leg, and 60 s or more on the slowest runner: four workers and the child processes of the multi-process tests (test_dist, the three test_*data_parallel* autoquant tests, the tests that share cpu_dist_workers) share the four vCPUs of a CI runner. A hung test is still reported, after 2 min instead of 1. Windows already had 120 s. Testing: the durations above come from the job logs of that CI run and of the serial run before it. Not re-run on a runner: the cap only adds headroom. Signed-off-by: Shengliang Xu --- tests/conftest.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index af77546a244..b1685e3fe63 100644 --- a/tests/conftest.py +++ b/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, } From 4cbc40ac3a1bd336c612eb6dbbf0bcf6d70ce239 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 20 Oct 2026 18:03:09 +0000 Subject: [PATCH 4/4] Run test_overfits_a_single_batch in float32 The tiny model is built in bf16, and a CPU without native bf16 support emulates its matmuls: one step costs 0.9 s, against 0.1 s in float32. Build it in float32, so the 30 steps main runs take 3.4 s instead of 27 s on one thread. The loop and both assertions are unchanged, and the test needs no timeout override either way (27 s is below the 60 s cap). Signed-off-by: Shengliang Xu --- tests/unit/torch/speculative/plugins/test_hf_dflash2.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py index 97def429a1e..79dcb2a894c 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py +++ b/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()