feature: EMR step runtime role support (#3703)

Co-authored-by: Thomas Brooks <thobrok@amazon.com>

Author: thbrooks22

src/sagemaker/workflow/emr_step.py

@@ -97,6 +97,14 @@ def to_request(self) -> RequestType:
     "must be explicitly set to None."
 )
 
+ERR_STR_WITH_EXEC_ROLE_ARN_AND_WITHOUT_CLUSTER_ID = (
+    "EMRStep {step_name} cannot have execution_role_arn"
+    "without cluster_id."
+    "To use EMRStep with "
+    "execution_role_arn, cluster_id "
+    "must not be None."
+)
+
 ERR_STR_WITHOUT_CLUSTER_ID_AND_CLUSTER_CFG = (
     "EMRStep {step_name} must have either cluster_id or cluster_config"
 )
@@ -155,6 +163,7 @@ def __init__(
         depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
         cache_config: CacheConfig = None,
         cluster_config: Dict[str, Any] = None,
+        execution_role_arn: str = None,
     ):
         """Constructs an `EMRStep`.
 
@@ -185,7 +194,11 @@ def __init__(
                 https://docs.aws.amazon.com/emr/latest/APIReference/API_RunJobFlow.html.
                 Note that if you want to use ``cluster_config``, then you have to set
                 ``cluster_id`` as None.
-
+            execution_role_arn(str): The ARN of the runtime role assumed by this `EMRStep`. The
+                job submitted to your EMR cluster uses this role to access AWS resources. This
+                value is passed as ExecutionRoleArn to the AddJobFlowSteps request (an EMR request)
+                called on the cluster specified by ``cluster_id``, so you can only include this
+                field if ``cluster_id`` is not None.
         """
         super(EMRStep, self).__init__(name, display_name, description, StepTypeEnum.EMR, depends_on)
 
@@ -198,9 +211,18 @@ def __init__(
         if cluster_id is not None and cluster_config is not None:
             raise ValueError(ERR_STR_WITH_BOTH_CLUSTER_ID_AND_CLUSTER_CFG.format(step_name=name))
 
+        if execution_role_arn is not None and cluster_id is None:
+            raise ValueError(
+                ERR_STR_WITH_EXEC_ROLE_ARN_AND_WITHOUT_CLUSTER_ID.format(step_name=name)
+            )
+
         if cluster_id is not None:
             emr_step_args["ClusterId"] = cluster_id
             root_property.__dict__["ClusterId"] = cluster_id
+
+            if execution_role_arn is not None:
+                emr_step_args["ExecutionRoleArn"] = execution_role_arn
+                root_property.__dict__["ExecutionRoleArn"] = execution_role_arn
         elif cluster_config is not None:
             self._validate_cluster_config(cluster_config, name)
             emr_step_args["ClusterConfig"] = cluster_config

tests/unit/sagemaker/workflow/test_emr_step.py

@@ -24,14 +24,16 @@
     ERR_STR_BOTH_OR_NONE_INSTANCEGROUPS_OR_INSTANCEFLEETS,
     ERR_STR_WITH_BOTH_CLUSTER_ID_AND_CLUSTER_CFG,
     ERR_STR_WITHOUT_CLUSTER_ID_AND_CLUSTER_CFG,
+    ERR_STR_WITH_EXEC_ROLE_ARN_AND_WITHOUT_CLUSTER_ID,
 )
 from sagemaker.workflow.steps import CacheConfig
 from sagemaker.workflow.pipeline import Pipeline, PipelineGraph
 from sagemaker.workflow.parameters import ParameterString
 from tests.unit.sagemaker.workflow.helpers import CustomStep, ordered
 
 
-def test_emr_step_with_one_step_config(sagemaker_session):
+@pytest.mark.parametrize("execution_role_arn", [None, "arn:aws:iam:000000000000:role/runtime-role"])
+def test_emr_step_with_one_step_config(sagemaker_session, execution_role_arn):
     emr_step_config = EMRStepConfig(
         jar="s3:/script-runner/script-runner.jar",
         args=["--arg_0", "arg_0_value"],
@@ -47,9 +49,11 @@ def test_emr_step_with_one_step_config(sagemaker_session):
         step_config=emr_step_config,
         depends_on=["TestStep"],
         cache_config=CacheConfig(enable_caching=True, expire_after="PT1H"),
+        execution_role_arn=execution_role_arn,
     )
     emr_step.add_depends_on(["SecondTestStep"])
-    assert emr_step.to_request() == {
+
+    expected_request = {
         "Name": "MyEMRStep",
         "Type": "EMR",
         "Arguments": {
@@ -72,7 +76,16 @@ def test_emr_step_with_one_step_config(sagemaker_session):
         "CacheConfig": {"Enabled": True, "ExpireAfter": "PT1H"},
     }
 
+    if execution_role_arn is not None:
+        expected_request["Arguments"]["ExecutionRoleArn"] = execution_role_arn
+
+    assert emr_step.to_request() == expected_request
     assert emr_step.properties.ClusterId == "MyClusterID"
+    assert (
+        emr_step.properties.ExecutionRoleArn == execution_role_arn
+        if execution_role_arn is not None
+        else True
+    )
     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"}
@@ -239,6 +252,27 @@ def test_emr_step_throws_exception_when_both_cluster_id_and_cluster_config_are_n
     assert actual_error_msg == expected_error_msg
 
 
+def test_emr_step_throws_exception_when_both_execution_role_arn_and_cluster_config_are_present():
+    with pytest.raises(ValueError) as exceptionInfo:
+        EMRStep(
+            name=g_emr_step_name,
+            display_name="MyEMRStep",
+            description="MyEMRStepDescription",
+            step_config=g_emr_step_config,
+            cluster_id=None,
+            cluster_config=g_cluster_config,
+            depends_on=["TestStep"],
+            cache_config=CacheConfig(enable_caching=True, expire_after="PT1H"),
+            execution_role_arn="arn:aws:iam:000000000000:role/some-role",
+        )
+    expected_error_msg = ERR_STR_WITH_EXEC_ROLE_ARN_AND_WITHOUT_CLUSTER_ID.format(
+        step_name=g_emr_step_name
+    )
+    actual_error_msg = exceptionInfo.value.args[0]
+
+    assert actual_error_msg == expected_error_msg
+
+
 def test_emr_step_with_valid_cluster_config():
     emr_step = EMRStep(
         name=g_emr_step_name,