From 7737a56c46f616f1749b0f6b0886d92831438d94 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= Date: Sat, 18 Dec 2021 02:40:31 +0100 Subject: [PATCH 1/6] Refactor tests that don't need to run assertions in `on_fit_start` --- .../test_accelerator_connector.py | 387 ++++++------------ 1 file changed, 117 insertions(+), 270 deletions(-) diff --git a/tests/accelerators/test_accelerator_connector.py b/tests/accelerators/test_accelerator_connector.py index 0ef10b4eb2a9f..3a09f3989a9cf 100644 --- a/tests/accelerators/test_accelerator_connector.py +++ b/tests/accelerators/test_accelerator_connector.py @@ -100,23 +100,15 @@ def test_accelerator_choice_ddp_spawn(cuda_available_mock, device_count_mock): @mock.patch("torch.cuda.set_device") @mock.patch("torch.cuda.device_count", return_value=2) @mock.patch("pytorch_lightning.plugins.DDPPlugin.setup_distributed", autospec=True) -def test_accelerator_choice_ddp_slurm(set_device_mock, device_count_mock, setup_distributed_mock): - class CB(Callback): - def on_fit_start(self, trainer, pl_module): - assert trainer._accelerator_connector._is_slurm_managing_tasks() - assert isinstance(trainer.accelerator, GPUAccelerator) - assert isinstance(trainer.training_type_plugin, DDPPlugin) - assert isinstance(trainer.training_type_plugin.cluster_environment, SLURMEnvironment) - assert trainer.training_type_plugin.cluster_environment.local_rank() == 1 - assert trainer.training_type_plugin.local_rank == 1 - raise SystemExit() - - model = BoringModel() +def test_accelerator_choice_ddp_slurm(*_): with pytest.deprecated_call(match=r"accelerator='ddp'\)` has been deprecated in v1.5"): - trainer = Trainer(fast_dev_run=True, accelerator="ddp", gpus=2, callbacks=[CB()]) - - with pytest.raises(SystemExit): - trainer.fit(model) + trainer = Trainer(fast_dev_run=True, accelerator="ddp", gpus=2) + assert trainer._accelerator_connector._is_slurm_managing_tasks() + assert isinstance(trainer.accelerator, GPUAccelerator) + assert isinstance(trainer.training_type_plugin, DDPPlugin) + assert isinstance(trainer.training_type_plugin.cluster_environment, SLURMEnvironment) + assert trainer.training_type_plugin.cluster_environment.local_rank() == 1 + assert trainer.training_type_plugin.local_rank == 1 @mock.patch.dict( @@ -133,25 +125,15 @@ def on_fit_start(self, trainer, pl_module): @mock.patch("torch.cuda.set_device") @mock.patch("torch.cuda.device_count", return_value=2) @mock.patch("pytorch_lightning.plugins.DDPPlugin.setup_distributed", autospec=True) -def test_accelerator_choice_ddp2_slurm(set_device_mock, device_count_mock, setup_distributed_mock): - class CB(Callback): - def on_fit_start(self, trainer, pl_module): - assert trainer._accelerator_connector._is_slurm_managing_tasks() - assert isinstance(trainer.accelerator, GPUAccelerator) - assert isinstance(trainer.training_type_plugin, DDP2Plugin) - assert isinstance(trainer.training_type_plugin.cluster_environment, SLURMEnvironment) - assert trainer.training_type_plugin.cluster_environment.local_rank() == 1 - assert trainer.training_type_plugin.local_rank == 1 - raise SystemExit() - - model = BoringModel() +def test_accelerator_choice_ddp2_slurm(*_): with pytest.deprecated_call(match=r"accelerator='ddp2'\)` has been deprecated in v1.5"): - trainer = Trainer(fast_dev_run=True, accelerator="ddp2", gpus=2, callbacks=[CB()]) - - with pytest.raises(SystemExit): - trainer.fit(model) - - set_device_mock.assert_called_once() + trainer = Trainer(fast_dev_run=True, accelerator="ddp2", gpus=2) + assert trainer._accelerator_connector._is_slurm_managing_tasks() + assert isinstance(trainer.accelerator, GPUAccelerator) + assert isinstance(trainer.training_type_plugin, DDP2Plugin) + assert isinstance(trainer.training_type_plugin.cluster_environment, SLURMEnvironment) + assert trainer.training_type_plugin.cluster_environment.local_rank() == 1 + assert trainer.training_type_plugin.local_rank == 1 @mock.patch.dict( @@ -168,24 +150,14 @@ def on_fit_start(self, trainer, pl_module): @mock.patch("torch.cuda.set_device") @mock.patch("torch.cuda.device_count", return_value=1) @mock.patch("pytorch_lightning.plugins.DDPPlugin.setup_distributed", autospec=True) -def test_accelerator_choice_ddp_te(set_device_mock, device_count_mock, setup_distributed_mock): - class CB(Callback): - def on_fit_start(self, trainer, pl_module): - assert isinstance(trainer.accelerator, GPUAccelerator) - assert isinstance(trainer.training_type_plugin, DDPPlugin) - assert isinstance(trainer.training_type_plugin.cluster_environment, TorchElasticEnvironment) - assert trainer.training_type_plugin.cluster_environment.local_rank() == 1 - assert trainer.training_type_plugin.local_rank == 1 - raise SystemExit() - - model = BoringModel() +def test_accelerator_choice_ddp_te(*_): with pytest.deprecated_call(match=r"accelerator='ddp'\)` has been deprecated in v1.5"): - trainer = Trainer(fast_dev_run=True, accelerator="ddp", gpus=2, callbacks=[CB()]) - - with pytest.raises(SystemExit): - trainer.fit(model) - - set_device_mock.assert_called_once() + trainer = Trainer(fast_dev_run=True, accelerator="ddp", gpus=2) + assert isinstance(trainer.accelerator, GPUAccelerator) + assert isinstance(trainer.training_type_plugin, DDPPlugin) + assert isinstance(trainer.training_type_plugin.cluster_environment, TorchElasticEnvironment) + assert trainer.training_type_plugin.cluster_environment.local_rank() == 1 + assert trainer.training_type_plugin.local_rank == 1 @mock.patch.dict( @@ -202,24 +174,14 @@ def on_fit_start(self, trainer, pl_module): @mock.patch("torch.cuda.set_device") @mock.patch("torch.cuda.device_count", return_value=1) @mock.patch("pytorch_lightning.plugins.DDPPlugin.setup_distributed", autospec=True) -def test_accelerator_choice_ddp2_te(set_device_mock, device_count_mock, setup_distributed_mock): - class CB(Callback): - def on_fit_start(self, trainer, pl_module): - assert isinstance(trainer.accelerator, GPUAccelerator) - assert isinstance(trainer.training_type_plugin, DDP2Plugin) - assert isinstance(trainer.training_type_plugin.cluster_environment, TorchElasticEnvironment) - assert trainer.training_type_plugin.cluster_environment.local_rank() == 1 - assert trainer.training_type_plugin.local_rank == 1 - raise SystemExit() - - model = BoringModel() +def test_accelerator_choice_ddp2_te(*_): with pytest.deprecated_call(match=r"accelerator='ddp2'\)` has been deprecated in v1.5"): - trainer = Trainer(fast_dev_run=True, accelerator="ddp2", gpus=2, callbacks=[CB()]) - - with pytest.raises(SystemExit): - trainer.fit(model) - - set_device_mock.assert_called_once() + trainer = Trainer(fast_dev_run=True, accelerator="ddp2", gpus=2) + assert isinstance(trainer.accelerator, GPUAccelerator) + assert isinstance(trainer.training_type_plugin, DDP2Plugin) + assert isinstance(trainer.training_type_plugin.cluster_environment, TorchElasticEnvironment) + assert trainer.training_type_plugin.cluster_environment.local_rank() == 1 + assert trainer.training_type_plugin.local_rank == 1 @mock.patch.dict( @@ -227,21 +189,13 @@ def on_fit_start(self, trainer, pl_module): ) @mock.patch("torch.cuda.device_count", return_value=0) @mock.patch("pytorch_lightning.plugins.DDPPlugin.setup_distributed", autospec=True) -def test_accelerator_choice_ddp_cpu_te(device_count_mock, setup_distributed_mock): - class CB(Callback): - def on_fit_start(self, trainer, pl_module): - assert isinstance(trainer.accelerator, CPUAccelerator) - assert isinstance(trainer.training_type_plugin, DDPPlugin) - assert isinstance(trainer.training_type_plugin.cluster_environment, TorchElasticEnvironment) - assert trainer.training_type_plugin.cluster_environment.local_rank() == 1 - assert trainer.training_type_plugin.local_rank == 1 - raise SystemExit() - - model = BoringModel() - trainer = Trainer(fast_dev_run=True, accelerator="ddp_cpu", num_processes=2, callbacks=[CB()]) - - with pytest.raises(SystemExit): - trainer.fit(model) +def test_accelerator_choice_ddp_cpu_te(*_): + trainer = Trainer(fast_dev_run=True, accelerator="ddp_cpu", num_processes=2) + assert isinstance(trainer.accelerator, CPUAccelerator) + assert isinstance(trainer.training_type_plugin, DDPPlugin) + assert isinstance(trainer.training_type_plugin.cluster_environment, TorchElasticEnvironment) + assert trainer.training_type_plugin.cluster_environment.local_rank() == 1 + assert trainer.training_type_plugin.local_rank == 1 @mock.patch.dict( @@ -258,24 +212,14 @@ def on_fit_start(self, trainer, pl_module): @mock.patch("torch.cuda.set_device") @mock.patch("torch.cuda.device_count", return_value=1) @mock.patch("pytorch_lightning.plugins.DDPPlugin.setup_distributed", autospec=True) -def test_accelerator_choice_ddp_kubeflow(set_device_mock, device_count_mock, setup_distributed_mock): - class CB(Callback): - def on_fit_start(self, trainer, pl_module): - assert isinstance(trainer.accelerator, GPUAccelerator) - assert isinstance(trainer.training_type_plugin, DDPPlugin) - assert isinstance(trainer.training_type_plugin.cluster_environment, KubeflowEnvironment) - assert trainer.training_type_plugin.cluster_environment.local_rank() == 0 - assert trainer.training_type_plugin.local_rank == 0 - raise SystemExit() - - model = BoringModel() +def test_accelerator_choice_ddp_kubeflow(*_): with pytest.deprecated_call(match=r"accelerator='ddp'\)` has been deprecated in v1.5"): - trainer = Trainer(fast_dev_run=True, accelerator="ddp", gpus=1, callbacks=[CB()]) - - with pytest.raises(SystemExit): - trainer.fit(model) - - set_device_mock.assert_called_once() + trainer = Trainer(fast_dev_run=True, accelerator="ddp", gpus=1) + assert isinstance(trainer.accelerator, GPUAccelerator) + assert isinstance(trainer.training_type_plugin, DDPPlugin) + assert isinstance(trainer.training_type_plugin.cluster_environment, KubeflowEnvironment) + assert trainer.training_type_plugin.cluster_environment.local_rank() == 0 + assert trainer.training_type_plugin.local_rank == 0 @mock.patch.dict( @@ -290,21 +234,13 @@ def on_fit_start(self, trainer, pl_module): ) @mock.patch("torch.cuda.device_count", return_value=0) @mock.patch("pytorch_lightning.plugins.DDPPlugin.setup_distributed", autospec=True) -def test_accelerator_choice_ddp_cpu_kubeflow(device_count_mock, setup_distributed_mock): - class CB(Callback): - def on_fit_start(self, trainer, pl_module): - assert isinstance(trainer.accelerator, CPUAccelerator) - assert isinstance(trainer.training_type_plugin, DDPPlugin) - assert isinstance(trainer.training_type_plugin.cluster_environment, KubeflowEnvironment) - assert trainer.training_type_plugin.cluster_environment.local_rank() == 0 - assert trainer.training_type_plugin.local_rank == 0 - raise SystemExit() - - model = BoringModel() - trainer = Trainer(fast_dev_run=True, accelerator="ddp_cpu", num_processes=1, callbacks=[CB()]) - - with pytest.raises(SystemExit): - trainer.fit(model) +def test_accelerator_choice_ddp_cpu_kubeflow(*_): + trainer = Trainer(fast_dev_run=True, accelerator="ddp_cpu", num_processes=1) + assert isinstance(trainer.accelerator, CPUAccelerator) + assert isinstance(trainer.training_type_plugin, DDPPlugin) + assert isinstance(trainer.training_type_plugin.cluster_environment, KubeflowEnvironment) + assert trainer.training_type_plugin.cluster_environment.local_rank() == 0 + assert trainer.training_type_plugin.local_rank == 0 @mock.patch.dict( @@ -320,21 +256,13 @@ def on_fit_start(self, trainer, pl_module): ) @mock.patch("torch.cuda.device_count", return_value=0) @mock.patch("pytorch_lightning.plugins.DDPPlugin.setup_distributed", autospec=True) -def test_accelerator_choice_ddp_cpu_slurm(device_count_mock, setup_distributed_mock): - class CB(Callback): - def on_fit_start(self, trainer, pl_module): - assert trainer._accelerator_connector._is_slurm_managing_tasks() - assert isinstance(trainer.accelerator, CPUAccelerator) - assert isinstance(trainer.training_type_plugin, DDPPlugin) - assert isinstance(trainer.training_type_plugin.cluster_environment, SLURMEnvironment) - assert trainer.training_type_plugin.local_rank == 0 - raise SystemExit() - - model = BoringModel() - trainer = Trainer(fast_dev_run=True, accelerator="ddp_cpu", num_processes=2, callbacks=[CB()]) - - with pytest.raises(SystemExit): - trainer.fit(model) +def test_accelerator_choice_ddp_cpu_slurm(*_): + trainer = Trainer(fast_dev_run=True, accelerator="ddp_cpu", num_processes=2) + assert trainer._accelerator_connector._is_slurm_managing_tasks() + assert isinstance(trainer.accelerator, CPUAccelerator) + assert isinstance(trainer.training_type_plugin, DDPPlugin) + assert isinstance(trainer.training_type_plugin.cluster_environment, SLURMEnvironment) + assert trainer.training_type_plugin.local_rank == 0 @RunIf(skip_windows=True, standalone=True) @@ -449,19 +377,11 @@ class DistributedPlugin(DDPPlugin): ) @mock.patch("torch.cuda.device_count", return_value=0) @mock.patch("pytorch_lightning.plugins.DDPPlugin.setup_distributed", autospec=True) -def test_dist_backend_accelerator_mapping(device_count_mock, setup_distributed_mock): - class CB(Callback): - def on_fit_start(self, trainer, pl_module): - assert isinstance(trainer.accelerator, CPUAccelerator) - assert isinstance(trainer.training_type_plugin, DDPPlugin) - assert trainer.training_type_plugin.local_rank == 0 - raise SystemExit() - - model = BoringModel() - trainer = Trainer(fast_dev_run=True, strategy="ddp_spawn", num_processes=2, callbacks=[CB()]) - - with pytest.raises(SystemExit): - trainer.fit(model) +def test_dist_backend_accelerator_mapping(*_): + trainer = Trainer(fast_dev_run=True, strategy="ddp_spawn", num_processes=2) + assert isinstance(trainer.accelerator, CPUAccelerator) + assert isinstance(trainer.training_type_plugin, DDPPlugin) + assert trainer.training_type_plugin.local_rank == 0 @mock.patch("pytorch_lightning.utilities._IS_INTERACTIVE", return_value=True) @@ -792,21 +712,13 @@ def test_strategy_choice_ddp_spawn(cuda_available_mock, device_count_mock): @mock.patch("pytorch_lightning.plugins.DDPPlugin.setup_distributed", autospec=True) @pytest.mark.parametrize("strategy", ["ddp", DDPPlugin()]) def test_strategy_choice_ddp_slurm(setup_distributed_mock, strategy): - class CB(Callback): - def on_fit_start(self, trainer, pl_module): - assert trainer._accelerator_connector._is_slurm_managing_tasks() - assert isinstance(trainer.accelerator, GPUAccelerator) - assert isinstance(trainer.training_type_plugin, DDPPlugin) - assert isinstance(trainer.training_type_plugin.cluster_environment, SLURMEnvironment) - assert trainer.training_type_plugin.cluster_environment.local_rank() == 1 - assert trainer.training_type_plugin.local_rank == 1 - raise SystemExit() - - model = BoringModel() - trainer = Trainer(fast_dev_run=True, strategy=strategy, gpus=2, callbacks=[CB()]) - - with pytest.raises(SystemExit): - trainer.fit(model) + trainer = Trainer(fast_dev_run=True, strategy=strategy, gpus=2) + assert trainer._accelerator_connector._is_slurm_managing_tasks() + assert isinstance(trainer.accelerator, GPUAccelerator) + assert isinstance(trainer.training_type_plugin, DDPPlugin) + assert isinstance(trainer.training_type_plugin.cluster_environment, SLURMEnvironment) + assert trainer.training_type_plugin.cluster_environment.local_rank() == 1 + assert trainer.training_type_plugin.local_rank == 1 @mock.patch.dict( @@ -825,23 +737,13 @@ def on_fit_start(self, trainer, pl_module): @mock.patch("pytorch_lightning.plugins.DDPPlugin.setup_distributed", autospec=True) @pytest.mark.parametrize("strategy", ["ddp2", DDP2Plugin()]) def test_strategy_choice_ddp2_slurm(set_device_mock, device_count_mock, setup_distributed_mock, strategy): - class CB(Callback): - def on_fit_start(self, trainer, pl_module): - assert trainer._accelerator_connector._is_slurm_managing_tasks() - assert isinstance(trainer.accelerator, GPUAccelerator) - assert isinstance(trainer.training_type_plugin, DDP2Plugin) - assert isinstance(trainer.training_type_plugin.cluster_environment, SLURMEnvironment) - assert trainer.training_type_plugin.cluster_environment.local_rank() == 1 - assert trainer.training_type_plugin.local_rank == 1 - raise SystemExit() - - model = BoringModel() - trainer = Trainer(fast_dev_run=True, strategy=strategy, gpus=2, callbacks=[CB()]) - - with pytest.raises(SystemExit): - trainer.fit(model) - - set_device_mock.assert_called_once() + trainer = Trainer(fast_dev_run=True, strategy=strategy, gpus=2) + assert trainer._accelerator_connector._is_slurm_managing_tasks() + assert isinstance(trainer.accelerator, GPUAccelerator) + assert isinstance(trainer.training_type_plugin, DDP2Plugin) + assert isinstance(trainer.training_type_plugin.cluster_environment, SLURMEnvironment) + assert trainer.training_type_plugin.cluster_environment.local_rank() == 1 + assert trainer.training_type_plugin.local_rank == 1 @mock.patch.dict( @@ -858,23 +760,13 @@ def on_fit_start(self, trainer, pl_module): @mock.patch("torch.cuda.set_device") @mock.patch("torch.cuda.device_count", return_value=2) @mock.patch("pytorch_lightning.plugins.DDPPlugin.setup_distributed", autospec=True) -def test_strategy_choice_ddp_te(set_device_mock, device_count_mock, setup_distributed_mock): - class CB(Callback): - def on_fit_start(self, trainer, pl_module): - assert isinstance(trainer.accelerator, GPUAccelerator) - assert isinstance(trainer.training_type_plugin, DDPPlugin) - assert isinstance(trainer.training_type_plugin.cluster_environment, TorchElasticEnvironment) - assert trainer.training_type_plugin.cluster_environment.local_rank() == 1 - assert trainer.training_type_plugin.local_rank == 1 - raise SystemExit() - - model = BoringModel() - trainer = Trainer(fast_dev_run=True, strategy="ddp", gpus=2, callbacks=[CB()]) - - with pytest.raises(SystemExit): - trainer.fit(model) - - set_device_mock.assert_called_once() +def test_strategy_choice_ddp_te(*_): + trainer = Trainer(fast_dev_run=True, strategy="ddp", gpus=2) + assert isinstance(trainer.accelerator, GPUAccelerator) + assert isinstance(trainer.training_type_plugin, DDPPlugin) + assert isinstance(trainer.training_type_plugin.cluster_environment, TorchElasticEnvironment) + assert trainer.training_type_plugin.cluster_environment.local_rank() == 1 + assert trainer.training_type_plugin.local_rank == 1 @mock.patch.dict( @@ -891,23 +783,13 @@ def on_fit_start(self, trainer, pl_module): @mock.patch("torch.cuda.set_device") @mock.patch("torch.cuda.device_count", return_value=2) @mock.patch("pytorch_lightning.plugins.DDPPlugin.setup_distributed", autospec=True) -def test_strategy_choice_ddp2_te(set_device_mock, device_count_mock, setup_distributed_mock): - class CB(Callback): - def on_fit_start(self, trainer, pl_module): - assert isinstance(trainer.accelerator, GPUAccelerator) - assert isinstance(trainer.training_type_plugin, DDP2Plugin) - assert isinstance(trainer.training_type_plugin.cluster_environment, TorchElasticEnvironment) - assert trainer.training_type_plugin.cluster_environment.local_rank() == 1 - assert trainer.training_type_plugin.local_rank == 1 - raise SystemExit() - - model = BoringModel() - trainer = Trainer(fast_dev_run=True, strategy="ddp2", gpus=2, callbacks=[CB()]) - - with pytest.raises(SystemExit): - trainer.fit(model) - - set_device_mock.assert_called_once() +def test_strategy_choice_ddp2_te(*_): + trainer = Trainer(fast_dev_run=True, strategy="ddp2", gpus=2) + assert isinstance(trainer.accelerator, GPUAccelerator) + assert isinstance(trainer.training_type_plugin, DDP2Plugin) + assert isinstance(trainer.training_type_plugin.cluster_environment, TorchElasticEnvironment) + assert trainer.training_type_plugin.cluster_environment.local_rank() == 1 + assert trainer.training_type_plugin.local_rank == 1 @mock.patch.dict( @@ -915,21 +797,13 @@ def on_fit_start(self, trainer, pl_module): ) @mock.patch("torch.cuda.device_count", return_value=0) @mock.patch("pytorch_lightning.plugins.DDPPlugin.setup_distributed", autospec=True) -def test_strategy_choice_ddp_cpu_te(device_count_mock, setup_distributed_mock): - class CB(Callback): - def on_fit_start(self, trainer, pl_module): - assert isinstance(trainer.accelerator, CPUAccelerator) - assert isinstance(trainer.training_type_plugin, DDPPlugin) - assert isinstance(trainer.training_type_plugin.cluster_environment, TorchElasticEnvironment) - assert trainer.training_type_plugin.cluster_environment.local_rank() == 1 - assert trainer.training_type_plugin.local_rank == 1 - raise SystemExit() - - model = BoringModel() - trainer = Trainer(fast_dev_run=True, strategy="ddp_spawn", num_processes=2, callbacks=[CB()]) - - with pytest.raises(SystemExit): - trainer.fit(model) +def test_strategy_choice_ddp_cpu_te(*_): + trainer = Trainer(fast_dev_run=True, strategy="ddp_spawn", num_processes=2) + assert isinstance(trainer.accelerator, CPUAccelerator) + assert isinstance(trainer.training_type_plugin, DDPPlugin) + assert isinstance(trainer.training_type_plugin.cluster_environment, TorchElasticEnvironment) + assert trainer.training_type_plugin.cluster_environment.local_rank() == 1 + assert trainer.training_type_plugin.local_rank == 1 @mock.patch.dict( @@ -946,23 +820,13 @@ def on_fit_start(self, trainer, pl_module): @mock.patch("torch.cuda.set_device") @mock.patch("torch.cuda.device_count", return_value=1) @mock.patch("pytorch_lightning.plugins.DDPPlugin.setup_distributed", autospec=True) -def test_strategy_choice_ddp_kubeflow(set_device_mock, device_count_mock, setup_distributed_mock): - class CB(Callback): - def on_fit_start(self, trainer, pl_module): - assert isinstance(trainer.accelerator, GPUAccelerator) - assert isinstance(trainer.training_type_plugin, DDPPlugin) - assert isinstance(trainer.training_type_plugin.cluster_environment, KubeflowEnvironment) - assert trainer.training_type_plugin.cluster_environment.local_rank() == 0 - assert trainer.training_type_plugin.local_rank == 0 - raise SystemExit() - - model = BoringModel() - trainer = Trainer(fast_dev_run=True, strategy="ddp", gpus=1, callbacks=[CB()]) - - with pytest.raises(SystemExit): - trainer.fit(model) - - set_device_mock.assert_called_once() +def test_strategy_choice_ddp_kubeflow(*_): + trainer = Trainer(fast_dev_run=True, strategy="ddp", gpus=1) + assert isinstance(trainer.accelerator, GPUAccelerator) + assert isinstance(trainer.training_type_plugin, DDPPlugin) + assert isinstance(trainer.training_type_plugin.cluster_environment, KubeflowEnvironment) + assert trainer.training_type_plugin.cluster_environment.local_rank() == 0 + assert trainer.training_type_plugin.local_rank == 0 @mock.patch.dict( @@ -977,21 +841,13 @@ def on_fit_start(self, trainer, pl_module): ) @mock.patch("torch.cuda.device_count", return_value=0) @mock.patch("pytorch_lightning.plugins.DDPPlugin.setup_distributed", autospec=True) -def test_strategy_choice_ddp_cpu_kubeflow(device_count_mock, setup_distributed_mock): - class CB(Callback): - def on_fit_start(self, trainer, pl_module): - assert isinstance(trainer.accelerator, CPUAccelerator) - assert isinstance(trainer.training_type_plugin, DDPPlugin) - assert isinstance(trainer.training_type_plugin.cluster_environment, KubeflowEnvironment) - assert trainer.training_type_plugin.cluster_environment.local_rank() == 0 - assert trainer.training_type_plugin.local_rank == 0 - raise SystemExit() - - model = BoringModel() - trainer = Trainer(fast_dev_run=True, strategy="ddp_spawn", num_processes=2, callbacks=[CB()]) - - with pytest.raises(SystemExit): - trainer.fit(model) +def test_strategy_choice_ddp_cpu_kubeflow(*_): + trainer = Trainer(fast_dev_run=True, strategy="ddp_spawn", num_processes=2) + assert isinstance(trainer.accelerator, CPUAccelerator) + assert isinstance(trainer.training_type_plugin, DDPPlugin) + assert isinstance(trainer.training_type_plugin.cluster_environment, KubeflowEnvironment) + assert trainer.training_type_plugin.cluster_environment.local_rank() == 0 + assert trainer.training_type_plugin.local_rank == 0 @mock.patch.dict( @@ -1009,20 +865,11 @@ def on_fit_start(self, trainer, pl_module): @mock.patch("pytorch_lightning.plugins.DDPPlugin.setup_distributed", autospec=True) @pytest.mark.parametrize("strategy", ["ddp", DDPPlugin()]) def test_strategy_choice_ddp_cpu_slurm(device_count_mock, setup_distributed_mock, strategy): - class CB(Callback): - def on_fit_start(self, trainer, pl_module): - assert trainer._accelerator_connector._is_slurm_managing_tasks() - assert isinstance(trainer.accelerator, CPUAccelerator) - assert isinstance(trainer.training_type_plugin, DDPPlugin) - assert isinstance(trainer.training_type_plugin.cluster_environment, SLURMEnvironment) - assert trainer.training_type_plugin.local_rank == 0 - raise SystemExit() - - model = BoringModel() - trainer = Trainer(fast_dev_run=True, strategy=strategy, num_processes=2, callbacks=[CB()]) - - with pytest.raises(SystemExit): - trainer.fit(model) + trainer = Trainer(fast_dev_run=True, strategy=strategy, num_processes=2) + assert isinstance(trainer.accelerator, CPUAccelerator) + assert isinstance(trainer.training_type_plugin, DDPPlugin) + assert isinstance(trainer.training_type_plugin.cluster_environment, SLURMEnvironment) + assert trainer.training_type_plugin.local_rank == 0 def test_unsupported_tpu_choice(monkeypatch): From 1fcf150ee21d5c30cd1fabf9312ed62944ece25c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= Date: Sat, 18 Dec 2021 02:45:40 +0100 Subject: [PATCH 2/6] add test for set_device --- tests/accelerators/test_gpu.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/accelerators/test_gpu.py b/tests/accelerators/test_gpu.py index ece78da24972d..bf880b12e6923 100644 --- a/tests/accelerators/test_gpu.py +++ b/tests/accelerators/test_gpu.py @@ -1,6 +1,10 @@ +from unittest import mock + import torch +from pytorch_lightning import Trainer from pytorch_lightning.accelerators import GPUAccelerator +from tests.helpers import BoringModel from tests.helpers.runif import RunIf @@ -26,3 +30,12 @@ def test_get_nvidia_gpu_stats(tmpdir): for f in fields: assert any(f in h for h in gpu_stats.keys()) + + +@RunIf(gpus=1) +@mock.patch("torch.cuda.set_device") +def test_set_cuda_device(set_device_mock, tmpdir): + model = BoringModel() + trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True, accelerator="gpu", devices=1, enable_checkpointing=False, enable_model_summary=False, enable_progress_bar=False) + trainer.fit(model) + set_device_mock.assert_called_once() From 9426db2bc4af58906a46758dcf10e56678b26fee Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= Date: Sat, 18 Dec 2021 02:46:07 +0100 Subject: [PATCH 3/6] update formatting --- tests/accelerators/test_gpu.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/accelerators/test_gpu.py b/tests/accelerators/test_gpu.py index bf880b12e6923..4a96ecbc28b89 100644 --- a/tests/accelerators/test_gpu.py +++ b/tests/accelerators/test_gpu.py @@ -36,6 +36,14 @@ def test_get_nvidia_gpu_stats(tmpdir): @mock.patch("torch.cuda.set_device") def test_set_cuda_device(set_device_mock, tmpdir): model = BoringModel() - trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True, accelerator="gpu", devices=1, enable_checkpointing=False, enable_model_summary=False, enable_progress_bar=False) + trainer = Trainer( + default_root_dir=tmpdir, + fast_dev_run=True, + accelerator="gpu", + devices=1, + enable_checkpointing=False, + enable_model_summary=False, + enable_progress_bar=False, + ) trainer.fit(model) set_device_mock.assert_called_once() From 2b6fdcb913baf278d8fc1caa6b9f830cdee8289a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= Date: Sat, 18 Dec 2021 02:56:27 +0100 Subject: [PATCH 4/6] add missing license --- tests/accelerators/test_gpu.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/accelerators/test_gpu.py b/tests/accelerators/test_gpu.py index 4a96ecbc28b89..d57a04ac3aa78 100644 --- a/tests/accelerators/test_gpu.py +++ b/tests/accelerators/test_gpu.py @@ -1,3 +1,16 @@ +# 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. from unittest import mock import torch From 097136df71fb6bb0ebd040b838d0d706431cd3d4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= Date: Sat, 18 Dec 2021 03:13:47 +0100 Subject: [PATCH 5/6] unused imports --- tests/accelerators/test_accelerator_connector.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/accelerators/test_accelerator_connector.py b/tests/accelerators/test_accelerator_connector.py index 3a09f3989a9cf..a6e9115ef4e74 100644 --- a/tests/accelerators/test_accelerator_connector.py +++ b/tests/accelerators/test_accelerator_connector.py @@ -24,7 +24,6 @@ from pytorch_lightning.accelerators.accelerator import Accelerator from pytorch_lightning.accelerators.cpu import CPUAccelerator from pytorch_lightning.accelerators.gpu import GPUAccelerator -from pytorch_lightning.callbacks import Callback from pytorch_lightning.plugins import ( DataParallelPlugin, DDP2Plugin, @@ -45,7 +44,6 @@ ) from pytorch_lightning.utilities import _AcceleratorType, _StrategyType from pytorch_lightning.utilities.exceptions import MisconfigurationException -from tests.helpers.boring_model import BoringModel from tests.helpers.runif import RunIf From b1b0cdd0c934d9c1fa14533a9df272033529aafe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= Date: Sat, 18 Dec 2021 03:14:56 +0100 Subject: [PATCH 6/6] set RunIf condition --- tests/accelerators/test_gpu.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/accelerators/test_gpu.py b/tests/accelerators/test_gpu.py index d57a04ac3aa78..110ba1be9a82c 100644 --- a/tests/accelerators/test_gpu.py +++ b/tests/accelerators/test_gpu.py @@ -45,7 +45,7 @@ def test_get_nvidia_gpu_stats(tmpdir): assert any(f in h for h in gpu_stats.keys()) -@RunIf(gpus=1) +@RunIf(min_gpus=1) @mock.patch("torch.cuda.set_device") def test_set_cuda_device(set_device_mock, tmpdir): model = BoringModel()