diff --git a/airflow/jobs/scheduler_job.py b/airflow/jobs/scheduler_job.py index 8452950e0467f..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) @@ -1374,6 +1375,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()