change: Add PipelineVariable annotation to composite argument of training

Author: Dewen Qi

src/sagemaker/amazon/amazon_estimator.py

@@ -243,7 +243,6 @@ def fit(
                 * `TrialComponentDisplayName` is used for display in Studio.
         """
         self._prepare_for_training(records, job_name=job_name, mini_batch_size=mini_batch_size)
-
         self.latest_training_job = _TrainingJob.start_new(
             self, records, experiment_config=experiment_config
         )
@@ -304,7 +303,12 @@ class RecordSet(object):
     """Placeholder docstring"""
 
     def __init__(
-        self, s3_data, num_records, feature_dim, s3_data_type="ManifestFile", channel="train"
+        self,
+        s3_data: Union[str, PipelineVariable],
+        num_records: int,
+        feature_dim: int,
+        s3_data_type: Union[str, PipelineVariable] = "ManifestFile",
+        channel: Union[str, PipelineVariable] = "train",
     ):
         """A collection of Amazon :class:~`Record` objects serialized and stored in S3.
 

src/sagemaker/debugger/debugger.py

@@ -24,6 +24,9 @@
 
 from abc import ABC
 
+from sagemaker.workflow.entities import PipelineVariable
+from typing import Union, Optional, List, Dict
+
 import attr
 
 import smdebug_rulesconfig as rule_configs
@@ -311,17 +314,17 @@ def sagemaker(
     @classmethod
     def custom(
         cls,
-        name,
-        image_uri,
-        instance_type,
-        volume_size_in_gb,
-        source=None,
-        rule_to_invoke=None,
-        container_local_output_path=None,
-        s3_output_path=None,
-        other_trials_s3_input_paths=None,
-        rule_parameters=None,
-        collections_to_save=None,
+        name: str,
+        image_uri: Union[str, PipelineVariable],
+        instance_type: Union[str, PipelineVariable],
+        volume_size_in_gb: Union[int, PipelineVariable],
+        source: Optional[str] = None,
+        rule_to_invoke: Optional[Union[str, PipelineVariable]] = None,
+        container_local_output_path: Optional[Union[str, PipelineVariable]] = None,
+        s3_output_path: Optional[Union[str, PipelineVariable]] = None,
+        other_trials_s3_input_paths: Optional[List[Union[str, PipelineVariable]]] = None,
+        rule_parameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
+        collections_to_save: Optional[List["CollectionConfig"]] = None,
         actions=None,
     ):
         """Initialize a ``Rule`` object for a *custom* debugging rule.
@@ -610,10 +613,10 @@ class DebuggerHookConfig(object):
 
     def __init__(
         self,
-        s3_output_path=None,
-        container_local_output_path=None,
-        hook_parameters=None,
-        collection_configs=None,
+        s3_output_path: Optional[Union[str, PipelineVariable]] = None,
+        container_local_output_path: Optional[Union[str, PipelineVariable]] = None,
+        hook_parameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
+        collection_configs: Optional[List["CollectionConfig"]] = None,
     ):
         """Initialize the DebuggerHookConfig instance.
 
@@ -679,7 +682,11 @@ def _to_request_dict(self):
 class TensorBoardOutputConfig(object):
     """Create a tensor ouput configuration object for debugging visualizations on TensorBoard."""
 
-    def __init__(self, s3_output_path, container_local_output_path=None):
+    def __init__(
+        self,
+        s3_output_path: Union[str, PipelineVariable],
+        container_local_output_path: Optional[Union[str, PipelineVariable]] = None,
+    ):
         """Initialize the TensorBoardOutputConfig instance.
 
         Args:
@@ -708,7 +715,11 @@ def _to_request_dict(self):
 class CollectionConfig(object):
     """Creates tensor collections for SageMaker Debugger."""
 
-    def __init__(self, name, parameters=None):
+    def __init__(
+        self,
+        name: Union[str, PipelineVariable],
+        parameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
+    ):
         """Constructor for collection configuration.
 
         Args:

src/sagemaker/debugger/profiler_config.py

@@ -13,7 +13,10 @@
 """Configuration for collecting system and framework metrics in SageMaker training jobs."""
 from __future__ import absolute_import
 
+from typing import Optional, Union
+
 from sagemaker.debugger.framework_profile import FrameworkProfile
+from sagemaker.workflow.entities import PipelineVariable
 
 
 class ProfilerConfig(object):
@@ -26,9 +29,9 @@ class ProfilerConfig(object):
 
     def __init__(
         self,
-        s3_output_path=None,
-        system_monitor_interval_millis=None,
-        framework_profile_params=None,
+        s3_output_path: Optional[Union[str, PipelineVariable]] = None,
+        system_monitor_interval_millis: Optional[Union[int, PipelineVariable]] = None,
+        framework_profile_params: Optional[FrameworkProfile] = None,
     ):
         """Initialize a ``ProfilerConfig`` instance.
 

src/sagemaker/estimator.py

@@ -2747,8 +2747,9 @@ def _validate_and_set_debugger_configs(self):
                 if not is_pipeline_variable(self.instance_count) and self.instance_count > 1:
                     disable_debugger_hook_config = True
                 if (
-                    hasattr(self, "distribution") and self.distribution is not None
-                ):  # pylint: disable=no-member
+                    hasattr(self, "distribution")
+                    and self.distribution is not None  # pylint: disable=no-member
+                ):
                     disable_debugger_hook_config = True
                 if disable_debugger_hook_config:
                     logger.info(

src/sagemaker/huggingface/training_compiler/config.py

@@ -13,8 +13,10 @@
 """Configuration for the SageMaker Training Compiler."""
 from __future__ import absolute_import
 import logging
+from typing import Union
 
 from sagemaker.training_compiler.config import TrainingCompilerConfig as BaseConfig
+from sagemaker.workflow.entities import PipelineVariable
 
 logger = logging.getLogger(__name__)
 
@@ -26,8 +28,8 @@ class TrainingCompilerConfig(BaseConfig):
 
     def __init__(
         self,
-        enabled=True,
-        debug=False,
+        enabled: Union[bool, PipelineVariable] = True,
+        debug: Union[bool, PipelineVariable] = False,
     ):
         """This class initializes a ``TrainingCompilerConfig`` instance.
 

src/sagemaker/inputs.py

@@ -13,6 +13,7 @@
 """Amazon SageMaker channel configurations for S3 data sources and file system data sources"""
 from __future__ import absolute_import, print_function
 
+from typing import Union, Optional, List
 import attr
 
 FILE_SYSTEM_TYPES = ["FSxLustre", "EFS"]
@@ -29,16 +30,16 @@ class TrainingInput(object):
 
     def __init__(
         self,
-        s3_data,
-        distribution=None,
-        compression=None,
-        content_type=None,
-        record_wrapping=None,
-        s3_data_type="S3Prefix",
-        input_mode=None,
-        attribute_names=None,
-        target_attribute_name=None,
-        shuffle_config=None,
+        s3_data: Union[str, PipelineVariable],
+        distribution: Optional[Union[str, PipelineVariable]] = None,
+        compression: Optional[Union[str, PipelineVariable]] = None,
+        content_type: Optional[Union[str, PipelineVariable]] = None,
+        record_wrapping: Optional[Union[str, PipelineVariable]] = None,
+        s3_data_type: Union[str, PipelineVariable] = "S3Prefix",
+        input_mode: Optional[Union[str, PipelineVariable]] = None,
+        attribute_names: Optional[List[Union[str, PipelineVariable]]] = None,
+        target_attribute_name: Optional[Union[str, PipelineVariable]] = None,
+        shuffle_config: Optional["ShuffleConfig"] = None,
     ):
         """Create a definition for input data used by an SageMaker training job.