Repository: facebookresearch/fairscale

Issue:
[refactor] unit tests use relative imports, which break with some CI / toolchains
## 🚀 Feature
Refactor unit tests to move testing utils in a common place within the fairscale library, instead of being ad-hoc relative imports

## Motivation
Cleaner (there are probably utils which could be shared in between different unit test families), and makes it easily compatible with some toolchains, such as Facebook's (see D25223316)

## Pitch
Remove all relative imports in unit tests, consolidate utils in a testing utils part of the fairscale lib

## Alternatives
Keep duplications
Not Facebook-CI compatible

## Additional context
D25223316 (FB internal)



Discussion:
cc @min-xu-ai @myleott @msbaines 

Resolving patch:
diff --git a/benchmarks/pipe.py b/benchmarks/pipe.py
--- a/benchmarks/pipe.py
+++ b/benchmarks/pipe.py
@@ -23,7 +23,7 @@
 from fairscale.nn.pipe import LazyModule, pipe
 from fairscale.optim import GradScaler
 from fairscale.optim.oss import OSS
-from tests.nn.model_parallel.commons import dist_init, get_worker_map
+from fairscale.utils.testing import dist_init, get_worker_map
 
 try:
     from fairscale.optim import Adam  # type: ignore


Test patch:
diff --git a/tests/nn/model_parallel/commons.py b/fairscale/utils/testing.py
similarity index 74%
rename from tests/nn/model_parallel/commons.py
rename to fairscale/utils/testing.py
--- a/tests/nn/model_parallel/commons.py
+++ b/fairscale/utils/testing.py
@@ -1,5 +1,3 @@
-# coding=utf-8
-
 # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
 #
 # This source code is licensed under the BSD license found in the
@@ -11,7 +9,7 @@
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
-#     http://www.apache.org/licenses/LICENSE-2.0
+#   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
@@ -19,14 +17,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# We're not responsible for pytest decorators
+# mypy: disallow_untyped_decorators = False
+
+"""
+Collection of some testing utilities for the Fairscale library. Please complement as you see fit, but refrain from ad-hoc test utils
+within the different feature sets and relative imports.
+"""
+
 import functools
 import inspect
+import logging
 import multiprocessing
 import os
 import random
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import numpy
-from packaging import version
 import pytest
 import torch
 import torch.distributed as dist
@@ -38,11 +45,11 @@
 
 
 class IdentityLayer(torch.nn.Module):
-    def __init__(self, size, scale=1.0):
+    def __init__(self, size: int, scale: float = 1.0) -> None:
         super(IdentityLayer, self).__init__()
         self.weight = torch.nn.Parameter(scale * torch.randn(size))
 
-    def forward(self):
+    def forward(self, *_: Any, **__: Any) -> Any:
         return self.weight
 
 
@@ -54,7 +61,26 @@ def set_random_seed(seed: int) -> None:
     model_parallel_cuda_manual_seed(seed)
 
 
-def dist_init(rank, world_size, hostname=None):
+def torch_version() -> Tuple[int, ...]:
+    numbering = torch.__version__.split(".")
+
+    assert len(numbering) == 3
+
+    # Catch torch version if run against internal pre-releases, like `1.8.0a0fb`,
+    if not numbering[2].isnumeric():
+        # Two options here:
+        # - either skip this version (minor number check is not relevant)
+        # - or check that our codebase is not broken by this ongoing development.
+
+        # Assuming that we're interested in the second usecase more than the first,
+        # return the pre-release or dev numbering
+        logging.warning(f"Pytorch pre-relase version {torch.__version__} - assuming intent to test it")
+        numbering[2] = "0"
+
+    return tuple(int(n) for n in numbering)
+
+
+def dist_init(rank: int, world_size: int, hostname: Optional[str] = None) -> None:
     if hostname is None:
         hostname = "localhost"
     print(f"dist init r={rank}, world={world_size}, host={hostname}")
@@ -63,7 +89,7 @@ def dist_init(rank, world_size, hostname=None):
     os.environ["WORLD_SIZE"] = str(world_size)
     os.environ["RANK"] = str(rank)
 
-    if version.parse(torch.__version__).release >= (1, 6, 0):
+    if torch_version() >= (1, 6, 0):
         init_method = f"tcp://{os.environ['MASTER_ADDR']}:{os.environ['MASTER_PORT']}"
         backend = "nccl" if torch.cuda.is_available() else "gloo"
         torch.distributed.init_process_group(backend=backend, rank=rank, world_size=world_size, init_method=init_method)
@@ -77,6 +103,7 @@ def dist_init(rank, world_size, hostname=None):
             backend=rpc.BackendType.TENSORPIPE,
             rpc_backend_options=rpc.TensorPipeRpcBackendOptions(init_method=init_method),
         )
+
     else:
         if world_size > 1:
             rpc.init_rpc(f"Test{rank}", rank=rank, world_size=world_size)
@@ -87,21 +114,21 @@ def dist_init(rank, world_size, hostname=None):
         torch.cuda.set_device(rank % torch.cuda.device_count())
 
 
-def get_worker_map():
+def get_worker_map() -> Dict[Any, Any]:
     return {rank: f"Test{rank}" for rank in range(dist.get_world_size())}
 
 
-def get_world_sizes():
+def get_world_sizes() -> List[int]:
     limit = torch.cuda.device_count()
     return [x for x in [1, 2, 4, 8] if x <= limit]
 
 
-def spawn_for_all_world_sizes(test_func, world_sizes=get_world_sizes(), args=[]):
+def spawn_for_all_world_sizes(test_func: Callable, world_sizes: List[int] = get_world_sizes(), args: Any = []) -> None:
     for world_size in world_sizes:
-        mp.spawn(test_func, args=(world_size, *args), nprocs=world_size, join=True)
+        mp.spawn(test_func, args=(world_size, *args), nprocs=world_size, join=True)  # type: ignore
 
 
-def worker_process(rank, world_size, func, args, error_queue):
+def worker_process(rank: int, world_size: int, func: Callable, args: Any, error_queue: Any) -> None:
     """Main function for unit tests launced with torch_spawn"""
 
     dist_init(rank, world_size)
@@ -120,11 +147,11 @@ def worker_process(rank, world_size, func, args, error_queue):
         raise e
 
 
-def torch_spawn(world_sizes=None):
+def torch_spawn(world_sizes: Optional[List[int]] = None) -> Callable:
     if world_sizes is None:
         world_sizes = get_world_sizes()
 
-    def prepare_test(func):
+    def prepare_test(func: Callable) -> Callable:
         """Function called with the test function as the argument. Generates a
         replacement which serves as the actual test function."""
 
@@ -138,8 +165,10 @@ def prepare_test(func):
             )
 
         @functools.wraps(func)
-        def replacement(*args, **kwargs):
+        def replacement(*args: Any, **kwargs: Any) -> None:
             assert args == tuple()
+            assert world_sizes is not None  # mypy crutch
+
             args = tuple(
                 kwargs[p] for p in parameters if p != "rank"
             )  # converting named parameters to positional parameters to pass to `spawn`
@@ -174,7 +203,9 @@ def replacement(*args, **kwargs):
 
         # Register a function with the same name, prefixed with "test_" in the
         # calling module, so it will be picked up by pytest
-        caller_module = inspect.getmodule(inspect.currentframe().f_back)
+        current_frame = inspect.currentframe()
+        assert current_frame is not None
+        caller_module = inspect.getmodule(current_frame.f_back)
         setattr(caller_module, f"test_{name}", replacement)
 
         return func
diff --git a/tests/nn/model_parallel/test_cross_entropy.py b/tests/nn/model_parallel/test_cross_entropy.py
--- a/tests/nn/model_parallel/test_cross_entropy.py
+++ b/tests/nn/model_parallel/test_cross_entropy.py
@@ -25,7 +25,7 @@
 from fairscale.nn.model_parallel import initialize as mpu
 from fairscale.nn.model_parallel.cross_entropy import vocab_parallel_cross_entropy
 from fairscale.nn.model_parallel.mappings import scatter_to_model_parallel_region
-from tests.nn.model_parallel.commons import IdentityLayer, dist_init, set_random_seed, spawn_for_all_world_sizes
+from fairscale.utils.testing import IdentityLayer, dist_init, set_random_seed, spawn_for_all_world_sizes
 
 
 def torch_cross_entropy(batch_size, seq_length, vocab_size, logits_scale, seed):
diff --git a/tests/nn/model_parallel/test_initialize.py b/tests/nn/model_parallel/test_initialize.py
--- a/tests/nn/model_parallel/test_initialize.py
+++ b/tests/nn/model_parallel/test_initialize.py
@@ -22,7 +22,7 @@
 import torch
 
 from fairscale.nn.model_parallel import initialize as mpu
-from tests.nn.model_parallel.commons import dist_init, spawn_for_all_world_sizes
+from fairscale.utils.testing import dist_init, spawn_for_all_world_sizes
 
 
 def run_test_initialize_model_parallel(rank, model_parallel_size):
diff --git a/tests/nn/model_parallel/test_layers.py b/tests/nn/model_parallel/test_layers.py
--- a/tests/nn/model_parallel/test_layers.py
+++ b/tests/nn/model_parallel/test_layers.py
@@ -31,13 +31,7 @@
 from fairscale.nn.model_parallel import initialize as mpu
 from fairscale.nn.model_parallel import layers
 from fairscale.nn.pipe import Pipe
-from tests.nn.model_parallel.commons import (
-    dist_init,
-    get_world_sizes,
-    set_random_seed,
-    spawn_for_all_world_sizes,
-    torch_spawn,
-)
+from fairscale.utils.testing import dist_init, get_world_sizes, set_random_seed, spawn_for_all_world_sizes, torch_spawn
 
 
 def run_test_parallel_embedding(rank, model_parallel_size):
diff --git a/tests/nn/model_parallel/test_random.py b/tests/nn/model_parallel/test_random.py
--- a/tests/nn/model_parallel/test_random.py
+++ b/tests/nn/model_parallel/test_random.py
@@ -24,7 +24,7 @@
 from fairscale.nn.model_parallel import initialize as mpu
 from fairscale.nn.model_parallel import random
 from fairscale.nn.model_parallel.random import get_cuda_rng_tracker, model_parallel_cuda_manual_seed
-from tests.nn.model_parallel.commons import dist_init, spawn_for_all_world_sizes
+from fairscale.utils.testing import dist_init, spawn_for_all_world_sizes
 
 
 def run_test_set_cuda_rng_state(rank, model_parallel_size):
diff --git a/tests/nn/pipe/conftest.py b/tests/nn/pipe/conftest.py
--- a/tests/nn/pipe/conftest.py
+++ b/tests/nn/pipe/conftest.py
@@ -17,17 +17,27 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import functools
+import os
+from typing import Any, Callable
+
 import pytest
 import torch
 
+from fairscale.nn.model_parallel import destroy_model_parallel
+
 
 @pytest.fixture(autouse=True)
-def manual_seed_zero():
+def manual_seed_zero() -> None:
     torch.manual_seed(0)
 
 
+def cuda_sleep_impl(seconds, cycles_per_ms):
+    torch.cuda._sleep(int(seconds * cycles_per_ms * 1000))
+
+
 @pytest.fixture(scope="session")
-def cuda_sleep():
+def cuda_sleep() -> Callable:
     # Warm-up CUDA.
     torch.empty(1, device="cuda")
 
@@ -40,11 +50,23 @@ def cuda_sleep():
     end.synchronize()
     cycles_per_ms = 1000000 / start.elapsed_time(end)
 
-    def cuda_sleep(seconds):
-        torch.cuda._sleep(int(seconds * cycles_per_ms * 1000))
-
-    return cuda_sleep
+    return functools.partial(cuda_sleep_impl, cycles_per_ms=cycles_per_ms)
 
 
-def pytest_report_header():
+def pytest_report_header() -> str:
     return f"torch: {torch.__version__}"
+
+
+def pytest_runtest_setup(item: Any) -> None:
+    print(f"setup mpi function called")
+
+
+def pytest_runtest_teardown(item: Any) -> None:
+    if "OMPI_COMM_WORLD_RANK" in os.environ:
+        destroy_model_parallel()
+        if torch.distributed.is_initialized():
+            torch.distributed.destroy_process_group()
+        try:
+            torch.distributed.rpc.shutdown()
+        except Exception:
+            pass
diff --git a/tests/nn/pipe_process/conftest.py b/tests/nn/pipe_process/conftest.py
--- a/tests/nn/pipe_process/conftest.py
+++ b/tests/nn/pipe_process/conftest.py
@@ -19,6 +19,7 @@
 
 import functools
 import os
+from typing import Any, Callable
 
 import pytest
 import torch
@@ -27,7 +28,7 @@
 
 
 @pytest.fixture(autouse=True)
-def manual_seed_zero():
+def manual_seed_zero() -> None:
     torch.manual_seed(0)
 
 
@@ -36,7 +37,7 @@ def cuda_sleep_impl(seconds, cycles_per_ms):
 
 
 @pytest.fixture(scope="session")
-def cuda_sleep():
+def cuda_sleep() -> Callable:
     # Warm-up CUDA.
     torch.empty(1, device="cuda")
 
@@ -52,15 +53,15 @@ def cuda_sleep():
     return functools.partial(cuda_sleep_impl, cycles_per_ms=cycles_per_ms)
 
 
-def pytest_report_header():
+def pytest_report_header() -> str:
     return f"torch: {torch.__version__}"
 
 
-def pytest_runtest_setup(item):
+def pytest_runtest_setup(item: Any) -> None:
     print(f"setup mpi function called")
 
 
-def pytest_runtest_teardown(item):
+def pytest_runtest_teardown(item: Any) -> None:
     if "OMPI_COMM_WORLD_RANK" in os.environ:
         destroy_model_parallel()
         if torch.distributed.is_initialized():
diff --git a/tests/nn/pipe_process/skip/test_gpipe.py b/tests/nn/pipe_process/skip/test_gpipe.py
--- a/tests/nn/pipe_process/skip/test_gpipe.py
+++ b/tests/nn/pipe_process/skip/test_gpipe.py
@@ -26,7 +26,7 @@
 from fairscale.nn.pipe import LazyModule, Pipe
 from fairscale.nn.pipe.skip import pop, skippable, stash
 from fairscale.nn.pipe.skip.portal import PortalBlue, PortalCopy, PortalOrange
-from tests.nn.model_parallel.commons import get_worker_map, torch_spawn
+from fairscale.utils.testing import get_worker_map, torch_spawn
 
 
 @torch_spawn([3])
diff --git a/tests/nn/pipe_process/skip/test_leak.py b/tests/nn/pipe_process/skip/test_leak.py
--- a/tests/nn/pipe_process/skip/test_leak.py
+++ b/tests/nn/pipe_process/skip/test_leak.py
@@ -26,7 +26,7 @@
 from fairscale.nn.pipe import Pipe, is_checkpointing, is_recomputing
 from fairscale.nn.pipe.skip import pop, skippable, stash
 from fairscale.nn.pipe.skip.tracker import current_skip_tracker
-from tests.nn.model_parallel.commons import get_worker_map, torch_spawn
+from fairscale.utils.testing import get_worker_map, torch_spawn
 
 
 @skippable(stash=["skip"])
diff --git a/tests/nn/pipe_process/test_bugs.py b/tests/nn/pipe_process/test_bugs.py
--- a/tests/nn/pipe_process/test_bugs.py
+++ b/tests/nn/pipe_process/test_bugs.py
@@ -23,7 +23,7 @@
 import torch.nn.functional as F
 
 from fairscale.nn.pipe import Pipe
-from tests.nn.model_parallel.commons import get_worker_map, torch_spawn
+from fairscale.utils.testing import get_worker_map, torch_spawn
 
 
 @torch_spawn([2])
diff --git a/tests/nn/pipe_process/test_inplace.py b/tests/nn/pipe_process/test_inplace.py
--- a/tests/nn/pipe_process/test_inplace.py
+++ b/tests/nn/pipe_process/test_inplace.py
@@ -22,7 +22,7 @@
 from torch import nn
 
 from fairscale.nn.pipe import Pipe
-from tests.nn.model_parallel.commons import get_worker_map, torch_spawn
+from fairscale.utils.testing import get_worker_map, torch_spawn
 
 
 @torch_spawn([2])
diff --git a/tests/nn/pipe_process/test_pipe.py b/tests/nn/pipe_process/test_pipe.py
--- a/tests/nn/pipe_process/test_pipe.py
+++ b/tests/nn/pipe_process/test_pipe.py
@@ -21,9 +21,7 @@
 from copy import deepcopy
 import os
 import time
-from typing import Tuple
 
-from packaging import version
 import pytest
 import torch
 from torch import nn
@@ -34,7 +32,7 @@
     initialize_model_parallel,
 )
 from fairscale.nn.pipe import LazyModule, Pipe
-from tests.nn.model_parallel.commons import get_worker_map, set_random_seed, torch_spawn
+from fairscale.utils.testing import get_worker_map, set_random_seed, torch_spawn, torch_version
 
 
 @torch_spawn([2])
@@ -373,24 +371,6 @@ def find_grad_fn(grad_fn, name):
     assert not find_grad_fn(eval_output.grad_fn, "RecomputeBackward")
 
 
-def torch_version() -> Tuple[int, ...]:
-    result = version.parse(torch.__version__).release
-
-    # Catch torch version if run against internal pre-releases, like `1.8.0a0fb`,
-    # for which version.parse().release will return None (version becomes of LegacyVersion type)
-    if result is None:
-        # Two options here:
-        # - either skip this version,
-        # - or check that Pipe is not broken by this ongoing development.
-
-        # Assuming that we're interested in the second usecase more than the first,
-        # return the pre-release or dev numbering
-        numbering = torch.__version__.split(".")
-        result = (int(numbering[0]), int(numbering[1]), 0)
-    assert result
-    return result
-
-
 @torch_spawn([2])
 @pytest.mark.xfail(torch_version() < (1, 6, 0), reason="Doesn't work on torch < 1.6.0", strict=True)
 @pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
diff --git a/tests/nn/pipe_process/test_rpc.py b/tests/nn/pipe_process/test_rpc.py
--- a/tests/nn/pipe_process/test_rpc.py
+++ b/tests/nn/pipe_process/test_rpc.py
@@ -8,7 +8,7 @@
 
 from fairscale.nn.model_parallel.initialize import get_pipeline_parallel_group
 from fairscale.nn.pipe import PipeRPCWrapper
-from tests.nn.model_parallel.commons import get_worker_map, torch_spawn
+from fairscale.utils.testing import get_worker_map, torch_spawn
 
 
 def init_rpc():
diff --git a/tests/nn/pipe_process/test_transparency.py b/tests/nn/pipe_process/test_transparency.py
--- a/tests/nn/pipe_process/test_transparency.py
+++ b/tests/nn/pipe_process/test_transparency.py
@@ -22,7 +22,7 @@
 from torch import nn
 
 from fairscale.nn import Pipe
-from tests.nn.model_parallel.commons import get_worker_map, set_random_seed, torch_spawn
+from fairscale.utils.testing import get_worker_map, set_random_seed, torch_spawn
 
 
 @torch_spawn([2])