Add model trainer documentation (#1639)

Author: benieric

doc/api/training/index.rst

@@ -5,6 +5,7 @@ Training APIs
 .. toctree::
    :maxdepth: 4
 
+   model_trainer
    algorithm
    analytics
    automl

doc/api/training/model_trainer.rst

@@ -0,0 +1,17 @@
+ModelTrainer
+------------
+
+.. autoclass:: sagemaker.modules.train.model_trainer.ModelTrainer
+    :members:
+
+Configs
+~~~~~~~
+
+.. automodule:: sagemaker.modules.configs
+    :members:
+
+Distributed
+~~~~~~~~~~~
+
+.. automodule:: sagemaker.modules.distributed
+    :members:

doc/overview.rst

@@ -4,6 +4,7 @@ Using the SageMaker Python SDK
 
 SageMaker Python SDK provides several high-level abstractions for working with Amazon SageMaker. These are:
 
+- **ModelTrainer**: New interface encapsulating training on SageMaker.
 - **Estimators**: Encapsulate training on SageMaker.
 - **Models**: Encapsulate built ML models.
 - **Predictors**: Provide real-time inference and transformation using Python data-types against a SageMaker endpoint.
@@ -24,8 +25,8 @@ Train a Model with the SageMaker Python SDK
 To train a model by using the SageMaker Python SDK, you:
 
 1. Prepare a training script
-2. Create an estimator
-3. Call the ``fit`` method of the estimator
+2. Create a ModelTrainer or Estimator
+3. Call the ``train`` method of the ModelTrainer or the ``fit`` method of the Estimator
 
 After you train a model, you can save it, and then serve the model as an endpoint to get real-time inferences or get inferences for an entire dataset by using batch transform.
 
@@ -85,6 +86,46 @@ If you want to use, for example, boolean hyperparameters, you need to specify ``
 For more on training environment variables, please visit `SageMaker Containers <https://github.com/aws/sagemaker-containers>`_.
 
 
+Using ModelTrainer
+==================
+
+To use the ModelTrainer class, you need to provide a few essential parameters such as the training image URI and the source code configuration. The class allows you to spin up a SageMaker training job with minimal parameters, particularly by specifying the source code and training image.
+
+For more information about class definitions see `ModelTrainer <https://sagemaker.readthedocs.io/en/stable/api/training/model_trainer.html>`_.
+
+Example: Launching a Training Job with Custom Script
+
+.. code:: python
+
+    from sagemaker.modules.train import ModelTrainer
+    from sagemaker.modules.configs import SourceCode, InputData
+
+    # Image URI for the training job
+    pytorch_image = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310"
+
+    # Define the script to be run
+    source_code = SourceCode(
+        source_dir="basic-script-mode",
+        requirements="requirements.txt",
+        entry_script="custom_script.py",
+    )
+
+    # Define the ModelTrainer
+    model_trainer = ModelTrainer(
+        training_image=pytorch_image,
+        source_code=source_code,
+        base_job_name="script-mode",
+    )
+
+    # Pass the input data
+    input_data = InputData(
+        channel_name="train",
+        data_source=training_input_path, # S3 path where training data is stored
+    )
+
+    # Start the training job
+    model_trainer.train(input_data_config=[input_data], wait=False)
+
 Using Estimators
 ================
 

src/sagemaker/modules/configs.py

@@ -10,12 +10,12 @@
 # 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.
-"""This module provides the configuration classes used in `sagemaker.modules`.
+"""This module provides the configuration classes used in ``sagemaker.modules``.
 
-Some of these classes are re-exported from `sagemaker_core.shapes`. For convinence,
-users can import these classes directly from `sagemaker.modules.configs`.
+Some of these classes are re-exported from ``sagemaker_core.shapes``. For convinence,
+users can import these classes directly from ``sagemaker.modules.configs``.
 
-For more documentation on `sagemaker_core.shapes`, see:
+For more documentation on ``sagemaker_core.shapes``, see:
     - https://sagemaker-core.readthedocs.io/en/stable/#sagemaker-core-shapes
 """
 
@@ -80,14 +80,14 @@ class SourceCode(BaseModel):
     The SourceCode class allows the user to specify the source code location, dependencies,
     entry script, or commands to be executed in the training job container.
 
-    Attributes:
+    Parameters:
         source_dir (Optional[str]):
             The local directory containing the source code to be used in the training job container.
         requirements (Optional[str]):
-            The path within `source_dir` to a `requirements.txt` file. If specified, the listed
+            The path within ``source_dir`` to a ``requirements.txt`` file. If specified, the listed
             requirements will be installed in the training job container.
         entry_script (Optional[str]):
-            The path within `source_dir` to the entry script that will be executed in the training
+            The path within ``source_dir`` to the entry script that will be executed in the training
             job container. If not specified, command must be provided.
         command (Optional[str]):
             The command(s) to execute in the training job container. Example: "python my_script.py".
@@ -103,10 +103,10 @@ class SourceCode(BaseModel):
 class Compute(shapes.ResourceConfig):
     """Compute.
 
-    The Compute class is a subclass of `sagemaker_core.shapes.ResourceConfig`
+    The Compute class is a subclass of ``sagemaker_core.shapes.ResourceConfig``
     and allows the user to specify the compute resources for the training job.
 
-    Attributes:
+    Parameters:
         instance_type (Optional[str]):
             The ML compute instance type. For information about available instance types,
             see https://aws.amazon.com/sagemaker/pricing/.
@@ -152,10 +152,10 @@ def _to_resource_config(self) -> shapes.ResourceConfig:
 class Networking(shapes.VpcConfig):
     """Networking.
 
-    The Networking class is a subclass of `sagemaker_core.shapes.VpcConfig ` and
+    The Networking class is a subclass of ``sagemaker_core.shapes.VpcConfig`` and
     allows the user to specify the networking configuration for the training job.
 
-    Attributes:
+    Parameters:
         security_group_ids (Optional[List[str]]):
             The VPC security group IDs, in the form sg-xxxxxxxx. Specify the
             security groups for the VPC that is specified in the Subnets field.
@@ -199,15 +199,15 @@ class InputData(BaseModel):
 
     This config allows the user to specify an input data source for the training job.
 
-    Will be found at `/opt/ml/input/data/<channel_name>` within the training container.
+    Will be found at ``/opt/ml/input/data/<channel_name>`` within the training container.
     For convience, can be referenced inside the training container like:
 
-    ```python
-    import os
-    input_data_dir = os.environ['SM_CHANNEL_<channel_name>']
-    ```
+    .. code:: python
 
-    Attributes:
+        import os
+        input_data_dir = os.environ['SM_CHANNEL_<channel_name>']
+
+    Parameters:
         channel_name (str):
             The name of the input data source channel.
         data_source (Union[str, S3DataSource, FileSystemDataSource]):

src/sagemaker/modules/distributed.py

@@ -25,7 +25,7 @@ class SMP(BaseModel):
     For more information on the model parallelism parameters, see:
     https://docs.aws.amazon.com/sagemaker/latest/dg/distributed-model-parallel-v2-reference.html#distributed-model-parallel-v2-reference-init-config
 
-    Attributes:
+    Parameters:
         hybrid_shard_degree (Optional[int]):
             Specifies a sharded parallelism degree for the model.
         sm_activation_offloading (Optional[bool]):
@@ -87,10 +87,10 @@ def model_dump(self, *args, **kwargs):
 class Torchrun(DistributedConfig):
     """Torchrun.
 
-    The Torchrun class configures a job that uses `torchrun` or
-    `torch.distributed.launch` in the backend to launch distributed training.
+    The Torchrun class configures a job that uses ``torchrun`` or
+    ``torch.distributed.launch`` in the backend to launch distributed training.
 
-    Attributes:
+    Parameters:
         process_count_per_node (int):
             The number of processes to run on each node in the training job.
             Will default to the number of GPUs available in the container.
@@ -107,10 +107,10 @@ class Torchrun(DistributedConfig):
 class MPI(DistributedConfig):
     """MPI.
 
-    The MPI class configures a job that uses `mpirun` in the backend to launch
+    The MPI class configures a job that uses ``mpirun`` in the backend to launch
     distributed training.
 
-    Attributes:
+    Parameters:
         process_count_per_node (int):
             The number of processes to run on each node in the training job.
             Will default to the number of GPUs available in the container.

src/sagemaker/modules/train/model_trainer.py

@@ -115,25 +115,31 @@ class ModelTrainer(BaseModel):
     """Class that trains a model using AWS SageMaker.
 
     Example:
-    ```python
-    from sagemaker.modules.train import ModelTrainer
-    from sagemaker.modules.configs import SourceCode, Compute, InputData
-
-    source_code = SourceCode(source_dir="source", entry_script="train.py")
-    training_image = "123456789012.dkr.ecr.us-west-2.amazonaws.com/my-training-image"
-    model_trainer = ModelTrainer(
-        training_image=training_image,
-        source_code=source_code,
-    )
-
-    train_data = InputData(channel_name="train", data_source="s3://bucket/train")
-    model_trainer.train(input_data_config=[train_data])
-    ```
-
-    Attributes:
+
+    .. code:: python
+
+        from sagemaker.modules.train import ModelTrainer
+        from sagemaker.modules.configs import SourceCode, Compute, InputData
+
+        source_code = SourceCode(source_dir="source", entry_script="train.py")
+        training_image = "123456789012.dkr.ecr.us-west-2.amazonaws.com/my-training-image"
+        model_trainer = ModelTrainer(
+            training_image=training_image,
+            source_code=source_code,
+        )
+
+        train_data = InputData(channel_name="train", data_source="s3://bucket/train")
+        model_trainer.train(input_data_config=[train_data])
+
+        training_job = model_trainer._latest_training_job
+
+    Parameters:
+        training_mode (Mode):
+            The training mode. Valid values are "Mode.LOCAL_CONTAINER" or
+            "Mode.SAGEMAKER_TRAINING_JOB".
         sagemaker_session (Optiona(Session)):
             The SageMakerCore session. For convinience, can be imported like:
-            `from sagemaker.modules import Session`.
+            ``from sagemaker.modules import Session``.
             If not specified, a new session will be created.
             If the default bucket for the artifacts needs to be updated, it can be done by
             passing it in the Session object.
@@ -149,7 +155,7 @@ class ModelTrainer(BaseModel):
             running the training job.
         distributed (Optional[Union[MPI, Torchrun]]):
             The distributed runner for the training job. This is used to configure
-            a distributed training job. If specifed, `source_code` must also
+            a distributed training job. If specifed, ``source_code`` must also
             be provided.
         compute (Optional[Compute]):
             The compute configuration. This is used to specify the compute resources for
@@ -176,7 +182,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
@@ -194,6 +200,9 @@ class ModelTrainer(BaseModel):
         tags (Optional[List[Tag]]):
             An array of key-value pairs. You can use tags to categorize your AWS resources
             in different ways, for example, by purpose, owner, or environment.
+        local_container_root (Optional[str]):
+            The local root directory to store artifacts from a training job launched in
+            "LOCAL_CONTAINER" mode.
     """
 
     model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
@@ -652,9 +661,10 @@ def create_input_data_channel(
             key_prefix (Optional[str]): The key prefix to use when uploading data to S3.
                 Only applicable when data_source is a local file path string.
                 If not specified, local data will be uploaded to:
-                    s3://<default_bucket_path>/<base_job_name>/input/<channel_name>/
+                ``s3://<default_bucket_path>/<base_job_name>/input/<channel_name>/``
+
                 If specified, local data will be uploaded to:
-                    s3://<default_bucket_path>/<key_prefix>/<channel_name>/
+                ``s3://<default_bucket_path>/<key_prefix>/<channel_name>/``
         """
         channel = None
         if isinstance(data_source, str):
@@ -881,7 +891,7 @@ def from_recipe(
             output_data_config (Optional[OutputDataConfig]):
                 The output data configuration. This is used to specify the output data location
                 for the training job.
-                If not specified, will default to `s3://<default_bucket>/<base_job_name>/output/`.
+                If not specified, will default to ``s3://<default_bucket>/<base_job_name>/output/``.
             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
@@ -910,7 +920,7 @@ def from_recipe(
         """
         if compute.instance_type is None:
             raise ValueError(
-                "Must set `instance_type` in compute_config when using training recipes."
+                "Must set ``instance_type`` in compute_config when using training recipes."
             )
         device_type = _determine_device_type(compute.instance_type)
         if device_type == "cpu":
@@ -970,7 +980,7 @@ def with_tensorboard_output_config(
         """Set the TensorBoard output configuration.
 
         Args:
-            tensorboard_output_config (TensorBoardOutputConfig):
+            tensorboard_output_config (sagemaker.modules.configs.TensorBoardOutputConfig):
                 The TensorBoard output configuration.
         """
         self._tensorboard_output_config = tensorboard_output_config