Repository: pfnet/pytorch-pfn-extras

Issue:
DistributedSnapshot with NCCL
I want to use DistributedSnapshot with NCCL. However, current `_get_ranks_from_env` only accept environment variables for MPI and MV1. 
In pytorch documentation, they use the environment variables: `WORLD_SIZE`, `RANK`, and `LOCAL_RANK`, for specifying a process.

https://github.com/pfnet/pytorch-pfn-extras/blob/8b16df9433a024c63d786c76f1e7ccbb88847283/pytorch_pfn_extras/training/extensions/_snapshot.py#L473-L488


Discussion:
Hi, thanks for the feedback, this is because use to run our training directly with mpirun so we rely in these variables.
I agree this is not correct and I am going to fix it now.

Resolving patch:
diff --git a/pytorch_pfn_extras/training/extensions/_snapshot.py b/pytorch_pfn_extras/training/extensions/_snapshot.py
--- a/pytorch_pfn_extras/training/extensions/_snapshot.py
+++ b/pytorch_pfn_extras/training/extensions/_snapshot.py
@@ -454,8 +454,12 @@ def __init__(
                          autoload, savefun, transform_models,
                          autoload_transform_models)
         # To support distributed snapshots
+        if not torch.distributed.is_initialized():
+            raise RuntimeError('The Distributed Snapshot extension',
+                               ' requires torch.distributed to be initialized')
         self._saver_rank = saver_rank
-        self._size, self._rank, self._local_rank = _get_ranks_from_env()
+        self._size = torch.distributed.get_world_size()
+        self._rank = torch.distributed.get_rank()
         if not (0 <= saver_rank < self._size):
             raise ValueError('Distributed snapshot requires a saver rank'
                              ' in the range [0-{})'.format(self._size))
@@ -468,21 +472,3 @@ def __call__(self, trainer):
                 self._make_snapshot(trainer)
             if self._size > 1:
                 torch.distributed.barrier()
-
-
-def _get_ranks_from_env():
-    if 'OMPI_COMM_WORLD_SIZE' in os.environ:
-        # We are running Open MPI
-        comm_world_size = int(os.environ["OMPI_COMM_WORLD_SIZE"])
-        comm_rank = int(os.environ["OMPI_COMM_WORLD_RANK"])
-        comm_local_rank = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK'])
-    elif 'MV2_COMM_WORLD_SIZE' in os.environ:
-        comm_world_size = int(os.environ['MV2_COMM_WORLD_SIZE'])
-        comm_rank = int(os.environ['MV2_COMM_WORLD_RANK'])
-        comm_local_rank = int(os.environ['MV2_COMM_WORLD_LOCAL_RANK'])
-    else:
-        comm_world_size = 1
-        comm_rank = 0
-        comm_local_rank = 0
-
-    return comm_world_size, comm_rank, comm_local_rank


Test patch:
diff --git a/tests/pytorch_pfn_extras_tests/training_tests/extensions_tests/test_distributed_snapshot.py b/tests/pytorch_pfn_extras_tests/training_tests/extensions_tests/test_distributed_snapshot.py
--- a/tests/pytorch_pfn_extras_tests/training_tests/extensions_tests/test_distributed_snapshot.py
+++ b/tests/pytorch_pfn_extras_tests/training_tests/extensions_tests/test_distributed_snapshot.py
@@ -9,6 +9,23 @@
 from pytorch_pfn_extras.training import extensions
 
 
+def get_ranks_from_env():
+    if 'OMPI_COMM_WORLD_SIZE' in os.environ:
+        # We are running Open MPI
+        comm_world_size = int(os.environ["OMPI_COMM_WORLD_SIZE"])
+        comm_rank = int(os.environ["OMPI_COMM_WORLD_RANK"])
+        comm_local_rank = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK'])
+    elif 'MV2_COMM_WORLD_SIZE' in os.environ:
+        comm_world_size = int(os.environ['MV2_COMM_WORLD_SIZE'])
+        comm_rank = int(os.environ['MV2_COMM_WORLD_RANK'])
+        comm_local_rank = int(os.environ['MV2_COMM_WORLD_LOCAL_RANK'])
+    else:
+        comm_world_size = 1
+        comm_rank = 0
+        comm_local_rank = 0
+    return comm_world_size, comm_rank, comm_local_rank
+
+
 def _create_distributed_model(gpu=True):
     comm_size, comm_rank, comm_local_rank, device = _init_distributed(True)
 
@@ -42,11 +59,12 @@ def _init_distributed(use_cuda):
             or 'MV2_COMM_WORLD_SIZE' in os.environ):
         # This test assumes NCCL backend.
         size, rank, local_rank = (
-            extensions._snapshot._get_ranks_from_env())
+            get_ranks_from_env())
 
         if not torch.distributed.is_initialized():
             os.environ["WORLD_SIZE"] = str(size)
             os.environ["RANK"] = str(rank)
+            os.environ["LOCAL_RANK"] = str(local_rank)
 
             torch.distributed.init_process_group(backend='nccl',
                                                  init_method='env://')