Repository navigation
[OMNIML-6179] Run the HF checkpoint multi-rank tests on the shared CPU worker pool - #2923
Conversation
|
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. |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 28.57% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 7 functions across 1 files. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly and concisely describes the main change: moving HF checkpoint multi-rank tests to the shared CPU worker pool. Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Security Anti-Patterns PASS. The PR changes only tests/gpu/torch/utils/test_hf_checkpoint_utils_dist.py, renamed totests/unit/torch/utils/test_hf_checkpoint_utils_dist.py. It makes no changes to themodeloptpackage,…Full details: Security Anti-Patterns
-
Explanation
PASS. The PR changes only
tests/gpu/torch/utils/test_hf_checkpoint_utils_dist.py, renamed totests/unit/torch/utils/test_hf_checkpoint_utils_dist.py. It makes no changes to themodeloptpackage, examples,pyproject.toml, orrequirements.txt, so none of the specified security anti-patterns apply.
-
✨ Finishing Touches: 💡 1
-
📝 Generate docstrings: 💡
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
|
/ok to test 30730c1 |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2923 +/- ##
==========================================
- Coverage 79.66% 79.65% -0.01%
==========================================
Files 659 659
Lines 74400 74408 +8
==========================================
Hits 59272 59272
- Misses 15128 15136 +8
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
The six ensure_local_checkpoint cases are gloo/CPU tests with a fake HfApi, but each one spawned a fresh 2-process job from tests/gpu (spawn_multiprocess_job), paying about 11 s of imports per case on a GPU runner (6 cases, about 69 s of the gpu job on a pull request, about 85 s nightly). They move next to the module's other unit tests and run on the session-scoped cpu_dist_workers pool, where a case takes 0.01-0.02 s once the pool is up. The workers outlive a case, so each job now undoes what it changes: snapshot_download, HfApi and os.path.isdir are patched with mock.patch.object, HF_HUB_OFFLINE is restored with mock.patch.dict(os.environ), and the group a job creates is destroyed. Cases and assertions are unchanged. Testing: the moved file alone, the whole tests/unit/torch/utils directory under xdist, and that directory together with the other cpu_dist_workers users in default, reversed and shuffled file order, on torch 2.10 and 2.14 (CPU only); a probe that runs the six jobs on one pool and compares the workers' state before and after (HF_HUB_OFFLINE, the patched names, the process-group registry) passes, and fails on the old job bodies. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
30730c1 to
d836289
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Scope: full review. 1 file changed (48 additions, 28 deletions): tests/gpu/torch/utils/test_hf_checkpoint_utils_dist.py → tests/unit/torch/utils/test_hf_checkpoint_utils_dist.py. Reviewed the whole moved file, plus ensure_local_checkpoint/_make_local/_snapshot_download in modelopt/torch/utils/plugins/hf_checkpoint_utils.py, DistributedWorkerPool/get_live_pool/default_worker_teardown in tests/_test_utils/torch/distributed/utils.py, and the cpu_dist_workers fixture and timeout mapping in tests/unit/conftest.py / tests/conftest.py. No modelopt/ or examples/ source changed, so no mode-registration, state-dict or export surface is in play.
Findings: CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1
What I checked and what holds
- Collective symmetry on a shared pool. This is the real risk of the move: a persistent pool can be poisoned by a job that leaves ranks at different points in the collective sequence. All six cases are symmetric.
dist.new_group(group_ranks, ...)is called by both ranks with identical arguments before the member check; forgroup_ranks=Noneboth ranks enterensure_local_checkpoint, which issues its ownnew_groupon both; for[0, 1]both are members; for[1]the non-member returns before any further collective and the single member broadcasts inside a one-rank subgroup, so the default group's sequence stays aligned either way. - Failure paths are symmetric too. With
fail=True, rank 0'sOSErroris caught insideensure_local_checkpointand sent throughbroadcast_object_list, so both ranks complete the same collectives and then both raiseRuntimeError;pytest.raisesabsorbs it on both, so the pool is not markedbroken. Same forhidden_from_rank1: rank 1's "this rank cannot see" fires after the broadcast, on theos.path.isdircheck. - Cleanup is complete relative to the fixture's documented contract ("Jobs must leave the process group intact and undo any global state they set").
snapshot_download,HfApiandos.path.isdirmove from unrestored global assignment tomock.patch.object;HF_HUB_OFFLINEtomock.patch.dict(os.environ)(whoseclear()/update()restoration does propagate throughos.environto the real environment); every subgroup the job creates is destroyed in afinally, andensure_local_checkpointalready destroys the one it makes itself. The default group is never touched.default_worker_teardownoffers nothing to reuse here (the CPU pool passesteardown_fn=None, and that helper only resetsAcceleratorState), so the self-cleaning jobs are the right mechanism rather than a parallel one. - Hermeticity under
tests/unit. PoppingHF_HUB_OFFLINEsends_snapshot_downloaddown its online branch, but that branch touches onlyHfApi().repo_infoandsnapshot_download, both patched — no path reaches the network, so the unit suite stays offline. - A hang still fails fast.
tests/conftest.pygivestests/unita 120 s default cap, well inside the pool's 600 sjob_timeout; pytest-timeout fires first,runaborts the workers, andget_live_poolrebuilds for the next test, so a regression in the very behavior these tests guard costs one test rather than the lane. - Calling convention and assertions.
partial(fn, tmp_path, ...)plus the pool'sfn(rank, world_size, *args)reconstructs the original(..., rank, size)signature unchanged; workers write thedownloaded-by-*markers before reporting, so the post-runtmp_path.globassertions are well ordered. Cases and assertions are otherwise identical to thetests/gpuversion, and the author's mutation check (3 of 6 fail with the rank-0 check stubbed toif True) backs that they still bite. - No basename collision with the sibling
test_hf_checkpoint_utils.py, and theimportorskip-after-module-import ordering matches what that already-in-tests/unitsibling does, so it is a pre-existing convention rather than something this move introduces.
The one suggestion is non-blocking: the finally guard leans on destroy_process_group early-returning for the NON_GROUP_MEMBER sentinel a non-member rank receives, which the adjacent comment doesn't mention.
Risk: low. Tests-only, no public surface, no behavior change for users; the move is well matched to CONTRIBUTING's "reuse the session-scoped cpu_dist_workers fixture for multi-rank CPU tests", and the net effect is removing ~69-85 s of per-case process spawning from the gpu lane.
🤖 Generated with Claude Code
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
LGTM: the test migration preserves all six cases and assertions, with cleanup appropriate for persistent workers; justified existing-test edits require human sign-off.
Needs action:
- Sign off on the justified migration in
tests/unit/torch/utils/test_hf_checkpoint_utils_dist.py: CPU/gloo cases reuse the existing pool without reducing coverage.
No action needed:
- Mock contexts restore patched globals and environment; explicit groups are destroyed without touching the default group.
- Tests were inspected, not executed during this review.
A rank outside the group gets the NON_GROUP_MEMBER sentinel from new_group, which destroy_process_group happens to ignore. Skip it explicitly instead of relying on that. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
LGTM: the prior cleanup suggestion is resolved, and all six cases retain their assertions; the justified existing-test migration requires human sign-off.
Needs action:
- Sign off on the justified test migration in
tests/unit/torch/utils/test_hf_checkpoint_utils_dist.py: CPU/gloo cases reuse the shared pool without reducing coverage.
No action needed:
- ✔️ Resolved since the last review: cleanup explicitly skips both the default group and the non-member sentinel using
isinstance(group, dist.ProcessGroup). - Mock contexts restore patched globals and environment variables; collective ordering and process-group cleanup remain appropriate for persistent workers.
- Tests were inspected, not executed during this review.
|
What does this PR do?
Type of change: CI (test runtime, tests only)
The six
ensure_local_checkpointmulti-rank cases are gloo/CPU tests with a fakeHfApi, but they lived intests/gpuand each one spawned a fresh 2-process job (spawn_multiprocess_job). The spawn re-imports torch, transformers and modelopt in both children: every case took 11.2-11.8 s in thegpujob of a pull request (13.7-14.3 s on the nightly's 2-GPU runner), 69 s and 85 s per run, to test a few lines of rank-0 logic.The file moves to
tests/unit/torch/utils/, next to the module's other unit tests, and runs on the session-scopedcpu_dist_workerspool that other multi-rank unit tests already share (CONTRIBUTING: "reuse the session-scopedcpu_dist_workersfixture for multi-rank CPU tests"). Cases and assertions are unchanged. A case takes 0.01-0.02 s once the pool is up.The pool's workers outlive a case, so each job now undoes what it changes instead of relying on a throwaway process:
snapshot_download,HfApiandos.path.isdirare patched withmock.patch.object,HF_HUB_OFFLINE(whichtests/unit/conftest.pysets for every unit test) is restored withmock.patch.dict(os.environ), and the process group a job creates is destroyed. The default group is never touched.Time improvement (measured)
gpujob, per-case time from the result-line gaps of two CI logs of the base:--run-manual, the nightly's setup (run 39668011982)unitjob: the six cases take 0.01-0.02 s each on a warm pool (CPU host, torch 2.14.1). The first test of an xdist worker that touches the pool pays the pool start (about 7 s on this host), but that cost is already paid by the six existing pool-using files, so the unit job gains about 0.1 s.Memory (the pool question from the xdist work): this is a seventh file that uses
cpu_dist_workers; the six existing ones aretest_fsdp2_parallel_export.py,test_autoquant.py,test_autoquant_shapley.py,quantization/test_dist.py,test_quantize_cpu.pyandquantization/test_utils.py. A pool is one per xdist worker (two gloo processes, about 0.62 GB each after their first job), so the ceiling is still four pools. The jobs add about 5 MiB to each pool process (623-624 MiB after the six jobs, against 619 MiB after the imports of a typical existing user). The only way this file could raise the unit job's peak is by giving a pool to an xdist worker that would otherwise have none, at most about 1.2-1.4 GB on that worker, which is unlikely with six other users. A separate change to the unit job's xdist grouping (every file with acpu_dist_workerstest goes to one xdist group, chosen fromitem.fixturenames) would pick this file up with no change.Usage
Nothing user-facing.
Testing
All on CPU with
CUDA_VISIBLE_DEVICES="", in the CI-like venv (Python 3.12, torch 2.14.1 CPU, xdist 3.8.0) unless noted:tests/unit/torch/utilsunder-n 4 --dist loadfile: 371 passed.cpu_dist_workersusers (quantization/test_dist.py,export/test_fsdp2_parallel_export.py,quantization/test_utils.py) in one process in default, reversed and two shuffled file orders (seeds 11 and 12): 403 passed each.HF_HUB_OFFLINE,snapshot_download,HfApi,os.path.isdir, the default group and the process-group registry): identical on torch 2.10 and 2.14.1. The same probe on the old job bodies, which patch globals and never restore them, reportsHF_HUB_OFFLINElost, the two names still patched,os.path.isdirpatched on rank 1 and extra process groups left behind.ensure_local_checkpointreplaced byif True, 3 of the 6 fail; with the "this rank cannot see" check removed, 1 fails.This PR's own unit job on a hosted runner passed. The
gpuworkflow was skipped by path gating (the file only leavestests/gpu), so the saving shows in CI on the first gpu run after this merges; the old times come from the CI logs above. Not run: Windows (cpu_dist_workersskips there).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.).CONTRIBUTING.md: no new dependency/claude reviewapproved d836289; its one suggestion is applied in 9ce78ae)Additional Information
Found by a survey of the GPU lanes for per-case process starts (the pattern #2855 removed from
gpu_megatron). Jira: OMNIML-6179, under the CI runtime epic OMNIML-6177.Summary by CodeRabbit