From d836289045bd233044aaee362140c2db9b0ab1a8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 21 Oct 2026 22:09:25 +0000 Subject: [PATCH 1/2] Run the HF checkpoint multi-rank tests on the shared CPU worker pool 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 --- .../utils/test_hf_checkpoint_utils_dist.py | 76 ++++++++++++------- 1 file changed, 48 insertions(+), 28 deletions(-) rename tests/{gpu => unit}/torch/utils/test_hf_checkpoint_utils_dist.py (57%) diff --git a/tests/gpu/torch/utils/test_hf_checkpoint_utils_dist.py b/tests/unit/torch/utils/test_hf_checkpoint_utils_dist.py similarity index 57% rename from tests/gpu/torch/utils/test_hf_checkpoint_utils_dist.py rename to tests/unit/torch/utils/test_hf_checkpoint_utils_dist.py index 72322ad50b4..d99b935e487 100644 --- a/tests/gpu/torch/utils/test_hf_checkpoint_utils_dist.py +++ b/tests/unit/torch/utils/test_hf_checkpoint_utils_dist.py @@ -13,14 +13,17 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Multi-rank tests for ``ensure_local_checkpoint`` (each case spawns a 2-process gloo job).""" +"""Multi-rank tests for ``ensure_local_checkpoint`` (the shared 2-rank gloo pool runs each case). + +The workers outlive a case, so every job undoes what it patches and destroys the groups it makes. +""" import os from functools import partial +from unittest import mock import pytest import torch.distributed as dist -from _test_utils.torch.distributed.utils import spawn_multiprocess_job from modelopt.torch.utils.plugins import hf_checkpoint_utils @@ -46,19 +49,26 @@ def fake_snapshot_download(repo_id, **kwargs): snapshot.mkdir(exist_ok=True) return str(snapshot) - utils.snapshot_download = fake_snapshot_download - utils.HfApi = _FakeHfApi - os.environ.pop("HF_HUB_OFFLINE", None) # inherited from tests/unit/conftest.py - # new_group is collective over every rank, members or not. - group = None if group_ranks is None else dist.new_group(group_ranks, backend="gloo") - if group_ranks is not None and rank not in group_ranks: - return - if fail: - with pytest.raises(RuntimeError, match="hub unreachable"): - utils.ensure_local_checkpoint("org/model", group=group) - else: - result = utils.ensure_local_checkpoint("org/model", group=group) - assert result == ("org/model", str(snapshot)) + with ( + mock.patch.object(utils, "snapshot_download", fake_snapshot_download), + mock.patch.object(utils, "HfApi", _FakeHfApi), + mock.patch.dict(os.environ), # restores the variable popped below + ): + os.environ.pop("HF_HUB_OFFLINE", None) # set for every unit test by tests/unit/conftest.py + # new_group is collective over every rank, members or not. + group = None if group_ranks is None else dist.new_group(group_ranks, backend="gloo") + try: + if group_ranks is not None and rank not in group_ranks: + return + if fail: + with pytest.raises(RuntimeError, match="hub unreachable"): + utils.ensure_local_checkpoint("org/model", group=group) + else: + result = utils.ensure_local_checkpoint("org/model", group=group) + assert result == ("org/model", str(snapshot)) + finally: + if group is not None: # None would mean the default group + dist.destroy_process_group(group) def _ensure_local_checkpoint_local_path_job(tmp_path, hidden_from_rank1, rank, size): @@ -67,25 +77,33 @@ def _ensure_local_checkpoint_local_path_job(tmp_path, hidden_from_rank1, rank, s def fail(*args, **kwargs): raise AssertionError("a local path must not be downloaded") - utils.snapshot_download = fail checkpoint = tmp_path / "checkpoint" checkpoint.mkdir(exist_ok=True) - if hidden_from_rank1 and rank == 1: - # As if the directory sat on rank 0's node-local disk only. - isdir = os.path.isdir - utils.os.path.isdir = lambda path: path != str(checkpoint) and isdir(path) - with pytest.raises(RuntimeError, match="this rank cannot see"): - utils.ensure_local_checkpoint(str(checkpoint)) - else: - assert utils.ensure_local_checkpoint(str(checkpoint)) == (None, str(checkpoint)) + with mock.patch.object(utils, "snapshot_download", fail): + if hidden_from_rank1 and rank == 1: + # As if the directory sat on rank 0's node-local disk only. + isdir = os.path.isdir + + def hide(path): + return path != str(checkpoint) and isdir(path) + + with ( + mock.patch.object(os.path, "isdir", hide), + pytest.raises(RuntimeError, match="this rank cannot see"), + ): + utils.ensure_local_checkpoint(str(checkpoint)) + else: + assert utils.ensure_local_checkpoint(str(checkpoint)) == (None, str(checkpoint)) @pytest.mark.parametrize("hidden_from_rank1", [False, True]) -def test_ensure_local_checkpoint_lets_rank_0_resolve_a_local_path(tmp_path, hidden_from_rank1): +def test_ensure_local_checkpoint_lets_rank_0_resolve_a_local_path( + tmp_path, hidden_from_rank1, cpu_dist_workers +): """Ranks that each checked a local path could disagree about it and strand one another in the collective; rank 0 decides, and a rank that cannot see its answer raises instead.""" job = partial(_ensure_local_checkpoint_local_path_job, tmp_path, hidden_from_rank1) - spawn_multiprocess_job(2, job, backend="gloo") + cpu_dist_workers.run(job) @pytest.mark.parametrize( @@ -98,10 +116,12 @@ def test_ensure_local_checkpoint_lets_rank_0_resolve_a_local_path(tmp_path, hidd (False, [1], "rank1"), ], ) -def test_ensure_local_checkpoint_downloads_once_per_group(tmp_path, fail, group_ranks, downloader): +def test_ensure_local_checkpoint_downloads_once_per_group( + tmp_path, fail, group_ranks, downloader, cpu_dist_workers +): """Rank 0 of the group downloads once and every member gets its result -- including its failure, so the waiting ranks raise instead of hanging.""" job = partial(_ensure_local_checkpoint_job, tmp_path, fail, group_ranks) - spawn_multiprocess_job(2, job, backend="gloo") + cpu_dist_workers.run(job) assert [p.name for p in tmp_path.glob("downloaded-by-*")] == [f"downloaded-by-{downloader}"] assert (tmp_path / f"downloaded-by-{downloader}").read_text().splitlines() == ["org/model"] From 9ce78ae735bd9f0188be25ac1f3fba4842f88a87 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 22 Oct 2026 17:10:22 +0000 Subject: [PATCH 2/2] Destroy only real process groups in the checkpoint test cleanup 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 --- tests/unit/torch/utils/test_hf_checkpoint_utils_dist.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit/torch/utils/test_hf_checkpoint_utils_dist.py b/tests/unit/torch/utils/test_hf_checkpoint_utils_dist.py index d99b935e487..ace69b6c4b8 100644 --- a/tests/unit/torch/utils/test_hf_checkpoint_utils_dist.py +++ b/tests/unit/torch/utils/test_hf_checkpoint_utils_dist.py @@ -67,7 +67,8 @@ def fake_snapshot_download(repo_id, **kwargs): result = utils.ensure_local_checkpoint("org/model", group=group) assert result == ("org/model", str(snapshot)) finally: - if group is not None: # None would mean the default group + # Skip None (the default group) and the sentinel a non-member rank gets from new_group. + if isinstance(group, dist.ProcessGroup): dist.destroy_process_group(group)