pytorch · kwen2501 · May 15, 2025 · May 16, 2025 · May 16, 2025 · May 17, 2025
diff --git a/test/distributed/test_c10d_ops_nccl.py b/test/distributed/test_c10d_ops_nccl.py
@@ -11,7 +11,6 @@
 import math
 import os
 import sys
-import tempfile
 
 import torch
 import torch.distributed as c10d
@@ -30,9 +29,9 @@
     requires_nccl,
     requires_nccl_version,
     sm_is_or_higher_than,
-    TEST_SKIPS,
 )
 from torch.testing._internal.common_utils import (
+    run_tests,
     skip_but_pass_in_sandcastle_if,
     skipIfRocm,
     TEST_WITH_DEV_DBG_ASAN,
@@ -1044,24 +1043,4 @@ def allgather_base(output_t, input_t):
 
 
 if __name__ == "__main__":
-    if not torch.cuda.is_available():
-        sys.exit(TEST_SKIPS["no_cuda"].exit_code)
-
-    rank = int(os.getenv("RANK", -1))
-    world_size = int(os.getenv("WORLD_SIZE", -1))
-
-    if world_size == -1:  # Not set by external launcher
-        world_size = torch.cuda.device_count()
-
-    if rank != -1:
-        # Launched with torchrun or other multi-proc launchers. Directly run the test.
-        ProcessGroupNCCLOpTest.run_rank(rank, world_size)
-    else:
-        # Launched as a single process. Spawn subprocess to run the tests.
-        # Also need a rendezvous file for `init_process_group` purpose.
-        rdvz_file = tempfile.NamedTemporaryFile(delete=False).name
-        torch.multiprocessing.spawn(
-            ProcessGroupNCCLOpTest.run_rank,
-            nprocs=world_size,
-            args=(world_size, rdvz_file),
-        )
+    run_tests()
diff --git a/test/distributed/test_nvshmem.py b/test/distributed/test_nvshmem.py
@@ -7,16 +7,13 @@
 
 import os
 import sys
-import tempfile
 
 import torch
 import torch.distributed as dist
 import torch.distributed._symmetric_memory as symm_mem
-from torch.testing._internal.common_distributed import (
-    MultiProcContinousTest,
-    TEST_SKIPS,
-)
+from torch.testing._internal.common_distributed import MultiProcContinousTest
 from torch.testing._internal.common_utils import (
+    run_tests,
     skip_but_pass_in_sandcastle_if,
     skipIfRocm,
 )
@@ -47,28 +44,20 @@ def requires_nvshmem():
 
 @requires_nvshmem()
 class NVSHMEMSymmetricMemoryTest(MultiProcContinousTest):
-    def setUp(self) -> None:
-        super().setUp()
+    def _init_device(self) -> None:
         # TODO: relieve this (seems to hang if without)
         device_module.set_device(self.device)
         # NOTE: required for nvshmem allocation
         torch.empty(1, device=self.device)
 
-    # Required by MultiProcContinousTest
-    @classmethod
-    def backend_str(cls) -> str:
-        return "nccl"
-
-    @property
-    def world_size(self) -> int:
-        return device_module.device_count()
-
     @property
     def device(self) -> torch.device:
         return torch.device(device_type, self.rank)
 
     @skipIfRocm
     def test_nvshmem_all_to_all(self) -> None:
+        self._init_device()
+
         group_name = dist.group.WORLD.group_name
         symm_mem.enable_symm_mem_for_group(group_name)
 
@@ -92,6 +81,8 @@ def test_nvshmem_all_to_all(self) -> None:
 
     @skipIfRocm
     def test_nvshmem_all_to_all_vdev(self) -> None:
+        self._init_device()
+
         group_name = dist.group.WORLD.group_name
         symm_mem.enable_symm_mem_for_group(group_name)
 
@@ -139,24 +130,4 @@ def test_nvshmem_all_to_all_vdev(self) -> None:
 
 
 if __name__ == "__main__":
-    if not device_module.is_available():
-        sys.exit(TEST_SKIPS["no_cuda"].exit_code)
-
-    # If launched by torchrun, these values would have been set
-    rank = int(os.getenv("RANK", "-1"))
-    world_size = int(os.getenv("WORLD_SIZE", "-1"))
-
-    if rank != -1:
-        # Launched with torchrun or other multi-proc launchers. Directly run the test.
-        NVSHMEMSymmetricMemoryTest.run_rank(rank, world_size)
-    else:
-        # No external launcher, spawn N processes
-        world_size = device_module.device_count()
-        # Launched as a single process. Spawn subprocess to run the tests.
-        # Also need a rendezvous file for `init_process_group` purpose.
-        rdvz_file = tempfile.NamedTemporaryFile(delete=False).name
-        torch.multiprocessing.spawn(
-            NVSHMEMSymmetricMemoryTest.run_rank,
-            nprocs=world_size,
-            args=(world_size, rdvz_file),
-        )
+    run_tests()
diff --git a/test/distributed/test_template.py b/test/distributed/test_template.py
@@ -0,0 +1,16 @@
+# Owner(s): ["oncall: distributed"]
+
+from torch.testing._internal.common_distributed import MultiProcContinousTest
+from torch.testing._internal.common_utils import run_tests
+
+
+class TestTemplate(MultiProcContinousTest):
+    def testABC(self):
+        print(f"rank {self.rank} of {self.world_size} testing ABC")
+
+    def testDEF(self):
+        print(f"rank {self.rank} of {self.world_size} testing DEF")
+
+
+if __name__ == "__main__":
+    run_tests()