test: unit suites for seven untested torch/utils modules (109 tests)

tests/unit/torch/utils/test_distributed.py

@@ -0,0 +1,143 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Single-process contracts only: everything here must hold without an initialized
+# torch.distributed process group.
+
+import os
+
+import pytest
+import torch
+from torch import nn
+
+import modelopt.torch.utils.distributed as dist
+
+
+def test_single_process_defaults():
+    assert not dist.is_initialized()
+    assert dist.backend() is None
+    assert dist.size() == 1
+    assert dist.rank() == 0
+    assert dist.is_master()
+    assert dist.is_last_process()
+
+
+def test_local_rank_from_env(monkeypatch):
+    monkeypatch.setenv("LOCAL_RANK", "3")
+    assert dist.local_rank() == 3
+
+
+def test_local_rank_fallback_warns_and_uses_global_rank(monkeypatch):
+    monkeypatch.delenv("LOCAL_RANK", raising=False)
+    with pytest.warns(UserWarning, match="LOCAL_RANK"):
+        assert dist.local_rank() == 0
+
+
+def test_broadcast_is_identity_in_single_process():
+    obj = {"a": [1, 2], "b": "x"}
+    assert dist.broadcast(obj) is obj
+
+
+def test_allgather_returns_singleton_list():
+    assert dist.allgather(5) == [5]
+    obj = {"k": 1}
+    assert dist.allgather(obj) == [obj]
+
+
+def test_allreduce_sum_and_unsupported_reduction():
+    assert dist.allreduce(5) == 5
+    with pytest.raises(NotImplementedError):
+        dist.allreduce(5, reduction="max")
+
+
+def test_barrier_is_noop():
+    assert dist.barrier() is None
+
+
+def test_master_only_runs_and_returns_result():
+    @dist.master_only
+    def compute():
+        return {"x": 1}
+
+    assert compute() == {"x": 1}
+
+
+def test_serialize_deserialize_round_trip():
+    obj = {"w": torch.arange(4, dtype=torch.float32), "meta": "hi"}
+    tensor = dist._serialize(obj)
+    assert tensor.dtype == torch.uint8
+    restored = dist._deserialize(tensor)
+    assert restored["meta"] == "hi"
+    assert torch.equal(restored["w"], obj["w"])
+    # The size argument slices the buffer back to the payload; note the
+    # deserializer tolerates trailing bytes either way (pickle stops at its
+    # own payload), so this pins correct size handling, not strict trimming.
+    padded = torch.cat([tensor, torch.zeros(8, dtype=torch.uint8)])
+    restored2 = dist._deserialize(padded, size=tensor.numel())
+    assert restored2["meta"] == "hi"
+    assert torch.equal(restored2["w"], obj["w"])
+
+
+@pytest.mark.parametrize("group", [None, -1])
+def test_process_group_uninitialized_defaults(group):
+    pg = dist.DistributedProcessGroup(group)
+    assert not pg.is_initialized()
+    assert pg.rank() == -1
+    assert pg.world_size() == -1
+    assert "initialized: False" in repr(pg)
+
+
+def test_get_dist_syncd_obj_passthrough_when_uninitialized():
+    pg = dist.DistributedProcessGroup(None)
+    # op must NOT be applied for uninitialized groups
+    assert dist.DistributedProcessGroup.get_dist_syncd_obj(7, pg, sum) == 7
+    assert dist.DistributedProcessGroup.get_dist_syncd_obj(7, [pg, pg], max) == 7
+
+
+def test_parallel_state_defaults_and_repr():
+    state = dist.ParallelState()
+    assert not state.data_parallel_group.is_initialized()
+    assert not state.tensor_parallel_group.is_initialized()
+    assert not state.expert_model_parallel_group.is_initialized()
+    assert state.data_parallel_group.world_size() == -1
+    for name in ("data_parallel_group", "tensor_parallel_group", "expert_model_parallel_group"):
+        assert name in repr(state)
+
+
+def test_get_group_returns_none_when_uninitialized():
+    assert dist.get_group([0]) is None
+
+
+def test_is_dtensor_sharded_false_for_plain_model():
+    assert not dist.is_dtensor_sharded(nn.Linear(2, 2))
+
+
+def test_filelock_context_creates_and_removes_lockfile(tmp_path):
+    lock_path = str(tmp_path / "my.lock")
+    with dist.FileLock(lock_path):
+        assert os.path.exists(lock_path)
+    assert not os.path.exists(lock_path)
+
+
+def test_filelock_try_acquire_conflict(tmp_path):
+    lock_path = str(tmp_path / "my.lock")
+    first = dist.FileLock(lock_path)
+    second = dist.FileLock(lock_path)
+    assert first.try_acquire()
+    assert not second.try_acquire()  # already held
+    first.release()
+    assert second.try_acquire()
+    second.release()
+    assert not os.path.exists(lock_path)

tests/unit/torch/utils/test_graph.py

@@ -0,0 +1,129 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch.nn.functional as F
+from torch import nn
+
+from modelopt.torch.utils.graph import match
+
+
+class LinearRelu(nn.Module):
+    def __init__(self, features=8):
+        super().__init__()
+        self.fc = nn.Linear(features, features)
+        self.act = nn.ReLU()
+
+    def forward(self, x):
+        return self.act(self.fc(x))
+
+
+class LinearSigmoid(nn.Module):
+    def __init__(self, features=8):
+        super().__init__()
+        self.fc = nn.Linear(features, features)
+        self.act = nn.Sigmoid()
+
+    def forward(self, x):
+        return self.act(self.fc(x))
+
+
+class LinearFuncRelu(nn.Module):
+    """Same math as LinearRelu but relu as call_function instead of call_module."""
+
+    def __init__(self, features=8):
+        super().__init__()
+        self.fc = nn.Linear(features, features)
+
+    def forward(self, x):
+        return F.relu(self.fc(x))
+
+
+class LinearMethodRelu(nn.Module):
+    """relu invoked as a tensor method (call_method node)."""
+
+    def __init__(self, features=8):
+        super().__init__()
+        self.fc = nn.Linear(features, features)
+
+    def forward(self, x):
+        return self.fc(x).relu()
+
+
+class JustLinear(nn.Module):
+    def __init__(self, features=8):
+        super().__init__()
+        self.fc = nn.Linear(features, features)
+
+    def forward(self, x):
+        return self.fc(x)
+
+
+class Untraceable(nn.Module):
+    """Data-dependent control flow makes this module untraceable by torch.fx."""
+
+    def forward(self, x):
+        if x.sum() > 0:
+            return x
+        return -x
+
+
+def test_match_same_structure():
+    assert match(LinearRelu(), [LinearRelu()])
+
+
+def test_match_compares_module_types_not_instances():
+    # different layer sizes must still match since submodules are compared by type
+    assert match(LinearRelu(features=4), [LinearRelu(features=16)])
+
+
+def test_no_match_different_activation():
+    assert not match(LinearRelu(), [LinearSigmoid()])
+
+
+def test_no_match_different_node_count():
+    assert not match(JustLinear(), [LinearRelu()])
+    assert not match(LinearRelu(), [JustLinear()])
+
+
+def test_empty_patterns_never_match():
+    assert not match(LinearRelu(), [])
+
+
+def test_match_any_of_multiple_patterns():
+    assert match(LinearRelu(), [LinearSigmoid(), JustLinear(), LinearRelu()])
+
+
+def test_call_function_nodes_match():
+    assert match(LinearFuncRelu(), [LinearFuncRelu()])
+
+
+def test_call_module_does_not_match_call_function():
+    # nn.ReLU (call_module) vs F.relu (call_function) are different graph ops
+    assert not match(LinearRelu(), [LinearFuncRelu()])
+
+
+def test_call_method_nodes_match():
+    assert match(LinearMethodRelu(), [LinearMethodRelu()])
+    assert not match(LinearMethodRelu(), [LinearFuncRelu()])
+
+
+def test_untraceable_module_returns_false():
+    assert not match(Untraceable(), [LinearRelu()])
+
+
+def test_sequential_matches_equivalent_custom_module():
+    seq = nn.Sequential(nn.Linear(8, 8), nn.ReLU())
+    assert match(seq, [LinearRelu()])
+    assert match(LinearRelu(), [seq])

tests/unit/torch/utils/test_import_utils.py

@@ -0,0 +1,88 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import warnings
+from contextlib import contextmanager
+
+import pytest
+
+from modelopt.torch.utils.import_utils import import_plugin
+
+
+@contextmanager
+def no_warnings():
+    """Fail the test if any warning is emitted inside the block."""
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        yield
+
+
+def _import_missing_module(tracker=None, **kwargs):
+    """Run import_plugin around an import that always fails with ModuleNotFoundError."""
+    with import_plugin("foo", **kwargs):
+        __import__("definitely_not_a_real_module_xyz")
+        if tracker is not None:
+            tracker.append("ran past import")
+
+
+def _raise_runtime_error(**kwargs):
+    with import_plugin("myplugin", **kwargs):
+        raise RuntimeError("boom")
+
+
+def test_no_error_and_no_success_msg_is_silent():
+    with no_warnings(), import_plugin("foo"):
+        pass
+
+
+def test_success_msg_emitted_on_clean_import():
+    with (
+        pytest.warns(UserWarning, match="foo plugin loaded"),
+        import_plugin("foo", success_msg="foo plugin loaded"),
+    ):
+        pass
+
+
+def test_module_not_found_suppressed_with_message():
+    tracker = []
+    with pytest.warns(UserWarning, match="foo is not installed"):
+        _import_missing_module(tracker=tracker, msg_if_missing="foo is not installed")
+    assert not tracker  # code after the failing import must not run
+
+
+def test_module_not_found_without_message_is_silent():
+    with no_warnings():
+        _import_missing_module()
+
+
+def test_module_not_found_verbose_false_is_silent():
+    with no_warnings():
+        _import_missing_module(msg_if_missing="should not appear", verbose=False)
+
+
+def test_generic_exception_suppressed_with_warning():
+    with pytest.warns(UserWarning, match="Failed to import modelopt myplugin plugin") as record:
+        _raise_runtime_error()
+    assert "RuntimeError('boom')" in str(record[0].message)
+
+
+def test_generic_exception_verbose_false_is_silent():
+    with no_warnings():
+        _raise_runtime_error(verbose=False)
+
+
+def test_success_msg_not_emitted_when_import_fails():
+    with no_warnings():
+        _import_missing_module(success_msg="should not appear", verbose=True)