[shengliangxu]

What does this PR do?

Type of change: CI improvement (memory of the unit lane); not a bug fix, and no speed gain was measurable on hosted runners

Puts all unit test files that use the cpu_dist_workers pool into one xdist group (--dist loadgroup), so that one worker instead of up to four holds the pool. The fixture says its workers "start once" and are shared by every multi-rank test; under xdist (#2845) each worker started its own pool.

Change. --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 its own group (which is 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.

What it improves.

• Memory, measured on a CI-like run of the whole tests/unit (Python 3.12, torch 2.14.1 CPU, -n 4, coverage on): the peak of the summed RSS falls from 14,024-14,126 MiB to 10,294-10,832 MiB on main and from 15,633-15,823 MiB to 10,914-11,089 MiB on the hfep stack's top (two runs each, the same tests pass). As physical memory (PSS, shared pages divided among the processes that map them; whole tests/unit at c00ac43, -n 4, no coverage, torch 2.10 CUDA build): 12,984 MiB with 20 processes and 29 s above 10 GB become 9,836 MiB with 12 processes, never above 10 GB.
• Time in the lab: pytest time 217 / 224 s (loadfile) against 201 / 197 s (this PR) on main on a loaded 96-CPU host, 158 s against 144 s in another lab run; fewer pool startups, each importing torch and transformers in two processes.

What it does not show: a speed gain on hosted runners. linux job, pytest's own time (the summary line), all on runner image ubuntu-24.04 20261012.334.1: this PR 245 s median over 6 runs (190-291 s), other trees since #2950 242 s over 12 runs, the hfep/epdp stack trees 296 s over 49 runs (they have 9% more tests and three times as many pool tests). The Run unit tests step is shorter (312 s against 382 s for other trees), but it includes the environment install, which this change does not touch. Over the 12 unit jobs the per-job step times are 0.89 of the other trees' median (9 faster), and one run varies by 20% on its own, so a gain below about 10% is not visible with 6 runs.

Context: the runner stops, which this does not fix. The multi-version torch_28 unit job was stopped mid-run by its runner ("The runner has received a shutdown signal") on 94 of 489 attempts between 10-20 and 10-23; no other unit job was. The ONNX conv/IsInf tests were in flight in 93 of the 94. #2947 shrank their input and merged on 10-23 at 00:47 UTC; since #2950 let CI run again (05:23 UTC), all 32 finished torch_28 attempts passed (23 on the hfep/epdp stack, where 32% were stopped before on the same runner image; 4 on other branches; 5 of this PR's own). The pools raised the risk without being the trigger:

pools alive when the first conv/IsInf test starts	stopped / attempts
2	5 / 27 (19%)
3	50 / 270 (19%)
4	39 / 114 (34%)

So this PR is optional memory headroom (the margin of a 16 GB runner, about 14.5 GB usable, is thin through the whole run: PSS keeps climbing because the pools accumulate), and is a draft. Not run: a torch 2.8 environment.

Testing

Full tests/unit, CI-like, peak of the summed RSS of the whole process tree sampled every 0.1 s, in MiB, two runs per cell; every run passes the same tests with and without the change.

tree	--dist loadfile	--dist loadgroup (this PR)
main 3a34d65 (5,353 passed, 31 skipped, 12 xfailed)	14,024 / 14,126	10,832 / 10,294
hfep stack top (#2754 at a3370d2) + this change (5,830 passed, 35 skipped, 12 xfailed)	15,633 / 15,823	11,089 / 10,914

The committed tree itself (main c00ac43 + this commit; 5,355 passed, 31 skipped, 12 xfailed), one run: 10,232 MiB. The pool files ran on one worker (243 tests on main, 743 on the hfep tree) and no file was split across workers; the group takes 53 s of test time on main (68 s on the hfep tree), so it is not on the critical path. RSS counts shared pages, so the absolute numbers overstate the real use; the comparison is what matters.

On the hosted runners the unit workflow of this PR ran 5 times (attempt 1 and 4 re-runs of the same commit) plus one run of the commit before the rebase; all passed.

No test was added: the hook only sets scheduling markers, and the whole unit suite runs under it.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: N/A (CI and test scheduling)
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: no new dependency
• Did you write any new necessary tests?: N/A (scheduling hook; the unit suite exercises it)
• Did you update Changelog?: N/A (CI)
• Did you get Claude approval on this PR?: N/A (not run yet)

Additional Information

Jira: [OMNIML-6178]. Only files that use cpu_dist_workers are grouped: another multi-process fixture or memory-heavy file added later is not covered.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.60%. Comparing base (659d317) to head (3d3cbea).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2953   +/-   ##
=======================================
  Coverage   79.60%   79.60%           
=======================================
  Files         660      660           
  Lines       74667    74667           
=======================================
  Hits        59435    59435           
  Misses      15232    15232           

Flag	Coverage Δ
unit	61.11% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[shengliangxu]

Closing this draft.

• The runner stops it was meant to address ended with [OMNIML-6228] Shrink the input of the conv/IsInf ONNX test model #2947: all 32 finished torch_28 attempts since [OMNIML-6260] Pin the hosted Linux runners to ubuntu-24.04 #2950 let CI run again passed (23 on the hfep/epdp stack, where 32% were stopped before, on the same runner image), and the conv/IsInf ONNX tests were in flight in 93 of the 94 earlier stops.
• On hosted runners the pool group shows no measurable speed gain: linux pytest time 245 s median over 6 runs of this PR against 242 s for other trees over 12 runs (a single run varies by about 20%).
• The memory reduction is real (PSS 12,984 to 9,836 MiB, 20 to 12 processes), but nothing fails for lack of it now.

The branch stays (shengliangx/ci-xdist-pool-group, 3d3cbea): reopen the PR if memory headroom of the unit lane matters again.

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-10-23 16:20 UTC