This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new f2403ccb58b Fix scheduler firing on_failure_callback for
heartbeat-timed-out retries (#66767)
f2403ccb58b is described below
commit f2403ccb58b32ae9a073b70d6648c11447d6c9b0
Author: Stefan Wang <[email protected]>
AuthorDate: Mon Jul 13 07:25:24 2026 -0400
Fix scheduler firing on_failure_callback for heartbeat-timed-out retries
(#66767)
When a worker stops heartbeating (OOMKill, node eviction), the scheduler's
``_purge_task_instances_without_heartbeats`` built a ``TaskCallbackRequest``
without ``task_callback_type``. The Dag processor's task-callback dispatch
branches on that field: ``UP_FOR_RETRY`` runs ``on_retry_callback``,
anything
else (including ``None``) runs ``on_failure_callback``. So heartbeat-timeout
cleanup always fired ``on_failure_callback`` even when the task still had
retries remaining, producing spurious failure alerts for tasks that
ultimately succeeded on retry.
Set ``task_callback_type`` from ``ti.is_eligible_to_retry()``, the canonical
retry-eligibility predicate, guarded by ``max_tries > 0``. The guard covers
the one gap the predicate has here: this path doesn't load ``ti.task``, so
the
predicate falls back to ``try_number <= max_tries`` and drops the
retries-configured check its task-loaded branch applies. Deferring to the
predicate also keeps a ``RESTARTING`` task (cleared while running) retry-
eligible past ``max_tries``, where a hand-rolled ``try_number <= max_tries``
check would have fired ``on_failure_callback``.
closes: #65400
Signed-off-by: 1fanwang <[email protected]>
Co-authored-by: kimhaggie <[email protected]>
---
.../src/airflow/jobs/scheduler_job_runner.py | 9 ++
airflow-core/tests/unit/jobs/test_scheduler_job.py | 115 +++++++++++++++++++++
2 files changed, 124 insertions(+)
diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
index b941a5b3580..7d5f7f58020 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -3581,6 +3581,14 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
# 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,
@@ -3588,6 +3596,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
version_data=_hb_version_data,
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 59bc50a7f12..1f3b6992c3e 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -8677,6 +8677,121 @@ class TestSchedulerJob:
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"),
[
