amoghrajesh opened a new pull request, #69174:
URL: https://github.com/apache/airflow/pull/69174
<!-- 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
-->
### Why we are doing this
`DatabricksRunNowOperator` triggers a run of an existing Databricks job,
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, a deploy, a spot reclaim), Airflow
retries the task in a fresh process with no memory of the `run_id`, so it calls
`run-now` again, triggering a second run of the same job. The original run
keeps executing on Databricks, orphaned, while the retry runs a duplicate.
This is the same crash-safety gap closed for `DatabricksSubmitRunOperator`
in #68974, applied to the run-now path. For long-running jobs it means paying
twice for the same work, and the only existing mitigation
(`cancel_previous_runs=True`) cancels rather than reconnects.
### Benefits this will bring in
- A worker crash and retry reconnects to the run already executing on
Databricks instead of triggering a duplicate, so you do not pay twice for the
same job.
- If the prior run already finished successfully, the retry returns
immediately without triggering a new run.
- If the stored run no longer exists because its history expired, the retry
degrades to a fresh run instead of failing the task.
- Works on the existing synchronous operator with no Triggerer required, and
is enabled by default 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 means reconnect and keep polling, already succeeded means return
immediately, and terminally failed (or a run whose history has expired) means
trigger a fresh run. Reconnect keys off the persisted `run_id`, so it works the
same whether the job was identified by `job_id` or `job_name`. Deferrable mode
is unchanged and takes precedence when enabled.
### Changes of note
- The run-state decoding helpers (`get_job_status` / `is_job_active` /
`is_job_succeeded`) are intentionally duplicated from
`DatabricksSubmitRunOperator` rather than extracted into a shared base, to keep
each operator self-contained and this PR scoped to RunNow. They are small and
each operator has its own tests covering them.
- The expired-run degradation reuses the typed `DatabricksApiError`
(carrying `http_status_code`) introduced alongside the SubmitRun work, so it is
a structured `404` check rather than message-string matching.
- The run-now payload build is split into a side-effect-free
`_build_run_now_payload()` (merge, validate, resolve `job_name` to `job_id`,
inject params) and the `cancel_previous_runs` step. On reconnect the payload is
rebuilt via the side-effect-free path, so `repair_run` has the resolved
`job_id` it needs and works across the reconnect boundary for both `job_id` and
`job_name`, without re-firing `cancel_previous_runs` against the run being
reconnected to.
### Backcompat
- Fully backward compatible. On a clean run the observable trigger-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 to exactly the old behavior (always triggers a
fresh run on retry) through a compatibility shim, so the provider keeps working
on older Airflow.
- If the task state store is unavailable at runtime, the operator logs that
crash recovery is disabled and behaves as before.
- One new optional flag is added; no existing parameters change and no
migration is needed.
### How to opt out
Set `durable=False` on the operator:
```python
DatabricksRunNowOperator(..., durable=False)
```
This restores the previous behavior: always trigger 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
DAG used:
```python
from datetime import timedelta
import pendulum
from airflow.providers.databricks.operators.databricks import
DatabricksRunNowOperator
from airflow.sdk import DAG
JOB_ID = 480676095338270
with DAG(
dag_id="databricks_runnow_repro",
schedule=None,
start_date=pendulum.datetime(2025, 1, 1, tz="UTC"),
catchup=False,
):
DatabricksRunNowOperator(
task_id="run_now_sleepy",
databricks_conn_id="databricks_default",
job_id=JOB_ID,
deferrable=False,
do_xcom_push=True,
retries=1,
retry_delay=timedelta(seconds=10),
)
```
#### Before my changes
Kill mid run:
<img width="1725" height="944" alt="image"
src="https://github.com/user-attachments/assets/6b26a26f-1a9a-461e-98d6-2332be260b30";
/>
Worker comes back up:
<img width="1725" height="944" alt="image"
src="https://github.com/user-attachments/assets/b045db3b-79ca-4ace-ba97-ab85eba2cc25";
/>
Duplicate submissions:
<img width="1725" height="944" alt="image"
src="https://github.com/user-attachments/assets/ef663058-48f3-46c2-ad1f-97952b56fbed";
/>
#### After my changes
Killed mid run:
<img width="1725" height="944" alt="image"
src="https://github.com/user-attachments/assets/bd287457-5359-4d25-bdfa-86ef397593b3";
/>
Job on dbx:
<img width="1725" height="944" alt="image"
src="https://github.com/user-attachments/assets/4f097f7b-10d7-4740-9b0c-aefe1365dfeb";
/>
Restored worker:
<img width="1725" height="944" alt="image"
src="https://github.com/user-attachments/assets/651fae44-0f07-4b8c-838d-40969d98c660";
/>
No duplicate job:
<img width="1725" height="944" alt="image"
src="https://github.com/user-attachments/assets/be8b0a3d-2265-4a2d-8823-73cb0f2d3602";
/>
Killed the job from airflow:
<img width="1725" height="944" alt="image"
src="https://github.com/user-attachments/assets/5f4e7493-e0ed-4fd7-b462-73ce66f286aa";
/>
<img width="1725" height="944" alt="image"
src="https://github.com/user-attachments/assets/07e67cdc-a4f1-4c2c-ac47-8aa5ad389577";
/>
---
##### 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".
-->
- [ ] Yes (please specify the tool below)
<!--
Generated-by: [Tool Name] following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
---
* 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]
