From 3d3cbea4249d9ce28365866a9bb27a3d6798d2ab Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 23 Oct 2026 02:16:35 +0000 Subject: [PATCH] Keep the cpu_dist_workers pool in one worker of the xdist unit run The multi-version torch_28 unit job is stopped by its runner ("The runner has received a shutdown signal") on 94 of 489 attempts between 10-20 and 10-23, none with a failed test and no other job affected. The ONNX conv/IsInf tests were in flight in 93 of the 94 stops, and the stop rate rises with the number of xdist workers that hold a cpu_dist_workers pool when those tests start: 5 of 27 attempts with 2 pools, 50 of 270 with 3, 39 of 114 with 4. The pool (two gloo processes of about 0.7 GB) is session-scoped, so under `--dist loadfile` every worker that ran one of the files using it keeps its own until the end. `--dist loadgroup` schedules by the xdist_group marker and treats an unmarked test as a group of its own. A tryfirst hook in tests/unit/conftest.py marks every file as a group (what loadfile does), except that all files with a cpu_dist_workers test share one group, so one worker holds the pool. unit_tests.yml passes `--dist loadgroup` in both places. Testing: full tests/unit run, CI-like (Python 3.12, torch 2.14.1 CPU, xdist 3.8.0, -n 4, coverage on, OMP_NUM_THREADS=1), peak of the summed RSS of the process tree sampled every 0.1 s, two runs each, same tests passing in both modes. main 3a34d65059 (5,353 passed): loadfile 14,024 / 14,126 MiB, loadgroup 10,832 / 10,294. hfep stack top a3370d2c0a4 with this change (5,830 passed): loadfile 15,633 / 15,823, loadgroup 11,089 / 10,914. The pool files ran on one worker (243 tests on main, 743 on the hfep tree) and no file was split. RSS counts shared pages, so the absolute numbers overstate the real use. Not run on a torch 2.8 environment or a hosted runner. Signed-off-by: Shengliang Xu --- .github/workflows/unit_tests.yml | 4 ++-- tests/unit/conftest.py | 19 +++++++++++++++++++ 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml index 5c4cd998c1b..68600a0ce0b 100644 --- a/.github/workflows/unit_tests.yml +++ b/.github/workflows/unit_tests.yml @@ -73,7 +73,7 @@ jobs: COVERAGE_CORE: sysmon # 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 + PYTEST_ADDOPTS: -n 4 --dist loadgroup 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 @@ -153,7 +153,7 @@ jobs: python-version: ${{ matrix.python_version }} - name: Run unit tests env: - PYTEST_ADDOPTS: -n 4 --dist loadfile + PYTEST_ADDOPTS: -n 4 --dist loadgroup OMP_NUM_THREADS: "1" run: pip install nox uv && nox -s "${{ matrix.nox_session }}" partial-install: diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index c8ee193c624..0ef5e19afd6 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -55,3 +55,22 @@ def cpu_dist_workers(_cpu_pools): return get_live_pool( _cpu_pools, 2, lambda: DistributedWorkerPool(world_size=2, backend="gloo", teardown_fn=None) ) + + +@pytest.hookimpl(tryfirst=True) # xdist reads the marker in a tryfirst hook of its own +def pytest_collection_modifyitems(config, items): + """Group the tests for ``--dist loadgroup`` so that one xdist worker holds the ``cpu_dist_workers`` pool. + + ``loadgroup`` treats an unmarked test as a group of its own, so every file becomes a group (what + ``loadfile`` does), except that all files with a ``cpu_dist_workers`` test share one group. The + pool (two gloo processes of about 0.7 GB) outlives its file, so with ``loadfile`` up to four + workers held one at once, which on a 16 GB runner leaves the memory-hungry tests too little. + Summed RSS of a CI-like unit run: 14-16 GB with ``loadfile``, about 11 GB with this group. + """ + if not config.pluginmanager.hasplugin("xdist"): + return + pool_files = {item.path for item in items if "cpu_dist_workers" in item.fixturenames} + for item in items: + if item.get_closest_marker("xdist_group") is None: + group = "cpu_pool" if item.path in pool_files else item.nodeid.split("::")[0] + item.add_marker(pytest.mark.xdist_group(group))