Zombie tasks in RESTARTING state are not cleaned

[Bisk1]

Apache Airflow version

2.7.0

Also reproduced on 2.5.0

What happened

Recently we added some automation to restarting Airflow tasks with "clear" command so we use this feature a lot. We often clear tasks in RUNNING state, which means that they go into RESTARTING state. We noticed that a lot of those tasks get stuck in RESTARTING state. Our Airflow infrastructure runs in an environment where any process can get suddenly killed without graceful shutdown.

We run Airflow on GKE but I managed to reproduce this behaviour on local environment with SequentialExecutor. See "How to reproduce" below for details.

What you think should happen instead

Tasks should get cleaned after scheduler restart and eventually get scheduled and executed.

How to reproduce

After some code investigation, I reproduced this kind of behaviour on local environment and it seems that RESTARTING tasks are only properly handled if the original restarting task is gracefully shut down so it can mark task as UP_FOR_RETRY or at least there is a healthy scheduler to do it if they fail for any other reason. The problem is with the following scenario:

1. Task is initially in RUNNING state.
2. Scheduler process dies suddenly.
3. The task process also dies suddenly.
4. "clear" command is executed on the task so the state is changed to RESTARTING state by webserver process.
5. From now on, even if we restart scheduler, the task will never get scheduled or change its state. It needs to have its state manually fixed, e.g. by clearing it again.

A recording of steps to reproduce on local environment:
https://vimeo.com/857192666?share=copy

Operating System

MacOS Ventura 13.4.1

Versions of Apache Airflow Providers

N/A

Deployment

Official Apache Airflow Helm Chart

Deployment details

N/A

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.

[Bisk1]

In the current implementation, the task process (not the scheduler process) is responsible for handling RESTARTING state. Handling RESTARTING state cannot be done reliably in the task process because it depends on correct propagation of SIGTERM and graceful shutdown which is not guaranteed.

This kind of situation could be handled by adopt_or_reset_orphaned_tasks() part of scheduler, but apparently that piece of code is meant to handle only tasks that were started by scheduler - that's why I would move that responsibility to scheduler first. Then scheduler would be able to handle orphaned RESTARTING tasks. The task process itself would only handle SIGTERM on best-effort basis to interrupt the restarted task.

My only concern is that separating these 2 responsibilities (terminating the old task attempt and starting a new one) would cause race condition - what if the scheduler attempts to reschedule restarted task before the old task process managed to handle SIGTERM? We would end up starting second process for the same task even before the first one was properly cleaned.

I'm willing to contribute a fix but scheduler's logic is quite complex so I'm interested in an opinion from a more experienced contributor first.

[potiuk]

it cannot be moved to scheduler, because of the support of various executors we have - namely Celery executor. Schduler does NOT control nor even really cares which workers are running the tasks. Celery executors work with QUEUE based mechanism, where Scheduler puts task to the queue, and workers pull tasks from the queue and once task is acknowledged further communication happens via the database - there is no other communication channel from scheduler to running tasks - it cannot simply order specifc tasks running on specific celery worker to restart.

Changing it woudl be a major change in how the whole architecture of scheduler -> executor -> worker workes. You would have to introduce a communication channel where scheduler could somehow signal specific celery worker to kill and restart task that runs on it. Currently it works in the way that there is a separate forked process that heartbeats the database and checks if the task state has been changed externally and kills child task running in order to restart it.

I am afraid the only option is to add a possibility of adopting tasks that are staying for too long in restarting state.

[Bisk1]

@potiuk ok, thanks for explanation.

I looked at the code again and now I think we can handle it with a change much simpler than I initially thought. Just adding RESTARTING state to the list of resettable states that are used by scheduler in adopt_or_reset_orphaned_tasks() and no other changes in scheduling responsibilities. (See #33706). RUNNING tasks are also handled this way and it doesn't seem to cause any problems with race conditions because by then they are either dead or can be successfully adopted - same case with RESTARTING tasks I assume.

You mention that we should check if these tasks

are staying for too long in restarting state

but we don't have such validation when resetting RUNNING or QUEUED tasks, so in line with the current design maybe there is no reason to treat RESETTING tasks in any special way different from those 2 states when it comes to adopting them?

[potiuk]

Yep. Likely that should do the trick. Opening PR and dscussion there where we can drag attention of others might be a good idea.