Merge branch 'master' into deprecation

Author: mariumof

src/sagemaker/image_uri_config/djl-fastertransformer.json

@@ -0,0 +1,35 @@
+{
+    "scope": ["inference"],
+    "versions": {
+        "0.21.0": {
+            "registries": {
+                "af-south-1": "626614931356",
+                "ap-east-1": "871362719292",
+                "ap-northeast-1": "763104351884",
+                "ap-northeast-2": "763104351884",
+                "ap-northeast-3": "364406365360",
+                "ap-south-1": "763104351884",
+                "ap-southeast-1": "763104351884",
+                "ap-southeast-2": "763104351884",
+                "ap-southeast-3": "907027046896",
+                "ca-central-1": "763104351884",
+                "cn-north-1": "727897471807",
+                "cn-northwest-1": "727897471807",
+                "eu-central-1": "763104351884",
+                "eu-north-1": "763104351884",
+                "eu-west-1": "763104351884",
+                "eu-west-2": "763104351884",
+                "eu-west-3": "763104351884",
+                "eu-south-1": "692866216735",
+                "me-south-1": "217643126080",
+                "sa-east-1": "763104351884",
+                "us-east-1": "763104351884",
+                "us-east-2": "763104351884",
+                "us-west-1": "763104351884",
+                "us-west-2": "763104351884"
+            },
+            "repository": "djl-inference",
+            "tag_prefix": "0.21.0-fastertransformer5.3.0-cu117"
+        }
+    }
+}

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/image_uris/test_djl.py

@@ -41,19 +41,31 @@
     "us-west-1": "763104351884",
     "us-west-2": "763104351884",
 }
-VERSIONS = ["0.21.0", "0.20.0", "0.19.0"]
-DJL_FRAMEWORKS = ["djl-deepspeed"]
+DJL_DEEPSPEED_VERSIONS = ["0.21.0", "0.20.0", "0.19.0"]
+DJL_FASTERTRANSFORMER_VERSIONS = ["0.21.0"]
 DJL_VERSIONS_TO_FRAMEWORK = {
     "0.19.0": {"djl-deepspeed": "deepspeed0.7.3-cu113"},
     "0.20.0": {"djl-deepspeed": "deepspeed0.7.5-cu116"},
-    "0.21.0": {"djl-deepspeed": "deepspeed0.8.0-cu117"},
+    "0.21.0": {
+        "djl-deepspeed": "deepspeed0.8.0-cu117",
+        "djl-fastertransformer": "fastertransformer5.3.0-cu117",
+    },
 }
 
 
 @pytest.mark.parametrize("region", ACCOUNTS.keys())
-@pytest.mark.parametrize("version", VERSIONS)
-@pytest.mark.parametrize("djl_framework", DJL_FRAMEWORKS)
-def test_djl_uris(region, version, djl_framework):
+@pytest.mark.parametrize("version", DJL_DEEPSPEED_VERSIONS)
+def test_djl_deepspeed(region, version):
+    _test_djl_uris(region, version, "djl-deepspeed")
+
+
+@pytest.mark.parametrize("region", ACCOUNTS.keys())
+@pytest.mark.parametrize("version", DJL_FASTERTRANSFORMER_VERSIONS)
+def test_djl_fastertransformer(region, version):
+    _test_djl_uris(region, version, "djl-fastertransformer")
+
+
+def _test_djl_uris(region, version, djl_framework):
     uri = image_uris.retrieve(framework=djl_framework, region=region, version=version)
     expected = expected_uris.djl_framework_uri(
         "djl-inference",

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,