amoghrajesh commented on code in PR #76908:
URL: https://github.com/apache/airflow/pull/76908#discussion_r4537464886
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3554,6 +3554,22 @@ def _handle_tasks_stuck_in_queued(self, *, session:
Session = NEW_SESSION) -> No
for executor, stuck_tis in
self._executor_to_workloads(tasks_stuck_in_queued, session).items():
try:
for ti in stuck_tis:
+ # The scan takes no lock; skip TIs a worker started since,
and hold the row lock
+ # through the commit below.
+ still_queued = session.scalar(
Review Comment:
Mapped and loop TIs are regional (`regional_expression` matches
`region_index >= 0`). Per #74349, paths that lock a regional TI must lock its
DagRun first, so they can't deadlock with a clear, like
`process_executor_events` and the heartbeat scan do. Could this lock the DagRun
first (skip_locked) for regional TIs and skip if it's held?
##########
airflow-core/tests/unit/jobs/test_scheduler_job.py:
##########
@@ -11840,6 +11844,133 @@ def
test_stuck_in_queued_failure_callback_for_removed_mapped_task(
request = mock_executors[0].send_callback.call_args[0][0]
assert request.ti.map_index == 1
+ @pytest.mark.parametrize(
+ ("retries", "expected_state", "expected_email_type"),
+ [
+ pytest.param(0, TaskInstanceState.FAILED, "failure",
id="no_retries_left"),
+ pytest.param(1, TaskInstanceState.UP_FOR_RETRY, "retry",
id="retries_left"),
+ ],
+ )
+ def test_stuck_in_queued_tries_exceeded_honours_retries(
+ self, dag_maker, session, mock_executors, retries, expected_state,
expected_email_type
+ ):
+ """A task failed for being stuck in queued must retry like any other
failure."""
+ with dag_maker(dag_id="stuck_retries", serialized=True):
+ upstream = EmptyOperator(
+ task_id="upstream",
+ retries=retries,
+ email="[email protected]",
+ email_on_failure=True,
+ email_on_retry=True,
+ on_failure_callback=TestSchedulerJob.mock_failure_callback,
+ )
+ upstream >> EmptyOperator(task_id="downstream")
+ dr = dag_maker.create_dagrun(state=DagRunState.RUNNING)
+ ti = dr.get_task_instance("upstream", session=session)
+ failed_id, failed_try_number = ti.id, ti.try_number
+ ti.state = TaskInstanceState.QUEUED
+ ti.queued_dttm = timezone.utcnow()
+ session.commit()
+
+ scheduler = SchedulerJobRunner(job=Job(), num_runs=0)
+ scheduler._task_queued_timeout = -300
+ scheduler._num_stuck_queued_retries = 0
+ executor = mock_executors[0]
+ executor.get_task_key.side_effect = lambda ti: TaskInstanceUuid(ti.id)
+ with _loader_mock(mock_executors):
+ scheduler._handle_tasks_stuck_in_queued()
+
+ session.expire_all()
+ ti = dr.get_task_instance("upstream", session=session)
+ assert ti.state == expected_state
+
+ sent = [call.args[0] for call in executor.send_callback.call_args_list]
+ callback, email = sent
+ assert isinstance(callback, TaskCallbackRequest)
+ assert callback.task_callback_type == expected_state
+ assert isinstance(email, EmailRequest)
+ assert email.email_type == expected_email_type
+ for request in sent:
+ assert request.ti.id == failed_id
+ assert request.ti.try_number == failed_try_number
+ executor.fail.assert_called_once_with(TaskInstanceUuid(failed_id))
+
+ if retries:
+ # The retry is a new attempt; downstream must still be waiting for
it.
+ assert ti.id != failed_id
+ assert ti.try_number == failed_try_number + 1
+ assert ti.retry_at is not None
+ dr.task_instance_scheduling_decisions(session=session)
+ assert dr.get_task_instance("downstream", session=session).state
is None
+
+ @pytest.mark.parametrize("dag_loads", [True, False], ids=["no_callbacks",
"dag_not_loadable"])
+ def test_stuck_in_queued_tries_exceeded_retries_without_callbacks(
+ self, dag_maker, session, mock_executors, dag_loads
+ ):
+ """Retries do not depend on the task having a callback, or on the DAG
being loadable."""
+ with dag_maker(dag_id="stuck_no_callbacks", serialized=True):
+ EmptyOperator(task_id="op", retries=1)
+ dr = dag_maker.create_dagrun(state=DagRunState.RUNNING)
+ ti = dr.get_task_instance("op", session=session)
+ ti.state = TaskInstanceState.QUEUED
+ ti.queued_dttm = timezone.utcnow()
+ session.commit()
+
+ scheduler = SchedulerJobRunner(job=Job(), num_runs=0)
+ scheduler._task_queued_timeout = -300
+ scheduler._num_stuck_queued_retries = 0
+ with _loader_mock(mock_executors), ExitStack() as stack:
+ if not dag_loads:
+ stack.enter_context(
+ mock.patch.object(
+ scheduler.scheduler_dag_bag, "get_dag_for_run",
side_effect=RuntimeError("gone")
Review Comment:
Please add `autospec=True`. Also, `dag_not_loadable` could be its own test
instead of the `ExitStack` branch, and the ids could be `["dag_loads",
"dag_not_loadable"]`.
##########
airflow-core/tests/unit/jobs/test_scheduler_job.py:
##########
@@ -11840,6 +11844,133 @@ def
test_stuck_in_queued_failure_callback_for_removed_mapped_task(
request = mock_executors[0].send_callback.call_args[0][0]
assert request.ti.map_index == 1
+ @pytest.mark.parametrize(
+ ("retries", "expected_state", "expected_email_type"),
+ [
+ pytest.param(0, TaskInstanceState.FAILED, "failure",
id="no_retries_left"),
+ pytest.param(1, TaskInstanceState.UP_FOR_RETRY, "retry",
id="retries_left"),
+ ],
+ )
+ def test_stuck_in_queued_tries_exceeded_honours_retries(
+ self, dag_maker, session, mock_executors, retries, expected_state,
expected_email_type
+ ):
+ """A task failed for being stuck in queued must retry like any other
failure."""
+ with dag_maker(dag_id="stuck_retries", serialized=True):
+ upstream = EmptyOperator(
+ task_id="upstream",
+ retries=retries,
+ email="[email protected]",
+ email_on_failure=True,
+ email_on_retry=True,
+ on_failure_callback=TestSchedulerJob.mock_failure_callback,
+ )
+ upstream >> EmptyOperator(task_id="downstream")
+ dr = dag_maker.create_dagrun(state=DagRunState.RUNNING)
+ ti = dr.get_task_instance("upstream", session=session)
+ failed_id, failed_try_number = ti.id, ti.try_number
+ ti.state = TaskInstanceState.QUEUED
+ ti.queued_dttm = timezone.utcnow()
+ session.commit()
+
+ scheduler = SchedulerJobRunner(job=Job(), num_runs=0)
+ scheduler._task_queued_timeout = -300
+ scheduler._num_stuck_queued_retries = 0
+ executor = mock_executors[0]
+ executor.get_task_key.side_effect = lambda ti: TaskInstanceUuid(ti.id)
+ with _loader_mock(mock_executors):
+ scheduler._handle_tasks_stuck_in_queued()
+
+ session.expire_all()
+ ti = dr.get_task_instance("upstream", session=session)
+ assert ti.state == expected_state
+
+ sent = [call.args[0] for call in executor.send_callback.call_args_list]
+ callback, email = sent
+ assert isinstance(callback, TaskCallbackRequest)
+ assert callback.task_callback_type == expected_state
+ assert isinstance(email, EmailRequest)
+ assert email.email_type == expected_email_type
+ for request in sent:
+ assert request.ti.id == failed_id
+ assert request.ti.try_number == failed_try_number
+ executor.fail.assert_called_once_with(TaskInstanceUuid(failed_id))
+
+ if retries:
+ # The retry is a new attempt; downstream must still be waiting for
it.
+ assert ti.id != failed_id
+ assert ti.try_number == failed_try_number + 1
+ assert ti.retry_at is not None
+ dr.task_instance_scheduling_decisions(session=session)
+ assert dr.get_task_instance("downstream", session=session).state
is None
+
+ @pytest.mark.parametrize("dag_loads", [True, False], ids=["no_callbacks",
"dag_not_loadable"])
+ def test_stuck_in_queued_tries_exceeded_retries_without_callbacks(
+ self, dag_maker, session, mock_executors, dag_loads
+ ):
+ """Retries do not depend on the task having a callback, or on the DAG
being loadable."""
+ with dag_maker(dag_id="stuck_no_callbacks", serialized=True):
+ EmptyOperator(task_id="op", retries=1)
+ dr = dag_maker.create_dagrun(state=DagRunState.RUNNING)
+ ti = dr.get_task_instance("op", session=session)
+ ti.state = TaskInstanceState.QUEUED
+ ti.queued_dttm = timezone.utcnow()
+ session.commit()
+
+ scheduler = SchedulerJobRunner(job=Job(), num_runs=0)
+ scheduler._task_queued_timeout = -300
+ scheduler._num_stuck_queued_retries = 0
+ with _loader_mock(mock_executors), ExitStack() as stack:
+ if not dag_loads:
+ stack.enter_context(
+ mock.patch.object(
+ scheduler.scheduler_dag_bag, "get_dag_for_run",
side_effect=RuntimeError("gone")
+ )
+ )
+ scheduler._handle_tasks_stuck_in_queued()
+
+ session.expire_all()
+ assert dr.get_task_instance("op", session=session).state ==
TaskInstanceState.UP_FOR_RETRY
+ mock_executors[0].send_callback.assert_not_called()
+ mock_executors[0].fail.assert_called_once()
+
+ @pytest.mark.parametrize("num_stuck_queued_retries", [0, 2],
ids=["would_fail", "would_requeue"])
+ def test_stuck_in_queued_skips_task_started_after_scan(
+ self, dag_maker, session, mock_executors, num_stuck_queued_retries
+ ):
+ """A worker can start the attempt between the scan and handling it;
leave that attempt alone."""
+ with dag_maker(dag_id="stuck_started_after_scan", serialized=True):
+ EmptyOperator(task_id="op", retries=1)
+ dr = dag_maker.create_dagrun(state=DagRunState.RUNNING)
+ ti = dr.get_task_instance("op", session=session)
+ ti.state = TaskInstanceState.QUEUED
+ ti.queued_dttm = timezone.utcnow()
+ session.commit()
+
+ scheduler = SchedulerJobRunner(job=Job(), num_runs=0)
+ scheduler._task_queued_timeout = -300
+ scheduler._num_stuck_queued_retries = num_stuck_queued_retries
+ stale_tis = list(scheduler._get_tis_stuck_in_queued(session))
+ assert [t.id for t in stale_tis] == [ti.id]
+ # The worker starts the attempt after the scan; the scheduler still
holds the queued copy.
+ session.execute(
+ update(TaskInstance).where(TaskInstance.id ==
ti.id).values(state=TaskInstanceState.RUNNING)
+ )
+ session.commit()
+
+ with (
+ _loader_mock(mock_executors),
+ mock.patch.object(scheduler, "_get_tis_stuck_in_queued",
return_value=stale_tis),
Review Comment:
`autospec=True` here too.
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3617,38 +3635,64 @@ def _maybe_requeue_stuck_ti(self, *, ti, session,
executor):
"The DAG or task could not be found. If a failure callback
exists, it will not be run.",
exc_info=True,
)
- else:
- if task.has_on_failure_callback:
- if inspect(ti).detached:
- ti = session.merge(ti)
- _stuck_bundle_name, _stuck_bundle_version,
_stuck_version_data = (
- _resolve_ti_callback_bundle_info(ti)
+ task = None
+ # handle_failure() reads retries and fail_fast from ti.task.
+ ti.task = task
+ task_callback_type = (
+ TaskInstanceState.UP_FOR_RETRY if ti.is_eligible_to_retry()
else TaskInstanceState.FAILED
+ )
+ needs_callback = task is not None and
(task.has_on_failure_callback or task.has_on_retry_callback)
+ needs_email = (
+ task is not None and bool(task.email) and
(task.email_on_failure or task.email_on_retry)
+ )
+
+ # Built before handle_failure(), which archives this attempt when
it allocates a retry.
+ # Pydantic requires a dag_version_id UUID.
+ if (needs_callback or needs_email) and
_ensure_ti_has_dag_version_id(ti, session, self.log):
+ bundle_name, bundle_version, version_data =
_resolve_ti_callback_bundle_info(ti)
Review Comment:
This callback/email block now exists in 3 places (here, heartbeat purge,
`process_executor_events`), and they already differ. Worth a shared helper?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
