This is an automated email from the ASF dual-hosted git repository.
shahar1 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 976c1d08f21 Fix scheduler crash when retrying a task that timed out or
was killed externally (#75472)
976c1d08f21 is described below
commit 976c1d08f211f1d070c2c39d692a56ecbe4dd67e
Author: Kaxil Naik <[email protected]>
AuthorDate: Sat Oct 24 20:54:03 2026 +0100
Fix scheduler crash when retrying a task that timed out or was killed
externally (#75472)
The database callback sink committed and closed the scheduler's scoped
session when sending a callback, detaching the task instance the
scheduler was about to retry. Archiving the attempt then wrote nothing,
and inserting the next attempt hit the task_instance_current_key unique
constraint, which stopped the scheduler.
The sink now adds the request to the caller's session, and the
scheduler commits it. The Dag-lookup-failure branch of executor event
processing passes its session to set_state for the same reason.
---
.../airflow/callbacks/database_callback_sink.py | 14 ++-
.../src/airflow/executors/base_executor.py | 2 +
.../src/airflow/jobs/scheduler_job_runner.py | 6 +-
airflow-core/tests/unit/jobs/test_scheduler_job.py | 115 +++++++++++++++++++++
4 files changed, 131 insertions(+), 6 deletions(-)
diff --git a/airflow-core/src/airflow/callbacks/database_callback_sink.py
b/airflow-core/src/airflow/callbacks/database_callback_sink.py
index d29ec81ead8..cdf6ce02e6b 100644
--- a/airflow-core/src/airflow/callbacks/database_callback_sink.py
+++ b/airflow-core/src/airflow/callbacks/database_callback_sink.py
@@ -19,9 +19,9 @@ from __future__ import annotations
from typing import TYPE_CHECKING
+from airflow import settings
from airflow.callbacks.base_callback_sink import BaseCallbackSink
from airflow.models.db_callback_request import DbCallbackRequest
-from airflow.utils.session import NEW_SESSION, provide_session
if TYPE_CHECKING:
from sqlalchemy.orm import Session
@@ -32,8 +32,14 @@ if TYPE_CHECKING:
class DatabaseCallbackSink(BaseCallbackSink):
"""Sends callbacks to database."""
- @provide_session
- def send(self, callback: CallbackRequest, *, session: Session =
NEW_SESSION) -> None:
- """Send callback for execution."""
+ def send(self, callback: CallbackRequest, *, session: Session | None =
None) -> None:
+ """
+ Send callback for execution.
+
+ The request is added to ``session``, or to the thread's scoped session
when none is given,
+ without committing: the caller commits it with the rest of its
transaction.
+ """
+ if session is None:
+ session = settings.get_session()()
db_callback = DbCallbackRequest(callback=callback, priority_weight=1)
session.add(db_callback)
diff --git a/airflow-core/src/airflow/executors/base_executor.py
b/airflow-core/src/airflow/executors/base_executor.py
index 8d7d3b9b76b..f47410be70f 100644
--- a/airflow-core/src/airflow/executors/base_executor.py
+++ b/airflow-core/src/airflow/executors/base_executor.py
@@ -836,6 +836,8 @@ class BaseExecutor(LoggingMixin):
Send callback for execution.
Provides a default implementation which sends the callback to the
`callback_sink` object.
+ The scheduler's database sink adds the request to the calling thread's
scoped session
+ without committing, so the caller must commit that session.
:param request: Callback request to be executed.
"""
diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
index 198ef4bdc0a..82ea45fcb03 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -1740,7 +1740,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
ti.handle_failure(error=msg, session=session)
continue
cls.logger().exception("Marking task instance %s as %s",
ti, state)
- ti.set_state(state)
+ ti.set_state(state, session=session)
continue
ti.task = task
if task.has_on_retry_callback or task.has_on_failure_callback:
@@ -2212,8 +2212,10 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
for dag_run, callback_to_run in callback_tuples:
dag = cached_get_dag(dag_run)
if dag:
- # Sending callbacks to the database, so it must be done
outside of prohibit_commit.
self._send_dag_callbacks_to_processor(dag, callback_to_run)
+ # The sink only adds the request to this session. Commit each
one so a later Dag
+ # lookup error, or the critical section below rolling back,
can't lose it.
+ session.commit()
else:
self.log.error("DAG '%s' not found in serialized_dag table",
dag_run.dag_id)
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py
b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index 9f7d83b9787..13399201cf3 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -225,6 +225,16 @@ def get_callback_ti(ti):
)
+def _attempt_states(ti, session):
+ attempts = session.scalars(
+ select(TaskInstance)
+ .where(TaskInstance.dag_id == ti.dag_id, TaskInstance.task_id ==
ti.task_id)
+ .execution_options(include_all_attempts=True)
+ .order_by(TaskInstance.try_number)
+ )
+ return [(attempt.try_number, attempt.state, attempt.working_set) for
attempt in attempts]
+
+
@pytest.fixture(scope="class")
def disable_load_example():
with conf_vars({("core", "load_examples"): "false"}):
@@ -1247,6 +1257,84 @@ class TestSchedulerJob:
},
)
+ @pytest.mark.parametrize(
+ "callback_kwargs",
+ [
+ pytest.param({"on_failure_callback": lambda context: None},
id="on_failure_callback"),
+ pytest.param({"email": "[email protected]"}, id="email"),
+ ],
+ )
+ def
test_process_executor_events_callback_for_running_ti_with_retries_moves_it_to_retry(
+ self, callback_kwargs, dag_maker, session
+ ):
+ """Sending the callback must not close the session and detach the TI
before its retry is set up."""
+ with dag_maker(dag_id="executor_event_running_ti_with_retries",
fileloc="/test_path1/"):
+ EmptyOperator(task_id="task", retries=1, **callback_kwargs)
+ dr = dag_maker.create_dagrun()
+ ti = dr.task_instances[0]
+ executor = MockExecutor(do_update=False)
+ executor.callback_sink = DatabaseCallbackSink()
+ scheduler_job = Job()
+ session.add(scheduler_job)
+ session.flush()
+ self.job_runner = SchedulerJobRunner(scheduler_job,
executors=[executor])
+ ti.state = State.RUNNING
+ ti.try_number = 1
+ ti.queued_by_job_id = scheduler_job.id
+ session.merge(ti)
+ session.commit()
+ executor.event_buffer[TaskInstanceUuid(ti.id)] = State.SUCCESS, None
+
+ self.job_runner._process_executor_events(executor=executor,
session=session)
+ session.commit()
+
+ session.expire_all()
+ assert _attempt_states(ti, session) == [
+ (1, State.FAILED, None),
+ (2, State.UP_FOR_RETRY, True),
+ ]
+ assert
session.scalar(select(func.count()).select_from(DbCallbackRequest)) == 1
+
+ def
test_process_executor_events_dag_lookup_failure_keeps_session_for_rest_of_batch(
+ self, dag_maker, session
+ ):
+ """A failed Dag lookup must not close the session and detach the rest
of the batch."""
+ with dag_maker(dag_id="executor_event_dag_lookup_failure"):
+ EmptyOperator(task_id="first", retries=1)
+ EmptyOperator(task_id="second", retries=1)
+ dr = dag_maker.create_dagrun()
+ executor = MockExecutor(do_update=False)
+ scheduler_job = Job()
+ session.add(scheduler_job)
+ session.flush()
+ self.job_runner = SchedulerJobRunner(scheduler_job,
executors=[executor])
+ for ti in dr.task_instances:
+ ti.state = State.RUNNING
+ ti.try_number = 1
+ ti.queued_by_job_id = scheduler_job.id
+ session.merge(ti)
+ executor.event_buffer[TaskInstanceUuid(ti.id)] = State.SUCCESS,
None
+ session.commit()
+ real_get_dag_for_run =
self.job_runner.scheduler_dag_bag.get_dag_for_run
+ lookup_errors = [AirflowException("Dag lookup failed")]
+
+ def get_dag_for_run(*, dag_run, session):
+ if lookup_errors:
+ raise lookup_errors.pop()
+ return real_get_dag_for_run(dag_run=dag_run, session=session)
+
+ with mock.patch.object(
+ self.job_runner.scheduler_dag_bag, "get_dag_for_run",
autospec=True, side_effect=get_dag_for_run
+ ):
+ self.job_runner._process_executor_events(executor=executor,
session=session)
+ session.commit()
+
+ session.expire_all()
+ live_states = sorted(
+
session.scalars(select(TaskInstance.state).where(TaskInstance.dag_id ==
dr.dag_id))
+ )
+ assert live_states == [State.SUCCESS, State.UP_FOR_RETRY]
+
@mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest")
def test_process_executor_events_callback_for_removed_mapped_task(
self, mock_task_callback, dag_maker, session
@@ -10632,6 +10720,33 @@ class TestSchedulerJob:
)
assert dispatched_callback == expected_dispatched_callback
+ def test_heartbeat_timeout_with_retries_moves_ti_to_retry(self, dag_maker,
session):
+ """Sending the heartbeat-timeout callback must not close the purge's
session and detach the TI."""
+ with dag_maker(dag_id="hb_timeout_with_retries", session=session):
+ EmptyOperator(task_id="test_task", retries=1)
+ dag_run = dag_maker.create_dagrun(run_id="test_run",
state=DagRunState.RUNNING)
+ executor = MockExecutor(do_update=False)
+ executor.callback_sink = DatabaseCallbackSink()
+ scheduler_job = Job()
+ session.add(scheduler_job)
+ session.flush()
+ self.job_runner = SchedulerJobRunner(scheduler_job,
executors=[executor])
+ ti = dag_run.get_task_instance(task_id="test_task", session=session)
+ ti.state = TaskInstanceState.RUNNING
+ ti.try_number = 1
+ ti.queued_by_job_id = scheduler_job.id
+ ti.last_heartbeat_at = timezone.utcnow() - timedelta(seconds=600)
+ session.commit()
+
+ self.job_runner._find_and_purge_task_instances_without_heartbeats()
+
+ session.expire_all()
+ assert _attempt_states(ti, session) == [
+ (1, TaskInstanceState.FAILED, None),
+ (2, TaskInstanceState.UP_FOR_RETRY, True),
+ ]
+ assert
session.scalar(select(func.count()).select_from(DbCallbackRequest)) == 1
+
@pytest.mark.parametrize(
("retries", "callback_kind", "expected"),
[
