amoghrajesh opened a new pull request, #77015:
URL: https://github.com/apache/airflow/pull/77015
<!-- 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:
https://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".
-->
- [ ] 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)
-->
### What
Durable execution already survives a worker crash. The aim here is to make
durable execution end-to-end: if the external job itself fails and has the
ability to recover from a stage, the retry resumes it instead of starting over.
Today, when a durable operator's external job fails, the task retry throws
that job away and starts a new one, so every step runs again. Many execution
engines can rerun a failed job from where it stopped, for example, a Databricks
repair or a dbt Cloud "rerun from failure". Durable operators already store the
job ID across retries, so a retry has everything it needs to use that.
Each operator also invents its own switch for this today (`repair_run`,
`retry_from_failure`, …), with different meanings. This PR adds one switch,
owned by `ResumableJobMixin`, that means the same thing for every operator.
### Current behaviour
On a retry, `ResumableJobMixin` reads the stored job id and checks the job's
status:
- If the job is running, the retry reconnects to it.
- If the job succeeded, the retry returns its result.
- If the job failed, the retry submits a fresh job.
Example: a Databricks job runs `clean`, `join` and `report`. `report` fails.
The retry starts a new run, so `clean` and `join` run again.
### Proposed change
- `ResumableJobMixin` gets an optional hook, `resume_failed_job(external_id,
status, context)`.
- When a retry finds the stored job failed, the mixin calls it.
- If it returns an ID, the mixin polls that ID instead of submitting
fresh. If the id is new, the mixin stores it.
- If it returns `None` or raises, the mixin submits fresh, as today.
- A new mixin parameter, `resume_on_retry: bool = True`, turns that call on
or off per task. `default_args` and `.partial()` work the same way they do for
`durable`.
- `DatabricksRunNowOperator` implements the hook. It repairs the stored run
in place and reruns only the failed tasks. It honours
`databricks_repair_reason_new_settings`, and skips the repair when the run no
longer exists.
What a retry does:
| `durable` | `resume_on_retry` | Stored job running | Stored job succeeded
| Stored job failed |
|---|---|---|---|---|
| `True` (default) | `True` (default) | reconnect | return result | **resume
it** (Databricks: repair the same run) |
| `True` | `False` | reconnect | return result | submit a fresh job |
| `False` | any | submit a fresh job | submit a fresh job | submit a fresh
job (no id is stored, so there is nothing to resume) |
For Databricks only:
| Setting | When it repairs |
|---|---|
| `repair_run=True` | once inside a try, while Airflow watches the run
(unchanged) |
| `resume_on_retry=True` | on the next retry, when the stored run already
failed |
Both can be on. Then a retry repairs the stored run, and if that repaired
run fails again, `repair_run` repairs it once more in the same try. Repairs
stay bounded by `retries`.
### Changes of Note
- **Kept separate from `durable` and from `repair_run`.**
- `durable` already means "reconnect after a crash", and users rely on it.
- `repair_run` repairs inside a try.
- Using either for this would change what an existing setting does.
- **The provider declares the parameter itself.**
- `DatabricksRunNowOperator` sets `self.resume_on_retry` after
`super().__init__()` instead of forwarding it. The mixin in Airflow 3.3 and 3.4
has no such parameter and would reject it.
- This keeps the parameter valid on every version, and `default_args`
reaches it on every version.
- It warns only when `resume_on_retry=True` is set explicitly below 3.5.
`False` and the default stay silent.
- **A failing resume never makes a retry worse.** If `resume_failed_job`
raises, the mixin logs a warning and submits fresh.
- **Observability:** a new `resume_failed` decision on the
`resumable_job.resume_decision` span, and a `resumable_job.resume_failed`
counter, next to the existing `terminal_resubmit`.
### User implications / backcompat
- **Airflow 3.5 with this Databricks provider:** durable
`DatabricksRunNowOperator` tasks repair a failed run on retry instead of
starting a new one. Set `resume_on_retry=False` to keep the old behaviour.
- **Airflow 3.3 and 3.4:** no change. The parameter is accepted and does
nothing. Setting it to `True` explicitly warns.
- **Clearing still starts fresh.** Clearing discards the stored job id, so
the next try submits a new job. This relies on #76806, which makes every clear
path discard task state, including clearing a whole Dag run. With
`keep_task_state`, a clear keeps the id, and the next try treats it like a
retry.
- **Operators without a `resume_failed_job` override:** no change. The
default returns `None`.
### Testing
**Unit tests**
- **Mixin:**
- resumes in place;
- resumes under a new id and stores it;
- submits fresh when the resume returns `None` or raises;
- doesn't call the resume when the job is running or succeeded;
- `resume_on_retry=False` submits fresh without calling the resume.
- **Databricks:**
- repairs the stored failed run with default flags, without `repair_run`;
- submits fresh when the failure doesn't match
`databricks_repair_reason_new_settings`;
- submits fresh with `resume_on_retry=False`;
- the `resume_on_retry` warning only fires for an explicit `True` below
3.5.
- **Runs:**
- `uv run --project task-sdk pytest
task-sdk/tests/task_sdk/bases/test_resumablejobmixin.py`: 50 passed.
- `uv run --project providers/databricks pytest
providers/databricks/tests`: 1312 passed, 12 skipped.
- `prek run mypy-task-sdk` and `breeze run mypy` on the Databricks files:
clean.
#### End to end against real Databricks and dbt Cloud
Airflow from this branch in Breeze. Each Dag sets `retries=1` and its job
fails on try 1 on purpose. The "before" Dag sets `resume_on_retry=False`, and
the "after" Dag uses the default.
Demo Dags:
```python
for resume in (False, True):
label = "after" if resume else "before"
with DAG(f"dbx_resume_{label}", schedule=None, default_args={"retries":
1, "retry_delay": timedelta(seconds=30)}):
DatabricksRunNowOperator(
task_id="run_sales_job",
databricks_conn_id="databricks_default",
job_id=<job_id>,
polling_period_seconds=10,
resume_on_retry=resume,
)
with DAG(f"dbt_resume_{label}", schedule=None, default_args={"retries":
1, "retry_delay": timedelta(seconds=90)}):
DbtCloudRunJobOperator(
task_id="run_dbt_job",
dbt_cloud_conn_id="dbt_cloud_default",
account_id=<account_id>,
job_id=<job_id>,
check_interval=10,
resume_on_retry=resume,
)
```
##### Databricks
- **Setup:**
- Job `repair_demo` runs `clean`, `join` and `report`.
- The `report` notebook keeps a counter and raises while it is 2 or less.
- Setting the counter to 1 makes `report` fail exactly once.
- **Steps:**
1. Set the counter to 1, then trigger `dbx_resume_before`.
2. Set the counter to 1, then trigger `dbx_resume_after`.
- **Before:** try 1 run `1047113470511919` failed on `report`. Try 2 logged
`Prior job in terminal state, resubmitting fresh` and started **new run
`197777498726462`**. All three tasks ran again, and try 2 took 2m 11s.
<img width="3200" height="2000" alt="airflow_dbx_resume_before_try1"
src="https://github.com/user-attachments/assets/fe1b2e64-75e5-488a-b643-970d1d2d2c55";
/>
<img width="3200" height="2000" alt="airflow_dbx_resume_before_try2"
src="https://github.com/user-attachments/assets/f08915f2-838d-4e22-b050-5b0dd29e00a8";
/>
<img width="3200" height="2000" alt="dbx_before_try1"
src="https://github.com/user-attachments/assets/11559620-7772-419a-9973-bda887da7a82";
/>
<img width="3200" height="2000" alt="dbx_before_try2"
src="https://github.com/user-attachments/assets/ddf0dca5-0d7d-4b0c-9fdb-8878417cb8fa";
/>
- **After:** try 1 run `221486698154948` failed on `report`. Try 2 logged
`Resuming failed job from where it stopped` and **repaired the same run**.
Databricks shows `clean` and `join` with one attempt each, and `report` with 2
attempts. Try 2 took 29.8s.
<img width="3200" height="2000" alt="airflow_dbx_resume_after_try1"
src="https://github.com/user-attachments/assets/1e1d6ecb-8b19-4b92-9380-9c433ae9370e";
/>
<img width="3200" height="2000" alt="airflow_dbx_resume_after_try2"
src="https://github.com/user-attachments/assets/3ff27011-20c9-4775-944a-b954ad39f2a0";
/>
<img width="3200" height="2000" alt="dbx_after_repaired"
src="https://github.com/user-attachments/assets/a0da9c57-89bc-460a-bca9-64ecaf4de041";
/>
<img width="3200" height="2000" alt="dbx_job_runs"
src="https://github.com/user-attachments/assets/08b0bf7e-3339-4a2a-a29d-162fa5f5b93b";
/>
##### dbt Cloud (a temporary `resume_failed_job` that calls dbt Cloud's
rerun endpoint, used only for this test and not part of this PR, dbt is being
made durable here: https://github.com/apache/airflow/pull/76918)
- **Setup:**
- The job runs `dbt run --select slow_model flaky_model`.
- `flaky_model` calls `raise_error()` while a one row Databricks table
`dbt_demo_flag` holds 1.
- **Steps, for each Dag:**
1. Set the flag to 1, then trigger the Dag.
2. When try 1 fails, set the flag to 0 during the 90s retry delay.
- **Before:** try 1 run `70506184200326` failed on `flaky_model`. Try 2
logged `Prior job in terminal state, resubmitting fresh` and triggered **new
run `70506184200349`**, which ran both models again (17s).
<img width="3200" height="2000" alt="airflow_dbt_resume_before_try1"
src="https://github.com/user-attachments/assets/4ccb6916-3cf1-49a1-910c-f5c9682e4f59";
/>
<img width="3200" height="2000" alt="airflow_dbt_resume_before_try2"
src="https://github.com/user-attachments/assets/fc23bc34-3029-4cd1-84d7-744889006fcc";
/>
<img width="3200" height="2000" alt="dbt_before_try1"
src="https://github.com/user-attachments/assets/b81fd01d-ed0a-4e23-9d32-ff230c5e7eb9";
/>
<img width="3200" height="2000" alt="dbt_before_try2"
src="https://github.com/user-attachments/assets/5809e944-4e4f-4da8-ad9b-1c6269b7906f";
/>
- **After:** try 1 run `70506184200356` failed. Try 2 logged `Resuming
failed job from where it stopped` and polled run `70506184200371`. dbt Cloud
shows "This run resumed execution from last failed step in run #70506184200356"
and ran `dbt retry` (5s).
<img width="3200" height="2000" alt="airflow_dbt_resume_after_try1"
src="https://github.com/user-attachments/assets/2b3395fb-cd90-4945-993e-b97f17ce0332";
/>
<img width="3200" height="2000" alt="airflow_dbt_resume_after_try2"
src="https://github.com/user-attachments/assets/24a24741-73c6-48cb-9f96-d57e1cfe8d1c";
/>
<img width="3200" height="2000" alt="dbt_after_try1"
src="https://github.com/user-attachments/assets/f50369a4-69fb-42f0-bb83-88e79f68efa3";
/>
<img width="3200" height="2000" alt="dbt_after_try2"
src="https://github.com/user-attachments/assets/cb3ae623-8e3d-43e0-84c4-b1a439f9086a";
/>
<img width="3200" height="2000" alt="dbt_job_history"
src="https://github.com/user-attachments/assets/14903ca9-5ebc-4165-af33-6cfabf99a9f6";
/>
### What's next
- dbt Cloud support.
- It needs the durable dbt operator PR first, worked on in:
https://github.com/apache/airflow/pull/76918
- dbt Cloud's rerun endpoint reruns the job's *latest* run and takes no
run id. So if anything else triggers that job between two tries, the retry
reruns the wrong run. That needs a guard before it ships.
- **More engines with a resume API:** Step Functions redrive, ADF rerun from
failure, and others.
- **A lost resume reply.** If `resume_failed_job` raises after the engine
accepted the request (for example, a timeout), the mixin submits fresh, and two
jobs run. I have a follow-up check that checks the job's status again before
falling back.
---
* 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]
