diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 355f693db0a02..dce7c88bad71c 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -3548,12 +3548,21 @@ def _purge_task_instances_without_heartbeats( # Backfill dag_version_id for legacy tasks (Pydantic requires uuid.UUID). if not _ensure_ti_has_dag_version_id(ti, session, self.log): continue + # ti.task isn't loaded in this purge path, so is_eligible_to_retry() uses its + # no-task fallback (``try_number <= max_tries``), which skips the retries-configured + # check its task-loaded branch applies; guard with ``max_tries > 0`` so a task + # declared with retries=0 isn't treated as retry-eligible here. + if ti.max_tries > 0 and ti.is_eligible_to_retry(): + task_callback_type = TaskInstanceState.UP_FOR_RETRY + else: + task_callback_type = TaskInstanceState.FAILED request = TaskCallbackRequest( filepath=ti.dag_model.relative_fileloc or "", bundle_name=_hb_bundle_name, bundle_version=_hb_bundle_version, ti=ti, msg=str(task_instance_heartbeat_timeout_message_details), + task_callback_type=task_callback_type, context_from_server=TIRunContext( dag_run=DRDataModel.model_validate(ti.dag_run, from_attributes=True), max_tries=ti.max_tries, diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index d7b08e662b19b..f5aeceb34b39c 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -8608,6 +8608,121 @@ def test_scheduler_passes_context_from_server_on_heartbeat_timeout(self, dag_mak assert callback_request.context_from_server.dag_run.logical_date == dag_run.logical_date assert callback_request.context_from_server.max_tries == ti.max_tries + @pytest.mark.parametrize( + ("state", "retries", "try_number", "expected_callback_type", "expected_dispatched_callback"), + [ + pytest.param( + TaskInstanceState.RUNNING, + 0, + 1, + TaskInstanceState.FAILED, + "on_failure_callback", + id="no_retries", + ), + pytest.param( + TaskInstanceState.RUNNING, + 2, + 1, + TaskInstanceState.UP_FOR_RETRY, + "on_retry_callback", + id="retries_available_first_attempt", + ), + pytest.param( + TaskInstanceState.RUNNING, + 2, + 2, + TaskInstanceState.UP_FOR_RETRY, + "on_retry_callback", + id="retries_available_mid_chain", + ), + pytest.param( + TaskInstanceState.RUNNING, + 2, + 3, + TaskInstanceState.FAILED, + "on_failure_callback", + id="retries_exhausted", + ), + pytest.param( + TaskInstanceState.RESTARTING, + 1, + 5, + TaskInstanceState.UP_FOR_RETRY, + "on_retry_callback", + id="restarting_stays_eligible_past_max_tries", + ), + ], + ) + def test_heartbeat_timeout_sets_callback_type_by_retry_eligibility( + self, + dag_maker, + session, + state, + retries, + try_number, + expected_callback_type, + expected_dispatched_callback, + ): + """Heartbeat-timeout cleanup must populate ``task_callback_type`` so the Dag processor + fires ``on_retry_callback`` when the task still has retries left, not + ``on_failure_callback``. + + Reproduces the bug end-to-end through the actual scheduler purge path: + + 1. A TI is ``RUNNING`` (or ``RESTARTING``) with a stale ``last_heartbeat_at`` (worker + OOMKilled, node evicted, scheduler restarted, etc.). + 2. ``_find_and_purge_task_instances_without_heartbeats`` builds a + ``TaskCallbackRequest`` and hands it to the executor's ``send_callback``. + 3. The Dag processor branches on ``request.task_callback_type``: + ``UP_FOR_RETRY`` -> ``task.on_retry_callback``; anything else (including ``None``) + -> ``task.on_failure_callback``. See + ``airflow-core/src/airflow/dag_processing/processor.py``::``_execute_task_callbacks``. + + Before the fix, step 2 left ``task_callback_type`` as ``None``, so step 3 always fell + into the ``else`` branch and ``on_failure_callback`` fired even when the task still had + retries left -- producing spurious failure alerts for tasks that ultimately succeeded on + retry. + + The parametrized cases cover the full ``max_tries`` / ``try_number`` matrix for a + ``RUNNING`` TI -- no retries, retries available (first attempt and mid-chain), and + retries exhausted (``try_number > max_tries``) -- plus a ``RESTARTING`` TI (cleared + while running), which ``is_eligible_to_retry`` keeps retry-eligible even past + ``max_tries``. The ``expected_dispatched_callback`` column mirrors the Dag processor's + branch so the assertion captures the user-visible outcome, not just the field value. + """ + with dag_maker(dag_id=f"hb_timeout_r{retries}_t{try_number}", session=session): + EmptyOperator(task_id="test_task", retries=retries) + + dag_run = dag_maker.create_dagrun(run_id="test_run", state=DagRunState.RUNNING) + + mock_executor = MagicMock() + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(scheduler_job, executors=[mock_executor]) + + ti = dag_run.get_task_instance(task_id="test_task") + ti.state = state + ti.try_number = try_number + ti.queued_by_job_id = scheduler_job.id + ti.last_heartbeat_at = timezone.utcnow() - timedelta(seconds=600) + session.merge(ti) + session.commit() + + self.job_runner._find_and_purge_task_instances_without_heartbeats() + + mock_executor.send_callback.assert_called_once() + request = mock_executor.send_callback.call_args[0][0] + assert isinstance(request, TaskCallbackRequest) + assert request.task_callback_type == expected_callback_type + # Mirror processor._execute_task_callbacks: UP_FOR_RETRY -> on_retry_callback, else + # on_failure_callback. Asserting the dispatched callback closes the loop on the + # user-visible behaviour, not just the field value. + dispatched_callback = ( + "on_retry_callback" + if request.task_callback_type is TaskInstanceState.UP_FOR_RETRY + else "on_failure_callback" + ) + assert dispatched_callback == expected_dispatched_callback + @pytest.mark.parametrize( ("retries", "callback_kind", "expected"), [