Fix TPU tests (#16628)

Author: carmocca

tests/tests_fabric/strategies/launchers/test_xla.py

@@ -14,7 +14,7 @@
 import os
 from functools import partial
 from unittest import mock
-from unittest.mock import Mock
+from unittest.mock import MagicMock, Mock
 
 import pytest
 import torch
@@ -84,14 +84,26 @@ def test_tpu_reduce():
 @RunIf(tpu=True)
 @mock.patch("lightning.fabric.strategies.xla.XLAStrategy.root_device")
 def test_xla_mp_device_dataloader_attribute(_, monkeypatch):
+    dataset = RandomDataset(32, 64)
+    dataloader = DataLoader(dataset)
+    strategy = XLAStrategy()
+    isinstance_return = True
+
     import torch_xla.distributed.parallel_loader as parallel_loader
 
-    mp_loader_mock = Mock()
+    class MpDeviceLoaderMock(MagicMock):
+        def __instancecheck__(self, instance):
+            # to make `isinstance(dataloader, MpDeviceLoader)` pass with a mock as class
+            return isinstance_return
+
+    mp_loader_mock = MpDeviceLoaderMock()
     monkeypatch.setattr(parallel_loader, "MpDeviceLoader", mp_loader_mock)
 
-    dataset = RandomDataset(32, 64)
-    dataloader = DataLoader(dataset)
-    strategy = XLAStrategy()
+    processed_dataloader = strategy.process_dataloader(dataloader)
+    assert processed_dataloader is dataloader
+    mp_loader_mock.assert_not_called()  # no-op
+
+    isinstance_return = False
     processed_dataloader = strategy.process_dataloader(dataloader)
     mp_loader_mock.assert_called_with(dataloader, strategy.root_device)
     assert processed_dataloader.dataset == processed_dataloader._loader.dataset

tests/tests_fabric/strategies/test_xla.py

@@ -15,7 +15,7 @@
 import os
 from copy import deepcopy
 from unittest import mock
-from unittest.mock import Mock, patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 import torch
@@ -267,16 +267,28 @@ def test_xla_checkpoint_plugin_being_default(tpu_available):
 
 
 @RunIf(tpu=True)
-@patch("pytorch.lightning.strategies.tpu_spawn.TPUSpawnStrategy.root_device")
+@patch("lightning.pytorch.strategies.tpu_spawn.TPUSpawnStrategy.root_device")
 def test_xla_mp_device_dataloader_attribute(_, monkeypatch):
+    dataset = RandomDataset(32, 64)
+    dataloader = DataLoader(dataset)
+    strategy = TPUSpawnStrategy()
+    isinstance_return = True
+
     import torch_xla.distributed.parallel_loader as parallel_loader
 
-    mp_loader_mock = Mock()
+    class MpDeviceLoaderMock(MagicMock):
+        def __instancecheck__(self, instance):
+            # to make `isinstance(dataloader, MpDeviceLoader)` pass with a mock as class
+            return isinstance_return
+
+    mp_loader_mock = MpDeviceLoaderMock()
     monkeypatch.setattr(parallel_loader, "MpDeviceLoader", mp_loader_mock)
 
-    dataset = RandomDataset(32, 64)
-    dataloader = DataLoader(dataset)
-    strategy = TPUSpawnStrategy()
+    processed_dataloader = strategy.process_dataloader(dataloader)
+    assert processed_dataloader is dataloader
+    mp_loader_mock.assert_not_called()  # no-op
+
+    isinstance_return = False
     processed_dataloader = strategy.process_dataloader(dataloader)
     mp_loader_mock.assert_called_with(dataloader, strategy.root_device)
     assert processed_dataloader.dataset == processed_dataloader._loader.dataset