Refactor: skipif for multi - gpus 1/n (#6266)

* ngpus

* gpu

* isort

* pt

* flake8

Author: Borda

tests/accelerators/test_accelerator_connector.py

@@ -34,6 +34,7 @@
 from pytorch_lightning.plugins.environments import ClusterEnvironment, SLURMEnvironment, TorchElasticEnvironment
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.helpers.boring_model import BoringModel
+from tests.helpers.skipif import RunIf
 
 
 def test_accelerator_choice_cpu(tmpdir):
@@ -119,7 +120,7 @@ def on_fit_start(self, trainer, pl_module):
         trainer.fit(model)
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU")
+@RunIf(min_gpus=1)
 @mock.patch.dict(
     os.environ, {
         "CUDA_VISIBLE_DEVICES": "0,1",
@@ -157,7 +158,7 @@ def on_fit_start(self, trainer, pl_module):
         trainer.fit(model)
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU")
+@RunIf(min_gpus=1)
 @mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1", "WORLD_SIZE": "2", "LOCAL_RANK": "10", "NODE_RANK": "0"})
 @mock.patch('torch.cuda.device_count', return_value=2)
 def test_accelerator_choice_ddp_te(device_count_mock):
@@ -185,7 +186,7 @@ def on_fit_start(self, trainer, pl_module):
         trainer.fit(model)
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU")
+@RunIf(min_gpus=1)
 @mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1", "WORLD_SIZE": "2", "LOCAL_RANK": "10", "NODE_RANK": "0"})
 @mock.patch('torch.cuda.device_count', return_value=2)
 def test_accelerator_choice_ddp2_te(device_count_mock):

tests/accelerators/test_ddp.py

@@ -21,12 +21,13 @@
 from pytorch_lightning import Trainer
 from tests.accelerators import ddp_model, DDPLauncher
 from tests.helpers.boring_model import BoringModel
+from tests.helpers.skipif import RunIf
 from tests.utilities.distributed import call_training_script
 
 CLI_ARGS = '--max_epochs 1 --gpus 2 --accelerator ddp'
 
 
-@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+@RunIf(min_gpus=2)
 def test_multi_gpu_model_ddp_fit_only(tmpdir):
     # call the script
     call_training_script(ddp_model, CLI_ARGS, 'fit', tmpdir, timeout=120)
@@ -39,7 +40,7 @@ def test_multi_gpu_model_ddp_fit_only(tmpdir):
     assert result['status'] == 'complete'
 
 
-@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+@RunIf(min_gpus=2)
 def test_multi_gpu_model_ddp_test_only(tmpdir):
     # call the script
     call_training_script(ddp_model, CLI_ARGS, 'test', tmpdir)
@@ -52,7 +53,7 @@ def test_multi_gpu_model_ddp_test_only(tmpdir):
     assert result['status'] == 'complete'
 
 
-@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+@RunIf(min_gpus=2)
 def test_multi_gpu_model_ddp_fit_test(tmpdir):
     # call the script
     call_training_script(ddp_model, CLI_ARGS, 'fit_test', tmpdir, timeout=20)
@@ -69,7 +70,7 @@ def test_multi_gpu_model_ddp_fit_test(tmpdir):
         assert out['test_acc'] > 0.7
 
 
-@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+@RunIf(min_gpus=2)
 @DDPLauncher.run(
     "--max_epochs [max_epochs] --gpus 2 --accelerator [accelerator]",
     max_epochs=["1"],

tests/accelerators/test_ddp_spawn.py

@@ -11,8 +11,6 @@
 # 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 pytest
-import torch
 
 import tests.helpers.pipelines as tpipes
 import tests.helpers.utils as tutils
@@ -23,9 +21,10 @@
 from tests.helpers import BoringModel
 from tests.helpers.datamodules import ClassifDataModule
 from tests.helpers.simple_models import ClassificationModel
+from tests.helpers.skipif import RunIf
 
 
-@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+@RunIf(min_gpus=2)
 def test_multi_gpu_early_stop_ddp_spawn(tmpdir):
     tutils.set_random_master_port()
 
@@ -44,7 +43,7 @@ def test_multi_gpu_early_stop_ddp_spawn(tmpdir):
     tpipes.run_model_test(trainer_options, model, dm)
 
 
-@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+@RunIf(min_gpus=2)
 def test_multi_gpu_model_ddp_spawn(tmpdir):
     tutils.set_random_master_port()
 
@@ -66,7 +65,7 @@ def test_multi_gpu_model_ddp_spawn(tmpdir):
     memory.get_memory_profile('min_max')
 
 
-@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+@RunIf(min_gpus=2)
 def test_ddp_all_dataloaders_passed_to_fit(tmpdir):
     """Make sure DDP works with dataloaders passed to fit()"""
     tutils.set_random_master_port()

tests/accelerators/test_dp.py

@@ -11,10 +11,6 @@
 # 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
-from unittest import mock
-
-import pytest
 import torch
 import torch.nn.functional as F
 
@@ -26,6 +22,7 @@
 from tests.helpers import BoringModel
 from tests.helpers.datamodules import ClassifDataModule
 from tests.helpers.simple_models import ClassificationModel
+from tests.helpers.skipif import RunIf
 
 PRETEND_N_OF_GPUS = 16
 
@@ -55,7 +52,7 @@ def test_step_end(self, outputs):
         self.log('test_acc', self.test_acc(outputs['logits'], outputs['y']))
 
 
-@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+@RunIf(min_gpus=2)
 def test_multi_gpu_early_stop_dp(tmpdir):
     """Make sure DDP works. with early stopping"""
     tutils.set_random_master_port()
@@ -76,7 +73,7 @@ def test_multi_gpu_early_stop_dp(tmpdir):
     tpipes.run_model_test(trainer_options, model, dm)
 
 
-@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+@RunIf(min_gpus=2)
 def test_multi_gpu_model_dp(tmpdir):
     tutils.set_random_master_port()
 
@@ -98,8 +95,7 @@ def test_multi_gpu_model_dp(tmpdir):
     memory.get_memory_profile('min_max')
 
 
-@mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1"})
-@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+@RunIf(min_gpus=2)
 def test_dp_test(tmpdir):
     tutils.set_random_master_port()
 

tests/callbacks/test_gpu_stats_monitor.py

@@ -24,9 +24,10 @@
 from pytorch_lightning.trainer.states import TrainerState
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.helpers import BoringModel
+from tests.helpers.skipif import RunIf
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
+@RunIf(min_gpus=1)
 def test_gpu_stats_monitor(tmpdir):
     """
     Test GPU stats are logged using a logger.
@@ -76,7 +77,7 @@ def test_gpu_stats_monitor_cpu_machine(tmpdir):
         GPUStatsMonitor()
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
+@RunIf(min_gpus=1)
 def test_gpu_stats_monitor_no_logger(tmpdir):
     """
     Test GPUStatsMonitor with no logger in Trainer.
@@ -96,7 +97,7 @@ def test_gpu_stats_monitor_no_logger(tmpdir):
         trainer.fit(model)
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
+@RunIf(min_gpus=1)
 def test_gpu_stats_monitor_no_gpu_warning(tmpdir):
     """
     Test GPUStatsMonitor raises a warning when not training on GPU device.

tests/callbacks/test_pruning.py

@@ -27,6 +27,7 @@
 from pytorch_lightning.callbacks import ModelPruning
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.helpers import BoringModel
+from tests.helpers.skipif import RunIf
 
 
 class TestModel(BoringModel):
@@ -168,7 +169,7 @@ def test_pruning_callback_ddp(tmpdir, use_global_unstructured, parameters_to_pru
     )
 
 
-@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+@RunIf(min_gpus=2)
 @pytest.mark.skipif(platform.system() == "Windows", reason="Distributed training is not supported on Windows")
 def test_pruning_callback_ddp_spawn(tmpdir):
     train_with_pruning_callback(tmpdir, use_global_unstructured=True, accelerator="ddp_spawn", gpus=2)

tests/callbacks/test_quantization.py

@@ -22,15 +22,12 @@
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.helpers.datamodules import RegressDataModule
 from tests.helpers.simple_models import RegressionModel
-from tests.helpers.skipif import SkipIf
+from tests.helpers.skipif import RunIf
 
 
-@pytest.mark.parametrize(
-    "observe",
-    ['average', pytest.param('histogram', marks=SkipIf(min_torch="1.5"))]
-)
+@pytest.mark.parametrize("observe", ['average', pytest.param('histogram', marks=RunIf(min_torch="1.5"))])
 @pytest.mark.parametrize("fuse", [True, False])
-@SkipIf(quantization=True)
+@RunIf(quantization=True)
 def test_quantization(tmpdir, observe, fuse):
     """Parity test for quant model"""
     seed_everything(42)
@@ -65,7 +62,7 @@ def test_quantization(tmpdir, observe, fuse):
     assert torch.allclose(org_score, quant_score, atol=0.45)
 
 
-@SkipIf(quantization=True)
+@RunIf(quantization=True)
 def test_quantize_torchscript(tmpdir):
     """Test converting to torchscipt """
     dm = RegressDataModule()
@@ -81,7 +78,7 @@ def test_quantize_torchscript(tmpdir):
     tsmodel(tsmodel.quant(batch[0]))
 
 
-@SkipIf(quantization=True)
+@RunIf(quantization=True)
 def test_quantization_exceptions(tmpdir):
     """Test wrong fuse layers"""
     with pytest.raises(MisconfigurationException, match='Unsupported qconfig'):
@@ -124,7 +121,7 @@ def custom_trigger_last(trainer):
         (custom_trigger_last, 2),
     ]
 )
-@SkipIf(quantization=True)
+@RunIf(quantization=True)
 def test_quantization_triggers(tmpdir, trigger_fn, expected_count):
     """Test  how many times the quant is called"""
     dm = RegressDataModule()

tests/callbacks/test_stochastic_weight_avg.py

@@ -24,6 +24,7 @@
 from pytorch_lightning.utilities import _TORCH_GREATER_EQUAL_1_6
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.helpers import BoringModel, RandomDataset
+from tests.helpers.skipif import RunIf
 
 if _TORCH_GREATER_EQUAL_1_6:
     from pytorch_lightning.callbacks import StochasticWeightAveraging
@@ -114,40 +115,37 @@ def train_with_swa(tmpdir, batchnorm=True, accelerator=None, gpus=None, num_proc
     assert trainer.lightning_module == model
 
 
-@pytest.mark.skipif(not _TORCH_GREATER_EQUAL_1_6, reason="SWA available from PyTorch 1.6.0")
-@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+@RunIf(min_gpus=2, min_torch="1.6.0")
 @pytest.mark.skipif(
     not os.getenv("PL_RUNNING_SPECIAL_TESTS", '0') == '1', reason="test should be run outside of pytest"
 )
 def test_swa_callback_ddp(tmpdir):
     train_with_swa(tmpdir, accelerator="ddp", gpus=2)
 
 
-@pytest.mark.skipif(not _TORCH_GREATER_EQUAL_1_6, reason="SWA available from PyTorch 1.6.0")
-@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+@RunIf(min_gpus=2, min_torch="1.6.0")
 def test_swa_callback_ddp_spawn(tmpdir):
     train_with_swa(tmpdir, accelerator="ddp_spawn", gpus=2)
 
 
-@pytest.mark.skipif(not _TORCH_GREATER_EQUAL_1_6, reason="SWA available from PyTorch 1.6.0")
+@RunIf(min_torch="1.6.0")
 @pytest.mark.skipif(platform.system() == "Windows", reason="ddp_cpu is not available on Windows")
 def test_swa_callback_ddp_cpu(tmpdir):
     train_with_swa(tmpdir, accelerator="ddp_cpu", num_processes=2)
 
 
-@pytest.mark.skipif(not _TORCH_GREATER_EQUAL_1_6, reason="SWA available from PyTorch 1.6.0")
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires a GPU machine")
+@RunIf(min_gpus=1, min_torch="1.6.0")
 def test_swa_callback_1_gpu(tmpdir):
     train_with_swa(tmpdir, gpus=1)
 
 
-@pytest.mark.skipif(not _TORCH_GREATER_EQUAL_1_6, reason="SWA available from PyTorch 1.6.0")
+@RunIf(min_torch="1.6.0")
 @pytest.mark.parametrize("batchnorm", (True, False))
 def test_swa_callback(tmpdir, batchnorm):
     train_with_swa(tmpdir, batchnorm=batchnorm)
 
 
-@pytest.mark.skipif(not _TORCH_GREATER_EQUAL_1_6, reason="SWA available from PyTorch 1.6.0")
+@RunIf(min_torch="1.6.0")
 def test_swa_raises():
     with pytest.raises(MisconfigurationException, match=">0 integer or a float between 0 and 1"):
         StochasticWeightAveraging(swa_epoch_start=0, swa_lrs=0.1)
@@ -161,7 +159,7 @@ def test_swa_raises():
 
 @pytest.mark.parametrize('stochastic_weight_avg', [False, True])
 @pytest.mark.parametrize('use_callbacks', [False, True])
-@pytest.mark.skipif(not _TORCH_GREATER_EQUAL_1_6, reason="SWA available from PyTorch 1.6.0")
+@RunIf(min_torch="1.6.0")
 def test_trainer_and_stochastic_weight_avg(tmpdir, use_callbacks, stochastic_weight_avg):
     """Test to ensure SWA Callback is injected when `stochastic_weight_avg` is provided to the Trainer"""
 

tests/checkpointing/test_torch_saving.py

@@ -19,6 +19,7 @@
 
 from pytorch_lightning import Trainer
 from tests.helpers import BoringModel
+from tests.helpers.skipif import RunIf
 
 
 def test_model_torch_save(tmpdir):
@@ -57,7 +58,7 @@ def test_model_torch_save_ddp_cpu(tmpdir):
     torch.save(trainer, temp_path)
 
 
-@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+@RunIf(min_gpus=2)
 def test_model_torch_save_ddp_cuda(tmpdir):
     """Test to ensure torch save does not fail for model and trainer using gpu ddp."""
     model = BoringModel()

tests/core/test_datamodules.py

@@ -17,7 +17,6 @@
 from unittest import mock
 from unittest.mock import PropertyMock
 
-import pytest
 import torch
 import torch.nn.functional as F
 
@@ -28,6 +27,7 @@
 from tests.helpers import BoringDataModule, BoringModel
 from tests.helpers.datamodules import ClassifDataModule
 from tests.helpers.simple_models import ClassificationModel
+from tests.helpers.skipif import RunIf
 from tests.helpers.utils import reset_seed, set_random_master_port
 
 
@@ -348,7 +348,7 @@ def test_trainer_attached_to_dm(tmpdir):
     assert dm.trainer is not None
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
+@RunIf(min_gpus=1)
 def test_full_loop_single_gpu(tmpdir):
     reset_seed()
 
@@ -373,7 +373,7 @@ def test_full_loop_single_gpu(tmpdir):
     assert result[0]['test_acc'] > 0.6
 
 
-@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+@RunIf(min_gpus=2)
 def test_full_loop_dp(tmpdir):
     set_random_master_port()
 
@@ -420,7 +420,7 @@ def test_step_end(self, outputs):
     assert result[0]['test_acc'] > 0.6
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
+@RunIf(min_gpus=1)
 @mock.patch("pytorch_lightning.accelerators.accelerator.Accelerator.lightning_module", new_callable=PropertyMock)
 def test_dm_apply_batch_transfer_handler(get_module_mock):
     expected_device = torch.device('cuda', 0)