Triggerer timeout exception is not handled properly

[arorasachin9]

Apache Airflow version

Other Airflow 3 version (please specify below)

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

3.1.3

What happened?

We are using the EMRContainers Operators and running it in deferred mode. We have a timeout of 20 minutes on the task. Sometimes the task takes more than 20 minutes. The control goes back to the worker pod from triggerer after the timeout but we are not caching the deferred timeout exception while getting the control back to worker. This is leading to the job keep on running on EMR instead of cancelling which happens in case of normal TaskTimeoutException. IN case of retry this is leading to parallel execution of the workloads which is not expected.
In file task_runner.py

        from airflow.sdk.execution_time.timeout import timeout

        # TODO: handle timeout in case of deferral
        timeout_seconds = task.execution_timeout.total_seconds()
        try:
            # It's possible we're already timed out, so fast-fail if true
            if timeout_seconds <= 0:
                raise AirflowTaskTimeout()
            # Run task in timeout wrapper
            with timeout(timeout_seconds):
                result = ctx.run(execute, context=context)
        except AirflowTaskTimeout:
            task.on_kill()
            raise
    else:
        result = ctx.run(execute, context=context)

We are only caching the AirflowTaskTimeout Exception but he exception that is thrown while resuming the taks on the worker is TaskDeferralTimeout
File: operator.py

    def resume_execution(self, next_method: str, next_kwargs: dict[str, Any] | None, context: Context):
        """Entrypoint method called by the Task Runner (instead of execute) when this task is resumed."""
        from airflow.sdk.exceptions import TaskDeferralError, TaskDeferralTimeout

        if next_kwargs is None:
            next_kwargs = {}
        # __fail__ is a special signal value for next_method that indicates
        # this task was scheduled specifically to fail.

        if next_method == TRIGGER_FAIL_REPR:
            next_kwargs = next_kwargs or {}
            traceback = next_kwargs.get("traceback")
            if traceback is not None:
                self.log.error("Trigger failed:\n%s", "\n".join(traceback))
            if (error := next_kwargs.get("error", "Unknown")) == TriggerFailureReason.TRIGGER_TIMEOUT:
                raise TaskDeferralTimeout(error)
            raise TaskDeferralError(error)
        # Grab the callable off the Operator/Task and add in any kwargs
        execute_callable = getattr(self, next_method)
        return execute_callable(context, **next_kwargs)

What you think should happen instead?

Expectation is the EMR job gets cancelled in case of the task timeout

How to reproduce

1. Add a EMRContainersOperator Task in airflow DAG with deferrable True
2. Let the task timeout with execution timeout
3. The job will keep on running on EMR.

Operating System

Linux

Versions of Apache Airflow Providers

No response

Deployment

Official Apache Airflow Helm Chart

Deployment details

No response

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

[arorasachin9]

@potiuk @ramitkataria ,
I think if we catch the exception when getting the control back to the worker will be a generalised solution for every operator.

        from airflow.sdk.execution_time.timeout import timeout

        # TODO: handle timeout in case of deferral
        timeout_seconds = task.execution_timeout.total_seconds()
        try:
            # It's possible we're already timed out, so fast-fail if true
            if timeout_seconds <= 0:
                raise AirflowTaskTimeout()
            # Run task in timeout wrapper
            with timeout(timeout_seconds):
                result = ctx.run(execute, context=context)
        except (AirflowTaskTimeout, TaskDeferralTimeout):
            task.on_kill()
            raise
    else:
        result = ctx.run(execute, context=context)

In my opinion the above fix should work. Let me know if seems any issues else I can raise a PR with this fix.

[Shally-Katariya]

Hi, I’d like to work on this issue. I have previously contributed to Airflow (PR #59824). I’ll start looking into the code and propose a fix soon.

[arorasachin9]

@Shally-Katariya The proposed fix and issue is already mentioned in the issue itself. Waiting for comment from the maintainer if it works. PR is also ready with me for these changes if it works.

[kaxil]

This conflates two different timeout concepts:

1. Execution timeout (execution_timeout on a task) - limits how long a task process can run. When exceeded, AirflowTaskTimeout is raised and on_kill() is called because the task process is actively running.

2. Deferral/Sensor timeout (passed to defer()) - limits how long a trigger waits for a condition. When exceeded, TaskDeferralTimeout is raised. The task process isn't running at this point — it's deferred. The triggerer fires an event and the task resumes just to raise the exception.

TaskDeferralTimeout was introduced in #33718 specifically to give deferrable sensors the same behavior as non-deferrable sensors w.r.t. sensor timeouts — i.e., the task fails without retry (via AirflowSensorTimeout). It was never intended to trigger on_kill().

If you need to cancel the EMR job when a deferral timeout occurs, that's operator-specific cleanup logic. The EMRContainers operator should handle this in its callback method or by overriding resume_execution to perform cleanup before re-raising. This isn't something the framework should do generically via on_kill() — on_kill() is for when a running task process is being terminated.

The fix here would be to update the EMR operator to handle cleanup on deferral timeout, not to change the task runner.

[arorasachin9]

@kaxil Can we have a genelized implementation just like the waiter implenentation is there for the cleanup also in whole AWS provider. It will help in others also like athena also needs job cancellation after Trigger timeout.
If you have context let me know or please tag someone who has context/maintainer on the AWS porovider.
And could you please change the Tag's also. I am unbale to change them.
cc @o-nikolas

[kaxil]

cc @vincbeck -- who is one of my our fellow PMC member but also AWS provider maintainer

[vincbeck]

@kaxil Can we have a genelized implementation just like the waiter implenentation is there for the cleanup also in whole AWS provider. It will help in others also like athena also needs job cancellation after Trigger timeout. If you have context let me know or please tag someone who has context/maintainer on the AWS porovider. And could you please change the Tag's also. I am unbale to change them. cc @o-nikolas

I think a general approach would be difficult and maybe dangerous (and over engineered?). I think I rather an operator to operator approach where the cancelation of the job happens in the operator. As a maintainer and user perspective, that also simplify the review and understanding of an operator. Having a global approach would look like job cancellation happening as "magic" instead of being explicit in the operator.

FYI, a PR resolving a similar issue has been merged recently: #51883

[arorasachin9]

@vincbeck ,
I think the method execute_complete will never get called as the exception is thrown before hand only. So handling the way it is handled in the attached PR might not work. Let me know if my understanding is wrong

def resume_execution(self, next_method: str, next_kwargs: dict[str, Any] | None, context: Context):
        """Entrypoint method called by the Task Runner (instead of execute) when this task is resumed."""
        from airflow.sdk.exceptions import TaskDeferralError, TaskDeferralTimeout

        if next_kwargs is None:
            next_kwargs = {}
        # __fail__ is a special signal value for next_method that indicates
        # this task was scheduled specifically to fail.

        if next_method == TRIGGER_FAIL_REPR:
            next_kwargs = next_kwargs or {}
            traceback = next_kwargs.get("traceback")
            if traceback is not None:
                self.log.error("Trigger failed:\n%s", "\n".join(traceback))
            if (error := next_kwargs.get("error", "Unknown")) == TriggerFailureReason.TRIGGER_TIMEOUT:
                raise TaskDeferralTimeout(error)
            raise TaskDeferralError(error)
        # Grab the callable off the Operator/Task and add in any kwargs
        execute_callable = getattr(self, next_method)
        return execute_callable(context, **next_kwargs)

This function gets called on the resume of the task after timeout. Here TaskDeferralTimeout(error) is raised before calling execute_complete. So It will not help if I handle here.
Every trigger in the provider implements AwsBaseWaiterTrigger. I need to paas the cancel method to this particular class to cleanup the task by putting the run function inside try catch or I can implement cleanup method in class AwsBaseWaiterTrigger.
Let me know what you think.

[vincbeck]

Yep your arguments seem valid. What is your proposal? Maybe you can create a PR so that we can discuss that directly there?