test and flake

Author: ant0nsc

InnerEye/ML/SSL/lightning_modules/ssl_online_evaluator.py

@@ -86,8 +86,6 @@ def on_pretrain_routine_start(self, trainer: pl.Trainer, pl_module: pl.Lightning
         if accelerator.is_distributed:
             if accelerator.use_ddp:
                 self.evaluator = DistributedDataParallel(self.evaluator, device_ids=[pl_module.device])  # type: ignore
-            elif accelerator.use_dp:
-                self.evaluator = DataParallel(self.evaluator, device_ids=[pl_module.device])  # type: ignore
             else:
                 rank_zero_warn("This type of distributed accelerator is not supported. "
                                "The online evaluator will not synchronize across GPUs.")

Tests/SSL/test_ssl_containers.py

@@ -12,7 +12,7 @@
 import pytest
 import torch
 from pl_bolts.models.self_supervised.resnets import ResNet
-from pytorch_lightning import Trainer
+from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.callbacks import ModelCheckpoint
 from torch.nn import Module
 from torch.nn.parallel import DistributedDataParallel
@@ -344,13 +344,11 @@ def test_online_evaluator_recovery(test_output_dirs: OutputFolderForTests) -> No
 
 
 @pytest.mark.gpu
-def test_online_evaluator_distributed() -> None:
+def test_online_evaluator_not_distributed() -> None:
     """
-    A very basic test to check if the online evaluator uses the DDP flag correctly.
+    Check if the online evaluator uses the DDP flag correctly when running not distributed
     """
-    mock_ddp_result = "mock_ddp_result"
-    with mock.patch("InnerEye.ML.SSL.lightning_modules.ssl_online_evaluator.DistributedDataParallel",
-                    return_value=mock_ddp_result) as mock_ddp:
+    with mock.patch("InnerEye.ML.SSL.lightning_modules.ssl_online_evaluator.DistributedDataParallel") as mock_ddp:
         callback = SSLOnlineEvaluatorInnerEye(class_weights=None,
                                               z_dim=1,
                                               num_classes=2,
@@ -361,20 +359,38 @@ def test_online_evaluator_distributed() -> None:
 
         # Standard trainer without DDP
         trainer = Trainer()
+        # Test the flag that the internal logic of on_pretrain_routine_start uses
+        assert not trainer.accelerator_connector.is_distributed
         mock_module = mock.MagicMock(device=torch.device("cpu"))
         callback.on_pretrain_routine_start(trainer, mock_module)
         assert isinstance(callback.evaluator, Module)
         mock_ddp.assert_not_called()
 
+
+@pytest.mark.gpu
+def test_online_evaluator_distributed() -> None:
+    """
+    Check if the online evaluator uses the DDP flag correctly when running distributed.
+    """
+    mock_ddp_result = "mock_ddp_result"
+    with mock.patch("InnerEye.ML.SSL.lightning_modules.ssl_online_evaluator.DistributedDataParallel",
+                    return_value=mock_ddp_result) as mock_ddp:
+        callback = SSLOnlineEvaluatorInnerEye(class_weights=None,
+                                              z_dim=1,
+                                              num_classes=2,
+                                              dataset="foo",
+                                              drop_p=0.2,
+                                              learning_rate=1e-5)
+
         # Trainer with DDP
-        mock_device = "fake_device"
-        mock_module = mock.MagicMock(device=mock_device)
+        device = torch.device("cuda:0")
+        mock_module = mock.MagicMock(device=device)
         trainer = Trainer(accelerator="ddp", gpus=2)
         # Test the two flags that the internal logic of on_pretrain_routine_start uses
         assert trainer.accelerator_connector.is_distributed
         assert trainer.accelerator_connector.use_ddp
         callback.on_pretrain_routine_start(trainer, mock_module)
         # Check that the evaluator has been turned into a DDP object
         # We still need to mock DDP here because the constructor relies on having a process group available
-        mock_ddp.assert_called_once_with(callback.evaluator, device_ids=[mock_device])
+        mock_ddp.assert_called_once_with(callback.evaluator, device_ids=[device])
         assert callback.evaluator == mock_ddp_result