add tests & update docs

Author: benieric

pyproject.toml

@@ -55,7 +55,8 @@ dependencies = [
   "tblib>=1.7.0,<4",
   "tqdm",
   "urllib3>=1.26.8,<3.0.0",
-  "uvicorn"
+  "uvicorn",
+  "graphene>=3,<4"
 ]
 
 [project.scripts]

src/sagemaker/modules/configs.py

@@ -258,8 +258,9 @@ class TensorBoardOutputConfig(shapes.TensorBoardOutputConfig):
 
     Parameters:
         s3_output_path (Optional[str]):
-            Path to Amazon S3 storage location for TensorBoard output. If not specified, will default to the default artifact location for the training job.
-            ``s3://<default_bucket>/<default_prefix>/<base_job_name>/<job_name>/``
+            Path to Amazon S3 storage location for TensorBoard output. If not specified, will
+            default to
+            ``s3://<default_bucket>/<default_prefix>/<base_job_name>/<job_name>/tensorboard-output``
         local_path (Optional[str]):
             Path to local storage location for tensorBoard output. Defaults to /opt/ml/output/tensorboard.
     """
@@ -276,8 +277,9 @@ class CheckpointConfig(shapes.CheckpointConfig):
 
     Parameters:
         s3_uri (Optional[str]):
-            Path to Amazon S3 storage location for the Checkpoint data. If not specified, will default to the default artifact location for the training job.
-            ``s3://<default_bucket>/<default_prefix>/<base_job_name>/<job_name>/``
+            Path to Amazon S3 storage location for the Checkpoint data. If not specified, will
+            default to
+            ``s3://<default_bucket>/<default_prefix>/<base_job_name>/<job_name>/checkpoints``
         local_path (Optional[str]):
             The local directory where checkpoints are written. The default directory is /opt/ml/checkpoints.
     """

src/sagemaker/modules/train/model_trainer.py

@@ -181,7 +181,7 @@ class ModelTrainer(BaseModel):
             The output data configuration. This is used to specify the output data location
             for the training job.
             If not specified in the session, will default to
-            s3://<default_bucket>/<default_prefix>/<base_job_name>/.
+            ``s3://<default_bucket>/<default_prefix>/<base_job_name>/``.
         input_data_config (Optional[List[Union[Channel, InputData]]]):
             The input data config for the training job.
             Takes a list of Channel or InputData objects. An InputDataSource can be an S3 URI
@@ -595,8 +595,7 @@ def train(
         """
         self._populate_intelligent_defaults()
         current_training_job_name = _get_unique_name(self.base_job_name)
-        default_artifact_path = f"{self.base_job_name}/{current_training_job_name}"
-        input_data_key_prefix = f"{default_artifact_path}/input"
+        input_data_key_prefix = f"{self.base_job_name}/{current_training_job_name}/input"
         if input_data_config and self.input_data_config:
             self.input_data_config = input_data_config
             # Add missing input data channels to the existing input_data_config
@@ -613,9 +612,15 @@ def train(
             )
 
         if self.checkpoint_config and not self.checkpoint_config.s3_uri:
-            self.checkpoint_config.s3_uri = f"s3://{self._fetch_bucket_name_and_prefix(self.sagemaker_session)}/{default_artifact_path}"
-        if self._tensorboard_output_config and not self._tensorboard_output_config.s3_uri:
-            self._tensorboard_output_config.s3_uri = f"s3://{self._fetch_bucket_name_and_prefix(self.sagemaker_session)}/{default_artifact_path}"
+            self.checkpoint_config.s3_uri = (
+                f"s3://{self._fetch_bucket_name_and_prefix(self.sagemaker_session)}/"
+                f"{self.base_job_name}/{current_training_job_name}/checkpoints"
+            )
+        if self._tensorboard_output_config and not self._tensorboard_output_config.s3_output_path:
+            self._tensorboard_output_config.s3_output_path = (
+                f"s3://{self._fetch_bucket_name_and_prefix(self.sagemaker_session)}/"
+                f"{self.base_job_name}"
+            )
 
         string_hyper_parameters = {}
         if self.hyperparameters:
@@ -646,7 +651,7 @@ def train(
                     data_source=self.source_code.source_dir,
                     key_prefix=input_data_key_prefix,
                 )
-                input_data_config.append(source_code_channel)
+                self.input_data_config.append(source_code_channel)
 
             self._prepare_train_script(
                 tmp_dir=tmp_dir,
@@ -667,7 +672,7 @@ def train(
                 data_source=tmp_dir.name,
                 key_prefix=input_data_key_prefix,
             )
-            input_data_config.append(sm_drivers_channel)
+            self.input_data_config.append(sm_drivers_channel)
 
             # If source_code is provided, we will always use
             # the default container entrypoint and arguments
@@ -970,10 +975,43 @@ def from_recipe(
     ) -> "ModelTrainer":
         """Create a ModelTrainer from a training recipe.
 
+        Example:
+
+        .. code:: python
+
+            from sagemaker.modules.train import ModelTrainer
+            from sagemaker.modules.configs import Compute
+
+            recipe_overrides = {
+                "run": {
+                    "results_dir": "/opt/ml/model",
+                },
+                "model": {
+                    "data": {
+                        "use_synthetic_data": True
+                    }
+                }
+            }
+
+            compute = Compute(
+                instance_type="ml.p5.48xlarge",
+                keep_alive_period_in_seconds=3600
+            )
+
+            model_trainer = ModelTrainer.from_recipe(
+                training_recipe="fine-tuning/deepseek/hf_deepseek_r1_distilled_llama_8b_seq8k_gpu_fine_tuning",
+                recipe_overrides=recipe_overrides,
+                compute=compute,
+            )
+
+            model_trainer.train(wait=False)
+
+
         Args:
             training_recipe (str):
                 The training recipe to use for training the model. This must be the name of
                 a sagemaker training recipe or a path to a local training recipe .yaml file.
+                For available training recipes, see: https://github.com/aws/sagemaker-hyperpod-recipes/
             compute (Compute):
                 The compute configuration. This is used to specify the compute resources for
                 the training job. If not specified, will default to 1 instance of ml.m5.xlarge.
@@ -1081,55 +1119,116 @@ def from_recipe(
         return model_trainer
 
     def with_tensorboard_output_config(
-        self, tensorboard_output_config: TensorBoardOutputConfig
+        self, tensorboard_output_config: Optional[TensorBoardOutputConfig] = None
     ) -> "ModelTrainer":
         """Set the TensorBoard output configuration.
 
+        Example:
+
+        .. code:: python
+
+            from sagemaker.modules.train import ModelTrainer
+
+            model_trainer = ModelTrainer(
+                ...
+            ).with_tensorboard_output_config()
+
         Args:
             tensorboard_output_config (sagemaker.modules.configs.TensorBoardOutputConfig):
                 The TensorBoard output configuration.
         """
-        self._tensorboard_output_config = tensorboard_output_config
+        self._tensorboard_output_config = tensorboard_output_config or TensorBoardOutputConfig()
         return self
 
     def with_retry_strategy(self, retry_strategy: RetryStrategy) -> "ModelTrainer":
         """Set the retry strategy for the training job.
 
+        Example:
+
+        .. code:: python
+
+            from sagemaker.modules.train import ModelTrainer
+            from sagemaker.modules.configs import RetryStrategy
+
+            retry_strategy = RetryStrategy(maximum_retry_attempts=3)
+
+            model_trainer = ModelTrainer(
+                ...
+            ).with_retry_strategy(retry_strategy)
+
         Args:
-            retry_strategy (RetryStrategy):
+            retry_strategy (sagemaker.modules.configs.RetryStrategy):
                 The retry strategy for the training job.
         """
         self._retry_strategy = retry_strategy
         return self
 
-    def with_infra_check_config(self, infra_check_config: InfraCheckConfig) -> "ModelTrainer":
+    def with_infra_check_config(
+        self, infra_check_config: Optional[InfraCheckConfig] = None
+    ) -> "ModelTrainer":
         """Set the infra check configuration for the training job.
 
+        Example:
+
+        .. code:: python
+
+            from sagemaker.modules.train import ModelTrainer
+
+            model_trainer = ModelTrainer(
+                ...
+            ).with_infra_check_config()
+
         Args:
-            infra_check_config (InfraCheckConfig):
+            infra_check_config (sagemaker.modules.configs.InfraCheckConfig):
                 The infra check configuration for the training job.
         """
-        self._infra_check_config = infra_check_config
+        self._infra_check_config = infra_check_config or InfraCheckConfig(enable_infra_check=True)
         return self
 
     def with_session_chaining_config(
-        self, session_chaining_config: SessionChainingConfig
+        self, session_chaining_config: Optional[SessionChainingConfig] = None
     ) -> "ModelTrainer":
         """Set the session chaining configuration for the training job.
 
+        Example:
+
+        .. code:: python
+
+            from sagemaker.modules.train import ModelTrainer
+
+            model_trainer = ModelTrainer(
+                ...
+            ).with_session_chaining_config()
+
         Args:
-            session_chaining_config (SessionChainingConfig):
+            session_chaining_config (sagemaker.modules.configs.SessionChainingConfig):
                 The session chaining configuration for the training job.
         """
-        self._session_chaining_config = session_chaining_config
+        self._session_chaining_config = session_chaining_config or SessionChainingConfig(
+            enable_session_tag_chaining=True
+        )
         return self
 
-    def with_remote_debug_config(self, remote_debug_config: RemoteDebugConfig) -> "ModelTrainer":
+    def with_remote_debug_config(
+        self, remote_debug_config: Optional[RemoteDebugConfig] = None
+    ) -> "ModelTrainer":
         """Set the remote debug configuration for the training job.
 
+        Example:
+
+        .. code:: python
+
+            from sagemaker.modules.train import ModelTrainer
+
+            model_trainer = ModelTrainer(
+                ...
+            ).with_remote_debug_config()
+
         Args:
-            remote_debug_config (RemoteDebugConfig):
+            remote_debug_config (sagemaker.modules.configs.RemoteDebugConfig):
                 The remote debug configuration for the training job.
         """
-        self._remote_debug_config = remote_debug_config
+        self._remote_debug_config = remote_debug_config or RemoteDebugConfig(
+            enable_remote_debug=True
+        )
         return self

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

@@ -1228,3 +1228,37 @@ def test_hyperparameters_invalid(mock_exists, modules_session):
                 compute=DEFAULT_COMPUTE_CONFIG,
                 hyperparameters="hyperparameters.yaml",
             )
+
+
+@patch("sagemaker.modules.train.model_trainer._get_unique_name")
+@patch("sagemaker.modules.train.model_trainer.TrainingJob")
+def test_model_trainer_default_paths(mock_training_job, mock_unique_name, modules_session):
+    def mock_upload_data(path, bucket, key_prefix):
+        return f"s3://{bucket}/{key_prefix}"
+
+    unique_name = "base-job-0123456789"
+    base_name = "base-job"
+
+    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.train()
+
+    _, kwargs = mock_training_job.create.call_args
+
+    default_base_path = f"s3://{DEFAULT_BUCKET}/{DEFAULT_BUCKET_PREFIX}/{base_name}"
+
+    assert kwargs["output_data_config"].s3_output_path == default_base_path
+    assert kwargs["output_data_config"].compression_type == "GZIP"
+
+    assert kwargs["checkpoint_config"].s3_uri == f"{default_base_path}/{unique_name}/checkpoints"
+    assert kwargs["checkpoint_config"].local_path == "/opt/ml/checkpoints"
+
+    assert kwargs["tensor_board_output_config"].s3_output_path == default_base_path
+    assert kwargs["tensor_board_output_config"].local_path == "/opt/ml/output/tensorboard"