EmrStartNotebookExecutionOperator seems to be calling NON-existing waiter

[june109109]

Apache Airflow Provider(s)

amazon

Versions of Apache Airflow Providers

apache-airflow-providers-amazon==9.7.0

Apache Airflow version

2.10.3

Operating System

Deployment

Amazon (AWS) MWAA

Deployment details

No response

What happened

airflow.providers.amazon.aws.operators.emr

    if self.wait_for_completion:
        self.hook.get_waiter("notebook_running").wait(
            NotebookExecutionId=notebook_execution_id,
            WaiterConfig=prune_dict(
                {
                    "Delay": self.waiter_delay,
                    "MaxAttempts": self.waiter_max_attempts,
                }
            ),
        )

EMR Hook doesn't seem to have waiter "notebook_running"

What you think should happen instead

No response

How to reproduce

airflow.providers.amazon.aws.operators.emr

    if self.wait_for_completion:
        self.hook.get_waiter("notebook_running").wait(
            NotebookExecutionId=notebook_execution_id,
            WaiterConfig=prune_dict(
                {
                    "Delay": self.waiter_delay,
                    "MaxAttempts": self.waiter_max_attempts,
                }
            ),
        )

EMR Hook doesn't seem to have waiter "notebook_running"

Anything else

No response

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.

[eladkal]

cc @ellisms can you look into it? This was refactored in #49486

[ellisms]

Yes, let me dig into this.

notebook_running definitely exists:

airflow/providers/amazon/src/airflow/providers/amazon/aws/waiters/emr.json

Line 4 in a4a51a0

"notebook_running": {

[ellisms]

@june109109 I noticed the deployment is MWAA. The constraints file used in 2.10.3 specify the Amazon Provider Package is 9.0.0. Are you overriding this to install a later version and have you ensured there are no other package conflicts?

There is a unit test which checks the type of waiter called and it passes for this operator:
assert_expected_waiter_type(mock_waiter, "notebook_running")

Is there an error when you run a DAG with this task? I haven't been able to spin up an EMR environment to test an actual DAG yet.

[june109109]

@ellisms
I'm sorry for late response.

~ $ aws emr describe-notebook-execution --notebook-execution-id ex-J0BYN0I0N06S4UB8DAAFKBJXKKWO5
{
    "NotebookExecution": {
        "NotebookExecutionId": "ex-J0BYN0I0N06S4UB8DAAFKBJXKKWO5",
        "EditorId": "e-AH8OW5O68ZRXGXYFFLF14HIAG",
        "ExecutionEngine": {
            "Id": "j-19O4HGHRC48IE",
            "Type": "EMR",
            "MasterInstanceSecurityGroupId": "sg-0c2a01931b0a299fd"
        },
        "NotebookExecutionName": "user_operator_test-user_nb",
        "NotebookParams": "{\"query\": \"SELECT * FROM tos.tcic_nrtm_integ_cont_prod WHERE strd_dt = '250702'\", \"key2\": 123, \"default_datestrings\": {\"yyyymmdd\": \"20250703\", \"yyyy-mm-dd\": \"2025-07-03\"}}",
        "Status": "FINISHED",
        "StartTime": "2025-07-03T08:09:36.552000+00:00",
        "EndTime": "2025-07-03T08:10:52.757000+00:00",
        "Arn": "arn:aws:elasticmapreduce:ap-northeast-2:302263049595:notebook-execution/ex-J0BYN0I0N06S4UB8DAAFKBJXKKWO5",
        "OutputNotebookURI": "s3://cosmos-computing-test-emr-20250609234141565800000008/poc/e-AH8OW5O68ZRXGXYFFLF14HIAG/executions/ex-J0BYN0I0N06S4UB8DAAFKBJXKKWO5/query.ipynb",
        "LastStateChangeReason": "Execution is finished for cluster j-19O4HGHRC48IE.",
        "NotebookInstanceSecurityGroupId": "",
        "Tags": [],
        "OutputNotebookS3Location": {
            "Bucket": "cosmos-computing-test-emr-20250609234141565800000008",
            "Key": "poc/e-AH8OW5O68ZRXGXYFFLF14HIAG/executions/ex-J0BYN0I0N06S4UB8DAAFKBJXKKWO5/query.ipynb"
        },
        "OutputNotebookFormat": "null"
    }
}

It seems issue occurs because the EMR waiter's internal argument path for status is set to "notebookexecution.status",
but the actual response from describe_notebook_execution uses "NotebookExecution.Status" (with upper case letter). The waiter may looks for response["notebookexecution"]["status"], but actual response is response["NotebookExecution"]["Status"], so waits while notebook is finished.

I tried custom waiter definition and found out it works well with my case.

NOTEBOOK_WAITER_MODEL = {
    "version": 2,
    "waiters": {
        "notebook_execution_terminal": {
            "operation": "DescribeNotebookExecution",
            "delay": 30,
            "maxAttempts": 120,
            "acceptors": [
                {
                    "matcher": "path",
                    "argument": "NotebookExecution.Status",
                    "expected": "FINISHED",
                    "state": "success"
                },
                {
                    "matcher": "path",
                    "argument": "NotebookExecution.Status",
                    "expected": "FAILED",
                    "state": "failure"
                },

            ]
        }
    }
}

    emr_hook = EmrHook(aws_conn_id=self.aws_conn_id)
    emr_client = emr_hook.get_conn()
    waiter = BaseBotoWaiter(emr_client, NOTEBOOK_WAITER_MODEL).waiter("notebook_execution_terminal")
    waiter.wait(
        NotebookExecutionId=notebook_execution_id,
        WaiterConfig=prune_dict({
            "Delay": self.waiter_delay,
            "MaxAttempts": self.waiter_max_attempts,
        }),
    )