Merge branch 'master' into feat/enhance-bucket-override-support

Author: evakravi

src/sagemaker/spark/processing.py

@@ -31,13 +31,20 @@
 from io import BytesIO
 from urllib.parse import urlparse
 
+from typing import Union, List, Dict, Optional
+
 from sagemaker import image_uris
 from sagemaker.local.image import _ecr_login_if_needed, _pull_image
 from sagemaker.processing import ProcessingInput, ProcessingOutput, ScriptProcessor
 from sagemaker.s3 import S3Uploader
 from sagemaker.session import Session
+from sagemaker.network import NetworkConfig
 from sagemaker.spark import defaults
 
+from sagemaker.workflow import is_pipeline_variable
+from sagemaker.workflow.entities import PipelineVariable
+from sagemaker.workflow.functions import Join
+
 logger = logging.getLogger(__name__)
 
 
@@ -249,6 +256,12 @@ def run(
         """
         self._current_job_name = self._generate_current_job_name(job_name=job_name)
 
+        if is_pipeline_variable(submit_app):
+            raise ValueError(
+                "submit_app argument has to be a valid S3 URI or local file path "
+                + "rather than a pipeline variable"
+            )
+
         return super().run(
             submit_app,
             inputs,
@@ -437,9 +450,14 @@ def _stage_submit_deps(self, submit_deps, input_channel_name):
 
         use_input_channel = False
         spark_opt_s3_uris = []
+        spark_opt_s3_uris_has_pipeline_var = False
 
         with tempfile.TemporaryDirectory() as tmpdir:
             for dep_path in submit_deps:
+                if is_pipeline_variable(dep_path):
+                    spark_opt_s3_uris.append(dep_path)
+                    spark_opt_s3_uris_has_pipeline_var = True
+                    continue
                 dep_url = urlparse(dep_path)
                 # S3 URIs are included as-is in the spark-submit argument
                 if dep_url.scheme in ["s3", "s3a"]:
@@ -482,11 +500,19 @@ def _stage_submit_deps(self, submit_deps, input_channel_name):
                 destination=f"{self._conf_container_base_path}{input_channel_name}",
                 input_name=input_channel_name,
             )
-            spark_opt = ",".join(spark_opt_s3_uris + [input_channel.destination])
+            spark_opt = (
+                Join(on=",", values=spark_opt_s3_uris + [input_channel.destination])
+                if spark_opt_s3_uris_has_pipeline_var
+                else ",".join(spark_opt_s3_uris + [input_channel.destination])
+            )
         # If no local files were uploaded, form the spark-submit option from a list of S3 URIs
         else:
             input_channel = None
-            spark_opt = ",".join(spark_opt_s3_uris)
+            spark_opt = (
+                Join(on=",", values=spark_opt_s3_uris)
+                if spark_opt_s3_uris_has_pipeline_var
+                else ",".join(spark_opt_s3_uris)
+            )
 
         return input_channel, spark_opt
 
@@ -592,6 +618,9 @@ def _validate_s3_uri(self, spark_output_s3_path):
         Args:
             spark_output_s3_path (str): The URI of the Spark output S3 Path.
         """
+        if is_pipeline_variable(spark_output_s3_path):
+            return
+
         if urlparse(spark_output_s3_path).scheme != "s3":
             raise ValueError(
                 f"Invalid s3 path: {spark_output_s3_path}. Please enter something like "
@@ -650,22 +679,22 @@ class PySparkProcessor(_SparkProcessorBase):
 
     def __init__(
         self,
-        role,
-        instance_type,
-        instance_count,
-        framework_version=None,
-        py_version=None,
-        container_version=None,
-        image_uri=None,
-        volume_size_in_gb=30,
-        volume_kms_key=None,
-        output_kms_key=None,
-        max_runtime_in_seconds=None,
-        base_job_name=None,
-        sagemaker_session=None,
-        env=None,
-        tags=None,
-        network_config=None,
+        role: str,
+        instance_type: Union[int, PipelineVariable],
+        instance_count: Union[str, PipelineVariable],
+        framework_version: Optional[str] = None,
+        py_version: Optional[str] = None,
+        container_version: Optional[str] = None,
+        image_uri: Optional[Union[str, PipelineVariable]] = None,
+        volume_size_in_gb: Union[int, PipelineVariable] = 30,
+        volume_kms_key: Optional[Union[str, PipelineVariable]] = None,
+        output_kms_key: Optional[Union[str, PipelineVariable]] = None,
+        max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None,
+        base_job_name: Optional[str] = None,
+        sagemaker_session: Optional[Session] = None,
+        env: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
+        tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None,
+        network_config: Optional[NetworkConfig] = None,
     ):
         """Initialize an ``PySparkProcessor`` instance.
 
@@ -795,20 +824,20 @@ def get_run_args(
 
     def run(
         self,
-        submit_app,
-        submit_py_files=None,
-        submit_jars=None,
-        submit_files=None,
-        inputs=None,
-        outputs=None,
-        arguments=None,
-        wait=True,
-        logs=True,
-        job_name=None,
-        experiment_config=None,
-        configuration=None,
-        spark_event_logs_s3_uri=None,
-        kms_key=None,
+        submit_app: str,
+        submit_py_files: Optional[List[Union[str, PipelineVariable]]] = None,
+        submit_jars: Optional[List[Union[str, PipelineVariable]]] = None,
+        submit_files: Optional[List[Union[str, PipelineVariable]]] = None,
+        inputs: Optional[List[ProcessingInput]] = None,
+        outputs: Optional[List[ProcessingOutput]] = None,
+        arguments: Optional[List[Union[str, PipelineVariable]]] = None,
+        wait: bool = True,
+        logs: bool = True,
+        job_name: Optional[str] = None,
+        experiment_config: Optional[Dict[str, str]] = None,
+        configuration: Optional[Union[List[Dict], Dict]] = None,
+        spark_event_logs_s3_uri: Optional[Union[str, PipelineVariable]] = None,
+        kms_key: Optional[str] = None,
     ):
         """Runs a processing job.
 
@@ -907,22 +936,22 @@ class SparkJarProcessor(_SparkProcessorBase):
 
     def __init__(
         self,
-        role,
-        instance_type,
-        instance_count,
-        framework_version=None,
-        py_version=None,
-        container_version=None,
-        image_uri=None,
-        volume_size_in_gb=30,
-        volume_kms_key=None,
-        output_kms_key=None,
-        max_runtime_in_seconds=None,
-        base_job_name=None,
-        sagemaker_session=None,
-        env=None,
-        tags=None,
-        network_config=None,
+        role: str,
+        instance_type: Union[int, PipelineVariable],
+        instance_count: Union[str, PipelineVariable],
+        framework_version: Optional[str] = None,
+        py_version: Optional[str] = None,
+        container_version: Optional[str] = None,
+        image_uri: Optional[Union[str, PipelineVariable]] = None,
+        volume_size_in_gb: Union[int, PipelineVariable] = 30,
+        volume_kms_key: Optional[Union[str, PipelineVariable]] = None,
+        output_kms_key: Optional[Union[str, PipelineVariable]] = None,
+        max_runtime_in_seconds: Optional[Union[int, PipelineVariable]] = None,
+        base_job_name: Optional[str] = None,
+        sagemaker_session: Optional[Session] = None,
+        env: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
+        tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None,
+        network_config: Optional[NetworkConfig] = None,
     ):
         """Initialize a ``SparkJarProcessor`` instance.
 
@@ -1052,20 +1081,20 @@ def get_run_args(
 
     def run(
         self,
-        submit_app,
-        submit_class=None,
-        submit_jars=None,
-        submit_files=None,
-        inputs=None,
-        outputs=None,
-        arguments=None,
-        wait=True,
-        logs=True,
-        job_name=None,
-        experiment_config=None,
-        configuration=None,
-        spark_event_logs_s3_uri=None,
-        kms_key=None,
+        submit_app: str,
+        submit_class: Union[str, PipelineVariable],
+        submit_jars: Optional[List[Union[str, PipelineVariable]]] = None,
+        submit_files: Optional[List[Union[str, PipelineVariable]]] = None,
+        inputs: Optional[List[ProcessingInput]] = None,
+        outputs: Optional[List[ProcessingOutput]] = None,
+        arguments: Optional[List[Union[str, PipelineVariable]]] = None,
+        wait: bool = True,
+        logs: bool = True,
+        job_name: Optional[str] = None,
+        experiment_config: Optional[Dict[str, str]] = None,
+        configuration: Optional[Union[List[Dict], Dict]] = None,
+        spark_event_logs_s3_uri: Optional[Union[str, PipelineVariable]] = None,
+        kms_key: Optional[str] = None,
     ):
         """Runs a processing job.
 

src/sagemaker/workflow/condition_step.py

@@ -15,20 +15,17 @@
 
 from typing import List, Union, Optional
 
-import attr
 
 from sagemaker.deprecations import deprecated_class
 from sagemaker.workflow.conditions import Condition
 from sagemaker.workflow.step_collections import StepCollection
+from sagemaker.workflow.functions import JsonGet as NewJsonGet
 from sagemaker.workflow.steps import (
     Step,
     StepTypeEnum,
 )
 from sagemaker.workflow.utilities import list_to_request
-from sagemaker.workflow.entities import (
-    RequestType,
-    PipelineVariable,
-)
+from sagemaker.workflow.entities import RequestType
 from sagemaker.workflow.properties import (
     Properties,
     PropertyFile,
@@ -93,16 +90,15 @@ def arguments(self) -> RequestType:
     @property
     def step_only_arguments(self):
         """Argument dict pertaining to the step only, and not the `if_steps` or `else_steps`."""
-        return self.conditions
+        return [condition.to_request() for condition in self.conditions]
 
     @property
     def properties(self):
         """A simple Properties object with `Outcome` as the only property"""
         return self._properties
 
 
-@attr.s
-class JsonGet(PipelineVariable):  # pragma: no cover
+class JsonGet(NewJsonGet):  # pragma: no cover
     """Get JSON properties from PropertyFiles.
 
     Attributes:
@@ -112,28 +108,8 @@ class JsonGet(PipelineVariable):  # pragma: no cover
         json_path (str): The JSON path expression to the requested value.
     """
 
-    step: Step = attr.ib()
-    property_file: Union[PropertyFile, str] = attr.ib()
-    json_path: str = attr.ib()
-
-    @property
-    def expr(self):
-        """The expression dict for a `JsonGet` function."""
-        if isinstance(self.property_file, PropertyFile):
-            name = self.property_file.name
-        else:
-            name = self.property_file
-        return {
-            "Std:JsonGet": {
-                "PropertyFile": {"Get": f"Steps.{self.step.name}.PropertyFiles.{name}"},
-                "Path": self.json_path,
-            }
-        }
-
-    @property
-    def _referenced_steps(self) -> List[str]:
-        """List of step names that this function depends on."""
-        return [self.step.name]
+    def __init__(self, step: Step, property_file: Union[PropertyFile, str], json_path: str):
+        super().__init__(step_name=step.name, property_file=property_file, json_path=json_path)
 
 
 JsonGet = deprecated_class(JsonGet, "JsonGet")

src/sagemaker/workflow/conditions.py

@@ -20,15 +20,13 @@
 import abc
 
 from enum import Enum
-from typing import Dict, List, Union
+from typing import List, Union
 
 import attr
 
-from sagemaker.workflow import is_pipeline_variable
 from sagemaker.workflow.entities import (
     DefaultEnumMeta,
     Entity,
-    Expression,
     PrimitiveType,
     RequestType,
 )
@@ -88,8 +86,8 @@ def to_request(self) -> RequestType:
         """Get the request structure for workflow service calls."""
         return {
             "Type": self.condition_type.value,
-            "LeftValue": primitive_or_expr(self.left),
-            "RightValue": primitive_or_expr(self.right),
+            "LeftValue": self.left,
+            "RightValue": self.right,
         }
 
     @property
@@ -227,8 +225,8 @@ def to_request(self) -> RequestType:
         """Get the request structure for workflow service calls."""
         return {
             "Type": self.condition_type.value,
-            "QueryValue": self.value.expr,
-            "Values": [primitive_or_expr(in_value) for in_value in self.in_values],
+            "QueryValue": self.value,
+            "Values": self.in_values,
         }
 
     @property
@@ -291,19 +289,3 @@ def _referenced_steps(self) -> List[str]:
         for condition in self.conditions:
             steps.extend(condition._referenced_steps)
         return steps
-
-
-def primitive_or_expr(
-    value: Union[ExecutionVariable, Expression, PrimitiveType, Parameter, Properties]
-) -> Union[Dict[str, str], PrimitiveType]:
-    """Provide the expression of the value or return value if it is a primitive.
-
-    Args:
-        value (Union[ConditionValueType, PrimitiveType]): The value to evaluate.
-
-    Returns:
-        Either the expression of the value or the primitive value.
-    """
-    if is_pipeline_variable(value):
-        return value.expr
-    return value

src/sagemaker/workflow/step_collections.py

@@ -57,6 +57,9 @@ def properties(self):
 class RegisterModel(StepCollection):  # pragma: no cover
     """Register Model step collection for workflow."""
 
+    _REGISTER_MODEL_NAME_BASE = "RegisterModel"
+    _REPACK_MODEL_NAME_BASE = "RepackModel"
+
     def __init__(
         self,
         name: str,
@@ -168,7 +171,7 @@ def __init__(
             kwargs = dict(**kwargs, output_kms_key=kwargs.pop("model_kms_key", None))
 
             repack_model_step = _RepackModelStep(
-                name=f"{name}RepackModel",
+                name="{}-{}".format(self.name, self._REPACK_MODEL_NAME_BASE),
                 depends_on=depends_on,
                 retry_policies=repack_model_step_retry_policies,
                 sagemaker_session=estimator.sagemaker_session,
@@ -212,7 +215,7 @@ def __init__(
                     model_name = model_entity.name or model_entity._framework_name
 
                     repack_model_step = _RepackModelStep(
-                        name=f"{model_name}RepackModel",
+                        name="{}-{}".format(model_name, self._REPACK_MODEL_NAME_BASE),
                         depends_on=depends_on,
                         retry_policies=repack_model_step_retry_policies,
                         sagemaker_session=sagemaker_session,
@@ -256,7 +259,7 @@ def __init__(
             )
 
         register_model_step = _RegisterModelStep(
-            name=name,
+            name="{}-{}".format(self.name, self._REGISTER_MODEL_NAME_BASE),
             estimator=estimator,
             model_data=model_data,
             content_types=content_types,