Repository: pfnet/pytorch-pfn-extras

Issue:
Adding pytest in an MPI environment
There are some functions that run in a distributed environment that have not been fully verified, so some pytests will be run in MPI to verify them.

https://github.com/pfnet/pytorch-pfn-extras/blob/266e8bde2c2a1a6aa3f8648d49e758975c8b436a/tests/pytorch_pfn_extras_tests/training_tests/extensions_tests/test_distributed_snapshot.py#L40-L46



Discussion:

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -16,27 +16,6 @@
     license="MIT License",
     install_requires=["numpy", "packaging", "torch", "typing-extensions>=3.10"],
     extras_require={
-        "test": [
-            "pytest",
-            "onnxruntime",
-            "torchvision",
-            "torchaudio",
-            "pysen",
-            "black==23.3.0",
-            "flake8==4.0.1",
-            "isort==5.10.1",
-            "mypy==1.3.0",
-            "types-PyYAML",
-            "types-setuptools",
-            "matplotlib",
-            "tensorboard",
-            "ipython",
-            "ipywidgets",
-            "pandas",
-            "optuna",
-            "onnx",
-            "pytorch-ignite",
-        ],
         "onnx": ["onnx"],
     },
     python_requires=">=3.6.0",


Test patch:
diff --git a/.flexci/windows/test.ps1 b/.flexci/windows/test.ps1
--- a/.flexci/windows/test.ps1
+++ b/.flexci/windows/test.ps1
@@ -57,7 +57,7 @@ if ($test -eq "torch110") {
 RunOrDie python -V
 
 # Install common requirements
-RunOrDie python -m pip install pytorch-ignite pytest flake8 matplotlib tensorboard onnx ipython ipywidgets pandas optuna cupy-cuda102 onnxruntime==1.15.1 slack_sdk
+RunOrDie python -m pip install -r tests/requirements.txt cupy-cuda102
 RunOrDie python -m pip list
 
 # Install
@@ -65,7 +65,7 @@ RunOrDie python -m pip install -e .
 
 # Unit Test
 $Env:JUPYTER_PLATFORM_DIRS = "1"
-RunOrDie python -m pytest tests
+RunOrDie python -m pytest -m "not mpi" tests
 
 # Examples
 .\.flexci\windows\download_mnist.ps1
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
@@ -2,8 +2,10 @@
 import os
 import tempfile
 
+import py
 import pytest
 import torch
+import torch.distributed
 from pytorch_pfn_extras import distributed, training
 from pytorch_pfn_extras.training import extensions
 
@@ -56,6 +58,22 @@ def path():
         yield t_path
 
 
+@pytest.fixture(scope="function")
+def mpi_tmp_path(tmpdir):
+    try:
+        from mpi4py import MPI
+    except ImportError:
+        pytest.fail("mpi4py needs to be installed to run this test")
+
+    comm = MPI.COMM_WORLD
+    rank = comm.Get_rank()
+
+    name = str(tmpdir) if rank == 0 else None
+    name = comm.bcast(name, root=0)
+    yield py.path.local(name)
+    comm.barrier()
+
+
 def _model_params_equal(model1, model2):
     for p1, p2 in zip(model1.parameters(), model2.parameters()):
         if p1.data.ne(p2.data).sum() > 0.0:
@@ -64,17 +82,18 @@ def _model_params_equal(model1, model2):
 
 
 @pytest.mark.gpu
-def test_distributed_snapshot(path):
+@pytest.mark.mpi
+@pytest.mark.parametrize("saver_rank", [0, 1])
+def test_distributed_snapshot(mpi_tmp_path, saver_rank):
     comm_size, comm_rank, comm_local_rank, device = _init_distributed(False)
 
     if comm_size > 1:
         torch.distributed.barrier()
 
-    saver_rank = 0
     fmt = "snapshot_iter_{.iteration}"
     snapshot = extensions.snapshot(filename=fmt, saver_rank=saver_rank)
 
-    trainer = get_trainer(path)
+    trainer = get_trainer(mpi_tmp_path)
     trainer.extend(snapshot, trigger=(1, "iteration"), priority=2)
     for _ in range(1):
         with trainer.run_iteration():
@@ -83,34 +102,84 @@ def test_distributed_snapshot(path):
     pattern = os.path.join(trainer.out, "snapshot_iter_*")
     found = [os.path.basename(path) for path in glob.glob(pattern)]
     # the snapshot is generated only for the saver rank
-    assert comm_rank == saver_rank and len(found) == 1 or len(found) == 0
+    assert len(found) == 1
 
-    if comm_rank == saver_rank:
-        new_trainer = get_trainer(path)
-        new_trainer.load_state_dict(torch.load(os.path.join(path, found[0])))
-        assert _model_params_equal(
-            trainer._models["main"], new_trainer._models["main"]
-        )
-
-    if comm_size > 1:
-        torch.distributed.barrier()
+    new_trainer = get_trainer(mpi_tmp_path)
+    assert not _model_params_equal(
+        trainer._models["main"], new_trainer._models["main"]
+    )
+    new_trainer.load_state_dict(
+        torch.load(os.path.join(mpi_tmp_path, found[0]))
+    )
+    assert _model_params_equal(
+        trainer._models["main"], new_trainer._models["main"]
+    )
 
 
 @pytest.mark.gpu
-def test_subprocess_distributed_snapshot(path):
+@pytest.mark.mpi
+@pytest.mark.parametrize("saver_rank", [0, 1])
+def test_distributed_snapshot_autoload(mpi_tmp_path, saver_rank):
     comm_size, comm_rank, comm_local_rank, device = _init_distributed(False)
 
     if comm_size > 1:
         torch.distributed.barrier()
-
-    saver_rank = 1
     fmt = "snapshot_iter_{.iteration}"
-    snapshot = extensions.snapshot(filename=fmt, saver_rank=saver_rank)
-
-    trainer = get_trainer(path)
+    snapshot = extensions.snapshot(
+        filename=fmt, saver_rank=saver_rank, autoload=True
+    )
+    trainer = get_trainer(mpi_tmp_path)
     trainer.extend(snapshot, trigger=(1, "iteration"), priority=2)
     for _ in range(1):
         with trainer.run_iteration():
             pass
-    assert snapshot.writer is not None
-    assert len(snapshot.writer._post_save_hooks) == 0
+    assert 1 == trainer.iteration
+    pattern = os.path.join(trainer.out, "snapshot_iter_*")
+    found = [os.path.basename(path) for path in glob.glob(pattern)]
+    assert len(found) == 1
+    new_trainer = get_trainer(mpi_tmp_path)
+    assert not _model_params_equal(
+        trainer._models["main"], new_trainer._models["main"]
+    )
+    snapshot.initialize(new_trainer)
+    assert _model_params_equal(
+        trainer._models["main"], new_trainer._models["main"]
+    )
+
+
+@pytest.mark.gpu
+@pytest.mark.mpi
+@pytest.mark.parametrize("saver_rank", [0, 1])
+def test_distributed_snapshot_on_error(mpi_tmp_path, saver_rank):
+    comm_size, comm_rank, comm_local_rank, device = _init_distributed(False)
+    if comm_size > 1:
+        torch.distributed.barrier()
+    trainer = get_trainer(mpi_tmp_path)
+
+    fmt = f"snapshot_iter_{comm_rank}_" "{.iteration}"
+    snapshot = extensions.snapshot(
+        filename=fmt,
+        snapshot_on_error=True,
+        saver_rank=saver_rank,
+    )
+    try:
+        raise RuntimeError
+    except RuntimeError as dummy_exception:
+        pass
+        dummy_tb = dummy_exception.__traceback__
+        snapshot.on_error(trainer, dummy_exception, dummy_tb)
+    pattern = os.path.join(trainer.out, f"snapshot_iter_{saver_rank}_*")
+    found = [os.path.basename(path) for path in glob.glob(pattern)]
+    # the snapshot is generated only for the saver rank
+    assert len(found) == 1
+
+    new_trainer = get_trainer(mpi_tmp_path)
+    assert not _model_params_equal(
+        trainer._models["main"], new_trainer._models["main"]
+    )
+    new_trainer.load_state_dict(
+        torch.load(os.path.join(mpi_tmp_path, found[0]))
+    )
+    assert _model_params_equal(
+        trainer._models["main"], new_trainer._models["main"]
+    )
diff --git a/tests/requirements.mpi.txt b/tests/requirements.mpi.txt
new file mode 100644
--- /dev/null
+++ b/tests/requirements.mpi.txt
@@ -0,0 +1 @@
+mpi4py
diff --git a/tests/requirements.txt b/tests/requirements.txt
new file mode 100644
--- /dev/null
+++ b/tests/requirements.txt
@@ -0,0 +1,21 @@
+pytest
+onnxruntime
+torchvision
+torchaudio
+pysen
+black==23.3.0
+flake8==4.0.1
+isort==5.10.1
+mypy==1.3.0
+types-PyYAML
+types-setuptools
+matplotlib
+tensorboard
+ipython
+ipywidgets
+pandas
+optuna
+onnx
+pytorch-ignite
+pytest-cov
+slack_sdk