EmrCreateJobFlowOperator returns early when `wait_for_completion=True`

[enewnham]

Apache Airflow version

2.10.4

If "Other Airflow 2 version" selected, which one?

No response

What happened?

The EmrCreateJobFlowOperator returns early when wait_for_completion=True due improperly treating Running as success in the acceptor. When a step is running that doesn't necessarily mean the job is complete, only waiting for new jobs or terminated should the cluster be considered idle and the job completed.

This was introduced in #37667, https://github.com/apache/airflow/pull/37667/files#diff-f4ade57b13bf58042c4171d225b9b28c7610072f8b33b722bb138db900feb7fcR65-R70

I added some printf debugging to find this issue

[2025-01-14, 17:22:34 UTC] {emr.py:706} INFO - Job flow with id j-YEXH6CUHAEDF created
[2025-01-14, 17:22:34 UTC] {emr.py:737} INFO - job_flow_waiting
[2025-01-14, 17:22:34 UTC] {waiter.py:394} INFO - here: wait
[2025-01-14, 17:23:34 UTC] {waiter.py:394} INFO - here: wait
[2025-01-14, 17:24:35 UTC] {waiter.py:394} INFO - here: wait
[2025-01-14, 17:25:35 UTC] {waiter.py:369} INFO - here: current_state=success

What you think should happen instead?

The EmrCreateJobFlowOperator should continue to wait_for_completion while there are steps still running in the EMR cluster

How to reproduce

Use the EmrCreateJobFlowOperator to create a cluster with one or more steps. Optionally set "KeepJobFlowAliveWhenNoSteps": False,

Operating System

Amazon Linux

Versions of Apache Airflow Providers

No response

Deployment

Amazon (AWS) MWAA

Deployment details

No response

Anything else?

Every time

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[boring-cyborg]

Thanks for opening your first issue here! Be sure to follow the issue template! If you are willing to raise PR to address this issue please do so, no need to wait for approval.

[enewnham]

Just realizing we relied on an undefined behavior of EmrCreateJobFlowOperator in that before it would wait for the cluster to terminate when KeepJobFlowAliveWhenNoSteps=False. The new behavior is that is no longer the case.

For others who may of stubbed their toe, we wrapped the operator with our own operator and added the following code.

    def wait_for_terminate(self, job_flow_id: str):
        # We need to wait for the job to complete before we can get the status of the group by step
        # Since KeepJobFlowAliveWhenNoSteps is set to False, the cluster will terminate after the job is done
        if self.deferrable:
            self.defer(
                trigger=EmrTerminateJobFlowTrigger(
                    job_flow_id=job_flow_id,
                    waiter_delay=self.waiter_delay,
                    waiter_max_attempts=self.waiter_max_attempts,
                    aws_conn_id=self.aws_conn_id,
                ),
                method_name="execute_complete",
                # timeout is set to ensure that if a trigger dies, the timeout does not restart
                # 60 seconds is added to allow the trigger to exit gracefully (i.e. yield TriggerEvent)
                timeout=timedelta(seconds=self.waiter_max_attempts * self.waiter_delay + 60),
            )
        else:
            super()._emr_hook.get_waiter("job_flow_terminated").wait(
                ClusterId=self._job_flow_id,
                WaiterConfig=prune_dict(
                    {
                        "Delay": self.waiter_delay,
                        "MaxAttempts": self.waiter_max_attempts,
                    }
                ),
            )