Add path to set Additional Settings in ModelTrainer (#1555)

Author: benieric

src/sagemaker/modules/configs.py

@@ -12,43 +12,84 @@
 # language governing permissions and limitations under the License.
 """This module provides the configuration classes used in `sagemaker.modules`.
 
-Some of these classes are re-exported from `sagemaker-core.shapes`. For convinence,
+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
 """
 
 from __future__ import absolute_import
 
-from typing import Optional, Dict, Any, List
+from typing import Optional, Union, Dict, Any, List
 from pydantic import BaseModel, model_validator
 
+import sagemaker_core.shapes as shapes
+
+# TODO: Can we add custom logic to some of these to set better defaults?
 from sagemaker_core.shapes import (
-    ResourceConfig,
     StoppingCondition,
+    RetryStrategy,
     OutputDataConfig,
     Channel,
+    ShuffleConfig,
     DataSource,
     S3DataSource,
     FileSystemDataSource,
     TrainingImageConfig,
-    VpcConfig,
+    TrainingRepositoryAuthConfig,
+    Tag,
+    MetricDefinition,
+    DebugHookConfig,
+    CollectionConfiguration,
+    DebugRuleConfiguration,
+    ExperimentConfig,
+    InfraCheckConfig,
+    ProfilerConfig,
+    ProfilerRuleConfiguration,
+    RemoteDebugConfig,
+    SessionChainingConfig,
+    InstanceGroup,
+    TensorBoardOutputConfig,
+    CheckpointConfig,
 )
 
 from sagemaker.modules import logger
+from sagemaker.modules.utils import convert_unassigned_to_none
 
 __all__ = [
     "SourceCodeConfig",
-    "ResourceConfig",
+    "TorchDistributionConfig",
+    "MPIDistributionConfig",
+    "SMDistributedSettings",
+    "DistributionConfig",
     "StoppingCondition",
+    "RetryStrategy",
     "OutputDataConfig",
     "Channel",
+    "ShuffleConfig",
     "DataSource",
     "S3DataSource",
     "FileSystemDataSource",
     "TrainingImageConfig",
-    "VpcConfig",
+    "TrainingRepositoryAuthConfig",
+    "Tag",
+    "MetricDefinition",
+    "DebugHookConfig",
+    "CollectionConfiguration",
+    "DebugRuleConfiguration",
+    "ExperimentConfig",
+    "InfraCheckConfig",
+    "ProfilerConfig",
+    "ProfilerRuleConfiguration",
+    "RemoteDebugConfig",
+    "SessionChainingConfig",
+    "InstanceGroup",
+    "TensorBoardOutputConfig",
+    "CheckpointConfig",
+    "ComputeConfig",
+    "NetworkingConfig",
+    "InputData",
 ]
 
 
@@ -161,14 +202,128 @@ class SourceCodeConfig(BaseModel):
         command (Optional[str]):
             The command(s) to execute in the training job container. Example: "python my_script.py".
             If not specified, entry_script must be provided.
-        distribution (Optional[Union[
-            MPIDistributionConfig,
-            TorchDistributionConfig,
-        ]]):
-            The distribution configuration for the training job.
     """
 
     source_dir: Optional[str] = None
     requirements: Optional[str] = None
     entry_script: Optional[str] = None
     command: Optional[str] = None
+
+
+class ComputeConfig(shapes.ResourceConfig):
+    """ComputeConfig.
+
+    The ComputeConfig is a subclass of `sagemaker_core.shapes.ResourceConfig`
+    and allows the user to specify the compute resources for the training job.
+
+    Attributes:
+        instance_type (Optional[str]):
+            The ML compute instance type. For information about available instance types,
+            see https://aws.amazon.com/sagemaker/pricing/. Default: ml.m5.xlarge
+        instance_count (Optional[int]): The number of ML compute instances to use. For distributed
+            training, provide a value greater than 1. Default: 1
+        volume_size_in_gb (Optional[int]):
+            The size of the ML storage volume that you want to provision.  ML storage volumes store
+            model artifacts and incremental states. Training algorithms might also use the ML
+            storage volume for scratch space. Default: 30
+        volume_kms_key_id (Optional[str]):
+            The Amazon Web Services KMS key that SageMaker uses to encrypt data on the storage
+            volume attached to the ML compute instance(s) that run the training job.
+        keep_alive_period_in_seconds (Optional[int]):
+            The duration of time in seconds to retain configured resources in a warm pool for
+            subsequent training jobs.
+        instance_groups (Optional[List[InstanceGroup]]):
+            A list of instance groups for heterogeneous clusters to be used in the training job.
+        enable_managed_spot_training (Optional[bool]):
+            To train models using managed spot training, choose True. Managed spot training
+            provides a fully managed and scalable infrastructure for training machine learning
+            models. this option is useful when training jobs can be interrupted and when there
+            is flexibility when the training job is run.
+    """
+
+    volume_size_in_gb: Optional[int] = 30
+    enable_managed_spot_training: Optional[bool] = None
+
+    @model_validator(mode="after")
+    def _model_validator(self) -> "ComputeConfig":
+        """Convert Unassigned values to None."""
+        return convert_unassigned_to_none(self)
+
+    def _to_resource_config(self) -> shapes.ResourceConfig:
+        """Convert to a sagemaker_core.shapes.ResourceConfig object."""
+        compute_config_dict = self.model_dump()
+        resource_config_fields = set(shapes.ResourceConfig.__annotations__.keys())
+        filtered_dict = {
+            k: v for k, v in compute_config_dict.items() if k in resource_config_fields
+        }
+        return shapes.ResourceConfig(**filtered_dict)
+
+
+class NetworkingConfig(shapes.VpcConfig):
+    """NetworkingConfig.
+
+    The NetworkingConifg is a subclass of `sagemaker_core.shapes.VpcConfig ` and
+    allows the user to specify the networking configuration for the training job.
+
+    Attributes:
+        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.
+        subnets (Optional[List[str]]):
+            The ID of the subnets in the VPC to which you want to connect your
+            training job or model.
+        enable_network_isolation (Optional[bool]):
+            Isolates the training container. No inbound or outbound network calls can be made,
+            except for calls between peers within a training cluster for distributed training.
+            If you enable network isolation for training jobs that are configured to use a VPC,
+            SageMaker downloads and uploads customer data and model artifacts through the
+            specified VPC, but the training container does not have network access.
+        enable_inter_container_traffic_encryption (Optional[bool]):
+            To encrypt all communications between ML compute instances in distributed training
+            choose True. Encryption provides greater security for distributed training, but
+            training might take longer. How long it takes depends on the amount of
+            communication between compute instances, especially if you use a deep learning
+            algorithm in distributed training.
+    """
+
+    enable_network_isolation: Optional[bool] = None
+    enable_inter_container_traffic_encryption: Optional[bool] = None
+
+    @model_validator(mode="after")
+    def _model_validator(self) -> "NetworkingConfig":
+        """Convert Unassigned values to None."""
+        return convert_unassigned_to_none(self)
+
+    def _to_vpc_config(self) -> shapes.VpcConfig:
+        """Convert to a sagemaker_core.shapes.VpcConfig object."""
+        compute_config_dict = self.model_dump()
+        resource_config_fields = set(shapes.VpcConfig.__annotations__.keys())
+        filtered_dict = {
+            k: v for k, v in compute_config_dict.items() if k in resource_config_fields
+        }
+        return shapes.VpcConfig(**filtered_dict)
+
+
+class InputData(BaseModel):
+    """InputData.
+
+    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.
+    For convience, can be referenced inside the training container like:
+
+    ```python
+    import os
+    input_data_dir = os.environ['SM_CHANNEL_<channel_name>']
+    ```
+
+    Attributes:
+        channel_name (str):
+            The name of the input data source channel.
+        data_source (Union[str, S3DataSource, FileSystemDataSource]):
+            The data source for the channel. Can be an S3 URI string, local file path string,
+            S3DataSource object, or FileSystemDataSource object.
+    """
+
+    channel_name: str = None
+    data_source: Union[str, FileSystemDataSource, S3DataSource] = None