hkc-8010 commented on code in PR #68008:
URL: https://github.com/apache/airflow/pull/68008#discussion_r3589726451
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3620,6 +3620,7 @@ def _purge_task_instances_without_heartbeats(
request,
)
self.executor.send_callback(request)
+
ti.handle_failure(error=str(task_instance_heartbeat_timeout_message_details),
session=session)
Review Comment:
Fixed. Loads `ti.task` from `self.scheduler_dag_bag` before calling
`handle_failure()`, so `fail_fast` (`ti.task.dag.fail_fast`) actually runs, and
sends an `EmailRequest` directly from this path since
`process_executor_events`'s external-kill branch can't reach this TI anymore
once `handle_failure()` moves it out of `RUNNING`/`RESTARTING`. Also replaced
the `max_tries > 0` guard with plain `ti.is_eligible_to_retry()` so the
callback type can't disagree with what `handle_failure()` actually persists for
a `RESTARTING` TI with `max_tries=0`. Covered by
`test_heartbeat_timeout_preserves_failure_email`,
`test_heartbeat_timeout_restarting_zero_max_tries_matches_final_state`, and
`test_heartbeat_timeout_honors_fail_fast`.
##########
airflow-core/tests/unit/jobs/test_scheduler_job.py:
##########
@@ -8677,6 +8677,82 @@ 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
+ def test_heartbeat_timeout_converges_ti_state_before_next_scan(self,
dag_maker, session):
+ """A heartbeat-timed-out TI should not be found again on the next
scheduler scan."""
+ with dag_maker(dag_id="test_heartbeat_timeout_dedupe",
session=session):
+ EmptyOperator(task_id="test_task", on_failure_callback=lambda
context: None)
+
+ 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 = TaskInstanceState.RUNNING
+ ti.try_number = 1
+ ti.max_tries = 0
+ ti.queued_by_job_id = scheduler_job.id
+ ti.start_date = timezone.utcnow() - timedelta(seconds=900)
+ ti.last_heartbeat_at = timezone.utcnow() - timedelta(seconds=600)
+ session.merge(ti)
+ session.commit()
+
+ self.job_runner._find_and_purge_task_instances_without_heartbeats()
+
+ session.expire_all()
+ ti.refresh_from_db(session=session)
+ assert ti.state != TaskInstanceState.RUNNING
+ assert
self.job_runner._find_task_instances_without_heartbeats(session=session) == []
+
+ self.job_runner._find_and_purge_task_instances_without_heartbeats()
+
+ mock_executor.send_callback.assert_called_once()
+
+ @pytest.mark.parametrize(
+ ("retries", "callback_kind", "expected"),
+ [
+ (1, "retry", TaskInstanceState.UP_FOR_RETRY),
+ (0, "failure", TaskInstanceState.FAILED),
+ ],
+ )
+ def test_heartbeat_timeout_sets_callback_type_param(
+ self, dag_maker, session, retries, callback_kind, expected
+ ):
Review Comment:
Added
`test_heartbeat_timeout_restarting_zero_max_tries_matches_final_state`: a
`RESTARTING` TI with `max_tries=0` now gets `task_callback_type=UP_FOR_RETRY`,
matching the actual persisted state.
--
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]
