ephraimbuddy commented on code in PR #68008:
URL: https://github.com/apache/airflow/pull/68008#discussion_r3588820843
##########
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:
This moves the TI out of `RUNNING` before the executor event is processed,
so `_process_executor_events()` skips the event at its `ti_queued` guard. That
also skips the `EmailRequest` path, and `ti.task` is not loaded here so
`handle_failure()` silently disables `fail_fast`; heartbeat-timeout failures
will therefore lose failure/retry emails and fail-fast handling. I think this
is a blocker.
##########
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:
`_find_task_instances_without_heartbeats()` also returns `RESTARTING` TIs.
With `max_tries=0`, the request is typed `FAILED`, while `handle_failure()`
treats `RESTARTING` as always retryable and persists `UP_FOR_RETRY`, so
`on_failure_callback` can run for a TI that is retried. Please use one retry
decision for both and add a `RESTARTING` case.
##########
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
+ ):
+ """Heartbeat timeout should mark callback type based on retry
eligibility."""
+ with dag_maker(dag_id=f"heartbeat_timeout_{callback_kind}",
session=session):
+ if callback_kind == "retry":
+ EmptyOperator(task_id="t1", retries=retries,
on_retry_callback=lambda ctx: None)
+ else:
+ EmptyOperator(task_id="t1", retries=retries,
on_failure_callback=lambda ctx: None)
+
+ dag_run = dag_maker.create_dagrun(run_id="test_run",
state=DagRunState.RUNNING)
+
+ executor = MockExecutor(do_update=False)
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(scheduler_job,
executors=[executor])
+
+ ti = dag_run.get_task_instance(task_id="t1")
+ ti.state = TaskInstanceState.RUNNING
+ ti.try_number = 1
+ ti.max_tries = retries
+ 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()
+
+ self.job_runner.executor.callback_sink.send.assert_called_once()
+ request = self.job_runner.executor.callback_sink.send.call_args[0][0]
+ assert isinstance(request, TaskCallbackRequest)
+ assert request.task_callback_type == expected
+
+ session.expire_all()
+ ti.refresh_from_db(session=session)
+ assert ti.state == expected
+
Review Comment:
Add this test:
```
@time_machine.travel(DEFAULT_DATE, tick=False)
def test_heartbeat_timeout_preserves_failure_email(self, dag_maker, session):
with dag_maker(dag_id="heartbeat_timeout_email", session=session):
EmptyOperator(
task_id="t1",
email="[email protected]",
email_on_failure=True,
)
dag_run = dag_maker.create_dagrun(run_id="test_run",
state=DagRunState.RUNNING)
executor = MockExecutor(do_update=False)
scheduler_job = Job()
self.job_runner = SchedulerJobRunner(scheduler_job, executors=[executor])
ti = dag_run.get_task_instance(task_id="t1")
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()
self.job_runner._process_executor_events(executor=executor,
session=session)
requests = [call.args[0] for call in
executor.callback_sink.send.call_args_list]
email_requests = [request for request in requests if isinstance(request,
EmailRequest)]
assert len(email_requests) == 1
assert email_requests[0].email_type == "failure"
```
--
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]
