feature: Add EMRStep support in Sagemaker pipeline

Author: chenxy

src/sagemaker/workflow/emr_step.py

@@ -68,11 +68,7 @@ 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,
@@ -83,11 +79,7 @@ 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
@@ -96,23 +88,13 @@ 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:

src/sagemaker/workflow/properties.py

@@ -25,32 +25,22 @@
 class PropertiesMeta(type):
     """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 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)
 
 
@@ -62,11 +52,7 @@ 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.
 
@@ -79,13 +65,7 @@ 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, {})

tests/unit/sagemaker/workflow/test_emr_step.py

@@ -19,10 +19,6 @@
 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
@@ -53,11 +49,7 @@ 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"),
@@ -67,11 +59,7 @@ 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"],
@@ -89,12 +77,7 @@ 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"}
@@ -109,7 +92,6 @@ def test_pipeline_interpolates_lambda_outputs(sagemaker_session):
 
 
 def test_pipeline_interpolates_emr_outputs(sagemaker_session):
->>>>>>> feature: Add EMRStep support in Sagemaker pipeline
     parameter = ParameterString("MyStr")
 
     emr_step_config_1 = EMRStepConfig(
@@ -121,11 +103,7 @@ def test_pipeline_interpolates_emr_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"],
@@ -136,11 +114,7 @@ def test_pipeline_interpolates_emr_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"],
@@ -167,11 +141,7 @@ def test_pipeline_interpolates_emr_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"],
@@ -192,11 +162,7 @@ def test_pipeline_interpolates_emr_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"}
                     },
@@ -207,15 +173,3 @@ def test_pipeline_interpolates_emr_outputs(sagemaker_session):
             },
         ],
     }
-<<<<<<< HEAD
-
-
-def test_properties_describe_step_output():
-    prop = Properties("Steps.MyStep", "DescribeStepOutput", external_service_name="emr")
-    some_prop_names = ["Step"]
-    print(prop.__dict__.keys())
-    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

tests/unit/sagemaker/workflow/test_properties.py

@@ -71,14 +71,6 @@ 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:
@@ -89,7 +81,6 @@ def test_properties_emr_step():
     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():