Repository: apache/airflow

Issue:
Zombie tasks in RESTARTING state are not cleaned
### 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?

- [X] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)



Discussion:
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.

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.
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.
@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 https://github.com/apache/airflow/pull/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?

Resolving patch:
diff --git a/airflow/jobs/scheduler_job_runner.py b/airflow/jobs/scheduler_job_runner.py
--- a/airflow/jobs/scheduler_job_runner.py
+++ b/airflow/jobs/scheduler_job_runner.py
@@ -1579,11 +1579,11 @@ def _emit_pool_metrics(self, session: Session = NEW_SESSION) -> None:
     @provide_session
     def adopt_or_reset_orphaned_tasks(self, session: Session = NEW_SESSION) -> int:
         """
-        Reset any TaskInstance in QUEUED or SCHEDULED state if its SchedulerJob is no longer running.
+        Adopt or reset any TaskInstance in resettable state if its SchedulerJob is no longer running.
 
         :return: the number of TIs reset
         """
-        self.log.info("Resetting orphaned tasks for active dag runs")
+        self.log.info("Adopting or resetting orphaned tasks for active dag runs")
         timeout = conf.getint("scheduler", "scheduler_health_check_threshold")
 
         for attempt in run_with_db_retries(logger=self.log):
@@ -1609,10 +1609,9 @@ def adopt_or_reset_orphaned_tasks(self, session: Session = NEW_SESSION) -> int:
                         self.log.info("Marked %d SchedulerJob instances as failed", num_failed)
                         Stats.incr(self.__class__.__name__.lower() + "_end", num_failed)
 
-                    resettable_states = [TaskInstanceState.QUEUED, TaskInstanceState.RUNNING]
                     query = (
                         select(TI)
-                        .where(TI.state.in_(resettable_states))
+                        .where(TI.state.in_(State.adoptable_states))
                         # outerjoin is because we didn't use to have queued_by_job
                         # set, so we need to pick up anything pre upgrade. This (and the
                         # "or queued_by_job_id IS NONE") can go as soon as scheduler HA is
@@ -1628,11 +1627,11 @@ def adopt_or_reset_orphaned_tasks(self, session: Session = NEW_SESSION) -> int:
                     )
 
                     # Lock these rows, so that another scheduler can't try and adopt these too
-                    tis_to_reset_or_adopt = with_row_locks(
+                    tis_to_adopt_or_reset = with_row_locks(
                         query, of=TI, session=session, **skip_locked(session=session)
                     )
-                    tis_to_reset_or_adopt = session.scalars(tis_to_reset_or_adopt).all()
-                    to_reset = self.job.executor.try_adopt_task_instances(tis_to_reset_or_adopt)
+                    tis_to_adopt_or_reset = session.scalars(tis_to_adopt_or_reset).all()
+                    to_reset = self.job.executor.try_adopt_task_instances(tis_to_adopt_or_reset)
 
                     reset_tis_message = []
                     for ti in to_reset:
@@ -1640,11 +1639,11 @@ def adopt_or_reset_orphaned_tasks(self, session: Session = NEW_SESSION) -> int:
                         ti.state = None
                         ti.queued_by_job_id = None
 
-                    for ti in set(tis_to_reset_or_adopt) - set(to_reset):
+                    for ti in set(tis_to_adopt_or_reset) - set(to_reset):
                         ti.queued_by_job_id = self.job.id
 
                     Stats.incr("scheduler.orphaned_tasks.cleared", len(to_reset))
-                    Stats.incr("scheduler.orphaned_tasks.adopted", len(tis_to_reset_or_adopt) - len(to_reset))
+                    Stats.incr("scheduler.orphaned_tasks.adopted", len(tis_to_adopt_or_reset) - len(to_reset))
 
                     if to_reset:
                         task_instance_str = "\n\t".join(reset_tis_message)
diff --git a/airflow/utils/state.py b/airflow/utils/state.py
--- a/airflow/utils/state.py
+++ b/airflow/utils/state.py
@@ -199,3 +199,11 @@ def color_fg(cls, state):
     """
     A list of states indicating that a task has been terminated.
     """
+
+    adoptable_states = frozenset(
+        [TaskInstanceState.QUEUED, TaskInstanceState.RUNNING, TaskInstanceState.RESTARTING]
+    )
+    """
+    A list of states indicating that a task can be adopted or reset by a scheduler job
+    if it was queued by another scheduler job that is not running anymore.
+    """


Test patch:
diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py
--- a/tests/jobs/test_scheduler_job.py
+++ b/tests/jobs/test_scheduler_job.py
@@ -3151,6 +3151,30 @@ def test_adopt_or_reset_orphaned_tasks_nothing(self):
         session = settings.Session()
         assert 0 == self.job_runner.adopt_or_reset_orphaned_tasks(session=session)
 
+    @pytest.mark.parametrize(
+        "adoptable_state",
+        State.adoptable_states,
+    )
+    def test_adopt_or_reset_resettable_tasks(self, dag_maker, adoptable_state):
+        dag_id = "test_adopt_or_reset_adoptable_tasks_" + adoptable_state.name
+        with dag_maker(dag_id=dag_id, schedule="@daily"):
+            task_id = dag_id + "_task"
+            EmptyOperator(task_id=task_id)
+
+        scheduler_job = Job()
+        self.job_runner = SchedulerJobRunner(job=scheduler_job, subdir=os.devnull)
+        session = settings.Session()
+
+        dr1 = dag_maker.create_dagrun(external_trigger=True)
+        ti = dr1.get_task_instances(session=session)[0]
+        ti.state = adoptable_state
+        session.merge(ti)
+        session.merge(dr1)
+        session.commit()
+
+        num_reset_tis = self.job_runner.adopt_or_reset_orphaned_tasks(session=session)
+        assert 1 == num_reset_tis
+
     def test_adopt_or_reset_orphaned_tasks_external_triggered_dag(self, dag_maker):
         dag_id = "test_reset_orphaned_tasks_external_triggered_dag"
         with dag_maker(dag_id=dag_id, schedule="@daily"):