Merge branch 'master' of https://github.com/rahven14/sagemaker-python-sdk into feature/infer-framework-and-version

Author: rahven14

requirements/extras/test_requirements.txt

@@ -12,9 +12,10 @@ awslogs==0.14.0
 black==22.3.0
 stopit==1.1.2
 apache-airflow==2.2.4
-apache-airflow-providers-amazon==3.0.0
+apache-airflow-providers-amazon==4.0.0
 attrs==20.3.0
 fabric==2.6.0
 requests==2.27.1
 sagemaker-experiments==0.1.35
 Jinja2==3.0.3
+pandas>=1.3.5,<1.5

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/utils.py

@@ -759,32 +759,55 @@ def update_container_with_inference_params(
         dict: dict with inference recommender params
     """
 
-    if (
-        framework is not None
-        and framework_version is not None
-        and nearest_model_name is not None
-        and data_input_configuration is not None
-    ):
+    if framework is not None and framework_version is not None and nearest_model_name is not None:
         if container_list is not None:
             for obj in container_list:
-                obj.update(
-                    {
-                        "Framework": framework,
-                        "FrameworkVersion": framework_version,
-                        "NearestModelName": nearest_model_name,
-                        "ModelInput": {
-                            "DataInputConfig": data_input_configuration,
-                        },
-                    }
+                construct_container_object(
+                    obj, data_input_configuration, framework, framework_version, nearest_model_name
                 )
         if container_obj is not None:
-            container_obj.update(
-                {
-                    "Framework": framework,
-                    "FrameworkVersion": framework_version,
-                    "NearestModelName": nearest_model_name,
-                    "ModelInput": {
-                        "DataInputConfig": data_input_configuration,
-                    },
-                }
+            construct_container_object(
+                container_obj,
+                data_input_configuration,
+                framework,
+                framework_version,
+                nearest_model_name,
             )
+
+
+def construct_container_object(
+    obj, data_input_configuration, framework, framework_version, nearest_model_name
+):
+    """Function to construct container object.
+
+    Args:
+        framework (str): Machine learning framework of the model package container image
+                (default: None).
+        framework_version (str): Framework version of the Model Package Container Image
+            (default: None).
+        nearest_model_name (str): Name of a pre-trained machine learning benchmarked by
+            Amazon SageMaker Inference Recommender (default: None).
+        data_input_configuration (str): Input object for the model (default: None).
+        container_obj (dict): object to be updated.
+        container_list (list): list to be updated.
+
+    Returns:
+        dict: container object
+    """
+
+    obj.update(
+        {
+            "Framework": framework,
+            "FrameworkVersion": framework_version,
+            "NearestModelName": nearest_model_name,
+        }
+    )
+
+    if data_input_configuration is not None:
+        obj.update(
+            {
+                "ModelInput": {
+                    "DataInputConfig": data_input_configuration,
+                },
+            }
+        )

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")