Error messages for removed DataModule hooks

src/pytorch_lightning/_graveyard/__init__.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import pytorch_lightning._graveyard.callbacks
+import pytorch_lightning._graveyard.core
 import pytorch_lightning._graveyard.loggers
 import pytorch_lightning._graveyard.trainer
 import pytorch_lightning._graveyard.training_type  # noqa: F401

src/pytorch_lightning/_graveyard/core.py

@@ -0,0 +1,37 @@
+# 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 typing import Any
+
+from pytorch_lightning import LightningDataModule
+
+
+def _on_save_checkpoint(_: LightningDataModule, __: Any) -> None:
+    # TODO: Remove in v2.0.0
+    raise NotImplementedError(
+        "`LightningDataModule.on_save_checkpoint` was deprecated in v1.6 and removed in v1.8."
+        " Use `state_dict` instead."
+    )
+
+
+def _on_load_checkpoint(_: LightningDataModule, __: Any) -> None:
+    # TODO: Remove in v2.0.0
+    raise NotImplementedError(
+        "`LightningDataModule.on_load_checkpoint` was deprecated in v1.6 and removed in v1.8."
+        " Use `load_state_dict` instead."
+    )
+
+
+# Methods
+LightningDataModule.on_save_checkpoint = _on_save_checkpoint
+LightningDataModule.on_load_checkpoint = _on_load_checkpoint

src/pytorch_lightning/_graveyard/trainer.py

@@ -40,154 +40,154 @@ def __init__(self) -> None:
 
 
 def _gpus(_: Trainer) -> None:
-    # Remove in v2.0.0
+    # TODO: Remove in v2.0.0
     raise AttributeError(
         "`Trainer.gpus` was deprecated in v1.6 and is no longer accessible as of v1.8."
         " Please use `Trainer.num_devices` or `Trainer.device_ids` to get device information instead."
     )
 
 
 def _root_gpu(_: Trainer) -> None:
-    # Remove in v2.0.0
+    # TODO: Remove in v2.0.0
     raise AttributeError(
         "`Trainer.root_gpu` was deprecated in v1.6 and is no longer accessible as of v1.8."
         " Please use `Trainer.strategy.root_device.index` instead."
     )
 
 
 def _tpu_cores(_: Trainer) -> None:
-    # Remove in v2.0.0
+    # TODO: Remove in v2.0.0
     raise AttributeError(
         "`Trainer.tpu_cores` was deprecated in v1.6 and is no longer accessible as of v1.8."
         " Please use `Trainer.num_devices` instead."
     )
 
 
 def _ipus(_: Trainer) -> None:
-    # Remove in v2.0.0
+    # TODO: Remove in v2.0.0
     raise AttributeError(
         "`Trainer.ipus` was deprecated in v1.6 and is no longer accessible as of v1.8."
         " Please use `Trainer.num_devices` instead."
     )
 
 
 def _num_gpus(_: Trainer) -> None:
-    # Remove in v2.0.0
+    # TODO: Remove in v2.0.0
     raise AttributeError(
         "`Trainer.num_gpus` was deprecated in v1.6 and is no longer accessible as of v1.8."
         " Please use `Trainer.num_devices` instead."
     )
 
 
 def _devices(_: Trainer) -> None:
-    # Remove in v2.0.0
+    # TODO: Remove in v2.0.0
     raise AttributeError(
         "`Trainer.devices` was deprecated in v1.6 and is no longer accessible as of v1.8."
         " Please use `Trainer.num_devices` or `Trainer.device_ids` to get device information instead."
     )
 
 
 def _use_amp(_: Trainer) -> None:
-    # Remove in v2.0.0
+    # TODO: Remove in v2.0.0
     raise AttributeError(
         "`Trainer.use_amp` was deprecated in v1.6 and is no longer accessible as of v1.8."
         " Please use `Trainer.amp_backend`.",
     )
 
 
 def _weights_save_path(_: Trainer) -> None:
-    # Remove in v2.0.0
+    # TODO: Remove in v2.0.0
     raise AttributeError("`Trainer.weights_save_path` was deprecated in v1.6 and is no longer accessible as of v1.8.")
 
 
 def _lightning_optimizers(_: Trainer) -> None:
-    # Remove in v2.0.0
+    # TODO: Remove in v2.0.0
     raise AttributeError(
         "`Trainer.lightning_optimizers` was deprecated in v1.6 and is no longer accessible as of v1.8."
     )
 
 
 def _should_rank_save_checkpoint(_: Trainer) -> None:
-    # Remove in v2.0.0
+    # TODO: Remove in v2.0.0
     raise AttributeError(
         "`Trainer.should_rank_save_checkpoint` was deprecated in v1.6 and is no longer accessible as of v1.8.",
     )
 
 
 def _validated_ckpt_path(_: Trainer) -> None:
-    # Remove in v2.0.0
+    # TODO: Remove in v2.0.0
     raise AttributeError(
         "The `Trainer.validated_ckpt_path` was deprecated in v1.6 and is no longer accessible as of v1.8."
         " Please use `Trainer.ckpt_path` instead."
     )
 
 
 def _validated_ckpt_path_setter(_: Trainer, __: Optional[str]) -> None:
-    # Remove in v2.0.0
+    # TODO: Remove in v2.0.0
     raise AttributeError(
         "The `Trainer.validated_ckpt_path` was deprecated in v1.6 and is no longer accessible as of v1.8."
         " Please use `Trainer.ckpt_path` instead."
     )
 
 
 def _tested_ckpt_path(_: Trainer) -> None:
-    # Remove in v2.0.0
+    # TODO: Remove in v2.0.0
     raise AttributeError(
         "The `Trainer.tested_ckpt_path` was deprecated in v1.6 and is no longer accessible as of v1.8."
         " Please use `Trainer.ckpt_path` instead."
     )
 
 
 def _tested_ckpt_path_setter(_: Trainer, __: Optional[str]) -> None:
-    # Remove in v2.0.0
+    # TODO: Remove in v2.0.0
     raise AttributeError(
         "The `Trainer.tested_ckpt_path` was deprecated in v1.6 and is no longer accessible as of v1.8."
         " Please use `Trainer.ckpt_path` instead."
     )
 
 
 def _predicted_ckpt_path(_: Trainer) -> None:
-    # Remove in v2.0.0
+    # TODO: Remove in v2.0.0
     raise AttributeError(
         "The `Trainer.predicted_ckpt_path` was deprecated in v1.6 and is no longer accessible as of v1.8."
         " Please use `Trainer.ckpt_path` instead."
     )
 
 
 def _predicted_ckpt_path_setter(_: Trainer, __: Optional[str]) -> None:
-    # Remove in v2.0.0
+    # TODO: Remove in v2.0.0
     raise AttributeError(
         "The `Trainer.predicted_ckpt_path` was deprecated in v1.6 and is no longer accessible as of v1.8."
         " Please use `Trainer.ckpt_path` instead."
     )
 
 
 def _verbose_evaluate(_: Trainer) -> None:
-    # Remove in v2.0.0
+    # TODO: Remove in v2.0.0
     raise AttributeError(
         "The `Trainer.verbose_evaluate` was deprecated in v1.6 and is no longer accessible as of v1.8."
         " Please use `trainer.{validate,test}_loop.verbose` instead.",
     )
 
 
 def _verbose_evaluate_setter(_: Trainer, __: bool) -> None:
-    # Remove in v2.0.0
+    # TODO: Remove in v2.0.0
     raise AttributeError(
         "The `Trainer.verbose_evaluate` was deprecated in v1.6 and is no longer accessible as of v1.8."
         " Please use `trainer.{validate,test}_loop.verbose` instead.",
     )
 
 
 def _run_stage(_: Trainer) -> None:
-    # Remove in v2.0.0
+    # TODO: Remove in v2.0.0
     raise NotImplementedError(
         "`Trainer.run_stage` was deprecated in v1.6 and is no longer supported as of v1.8."
         " Please use `Trainer.{fit,validate,test,predict}` instead."
     )
 
 
 def _call_hook(_: Trainer, *__: Any, **___: Any) -> Any:
-    # Remove in v2.0.0
+    # TODO: Remove in v2.0.0
     raise NotImplementedError("`Trainer.call_hook` was deprecated in v1.6 and is no longer supported as of v1.8.")
 
 

src/pytorch_lightning/trainer/configuration_validator.py

@@ -57,6 +57,8 @@ def verify_loop_configurations(trainer: "pl.Trainer") -> None:
     _check_on_pretrain_routine(model)
     # TODO: Delete this check in v2.0
     _check_deprecated_logger_methods(trainer)
+    # TODO: Delete this check in v2.0
+    _check_unsupported_datamodule_hooks(trainer)
 
 
 def __verify_train_val_loop_configuration(trainer: "pl.Trainer", model: "pl.LightningModule") -> None:
@@ -278,3 +280,19 @@ def _check_deprecated_logger_methods(trainer: "pl.Trainer") -> None:
                 f"`{type(logger).__name__}.agg_and_log_metrics` was deprecated in v1.6 and is no longer supported as of"
                 " v1.8."
             )
+
+
+def _check_unsupported_datamodule_hooks(trainer: "pl.Trainer") -> None:
+    datahook_selector = trainer._data_connector._datahook_selector
+    assert datahook_selector is not None
+
+    if is_overridden("on_save_checkpoint", datahook_selector.datamodule):
+        raise NotImplementedError(
+            "`LightningDataModule.on_save_checkpoint` was deprecated in v1.6 and removed in v1.8."
+            " Use `state_dict` instead."
+        )
+    if is_overridden("on_load_checkpoint", datahook_selector.datamodule):
+        raise NotImplementedError(
+            "`LightningDataModule.on_load_checkpoint` was deprecated in v1.6 and removed in v1.8."
+            " Use `load_state_dict` instead."
+        )

tests/tests_pytorch/graveyard/test_core.py

@@ -0,0 +1,55 @@
+# 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.
+import pytest
+
+from pytorch_lightning import Trainer
+from pytorch_lightning.demos.boring_classes import BoringDataModule, BoringModel
+
+
+def test_v2_0_0_unsupported_datamodule_on_save_load_checkpoint():
+    datamodule = BoringDataModule()
+    with pytest.raises(
+        NotImplementedError,
+        match="`LightningDataModule.on_save_checkpoint` was deprecated in v1.6 and removed in v1.8.",
+    ):
+        datamodule.on_save_checkpoint({})
+
+    with pytest.raises(
+        NotImplementedError,
+        match="`LightningDataModule.on_load_checkpoint` was deprecated in v1.6 and removed in v1.8.",
+    ):
+        datamodule.on_load_checkpoint({})
+
+    class OnSaveDataModule(BoringDataModule):
+        def on_save_checkpoint(self, checkpoint):
+            pass
+
+    class OnLoadDataModule(BoringDataModule):
+        def on_load_checkpoint(self, checkpoint):
+            pass
+
+    trainer = Trainer()
+    model = BoringModel()
+
+    with pytest.raises(
+        NotImplementedError,
+        match="`LightningDataModule.on_save_checkpoint` was deprecated in v1.6 and removed in v1.8.",
+    ):
+        trainer.fit(model, OnSaveDataModule())
+
+    with pytest.raises(
+        NotImplementedError,
+        match="`LightningDataModule.on_load_checkpoint` was deprecated in v1.6 and removed in v1.8.",
+    ):
+        trainer.fit(model, OnLoadDataModule())