Skip to content

Cancel the Dataflow job when a user kills the deferred task#69586

Open
steveahnahn wants to merge 1 commit into
apache:mainfrom
steveahnahn:dataflow-cancel-job-on-user-kill-deferred
Open

Cancel the Dataflow job when a user kills the deferred task#69586
steveahnahn wants to merge 1 commit into
apache:mainfrom
steveahnahn:dataflow-cancel-job-on-user-kill-deferred

Conversation

@steveahnahn

@steveahnahn steveahnahn commented Jul 8, 2026

Copy link
Copy Markdown
Contributor

When a user kills a deferred Dataflow task (mark failed, clear, or mark succeeded), the Dataflow job itself keeps running in Google Cloud. TemplateJobStartTrigger has no on_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 legacy asyncio.CancelledError path cancelled the job from inside the trigger, so for deferrable users this is also a regression compared to earlier releases.

This adds on_kill to TemplateJobStartTrigger, following the same pattern as the merged BigQuery, Dataproc and EMR (#64770) implementations. It checks cancel_on_kill, job_id and project_id, then cancels through the synchronous DataflowHook. The drain_pipeline setting 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 in on_kill. Unlike the AWS hooks, DataflowHook resolves 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 the try, so a connection error is logged rather than escaping.

End to end verification

Verified against a real Dataflow job on GCP:

  1. A deferrable DataflowTemplatedJobStartOperator launched a real job; the task reached deferred and the job reached JOB_STATE_RUNNING.
  2. Marking the deferred task failed drove the trigger through the user action cancel path. Trigger log: Trigger cancelled by user action, invoking on_kill then Stopping Dataflow job. Project ID: ... Job ID: 2026-07-07_23_40_11-..., drain: False, with no exception escaping on_kill.
  3. The Dataflow job transitioned JOB_STATE_RUNNING to JOB_STATE_CANCELLING within ~12s and then to JOB_STATE_CANCELLED.

Tests

Seven new trigger tests: on_kill cancels with the expected hook construction and arguments, drain_pipeline is honored, it is a no op when cancel_on_kill is False or job_id or project_id is missing, hook errors are swallowed and logged, and hook construction errors are swallowed. On the code before this change, the new tests plus the updated serialize expectation 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?
  • Yes - Claude Code (Opus 4.8)

Generated-by: Claude Code (Opus 4.8) following the guidelines

@boring-cyborg boring-cyborg Bot added area:providers provider:google Google (including GCP) related issues labels Jul 8, 2026
@steveahnahn steveahnahn force-pushed the dataflow-cancel-job-on-user-kill-deferred branch from e368463 to cadba9a Compare July 8, 2026 05:41
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.
@steveahnahn steveahnahn force-pushed the dataflow-cancel-job-on-user-kill-deferred branch from cadba9a to f4df591 Compare July 8, 2026 06:46
@steveahnahn steveahnahn marked this pull request as ready for review July 8, 2026 21:02
@steveahnahn steveahnahn requested a review from shahar1 as a code owner July 8, 2026 21:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:providers provider:google Google (including GCP) related issues

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant