Update save checkpoint on exception tests to use a shorter more precisly defined epoch lenght

Author: vsey

tests/tests_pytorch/checkpointing/test_model_checkpoint.py

@@ -864,12 +864,13 @@ def on_train_epoch_start(self, trainer, pl_module):
                 raise RuntimeError("Trouble!")
 
     model = BoringModel()
-    epoch_length = 64
+    epoch_length = 2
     checkpoint_callback = ModelCheckpoint(dirpath=tmp_path, filename="{step}", save_on_exception=True, every_n_epochs=4)
     trainer = Trainer(
         default_root_dir=tmp_path,
         callbacks=[checkpoint_callback, TroublemakerOnTrainEpochStart()],
         max_epochs=5,
+        limit_train_batches=epoch_length,
         logger=False,
         enable_progress_bar=False,
     )
@@ -887,12 +888,13 @@ def on_train_epoch_end(self, trainer, pl_module):
                 raise RuntimeError("Trouble!")
 
     model = BoringModel()
-    epoch_length = 64
+    epoch_length = 2
     checkpoint_callback = ModelCheckpoint(dirpath=tmp_path, filename="{step}", save_on_exception=True, every_n_epochs=4)
     trainer = Trainer(
         default_root_dir=tmp_path,
         callbacks=[checkpoint_callback, TroublemakerOnTrainEpochEnd()],
         max_epochs=5,
+        limit_train_batches=epoch_length,
         logger=False,
         enable_progress_bar=False,
     )
@@ -956,12 +958,13 @@ def on_validation_epoch_start(self, trainer, pl_module):
                 raise RuntimeError("Trouble!")
 
     model = BoringModel()
-    epoch_length = 64
+    epoch_length = 2
     checkpoint_callback = ModelCheckpoint(dirpath=tmp_path, filename="{step}", save_on_exception=True, every_n_epochs=4)
     trainer = Trainer(
         default_root_dir=tmp_path,
         callbacks=[checkpoint_callback, TroublemakerOnValidationEpochStart()],
         max_epochs=5,
+        limit_train_batches=epoch_length,
         logger=False,
         enable_progress_bar=False,
     )
@@ -979,12 +982,13 @@ def on_validation_epoch_end(self, trainer, pl_module):
                 raise RuntimeError("Trouble!")
 
     model = BoringModel()
-    epoch_length = 64
+    epoch_length = 2
     checkpoint_callback = ModelCheckpoint(dirpath=tmp_path, filename="{step}", save_on_exception=True, every_n_epochs=4)
     trainer = Trainer(
         default_root_dir=tmp_path,
         callbacks=[checkpoint_callback, TroublemakerOnValidationEpochEnd()],
         max_epochs=5,
+        limit_train_batches=epoch_length,
         logger=False,
         enable_progress_bar=False,
     )
@@ -1002,12 +1006,13 @@ def on_validation_start(self, trainer, pl_module):
                 raise RuntimeError("Trouble!")
 
     model = BoringModel()
-    epoch_length = 64
+    epoch_length = 2
     checkpoint_callback = ModelCheckpoint(dirpath=tmp_path, filename="{step}", save_on_exception=True, every_n_epochs=4)
     trainer = Trainer(
         default_root_dir=tmp_path,
         callbacks=[checkpoint_callback, TroublemakerOnValidationStart()],
         max_epochs=5,
+        limit_train_batches=epoch_length,
         logger=False,
         enable_progress_bar=False,
     )
@@ -1025,12 +1030,13 @@ def on_validation_end(self, trainer, pl_module):
                 raise RuntimeError("Trouble!")
 
     model = BoringModel()
-    epoch_length = 64
+    epoch_length = 2
     checkpoint_callback = ModelCheckpoint(dirpath=tmp_path, filename="{step}", save_on_exception=True, every_n_epochs=4)
     trainer = Trainer(
         default_root_dir=tmp_path,
         callbacks=[checkpoint_callback, TroublemakerOnValidationEnd()],
         max_epochs=5,
+        limit_train_batches=epoch_length,
         logger=False,
         enable_progress_bar=False,
     )