Improve default handling

Author: benieric

src/sagemaker/modules/configs.py

@@ -30,7 +30,6 @@
 from sagemaker_core.shapes import (
     StoppingCondition,
     RetryStrategy,
-    OutputDataConfig,
     Channel,
     ShuffleConfig,
     DataSource,
@@ -42,9 +41,7 @@
     InfraCheckConfig,
     RemoteDebugConfig,
     SessionChainingConfig,
-    InstanceGroup,
-    TensorBoardOutputConfig,
-    CheckpointConfig,
+    InstanceGroup
 )
 
 from sagemaker.modules.utils import convert_unassigned_to_none
@@ -131,6 +128,8 @@ class Compute(shapes.ResourceConfig):
             subsequent training jobs.
         instance_groups (Optional[List[InstanceGroup]]):
             A list of instance groups for heterogeneous clusters to be used in the training job.
+        training_plan_arn (Optional[str]):
+            The Amazon Resource Name (ARN) of the training plan to use for this resource configuration.
         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
@@ -224,3 +223,64 @@ class InputData(BaseConfig):
 
     channel_name: str = None
     data_source: Union[str, FileSystemDataSource, S3DataSource] = None
+
+
+class OutputDataConfig(shapes.OutputDataConfig):
+    """OutputDataConfig.
+
+    The OutputDataConfig class is a subclass of ``sagemaker_core.shapes.OutputDataConfig``
+    and allows the user to specify the output data configuration for the training job.
+
+    Parameters:
+        s3_output_path (Optional[str]):
+            The S3 URI where the output data will be stored. This is the location where the
+            training job will save its output data, such as model artifacts and logs.
+        kms_key_id (Optional[str]):
+            The Amazon Web Services Key Management Service (Amazon Web Services KMS) key that
+            SageMaker uses to encrypt the model artifacts at rest using Amazon S3 server-side
+            encryption.
+        compression_type (Optional[str]):
+            The model output compression type. Select None to output an uncompressed model,
+            recommended for large model outputs. Defaults to gzip.
+    """
+
+    s3_output_path: Optional[str] = None
+    kms_key_id: Optional[str] = None
+    compression_type: Optional[str] = None
+
+
+class TensorBoardOutputConfig(shapes.TensorBoardOutputConfig):
+    """TensorBoardOutputConfig.
+
+    The TensorBoardOutputConfig class is a subclass of ``sagemaker_core.shapes.TensorBoardOutputConfig``
+    and allows the user to specify the storage locations for the Amazon SageMaker
+    Debugger TensorBoard.
+
+    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>/``
+        local_path (Optional[str]):
+            Path to local storage location for tensorBoard output. Defaults to /opt/ml/output/tensorboard.
+    """
+
+    s3_output_path: Optional[str] = None
+    local_path: Optional[str] = "/opt/ml/output/tensorboard"
+
+
+class CheckpointConfig(shapes.CheckpointConfig):
+    """CheckpointConfig.
+
+    The CheckpointConfig class is a subclass of ``sagemaker_core.shapes.CheckpointConfig``
+    and allows the user to specify the checkpoint configuration for the training job.
+
+    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>/``
+        local_path (Optional[str]):
+            The local directory where checkpoints are written. The default directory is /opt/ml/checkpoints.
+    """
+
+    s3_uri: Optional[str] = None
+    local_path: Optional[str] = "/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
@@ -477,6 +477,20 @@ def model_post_init(self, __context: Any):
             )
             logger.warning(f"Compute not provided. Using default:\n{self.compute}")
 
+        if self.compute.instance_type is None:
+            self.compute.instance_type = DEFAULT_INSTANCE_TYPE
+            logger.warning(f"Instance type not provided. Using default:\n{DEFAULT_INSTANCE_TYPE}")
+        if self.compute.instance_count is None:
+            self.compute.instance_count = 1
+            logger.warning(
+                f"Instance count not provided. Using default:\n{self.compute.instance_count}"
+            )
+        if self.compute.volume_size_in_gb is None:
+            self.compute.volume_size_in_gb = 30
+            logger.warning(
+                f"Volume size not provided. Using default:\n{self.compute.volume_size_in_gb}"
+            )
+
         if self.stopping_condition is None:
             self.stopping_condition = StoppingCondition(
                 max_runtime_in_seconds=3600,
@@ -486,6 +500,12 @@ def model_post_init(self, __context: Any):
             logger.warning(
                 f"StoppingCondition not provided. Using default:\n{self.stopping_condition}"
             )
+        if self.stopping_condition.max_runtime_in_seconds is None:
+            self.stopping_condition.max_runtime_in_seconds = 3600
+            logger.info(
+                "Max runtime not provided. Using default:\n"
+                f"{self.stopping_condition.max_runtime_in_seconds}"
+            )
 
         if self.hyperparameters and isinstance(self.hyperparameters, str):
             if not os.path.exists(self.hyperparameters):
@@ -511,23 +531,40 @@ def model_post_init(self, __context: Any):
                         )
 
         if self.training_mode == Mode.SAGEMAKER_TRAINING_JOB and self.output_data_config is None:
-            session = self.sagemaker_session
-            base_job_name = self.base_job_name
-            self.output_data_config = OutputDataConfig(
-                s3_output_path=f"s3://{self._fetch_bucket_name_and_prefix(session)}"
-                f"/{base_job_name}",
-                compression_type="GZIP",
-                kms_key_id=None,
-            )
-            logger.warning(
-                f"OutputDataConfig not provided. Using default:\n{self.output_data_config}"
-            )
+            if self.output_data_config is None:
+                session = self.sagemaker_session
+                base_job_name = self.base_job_name
+                self.output_data_config = OutputDataConfig(
+                    s3_output_path=f"s3://{self._fetch_bucket_name_and_prefix(session)}"
+                    f"/{base_job_name}",
+                    compression_type="GZIP",
+                    kms_key_id=None,
+                )
+                logger.warning(
+                    f"OutputDataConfig not provided. Using default:\n{self.output_data_config}"
+                )
+            if self.output_data_config.s3_output_path is None:
+                session = self.sagemaker_session
+                base_job_name = self.base_job_name
+                self.output_data_config.s3_output_path = (
+                    f"s3://{self._fetch_bucket_name_and_prefix(session)}/{base_job_name}"
+                )
+                logger.warning(
+                    f"OutputDataConfig s3_output_path not provided. Using default:\n"
+                    f"{self.output_data_config.s3_output_path}"
+                )
+            if self.output_data_config.compression_type is None:
+                self.output_data_config.compression_type = "GZIP"
+                logger.warning(
+                    f"OutputDataConfig compression type not provided. Using default:\n"
+                    f"{self.output_data_config.compression_type}"
+                )
 
-        # TODO: Autodetect which image to use if source_code is provided
         if self.training_image:
             logger.info(f"Training image URI: {self.training_image}")
 
-    def _fetch_bucket_name_and_prefix(self, session: Session) -> str:
+    @staticmethod
+    def _fetch_bucket_name_and_prefix(session: Session) -> str:
         """Helper function to get the bucket name with the corresponding prefix if applicable"""
         if session.default_bucket_prefix is not None:
             return f"{session.default_bucket()}/{session.default_bucket_prefix}"
@@ -558,16 +595,28 @@ def train(
         """
         self._populate_intelligent_defaults()
         current_training_job_name = _get_unique_name(self.base_job_name)
-        input_data_key_prefix = f"{self.base_job_name}/{current_training_job_name}/input"
-        if input_data_config:
+        default_artifact_path = f"{self.base_job_name}/{current_training_job_name}"
+        input_data_key_prefix = f"{default_artifact_path}/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
+            final_input_channel_names = {i.channel_name for i in input_data_config}
+            for input_data in self.input_data_config:
+                if input_data.channel_name not in final_input_channel_names:
+                    input_data_config.append(input_data)
+
+        self.input_data_config = input_data_config or self.input_data_config or []
 
-        input_data_config = []
         if self.input_data_config:
             input_data_config = self._get_input_data_config(
                 self.input_data_config, input_data_key_prefix
             )
 
+        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}"
+
         string_hyper_parameters = {}
         if self.hyperparameters:
             for hyper_parameter, value in self.hyperparameters.items():