[OMNIML-6178] Keep the cpu_dist_workers pool in one worker of the xdist unit run

.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:

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))