ChickenBenny commented on code in PR #76544:
URL: https://github.com/apache/airflow/pull/76544#discussion_r4504731218
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -2844,7 +2844,6 @@ def save_to_db(ti: TaskInstance, *, session: Session =
NEW_SESSION):
ti.updated_at = timezone.utcnow()
session.merge(ti)
session.flush()
- session.commit()
Review Comment:
Before this, the first `handle_failure()` in `process_executor_events`
committed the whole scheduler transaction mid-loop. That also released the
`with_row_locks(..., skip_locked=True)` locks on the rest of the batch, so
another HA scheduler could pick those TIs up while this one was still
iterating. After this change, failures are committed once per batch with the
rest of the loop's work, and an exception raised outside the savepoint later in
the batch now rolls back the earlier TIs too, where previously they had already
been committed.
Both seem fine to me (and the first is a fix), but since the behaviour
changes, it would help reviewers and anyone bisecting later if this were
described, or kept as a separate commit.
##########
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:
In the executor-event path the failure/retry callback and the email request
are sent before this helper runs (L1826, L1858). `DatabaseCallbackSink.send()`
only does `session.add()`, and `begin_nested()` flushes that before the
savepoint opens, so rolling back the savepoint does not remove them.
For a task that still has retries left, the queued callback/email says
`UP_FOR_RETRY` (`task_callback_type` is computed from
`ti.is_eligible_to_retry()` at L1813), but this fallback then puts the TI into
a terminal `FAILED` and it never retries. So the user's `on_retry_callback`
fires for a retry that never happens.
The test only uses `max_tries = 0`, so it does not cover this. Either moving
the callback/email requests inside the same savepoint, or choosing a fallback
that respects retries, would keep them consistent. Whichever state is chosen,
it would help to state in the PR description that the fallback skips retries,
fail_fast and listeners, since the issue left the fallback state as an open
question.
##########
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 write the "state mismatch" `Log`, load the task,
and build and send the `TaskCallbackRequest` / `EmailRequest`
(`_resolve_ti_callback_bundle_info`, `task_instance_to_runtime`,
`DRDataModel.model_validate`, ...). The heartbeat purge 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:
```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 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:
...
with cls._isolate_failure_handling(ti, session=session):
session.add(Log(event="state mismatch", ...))
# load task, send callback / email requests (unchanged)
ti.handle_failure(error=msg, session=session)
```
--
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]
