Remove rank_zero_only on DataModule prepare_data (#7945)

Signed-off-by: Max Ehrlich <[email protected]>

Author: Queuecumber

pytorch_lightning/core/datamodule.py

@@ -20,7 +20,6 @@
 from torch.utils.data import DataLoader, Dataset, IterableDataset
 
 from pytorch_lightning.core.hooks import CheckpointHooks, DataHooks
-from pytorch_lightning.utilities import rank_zero_only
 from pytorch_lightning.utilities.argparse import add_argparse_args, from_argparse_args, get_init_arguments_and_types
 
 
@@ -329,7 +328,7 @@ def test_dataloader():
     def __new__(cls, *args: Any, **kwargs: Any) -> 'LightningDataModule':
         obj = super().__new__(cls)
         # track `DataHooks` calls and run `prepare_data` only on rank zero
-        obj.prepare_data = cls._track_data_hook_calls(obj, rank_zero_only(obj.prepare_data))
+        obj.prepare_data = cls._track_data_hook_calls(obj, obj.prepare_data)
         obj.setup = cls._track_data_hook_calls(obj, obj.setup)
         obj.teardown = cls._track_data_hook_calls(obj, obj.teardown)
         return obj

tests/core/test_datamodules.py

@@ -34,6 +34,7 @@
 @mock.patch("pytorch_lightning.trainer.trainer.Trainer.local_rank", new_callable=PropertyMock)
 def test_can_prepare_data(local_rank, node_rank):
 
+    model = BoringModel()
     dm = BoringDataModule()
     trainer = Trainer()
     trainer.datamodule = dm
@@ -43,30 +44,54 @@ def test_can_prepare_data(local_rank, node_rank):
     # local rank = 0   (True)
     trainer.prepare_data_per_node = True
 
+    dm.random_full = None
+    dm._has_prepared_data = False
     local_rank.return_value = 0
     assert trainer.local_rank == 0
     assert trainer.data_connector.can_prepare_data()
 
+    trainer.data_connector.prepare_data(model)
+    assert dm.random_full is not None
+
     # local rank = 1   (False)
+    dm.random_full = None
+    dm._has_prepared_data = False
     local_rank.return_value = 1
     assert trainer.local_rank == 1
     assert not trainer.data_connector.can_prepare_data()
 
+    trainer.data_connector.prepare_data(model)
+    assert dm.random_full is None
+
     # prepare_data_per_node = False (prepare across all nodes)
     # global rank = 0   (True)
+    dm.random_full = None
+    dm._has_prepared_data = False
     trainer.prepare_data_per_node = False
     node_rank.return_value = 0
     local_rank.return_value = 0
     assert trainer.data_connector.can_prepare_data()
 
+    trainer.data_connector.prepare_data(model)
+    assert dm.random_full is not None
+
     # global rank = 1   (False)
+    dm.random_full = None
+    dm._has_prepared_data = False
     node_rank.return_value = 1
     local_rank.return_value = 0
     assert not trainer.data_connector.can_prepare_data()
+
+    trainer.data_connector.prepare_data(model)
+    assert dm.random_full is None
+
     node_rank.return_value = 0
     local_rank.return_value = 1
     assert not trainer.data_connector.can_prepare_data()
 
+    trainer.data_connector.prepare_data(model)
+    assert dm.random_full is None
+
     # 2 dm
     # prepar per node = True
     # local rank = 0 (True)