feature: Add EMRStep support in Sagemaker pipeline

chenxy · EthanShouhanCheng · commit a71105b3f622 · 2022-01-11T17:27:00.000-08:00
diff --git a/src/sagemaker/workflow/emr_step.py b/src/sagemaker/workflow/emr_step.py
@@ -68,7 +68,11 @@ def __init__(
         name: str,
         display_name: str,
         description: str,
+<<<<<<< HEAD
         job_flow_id: str,
+=======
+        cluster_id: str,
+>>>>>>> feature: Add EMRStep support in Sagemaker pipeline
         step_config: EMRStepConfig,
         depends_on: List[str] = None,
         cache_config: CacheConfig = None,
@@ -79,7 +83,11 @@ def __init__(
             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.
+<<<<<<< HEAD
             job_flow_id(str): A string that uniquely identifies the job flow(cluster).
+=======
+            cluster_id(str): A string that uniquely identifies the cluster.
+>>>>>>> feature: Add EMRStep support in Sagemaker pipeline
             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
@@ -88,13 +96,23 @@ def __init__(
         """
         super(EMRStep, self).__init__(name, display_name, description, StepTypeEnum.EMR, depends_on)
 
+<<<<<<< HEAD
         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"
         )
+=======
+        emr_step_args = {"ClusterId": cluster_id, "StepConfig": step_config.to_request()}
+        self.args = emr_step_args
+        self.cache_config = cache_config
+
+        root_property = Properties(path=f"Steps.{name}", shape_name="Step", service_name="emr")
+        root_property.__dict__["ClusterId"] = cluster_id
+        self._properties = root_property
+>>>>>>> feature: Add EMRStep support in Sagemaker pipeline
 
     @property
     def arguments(self) -> RequestType:
diff --git a/src/sagemaker/workflow/properties.py b/src/sagemaker/workflow/properties.py
@@ -23,20 +23,34 @@
 
 
 class PropertiesMeta(type):
-    """Load an internal shapes attribute from the botocore sagemaker service model."""
+    """Load an internal shapes attribute from the botocore service model
 
+<<<<<<< HEAD
     _shapes = None
     _emr_shapes = None
+=======
+    for sagemaker and emr service.
+    """
+
+    _shapes_map = dict()
+>>>>>>> feature: Add EMRStep support in Sagemaker pipeline
     _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:
+        if len(mcs._shapes_map.keys()) == 0:
             loader = botocore.loaders.Loader()
+<<<<<<< HEAD
             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"]
+=======
+            sagemaker_model = loader.load_service_model("sagemaker", "service-2")
+            emr_model = loader.load_service_model("emr", "service-2")
+            mcs._shapes_map["sagemaker"] = sagemaker_model["shapes"]
+            mcs._shapes_map["emr"] = emr_model["shapes"]
+>>>>>>> feature: Add EMRStep support in Sagemaker pipeline
         return super().__new__(mcs, *args, **kwargs)
 
 
@@ -48,7 +62,11 @@ def __init__(
         path: str,
         shape_name: str = None,
         shape_names: List[str] = None,
+<<<<<<< HEAD
         external_service_name: str = None,
+=======
+        service_name: str = "sagemaker",
+>>>>>>> feature: Add EMRStep support in Sagemaker pipeline
     ):
         """Create a Properties instance representing the given shape.
 
@@ -61,9 +79,13 @@ 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
 
+<<<<<<< HEAD
         shapes = Properties._shapes
         if external_service_name == "emr":
             shapes = Properties._emr_shapes
+=======
+        shapes = Properties._shapes_map.get(service_name, {})
+>>>>>>> feature: Add EMRStep support in Sagemaker pipeline
 
         for name in self._shape_names:
             shape = shapes.get(name, {})
@@ -78,7 +100,9 @@ def __init__(
                     elif Properties._shapes.get(info["shape"], {}).get("type") == "map":
                         self.__dict__[key] = PropertiesMap(f"{path}.{key}", info["shape"])
                     else:
-                        self.__dict__[key] = Properties(f"{path}.{key}", info["shape"])
+                        self.__dict__[key] = Properties(
+                            f"{path}.{key}", info["shape"], service_name=service_name
+                        )
 
     @property
     def expr(self):
@@ -89,16 +113,17 @@ def expr(self):
 class PropertiesList(Properties):
     """PropertiesList 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 PropertiesList instance representing the given shape.
 
         Args:
             path (str): The parent path of the PropertiesList instance.
-            shape_name (str): The botocore sagemaker service model shape name.
-            root_shape_name (str): The botocore sagemaker service model shape name.
+            shape_name (str): The botocore service model shape name.
+            service_name (str): The botocore service name.
         """
         super(PropertiesList, 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]):
@@ -108,7 +133,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["member"]["shape"]
             if isinstance(item, str):
                 property_item = Properties(f"{self._path}['{item}']", member)
diff --git a/tests/integ/test_workflow.py b/tests/integ/test_workflow.py
@@ -1160,9 +1160,9 @@ def test_one_step_emr_pipeline(sagemaker_session, role, pipeline_name, region_na
     )
 
     step_emr = EMRStep(
-        name="emr-step1",
-        job_flow_id="MyClusterID",
-        display_name="emr_step_1",
+        name="emr-step",
+        cluster_id="MyClusterID",
+        display_name="emr_step",
         description="MyEMRStepDescription",
         step_config=emr_step_config,
     )
@@ -1177,18 +1177,19 @@ def test_one_step_emr_pipeline(sagemaker_session, role, pipeline_name, region_na
     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,
-        )
+        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()
@@ -1209,19 +1210,19 @@ def test_two_steps_emr_pipeline_without_nullable_config_fields(
     )
 
     step_emr_1 = EMRStep(
-        name="emr_step_1",
-        job_flow_id="MyClusterID",
-        display_name="emr_step_1",
+        name="emr-step-1",
+        cluster_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",
+        name="emr-step-2",
+        cluster_id="MyClusterID",
+        display_name="emr-step-2",
         description="MyEMRStepDescription",
         step_config=emr_step_config_2,
     )
@@ -1236,10 +1237,20 @@ 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)
+        except WaiterError:
+            pass
+
+        execution_steps = execution.list_steps()
+        assert len(execution_steps) == 2
+        assert execution_steps[0]["StepName"] == "emr-step-1"
+        assert execution_steps[1]["StepName"] == "emr-step-2"
 
         pipeline.parameters = [ParameterInteger(name="InstanceCount", default_value=1)]
         response = pipeline.update(role)
diff --git a/tests/unit/sagemaker/workflow/test_emr_step.py b/tests/unit/sagemaker/workflow/test_emr_step.py
@@ -19,7 +19,10 @@
 from mock import Mock
 
 from sagemaker.workflow.emr_step import EMRStep, EMRStepConfig
+<<<<<<< HEAD
 from sagemaker.workflow.properties import Properties
+=======
+>>>>>>> feature: Add EMRStep support in Sagemaker pipeline
 from sagemaker.workflow.steps import CacheConfig
 from sagemaker.workflow.pipeline import Pipeline
 from sagemaker.workflow.parameters import ParameterString
@@ -50,7 +53,11 @@ def test_emr_step_with_one_step_config(sagemaker_session):
         name="MyEMRStep",
         display_name="MyEMRStep",
         description="MyEMRStepDescription",
+<<<<<<< HEAD
         job_flow_id="MyClusterID",
+=======
+        cluster_id="MyClusterID",
+>>>>>>> feature: Add EMRStep support in Sagemaker pipeline
         step_config=emr_step_config,
         depends_on=["TestStep"],
         cache_config=CacheConfig(enable_caching=True, expire_after="PT1H"),
@@ -60,7 +67,11 @@ def test_emr_step_with_one_step_config(sagemaker_session):
         "Name": "MyEMRStep",
         "Type": "EMR",
         "Arguments": {
+<<<<<<< HEAD
             "JobFlowId": "MyClusterID",
+=======
+            "ClusterId": "MyClusterID",
+>>>>>>> feature: Add EMRStep support in Sagemaker pipeline
             "StepConfig": {
                 "HadoopJarStep": {
                     "Args": ["--arg_0", "arg_0_value"],
@@ -78,10 +89,27 @@ def test_emr_step_with_one_step_config(sagemaker_session):
         "Description": "MyEMRStepDescription",
         "CacheConfig": {"Enabled": True, "ExpireAfter": "PT1H"},
     }
+<<<<<<< HEAD
     assert emr_step.properties.Step.expr == {"Get": "Steps.MyEMRStep.Step"}
 
 
 def test_pipeline_interpolates_lambda_outputs(sagemaker_session):
+=======
+    assert emr_step.properties.ClusterId == "MyClusterID"
+    assert emr_step.properties.ActionOnFailure.expr == {"Get": "Steps.MyEMRStep.ActionOnFailure"}
+    assert emr_step.properties.Config.Args.expr == {"Get": "Steps.MyEMRStep.Config.Args"}
+    assert emr_step.properties.Config.Jar.expr == {"Get": "Steps.MyEMRStep.Config.Jar"}
+    assert emr_step.properties.Config.MainClass.expr == {"Get": "Steps.MyEMRStep.Config.MainClass"}
+    assert emr_step.properties.Id.expr == {"Get": "Steps.MyEMRStep.Id"}
+    assert emr_step.properties.Name.expr == {"Get": "Steps.MyEMRStep.Name"}
+    assert emr_step.properties.Status.State.expr == {"Get": "Steps.MyEMRStep.Status.State"}
+    assert emr_step.properties.Status.FailureDetails.Reason.expr == {
+        "Get": "Steps.MyEMRStep.Status.FailureDetails.Reason"
+    }
+
+
+def test_pipeline_interpolates_emr_outputs(sagemaker_session):
+>>>>>>> feature: Add EMRStep support in Sagemaker pipeline
     parameter = ParameterString("MyStr")
 
     emr_step_config_1 = EMRStepConfig(
@@ -93,7 +121,11 @@ def test_pipeline_interpolates_lambda_outputs(sagemaker_session):
 
     step_emr_1 = EMRStep(
         name="emr_step_1",
+<<<<<<< HEAD
         job_flow_id="MyClusterID",
+=======
+        cluster_id="MyClusterID",
+>>>>>>> feature: Add EMRStep support in Sagemaker pipeline
         display_name="emr_step_1",
         description="MyEMRStepDescription",
         depends_on=["TestStep"],
@@ -104,7 +136,11 @@ def test_pipeline_interpolates_lambda_outputs(sagemaker_session):
 
     step_emr_2 = EMRStep(
         name="emr_step_2",
+<<<<<<< HEAD
         job_flow_id="MyClusterID",
+=======
+        cluster_id="MyClusterID",
+>>>>>>> feature: Add EMRStep support in Sagemaker pipeline
         display_name="emr_step_2",
         description="MyEMRStepDescription",
         depends_on=["TestStep"],
@@ -131,7 +167,11 @@ def test_pipeline_interpolates_lambda_outputs(sagemaker_session):
                 "Name": "emr_step_1",
                 "Type": "EMR",
                 "Arguments": {
+<<<<<<< HEAD
                     "JobFlowId": "MyClusterID",
+=======
+                    "ClusterId": "MyClusterID",
+>>>>>>> feature: Add EMRStep support in Sagemaker pipeline
                     "StepConfig": {
                         "HadoopJarStep": {
                             "Args": ["--arg_0", "arg_0_value"],
@@ -152,7 +192,11 @@ def test_pipeline_interpolates_lambda_outputs(sagemaker_session):
                 "Name": "emr_step_2",
                 "Type": "EMR",
                 "Arguments": {
+<<<<<<< HEAD
                     "JobFlowId": "MyClusterID",
+=======
+                    "ClusterId": "MyClusterID",
+>>>>>>> feature: Add EMRStep support in Sagemaker pipeline
                     "StepConfig": {
                         "HadoopJarStep": {"Jar": "s3:/script-runner/script-runner_2.jar"}
                     },
@@ -163,6 +207,7 @@ def test_pipeline_interpolates_lambda_outputs(sagemaker_session):
             },
         ],
     }
+<<<<<<< HEAD
 
 
 def test_properties_describe_step_output():
@@ -172,3 +217,5 @@ def test_properties_describe_step_output():
     for name in some_prop_names:
         assert name in prop.__dict__.keys()
     assert prop.Step.expr == {"Get": "Steps.MyStep.Step"}
+=======
+>>>>>>> feature: Add EMRStep support in Sagemaker pipeline
diff --git a/tests/unit/sagemaker/workflow/test_properties.py b/tests/unit/sagemaker/workflow/test_properties.py
@@ -71,12 +71,25 @@ def test_properties_tuning_job():
 
 
 def test_properties_emr_step():
+<<<<<<< HEAD
     prop = Properties("Steps.MyStep", "DescribeStepOutput", external_service_name="emr")
     some_prop_names = ["Step"]
     for name in some_prop_names:
         assert name in prop.__dict__.keys()
 
     assert prop.Step.expr == {"Get": "Steps.MyStep.Step"}
+=======
+    prop = Properties("Steps.MyStep", "Step", service_name="emr")
+    some_prop_names = ["Id", "Name", "Config", "ActionOnFailure", "Status"]
+    for name in some_prop_names:
+        assert name in prop.__dict__.keys()
+
+    assert prop.Id.expr == {"Get": "Steps.MyStep.Id"}
+    assert prop.Name.expr == {"Get": "Steps.MyStep.Name"}
+    assert prop.ActionOnFailure.expr == {"Get": "Steps.MyStep.ActionOnFailure"}
+    assert prop.Config.Jar.expr == {"Get": "Steps.MyStep.Config.Jar"}
+    assert prop.Status.State.expr == {"Get": "Steps.MyStep.Status.State"}
+>>>>>>> feature: Add EMRStep support in Sagemaker pipeline
 
 
 def test_properties_describe_model_package_output():
