update checkpoint config

Author: benieric

src/sagemaker/modules/train/model_trainer.py

@@ -53,6 +53,7 @@
 
 from sagemaker.utils import resolve_value_from_config
 from sagemaker.modules import Session, get_execution_role
+from sagemaker.modules import configs
 from sagemaker.modules.configs import (
     Compute,
     StoppingCondition,
@@ -1132,7 +1133,9 @@ def with_tensorboard_output_config(
             tensorboard_output_config (sagemaker.modules.configs.TensorBoardOutputConfig):
                 The TensorBoard output configuration.
         """
-        self._tensorboard_output_config = tensorboard_output_config or TensorBoardOutputConfig()
+        self._tensorboard_output_config = (
+            tensorboard_output_config or configs.TensorBoardOutputConfig()
+        )
         return self
 
     def with_retry_strategy(self, retry_strategy: RetryStrategy) -> "ModelTrainer":  # noqa: D412
@@ -1227,3 +1230,25 @@ def with_remote_debug_config(
             enable_remote_debug=True
         )
         return self
+
+    def with_checkpoint_config(
+        self, checkpoint_config: Optional[CheckpointConfig] = None
+    ) -> "ModelTrainer":
+        """Set the checkpoint configuration for the training job.
+
+        Example:
+
+        .. code:: python
+
+            from sagemaker.modules.train import ModelTrainer
+
+            model_trainer = ModelTrainer(
+                ...
+            ).with_checkpoint_config()
+
+        Args:
+            checkpoint_config (sagemaker.modules.configs.CheckpointConfig):
+                The checkpoint configuration for the training job.
+        """
+        self.checkpoint_config = checkpoint_config or configs.CheckpointConfig()
+        return self

tests/unit/sagemaker/modules/train/test_model_trainer.py

@@ -1242,12 +1242,16 @@ def mock_upload_data(path, bucket, key_prefix):
     modules_session.upload_data.side_effect = mock_upload_data
     mock_unique_name.return_value = unique_name
 
-    model_trainer = ModelTrainer(
-        training_image=DEFAULT_IMAGE,
-        sagemaker_session=modules_session,
-        checkpoint_config=CheckpointConfig(),
-        base_job_name=base_name,
-    ).with_tensorboard_output_config(TensorBoardOutputConfig())
+    model_trainer = (
+        ModelTrainer(
+            training_image=DEFAULT_IMAGE,
+            sagemaker_session=modules_session,
+            base_job_name=base_name,
+        )
+        .with_tensorboard_output_config()
+        .with_checkpoint_config()
+    )
+
     model_trainer.train()
 
     _, kwargs = mock_training_job.create.call_args