feature: Add EMRStep support in Sagemaker pipeline

Author: chenxy

src/sagemaker/workflow/emr_step.py

@@ -0,0 +1,119 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""The step definitions for workflow."""
+from __future__ import absolute_import
+
+from typing import List
+
+from sagemaker.workflow.entities import (
+    RequestType,
+)
+from sagemaker.workflow.properties import (
+    Properties,
+)
+from sagemaker.workflow.steps import Step, StepTypeEnum, CacheConfig
+
+
+class EMRStepConfig:
+    """Config for a Hadoop Jar step."""
+
+    def __init__(
+        self, jar, args: List[str] = None, main_class: str = None, properties: List[dict] = None
+    ):
+        """Create a definition for input data used by an EMR cluster(job flow) step.
+
+        See AWS documentation on the ``StepConfig`` API for more details on the parameters.
+
+        Args:
+            args(List[str]):
+                A list of command line arguments passed to
+                the JAR file's main function when executed.
+            jar(str): A path to a JAR file run during the step.
+            main_class(str): The name of the main class in the specified Java file.
+            properties(List(dict)): A list of key-value pairs that are set when the step runs.
+        """
+        self.jar = jar
+        self.args = args
+        self.main_class = main_class
+        self.properties = properties
+
+    def to_request(self) -> RequestType:
+        """Convert EMRStepConfig object to request dict."""
+        config = {"HadoopJarStep": {"Jar": self.jar}}
+        if self.args is not None:
+            config["HadoopJarStep"]["Args"] = self.args
+        if self.main_class is not None:
+            config["HadoopJarStep"]["MainClass"] = self.main_class
+        if self.properties is not None:
+            config["HadoopJarStep"]["Properties"] = self.properties
+
+        return config
+
+
+class EMRStep(Step):
+    """EMR step for workflow."""
+
+    def __init__(
+        self,
+        name: str,
+        display_name: str,
+        description: str,
+        job_flow_id: str,
+        step_config: EMRStepConfig,
+        depends_on: List[str] = None,
+        cache_config: CacheConfig = None,
+    ):
+        """Constructs a LambdaStep.
+
+        Args:
+            name(str): The name of the EMR step.
+            display_name(str): The display name of the EMR step.
+            description(str): The description of the EMR step.
+            job_flow_id(str): A string that uniquely identifies the job flow(cluster).
+            step_config(EMRStepConfig): One StepConfig to be executed by the job flow.
+            depends_on(List[str]):
+                A list of step names this `sagemaker.workflow.steps.EMRStep` depends on
+            cache_config(CacheConfig):  A `sagemaker.workflow.steps.CacheConfig` instance.
+
+        """
+        super(EMRStep, self).__init__(name, display_name, description, StepTypeEnum.EMR, depends_on)
+
+        emr_step_args = {"JobFlowId": job_flow_id, "StepConfig": step_config.to_request()}
+        self.args = emr_step_args
+        self.cache_config = cache_config
+
+        self._properties = Properties(
+            path=f"Steps.{name}", external_service_name="emr", shape_name="DescribeStepOutput"
+        )
+
+    @property
+    def arguments(self) -> RequestType:
+        """The arguments dict that is used to call `AddJobFlowSteps`.
+
+        NOTE: The AddFlowJobSteps request is not quite the args list that workflow needs.
+        The Name attribute in AddJobFlowSteps cannot be passed; it will be set during runtime.
+        In addition to that, we will also need to include emr job inputs and output config.
+        """
+        return self.args
+
+    @property
+    def properties(self) -> RequestType:
+        """A Properties object representing the EMR DescribeStepResponse model"""
+        return self._properties
+
+    def to_request(self) -> RequestType:
+        """Updates the dictionary with cache configuration."""
+        request_dict = super().to_request()
+        if self.cache_config:
+            request_dict.update(self.cache_config.config)
+        return request_dict

src/sagemaker/workflow/properties.py

@@ -26,14 +26,17 @@ class PropertiesMeta(type):
     """Load an internal shapes attribute from the botocore sagemaker service model."""
 
     _shapes = None
+    _emr_shapes = None
     _primitive_types = {"string", "boolean", "integer", "float"}
 
     def __new__(mcs, *args, **kwargs):
         """Loads up the shapes from the botocore sagemaker service model."""
         if mcs._shapes is None:
             loader = botocore.loaders.Loader()
             model = loader.load_service_model("sagemaker", "service-2")
+            emr_model = loader.load_service_model("emr", "service-2")
             mcs._shapes = model["shapes"]
+            mcs._emr_shapes = emr_model["shapes"]
         return super().__new__(mcs, *args, **kwargs)
 
 
@@ -45,6 +48,7 @@ def __init__(
         path: str,
         shape_name: str = None,
         shape_names: List[str] = None,
+        external_service_name: str = None,
     ):
         """Create a Properties instance representing the given shape.
 
@@ -57,15 +61,19 @@ def __init__(
         shape_names = [] if shape_names is None else shape_names
         self._shape_names = shape_names if shape_name is None else [shape_name] + shape_names
 
+        shapes = Properties._shapes
+        if external_service_name == "emr":
+            shapes = Properties._emr_shapes
+
         for name in self._shape_names:
-            shape = Properties._shapes.get(name, {})
+            shape = shapes.get(name, {})
             shape_type = shape.get("type")
             if shape_type in Properties._primitive_types:
                 self.__str__ = name
             elif shape_type == "structure":
                 members = shape["members"]
                 for key, info in members.items():
-                    if Properties._shapes.get(info["shape"], {}).get("type") == "list":
+                    if shapes.get(info["shape"], {}).get("type") == "list":
                         self.__dict__[key] = PropertiesList(f"{path}.{key}", info["shape"])
                     elif Properties._shapes.get(info["shape"], {}).get("type") == "map":
                         self.__dict__[key] = PropertiesMap(f"{path}.{key}", info["shape"])

src/sagemaker/workflow/steps.py

@@ -59,6 +59,7 @@ class StepTypeEnum(Enum, metaclass=DefaultEnumMeta):
     LAMBDA = "Lambda"
     QUALITY_CHECK = "QualityCheck"
     CLARIFY_CHECK = "ClarifyCheck"
+    EMR = "EMR"
 
 
 @attr.s

tests/integ/test_workflow.py

@@ -70,6 +70,7 @@
 from sagemaker.workflow.callback_step import CallbackStep, CallbackOutput, CallbackOutputTypeEnum
 from sagemaker.workflow.lambda_step import LambdaStep, LambdaOutput, LambdaOutputTypeEnum
 from sagemaker.workflow.properties import PropertyFile
+from sagemaker.workflow.emr_step import EMRStep, EMRStepConfig
 from sagemaker.wrangler.processing import DataWranglerProcessor
 from sagemaker.dataset_definition.inputs import DatasetDefinition, AthenaDatasetDefinition
 from sagemaker.workflow.execution_variables import ExecutionVariables
@@ -1148,6 +1149,112 @@ def test_two_step_lambda_pipeline_with_output_reference(
             pass
 
 
+def test_one_step_emr_pipeline(sagemaker_session, role, pipeline_name, region_name):
+    instance_count = ParameterInteger(name="InstanceCount", default_value=2)
+
+    emr_step_config = EMRStepConfig(
+        jar="s3:/script-runner/script-runner.jar",
+        args=["--arg_0", "arg_0_value"],
+        main_class="com.my.main",
+        properties=[{"Key": "Foo", "Value": "Foo_value"}, {"Key": "Bar", "Value": "Bar_value"}],
+    )
+
+    step_emr = EMRStep(
+        name="emr-step1",
+        job_flow_id="MyClusterID",
+        display_name="emr_step_1",
+        description="MyEMRStepDescription",
+        step_config=emr_step_config,
+    )
+
+    pipeline = Pipeline(
+        name=pipeline_name,
+        parameters=[instance_count],
+        steps=[step_emr],
+        sagemaker_session=sagemaker_session,
+    )
+
+    try:
+        response = pipeline.create(role)
+        create_arn = response["PipelineArn"]
+        assert re.match(
+            fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
+            create_arn,
+        )
+
+        pipeline.parameters = [ParameterInteger(name="InstanceCount", default_value=1)]
+        response = pipeline.update(role)
+        update_arn = response["PipelineArn"]
+        assert re.match(
+            fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
+            update_arn,
+        )
+    finally:
+        try:
+            pipeline.delete()
+        except Exception:
+            pass
+
+
+def test_two_steps_emr_pipeline_without_nullable_config_fields(
+    sagemaker_session, role, pipeline_name, region_name
+):
+    instance_count = ParameterInteger(name="InstanceCount", default_value=2)
+
+    emr_step_config_1 = EMRStepConfig(
+        jar="s3:/script-runner/script-runner_1.jar",
+        args=["--arg_0", "arg_0_value"],
+        main_class="com.my.main",
+        properties=[{"Key": "Foo", "Value": "Foo_value"}, {"Key": "Bar", "Value": "Bar_value"}],
+    )
+
+    step_emr_1 = EMRStep(
+        name="emr_step_1",
+        job_flow_id="MyClusterID",
+        display_name="emr_step_1",
+        description="MyEMRStepDescription",
+        step_config=emr_step_config_1,
+    )
+
+    emr_step_config_2 = EMRStepConfig(jar="s3:/script-runner/script-runner_2.jar")
+
+    step_emr_2 = EMRStep(
+        name="emr_step_2",
+        job_flow_id="MyClusterID",
+        display_name="emr_step_2",
+        description="MyEMRStepDescription",
+        step_config=emr_step_config_2,
+    )
+
+    pipeline = Pipeline(
+        name=pipeline_name,
+        parameters=[instance_count],
+        steps=[step_emr_1, step_emr_2],
+        sagemaker_session=sagemaker_session,
+    )
+
+    try:
+        response = pipeline.create(role)
+        create_arn = response["PipelineArn"]
+        assert re.match(
+            fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
+            create_arn,
+        )
+
+        pipeline.parameters = [ParameterInteger(name="InstanceCount", default_value=1)]
+        response = pipeline.update(role)
+        update_arn = response["PipelineArn"]
+        assert re.match(
+            fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
+            update_arn,
+        )
+    finally:
+        try:
+            pipeline.delete()
+        except Exception:
+            pass
+
+
 def test_conditional_pytorch_training_model_registration(
     sagemaker_session,
     role,