change: Implement test mechanism for Pipeline variables

Author: Dewen Qi

src/sagemaker/debugger/debugger.py

@@ -23,13 +23,15 @@
 import time
 
 from abc import ABC
+from typing import Union, Optional
 
 import attr
 
 import smdebug_rulesconfig as rule_configs
 
 from sagemaker import image_uris
 from sagemaker.utils import build_dict
+from sagemaker.workflow.entities import PipelineVariable
 
 framework_name = "debugger"
 DEBUGGER_FLAG = "USE_SMDEBUG"
@@ -311,10 +313,10 @@ def sagemaker(
     @classmethod
     def custom(
         cls,
-        name,
-        image_uri,
-        instance_type,
-        volume_size_in_gb,
+        name: str,
+        image_uri: Union[str, PipelineVariable],
+        instance_type: Union[str, PipelineVariable],
+        volume_size_in_gb: Union[int, PipelineVariable],
         source=None,
         rule_to_invoke=None,
         container_local_output_path=None,
@@ -610,7 +612,7 @@ class DebuggerHookConfig(object):
 
     def __init__(
         self,
-        s3_output_path=None,
+        s3_output_path: Optional[Union[str, PipelineVariable]] = None,
         container_local_output_path=None,
         hook_parameters=None,
         collection_configs=None,
@@ -679,7 +681,9 @@ 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=None
+    ):
         """Initialize the TensorBoardOutputConfig instance.
 
         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 Union, Optional
+
 from sagemaker.debugger.framework_profile import FrameworkProfile
+from sagemaker.workflow.entities import PipelineVariable
 
 
 class ProfilerConfig(object):
@@ -27,7 +30,7 @@ class ProfilerConfig(object):
     def __init__(
         self,
         s3_output_path=None,
-        system_monitor_interval_millis=None,
+        system_monitor_interval_millis: Optional[Union[int, PipelineVariable]] = None,
         framework_profile_params=None,
     ):
         """Initialize a ``ProfilerConfig`` instance.

src/sagemaker/inputs.py

@@ -13,8 +13,12 @@
 """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
+
 import attr
 
+from sagemaker.workflow.entities import PipelineVariable
+
 FILE_SYSTEM_TYPES = ["FSxLustre", "EFS"]
 FILE_SYSTEM_ACCESS_MODES = ["ro", "rw"]
 
@@ -29,10 +33,10 @@ class TrainingInput(object):
 
     def __init__(
         self,
-        s3_data,
+        s3_data: Union[str, PipelineVariable],
         distribution=None,
         compression=None,
-        content_type=None,
+        content_type: Optional[Union[str, PipelineVariable]] = None,
         record_wrapping=None,
         s3_data_type="S3Prefix",
         input_mode=None,

src/sagemaker/serverless/serverless_inference_config.py

@@ -17,6 +17,10 @@
 """
 from __future__ import print_function, absolute_import
 
+from typing import Union
+
+from sagemaker.workflow.entities import PipelineVariable
+
 
 class ServerlessInferenceConfig(object):
     """Configuration object passed in when deploying models to Amazon SageMaker Endpoints.

src/sagemaker/tensorflow/training_compiler/config.py

@@ -13,10 +13,13 @@
 """Configuration for the SageMaker Training Compiler."""
 from __future__ import absolute_import
 import logging
+from typing import Union
+
 from packaging.specifiers import SpecifierSet
 from packaging.version import Version
 
 from sagemaker.training_compiler.config import TrainingCompilerConfig as BaseConfig
+from sagemaker.workflow.entities import PipelineVariable
 
 logger = logging.getLogger(__name__)
 
@@ -29,8 +32,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.