Add RunIf

Author: carmocca

tests/helpers/runif.py

@@ -0,0 +1,184 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# 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
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import sys
+from distutils.version import LooseVersion
+from typing import Optional
+
+import pytest
+import torch
+from pkg_resources import get_distribution
+
+from pytorch_lightning.utilities import (
+    _APEX_AVAILABLE,
+    _DEEPSPEED_AVAILABLE,
+    _FAIRSCALE_AVAILABLE,
+    _FAIRSCALE_PIPE_AVAILABLE,
+    _HOROVOD_AVAILABLE,
+    _NATIVE_AMP_AVAILABLE,
+    _RPC_AVAILABLE,
+    _TORCH_QUANTIZE_AVAILABLE,
+    _TPU_AVAILABLE,
+)
+
+try:
+    from horovod.common.util import nccl_built
+    nccl_built()
+except (ImportError, ModuleNotFoundError, AttributeError):
+    _HOROVOD_NCCL_AVAILABLE = False
+finally:
+    _HOROVOD_NCCL_AVAILABLE = True
+
+
+class RunIf:
+    """
+    RunIf wrapper for simple marking specific cases, fully compatible with pytest.mark::
+
+        @RunIf(min_torch="0.0")
+        @pytest.mark.parametrize("arg1", [1, 2.0])
+        def test_wrapper(arg1):
+            assert arg1 > 0.0
+    """
+
+    def __new__(
+        self,
+        *args,
+        min_gpus: int = 0,
+        min_torch: Optional[str] = None,
+        max_torch: Optional[str] = None,
+        min_python: Optional[str] = None,
+        quantization: bool = False,
+        amp_apex: bool = False,
+        amp_native: bool = False,
+        tpu: bool = False,
+        horovod: bool = False,
+        horovod_nccl: bool = False,
+        skip_windows: bool = False,
+        special: bool = False,
+        rpc: bool = False,
+        fairscale: bool = False,
+        fairscale_pipe: bool = False,
+        deepspeed: bool = False,
+        **kwargs
+    ):
+        """
+        Args:
+            args: native pytest.mark.skipif arguments
+            min_gpus: min number of gpus required to run test
+            min_torch: minimum pytorch version to run test
+            max_torch: maximum pytorch version to run test
+            min_python: minimum python version required to run test
+            quantization: if `torch.quantization` package is required to run test
+            amp_apex: NVIDIA Apex is installed
+            amp_native: if native PyTorch native AMP is supported
+            tpu: if TPU is available
+            horovod: if Horovod is installed
+            horovod_nccl: if Horovod is installed with NCCL support
+            skip_windows: skip test for Windows platform (typically fo some limited torch functionality)
+            special: running in special mode, outside pytest suit
+            rpc: requires Remote Procedure Call (RPC)
+            fairscale: if `fairscale` module is required to run the test
+            deepspeed: if `deepspeed` module is required to run the test
+            kwargs: native pytest.mark.skipif keyword arguments
+        """
+        conditions = []
+        reasons = []
+
+        if min_gpus:
+            conditions.append(torch.cuda.device_count() < min_gpus)
+            reasons.append(f"GPUs>={min_gpus}")
+
+        if min_torch:
+            torch_version = LooseVersion(get_distribution("torch").version)
+            conditions.append(torch_version < LooseVersion(min_torch))
+            reasons.append(f"torch>={min_torch}")
+
+        if max_torch:
+            torch_version = LooseVersion(get_distribution("torch").version)
+            conditions.append(torch_version >= LooseVersion(max_torch))
+            reasons.append(f"torch<{max_torch}")
+
+        if min_python:
+            py_version = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
+            conditions.append(py_version < LooseVersion(min_python))
+            reasons.append(f"python>={min_python}")
+
+        if quantization:
+            _miss_default = 'fbgemm' not in torch.backends.quantized.supported_engines
+            conditions.append(not _TORCH_QUANTIZE_AVAILABLE or _miss_default)
+            reasons.append("PyTorch quantization")
+
+        if amp_native:
+            conditions.append(not _NATIVE_AMP_AVAILABLE)
+            reasons.append("native AMP")
+
+        if amp_apex:
+            conditions.append(not _APEX_AVAILABLE)
+            reasons.append("NVIDIA Apex")
+
+        if skip_windows:
+            conditions.append(sys.platform == "win32")
+            reasons.append("unimplemented on Windows")
+
+        if tpu:
+            conditions.append(not _TPU_AVAILABLE)
+            reasons.append("TPU")
+
+        if horovod:
+            conditions.append(not _HOROVOD_AVAILABLE)
+            reasons.append("Horovod")
+
+        if horovod_nccl:
+            conditions.append(not _HOROVOD_NCCL_AVAILABLE)
+            reasons.append("Horovod with NCCL")
+
+        if special:
+            env_flag = os.getenv("PL_RUNNING_SPECIAL_TESTS", '0')
+            conditions.append(env_flag != '1')
+            reasons.append("Special execution")
+
+        if rpc:
+            conditions.append(not _RPC_AVAILABLE)
+            reasons.append("RPC")
+
+        if fairscale:
+            conditions.append(not _FAIRSCALE_AVAILABLE)
+            reasons.append("Fairscale")
+
+        if fairscale_pipe:
+            conditions.append(not _FAIRSCALE_PIPE_AVAILABLE)
+            reasons.append("Fairscale Pipe")
+
+        if deepspeed:
+            conditions.append(not _DEEPSPEED_AVAILABLE)
+            reasons.append("Deepspeed")
+
+        reasons = [rs for cond, rs in zip(conditions, reasons) if cond]
+        return pytest.mark.skipif(
+            *args,
+            condition=any(conditions),
+            reason=f"Requires: [{' + '.join(reasons)}]",
+            **kwargs,
+        )
+
+
+@RunIf(min_torch="99")
+def test_always_skip():
+    exit(1)
+
+
+@pytest.mark.parametrize("arg1", [0.5, 1.0, 2.0])
+@RunIf(min_torch="0.0")
+def test_wrapper(arg1: float):
+    assert arg1 > 0.0