feature: Add EMRStep support in Sagemaker pipeline

Author: chenxy

src/sagemaker/workflow/properties.py

@@ -76,9 +76,13 @@ def __init__(
                 members = shape["members"]
                 for key, info in members.items():
                     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"])
+                        self.__dict__[key] = PropertiesList(
+                            f"{path}.{key}", info["shape"], service_name
+                        )
+                    elif shapes.get(info["shape"], {}).get("type") == "map":
+                        self.__dict__[key] = PropertiesMap(
+                            f"{path}.{key}", info["shape"], service_name
+                        )
                     else:
                         self.__dict__[key] = Properties(
                             f"{path}.{key}", info["shape"], service_name=service_name
@@ -127,15 +131,17 @@ def __getitem__(self, item: Union[int, str]):
 class PropertiesMap(Properties):
     """PropertiesMap for use in workflow expressions."""
 
-    def __init__(self, path: str, shape_name: str = None):
+    def __init__(self, path: str, shape_name: str = None, service_name: str = "sagemaker"):
         """Create a PropertiesMap instance representing the given shape.
 
         Args:
             path (str): The parent path of the PropertiesMap instance.
             shape_name (str): The botocore sagemaker service model shape name.
+            service_name (str): The botocore service name.
         """
         super(PropertiesMap, self).__init__(path, shape_name)
         self.shape_name = shape_name
+        self.service_name = service_name
         self._items: Dict[Union[int, str], Properties] = dict()
 
     def __getitem__(self, item: Union[int, str]):
@@ -145,7 +151,7 @@ def __getitem__(self, item: Union[int, str]):
             item (Union[int, str]): The index of the item in sequence.
         """
         if item not in self._items.keys():
-            shape = Properties._shapes.get(self.shape_name)
+            shape = Properties._shapes_map.get(self.service_name, {}).get(self.shape_name)
             member = shape["value"]["shape"]
             if isinstance(item, str):
                 property_item = Properties(f"{self._path}['{item}']", member)

tests/data/workflow/emr-script.sh

@@ -0,0 +1,2 @@
+echo "This is emr test script..."
+sleep 15

tests/integ/test_workflow.py

@@ -76,6 +76,7 @@
 from sagemaker.feature_store.feature_group import FeatureGroup, FeatureDefinition, FeatureTypeEnum
 from tests.integ import DATA_DIR
 from tests.integ.kms_utils import get_or_create_kms_key
+from tests.integ.vpc_test_utils import get_or_create_vpc_resources
 
 
 def ordered(obj):
@@ -261,6 +262,75 @@ def build_jar():
     subprocess.run(["rm", os.path.join(jar_file_path, java_file_path, "HelloJavaSparkApp.class")])
 
 
+@pytest.fixture(scope="module")
+def emr_script_path(sagemaker_session):
+    input_path = sagemaker_session.upload_data(
+        path=os.path.join(DATA_DIR, "workflow", "emr-script.sh"),
+        key_prefix="integ-test-data/workflow",
+    )
+    return input_path
+
+
+@pytest.fixture(scope="module")
+def emr_cluster_id(sagemaker_session, role):
+    emr_client = sagemaker_session.boto_session.client("emr")
+    cluster_name = "emr-step-test-cluster"
+    cluster_id = get_existing_emr_cluster_id(emr_client, cluster_name)
+
+    if cluster_id is None:
+        create_new_emr_cluster(sagemaker_session, emr_client, cluster_name)
+    return cluster_id
+
+
+def get_existing_emr_cluster_id(emr_client, cluster_name):
+    try:
+        response = emr_client.list_clusters(ClusterStates=["RUNNING", "WAITING"])
+        for cluster in response["Clusters"]:
+            if cluster["Name"].startswith(cluster_name):
+                cluster_id = cluster["Id"]
+                print("Using existing cluster: {}".format(cluster_id))
+                return cluster_id
+    except Exception:
+        raise
+
+
+def create_new_emr_cluster(sagemaker_session, emr_client, cluster_name):
+    ec2_client = sagemaker_session.boto_session.client("ec2")
+    subnet_ids, security_group_id = get_or_create_vpc_resources(ec2_client)
+    try:
+        response = emr_client.run_job_flow(
+            Name="emr-step-test-cluster",
+            LogUri="s3://{}/{}".format(sagemaker_session.default_bucket(), "emr-test-logs"),
+            ReleaseLabel="emr-6.3.0",
+            Applications=[
+                {"Name": "Hadoop"},
+                {"Name": "Spark"},
+            ],
+            Instances={
+                "InstanceGroups": [
+                    {
+                        "Name": "Master nodes",
+                        "Market": "ON_DEMAND",
+                        "InstanceRole": "MASTER",
+                        "InstanceType": "m4.large",
+                        "InstanceCount": 1,
+                    }
+                ],
+                "KeepJobFlowAliveWhenNoSteps": True,
+                "TerminationProtected": False,
+                "Ec2SubnetId": subnet_ids[0],
+            },
+            VisibleToAllUsers=True,
+            JobFlowRole="EMR_EC2_DefaultRole",
+            ServiceRole="EMR_DefaultRole",
+        )
+        cluster_id = response["JobFlowId"]
+        print("Created new cluster: {}".format(cluster_id))
+        return cluster_id
+    except Exception:
+        raise
+
+
 def test_three_step_definition(
     sagemaker_session,
     region_name,
@@ -1129,82 +1199,30 @@ 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-step",
-        cluster_id="MyClusterID",
-        display_name="emr_step",
-        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"]
-
-        execution = pipeline.start()
-        response = execution.describe()
-        assert response["PipelineArn"] == create_arn
-
-        try:
-            execution.wait(delay=60, max_attempts=10)
-        except WaiterError:
-            pass
-
-        execution_steps = execution.list_steps()
-        assert len(execution_steps) == 1
-        assert execution_steps[0]["StepName"] == "emr-step"
-    finally:
-        try:
-            pipeline.delete()
-        except Exception:
-            pass
-
-
-def test_two_steps_emr_pipeline_without_nullable_config_fields(
-    sagemaker_session, role, pipeline_name, region_name
+def test_two_steps_emr_pipeline(
+    sagemaker_session, role, pipeline_name, region_name, emr_cluster_id, emr_script_path
 ):
     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"}],
+    emr_step_config = EMRStepConfig(
+        jar="s3://us-west-2.elasticmapreduce/libs/script-runner/script-runner.jar",
+        args=[emr_script_path],
     )
 
     step_emr_1 = EMRStep(
         name="emr-step-1",
-        cluster_id="MyClusterID",
-        display_name="emr-step-1",
+        cluster_id=emr_cluster_id,
+        display_name="emr_step_1",
         description="MyEMRStepDescription",
-        step_config=emr_step_config_1,
+        step_config=emr_step_config,
     )
 
-    emr_step_config_2 = EMRStepConfig(jar="s3:/script-runner/script-runner_2.jar")
-
     step_emr_2 = EMRStep(
         name="emr-step-2",
-        cluster_id="MyClusterID",
-        display_name="emr-step-2",
+        cluster_id=step_emr_1.properties.ClusterId,
+        display_name="emr_step_2",
         description="MyEMRStepDescription",
-        step_config=emr_step_config_2,
+        step_config=emr_step_config,
     )
 
     pipeline = Pipeline(
@@ -1217,20 +1235,24 @@ def test_two_steps_emr_pipeline_without_nullable_config_fields(
     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
+        )
 
         execution = pipeline.start()
-        response = execution.describe()
-        assert response["PipelineArn"] == create_arn
-
         try:
-            execution.wait(delay=60, max_attempts=10)
+            execution.wait(delay=60, max_attempts=5)
         except WaiterError:
             pass
 
         execution_steps = execution.list_steps()
         assert len(execution_steps) == 2
         assert execution_steps[0]["StepName"] == "emr-step-1"
+        assert execution_steps[0].get("FailureReason", "") == ""
+        assert execution_steps[0]["StepStatus"] == "Succeeded"
         assert execution_steps[1]["StepName"] == "emr-step-2"
+        assert execution_steps[1].get("FailureReason", "") == ""
+        assert execution_steps[1]["StepStatus"] == "Succeeded"
 
         pipeline.parameters = [ParameterInteger(name="InstanceCount", default_value=1)]
         response = pipeline.update(role)