Cancel the Dataflow job when a user kills the deferred task#69586
Open
steveahnahn wants to merge 1 commit into
Open
Cancel the Dataflow job when a user kills the deferred task#69586steveahnahn wants to merge 1 commit into
steveahnahn wants to merge 1 commit into
Conversation
e368463 to
cadba9a
Compare
Killing a deferred Dataflow task (mark-failed, clear or mark-succeeded) leaves the Dataflow job running: the trigger has no on_kill, so the user-kill cancellation the sibling BigQuery and Dataproc triggers already perform never happens, and a streaming pipeline keeps running and billing until someone notices in the Cloud console. Before the Airflow 3.3 trigger cancellation redesign the legacy path cancelled the job from inside the trigger, so this is also a behavior regression for deferrable users. Implement on_kill on TemplateJobStartTrigger following the merged BigQuery/Dataproc/EMR pattern, honoring the operator's drain_pipeline setting so streaming jobs are drained rather than cancelled when configured, exactly as the non-deferred kill path does. The hook is built and invoked inside a worker thread because the synchronous DataflowHook resolves its connection during construction, which must not run in the triggerer's event loop.
cadba9a to
f4df591
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
When a user kills a deferred Dataflow task (mark failed, clear, or mark succeeded), the Dataflow job itself keeps running in Google Cloud.
TemplateJobStartTriggerhas noon_kill, so the cancellation that the BigQuery and Dataproc triggers already perform in the same situation never happens for Dataflow. For a streaming pipeline this means the job runs and bills until someone notices it in the Cloud console. Before the Airflow 3.3 trigger cancellation redesign, the legacyasyncio.CancelledErrorpath cancelled the job from inside the trigger, so for deferrable users this is also a regression compared to earlier releases.This adds
on_killtoTemplateJobStartTrigger, following the same pattern as the merged BigQuery, Dataproc and EMR (#64770) implementations. It checkscancel_on_kill,job_idandproject_id, then cancels through the synchronousDataflowHook. Thedrain_pipelinesetting is honored, so a streaming job is drained instead of cancelled when the operator is configured that way, matching what the non deferred kill path does.The hook is constructed and invoked inside a worker thread (
sync_to_async), not directly inon_kill. Unlike the AWS hooks,DataflowHookresolves its Airflow connection eagerly during construction, and that synchronous work must not run in the triggerer's event loop. Building it in the thread also keeps construction inside thetry, so a connection error is logged rather than escaping.End to end verification
Verified against a real Dataflow job on GCP:
DataflowTemplatedJobStartOperatorlaunched a real job; the task reacheddeferredand the job reachedJOB_STATE_RUNNING.Trigger cancelled by user action, invoking on_killthenStopping Dataflow job. Project ID: ... Job ID: 2026-07-07_23_40_11-..., drain: False, with no exception escapingon_kill.JOB_STATE_RUNNINGtoJOB_STATE_CANCELLINGwithin ~12s and then toJOB_STATE_CANCELLED.Tests
Seven new trigger tests:
on_killcancels with the expected hook construction and arguments,drain_pipelineis honored, it is a no op whencancel_on_killisFalseorjob_idorproject_idis missing, hook errors are swallowed and logged, and hook construction errors are swallowed. On the code before this change, the new tests plus the updatedserializeexpectation fail. All 88 tests in the file pass with the change, and provider mypy is clean on both touched files.Was generative AI tooling used to co-author this PR?
Generated-by: Claude Code (Opus 4.8) following the guidelines