amoghrajesh opened a new pull request, #68974:
URL: https://github.com/apache/airflow/pull/68974
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull
request.
Write a good git commit message following this guide:
http://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the
documentation.
For user-facing UI changes, please attach before/after screenshots (or a
short
screen recording) so reviewers can assess the visual impact.
Feel free to ping (in general) for the review if you do not see reaction for
a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we
sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR,
please
change below checkbox to `[X]` followed by the name of the tool, uncomment
the "Generated-by".
-->
- [x] Yes - claude sonnet 4.6
<!--
Generated-by: [Tool Name] following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
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:
```python
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:
```python
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/[email protected]/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:
<img width="1726" height="940" alt="image"
src="https://github.com/user-attachments/assets/0a230cbe-9e00-42d8-8474-5e8b8eaba610";
/>
First run:
<img width="1726" height="940" alt="image"
src="https://github.com/user-attachments/assets/453a2585-f4e0-4849-9fea-7214f4c40bc4";
/>
Worker comes back up:
<img width="1726" height="940" alt="image"
src="https://github.com/user-attachments/assets/82fa05f1-8365-41a3-bc35-ccc087e94955";
/>
Extra run submitted now due to that:
<img width="1726" height="940" alt="image"
src="https://github.com/user-attachments/assets/66f19609-1af8-4ab0-8db6-71aabdd4f1fb";
/>
#### After my changes:
First run:
<img width="1726" height="940" alt="image"
src="https://github.com/user-attachments/assets/4b5a1377-f8c3-478d-8f47-1617bd6b1761";
/>
Worker comes back up:
<img width="1726" height="940" alt="image"
src="https://github.com/user-attachments/assets/f6f26366-cddb-4f84-b016-4d654b21ac8d";
/>
Just one job run:
<img width="1726" height="940" alt="image"
src="https://github.com/user-attachments/assets/e3812926-baff-42e2-8a31-58da11e3d055";
/>
Tried killing the job and it kills external job too:
<img width="1726" height="940" alt="image"
src="https://github.com/user-attachments/assets/173c8959-c9ba-43a8-995c-24b1b4760b3d";
/>
---
* Read the **[Pull Request
Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#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](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals))
is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License
Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment:
`{pr_number}.significant.rst`, in
[airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
You can add this file in a follow-up commit after the PR is created so you
know the PR number.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
