fix(tests): reuse real torch across collector test modules

tests/unit/collector/_real_torch.py

@@ -0,0 +1,38 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Borrow real torch without disturbing collector tests that require a mock."""
+
+import importlib
+import sys
+from collections.abc import Iterator
+from contextlib import contextmanager
+from types import ModuleType
+from unittest.mock import MagicMock
+
+import pytest
+
+_real_torch: ModuleType | None = None
+_MISSING = object()
+
+
+@contextmanager
+def real_torch() -> Iterator[ModuleType]:
+    """Borrow torch, restoring the original module state on exit."""
+    global _real_torch
+    previous = sys.modules.get("torch", _MISSING)
+    try:
+        if _real_torch is None:
+            if isinstance(previous, MagicMock):
+                sys.modules.pop("torch")
+            try:
+                _real_torch = importlib.import_module("torch")
+            except ImportError:
+                pytest.skip("real torch required for tensor operations", allow_module_level=True)
+        sys.modules["torch"] = _real_torch
+        yield _real_torch
+    finally:
+        if previous is _MISSING:
+            sys.modules.pop("torch", None)
+        else:
+            sys.modules["torch"] = previous

tests/unit/collector/test_dsv4_megamoe_workload.py

@@ -1,36 +1,26 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-import sys
 from argparse import Namespace
-from unittest.mock import MagicMock
 
 import pytest
 
-_saved_mock = sys.modules.get("torch")
-_restore_mock = isinstance(_saved_mock, MagicMock)
-if _restore_mock:
-    sys.modules.pop("torch")
+from tests.unit.collector._real_torch import real_torch
 
-try:
-    import torch as _real_torch
-except ImportError:
-    if _restore_mock:
-        sys.modules["torch"] = _saved_mock
-    pytest.skip("real torch required for tensor operations", allow_module_level=True)
-
-try:
+with real_torch() as torch:
     from collector.sglang.collect_dsv4_megamoe import build_cases
     from collector.sglang.dsv4_megamoe_workload import (
         _sampled_power_law_xmax,
         build_routing_plan,
         parse_distribution,
     )
-finally:
-    if _restore_mock:
-        sys.modules["torch"] = _saved_mock
 
-torch = _real_torch
+
+@pytest.fixture(autouse=True)
+def _use_real_torch():
+    # Routing helpers import torch lazily while the other collector tests need a mock.
+    with real_torch():
+        yield
 
 
 @pytest.mark.unit

tests/unit/collector/test_helper_moe_distribution.py

@@ -14,39 +14,16 @@
 """
 
 import sys
-from unittest.mock import MagicMock
 
 import pytest
 
-# test_collect_provenance_writer deliberately leaves a MagicMock cached as
-# sys.modules["torch"] (collect.py's fork-worker tests depend on it), so a
-# plain importorskip("torch") can "succeed" with the mock and every tensor
-# assertion below dies with `TypeError: '<' not supported between instances
-# of 'MagicMock' and 'int'` — whether that happens depends only on which
-# module pytest-xdist imports first on the worker (same pattern as
-# test_dsv4_megamoe_workload). Evict the mock, import the real torch (or
-# skip when it isn't installed), then put the mock back for the siblings
-# that rely on it. The restore lives in a `finally` so even an unexpected
-# import failure (e.g. an OSError loading torch's native libraries) cannot
-# leave the eviction in place; do NOT let collect.py see the real torch —
-# its fork-worker tests deadlock on macOS with a real torch cached.
-_saved_mock = sys.modules.get("torch")
-_restore_mock = isinstance(_saved_mock, MagicMock)
-if _restore_mock:
-    sys.modules.pop("torch")
-try:
-    import torch
-except ImportError:
-    # pytest.skip raises; the mock is restored by the finally during unwind.
-    pytest.skip("real torch required for tensor operations", allow_module_level=True)
-finally:
-    if _restore_mock:
-        sys.modules["torch"] = _saved_mock
-
-from collector.helper import (
-    _generate_power_law_distribution,
-    _round_robin_adjust_per_rank,
-)
+from tests.unit.collector._real_torch import real_torch
+
+with real_torch() as torch:
+    from collector.helper import (
+        _generate_power_law_distribution,
+        _round_robin_adjust_per_rank,
+    )
 
 pytestmark = pytest.mark.unit
 

tests/unit/collector/test_real_torch.py

@@ -0,0 +1,50 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Exercise real tensor operations across repeated mock-preserving imports."""
+
+import sys
+from unittest.mock import MagicMock
+
+import pytest
+
+from tests.unit.collector._real_torch import real_torch
+
+pytestmark = pytest.mark.unit
+
+
+@pytest.mark.parametrize("entry", [None, MagicMock()])
+def test_repeated_borrows_preserve_module_entry(monkeypatch, entry):
+    if entry is None:
+        monkeypatch.delitem(sys.modules, "torch", raising=False)
+    else:
+        monkeypatch.setitem(sys.modules, "torch", entry)
+
+    for _ in range(3):
+        with real_torch() as torch:
+            assert sys.modules["torch"] is torch
+            assert torch.arange(3).sum().item() == 3
+        if entry is None:
+            assert "torch" not in sys.modules
+        else:
+            assert sys.modules["torch"] is entry
+
+
+def test_borrow_restores_mock_when_body_raises(monkeypatch):
+    mock_torch = MagicMock()
+    monkeypatch.setitem(sys.modules, "torch", mock_torch)
+
+    with pytest.raises(RuntimeError, match="consumer failed"), real_torch() as torch:
+        assert torch.ones(2).sum().item() == 2
+        raise RuntimeError("consumer failed")
+
+    assert sys.modules["torch"] is mock_torch
+
+
+def test_borrow_preserves_existing_real_torch(monkeypatch):
+    with real_torch() as torch:
+        pass
+    monkeypatch.setitem(sys.modules, "torch", torch)
+    with real_torch() as borrowed:
+        assert borrowed is torch
+    assert sys.modules["torch"] is torch

tests/unit/collector/test_sglang_moe_ep_routing.py

@@ -4,51 +4,32 @@
 import importlib.util
 import sys
 from pathlib import Path
-from unittest.mock import MagicMock
 
 import pytest
 
+from tests.unit.collector._real_torch import real_torch
+
 
 def _import_helper_module():
     module_name = "collector.helper_test_copy"
     helper_path = Path(__file__).resolve().parents[3] / "collector" / "helper.py"
     spec = importlib.util.spec_from_file_location(module_name, helper_path)
     module = importlib.util.module_from_spec(spec)
+    sys.modules[module_name] = module
     spec.loader.exec_module(module)
     return module
 
 
-# test_parallel_run.py injects a MagicMock as "torch" into sys.modules so
-# collector code can be imported without CUDA.  This test needs real tensors
-# and a helper.py copy imported against real torch, but it must not permanently
-# replace the injected mock in sys.modules.
-_saved_mock = sys.modules.get("torch")
-_restore_mock = isinstance(_saved_mock, MagicMock)
-if _restore_mock:
-    sys.modules.pop("torch")
-
-try:
-    import torch as _real_torch
-except ImportError:
-    if _restore_mock:
-        sys.modules["torch"] = _saved_mock
-    pytest.skip("real torch required for tensor operations", allow_module_level=True)
-
-try:
+with real_torch() as torch:
     _HELPER_MODULE = _import_helper_module()
-finally:
-    if _restore_mock:
-        sys.modules["torch"] = _saved_mock
-
-torch = _real_torch
 
 
 @pytest.fixture(autouse=True)
 def _use_real_torch(monkeypatch):
     # At test execution time sys.modules["torch"] is still test_parallel_run.py's
     # MagicMock. helper.py functions do lazy `import torch`, so they pick up the
     # mock rather than the real module. Swap in real torch for each test's duration.
-    monkeypatch.setitem(sys.modules, "torch", _real_torch)
+    monkeypatch.setitem(sys.modules, "torch", torch)
 
 
 @pytest.mark.unit