amoghrajesh opened a new pull request, #76806:
URL: https://github.com/apache/airflow/pull/76806
<!-- 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)
-->
closes: #72929
### What
#72100 made clearing a task instance discard its `task_state_store` entries,
so a cleared task starts over instead of resuming from a stale checkpoint. That
only covered `POST /dags/{dag_id}/clearTaskInstances`. Every other way of
clearing a task still kept task state, so clearing a task and clearing its Dag
run gave different results.
### Current behaviour
These paths call `clear_task_instances()` directly and never discard:
- Clearing a Dag run, single or bulk (UI "Clear Run" dialog and REST API)
- `airflow dags clear` and `airflow tasks clear`
- Marking a task as success or failed, which clears failed downstream tasks
- Backfills that reprocess existing runs,
`TriggerDagRunOperator(reset_dag_run=True)`, and partition clears with
`--clear-task-instances`
After a Dag run clear, a checkpointing task resumes and mixes old output
with new. A `ResumableJobMixin` task whose job already succeeded returns the
old result in seconds and runs nothing.
### Proposed change
- The discard moves into `clear_task_instances()`, behind a new
`keep_task_state: bool = False` argument. Every caller now gets the same
default, and the TI clear route passes its existing flag down instead of
discarding on its own.
- `clear_task_instances_for_runs`, the loop clear helpers,
`SerializedDAG.clear()` and `clear_dags()` pass `keep_task_state` through.
- Opt-outs:
- REST: `keep_task_state` on `DAGRunClearBody` and `BulkDAGRunClearBody`
- UI: "Keep task state and resume" checkbox in the single and bulk Clear
Run dialogs, seeded from the same per-browser default #72100 added
- CLI: `--keep-task-state` on `airflow dags clear` and `airflow tasks
clear`
- Marking tasks, backfill, `reset_dag_run` and partition clears always
discard, with no opt-out.
### Changes of Note
- The discard is now one `DELETE ... WHERE (dag_id, run_id, task_id,
region_id, region_index) IN (...)` per batch of 500 task instances. The old
helper ran one `DELETE` per task instance, which turns a large bulk Dag run
clear into thousands of statements.
- The discard runs at the end of `clear_task_instances()`, in the same
transaction, so a failed clear rolls back the discard too. Dry runs never reach
it.
- As before, the discard deletes metadata DB rows only. It leaves a custom
`[state_store] backend` alone.
### User implications / backcompat
Behaviour change, covered by a `significant` newsfragment. After a Dag run
clear or a core CLI clear, tasks start over. They resume from a checkpoint or
reconnect to a stored external job only when you set `keep_task_state`.
The one case where the default costs something is a `ResumableJobMixin` job
still running behind a failed task. Discarding the stored id submits a second
job, and the first keeps running. Tick "keep task state" for that case. The
docs from #72100 already explain when.
`clear_task_instances()` and `SerializedDAG.clear()` gain a keyword argument
with a default, so existing callers keep working. Provider code that calls
`clear_task_instances()` directly, such as the Databricks repair plugin, now
discards task state by default.
### Testing
Unit tests:
- `clear_task_instances()` default vs keep, plus a run that crosses a batch
boundary
- Dag run clear, single and bulk endpoints, default vs keep
- `--keep-task-state` on `airflow dags clear` and `airflow tasks clear`
- Both UI dialogs send `keep_task_state` only when the box is ticked. These
tests fail without the change.
Manual testing on Breeze (LocalExecutor, Postgres) with a Dag running three
task types side by side:
- `regular`: a `BashOperator` that keeps no state
- `checkpoint`: a `@task` that stores its progress in `task_state_store` and
logs whether it starts fresh or resumes
- `durable_job`: a `ResumableJobMixin` operator whose "external job" lives
in an Airflow Variable, so job states can be set by hand
Results for each clear path, default vs `keep_task_state`:
| Clear path | Default | Keep task state |
|---|---|---|
| REST Dag run clear | checkpoint starts fresh, durable job submits a new
job | checkpoint resumes, durable job returns the stored result |
| REST bulk Dag run clear (2 runs) | same as above for both runs | same as
above for both runs |
| UI "Clear Run" dialog | same as above | same as above |
| `airflow dags clear` | same as above | same as above |
| `airflow tasks clear` | same as above | same as above |
| `airflowctl tasks clear` (regression check on the TI endpoint) | new job |
stored result |
| Mark upstream success, failed downstream cleared | downstream starts fresh
| n/a |
| `airflow backfill create --reprocess-behavior completed` | starts fresh |
n/a |
| `TriggerDagRunOperator(reset_dag_run=True)` | target run starts fresh |
n/a |
The `regular` task reran the same way in every case.
`ResumableJobMixin` job states at clear time:
| Job state | Task state | Default | Keep task state |
|---|---|---|---|
| Succeeded | succeeded | submits a new job | "Job already completed
successfully, skipping resubmission" |
| Failed | failed | submits a new job | "Prior job in terminal state,
resubmitting fresh" |
| Running | failed (worker crashed) | submits a second job, the first stays
running | "Reconnecting to existing job" |
| Running | running (cleared live) | restarts and submits a second job |
restarts and reconnects to the same job |
Batched discard, from `airflow dags clear`:
```
[info] Discarded 2 task state store entries for 3 cleared task instances
[airflow.models.taskinstance]
```
Screenshots:
**Clear Run dialog, default (unticked)**
<img width="1500" height="950" alt="ui-clear-run-single-default"
src="https://github.com/user-attachments/assets/7e664d21-5000-4043-ad2f-5414f6c4aabc";
/>
**Clear Run dialog, "Keep task state and resume" ticked**
<img width="1500" height="950" alt="ui-clear-run-single-keep"
src="https://github.com/user-attachments/assets/a4639899-3b5d-40e6-a35a-77cdb038836a";
/>
**Bulk "Clear Dag Runs" dialog**
<img width="1500" height="950" alt="ui-bulk-clear-runs"
src="https://github.com/user-attachments/assets/59a23521-6b35-41f8-b723-1f0df54de457";
/>
**Task logs after a default Dag run clear (starts fresh, new job)**
<!-- screenshot -->
**Task logs after a Dag run clear with keep task state (resumes, stored
result)**
<!-- screenshot -->
### What's next
- Clearing a running task discards its state while the old attempt is still
alive, until the supervisor kills it. If that attempt writes to
`task_state_store` in that window, the next attempt resumes from the write.
#72100 has the same gap on the TI clear endpoint.
- Discard still only reaches the metadata DB. A custom `[state_store]
backend` keeps its state after a clear.
---
* 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]
