diff --git a/airflow/jobs/scheduler_job.py b/airflow/jobs/scheduler_job.py index 89be0d7a181e9..1745c01e53072 100644 --- a/airflow/jobs/scheduler_job.py +++ b/airflow/jobs/scheduler_job.py @@ -1520,8 +1520,7 @@ def check_trigger_timeouts(self, session: Session = NEW_SESSION) -> None: if num_timed_out_tasks: self.log.info("Timed out %i deferred tasks without fired triggers", num_timed_out_tasks) - @provide_session - def _find_zombies(self, session: Session) -> None: + def _find_zombies(self) -> None: """ Find zombie task instances, which are tasks haven't heartbeated for too long or have a no-longer-running LocalTaskJob, and create a TaskCallbackRequest @@ -1532,30 +1531,30 @@ def _find_zombies(self, session: Session) -> None: self.log.debug("Finding 'running' jobs without a recent heartbeat") limit_dttm = timezone.utcnow() - timedelta(seconds=self._zombie_threshold_secs) - zombies = ( - session.query(TaskInstance, DagModel.fileloc) - .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql") - .join(LocalTaskJob, TaskInstance.job_id == LocalTaskJob.id) - .join(DagModel, TaskInstance.dag_id == DagModel.dag_id) - .filter(TaskInstance.state == TaskInstanceState.RUNNING) - .filter( - or_( - LocalTaskJob.state != State.RUNNING, - LocalTaskJob.latest_heartbeat < limit_dttm, + with create_session() as session: + zombies: list[tuple[TI, str, str]] = ( + session.query(TI, DM.fileloc, DM.processor_subdir) + .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql") + .join(LocalTaskJob, TaskInstance.job_id == LocalTaskJob.id) + .filter(TI.state == TaskInstanceState.RUNNING) + .filter( + or_( + LocalTaskJob.state != State.RUNNING, + LocalTaskJob.latest_heartbeat < limit_dttm, + ) ) + .filter(TI.queued_by_job_id == self.id) + .all() ) - .filter(TaskInstance.queued_by_job_id == self.id) - .all() - ) if zombies: self.log.warning("Failing (%s) jobs without heartbeat after %s", len(zombies), limit_dttm) - for ti, file_loc in zombies: + for ti, file_loc, processor_subdir in zombies: zombie_message_details = self._generate_zombie_message_details(ti) request = TaskCallbackRequest( full_filepath=file_loc, - processor_subdir=ti.dag_model.processor_subdir, + processor_subdir=processor_subdir, simple_task_instance=SimpleTaskInstance.from_ti(ti), msg=str(zombie_message_details), ) diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py index 0363ee30eb48c..b6858cf085774 100644 --- a/tests/jobs/test_scheduler_job.py +++ b/tests/jobs/test_scheduler_job.py @@ -4129,14 +4129,13 @@ def test_timeout_triggers(self, dag_maker): assert ti2.state == State.DEFERRED def test_find_zombies_nothing(self): - with create_session() as session: - executor = MockExecutor(do_update=False) - self.scheduler_job = SchedulerJob(executor=executor) - self.scheduler_job.processor_agent = mock.MagicMock() + executor = MockExecutor(do_update=False) + self.scheduler_job = SchedulerJob(executor=executor) + self.scheduler_job.processor_agent = mock.MagicMock() - self.scheduler_job._find_zombies(session=session) + self.scheduler_job._find_zombies() - self.scheduler_job.executor.callback_sink.send.assert_not_called() + self.scheduler_job.executor.callback_sink.send.assert_not_called() def test_find_zombies(self, load_examples): dagbag = DagBag(TEST_DAG_FOLDER, read_dags_from_db=False) @@ -4179,20 +4178,21 @@ def test_find_zombies(self, load_examples): ti.queued_by_job_id = self.scheduler_job.id session.flush() - self.scheduler_job._find_zombies(session=session) + self.scheduler_job._find_zombies() - self.scheduler_job.executor.callback_sink.send.assert_called_once() - requests = self.scheduler_job.executor.callback_sink.send.call_args[0] - assert 1 == len(requests) - assert requests[0].full_filepath == dag.fileloc - assert requests[0].msg == str(self.scheduler_job._generate_zombie_message_details(ti)) - assert requests[0].is_failure_callback is True - assert isinstance(requests[0].simple_task_instance, SimpleTaskInstance) - assert ti.dag_id == requests[0].simple_task_instance.dag_id - assert ti.task_id == requests[0].simple_task_instance.task_id - assert ti.run_id == requests[0].simple_task_instance.run_id - assert ti.map_index == requests[0].simple_task_instance.map_index + self.scheduler_job.executor.callback_sink.send.assert_called_once() + requests = self.scheduler_job.executor.callback_sink.send.call_args[0] + assert 1 == len(requests) + assert requests[0].full_filepath == dag.fileloc + assert requests[0].msg == str(self.scheduler_job._generate_zombie_message_details(ti)) + assert requests[0].is_failure_callback is True + assert isinstance(requests[0].simple_task_instance, SimpleTaskInstance) + assert ti.dag_id == requests[0].simple_task_instance.dag_id + assert ti.task_id == requests[0].simple_task_instance.task_id + assert ti.run_id == requests[0].simple_task_instance.run_id + assert ti.map_index == requests[0].simple_task_instance.map_index + with create_session() as session: session.query(TaskInstance).delete() session.query(LocalTaskJob).delete() @@ -4267,12 +4267,11 @@ def test_find_zombies_handle_failure_callbacks_are_correctly_passed_to_dag_proce Check that the same set of failure callback with zombies are passed to the dag file processors until the next zombie detection logic is invoked. """ - with conf_vars({("core", "load_examples"): "False"}): + with conf_vars({("core", "load_examples"): "False"}), create_session() as session: dagbag = DagBag( dag_folder=os.path.join(settings.DAGS_FOLDER, "test_example_bash_operator.py"), read_dags_from_db=False, ) - session = settings.Session() session.query(LocalTaskJob).delete() dag = dagbag.get_dag("test_example_bash_operator") dag.sync_to_db(processor_subdir=TEST_DAG_FOLDER) @@ -4301,7 +4300,7 @@ def test_find_zombies_handle_failure_callbacks_are_correctly_passed_to_dag_proce self.scheduler_job.executor = MockExecutor() self.scheduler_job.processor_agent = mock.MagicMock() - self.scheduler_job._find_zombies(session=session) + self.scheduler_job._find_zombies() self.scheduler_job.executor.callback_sink.send.assert_called_once()