fix: Support parameterized source code input for TrainingStep (#3202)

Co-authored-by: Dewen Qi <[email protected]>

Author: qidewenwhen

src/sagemaker/amazon/hyperparameter.py

@@ -15,6 +15,8 @@
 
 import json
 
+from sagemaker.workflow import is_pipeline_variable
+
 
 class Hyperparameter(object):
     """An algorithm hyperparameter with optional validation.
@@ -98,8 +100,14 @@ def serialize_all(obj):
         """
         if "_hyperparameters" not in dir(obj):
             return {}
-        return {
-            k: json.dumps(v) if isinstance(v, list) else str(v)
-            for k, v in obj._hyperparameters.items()
-            if v is not None
-        }
+        hps = {}
+        for k, v in obj._hyperparameters.items():
+            if v is not None:
+                if isinstance(v, list):
+                    v = json.dumps(v)
+                elif is_pipeline_variable(v):
+                    v = v.to_string()
+                else:
+                    v = str(v)
+                hps[k] = v
+        return hps

src/sagemaker/chainer/estimator.py

@@ -14,6 +14,7 @@
 from __future__ import absolute_import
 
 import logging
+from typing import Union, Optional
 
 from sagemaker.estimator import Framework, EstimatorBase
 from sagemaker.fw_utils import (
@@ -25,6 +26,7 @@
 from sagemaker.chainer import defaults
 from sagemaker.chainer.model import ChainerModel
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
+from sagemaker.workflow.entities import PipelineVariable
 
 logger = logging.getLogger("sagemaker")
 
@@ -42,12 +44,12 @@ class Chainer(Framework):
 
     def __init__(
         self,
-        entry_point,
+        entry_point: Union[str, PipelineVariable],
         use_mpi=None,
         num_processes=None,
         process_slots_per_host=None,
         additional_mpi_options=None,
-        source_dir=None,
+        source_dir: Optional[Union[str, PipelineVariable]] = None,
         hyperparameters=None,
         framework_version=None,
         py_version=None,

src/sagemaker/estimator.py

@@ -47,6 +47,7 @@
     get_mp_parameters,
     tar_and_upload_dir,
     validate_source_dir,
+    validate_source_code_input_against_pipeline_variables,
 )
 from sagemaker.inputs import TrainingInput, FileSystemInput
 from sagemaker.job import _Job
@@ -140,12 +141,12 @@ def __init__(
         disable_profiler: bool = False,
         environment: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
         max_retry_attempts: Optional[Union[int, PipelineVariable]] = None,
-        source_dir: Optional[str] = None,
+        source_dir: Optional[Union[str, PipelineVariable]] = None,
         git_config: Optional[Dict[str, str]] = None,
         hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
         container_log_level: Union[int, PipelineVariable] = logging.INFO,
         code_location: Optional[str] = None,
-        entry_point: Optional[str] = None,
+        entry_point: Optional[Union[str, PipelineVariable]] = None,
         dependencies: Optional[List[Union[str]]] = None,
         instance_groups: Optional[Dict[str, Union[str, int]]] = None,
         **kwargs,
@@ -461,6 +462,13 @@ def __init__(
             "train_volume_kms_key", "volume_kms_key", volume_kms_key, kwargs
         )
 
+        validate_source_code_input_against_pipeline_variables(
+            entry_point=entry_point,
+            source_dir=source_dir,
+            git_config=git_config,
+            enable_network_isolation=enable_network_isolation,
+        )
+
         self.role = role
         self.instance_count = instance_count
         self.instance_type = instance_type
@@ -663,7 +671,11 @@ def _prepare_for_training(self, job_name=None):
             # validate source dir will raise a ValueError if there is something wrong with
             # the source directory. We are intentionally not handling it because this is a
             # critical error.
-            if self.source_dir and not self.source_dir.lower().startswith("s3://"):
+            if (
+                self.source_dir
+                and not is_pipeline_variable(self.source_dir)
+                and not self.source_dir.lower().startswith("s3://")
+            ):
                 validate_source_dir(self.entry_point, self.source_dir)
 
             # if we are in local mode with local_code=True. We want the container to just
@@ -2151,11 +2163,11 @@ def __init__(
         disable_profiler: bool = False,
         environment: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
         max_retry_attempts: Optional[Union[int, PipelineVariable]] = None,
-        source_dir: Optional[str] = None,
+        source_dir: Optional[Union[str, PipelineVariable]] = None,
         git_config: Optional[Dict[str, str]] = None,
         container_log_level: Union[int, PipelineVariable] = logging.INFO,
         code_location: Optional[str] = None,
-        entry_point: Optional[str] = None,
+        entry_point: Optional[Union[str, PipelineVariable]] = None,
         dependencies: Optional[List[str]] = None,
         instance_groups: Optional[Dict[str, Union[str, int]]] = None,
         **kwargs,
@@ -2603,8 +2615,8 @@ class Framework(EstimatorBase):
 
     def __init__(
         self,
-        entry_point: str,
-        source_dir: Optional[str] = None,
+        entry_point: Union[str, PipelineVariable],
+        source_dir: Optional[Union[str, PipelineVariable]] = None,
         hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
         container_log_level: Union[int, PipelineVariable] = logging.INFO,
         code_location: Optional[str] = None,
@@ -2783,7 +2795,14 @@ def __init__(
         """
         super(Framework, self).__init__(enable_network_isolation=enable_network_isolation, **kwargs)
         image_uri = renamed_kwargs("image_name", "image_uri", image_uri, kwargs)
-        if entry_point.startswith("s3://"):
+
+        validate_source_code_input_against_pipeline_variables(
+            entry_point=entry_point,
+            source_dir=source_dir,
+            git_config=git_config,
+            enable_network_isolation=enable_network_isolation,
+        )
+        if not is_pipeline_variable(entry_point) and entry_point.startswith("s3://"):
             raise ValueError(
                 "Invalid entry point script: {}. Must be a path to a local file.".format(
                     entry_point

src/sagemaker/fw_utils.py

@@ -20,14 +20,15 @@
 import shutil
 import tempfile
 from collections import namedtuple
-from typing import Optional
+from typing import Optional, Union, Dict
 
 import sagemaker.image_uris
 from sagemaker.session_settings import SessionSettings
 import sagemaker.utils
 from sagemaker.workflow import is_pipeline_variable
 
 from sagemaker.deprecations import renamed_warning, renamed_kwargs
+from sagemaker.workflow.entities import PipelineVariable
 
 logger = logging.getLogger(__name__)
 
@@ -124,6 +125,58 @@ def validate_source_dir(script, directory):
     return True
 
 
+def validate_source_code_input_against_pipeline_variables(
+    entry_point: Optional[Union[str, PipelineVariable]] = None,
+    source_dir: Optional[Union[str, PipelineVariable]] = None,
+    git_config: Optional[Dict[str, str]] = None,
+    enable_network_isolation: Union[bool, PipelineVariable] = False,
+):
+    """Validate source code input against pipeline variables
+
+    Args:
+        entry_point (str, PipelineVariable): The path to the local Python source file that
+            should be executed as the entry point to training (default: None).
+        source_dir (str, PipelineVariable): The Path to a directory with any other
+            training source code dependencies aside from the entry point file (default: None).
+        git_config (Dict[str, str]): Git configurations used for cloning files (default: None).
+        enable_network_isolation (bool, PipelineVariable): Specifies whether container will run
+            in network isolation mode (default: False).
+    """
+    if is_pipeline_variable(enable_network_isolation) or enable_network_isolation is True:
+        if is_pipeline_variable(entry_point) or is_pipeline_variable(source_dir):
+            raise TypeError(
+                "entry_point, source_dir should not be pipeline variables "
+                "when enable_network_isolation is a pipeline variable or it is set to True."
+            )
+    if git_config:
+        if is_pipeline_variable(entry_point) or is_pipeline_variable(source_dir):
+            raise TypeError(
+                "entry_point, source_dir should not be pipeline variables when git_config is given."
+            )
+    if is_pipeline_variable(entry_point):
+        if not source_dir:
+            raise TypeError(
+                "The entry_point should not be a pipeline variable when source_dir is missing."
+            )
+        if not is_pipeline_variable(source_dir) and not source_dir.lower().startswith("s3://"):
+            raise TypeError(
+                "The entry_point should not be a pipeline variable when source_dir is a local path."
+            )
+        logger.warning(
+            "The entry_point is a pipeline variable: %s. During pipeline execution, "
+            "the interpreted value of entry_point has to be a local path in the container "
+            "pointing to a Python source file which is located at the root of source_dir.",
+            type(entry_point),
+        )
+    if is_pipeline_variable(source_dir):
+        logger.warning(
+            "The source_dir is a pipeline variable: %s. During pipeline execution, "
+            "the interpreted value of source_dir has to be an S3 URI and "
+            "must point to a tar.gz file",
+            type(source_dir),
+        )
+
+
 def get_mp_parameters(distribution):
     """Get the model parallelism parameters provided by the user.
 
@@ -265,7 +318,7 @@ def tar_and_upload_dir(
         sagemaker.fw_utils.UserCode: An object with the S3 bucket and key (S3 prefix) and
             script name.
     """
-    if directory and directory.lower().startswith("s3://"):
+    if directory and (is_pipeline_variable(directory) or directory.lower().startswith("s3://")):
         return UploadedCode(s3_prefix=directory, script_name=script)
 
     script_name = script if directory else os.path.basename(script)

src/sagemaker/huggingface/estimator.py

@@ -15,6 +15,7 @@
 
 import logging
 import re
+from typing import Optional, Union, Dict
 
 from sagemaker.deprecations import renamed_kwargs
 from sagemaker.estimator import Framework, EstimatorBase
@@ -27,6 +28,7 @@
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 
 from sagemaker.huggingface.training_compiler.config import TrainingCompilerConfig
+from sagemaker.workflow.entities import PipelineVariable
 
 logger = logging.getLogger("sagemaker")
 
@@ -38,16 +40,16 @@ class HuggingFace(Framework):
 
     def __init__(
         self,
-        py_version,
-        entry_point,
-        transformers_version=None,
-        tensorflow_version=None,
-        pytorch_version=None,
-        source_dir=None,
-        hyperparameters=None,
-        image_uri=None,
-        distribution=None,
-        compiler_config=None,
+        py_version: str,
+        entry_point: Union[str, PipelineVariable],
+        transformers_version: Optional[str] = None,
+        tensorflow_version: Optional[str] = None,
+        pytorch_version: Optional[str] = None,
+        source_dir: Optional[Union[str, PipelineVariable]] = None,
+        hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
+        image_uri: Optional[Union[str, PipelineVariable]] = None,
+        distribution: Optional[Dict] = None,
+        compiler_config: Optional[TrainingCompilerConfig] = None,
         **kwargs,
     ):
         """This estimator runs a Hugging Face training script in a SageMaker training environment.

src/sagemaker/job.py

@@ -18,6 +18,7 @@
 
 from sagemaker.inputs import FileSystemInput, TrainingInput
 from sagemaker.local import file_input
+from sagemaker.workflow import is_pipeline_variable
 
 
 class _Job(object):
@@ -168,14 +169,14 @@ def _format_string_uri_input(
         target_attribute_name=None,
     ):
         """Placeholder docstring"""
+        s3_input_result = TrainingInput(
+            uri_input,
+            content_type=content_type,
+            input_mode=input_mode,
+            compression=compression,
+            target_attribute_name=target_attribute_name,
+        )
         if isinstance(uri_input, str) and validate_uri and uri_input.startswith("s3://"):
-            s3_input_result = TrainingInput(
-                uri_input,
-                content_type=content_type,
-                input_mode=input_mode,
-                compression=compression,
-                target_attribute_name=target_attribute_name,
-            )
             return s3_input_result
         if isinstance(uri_input, str) and validate_uri and uri_input.startswith("file://"):
             return file_input(uri_input)
@@ -185,16 +186,11 @@ def _format_string_uri_input(
                 '"file://"'.format(uri_input)
             )
         if isinstance(uri_input, str):
-            s3_input_result = TrainingInput(
-                uri_input,
-                content_type=content_type,
-                input_mode=input_mode,
-                compression=compression,
-                target_attribute_name=target_attribute_name,
-            )
             return s3_input_result
         if isinstance(uri_input, (TrainingInput, file_input, FileSystemInput)):
             return uri_input
+        if is_pipeline_variable(uri_input):
+            return s3_input_result
 
         raise ValueError(
             "Cannot format input {}. Expecting one of str, TrainingInput, file_input or "

src/sagemaker/jumpstart/utils.py

@@ -30,7 +30,7 @@
     JumpStartModelSpecs,
     JumpStartVersionedModelId,
 )
-
+from sagemaker.workflow import is_pipeline_variable
 
 LOGGER = logging.getLogger(__name__)
 
@@ -271,26 +271,41 @@ def add_jumpstart_tags(
         training_script_uri (Optional[str]): S3 URI for training script tarball.
             (Default: None).
     """
-
+    warn_msg = (
+        "The URI (%s) is a pipeline variable which is only interpreted at execution time. "
+        "As a result, the JumpStart resources will not be tagged."
+    )
     if inference_model_uri:
-        tags = add_single_jumpstart_tag(
-            inference_model_uri, enums.JumpStartTag.INFERENCE_MODEL_URI, tags
-        )
+        if is_pipeline_variable(inference_model_uri):
+            logging.warning(warn_msg, "inference_model_uri")
+        else:
+            tags = add_single_jumpstart_tag(
+                inference_model_uri, enums.JumpStartTag.INFERENCE_MODEL_URI, tags
+            )
 
     if inference_script_uri:
-        tags = add_single_jumpstart_tag(
-            inference_script_uri, enums.JumpStartTag.INFERENCE_SCRIPT_URI, tags
-        )
+        if is_pipeline_variable(inference_script_uri):
+            logging.warning(warn_msg, "inference_script_uri")
+        else:
+            tags = add_single_jumpstart_tag(
+                inference_script_uri, enums.JumpStartTag.INFERENCE_SCRIPT_URI, tags
+            )
 
     if training_model_uri:
-        tags = add_single_jumpstart_tag(
-            training_model_uri, enums.JumpStartTag.TRAINING_MODEL_URI, tags
-        )
+        if is_pipeline_variable(training_model_uri):
+            logging.warning(warn_msg, "training_model_uri")
+        else:
+            tags = add_single_jumpstart_tag(
+                training_model_uri, enums.JumpStartTag.TRAINING_MODEL_URI, tags
+            )
 
     if training_script_uri:
-        tags = add_single_jumpstart_tag(
-            training_script_uri, enums.JumpStartTag.TRAINING_SCRIPT_URI, tags
-        )
+        if is_pipeline_variable(training_script_uri):
+            logging.warning(warn_msg, "training_script_uri")
+        else:
+            tags = add_single_jumpstart_tag(
+                training_script_uri, enums.JumpStartTag.TRAINING_SCRIPT_URI, tags
+            )
 
     return tags