Call listeners for running task instance when a Dag Run state is manually set#69874
Call listeners for running task instance when a Dag Run state is manually set#69874kacpermuda wants to merge 2 commits into
Conversation
kacpermuda
requested review from
bugraoz93,
choo121600,
ephraimbuddy,
henry3260,
jason810496,
pierrejeambrun,
rawwar and
shubhamraj-git
as code owners
kacpermuda
force-pushed
the
dag-run-ti-listeners
branch
from
a1ba686 to
d5ff798
Compare
amoghrajesh
left a comment
amoghrajesh
left a comment
There was a problem hiding this comment.
Good investigation but some comments need handling
There was a problem hiding this comment.
What about teardown task instances that are left running by design (so they can finish their own cleanup)? They are still included in the running_tis list returned by _set_dag_run_terminal_state, so the API server will fire a false "success"/"failed" terminal event for a teardown task that is, in fact, still executing.
This might've been an untested edge case even before this PR but still flagging.
There was a problem hiding this comment.
Good spot ! Thanks, excluded teardown tasks now.
| @pytest.mark.parametrize( | ||
| ("dag_run_state", "expected_ti_listener_state"), | ||
| [ | ||
| ("success", TaskInstanceState.SUCCESS), | ||
| ("failed", TaskInstanceState.FAILED), | ||
| ], | ||
| ) | ||
| @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle") | ||
| def test_patch_dag_run_notifies_ti_listeners_for_running_tasks( | ||
| self, | ||
| test_client, | ||
| dag_maker, | ||
| session, | ||
| listener_manager, | ||
| dag_run_state, | ||
| expected_ti_listener_state, | ||
| ): |
There was a problem hiding this comment.
We should also cover for teardown tasks, both here and in the other test.
There was a problem hiding this comment.
Added a separate tests for teardown case
pierrejeambrun
left a comment
pierrejeambrun
left a comment
There was a problem hiding this comment.
Small nit, in addition of Amogh's comments
| dag=dag, run_id=dag_run.run_id, commit=True, session=session | ||
| ) | ||
| try: | ||
| _emit_state_listener_hooks(running_tis, TaskInstanceState.FAILED) |
There was a problem hiding this comment.
_emit_state_listener_hooks already wraps everything in a broad try/except I don't think that's needed here
There was a problem hiding this comment.
Correct, I missed that. Thanks !
| try: | ||
| _emit_state_listener_hooks(running_tis, TaskInstanceState.SUCCESS) | ||
| except Exception: | ||
| log.exception("error calling listener") |
There was a problem hiding this comment.
Same here. This can never catch anything
There was a problem hiding this comment.
Correct, I missed that. Thanks !
kacpermuda
force-pushed
the
dag-run-ti-listeners
branch
from
d5ff798 to
b9f98d3
Compare
3 participants
When a task instance is individually set to success or failure via the API,
on_task_instance_success/on_task_instance_failedis called from the API server, giving listeners a terminal event to act on. The same does not happen when an entire Dag Run is force-set to success or failure: running tasks are bulk-transitioned in the database, but their task-instance listeners are never fired. Listeners that registered a start event (on_task_instance_running) for those tasks receive no completion event, leaving them — notably OpenLineage — with a job permanently stuck in RUNNING.This PR closes the gap:
_set_dag_run_terminal_statenow returns(all_updated_tis, running_tis)whererunning_tisis the snapshot of task instances that were in an active state (RUNNING, DEFERRED, UP_FOR_RESCHEDULE, AWAITING_INPUT) before the bulk transition, captured without an extra query.set_dag_run_state_to_success/set_dag_run_state_to_failedexpose the same tuple so callers have access to the running subset.patch_dag_run_statefireson_task_instance_success/on_task_instance_failedfor those running task instances before the dag-run-level hook. Each group has its owntry/exceptso a listener failure in one cannot suppress the other.Tasks that were only QUEUED or SCHEDULED (never started, no start event fired) are excluded — they transition to SKIPPED and receive no listener call, consistent with their lifecycle.
Why the terminal event is missing
The gap is a deliberate design decision that becomes a problem specifically for the dag run force-terminate case:
Result: on_task_instance_running fired from the task runner → task process killed → nothing on the other end. OpenLineage has a dangling RUNNING job with no terminal event. We can't extend the live of the worker process, but we can fire from the api server in that case.
Was generative AI tooling used to co-author this PR?
Generated-by: Claude Code (claude-sonnet-4-6) following the guidelines
{pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.