pr review feedback changes

Author: EthanShouhanCheng

src/sagemaker/workflow/emr_step.py

@@ -73,13 +73,13 @@ def __init__(
         depends_on: List[str] = None,
         cache_config: CacheConfig = None,
     ):
-        """Constructs a LambdaStep.
+        """Constructs a EMRStep.
 
         Args:
             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.
-            cluster_id(str): A string that uniquely identifies the cluster.
+            cluster_id(str): The ID of the running EMR cluster.
             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

tests/integ/test_workflow.py

@@ -298,7 +298,7 @@ def emr_cluster_id(sagemaker_session, role):
     cluster_id = get_existing_emr_cluster_id(emr_client, cluster_name)
 
     if cluster_id is None:
-        create_new_emr_cluster(sagemaker_session, emr_client, cluster_name)
+        cluster_id = create_new_emr_cluster(sagemaker_session, emr_client, cluster_name)
     return cluster_id
 
 
@@ -310,6 +310,7 @@ def get_existing_emr_cluster_id(emr_client, cluster_name):
                 cluster_id = cluster["Id"]
                 print("Using existing cluster: {}".format(cluster_id))
                 return cluster_id
+        return None
     except Exception:
         raise