Add durable execution to DatabricksSubmitRunOperator

[amoghrajesh]

---

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

• Yes - claude sonnet 4.6

Next application of resumablejobmixin!

Why we are doing this

DatabricksSubmitRunOperator submits a run to Databricks, gets back a run id, and then polls synchronously on the worker until the run finishes. That run id lives only in the worker process. If the worker crashes or is preempted mid-poll (eviction, OOM, or whatever reason), Airflow retries the task in a fresh process with no memory of the run id, so it submits a brand-new run. The original run keeps executing on Databricks, orphaned, while the retry runs a duplicate.

For long-running Databricks jobs this means paying twice (or more) for the same work, and it is a real operational pain for users running multi-hour jobs. Deferrable mode already protects the long wait (the Triggerer holds the run id), but a large share of users do not run a Triggerer, and deferrable optimizes the worker slot rather than the cost of the external job. This change makes the plain synchronous path crash safe with no new infrastructure.

Benefits this will bring in

• A worker crash and retry reconnects to the already-running Databricks run instead of submitting a duplicate, so you do not pay twice for the same job.
• If the prior run already finished successfully, the retry returns immediately without resubmitting or re-polling.
• Works on the existing synchronous operator with no Triggerer required.
• Enabled by default, so users get crash safety automatically (on Airflow 3.3+) with no Dag changes.

Approach

The operator now builds on the AIP-103 task state store. On the first run it persists the Databricks run id to the task state store before polling begins. On a retry it reads that id back and inspects the run's current state:

• still running: reconnect and keep polling, no new submission
• already succeeded: return immediately, no submission and no polling
• terminally failed: submit a fresh run

The task state store is scoped to the task instance and survives across retries, which is what makes the reconnect possible. Deferrable mode is unchanged and takes precedence when it is enabled.

Backcompat

• Fully backward compatible. On a clean run the observable submit-and-poll behavior is unchanged, the only addition is that the run id is also written to the task state store before polling.
• The task state store is an Airflow 3.3+ capability. On Airflow 2.x / pre-3.3 the operator degrades gracefully to exactly the old behavior (always submits fresh on retry) through a compatibility shim, so the provider keeps working on older Airflow.
• If the task state store is unavailable at runtime (for example, not configured), the operator logs that crash recovery is disabled and behaves exactly as before.
• No changes to existing parameters; one new optional flag is added and no migration is needed.

How to opt out

Set durable=False on the operator:

DatabricksSubmitRunOperator(..., durable=False)

This restores the previous behavior: always submit a fresh run on retry and never touch the task state store. It can also be set through default_args to opt out across a whole Dag or deployment.

Testing

Running this dag earlier and killing worker mid run would look like this:

from __future__ import annotations

from datetime import timedelta

import pendulum

from airflow.providers.databricks.operators.databricks import DatabricksSubmitRunOperator
from airflow.sdk import DAG

NOTEBOOK_PATH = "/Users/amoghdesai.oss@gmail.com/sleepy"

with DAG(
    dag_id="databricks_resumable_repro",
    schedule=None,
    start_date=pendulum.datetime(2025, 1, 1, tz="UTC"),
    catchup=False,
):
    DatabricksSubmitRunOperator(
        task_id="submit_sleepy",
        databricks_conn_id="databricks_default",
        tasks=[
            {
                "task_key": "sleepy",
                "notebook_task": {"notebook_path": NOTEBOOK_PATH},
            }
        ],
        deferrable=False,
        do_xcom_push=True,
        retries=1,
        retry_delay=timedelta(seconds=10),
    )

Before my changes

Killed mid run:

First run:

Worker comes back up:

Extra run submitted now due to that:

After my changes:

First run:

Worker comes back up:

Just one job run:

Tried killing the job and it kills external job too:

---

• Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
• For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {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.

[Copilot]

Pull request overview

This PR adds crash-safe (“durable”) synchronous execution to DatabricksSubmitRunOperator by integrating the Task SDK’s ResumableJobMixin / task state store so retries can reconnect to an existing Databricks run instead of submitting duplicates.

Changes:

• Extend DatabricksSubmitRunOperator with ResumableJobMixin hooks (submit_job, get_job_status, etc.) and add a durable option (default-enabled on Airflow 3.3+).
• Add/adjust unit tests to cover durable reconnect / short-circuit behavior and to keep legacy tests working by disabling durable where execute(None) is used.
• Document durable execution semantics and the durable=False opt-out.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 8 comments.

File	Description
providers/databricks/src/airflow/providers/databricks/operators/databricks.py	Implements durable synchronous execution via ResumableJobMixin and adds reconnect-aware polling/logging behavior.
providers/databricks/tests/unit/databricks/operators/test_databricks.py	Updates existing tests to pass durable=False where needed; adds a new durable-focused test suite.
providers/databricks/docs/operators/submit_run.rst	Documents durable execution behavior, version requirements, and opt-out configuration.

[amoghrajesh]

Thanks for your reviews, merging this.