[fix] mypy and flaky test (#624)

* [fix] mypy and flaky test

- CI didn't seem to catch this or maybe I merged incorrectly yesterday
- this should fix the mypy error on master
- also updated a test that seems to be flaky due to tcp port conflict

* another flaky test, hopefully more determinism helps

* CR

* skip 1.6

* fix

* minor
Co-authored-by: Min Xu <min.xu@acm.org>

fairscale/nn/data_parallel/fully_sharded_data_parallel.py

@@ -1662,9 +1662,13 @@ def auto_wrap_bn(module: nn.Module, single_rank_pg: bool = False, process_group:
     def wrap_bn_only_policy(module: nn.Module, recurse: bool, unwrapped_params: int) -> bool:
         is_bn = isinstance(module, torch.nn.modules.batchnorm._BatchNorm)
         if recurse:
-            return not isinstance(module, tuple(default_auto_wrap_policy.FORCE_LEAF_MODULES))  # type: ignore
+            return not isinstance(
+                module, tuple(default_auto_wrap_policy.FORCE_LEAF_MODULES)  # type: ignore
+            )
         else:
-            return is_bn and not isinstance(module, tuple(default_auto_wrap_policy.EXCLUDE_WRAP_MODULES))  # type: ignore
+            return is_bn and not isinstance(
+                module, tuple(default_auto_wrap_policy.EXCLUDE_WRAP_MODULES)  # type: ignore
+            )
 
     pg = None
     if single_rank_pg:

tests/nn/data_parallel/test_fsdp_input.py

@@ -9,8 +9,7 @@
 
 """ Test FSDP with different input types. """
 
-import os
-import random
+import tempfile
 
 import pytest
 import torch
@@ -19,62 +18,72 @@ from torch.optim import SGD
 
 from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
 from fairscale.nn.data_parallel import TrainingState
-from fairscale.utils.testing import skip_if_no_cuda, torch_version
+from fairscale.utils.testing import dist_init, rmf, skip_if_no_cuda, teardown, torch_version
 
 
-# We only test on GPU since mix-precision only really works on GPU.
+# A fixture to get tempfiles and ensure they are cleaned up.
+@pytest.fixture()
+def temp_files():
+    num = 2  # dist_init needs 2 files
+    files = [tempfile.mkstemp()[1] for _ in range(num)]
+
+    yield tuple(files)
+
+    # temp files could have been removed, so we use rmf.
+    for name in files:
+        rmf(name)
+
+
+# We only test on GPU since mix-precision only works on GPU.
 @skip_if_no_cuda
 @pytest.mark.parametrize(
     "fsdp_config", [{}, {"mixed_precision": True}],
 )
 @pytest.mark.parametrize("input_cls", [dict, list])
-def test_it(fsdp_config, input_cls):
+def test_input_type(temp_files, fsdp_config, input_cls):
     """Test FSDP with input being a list or a dict, only single GPU."""
-    if torch_version() < (1, 6, 0):
-        pytest.skip("older pytorch doesn't support reduce_scatter")
-
-    # Random port in case the next test run quickly, same port would cause conflict.
-    os.environ["MASTER_ADDR"] = "localhost"
-    os.environ["MASTER_PORT"] = str(random.randint(2000, 3000))
-    torch.distributed.init_process_group(backend="nccl", rank=0, world_size=1)
-
-    try:
-        assert isinstance(fsdp_config, dict), str(fsdp_config)
-
-        class Model(Module):
-            def __init__(self):
-                super().__init__()
-                self.layer = Linear(4, 4)
-
-            def forward(self, input):
-                if isinstance(input, list):
-                    input = input[0]
-                else:
-                    assert isinstance(input, dict), input
-                    input = input["in"]
-                return self.layer(input)
-
-        model = FSDP(Model(), **fsdp_config).cuda()
-        optim = SGD(model.parameters(), lr=0.1)
-
-        for _ in range(5):
-            in_data = torch.rand(64, 4).cuda()
-            in_data.requires_grad = True
-            if input_cls is list:
-                in_data = [in_data]
-            else:
-                assert input_cls is dict
-                in_data = {"in": in_data}
 
-            out = model(in_data)
-            out.sum().backward()
-            optim.step()
-            optim.zero_grad()
+    if torch_version() < (1, 7, 0):
+        # This test runs multiple test cases in a single process. On 1.6.0 it
+        # throw an error like this:
+        #     RuntimeError: Container is already initialized! Cannot initialize it twice!
+        pytest.skip("older pytorch doesn't work well with single process dist_init multiple times")
+
+    result = dist_init(rank=0, world_size=1, filename=temp_files[0], filename_rpc=temp_files[1])
+    assert result, "Dist init failed"
 
-        model.assert_state(TrainingState.IDLE)
+    assert isinstance(fsdp_config, dict), str(fsdp_config)
 
-    finally:
-        # Clean-up is important or the next test in this file may fail to init the PG.
-        torch.distributed.destroy_process_group()
-        del os.environ["MASTER_ADDR"]
-        del os.environ["MASTER_PORT"]
+    class Model(Module):
+        def __init__(self):
+            super().__init__()
+            self.layer = Linear(4, 4)
+
+        def forward(self, input):
+            if isinstance(input, list):
+                input = input[0]
+            else:
+                assert isinstance(input, dict), input
+                input = input["in"]
+            return self.layer(input)
+
+    model = FSDP(Model(), **fsdp_config).cuda()
+    optim = SGD(model.parameters(), lr=0.1)
+
+    for _ in range(5):
+        in_data = torch.rand(64, 4).cuda()
+        in_data.requires_grad = True
+        if input_cls is list:
+            in_data = [in_data]
+        else:
+            assert input_cls is dict
+            in_data = {"in": in_data}
+
+        out = model(in_data)
+        out.sum().backward()
+        optim.step()
+        optim.zero_grad()
+
+    model.assert_state(TrainingState.IDLE)
+
+    teardown()

tests/nn/data_parallel/test_fsdp_regnet.py

@@ -169,9 +169,9 @@ def ddp_ref():
         unused = tempfile.mkstemp()[1]
         rank_0_output = tempfile.mkstemp()[1]
         try:
-            fsdp_config = None  # This means we use DDP in _test_func.
+            fsdp_config = None  # This means we use DDP in _distributed_worker.
             mp.spawn(
-                _test_func,
+                _distributed_worker,
                 args=(
                     world_size,
                     fsdp_config,
@@ -214,7 +214,7 @@ def temp_files():
     rmf(unused)
 
 
-def _test_func(
+def _distributed_worker(
     rank,
     world_size,
     fsdp_config,
@@ -227,6 +227,8 @@ def _test_func(
     rank_0_output,
     state_after,
 ):
+    torch.backends.cudnn.deterministic = True
+
     result = dist_init(rank, world_size, tempfile_name, unused)
     assert result, "Dist init failed"
 
@@ -328,7 +330,7 @@ def _test_func(
 @skip_if_single_gpu
 @pytest.mark.parametrize("precision", ["full", "mixed"])
 @pytest.mark.parametrize("flatten", ["flatten", "no_flatten"])
-def test1(temp_files, ddp_ref, precision, flatten):
+def test_regnet(temp_files, ddp_ref, precision, flatten):
     if torch_version() < (1, 6, 0):
         pytest.skip("older pytorch doesn't support reduce_scatter")
 
@@ -356,7 +358,7 @@ def test1(temp_files, ddp_ref, precision, flatten):
 
     world_size = _world_size
     mp.spawn(
-        _test_func,
+        _distributed_worker,
         args=(
             world_size,
             fsdp_config,