namanjain24-sudo opened a new pull request, #73213: URL: https://github.com/apache/airflow/pull/73213
`SchedulerJobRunner.process_executor_events` calls `ti.set_state()` without `session` in two places: when a cleared (`RESTARTING`) task instance is reported as successfully terminated, and when the Dag of a finished task instance cannot be loaded. `TaskInstance.set_state` is `@provide_session` and `settings.Session` is scoped, so `create_session()` returns the scheduler's own session and commits and closes it on exit. This is the same mechanism as #67850 and #71968. `process_executor_events` runs inside `with create_session()` in `_run_scheduler_loop`, not under `prohibit_commit`, so nothing raises. Instead, in the middle of an executor-event batch: - the scheduler's transaction is committed early, which releases the `FOR UPDATE SKIP LOCKED` row locks taken on the batch's task instances; - `close()` detaches the task instances loaded for the batch, so changes made to them afterwards are never written. In a local run with one `RESTARTING` → `SUCCESS` event and four `QUEUED` events carrying an external executor id, none of the four `external_executor_id` values reached the database without this change, and all four did with it (SQLite and PostgreSQL 16). This passes the scheduler's session at both call sites, as `_enqueue_task_instances_with_queued_state` already does for its own `ti.set_state()` call. Tests: - `test_process_executor_events_sets_state_in_callers_transaction` covers both paths: the state change has to roll back with the caller's transaction. On main both cases fail (the task instance is already committed as `None` / `failed`); with this change they pass. - The rest of `test_scheduler_job.py` and prek, including mypy for airflow-core, pass locally on SQLite and PostgreSQL 16. **Not covered here.** In the same loop, `executor.send_callback()` reaches `DatabaseCallbackSink.send`, which is also `@provide_session` and gets no session. So when the task instance has an `on_failure_callback` / `on_retry_callback` (reproduced locally with this change applied) or email configured, the scheduler's transaction is still committed at that call. `_maybe_requeue_stuck_ti` and `_purge_task_instances_without_heartbeats` call `send_callback()` the same way. Fixing that means either adding a `session` argument to `BaseExecutor.send_callback` and `BaseCallbackSink.send`, which executors that override `send_callback` would then have to accept, or handing the session to the sink directly, as `Trigger` already does with `DatabaseCallbackSink().send(callback=request, session=session)`. I kept this PR to the `set_state` calls and am happy to follow up with whichever approach you prefer. --- ##### Was generative AI tooling used to co-author this PR? - [X] Yes (please specify the tool below) Generated-by: a Gen-AI coding assistant, following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions). I reviewed the change and ran the checks above. -- 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]
