hkc-8010 commented on code in PR #76908:
URL: https://github.com/apache/airflow/pull/76908#discussion_r4537367264


##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3617,38 +3619,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)
+                runtime_ti = task_instance_to_runtime(
+                    ti, model=TIDataModel, 
map_index=coordinates.public_map_index(ti)
+                )
+                context_from_server = TIRunContext(
+                    dag_run=DRDataModel.model_validate(ti.dag_run, 
from_attributes=True),
+                    max_tries=ti.max_tries,
+                    variables=[],
+                    connections=[],
+                    xcom_keys_to_clear=[],
+                )
+                if needs_callback:
+                    executor.send_callback(
+                        TaskCallbackRequest(
+                            filepath=ti.dag_model.relative_fileloc or "",
+                            bundle_name=bundle_name,
+                            bundle_version=bundle_version,
+                            version_data=version_data,
+                            ti=runtime_ti,
+                            msg=msg,
+                            task_callback_type=task_callback_type,
+                            context_from_server=context_from_server,
+                        )
                     )
-                    # Backfill dag_version_id for legacy tasks (Pydantic 
requires uuid.UUID).
-                    # Note: we cannot use `continue` here because this method 
is not
-                    # inside a loop.  If backfilling fails we simply skip the 
callback.
-                    if _ensure_ti_has_dag_version_id(ti, session, self.log):
-                        request = TaskCallbackRequest(
+                # The TI is no longer queued when its executor event arrives, 
so that path won't email.
+                if needs_email:
+                    executor.send_callback(
+                        EmailRequest(
                             filepath=ti.dag_model.relative_fileloc or "",
-                            bundle_name=_stuck_bundle_name,
-                            bundle_version=_stuck_bundle_version,
-                            version_data=_stuck_version_data,
-                            ti=task_instance_to_runtime(
-                                ti, model=TIDataModel, 
map_index=coordinates.public_map_index(ti)
-                            ),
+                            bundle_name=bundle_name,
+                            bundle_version=bundle_version,
+                            version_data=version_data,
+                            ti=runtime_ti,
                             msg=msg,
-                            context_from_server=TIRunContext(
-                                dag_run=ti.dag_run,
-                                max_tries=ti.max_tries,
-                                variables=[],
-                                connections=[],
-                                xcom_keys_to_clear=[],
+                            email_type=(
+                                "retry" if task_callback_type == 
TaskInstanceState.UP_FOR_RETRY else "failure"
                             ),
+                            context_from_server=context_from_server,
                         )
-                        executor.send_callback(request)
-            finally:
-                ti.set_state(TaskInstanceState.FAILED, session=session)
-                executor.fail(executor.get_task_key(ti))
+                    )
+
+            failed_key = executor.get_task_key(ti)
+            ti.handle_failure(error=msg, session=session)

Review Comment:
   Good catch, thanks. Added a locked recheck in 
`_handle_tasks_stuck_in_queued` (skip_locked, `state == queued`) so a TI a 
worker already picked up is skipped before we revoke, requeue or fail it. It 
covers the requeue branch too, since that had the same problem. Test is 
`test_stuck_in_queued_skips_task_started_after_scan`.



-- 
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]

Reply via email to