Add on_kill() to DatabricksTaskBaseOperator to cancel runs on task kill

[victorymakes]

Summary

DatabricksSubmitRunOperator and DatabricksRunNowOperator both implement on_kill() to cancel the Databricks run when an Airflow task is killed (SIGTERM or execution_timeout). DatabricksTaskBaseOperator — the base for DatabricksTaskOperator and DatabricksNotebookOperator — is missing the same implementation, so Databricks jobs continue running after the Airflow task is killed, orphaning compute resources and incurring unnecessary cloud spend.

DatabricksWorkflowTaskGroup received on_kill() in #42115; this PR closes the remaining gap for standalone task operators.

Changes

• Adds on_kill() to DatabricksTaskBaseOperator using self.databricks_run_id, which is:
  • initialised to None in __init__ (no AttributeError risk)
  • set by _launch_job() the moment the run is submitted — earlier than any polling or permission calls
• Adds unit tests covering both the cancel and no-op paths

Testing

pytest providers/databricks/tests/unit/databricks/operators/test_databricks.py -k "on_kill"

Was generative AI tooling used to co-author this PR?

• Yes — Claude Code (Sonnet 4.6)

[Vamsi-klu]

databricks_run_id is the shared parent workflow-run ID for a member of DatabricksWorkflowTaskGroup. Cancelling it here would therefore stop sibling tasks as well.

Could on_kill mirror monitor_databricks_job: use _get_current_databricks_task()["run_id"] for workflow members, while retaining self.databricks_run_id for standalone operators? A regression test should set parent run ID 1, return child attempt ID 999, and assert cancel_run(999).

The branch also contains two fix: commit subjects, which current Airflow commit checks reject; those will need rewriting when the branch is rebased.

---

Drafted-by: Codex (GPT-5)

[victorymakes]

Thanks for the review!

You're right — self.databricks_run_id is the shared parent workflow run ID for workflow members, so cancelling it would stop sibling tasks. Fixed in the latest commit.

on_kill() now mirrors monitor_databricks_job: for workflow members it calls _get_current_databricks_task()["run_id"] to get the child task's own run ID, while standalone operators continue to cancel via self.databricks_run_id directly.

Also added the regression test you suggested: parent run_id=1, _get_current_databricks_task returns child run_id=999, asserts cancel_run(999).

The branch has also been squashed to a single commit to fix the fix: subject check.

[Vamsi-klu]

Thanks for the quick turnaround. I re-checked HEAD (feb6783) end-to-end.

Confirmed fixed

• Workflow members cancel via _get_current_databricks_task()["run_id"] only (not the shared parent run id)
• On child-run resolve failure: log + return (no parent fallback)
• Unit coverage:
  • standalone cancel
  • no-op when databricks_run_id is None
  • workflow member child cancel (parent=1 -> cancel_run(999))
  • workflow member exception path (cancel_run not called)
• Single commit; subject no longer uses a rejected fix: prefix

Residual before this is ready

Both touched files drop the first two lines of the standard ASF license header:

#
# Licensed to the Apache Software Foundation (ASF) under one

Please restore the full header in:

1. providers/databricks/src/airflow/providers/databricks/operators/databricks.py
2. providers/databricks/tests/unit/databricks/operators/test_databricks.py

Testing bar

No live Databricks credentials/UI testing needed for merge. This matches how DatabricksSubmitRunOperator / DatabricksRunNowOperator on_kill is validated: mock cancel_run and assert the correct run id.

A rebase onto current main will still be needed before merge (branch is quite behind). That can land together with the header fix.

Logic otherwise looks good from my side.

[victorymakes]

Rebased onto current main and fixed the missing license header. Branch is now a single commit on top of 84e520a.

[eladkal]

cc @moomindani for Databricks team review

[moomindani]

LGTM. I verified the run-id semantics both against internal Databricks documentation and empirically on a live workspace, because that is the crux of the original objection about cancelling the shared parent run.

Docs: the Jobs CLI/API reference states runs/cancel "Cancels a job run or a task run", so passing a child task run id is a supported operation, not an accident that happens to work. The internal Jobs Task API design doc confirms the id model — job_run_id is the parent, and multitask runs have multiple task runs grouped under it.

Empirically, with a two-task job running both tasks concurrently:

Action	task_a	task_b	parent run
cancel task_a's child run (post-fix)	TERMINATED/CANCELED	still RUNNING	RUNNING
cancel the parent run (pre-fix)	CANCELED	CANCELED	CANCELED

So that concern was exactly right, and the fix does what it claims: a killed Airflow task now cancels only its own Databricks task run and leaves siblings alone. Cancelling the parent really does take the siblings down with it.

The rest checks out: on_kill sits on DatabricksTaskBaseOperator so both DatabricksTaskOperator and DatabricksNotebookOperator inherit it, it mirrors monitor_databricks_job's _get_current_databricks_task()["run_id"] pattern, and refusing to fall back to the parent id on resolution failure is the right call — that fallback is precisely the sibling-cancellation I measured. 204 tests pass in the file, prek --stage pre-commit clean.

One thing worth recording rather than changing: in deferrable mode this on_kill does not fire, since the operator has left the worker by then. That is not a gap — DatabricksExecutionTrigger has its own async on_kill, and monitor_databricks_job defers with run_id=current_task_run_id, i.e. the child run, so the deferred path cancels the correct run too. Both paths are covered.

---

Drafted-by: Claude Code (Opus 5)