EMR Sample DAG and Docs Update (#22189)

Author: ferruzzi

airflow/providers/amazon/aws/example_dags/example_emr_job_flow_automatic_steps.py

@@ -15,15 +15,16 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""
-This is an example dag for a AWS EMR Pipeline with auto steps.
-"""
-from datetime import datetime, timedelta
+import os
+from datetime import datetime
 
 from airflow import DAG
 from airflow.providers.amazon.aws.operators.emr import EmrCreateJobFlowOperator
 from airflow.providers.amazon.aws.sensors.emr import EmrJobFlowSensor
 
+JOB_FLOW_ROLE = os.getenv('EMR_JOB_FLOW_ROLE', 'EMR_EC2_DefaultRole')
+SERVICE_ROLE = os.getenv('EMR_SERVICE_ROLE', 'EMR_DefaultRole')
+
 # [START howto_operator_emr_automatic_steps_config]
 SPARK_STEPS = [
     {
@@ -44,38 +45,40 @@
         'InstanceGroups': [
             {
                 'Name': 'Primary node',
-                'Market': 'SPOT',
+                'Market': 'ON_DEMAND',
                 'InstanceRole': 'MASTER',
-                'InstanceType': 'm1.medium',
+                'InstanceType': 'm5.xlarge',
                 'InstanceCount': 1,
-            }
+            },
         ],
         'KeepJobFlowAliveWhenNoSteps': False,
         'TerminationProtected': False,
     },
     'Steps': SPARK_STEPS,
-    'JobFlowRole': 'EMR_EC2_DefaultRole',
-    'ServiceRole': 'EMR_DefaultRole',
+    'JobFlowRole': JOB_FLOW_ROLE,
+    'ServiceRole': SERVICE_ROLE,
 }
 # [END howto_operator_emr_automatic_steps_config]
 
+
 with DAG(
-    dag_id='emr_job_flow_automatic_steps_dag',
-    dagrun_timeout=timedelta(hours=2),
+    dag_id='example_emr_job_flow_automatic_steps',
+    schedule_interval=None,
     start_date=datetime(2021, 1, 1),
-    schedule_interval='0 3 * * *',
-    catchup=False,
     tags=['example'],
+    catchup=False,
 ) as dag:
 
-    # [START howto_operator_emr_automatic_steps_tasks]
+    # [START howto_operator_emr_create_job_flow]
     job_flow_creator = EmrCreateJobFlowOperator(
         task_id='create_job_flow',
         job_flow_overrides=JOB_FLOW_OVERRIDES,
     )
+    # [END howto_operator_emr_create_job_flow]
 
-    job_sensor = EmrJobFlowSensor(task_id='check_job_flow', job_flow_id=job_flow_creator.output)
-    # [END howto_operator_emr_automatic_steps_tasks]
-
-    # Task dependency created via `XComArgs`:
-    #   job_flow_creator >> job_sensor
+    # [START howto_sensor_emr_job_flow_sensor]
+    job_sensor = EmrJobFlowSensor(
+        task_id='check_job_flow',
+        job_flow_id=job_flow_creator.output,
+    )
+    # [END howto_sensor_emr_job_flow_sensor]

airflow/providers/amazon/aws/example_dags/example_emr_job_flow_manual_steps.py

@@ -15,22 +15,21 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""
-This is an example dag for a AWS EMR Pipeline.
-
-Starting by creating a cluster, adding steps/operations, checking steps and finally when finished
-terminating the cluster.
-"""
-from datetime import datetime, timedelta
+import os
+from datetime import datetime
 
 from airflow import DAG
+from airflow.models.baseoperator import chain
 from airflow.providers.amazon.aws.operators.emr import (
     EmrAddStepsOperator,
     EmrCreateJobFlowOperator,
     EmrTerminateJobFlowOperator,
 )
 from airflow.providers.amazon.aws.sensors.emr import EmrStepSensor
 
+JOB_FLOW_ROLE = os.getenv('EMR_JOB_FLOW_ROLE', 'EMR_EC2_DefaultRole')
+SERVICE_ROLE = os.getenv('EMR_SERVICE_ROLE', 'EMR_DefaultRole')
+
 SPARK_STEPS = [
     {
         'Name': 'calculate_pi',
@@ -50,54 +49,58 @@
         'InstanceGroups': [
             {
                 'Name': 'Primary node',
-                'Market': 'SPOT',
+                'Market': 'ON_DEMAND',
                 'InstanceRole': 'MASTER',
-                'InstanceType': 'm1.medium',
+                'InstanceType': 'm5.xlarge',
                 'InstanceCount': 1,
-            }
+            },
         ],
-        'KeepJobFlowAliveWhenNoSteps': True,
+        'KeepJobFlowAliveWhenNoSteps': False,
         'TerminationProtected': False,
     },
-    'JobFlowRole': 'EMR_EC2_DefaultRole',
-    'ServiceRole': 'EMR_DefaultRole',
+    'JobFlowRole': JOB_FLOW_ROLE,
+    'ServiceRole': SERVICE_ROLE,
 }
 
+
 with DAG(
-    dag_id='emr_job_flow_manual_steps_dag',
-    dagrun_timeout=timedelta(hours=2),
+    dag_id='example_emr_job_flow_manual_steps',
+    schedule_interval=None,
     start_date=datetime(2021, 1, 1),
-    schedule_interval='0 3 * * *',
-    catchup=False,
     tags=['example'],
+    catchup=False,
 ) as dag:
 
-    # [START howto_operator_emr_manual_steps_tasks]
     cluster_creator = EmrCreateJobFlowOperator(
         task_id='create_job_flow',
         job_flow_overrides=JOB_FLOW_OVERRIDES,
     )
 
+    # [START howto_operator_emr_add_steps]
     step_adder = EmrAddStepsOperator(
         task_id='add_steps',
         job_flow_id=cluster_creator.output,
         steps=SPARK_STEPS,
     )
+    # [END howto_operator_emr_add_steps]
 
+    # [START howto_sensor_emr_step_sensor]
     step_checker = EmrStepSensor(
         task_id='watch_step',
         job_flow_id=cluster_creator.output,
         step_id="{{ task_instance.xcom_pull(task_ids='add_steps', key='return_value')[0] }}",
     )
+    # [END howto_sensor_emr_step_sensor]
 
+    # [START howto_operator_emr_terminate_job_flow]
     cluster_remover = EmrTerminateJobFlowOperator(
-        task_id='remove_cluster', job_flow_id=cluster_creator.output
+        task_id='remove_cluster',
+        job_flow_id=cluster_creator.output,
     )
+    # [END howto_operator_emr_terminate_job_flow]
 
-    step_adder >> step_checker >> cluster_remover
-    # [END howto_operator_emr_manual_steps_tasks]
-
-    # Task dependencies created via `XComArgs`:
-    #   cluster_creator >> step_adder
-    #   cluster_creator >> step_checker
-    #   cluster_creator >> cluster_remover
+    chain(
+        step_adder,
+        step_checker,
+        cluster_remover,
+    )

airflow/providers/amazon/aws/operators/emr.py

@@ -42,6 +42,10 @@ class EmrAddStepsOperator(BaseOperator):
     """
     An operator that adds steps to an existing EMR job_flow.
 
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:EmrAddStepsOperator`
+
     :param job_flow_id: id of the JobFlow to add steps to. (templated)
     :param job_flow_name: name of the JobFlow to add steps to. Use as an alternative to passing
         job_flow_id. will search for id of JobFlow with matching name in one of the states in
@@ -264,6 +268,10 @@ class EmrCreateJobFlowOperator(BaseOperator):
     A dictionary of JobFlow overrides can be passed that override
     the config from the connection.
 
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:EmrCreateJobFlowOperator`
+
     :param aws_conn_id: aws connection to uses
     :param emr_conn_id: emr connection to use
     :param job_flow_overrides: boto3 style arguments or reference to an arguments file
@@ -320,6 +328,11 @@ def execute(self, context: 'Context') -> str:
 class EmrModifyClusterOperator(BaseOperator):
     """
     An operator that modifies an existing EMR cluster.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:EmrModifyClusterOperator`
+
     :param cluster_id: cluster identifier
     :param step_concurrency_level: Concurrency of the cluster
     :param aws_conn_id: aws connection to uses
@@ -364,6 +377,10 @@ class EmrTerminateJobFlowOperator(BaseOperator):
     """
     Operator to terminate EMR JobFlows.
 
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:EmrTerminateJobFlowOperator`
+
     :param job_flow_id: id of the JobFlow to terminate. (templated)
     :param aws_conn_id: aws connection to uses
     """

airflow/providers/amazon/aws/sensors/emr.py

@@ -122,6 +122,10 @@ class EmrContainerSensor(BaseSensorOperator):
     Asks for the state of the job run until it reaches a failure state or success state.
     If the job run fails, the task will fail.
 
+    .. seealso::
+        For more information on how to use this sensor, take a look at the guide:
+        :ref:`howto/sensor:EmrContainerSensor`
+
     :param job_id: job_id to check the state of
     :param max_retries: Number of times to poll for query state before
         returning the current state, defaults to None
@@ -189,6 +193,10 @@ class EmrJobFlowSensor(EmrBaseSensor):
     When target_states is set to ['RUNNING', 'WAITING'] sensor waits
     until job flow to be ready (after 'STARTING' and 'BOOTSTRAPPING' states)
 
+    .. seealso::
+        For more information on how to use this sensor, take a look at the guide:
+        :ref:`howto/sensor:EmrJobFlowSensor`
+
     :param job_flow_id: job_flow_id to check the state of
     :param target_states: the target states, sensor waits until
         job flow reaches any of these states
@@ -263,6 +271,10 @@ class EmrStepSensor(EmrBaseSensor):
 
     With the default target states, sensor waits step to be completed.
 
+    .. seealso::
+        For more information on how to use this sensor, take a look at the guide:
+        :ref:`howto/sensor:EmrStepSensor`
+
     :param job_flow_id: job_flow_id which contains the step check the state of
     :param step_id: step to check the state of
     :param target_states: the target states, sensor waits until