Addressed aoguo64's comments

(cherry picked from commit c2cb68e9b44d3600537daa992668a38f7529da72)

Author: Shegufta Ahsan

src/sagemaker/workflow/emr_step.py

@@ -61,62 +61,89 @@ def to_request(self) -> RequestType:
         return config
 
 
-def validate_cluster_config(cluster_config, step_name):
-    """Validates user provided cluster_config.
-
-    Args:
-        cluster_config(Union[Dict[str, Any], List[Dict[str, Any]]]):
-            user provided cluster configuration.
-        step_name: The name of the EMR step.
-    """
-
-    instances = "Instances"
-    instancegroups = "InstanceGroups"
-    instancefleets = "InstanceFleets"
-    prefix_with_in = "In EMRStep " + step_name + ", "
-
-    if (
-        "Name" in cluster_config
-        or "AutoTerminationPolicy" in cluster_config
-        or "Steps" in cluster_config
-    ):
-        raise Exception(
-            prefix_with_in + "cluster_config should not contain any of Name, "
-            "AutoTerminationPolicy and/or Steps"
-        )
+instances = "Instances"
+instancegroups = "InstanceGroups"
+instancefleets = "InstanceFleets"
+err_str_with_name_auto_termination_or_steps = (
+    "In EMRStep {step_name}, cluster_config "
+    "should not contain any of the Name, "
+    "AutoTerminationPolicy and/or Steps."
+)
 
-    if instances not in cluster_config:
-        raise Exception(prefix_with_in + "cluster_config must contain Instances")
+err_str_without_instance = "In EMRStep {step_name}, cluster_config must contain " + instances + "."
 
-    if (
-        "KeepJobFlowAliveWhenNoSteps" in cluster_config[instances]
-        or "TerminationProtected" in cluster_config[instances]
-    ):
-        raise Exception(
-            prefix_with_in + instances + " should not contain "
-            "KeepJobFlowAliveWhenNoSteps or "
-            "TerminationProtected"
-        )
-
-    if (
-        instancegroups in cluster_config[instances] and instancefleets in cluster_config[instances]
-    ) or (
-        instancegroups not in cluster_config[instances]
-        and instancefleets not in cluster_config[instances]
-    ):
-        raise Exception(
-            prefix_with_in
-            + instances
-            + " should contain either "
-            + instancegroups
-            + " or "
-            + instancefleets
-        )
+err_str_with_keepjobflow_or_terminationprotected = (
+    "In EMRStep {step_name}, " + instances + " should not contain "
+    "KeepJobFlowAliveWhenNoSteps or "
+    "TerminationProtected."
+)
+
+err_str_both_or_none_instancegroups_or_instancefleets = (
+    "In EMRStep {step_name}, "
+    + instances
+    + " should contain either "
+    + instancegroups
+    + " or "
+    + instancefleets
+    + "."
+)
+
+err_str_with_both_cluster_id_and_cluster_cfg = (
+    "EMRStep {step_name} can not have both cluster_id"
+    "or cluster_config."
+    "To use EMRStep with "
+    "cluster_config, cluster_id "
+    "must be explicitly set to None."
+)
+
+err_str_without_cluster_id_and_cluster_cfg = (
+    "EMRStep {step_name} must have either cluster_id or cluster_config"
+)
 
 
 class EMRStep(Step):
     """EMR step for workflow."""
 
+    def _validate_cluster_config(self, cluster_config, step_name):
+        """Validates user provided cluster_config.
+
+        Args:
+            cluster_config(Union[Dict[str, Any], List[Dict[str, Any]]]):
+                user provided cluster configuration.
+            step_name: The name of the EMR step.
+        """
+
+        if (
+            "Name" in cluster_config
+            or "AutoTerminationPolicy" in cluster_config
+            or "Steps" in cluster_config
+        ):
+            raise ValueError(
+                err_str_with_name_auto_termination_or_steps.format(step_name=step_name)
+            )
+
+        if instances not in cluster_config:
+            raise ValueError(err_str_without_instance.format(step_name=step_name))
+
+        if (
+            "KeepJobFlowAliveWhenNoSteps" in cluster_config[instances]
+            or "TerminationProtected" in cluster_config[instances]
+        ):
+            raise ValueError(
+                err_str_with_keepjobflow_or_terminationprotected.format(step_name=step_name)
+            )
+
+        if (
+            instancegroups in cluster_config[instances]
+            and instancefleets in cluster_config[instances]
+        ) or (
+            instancegroups not in cluster_config[instances]
+            and instancefleets not in cluster_config[instances]
+        ):
+            raise ValueError(
+                err_str_both_or_none_instancegroups_or_instancefleets.format(step_name=step_name)
+            )
+
     def __init__(
         self,
         name: str,
@@ -128,7 +155,7 @@ def __init__(
         cache_config: CacheConfig = None,
         cluster_config: RequestType = None,
     ):
-        """Constructs a EMRStep.
+        """Constructs an EMRStep.
 
         Args:
             name(str): The name of the EMR step.
@@ -141,40 +168,36 @@ def __init__(
                 depends on.
             cache_config(CacheConfig):  A `sagemaker.workflow.steps.CacheConfig` instance.
             cluster_config(Union[Dict[str, Any], List[Dict[str, Any]]]): The recipe of the
-                EMR Cluster. It is a dictionary.
-                The elements are defined in the Request Syntax Section:
+                EMR cluster, passed as a dictionary. The elements are defined in the request syntax
+                for RunJobFlow. However, the following elements are not recognized as part of the
+                cluster configuration and you should not include them in the dictionary:
+                    1. cluster_config[Name]
+                    2. cluster_config[Steps]
+                    3. cluster_config[AutoTerminationPolicy]
+                    4. cluster_config[Instances][KeepJobFlowAliveWhenNoSteps]
+                    5. cluster_config[Instances][TerminationProtected]
+                For more information about the fields you can include in your cluster
+                configuration, see
                 https://docs.aws.amazon.com/emr/latest/APIReference/API_RunJobFlow.html
-                However, the following five elements are restricted, and must not present
-                in the dictionary:
-                1. cluster_config[Name]
-                2. cluster_config[Steps]
-                3. cluster_config[AutoTerminationPolicy]
-                4. cluster_config[Instances][KeepJobFlowAliveWhenNoSteps]
-                5. cluster_config[Instances][TerminationProtected]
-                Note that, if user wants to use cluster_config, then they have to explicitly set
-                cluster_id as None
-
+                Note that if you want to use cluster_config,
+                then you have to set cluster_id as None.
         """
         super(EMRStep, self).__init__(name, display_name, description, StepTypeEnum.EMR, depends_on)
 
         emr_step_args = {"StepConfig": step_config.to_request()}
         root_property = Properties(step_name=name, shape_name="Step", service_name="emr")
 
         if cluster_id is None and cluster_config is None:
-            raise Exception("EMRStep " + name + " must have either cluster_id or cluster_config")
+            raise ValueError(err_str_without_cluster_id_and_cluster_cfg.format(step_name=name))
 
         if cluster_id is not None and cluster_config is not None:
-            raise Exception(
-                "EMRStep " + name + " can not have both cluster_id or cluster_config. "
-                "If user wants to use cluster_config, then they "
-                "have to explicitly set cluster_id as None"
-            )
+            raise ValueError(err_str_with_both_cluster_id_and_cluster_cfg.format(step_name=name))
 
         if cluster_id is not None:
             emr_step_args["ClusterId"] = cluster_id
             root_property.__dict__["ClusterId"] = cluster_id
         elif cluster_config is not None:
-            validate_cluster_config(cluster_config, name)
+            self._validate_cluster_config(cluster_config, name)
             emr_step_args["ClusterConfig"] = cluster_config
             root_property.__dict__["ClusterConfig"] = cluster_config
 

tests/unit/sagemaker/workflow/test_emr_step.py

@@ -15,7 +15,16 @@
 import json
 import pytest
 
-from sagemaker.workflow.emr_step import EMRStep, EMRStepConfig
+from sagemaker.workflow.emr_step import (
+    EMRStep,
+    EMRStepConfig,
+    err_str_with_name_auto_termination_or_steps,
+    err_str_without_instance,
+    err_str_with_keepjobflow_or_terminationprotected,
+    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,
+)
 from sagemaker.workflow.steps import CacheConfig
 from sagemaker.workflow.pipeline import Pipeline, PipelineGraph
 from sagemaker.workflow.parameters import ParameterString
@@ -172,8 +181,6 @@ def test_pipeline_interpolates_emr_outputs(sagemaker_session):
 
 g_emr_step_config = EMRStepConfig(jar="s3:/script-runner/script-runner.jar")
 g_emr_step_name = "MyEMRStep"
-g_prefix = "EMRStep " + g_emr_step_name + " "
-g_prefix_with_in = "In EMRStep " + g_emr_step_name + ", "
 g_cluster_config = {
     "Instances": {
         "InstanceGroups": [
@@ -194,7 +201,7 @@ def test_pipeline_interpolates_emr_outputs(sagemaker_session):
 
 
 def test_emr_step_throws_exception_when_both_cluster_id_and_cluster_config_are_present():
-    with pytest.raises(Exception) as exceptionInfo:
+    with pytest.raises(ValueError) as exceptionInfo:
         EMRStep(
             name=g_emr_step_name,
             display_name="MyEMRStep",
@@ -205,18 +212,16 @@ def test_emr_step_throws_exception_when_both_cluster_id_and_cluster_config_are_p
             depends_on=["TestStep"],
             cache_config=CacheConfig(enable_caching=True, expire_after="PT1H"),
         )
-    expected_error_msg = (
-        g_prefix + "can not have both cluster_id or cluster_config. "
-        "If user wants to use cluster_config, then they "
-        "have to explicitly set cluster_id as None"
+    expected_error_msg = err_str_with_both_cluster_id_and_cluster_cfg.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_throws_exception_when_both_cluster_id_and_cluster_config_are_none():
-    with pytest.raises(Exception) as exceptionInfo:
+    with pytest.raises(ValueError) as exceptionInfo:
         EMRStep(
             name=g_emr_step_name,
             display_name="MyEMRStep",
@@ -226,7 +231,9 @@ def test_emr_step_throws_exception_when_both_cluster_id_and_cluster_config_are_n
             depends_on=["TestStep"],
             cache_config=CacheConfig(enable_caching=True, expire_after="PT1H"),
         )
-    expected_error_msg = g_prefix + "must have either cluster_id or cluster_config"
+    expected_error_msg = err_str_without_cluster_id_and_cluster_cfg.format(
+        step_name=g_emr_step_name
+    )
     actual_error_msg = exceptionInfo.value.args[0]
 
     assert actual_error_msg == expected_error_msg
@@ -327,8 +334,7 @@ def test_emr_step_with_valid_cluster_config():
                     ],
                 },
             },
-            g_prefix_with_in + "cluster_config should not contain any of "
-            "Name, AutoTerminationPolicy and/or Steps",
+            err_str_with_name_auto_termination_or_steps.format(step_name=g_emr_step_name),
         ),
         (
             {
@@ -341,8 +347,7 @@ def test_emr_step_with_valid_cluster_config():
                     ],
                 },
             },
-            g_prefix_with_in + "cluster_config should not contain any of "
-            "Name, AutoTerminationPolicy and/or Steps",
+            err_str_with_name_auto_termination_or_steps.format(step_name=g_emr_step_name),
         ),
         (
             {
@@ -355,22 +360,20 @@ def test_emr_step_with_valid_cluster_config():
                     ],
                 },
             },
-            g_prefix_with_in + "cluster_config should not contain any of "
-            "Name, AutoTerminationPolicy and/or Steps",
+            err_str_with_name_auto_termination_or_steps.format(step_name=g_emr_step_name),
         ),
         (
             {
                 "AmiVersion": "3.8.0",
                 "AdditionalInfo": "MyAdditionalInfo",
             },
-            g_prefix_with_in + "cluster_config must contain Instances",
+            err_str_without_instance.format(step_name=g_emr_step_name),
         ),
         (
             {
                 "Instances": {},
             },
-            g_prefix_with_in + "Instances should contain either "
-            "InstanceGroups or InstanceFleets",
+            err_str_both_or_none_instancegroups_or_instancefleets.format(step_name=g_emr_step_name),
         ),
         (
             {
@@ -387,8 +390,7 @@ def test_emr_step_with_valid_cluster_config():
                     ],
                 },
             },
-            g_prefix_with_in + "Instances should contain either "
-            "InstanceGroups or InstanceFleets",
+            err_str_both_or_none_instancegroups_or_instancefleets.format(step_name=g_emr_step_name),
         ),
         (
             {
@@ -401,9 +403,7 @@ def test_emr_step_with_valid_cluster_config():
                     "KeepJobFlowAliveWhenNoSteps": True,
                 },
             },
-            g_prefix_with_in + "Instances should not contain "
-            "KeepJobFlowAliveWhenNoSteps or "
-            "TerminationProtected",
+            err_str_with_keepjobflow_or_terminationprotected.format(step_name=g_emr_step_name),
         ),
         (
             {
@@ -416,16 +416,14 @@ def test_emr_step_with_valid_cluster_config():
                     "TerminationProtected": True,
                 },
             },
-            g_prefix_with_in + "Instances should not contain "
-            "KeepJobFlowAliveWhenNoSteps or "
-            "TerminationProtected",
+            err_str_with_keepjobflow_or_terminationprotected.format(step_name=g_emr_step_name),
         ),
     ],
 )
 def test_emr_step_throws_exception_when_cluster_config_contains_restricted_entities(
     invalid_cluster_config, expected_error_msg
 ):
-    with pytest.raises(Exception) as exceptionInfo:
+    with pytest.raises(ValueError) as exceptionInfo:
         EMRStep(
             name=g_emr_step_name,
             display_name="MyEMRStep",