From 923aa16b099ede7cb8d054648eb2457dabe06ee6 Mon Sep 17 00:00:00 2001 From: Jed Cunningham Date: Thu, 7 Jul 2022 14:01:57 -0600 Subject: [PATCH 1/2] Fix zombie task handling with multiple schedulers Each scheduler was looking at all running tasks for zombies, leading to multiple schedulers handling the zombies. This causes problems with retries (e.g. being marked as FAILED instead of UP_FOR_RETRY) and callbacks (e.g. `on_failure_callback` being called multiple times). When the second scheduler tries to determine if the task is able to be retried, and it's already in UP_FOR_RETRY (the first scheduler already finished), it sees the "next" try_number (as it's no longer running), which then leads it to be FAILED instead. The easy fix is to simply restrict each scheduler to its own TIs, as orphaned running TIs will be adopted anyways. --- airflow/jobs/scheduler_job.py | 1 + tests/jobs/test_scheduler_job.py | 35 +++++++++++++++++++++----------- 2 files changed, 24 insertions(+), 12 deletions(-) diff --git a/airflow/jobs/scheduler_job.py b/airflow/jobs/scheduler_job.py index 8452950e0467f..66515d4f6b6d6 100644 --- a/airflow/jobs/scheduler_job.py +++ b/airflow/jobs/scheduler_job.py @@ -1374,6 +1374,7 @@ def _find_zombies(self, session): LocalTaskJob.latest_heartbeat < limit_dttm, ) ) + .filter(TaskInstance.queued_by_job_id == self.id) .all() ) diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py index 27931d460aa0c..7580d59c1cdd9 100644 --- a/tests/jobs/test_scheduler_job.py +++ b/tests/jobs/test_scheduler_job.py @@ -3890,7 +3890,6 @@ def test_find_zombies(self): session.query(LocalTaskJob).delete() dag = dagbag.get_dag('example_branch_operator') dag.sync_to_db() - task = dag.get_task(task_id='run_this_first') dag_run = dag.create_dagrun( state=DagRunState.RUNNING, @@ -3899,21 +3898,33 @@ def test_find_zombies(self): session=session, ) - ti = TaskInstance(task, run_id=dag_run.run_id, state=State.RUNNING) - local_job = LocalTaskJob(ti) - local_job.state = State.SHUTDOWN - - session.add(local_job) - session.flush() - - ti.job_id = local_job.id - session.add(ti) - session.flush() - self.scheduler_job = SchedulerJob(subdir=os.devnull) self.scheduler_job.executor = MockExecutor() self.scheduler_job.processor_agent = mock.MagicMock() + # We will provision 2 tasks so we can check we only find zombies from this scheduler + tasks_to_setup = ['branching', 'run_this_first'] + + for task_id in tasks_to_setup: + task = dag.get_task(task_id=task_id) + ti = TaskInstance(task, run_id=dag_run.run_id, state=State.RUNNING) + ti.queued_by_job_id = 999 + + local_job = LocalTaskJob(ti) + local_job.state = State.SHUTDOWN + + session.add(local_job) + session.flush() + + ti.job_id = local_job.id + session.add(ti) + session.flush() + + assert task.task_id == 'run_this_first' # Make sure we have the task/ti we expect + + ti.queued_by_job_id = self.scheduler_job.id + session.flush() + self.scheduler_job._find_zombies(session=session) self.scheduler_job.executor.callback_sink.send.assert_called_once() From f4460c0b93ea8d395004f15b8a3ca8896908bd9a Mon Sep 17 00:00:00 2001 From: Jed Cunningham Date: Thu, 7 Jul 2022 14:54:23 -0600 Subject: [PATCH 2/2] Update comment --- airflow/jobs/scheduler_job.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/airflow/jobs/scheduler_job.py b/airflow/jobs/scheduler_job.py index 66515d4f6b6d6..39aac294a88ba 100644 --- a/airflow/jobs/scheduler_job.py +++ b/airflow/jobs/scheduler_job.py @@ -1358,7 +1358,8 @@ def check_trigger_timeouts(self, session: Session = None): def _find_zombies(self, session): """ Find zombie task instances, which are tasks haven't heartbeated for too long - and update the current zombie list. + or have a no-longer-running LocalTaskJob, and send them off to the DAG processor + to be handled. """ self.log.debug("Finding 'running' jobs without a recent heartbeat") limit_dttm = timezone.utcnow() - timedelta(seconds=self._zombie_threshold_secs)