Distributed initialization process for single process

pfnet · Oct 30, 2023 · 5485a79 · 5485a79
1 parent c5b4d58
commit 5485a79
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/pytorch_pfn_extras/distributed/_initialize.py b/pytorch_pfn_extras/distributed/_initialize.py
@@ -1,7 +1,7 @@
 import os
 from typing import Tuple
 
-import torch
+import torch.distributed
 
 
 def initialize_ompi_environment(
@@ -60,7 +60,7 @@ def initialize_ompi_environment(
             "Invalid value for init_method, only 'env' and 'tcp' are supported"
         )
 
-    if world_size > 1 and not torch.distributed.is_initialized():  # type: ignore
+    if not torch.distributed.is_initialized():  # type: ignore
         torch.distributed.init_process_group(  # type: ignore
             backend, init_method=init_method, world_size=world_size, rank=rank
         )

diff --git a/tests/pytorch_pfn_extras_tests/distributed_tests/test_initialize.py b/tests/pytorch_pfn_extras_tests/distributed_tests/test_initialize.py
@@ -0,0 +1,12 @@
+import torch.distributed
+from pytorch_pfn_extras.distributed import initialize_ompi_environment
+
+
+def test_initialize_ompi_environment_with_single_process():
+    assert not torch.distributed.is_initialized()
+    world_size, rank, local_rank = initialize_ompi_environment(
+        backend="gloo", init_method="tcp"
+    )
+    assert torch.distributed.is_initialized()
+    torch.distributed.destroy_process_group()
+    assert not torch.distributed.is_initialized()