Removed the deprecated datamodule_checkpointhooks (#14909)

Co-authored-by: awaelchli <[email protected]>
Co-authored-by: otaj <[email protected]>

Author: 3 people

src/pytorch_lightning/CHANGELOG.md

@@ -274,6 +274,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Removed the deprecated `Trainer.{validated,tested,predicted}_ckpt_path` ([#14897](https://github.com/Lightning-AI/lightning/pull/14897))
 
+- Removed the deprecated `LightningDataModule.on_save/load_checkpoint` hooks ([#14909](https://github.com/Lightning-AI/lightning/pull/14909))
 
 ### Fixed
 

src/pytorch_lightning/core/datamodule.py

@@ -20,7 +20,7 @@
 
 import pytorch_lightning as pl
 from lightning_lite.utilities.types import _PATH
-from pytorch_lightning.core.hooks import CheckpointHooks, DataHooks
+from pytorch_lightning.core.hooks import DataHooks
 from pytorch_lightning.core.mixins import HyperparametersMixin
 from pytorch_lightning.core.saving import _load_from_checkpoint
 from pytorch_lightning.utilities.argparse import (
@@ -32,7 +32,7 @@
 from pytorch_lightning.utilities.types import _ADD_ARGPARSE_RETURN, EVAL_DATALOADERS, TRAIN_DATALOADERS
 
 
-class LightningDataModule(CheckpointHooks, DataHooks, HyperparametersMixin):
+class LightningDataModule(DataHooks, HyperparametersMixin):
     """A DataModule standardizes the training, val, test splits, data preparation and transforms. The main
     advantage is consistent data splits, data preparation and transforms across models.
 

src/pytorch_lightning/core/saving.py

@@ -222,8 +222,9 @@ def _load_state(
 
     obj = cls(**_cls_kwargs)
 
-    # give model a chance to load something
-    obj.on_load_checkpoint(checkpoint)
+    if isinstance(obj, pl.LightningModule):
+        # give model a chance to load something
+        obj.on_load_checkpoint(checkpoint)
 
     if isinstance(obj, pl.LightningDataModule):
         if obj.__class__.__qualname__ in checkpoint:

src/pytorch_lightning/trainer/configuration_validator.py

@@ -51,8 +51,6 @@ def verify_loop_configurations(trainer: "pl.Trainer") -> None:
     _check_on_epoch_start_end(model)
     # TODO: Delete on_pretrain_routine_start/end hooks in v1.8
     _check_on_pretrain_routine(model)
-    # TODO: Delete CheckpointHooks off LightningDataModule in v1.8
-    _check_datamodule_checkpoint_hooks(trainer)
 
 
 def __verify_train_val_loop_configuration(trainer: "pl.Trainer", model: "pl.LightningModule") -> None:
@@ -261,16 +259,3 @@ def _check_deprecated_callback_hooks(trainer: "pl.Trainer") -> None:
                     f"The `Callback.{hook}` hook has been deprecated in v1.6 and"
                     " will be removed in v1.8. Please use `Callback.on_fit_start` instead."
                 )
-
-
-def _check_datamodule_checkpoint_hooks(trainer: "pl.Trainer") -> None:
-    if is_overridden(method_name="on_save_checkpoint", instance=trainer.datamodule):
-        rank_zero_deprecation(
-            "`LightningDataModule.on_save_checkpoint` was deprecated in"
-            " v1.6 and will be removed in v1.8. Use `state_dict` instead."
-        )
-    if is_overridden(method_name="on_load_checkpoint", instance=trainer.datamodule):
-        rank_zero_deprecation(
-            "`LightningDataModule.on_load_checkpoint` was deprecated in"
-            " v1.6 and will be removed in v1.8. Use `load_state_dict` instead."
-        )

src/pytorch_lightning/trainer/connectors/checkpoint_connector.py

@@ -243,12 +243,10 @@ def restore_datamodule(self) -> None:
             return
 
         datamodule = self.trainer.datamodule
-        if datamodule is not None:
-            self.trainer._call_lightning_datamodule_hook("on_load_checkpoint", self._loaded_checkpoint)
-            if datamodule.__class__.__qualname__ in self._loaded_checkpoint:
-                self.trainer._call_lightning_datamodule_hook(
-                    "load_state_dict", self._loaded_checkpoint[datamodule.__class__.__qualname__]
-                )
+        if datamodule is not None and datamodule.__class__.__qualname__ in self._loaded_checkpoint:
+            self.trainer._call_lightning_datamodule_hook(
+                "load_state_dict", self._loaded_checkpoint[datamodule.__class__.__qualname__]
+            )
 
     def restore_model(self) -> None:
         """Restores a model's weights from a PyTorch Lightning checkpoint.
@@ -519,9 +517,6 @@ def dump_checkpoint(self, weights_only: bool = False) -> dict:
             # will be removed in v1.8
             self.trainer._call_callbacks_on_save_checkpoint(checkpoint)
         self.trainer._call_lightning_module_hook("on_save_checkpoint", checkpoint)
-        if datamodule is not None:
-            self.trainer._call_lightning_datamodule_hook("on_save_checkpoint", checkpoint)
-
         return checkpoint
 
     def save_checkpoint(

tests/tests_pytorch/core/test_datamodules.py

@@ -202,13 +202,6 @@ def state_dict(self) -> Dict[str, Any]:
         def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
             self.my_state_dict = state_dict
 
-        def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
-            checkpoint[self.__class__.__qualname__].update({"on_save": "update"})
-
-        def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
-            self.checkpoint_state = checkpoint.get(self.__class__.__qualname__).copy()
-            checkpoint[self.__class__.__qualname__].pop("on_save")
-
     reset_seed()
     dm = CustomBoringDataModule()
     model = CustomBoringModel()
@@ -223,21 +216,15 @@ def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
     )
 
     # fit model
-    with pytest.deprecated_call(
-        match="`LightningDataModule.on_save_checkpoint` was deprecated in"
-        " v1.6 and will be removed in v1.8. Use `state_dict` instead."
-    ):
-        trainer.fit(model, datamodule=dm)
-    assert trainer.state.finished, f"Training failed with {trainer.state}"
+    trainer.fit(model, datamodule=dm)
     checkpoint_path = list(trainer.checkpoint_callback.best_k_models.keys())[0]
     checkpoint = torch.load(checkpoint_path)
     assert dm.__class__.__qualname__ in checkpoint
-    assert checkpoint[dm.__class__.__qualname__] == {"my": "state_dict", "on_save": "update"}
+    assert checkpoint[dm.__class__.__qualname__] == {"my": "state_dict"}
 
     for trainer_fn in TrainerFn:
         trainer.state.fn = trainer_fn
         trainer._restore_modules_and_callbacks(checkpoint_path)
-        assert dm.checkpoint_state == {"my": "state_dict", "on_save": "update"}
         assert dm.my_state_dict == {"my": "state_dict"}
 
 
@@ -510,7 +497,6 @@ def state_dict(self):
         "[LightningDataModule]CustomBoringDataModule.prepare_data",
         "[LightningDataModule]CustomBoringDataModule.setup",
         "[LightningDataModule]CustomBoringDataModule.state_dict",
-        "[LightningDataModule]CustomBoringDataModule.on_save_checkpoint",
         "[LightningDataModule]CustomBoringDataModule.teardown",
     ]
     for key in keys:
@@ -527,7 +513,6 @@ def state_dict(self):
     keys = [
         "[LightningDataModule]CustomBoringDataModule.prepare_data",
         "[LightningDataModule]CustomBoringDataModule.setup",
-        "[LightningDataModule]CustomBoringDataModule.on_load_checkpoint",
         "[LightningDataModule]CustomBoringDataModule.load_state_dict",
         "[LightningDataModule]CustomBoringDataModule.teardown",
     ]

tests/tests_pytorch/deprecated_api/test_remove_1-8.py

@@ -13,14 +13,12 @@
 # limitations under the License.
 """Test deprecated functionality which will be removed in v1.8.0."""
 from unittest import mock
-from unittest.mock import Mock
 
 import pytest
 
 from pytorch_lightning import Callback, Trainer
 from pytorch_lightning.callbacks import ModelCheckpoint
-from pytorch_lightning.demos.boring_classes import BoringDataModule, BoringModel
-from pytorch_lightning.trainer.configuration_validator import _check_datamodule_checkpoint_hooks
+from pytorch_lightning.demos.boring_classes import BoringModel
 
 
 def test_v1_8_0_on_init_start_end(tmpdir):
@@ -230,32 +228,6 @@ def on_pretrain_routine_end(self, trainer, pl_module):
         trainer.fit(model)
 
 
-def test_v1_8_0_datamodule_checkpointhooks():
-    class CustomBoringDataModuleSave(BoringDataModule):
-        def on_save_checkpoint(self, checkpoint):
-            print("override on_save_checkpoint")
-
-    class CustomBoringDataModuleLoad(BoringDataModule):
-        def on_load_checkpoint(self, checkpoint):
-            print("override on_load_checkpoint")
-
-    trainer = Mock()
-
-    trainer.datamodule = CustomBoringDataModuleSave()
-    with pytest.deprecated_call(
-        match="`LightningDataModule.on_save_checkpoint` was deprecated in"
-        " v1.6 and will be removed in v1.8. Use `state_dict` instead."
-    ):
-        _check_datamodule_checkpoint_hooks(trainer)
-
-    trainer.datamodule = CustomBoringDataModuleLoad()
-    with pytest.deprecated_call(
-        match="`LightningDataModule.on_load_checkpoint` was deprecated in"
-        " v1.6 and will be removed in v1.8. Use `load_state_dict` instead."
-    ):
-        _check_datamodule_checkpoint_hooks(trainer)
-
-
 def test_deprecated_mc_save_checkpoint():
     mc = ModelCheckpoint()
     trainer = Trainer()

tests/tests_pytorch/models/test_hooks.py

@@ -955,7 +955,6 @@ def test_trainer_datamodule_hook_system(tmpdir):
         dict(name="val_dataloader"),
         dict(name="train_dataloader"),
         dict(name="state_dict"),
-        dict(name="on_save_checkpoint", args=(ANY,)),
         dict(name="teardown", kwargs=dict(stage="fit")),
     ]
     assert called == expected
@@ -1022,13 +1021,10 @@ def state_dict(self):
     }
 
     assert lm_called == [dict(name="on_save_checkpoint", args=(saved_ckpt,))]
-    assert ldm_called == [dict(name="state_dict"), dict(name="on_save_checkpoint", args=(saved_ckpt,))]
+    assert ldm_called == [dict(name="state_dict")]
 
     lm_called, ldm_called = [], []
-    model = HookedModel.load_from_checkpoint(ckpt_path, called=lm_called)
-    datamodule = CustomHookedDataModule.load_from_checkpoint(ckpt_path, called=ldm_called)
+    _ = HookedModel.load_from_checkpoint(ckpt_path, called=lm_called)
+    _ = CustomHookedDataModule.load_from_checkpoint(ckpt_path, called=ldm_called)
     assert lm_called == [dict(name="on_load_checkpoint", args=({**saved_ckpt, "hyper_parameters": ANY},))]
-    assert ldm_called == [
-        dict(name="on_load_checkpoint", args=({**saved_ckpt, "datamodule_hyper_parameters": ANY},)),
-        dict(name="load_state_dict", args=(saved_ckpt[datamodule_state_dict_key],)),
-    ]
+    assert ldm_called == [dict(name="load_state_dict", args=(saved_ckpt[datamodule_state_dict_key],))]