ChickenBenny commented on code in PR #76544:
URL: https://github.com/apache/airflow/pull/76544#discussion_r4504728264
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -1499,6 +1499,22 @@ def _emit_executor_events_batch_metrics(num_events: int)
-> None:
stats.gauge("scheduler.executor_events.batch_size", num_events)
stats.incr("scheduler.executor_events.processed", num_events)
+ @classmethod
+ def _handle_task_failure(cls, ti: TI, *, error: str, session: Session) ->
None:
+ task_instance_key = ti.key
+ # begin_nested flushes pending batch changes before opening the
savepoint.
+ savepoint = session.begin_nested()
+ try:
+ with savepoint:
+ ti.handle_failure(error=error, session=session)
+ except Exception as exc:
+ if not session.is_active or (isinstance(exc, DBAPIError) and
exc.connection_invalidated):
+ raise
+ cls.logger().exception(
+ "Could not handle task instance failure",
task_instance_key=task_instance_key
+ )
+ ti.set_state(TaskInstanceState.FAILED, session=session)
Review Comment:
Right now it only wraps the `handle_failure()` call. But when the scheduler
decides a TI has failed, `handle_failure()` is just the last step. Before it,
in the same branch, we also load the task and build and send the
`TaskCallbackRequest` / `EmailRequest` (`_resolve_ti_callback_bundle_info`,
`task_instance_to_runtime`, `DRDataModel.model_validate`, ...).
`_purge_task_instances_without_heartbeats` has a similar block. If anything in
there raises for one TI, we're back to the same crash this PR is trying to
prevent.
What do you think about moving the savepoint up so it covers that whole
"fail this TI" block? For example, turn the helper into a context manager,
keeping the guards you already have:
```python
@classmethod
@contextmanager
def _isolate_failure_handling(cls, ti: TI, *, session: Session) ->
Iterator[None]:
task_instance_key = ti.key
try:
with session.begin_nested():
yield
except Exception as exc:
if not session.is_active or (isinstance(exc, DBAPIError) and
exc.connection_invalidated):
raise
cls.logger().exception("Could not handle task instance failure",
task_instance_key=task_instance_key)
ti.set_state(TaskInstanceState.FAILED, session=session)
```
and use it around the existing code in `process_executor_events`:
```python
if ti_queued and not ti_requeued:
...
session.add(Log(event="state mismatch", ...)) # outside, so the audit
entry survives the fallback
with cls._isolate_failure_handling(ti, session=session):
# load task, send callback / email requests (unchanged)
ti.handle_failure(error=msg, session=session)
```
In the heartbeat purge, the same `with` would go from
`_resolve_ti_callback_bundle_info(ti)` down to the `handle_failure()` call,
leaving the `heartbeat timeout` `Log` outside.
This would also cover the callback/retry mismatch in the next comment, since
the queued requests would be rolled back together with the failed handling.
--
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]
