Call on_kill when clearing a defered task

[Max0u]

Is there a way to call on_kill method when clearing (or setting on failed or succeed) a task ?

[potiuk]

I think on_kill is called when the task itself fails or succeeds. I believe when the task gets cleared the attempt is to terminate task gracefully and the on_kill method should be called, but on_kill method is always executed in the same process that executes the main "execute() method, so there are cases when attempt to kill the processs gracefully fail and on_kill method might get not called because the task process is killed forcefully.

I think what you can do to get more resillient to it is to implement on_failure_callback or use the new features of Airlfow 2.5 - listeners if you want to be sure to receive the callback, regardless of whatever happened with the task process. Callbacks and listeners are handled in other processes in case the task process fails to be killed gracefully.

[Max0u]

Ok I totally understand what you're saying but I didn't give enough context and it's my bad. I'm trying to add defer to the Athena Operator, since we have a lot of long queries I would like to free the workers. But I do want to keep the feature that when I clear (or change the state) of my task, the operator will stop the query. But when the task is in defer mode, he is not in the worker, so when I clear it Airflow will just skip it. I tried to use the on_failure_callback but since I lost the context, I lost the query execution id. Is there a way to retrieves the query execution id ? Or another workaround ?

Thanks a lot

[potiuk]

I think it's a missing feature in deferred operators. opening a feature for that is a good idea.

In the meantime, I thinl you could leverage 2.5.0 listeners to perform such cleanup.

[thesuperzapper]

@potiuk did we ever find an elegant way to handle "manual state changes" for deferred tasks?

Probably the most common use case for deferred tasks is long-running queries/jobs, so it's pretty important we allow users to cancel that job/query when they change the state of the deferred task.

Because right now, on_failure_callback, on_kill (BaseOperator), and on_task_instance_failed (Listener) will NOT run when you manually set the state of ANY task (including deferred) to FAILED.

EDIT: I see that there is a related discussion in #36090