Revert "remove save_on_train_epoch_end from state_key"

This reverts commit 8ad08ca.

Author: awaelchli

pytorch_lightning/callbacks/model_checkpoint.py

@@ -197,7 +197,7 @@ class ModelCheckpoint(Callback):
     .. tip:: Saving and restoring multiple checkpoint callbacks at the same time is supported under variation in the
         following arguments:
 
-        *monitor, mode, every_n_train_steps, every_n_epochs, train_time_interval*
+        *monitor, mode, every_n_train_steps, every_n_epochs, train_time_interval, save_on_train_epoch_end*
 
         Read more: :ref:`Persisting Callback State`
     """
@@ -262,6 +262,7 @@ def state_key(self) -> str:
             every_n_train_steps=self._every_n_train_steps,
             every_n_epochs=self._every_n_epochs,
             train_time_interval=self._train_time_interval,
+            save_on_train_epoch_end=self._save_on_train_epoch_end,
         )
 
     def on_pretrain_routine_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:

tests/checkpointing/test_model_checkpoint.py

@@ -47,7 +47,7 @@ def test_model_checkpoint_state_key():
     early_stopping = ModelCheckpoint(monitor="val_loss")
     expected_id = (
         "ModelCheckpoint{'monitor': 'val_loss', 'mode': 'min', 'every_n_train_steps': 0, 'every_n_epochs': 1,"
-        " 'train_time_interval': None}"
+        " 'train_time_interval': None, 'save_on_train_epoch_end': None}"
     )
     assert early_stopping.state_key == expected_id
 
@@ -159,7 +159,7 @@ def on_validation_epoch_end(self):
 
         mc_specific_data = chk["callbacks"][
             f"ModelCheckpoint{{'monitor': '{monitor}', 'mode': 'min', 'every_n_train_steps': 0, 'every_n_epochs': 1,"
-            " 'train_time_interval': None}"
+            " 'train_time_interval': None, 'save_on_train_epoch_end': True}"
         ]
         assert mc_specific_data["dirpath"] == checkpoint.dirpath
         assert mc_specific_data["monitor"] == monitor
@@ -273,7 +273,7 @@ def _make_assertions(epoch, ix, version=""):
 
         mc_specific_data = chk["callbacks"][
             f"ModelCheckpoint{{'monitor': '{monitor}', 'mode': 'min', 'every_n_train_steps': 0, 'every_n_epochs': 1,"
-            " 'train_time_interval': None}"
+            " 'train_time_interval': None, 'save_on_train_epoch_end': False}"
         ]
         assert mc_specific_data["dirpath"] == checkpoint.dirpath
         assert mc_specific_data["monitor"] == monitor
@@ -875,7 +875,7 @@ def test_model_checkpoint_save_last_checkpoint_contents(tmpdir):
 
     ckpt_id = (
         "ModelCheckpoint{'monitor': 'early_stop_on', 'mode': 'min', 'every_n_train_steps': 0, 'every_n_epochs': 1,"
-        " 'train_time_interval': None}"
+        " 'train_time_interval': None, 'save_on_train_epoch_end': True}"
     )
     assert ckpt_last["callbacks"][ckpt_id] == ckpt_last_epoch["callbacks"][ckpt_id]
 
@@ -1118,7 +1118,7 @@ def training_step(self, *args):
     ckpts = [
         ckpt["callbacks"][
             "ModelCheckpoint{'monitor': 'foo', 'mode': 'min', 'every_n_train_steps': 0, 'every_n_epochs': 1,"
-            " 'train_time_interval': None}"
+            " 'train_time_interval': None, 'save_on_train_epoch_end': True}"
         ]
         for ckpt in ckpts
     ]

tests/trainer/connectors/test_callback_connector.py

@@ -111,7 +111,7 @@ def test_all_callback_states_saved_before_checkpoint_callback(tmpdir):
     assert "content1" in state2 and state2["content1"] == "two"
     assert (
         "ModelCheckpoint{'monitor': None, 'mode': 'min', 'every_n_train_steps': 0, 'every_n_epochs': 1,"
-        " 'train_time_interval': None}" in ckpt["callbacks"]
+        " 'train_time_interval': None, 'save_on_train_epoch_end': True}" in ckpt["callbacks"]
     )