You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "amoghrajesh (via GitHub)" <gi...@apache.org> on 2026/05/06 09:15:56 UTC
[PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
amoghrajesh opened a new pull request, #66463:
URL: https://github.com/apache/airflow/pull/66463
<!-- 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.
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
-->
closes: https://github.com/apache/airflow/issues/66459
### What?
Task state rows right now accumulate indefinitely without a cleanup mechanism. The AIP specifies that the backend will be responsible for enforcing retention via a periodic job to avoid ever-growing table size.
Two things needed:
1. time-based garbage collection (delete rows older than N days) and
2. early expiry (per-key override for things like short-lived job IDs).
3. Asset state also needs orphan cleanup as mentioned in: https://github.com/apache/airflow/pull/66073#discussion_r3175794462 when an asset is removed from all DAGs its `asset_active` entry is deleted, but `asset_state` rows stay behind silently.
### Proposed change
- `expires_at` column on `task_state` table — `updated_at` alone can't distinguish a 7-day key from a 30-day key. A per-row column is the only way to honour per-key retention overrides. `NULL` means override isn't present and fall back to the global default; if its set means delete after this timestamp regardless of `updated_at`.
- `BaseStateBackend.cleanup(*, older_than: datetime)` no-op default — custom backends can override this for programmatic garbage collection. Airflow calls it from the scheduler.
- New config options under `[state_store]`: `default_retention_days = 30` and `state_cleanup_interval = 86400`.
- `MetastoreStateBackend.cleanup()` runs three DELETEs in one pass: rows past `updated_at` cutoff, rows with `expires_at < now()`, and `asset_state` rows whose asset has no `asset_active` entry.
- `_cleanup_expired_task_state()` hooked into the scheduler via `call_regular_interval(state_cleanup_interval, ...)`, consistent with how the scheduler handles `_remove_unreferenced_triggers` and `_update_asset_orphanage`. Calls `get_state_backend().cleanup()` — respects the configured backend, works for both default and custom backends.
### Why choose scheduler and not `airflow db cleanup`?
- `airflow db cleanup` uses a single `--days-ago` flag for all tables, so operators would have to manually align it with `default_retention_days`.
- It also can't handle `expires_at < now()` since that is an absolute cutoff threshold, not relative to `--days-ago`.
- A dedicated CLI command was also considered but would require operators to schedule it separately. The scheduler already runs this kind of periodic cleanup for triggers and assets, so hooking in there keeps it automatic with no operator action needed.
**Does scheduler have any performance implication? (cc: @ashb)**
` _cleanup_expired_task_state()` runs every 24 hrs by default — same cadence as similar maintenance tasks. When it fires it opens its own session and runs three DELETEs. It can add some overhead but should be manageable since queries are simple?
The only real concern: if someone sets `state_cleanup_interval` very low and has a large task_state table, the repeated DELETEs could put load on the DB.
### User implications / backcompat
New config options under `[state_store]` with safe defaults — no action needed to maintain existing behaviour. The `expires_at` column is nullable; existing rows get `NULL` (global default retention applies).
### Testing
#### Test setup
1. Started breeze with these configurations:
```shell script
export AIRFLOW__STATE_STORE__STATE_CLEANUP_INTERVAL=300
export AIRFLOW__STATE_STORE__DEFAULT_RETENTION_DAYS=1
```
2. Created a dag run and pushed in 3 task states for it:
<img width="1699" height="890" alt="image (76)" src="https://github.com/user-attachments/assets/2e328b8f-4288-4571-8c8f-71a28e197110" />
#### Testing for `updated_at` / global cleanup
1. Ran this query:
```sql
UPDATE task_state
SET updated_at = '2026-05-04 00:00:00+00:00'
WHERE key = 'job_id_2';
```
2. Waited for the scheduler to run the GC job and rows to be deleted.
```shell script
2026-05-06T08:40:57.899773Z [info ] Adopting or resetting orphaned tasks for active dag runs [airflow.jobs.scheduler_job_runner.SchedulerJobRunner] loc=scheduler_job_runner.py:2786
2026-05-06T08:40:57.904527Z [info ] Running task state cleanup [airflow.jobs.scheduler_job_runner.SchedulerJobRunner] loc=scheduler_job_runner.py:3066 older_than=datetime.datetime(2026, 5, 5, 8, 40, 57, 904487, tzinfo=Timezone('UTC')) retention_days=1
```
<img width="1699" height="890" alt="image (77)" src="https://github.com/user-attachments/assets/5bfa9d7d-b0c0-4cc6-954d-e82c0f62506a" />
#### Testing for `expired_at`
1. Ran this query:
```sql
UPDATE task_state
SET expires_at = '2026-05-06 08:45:00+00:00'
WHERE key = 'job_id';
```
<img width="1699" height="890" alt="image (78)" src="https://github.com/user-attachments/assets/2559abd2-055e-4374-9ce1-5ded1b85d683" />
2. Waited for the scheduler to run the GC job and rows to be deleted (note here that the row with `expired_at` is cleared even though `updated_at` is still not reaching its cleanup interval)
<img width="1699" height="890" alt="image (79)" src="https://github.com/user-attachments/assets/85775ccb-5674-4dde-89a6-c474c6070846" />
### What's next
- `clear_on_success` hook: https://github.com/apache/airflow/issues/66460
- `task_state.set(retention_days=N)` API to populate `expires_at` at write time: https://github.com/apache/airflow/issues/66461
---
##### 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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3201644389
##########
airflow-core/src/airflow/cli/cli_config.py:
##########
@@ -1519,6 +1519,19 @@ class GroupCommand(NamedTuple):
args=(ARG_OUTPUT, ARG_VERBOSE),
),
)
+STATE_STORE_COMMANDS = (
+ ActionCommand(
+ name="cleanup",
+ help="Remove expired task state rows via the configured state backend",
Review Comment:
Nit:
```suggestion
help="Remove expired stored state via the configured state backend",
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3217755989
##########
airflow-core/tests/unit/state/test_metastore.py:
##########
@@ -234,6 +239,112 @@ def test_clear_with_all_map_indices_flag_wipes_wide(
assert backend.get(scope0, "job_id", session=session) is None
assert backend.get(scope1, "job_id", session=session) is None
+ def test_set_populates_expires_at(
+ self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+ ):
+ """set() always populates expires_at so cleanup has a single pass."""
+ scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+ backend.set(scope, "job_id", "app_1234", session=session)
+ session.flush()
+
+ row = session.scalar(select(TaskStateModel).where(TaskStateModel.key == "job_id"))
+ assert row is not None
Review Comment:
Done. Added assert `row.expires_at > row.updated_at` to actually verify
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3225541284
##########
airflow-core/src/airflow/cli/commands/state_store_command.py:
##########
@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+
+from airflow.state import get_state_backend
+from airflow.state.metastore import MetastoreStateBackend
+
+log = logging.getLogger(__name__)
+
+# Other state operations (list, get, delete per key) will be added here in the future.
+
+
+def cleanup(args) -> None:
+ """Remove expired task state rows via the configured state backend."""
+ backend = get_state_backend()
+
+ if args.dry_run:
+ if isinstance(backend, MetastoreStateBackend):
+ summary = backend._summary_dry_run_()
+ expired = summary["expired"]
+ if not expired:
+ print("Nothing to delete.")
+ return
+ print(f"Would delete {len(expired)} task state row(s):\n")
+ for dag_id, run_id, task_id, map_index, key in expired:
+ print(
+ f" Dag {dag_id!r}, run {run_id!r}, task {task_id!r}, map_index {map_index!r}, key {key!r}"
+ )
+ else:
+ print("Custom backend configured — cannot preview rows.")
Review Comment:
Not really, no, this command makes sense for MetadataStateBackend only. The CLI is mostly used for that cleanup only
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3217281033
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -33,7 +33,22 @@
from itertools import groupby
from typing import TYPE_CHECKING, Any, cast
-from sqlalchemy import CTE, and_, case, delete, exists, func, inspect, or_, select, text, tuple_, update
+from sqlalchemy import (
+ CTE,
+ and_,
+ case,
+ delete,
+ delete as _delete,
Review Comment:
Oh absolutely no reason, reverting those aliases.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3217106108
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +272,51 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows.
+
+ ``expires_at`` is set at write time on every ``set()`` call, so cleanup is a single
+ ``WHERE expires_at < now()`` pass. Rows with ``expires_at=NULL`` (default_retention_days=0)
+ are never deleted. Batching is configurable via ``[state_store] state_cleanup_batch_size``.
+ """
+ batch_size = conf.getint("state_store", "state_cleanup_batch_size")
+ now = timezone.utcnow()
+
+ def _delete_batched(where_clause) -> int:
+ total = 0
+ with create_session() as session:
+ while True:
+ id_query = select(TaskStateModel.id).where(where_clause)
+ if batch_size > 0:
+ id_query = id_query.limit(batch_size)
+ ids = session.scalars(id_query).all()
+ if not ids:
+ break
+ session.execute(delete(TaskStateModel).where(TaskStateModel.id.in_(ids)))
+ session.commit()
+ total += len(ids)
+ if batch_size <= 0 or len(ids) < batch_size:
+ break
+ return total
+
+ deleted = _delete_batched(TaskStateModel.expires_at < now)
Review Comment:
I wonder if it’d be a good idea if the actual expiration is calculated on the fly instead. If I’m understaing correctly, this currently relies on the `expires_at` column being correctly updated whenever `update_at` is updated (if the former is not set explicitly). This seems a bit fragile.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3217316017
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +272,51 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows.
+
+ ``expires_at`` is set at write time on every ``set()`` call, so cleanup is a single
+ ``WHERE expires_at < now()`` pass. Rows with ``expires_at=NULL`` (default_retention_days=0)
+ are never deleted. Batching is configurable via ``[state_store] state_cleanup_batch_size``.
+ """
+ batch_size = conf.getint("state_store", "state_cleanup_batch_size")
+ now = timezone.utcnow()
+
+ def _delete_batched(where_clause) -> int:
+ total = 0
+ with create_session() as session:
+ while True:
+ id_query = select(TaskStateModel.id).where(where_clause)
+ if batch_size > 0:
+ id_query = id_query.limit(batch_size)
+ ids = session.scalars(id_query).all()
+ if not ids:
+ break
+ session.execute(delete(TaskStateModel).where(TaskStateModel.id.in_(ids)))
+ session.commit()
+ total += len(ids)
+ if batch_size <= 0 or len(ids) < batch_size:
+ break
+ return total
+
+ deleted = _delete_batched(TaskStateModel.expires_at < now)
Review Comment:
`expires_at` is set once at write time in every `set()` call and is never updated independently of the row — there's no dependency on `updated_at` being in sync with it. If you call `set()` again on the same key, the upsert recalculates and overwrites both `updated_at` and `expires_at` together atomically.
One legitimate edge case you may be pointing at: if a user starts with `default_retention_days=0, then later raises it to 30 days, those old NULL rows won't be picked up by the current WHERE `expires_at < now()` pass. We can add a second pass WHERE `expires_at` IS NULL AND `updated_at < now - default_retention_days` for that case. How does that sound?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3217820469
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +272,51 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows.
+
+ ``expires_at`` is set at write time on every ``set()`` call, so cleanup is a single
+ ``WHERE expires_at < now()`` pass. Rows with ``expires_at=NULL`` (default_retention_days=0)
+ are never deleted. Batching is configurable via ``[state_store] state_cleanup_batch_size``.
+ """
+ batch_size = conf.getint("state_store", "state_cleanup_batch_size")
+ now = timezone.utcnow()
+
+ def _delete_batched(where_clause) -> int:
+ total = 0
+ with create_session() as session:
+ while True:
+ id_query = select(TaskStateModel.id).where(where_clause)
+ if batch_size > 0:
+ id_query = id_query.limit(batch_size)
+ ids = session.scalars(id_query).all()
+ if not ids:
+ break
+ session.execute(delete(TaskStateModel).where(TaskStateModel.id.in_(ids)))
+ session.commit()
+ total += len(ids)
+ if batch_size <= 0 or len(ids) < batch_size:
+ break
+ return total
+
+ deleted = _delete_batched(TaskStateModel.expires_at < now)
+ log.info("Deleted expired task_state rows", rows_deleted=deleted)
+
+ def _dry_run_summary(self) -> dict[str, list]:
Review Comment:
Fine by me!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3198954290
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +253,41 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows and orphaned asset state rows.
+
+ Reads ``[state_store] default_retention_days`` from config to determine the time-based
+ cutoff. Set to 0 to disable time-based cleanup (expires_at and orphan cleanup still run).
+
+ Two passes for task_state:
+ a. Rows where updated_at < now() - default_retention_days (global retention)
+ b. Rows where expires_at < now() (per-key early expiry set by the operator)
+
+ Asset state orphan cleanup: asset_state rows whose FK'd asset has no asset_active
+ entry are stale (asset removed from all DAGs) and are deleted here.
+ """
+ from datetime import timedelta
+
+ from airflow.configuration import conf
+
+ retention_days = conf.getint("state_store", "default_retention_days")
+ now = timezone.utcnow()
+ older_than = now - timedelta(days=retention_days) if retention_days > 0 else None
+ with create_session() as session:
+ if older_than:
+ session.execute(delete(TaskStateModel).where(TaskStateModel.updated_at < older_than))
+ session.execute(
+ delete(TaskStateModel).where(
+ TaskStateModel.expires_at.isnot(None),
+ TaskStateModel.expires_at < now,
+ )
+ )
+ active_asset_ids = select(AssetModel.id).join(
+ AssetActive, (AssetActive.name == AssetModel.name) & (AssetActive.uri == AssetModel.uri)
+ )
+ session.execute(delete(AssetStateModel).where(AssetStateModel.asset_id.not_in(active_asset_ids)))
+
Review Comment:
If it's a valid assumption that users might produce a large amount of state records here between the `state_cleanup_interval` time window. I have some concern regarding the single pass delete Transaction here.
I just double checked the concern with Claude:
---
Yes, this is a real problem. Several compounding issues:
1. Missing indexes — every cleanup will be a full table scan
The two predicates the cleanup filters on are not indexed:
- task_state.updated_at — no index (only task_state_pkey on (dag_run_id,
task_id, map_index, key) and idx_task_state_lookup on (dag_id, run_id,
task_id, map_index))
- task_state.expires_at — no index (just added in this PR)
So both DELETE WHERE updated_at < cutoff and DELETE WHERE expires_at < now()
do full sequential scans. On a deployment with millions of rows that's minutes
of scanning every 24h, plus the locks held for the whole duration.
2. No batching / no LIMIT
Compare to airflow db cleanup (utils/db_cleanup.py:217), which deletes in
configurable batches and commits between them. The new path runs three plain
bulk DELETEs in a single session. Long-running bulk DELETE means:
- Row locks held for the duration (writers calling task_state.set() upserts on
matching rows block — they queue behind the cleanup transaction).
- On Postgres: massive WAL churn, autovacuum can't keep up, table bloat.
- On MySQL/InnoDB at REPEATABLE READ (Airflow's default): next-key/gap locks
make conflicts even more likely.
3. All three DELETEs share one transaction
with create_session() as session: opens one session; each session.execute()
runs inside it; commit happens at exit. If pass 1 takes 90s, the locks from
pass 1 are held while pass 2 and pass 3 run. A failure in pass 3 rolls back
passes 1 and 2 (cleanup makes no forward progress at all).
4. Scheduler main loop is blocked
_cleanup_expired_task_state is registered via call_regular_interval, which is
synchronous in the scheduler loop. Same pattern as
_remove_unreferenced_triggers and _update_asset_orphanage — but those have
small cardinality. task_state is user-driven and unbounded (the AIP encourages
users to write a lot of it). With a multi-minute cleanup the scheduler is not
scheduling for those minutes.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3231894237
##########
airflow-core/tests/unit/state/test_metastore.py:
##########
Review Comment:
Added `test_custom_backend_is_skipped` which mocks a StateBackend, asserts the "Custom backend configured" message is printed, and verifies no cleanup() is called.
##########
airflow-core/src/airflow/cli/commands/state_store_command.py:
##########
@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+
+from airflow.state import get_state_backend
+from airflow.state.metastore import MetastoreStateBackend
+
+log = logging.getLogger(__name__)
+
+# Other state operations (list, get, delete per key) will be added here in the future.
+
+
+def cleanup(args) -> None:
+ """Remove expired task state rows via the configured state backend."""
+ backend = get_state_backend()
+
+ if args.dry_run:
+ if isinstance(backend, MetastoreStateBackend):
+ summary = backend._summary_dry_run_()
+ expired = summary["expired"]
+ if not expired:
+ print("Nothing to delete.")
+ return
+ print(f"Would delete {len(expired)} task state row(s):\n")
+ for dag_id, run_id, task_id, map_index, key in expired:
+ print(
+ f" Dag {dag_id!r}, run {run_id!r}, task {task_id!r}, map_index {map_index!r}, key {key!r}"
+ )
+ else:
+ print("Custom backend configured — cannot preview rows.")
Review Comment:
Cool, what do you think about this: I renamed the command to `cleanup-task-states` and updated the description to explicitly say "Only applies when MetastoreStateBackend is configured; custom backends are skipped." Also moved the backend check to a top-level early return so the limitation is enforced before any dry-run logic.
##########
airflow-core/src/airflow/cli/cli_config.py:
##########
@@ -2108,6 +2115,11 @@ class GroupCommand(NamedTuple):
help="Display providers",
subcommands=PROVIDERS_COMMANDS,
),
+ GroupCommand(
+ name="state-store",
+ help="Manage task and asset state storage",
Review Comment:
Renamed the command to `cleanup-task-states` so it's scoped correctly from the start. The group help is now accurate since we only have task state management. Asset state retention can be added as a separate `cleanup-asset-states` command when we get there.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3233259812
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +272,51 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows.
+
+ ``expires_at`` is set at write time on every ``set()`` call, so cleanup is a single
+ ``WHERE expires_at < now()`` pass. Rows with ``expires_at=NULL`` (default_retention_days=0)
+ are never deleted. Batching is configurable via ``[state_store] state_cleanup_batch_size``.
+ """
+ batch_size = conf.getint("state_store", "state_cleanup_batch_size")
+ now = timezone.utcnow()
+
+ def _delete_batched(where_clause) -> int:
+ total = 0
+ with create_session() as session:
+ while True:
+ id_query = select(TaskStateModel.id).where(where_clause)
+ if batch_size > 0:
+ id_query = id_query.limit(batch_size)
+ ids = session.scalars(id_query).all()
+ if not ids:
+ break
+ session.execute(delete(TaskStateModel).where(TaskStateModel.id.in_(ids)))
+ session.commit()
+ total += len(ids)
+ if batch_size <= 0 or len(ids) < batch_size:
+ break
+ return total
+
+ deleted = _delete_batched(TaskStateModel.expires_at < now)
Review Comment:
@uranusjr now that I think of it, they can also delete it by calling a `task_state.clear(all_map_indices=True)`, so its possible in case they want it
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3195564690
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -3025,6 +3014,29 @@ state_store:
type: string
example: "mypackage.state.CustomStateBackend"
default: "airflow.state.metastore.MetastoreStateBackend"
+ default_retention_days:
+ description: |
+ Number of days to retain task state rows after their last update.
Review Comment:
will it also affect asset state?
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3047,6 +3054,19 @@ def _remove_unreferenced_triggers(self, *, session: Session = NEW_SESSION) -> No
.execution_options(synchronize_session="fetch")
)
+ def _cleanup_expired_task_state(self) -> None:
+ """
+ Remove expired task state rows and orphaned asset state rows.
+
+ Reads ``[state_store] default_retention_days`` from config and delegates to the
+ configured state backend. Runs on the interval set by ``[state_store] state_cleanup_interval``.
+ """
+ retention_days = conf.getint("state_store", "default_retention_days")
Review Comment:
if `default_retention_days` is set to 0, then we'll delete everything older than now, it doesn't look right 🤔
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3217711609
##########
airflow-core/src/airflow/cli/commands/state_store_command.py:
##########
@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+
+log = logging.getLogger(__name__)
+
+# Other state operations (list, get, delete per key) will be added here in the future.
+
+
+def cleanup(args) -> None:
+ """Remove expired task state rows via the configured state backend."""
+ from airflow.state import get_state_backend
+ from airflow.state.metastore import MetastoreStateBackend
Review Comment:
No reason for it, good catch, handled
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3234087336
##########
shared/state/src/airflow_shared/state/__init__.py:
##########
@@ -122,3 +122,14 @@ async def aclear(self, scope: StateScope, *, all_map_indices: bool = False) -> N
scope are cleared. Pass ``all_map_indices=True`` to wipe state across every
mapped instance of the task. For ``AssetScope`` the flag has no effect.
"""
+
+ def cleanup(self) -> None:
+ """
+ Remove expired and orphaned state records.
+
+ This is a no-op by default. Custom backends override this to implement their own
+ retention policy. The backend is responsible for reading any relevant config (e.g.
+ ``[state_store] default_retention_days``) and deciding what to delete.
+ Airflow does not call this from any standard job — the scheduler triggers it via
+ ``call_regular_interval`` for the default backend.
Review Comment:
Yeah, earlier it was added to scheduler, but now we do not do that and have a CLI command specifically because scheduler running in its main loop can have some performance implications, which is better to avoid. Updated this in 77450e580c
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -3025,6 +3025,24 @@ state_store:
type: string
example: "mypackage.state.CustomStateBackend"
default: "airflow.state.metastore.MetastoreStateBackend"
+ default_retention_days:
+ description: |
+ Number of days to retain task state after their last update.
+ Rows older than this are removed by the scheduler's periodic cleanup.
Review Comment:
Yeah, i have pushed a fix for it: 77450e580c
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3231470966
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +272,51 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows.
+
+ ``expires_at`` is set at write time on every ``set()`` call, so cleanup is a single
+ ``WHERE expires_at < now()`` pass. Rows with ``expires_at=NULL`` (default_retention_days=0)
+ are never deleted. Batching is configurable via ``[state_store] state_cleanup_batch_size``.
+ """
+ batch_size = conf.getint("state_store", "state_cleanup_batch_size")
+ now = timezone.utcnow()
+
+ def _delete_batched(where_clause) -> int:
+ total = 0
+ with create_session() as session:
+ while True:
+ id_query = select(TaskStateModel.id).where(where_clause)
+ if batch_size > 0:
+ id_query = id_query.limit(batch_size)
+ ids = session.scalars(id_query).all()
+ if not ids:
+ break
+ session.execute(delete(TaskStateModel).where(TaskStateModel.id.in_(ids)))
+ session.commit()
+ total += len(ids)
+ if batch_size <= 0 or len(ids) < batch_size:
+ break
+ return total
+
+ deleted = _delete_batched(TaskStateModel.expires_at < now)
Review Comment:
I wonder if we should add a command to change retention of existing tis; I feel some would need it, either because they didn’t know about the feature previously or want to change policy entirely. Or maybe those people can just manually delete the tis in another way anyway?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3217777118
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +272,51 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows.
+
+ ``expires_at`` is set at write time on every ``set()`` call, so cleanup is a single
+ ``WHERE expires_at < now()`` pass. Rows with ``expires_at=NULL`` (default_retention_days=0)
+ are never deleted. Batching is configurable via ``[state_store] state_cleanup_batch_size``.
+ """
+ batch_size = conf.getint("state_store", "state_cleanup_batch_size")
+ now = timezone.utcnow()
+
+ def _delete_batched(where_clause) -> int:
+ total = 0
+ with create_session() as session:
+ while True:
+ id_query = select(TaskStateModel.id).where(where_clause)
+ if batch_size > 0:
+ id_query = id_query.limit(batch_size)
+ ids = session.scalars(id_query).all()
+ if not ids:
+ break
+ session.execute(delete(TaskStateModel).where(TaskStateModel.id.in_(ids)))
+ session.commit()
+ total += len(ids)
+ if batch_size <= 0 or len(ids) < batch_size:
+ break
+ return total
+
+ deleted = _delete_batched(TaskStateModel.expires_at < now)
+ log.info("Deleted expired task_state rows", rows_deleted=deleted)
+
+ def _dry_run_summary(self) -> dict[str, list]:
Review Comment:
`_dry_run_summary` looks like a property for me. But based on the DB interaction we're doing here, it's not a good idea to make it a property IMO. Thus, rename it to `_summary_dry_run` makes it more function like
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3201723567
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +259,84 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows.
+
+ Reads ``[state_store] default_retention_days`` and ``[state_store] state_cleanup_batch_size``
+ from config. Each pass runs in its own transaction so partial progress is committed even if a
+ later pass fails. Each pass is batched to avoid long-running locks on the table.
+
+ Two passes:
+ a. Rows where updated_at < now() - default_retention_days (global retention)
+ b. Rows where expires_at < now() (per-key early expiry set by the operator)
+ """
+ retention_days = conf.getint("state_store", "default_retention_days")
+ batch_size = conf.getint("state_store", "state_cleanup_batch_size")
+ now = timezone.utcnow()
+ older_than = now - timedelta(days=retention_days) if retention_days > 0 else None
+
+ pk_cols = (
+ TaskStateModel.dag_run_id,
+ TaskStateModel.task_id,
+ TaskStateModel.map_index,
+ TaskStateModel.key,
+ )
+
+ def _delete_batched(where_clause) -> int:
+ total = 0
+ while True:
+ with create_session() as session:
Review Comment:
I don't think this should be a new session object each time around the loop, but instead one session object that is explicitly `session.commit()`ed after each batch.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3201716728
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +259,84 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows.
+
+ Reads ``[state_store] default_retention_days`` and ``[state_store] state_cleanup_batch_size``
+ from config. Each pass runs in its own transaction so partial progress is committed even if a
+ later pass fails. Each pass is batched to avoid long-running locks on the table.
+
+ Two passes:
+ a. Rows where updated_at < now() - default_retention_days (global retention)
+ b. Rows where expires_at < now() (per-key early expiry set by the operator)
+ """
+ retention_days = conf.getint("state_store", "default_retention_days")
+ batch_size = conf.getint("state_store", "state_cleanup_batch_size")
+ now = timezone.utcnow()
+ older_than = now - timedelta(days=retention_days) if retention_days > 0 else None
+
+ pk_cols = (
+ TaskStateModel.dag_run_id,
+ TaskStateModel.task_id,
+ TaskStateModel.map_index,
+ TaskStateModel.key,
+ )
Review Comment:
Given how this is used, it might be time to add a single column `id` pk (either integer, or uuid)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3232978829
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +272,51 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows.
+
+ ``expires_at`` is set at write time on every ``set()`` call, so cleanup is a single
+ ``WHERE expires_at < now()`` pass. Rows with ``expires_at=NULL`` (default_retention_days=0)
+ are never deleted. Batching is configurable via ``[state_store] state_cleanup_batch_size``.
+ """
+ batch_size = conf.getint("state_store", "state_cleanup_batch_size")
+ now = timezone.utcnow()
+
+ def _delete_batched(where_clause) -> int:
+ total = 0
+ with create_session() as session:
+ while True:
+ id_query = select(TaskStateModel.id).where(where_clause)
+ if batch_size > 0:
+ id_query = id_query.limit(batch_size)
+ ids = session.scalars(id_query).all()
+ if not ids:
+ break
+ session.execute(delete(TaskStateModel).where(TaskStateModel.id.in_(ids)))
+ session.commit()
+ total += len(ids)
+ if batch_size <= 0 or len(ids) < batch_size:
+ break
+ return total
+
+ deleted = _delete_batched(TaskStateModel.expires_at < now)
+ log.info("Deleted expired task_state rows", rows_deleted=deleted)
+
+ def _summary_dry_run_(self) -> dict[str, list]:
Review Comment:
```suggestion
def _summary_dry_run(self) -> dict[str, list]:
```
typo?
##########
shared/state/src/airflow_shared/state/__init__.py:
##########
@@ -122,3 +122,14 @@ async def aclear(self, scope: StateScope, *, all_map_indices: bool = False) -> N
scope are cleared. Pass ``all_map_indices=True`` to wipe state across every
mapped instance of the task. For ``AssetScope`` the flag has no effect.
"""
+
+ def cleanup(self) -> None:
+ """
+ Remove expired and orphaned state records.
+
+ This is a no-op by default. Custom backends override this to implement their own
+ retention policy. The backend is responsible for reading any relevant config (e.g.
+ ``[state_store] default_retention_days``) and deciding what to delete.
+ Airflow does not call this from any standard job — the scheduler triggers it via
+ ``call_regular_interval`` for the default backend.
Review Comment:
do we still have this
##########
shared/state/src/airflow_shared/state/__init__.py:
##########
@@ -122,3 +122,14 @@ async def aclear(self, scope: StateScope, *, all_map_indices: bool = False) -> N
scope are cleared. Pass ``all_map_indices=True`` to wipe state across every
mapped instance of the task. For ``AssetScope`` the flag has no effect.
"""
+
+ def cleanup(self) -> None:
+ """
+ Remove expired and orphaned state records.
+
+ This is a no-op by default. Custom backends override this to implement their own
+ retention policy. The backend is responsible for reading any relevant config (e.g.
+ ``[state_store] default_retention_days``) and deciding what to delete.
+ Airflow does not call this from any standard job — the scheduler triggers it via
+ ``call_regular_interval`` for the default backend.
Review Comment:
i remember we have something in the scheduler in the previous review. did we remove it? or is it in another PR?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3231520604
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +272,51 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows.
+
+ ``expires_at`` is set at write time on every ``set()`` call, so cleanup is a single
+ ``WHERE expires_at < now()`` pass. Rows with ``expires_at=NULL`` (default_retention_days=0)
+ are never deleted. Batching is configurable via ``[state_store] state_cleanup_batch_size``.
+ """
+ batch_size = conf.getint("state_store", "state_cleanup_batch_size")
+ now = timezone.utcnow()
+
+ def _delete_batched(where_clause) -> int:
+ total = 0
+ with create_session() as session:
+ while True:
+ id_query = select(TaskStateModel.id).where(where_clause)
+ if batch_size > 0:
+ id_query = id_query.limit(batch_size)
+ ids = session.scalars(id_query).all()
+ if not ids:
+ break
+ session.execute(delete(TaskStateModel).where(TaskStateModel.id.in_(ids)))
+ session.commit()
+ total += len(ids)
+ if batch_size <= 0 or len(ids) < batch_size:
+ break
+ return total
+
+ deleted = _delete_batched(TaskStateModel.expires_at < now)
Review Comment:
TBH, i would hold off for now — people can run a direct sql statements if needed, and any new key(s) that gets written will automatically pick up the new `default_retention_days` configured. If there is a clear demand for it we can add something like `airflow state-store set-expiry` when needed, but feels premature before we know how common that use case is.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3233259812
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +272,51 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows.
+
+ ``expires_at`` is set at write time on every ``set()`` call, so cleanup is a single
+ ``WHERE expires_at < now()`` pass. Rows with ``expires_at=NULL`` (default_retention_days=0)
+ are never deleted. Batching is configurable via ``[state_store] state_cleanup_batch_size``.
+ """
+ batch_size = conf.getint("state_store", "state_cleanup_batch_size")
+ now = timezone.utcnow()
+
+ def _delete_batched(where_clause) -> int:
+ total = 0
+ with create_session() as session:
+ while True:
+ id_query = select(TaskStateModel.id).where(where_clause)
+ if batch_size > 0:
+ id_query = id_query.limit(batch_size)
+ ids = session.scalars(id_query).all()
+ if not ids:
+ break
+ session.execute(delete(TaskStateModel).where(TaskStateModel.id.in_(ids)))
+ session.commit()
+ total += len(ids)
+ if batch_size <= 0 or len(ids) < batch_size:
+ break
+ return total
+
+ deleted = _delete_batched(TaskStateModel.expires_at < now)
Review Comment:
@uranusjr now that I think of it, they can also delete it by calling a `context["task_state"].clear(all_map_indices=True)`, so its possible in case they want it
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh merged PR #66463:
URL: https://github.com/apache/airflow/pull/66463
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3201551462
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +253,41 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows and orphaned asset state rows.
+
+ Reads ``[state_store] default_retention_days`` from config to determine the time-based
+ cutoff. Set to 0 to disable time-based cleanup (expires_at and orphan cleanup still run).
+
+ Two passes for task_state:
+ a. Rows where updated_at < now() - default_retention_days (global retention)
+ b. Rows where expires_at < now() (per-key early expiry set by the operator)
+
+ Asset state orphan cleanup: asset_state rows whose FK'd asset has no asset_active
+ entry are stale (asset removed from all DAGs) and are deleted here.
+ """
+ from datetime import timedelta
+
+ from airflow.configuration import conf
+
+ retention_days = conf.getint("state_store", "default_retention_days")
+ now = timezone.utcnow()
+ older_than = now - timedelta(days=retention_days) if retention_days > 0 else None
+ with create_session() as session:
+ if older_than:
+ session.execute(delete(TaskStateModel).where(TaskStateModel.updated_at < older_than))
+ session.execute(
+ delete(TaskStateModel).where(
+ TaskStateModel.expires_at.isnot(None),
+ TaskStateModel.expires_at < now,
+ )
+ )
+ active_asset_ids = select(AssetModel.id).join(
+ AssetActive, (AssetActive.name == AssetModel.name) & (AssetActive.uri == AssetModel.uri)
+ )
+ session.execute(delete(AssetStateModel).where(AssetStateModel.asset_id.not_in(active_asset_ids)))
+
Review Comment:
Handled it in: cdc423703e
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3206773888
##########
airflow-core/src/airflow/cli/cli_config.py:
##########
@@ -1519,6 +1519,19 @@ class GroupCommand(NamedTuple):
args=(ARG_OUTPUT, ARG_VERBOSE),
),
)
+STATE_STORE_COMMANDS = (
+ ActionCommand(
+ name="cleanup",
+ help="Remove expired task state rows via the configured state backend",
Review Comment:
Handled in df379c50e6
##########
airflow-core/src/airflow/migrations/versions/0112_3_3_0_add_task_state_and_asset_state_tables.py:
##########
@@ -65,6 +65,11 @@ def upgrade():
sa.Column("run_id", StringID(), nullable=False),
sa.Column("value", sa.Text().with_variant(mysql.MEDIUMTEXT(), "mysql"), nullable=False),
sa.Column("updated_at", UtcDateTime(), nullable=False),
+ # Optional early-expiry override. When set, GC deletes this row when expires_at < now()
+ # even if updated_at is recent. NULL means no early expiry — the row is still cleaned
+ # up by the global updated_at + default_retention_days check. Populated via
Review Comment:
Handled in df379c50e6
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +259,84 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows.
+
+ Reads ``[state_store] default_retention_days`` and ``[state_store] state_cleanup_batch_size``
+ from config. Each pass runs in its own transaction so partial progress is committed even if a
+ later pass fails. Each pass is batched to avoid long-running locks on the table.
+
+ Two passes:
+ a. Rows where updated_at < now() - default_retention_days (global retention)
+ b. Rows where expires_at < now() (per-key early expiry set by the operator)
+ """
+ retention_days = conf.getint("state_store", "default_retention_days")
+ batch_size = conf.getint("state_store", "state_cleanup_batch_size")
+ now = timezone.utcnow()
+ older_than = now - timedelta(days=retention_days) if retention_days > 0 else None
+
+ pk_cols = (
+ TaskStateModel.dag_run_id,
+ TaskStateModel.task_id,
+ TaskStateModel.map_index,
+ TaskStateModel.key,
+ )
Review Comment:
Handled in df379c50e6
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +259,84 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows.
+
+ Reads ``[state_store] default_retention_days`` and ``[state_store] state_cleanup_batch_size``
+ from config. Each pass runs in its own transaction so partial progress is committed even if a
+ later pass fails. Each pass is batched to avoid long-running locks on the table.
+
+ Two passes:
+ a. Rows where updated_at < now() - default_retention_days (global retention)
+ b. Rows where expires_at < now() (per-key early expiry set by the operator)
+ """
+ retention_days = conf.getint("state_store", "default_retention_days")
+ batch_size = conf.getint("state_store", "state_cleanup_batch_size")
+ now = timezone.utcnow()
+ older_than = now - timedelta(days=retention_days) if retention_days > 0 else None
+
+ pk_cols = (
+ TaskStateModel.dag_run_id,
+ TaskStateModel.task_id,
+ TaskStateModel.map_index,
+ TaskStateModel.key,
+ )
+
+ def _delete_batched(where_clause) -> int:
+ total = 0
+ while True:
+ with create_session() as session:
Review Comment:
Handled in df379c50e6
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3206809313
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +253,41 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows and orphaned asset state rows.
+
+ Reads ``[state_store] default_retention_days`` from config to determine the time-based
+ cutoff. Set to 0 to disable time-based cleanup (expires_at and orphan cleanup still run).
+
+ Two passes for task_state:
+ a. Rows where updated_at < now() - default_retention_days (global retention)
+ b. Rows where expires_at < now() (per-key early expiry set by the operator)
+
+ Asset state orphan cleanup: asset_state rows whose FK'd asset has no asset_active
+ entry are stale (asset removed from all DAGs) and are deleted here.
+ """
+ from datetime import timedelta
+
+ from airflow.configuration import conf
+
+ retention_days = conf.getint("state_store", "default_retention_days")
+ now = timezone.utcnow()
+ older_than = now - timedelta(days=retention_days) if retention_days > 0 else None
+ with create_session() as session:
+ if older_than:
+ session.execute(delete(TaskStateModel).where(TaskStateModel.updated_at < older_than))
+ session.execute(
+ delete(TaskStateModel).where(
+ TaskStateModel.expires_at.isnot(None),
+ TaskStateModel.expires_at < now,
+ )
+ )
+ active_asset_ids = select(AssetModel.id).join(
+ AssetActive, (AssetActive.name == AssetModel.name) & (AssetActive.uri == AssetModel.uri)
+ )
+ session.execute(delete(AssetStateModel).where(AssetStateModel.asset_id.not_in(active_asset_ids)))
+
Review Comment:
Thanks!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3201662393
##########
airflow-core/src/airflow/migrations/versions/0112_3_3_0_add_task_state_and_asset_state_tables.py:
##########
@@ -65,6 +65,11 @@ def upgrade():
sa.Column("run_id", StringID(), nullable=False),
sa.Column("value", sa.Text().with_variant(mysql.MEDIUMTEXT(), "mysql"), nullable=False),
sa.Column("updated_at", UtcDateTime(), nullable=False),
+ # Optional early-expiry override. When set, GC deletes this row when expires_at < now()
+ # even if updated_at is recent. NULL means no early expiry — the row is still cleaned
+ # up by the global updated_at + default_retention_days check. Populated via
Review Comment:
This is confusing IMO -- an `expires_at` of None should mean it never expires.
We can pre-compute the expires_at value at update time by reading the default_retention config then (i.e. cleanup becomes a simpler "SELECT where expires_at < Now()`.
This possibly also removes the need for an index on udpated_at.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3217747543
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +272,51 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows.
+
+ ``expires_at`` is set at write time on every ``set()`` call, so cleanup is a single
+ ``WHERE expires_at < now()`` pass. Rows with ``expires_at=NULL`` (default_retention_days=0)
+ are never deleted. Batching is configurable via ``[state_store] state_cleanup_batch_size``.
+ """
+ batch_size = conf.getint("state_store", "state_cleanup_batch_size")
+ now = timezone.utcnow()
+
+ def _delete_batched(where_clause) -> int:
+ total = 0
+ with create_session() as session:
+ while True:
+ id_query = select(TaskStateModel.id).where(where_clause)
+ if batch_size > 0:
+ id_query = id_query.limit(batch_size)
+ ids = session.scalars(id_query).all()
+ if not ids:
+ break
+ session.execute(delete(TaskStateModel).where(TaskStateModel.id.in_(ids)))
+ session.commit()
+ total += len(ids)
+ if batch_size <= 0 or len(ids) < batch_size:
+ break
+ return total
+
+ deleted = _delete_batched(TaskStateModel.expires_at < now)
+ log.info("Deleted expired task_state rows", rows_deleted=deleted)
+
+ def _dry_run_summary(self) -> dict[str, list]:
Review Comment:
`_dry_run_summary` doesn't look like a property? It's a regular function, were you proposing just a rename? I am totally ok with that just checking for intent.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3219190096
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +272,51 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows.
+
+ ``expires_at`` is set at write time on every ``set()`` call, so cleanup is a single
+ ``WHERE expires_at < now()`` pass. Rows with ``expires_at=NULL`` (default_retention_days=0)
+ are never deleted. Batching is configurable via ``[state_store] state_cleanup_batch_size``.
+ """
+ batch_size = conf.getint("state_store", "state_cleanup_batch_size")
+ now = timezone.utcnow()
+
+ def _delete_batched(where_clause) -> int:
+ total = 0
+ with create_session() as session:
+ while True:
+ id_query = select(TaskStateModel.id).where(where_clause)
+ if batch_size > 0:
+ id_query = id_query.limit(batch_size)
+ ids = session.scalars(id_query).all()
+ if not ids:
+ break
+ session.execute(delete(TaskStateModel).where(TaskStateModel.id.in_(ids)))
+ session.commit()
+ total += len(ids)
+ if batch_size <= 0 or len(ids) < batch_size:
+ break
+ return total
+
+ deleted = _delete_batched(TaskStateModel.expires_at < now)
Review Comment:
What do you mean by a second pass? Where would this happen? (In abstract it sounds like a plan; it’s similar to how the next run needs to be recalculated when you change the dag schedule definition.)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3231626009
##########
airflow-core/src/airflow/cli/cli_config.py:
##########
@@ -2108,6 +2115,11 @@ class GroupCommand(NamedTuple):
help="Display providers",
subcommands=PROVIDERS_COMMANDS,
),
+ GroupCommand(
+ name="state-store",
+ help="Manage task and asset state storage",
Review Comment:
This is not true based on what we have now. We don't do asset state management. We can create an issue or add a command that prints the message, but does not yet do anything (`raise NotImplemented`, maybe?) to ensure that we do not forget and even if we forget it still kinda make sense
##########
airflow-core/src/airflow/cli/commands/state_store_command.py:
##########
@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+
+from airflow.state import get_state_backend
+from airflow.state.metastore import MetastoreStateBackend
+
+log = logging.getLogger(__name__)
+
+# Other state operations (list, get, delete per key) will be added here in the future.
+
+
+def cleanup(args) -> None:
+ """Remove expired task state rows via the configured state backend."""
+ backend = get_state_backend()
+
+ if args.dry_run:
+ if isinstance(backend, MetastoreStateBackend):
+ summary = backend._summary_dry_run_()
+ expired = summary["expired"]
+ if not expired:
+ print("Nothing to delete.")
+ return
+ print(f"Would delete {len(expired)} task state row(s):\n")
+ for dag_id, run_id, task_id, map_index, key in expired:
+ print(
+ f" Dag {dag_id!r}, run {run_id!r}, task {task_id!r}, map_index {map_index!r}, key {key!r}"
+ )
+ else:
+ print("Custom backend configured — cannot preview rows.")
Review Comment:
If that's the case, we'd better emphasize this (only support MetastoreStateBackend) in the doc, help text, etc.
##########
airflow-core/tests/unit/state/test_metastore.py:
##########
Review Comment:
let's also add a test for other store to ensure the message is printted and it's a no-op
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3206380715
##########
airflow-core/src/airflow/cli/cli_config.py:
##########
@@ -1519,6 +1519,19 @@ class GroupCommand(NamedTuple):
args=(ARG_OUTPUT, ARG_VERBOSE),
),
)
+STATE_STORE_COMMANDS = (
+ ActionCommand(
+ name="cleanup",
+ help="Remove expired task state rows via the configured state backend",
Review Comment:
Sure
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3206440788
##########
airflow-core/src/airflow/migrations/versions/0112_3_3_0_add_task_state_and_asset_state_tables.py:
##########
@@ -65,6 +65,11 @@ def upgrade():
sa.Column("run_id", StringID(), nullable=False),
sa.Column("value", sa.Text().with_variant(mysql.MEDIUMTEXT(), "mysql"), nullable=False),
sa.Column("updated_at", UtcDateTime(), nullable=False),
+ # Optional early-expiry override. When set, GC deletes this row when expires_at < now()
+ # even if updated_at is recent. NULL means no early expiry — the row is still cleaned
+ # up by the global updated_at + default_retention_days check. Populated via
Review Comment:
Good call. Making those changes
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3196104464
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3047,6 +3054,19 @@ def _remove_unreferenced_triggers(self, *, session: Session = NEW_SESSION) -> No
.execution_options(synchronize_session="fetch")
)
+ def _cleanup_expired_task_state(self) -> None:
+ """
+ Remove expired task state rows and orphaned asset state rows.
+
+ Reads ``[state_store] default_retention_days`` from config and delegates to the
+ configured state backend. Runs on the interval set by ``[state_store] state_cleanup_interval``.
+ """
+ retention_days = conf.getint("state_store", "default_retention_days")
Review Comment:
Oops, fixed it! `retention_days=0` now disables time based cleanup entirely. Refactored so `MetastoreStateBackend.cleanup()` reads `default_retention_days` from config directly rather than receiving it as an argument, felt cleaner since the backend is responsible for enforcing its own retention policy. Scheduler just calls cleanup() with no args.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3234082675
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +272,51 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows.
+
+ ``expires_at`` is set at write time on every ``set()`` call, so cleanup is a single
+ ``WHERE expires_at < now()`` pass. Rows with ``expires_at=NULL`` (default_retention_days=0)
+ are never deleted. Batching is configurable via ``[state_store] state_cleanup_batch_size``.
+ """
+ batch_size = conf.getint("state_store", "state_cleanup_batch_size")
+ now = timezone.utcnow()
+
+ def _delete_batched(where_clause) -> int:
+ total = 0
+ with create_session() as session:
+ while True:
+ id_query = select(TaskStateModel.id).where(where_clause)
+ if batch_size > 0:
+ id_query = id_query.limit(batch_size)
+ ids = session.scalars(id_query).all()
+ if not ids:
+ break
+ session.execute(delete(TaskStateModel).where(TaskStateModel.id.in_(ids)))
+ session.commit()
+ total += len(ids)
+ if batch_size <= 0 or len(ids) < batch_size:
+ break
+ return total
+
+ deleted = _delete_batched(TaskStateModel.expires_at < now)
+ log.info("Deleted expired task_state rows", rows_deleted=deleted)
+
+ def _summary_dry_run_(self) -> dict[str, list]:
Review Comment:
Yes it is, fixed in 77450e580c
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3217771553
##########
airflow-core/tests/unit/state/test_metastore.py:
##########
@@ -234,6 +239,112 @@ def test_clear_with_all_map_indices_flag_wipes_wide(
assert backend.get(scope0, "job_id", session=session) is None
assert backend.get(scope1, "job_id", session=session) is None
+ def test_set_populates_expires_at(
+ self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+ ):
+ """set() always populates expires_at so cleanup has a single pass."""
+ scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+ backend.set(scope, "job_id", "app_1234", session=session)
+ session.flush()
+
+ row = session.scalar(select(TaskStateModel).where(TaskStateModel.key == "job_id"))
+ assert row is not None
Review Comment:
Handled in 7427d04ed2
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3200,6 +3215,21 @@ def _activate_assets_generate_warnings() -> Iterator[tuple[str, str]]:
session.add(warning)
existing_warned_dag_ids.add(warning.dag_id)
+ @staticmethod
+ def _cleanup_orphaned_asset_state(*, session: Session) -> None:
+ """
+ Delete asset_state rows for assets no longer active in any DAG.
Review Comment:
Handled in 7427d04ed2
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -3025,6 +3025,24 @@ state_store:
type: string
example: "mypackage.state.CustomStateBackend"
default: "airflow.state.metastore.MetastoreStateBackend"
+ default_retention_days:
+ description: |
+ Number of days to retain task_state rows after their last update.
Review Comment:
Handled in 7427d04ed2
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3217734827
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -3025,6 +3025,24 @@ state_store:
type: string
example: "mypackage.state.CustomStateBackend"
default: "airflow.state.metastore.MetastoreStateBackend"
+ default_retention_days:
+ description: |
+ Number of days to retain task_state rows after their last update.
+ Rows older than this are removed by the scheduler's periodic cleanup.
+ This config does not affect asset_state rows.
Review Comment:
asset_state rows are cleaned up via foreign key cascade and are deleted when the asset is deregistered — their lifecycle is tied to the asset, not a time window. `default_retention_days` only applies to task_state because that's the only one with time based expiry semantics. The description already calls this out explicitly on the next line.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3217740886
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3200,6 +3215,21 @@ def _activate_assets_generate_warnings() -> Iterator[tuple[str, str]]:
session.add(warning)
existing_warned_dag_ids.add(warning.dag_id)
+ @staticmethod
+ def _cleanup_orphaned_asset_state(*, session: Session) -> None:
+ """
+ Delete asset_state rows for assets no longer active in any DAG.
Review Comment:
Done, applied.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3225541284
##########
airflow-core/src/airflow/cli/commands/state_store_command.py:
##########
@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+
+from airflow.state import get_state_backend
+from airflow.state.metastore import MetastoreStateBackend
+
+log = logging.getLogger(__name__)
+
+# Other state operations (list, get, delete per key) will be added here in the future.
+
+
+def cleanup(args) -> None:
+ """Remove expired task state rows via the configured state backend."""
+ backend = get_state_backend()
+
+ if args.dry_run:
+ if isinstance(backend, MetastoreStateBackend):
+ summary = backend._summary_dry_run_()
+ expired = summary["expired"]
+ if not expired:
+ print("Nothing to delete.")
+ return
+ print(f"Would delete {len(expired)} task state row(s):\n")
+ for dag_id, run_id, task_id, map_index, key in expired:
+ print(
+ f" Dag {dag_id!r}, run {run_id!r}, task {task_id!r}, map_index {map_index!r}, key {key!r}"
+ )
+ else:
+ print("Custom backend configured — cannot preview rows.")
Review Comment:
Not really, no, this command makes sense for MetadataStateBackend only. The CLI is mostly used for that cleanup only
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3231896356
##########
airflow-core/src/airflow/cli/commands/state_store_command.py:
##########
@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+
+from airflow.state import get_state_backend
+from airflow.state.metastore import MetastoreStateBackend
+
+log = logging.getLogger(__name__)
+
+# Other state operations (list, get, delete per key) will be added here in the future.
+
+
+def cleanup(args) -> None:
+ """Remove expired task state rows via the configured state backend."""
+ backend = get_state_backend()
+
+ if args.dry_run:
+ if isinstance(backend, MetastoreStateBackend):
+ summary = backend._summary_dry_run_()
+ expired = summary["expired"]
+ if not expired:
+ print("Nothing to delete.")
+ return
+ print(f"Would delete {len(expired)} task state row(s):\n")
+ for dag_id, run_id, task_id, map_index, key in expired:
+ print(
+ f" Dag {dag_id!r}, run {run_id!r}, task {task_id!r}, map_index {map_index!r}, key {key!r}"
+ )
+ else:
+ print("Custom backend configured — cannot preview rows.")
Review Comment:
Handled in 6b6968e189
##########
airflow-core/tests/unit/state/test_metastore.py:
##########
Review Comment:
Handled in 6b6968e189
##########
airflow-core/src/airflow/cli/cli_config.py:
##########
@@ -2108,6 +2115,11 @@ class GroupCommand(NamedTuple):
help="Display providers",
subcommands=PROVIDERS_COMMANDS,
),
+ GroupCommand(
+ name="state-store",
+ help="Manage task and asset state storage",
Review Comment:
Handled in 6b6968e189
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3231437569
##########
airflow-core/src/airflow/cli/commands/state_store_command.py:
##########
@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+
+from airflow.state import get_state_backend
+from airflow.state.metastore import MetastoreStateBackend
+
+log = logging.getLogger(__name__)
+
+# Other state operations (list, get, delete per key) will be added here in the future.
+
+
+def cleanup(args) -> None:
+ """Remove expired task state rows via the configured state backend."""
+ backend = get_state_backend()
+
+ if args.dry_run:
+ if isinstance(backend, MetastoreStateBackend):
+ summary = backend._summary_dry_run_()
+ expired = summary["expired"]
+ if not expired:
+ print("Nothing to delete.")
+ return
+ print(f"Would delete {len(expired)} task state row(s):\n")
+ for dag_id, run_id, task_id, map_index, key in expired:
+ print(
+ f" Dag {dag_id!r}, run {run_id!r}, task {task_id!r}, map_index {map_index!r}, key {key!r}"
+ )
+ else:
+ print("Custom backend configured — cannot preview rows.")
Review Comment:
Good thought, making `dry_run_summary()` part of `BaseStateBackend` would clean up the `isinstance` check and give custom server-side backends a hook to implement their own preview. That said, the current design is intentional: the return format ({"expired": list} of DB rows) is specific to `MetastoreStateBackend` storage model. A Redis or S3 backend would likely have nothing meaningful to report here, or a completely different representation.
I'd keep it scoped to `MetastoreStateBackend` for now and track it as a need-to-do basis because I do not know if the CLI cleanup is ideal to cleanup custom backends (i think not). If custom backends ever need dry-run support, we can design the interface with their semantics in mind rather than retrofitting the DB row format.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3239443134
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +272,51 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows.
+
+ ``expires_at`` is set at write time on every ``set()`` call, so cleanup is a single
+ ``WHERE expires_at < now()`` pass. Rows with ``expires_at=NULL`` (default_retention_days=0)
+ are never deleted. Batching is configurable via ``[state_store] state_cleanup_batch_size``.
+ """
+ batch_size = conf.getint("state_store", "state_cleanup_batch_size")
+ now = timezone.utcnow()
+
+ def _delete_batched(where_clause) -> int:
+ total = 0
+ with create_session() as session:
+ while True:
+ id_query = select(TaskStateModel.id).where(where_clause)
+ if batch_size > 0:
+ id_query = id_query.limit(batch_size)
+ ids = session.scalars(id_query).all()
+ if not ids:
+ break
+ session.execute(delete(TaskStateModel).where(TaskStateModel.id.in_(ids)))
+ session.commit()
+ total += len(ids)
+ if batch_size <= 0 or len(ids) < batch_size:
+ break
+ return total
+
+ deleted = _delete_batched(TaskStateModel.expires_at < now)
Review Comment:
Resolving this with the response above, if you still feel the need to add that, we can chat and I can do in a follow up
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3206523316
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +259,84 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows.
+
+ Reads ``[state_store] default_retention_days`` and ``[state_store] state_cleanup_batch_size``
+ from config. Each pass runs in its own transaction so partial progress is committed even if a
+ later pass fails. Each pass is batched to avoid long-running locks on the table.
+
+ Two passes:
+ a. Rows where updated_at < now() - default_retention_days (global retention)
+ b. Rows where expires_at < now() (per-key early expiry set by the operator)
+ """
+ retention_days = conf.getint("state_store", "default_retention_days")
+ batch_size = conf.getint("state_store", "state_cleanup_batch_size")
+ now = timezone.utcnow()
+ older_than = now - timedelta(days=retention_days) if retention_days > 0 else None
+
+ pk_cols = (
+ TaskStateModel.dag_run_id,
+ TaskStateModel.task_id,
+ TaskStateModel.map_index,
+ TaskStateModel.key,
+ )
+
+ def _delete_batched(where_clause) -> int:
+ total = 0
+ while True:
+ with create_session() as session:
Review Comment:
Oh yes, `_delete_batched` now opens a single session for the entire cleanup pass and commits explicitly after each batch delete.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3217080518
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -33,7 +33,22 @@
from itertools import groupby
from typing import TYPE_CHECKING, Any, cast
-from sqlalchemy import CTE, and_, case, delete, exists, func, inspect, or_, select, text, tuple_, update
+from sqlalchemy import (
+ CTE,
+ and_,
+ case,
+ delete,
+ delete as _delete,
Review Comment:
Why do these need to be aliases? (Also e.g. `_AssetActive`)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3196026872
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -3025,6 +3014,29 @@ state_store:
type: string
example: "mypackage.state.CustomStateBackend"
default: "airflow.state.metastore.MetastoreStateBackend"
+ default_retention_days:
+ description: |
+ Number of days to retain task state rows after their last update.
Review Comment:
`default_retention_days` only applies to `task_state` rows. `asset_state` has no time-based GC — it is cleared explicitly (as per AIP) or via orphan cleanup when the asset is deregistered. Updated the config description to make this clear.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "jason810496 (via GitHub)" <gi...@apache.org>.
jason810496 commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3198954290
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +253,41 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows and orphaned asset state rows.
+
+ Reads ``[state_store] default_retention_days`` from config to determine the time-based
+ cutoff. Set to 0 to disable time-based cleanup (expires_at and orphan cleanup still run).
+
+ Two passes for task_state:
+ a. Rows where updated_at < now() - default_retention_days (global retention)
+ b. Rows where expires_at < now() (per-key early expiry set by the operator)
+
+ Asset state orphan cleanup: asset_state rows whose FK'd asset has no asset_active
+ entry are stale (asset removed from all DAGs) and are deleted here.
+ """
+ from datetime import timedelta
+
+ from airflow.configuration import conf
+
+ retention_days = conf.getint("state_store", "default_retention_days")
+ now = timezone.utcnow()
+ older_than = now - timedelta(days=retention_days) if retention_days > 0 else None
+ with create_session() as session:
+ if older_than:
+ session.execute(delete(TaskStateModel).where(TaskStateModel.updated_at < older_than))
+ session.execute(
+ delete(TaskStateModel).where(
+ TaskStateModel.expires_at.isnot(None),
+ TaskStateModel.expires_at < now,
+ )
+ )
+ active_asset_ids = select(AssetModel.id).join(
+ AssetActive, (AssetActive.name == AssetModel.name) & (AssetActive.uri == AssetModel.uri)
+ )
+ session.execute(delete(AssetStateModel).where(AssetStateModel.asset_id.not_in(active_asset_ids)))
+
Review Comment:
If it's a valid assumption that users might produce a large amount of state records here between the `state_cleanup_interval` time window. I have some concern regarding the single pass delete Transaction here.
I just double checked the concern with Claude:
---
Yes, this is a real problem. Several compounding issues:
1. Missing indexes — every cleanup will be a full table scan
The two predicates the cleanup filters on are not indexed:
- task_state.updated_at — no index (only task_state_pkey on (dag_run_id,
task_id, map_index, key) and idx_task_state_lookup on (dag_id, run_id,
task_id, map_index))
- task_state.expires_at — no index (just added in this PR)
So both DELETE WHERE updated_at < cutoff and DELETE WHERE expires_at < now()
do full sequential scans. On a deployment with millions of rows that's minutes
of scanning every 24h, plus the locks held for the whole duration.
2. No batching / no LIMIT
Compare to airflow db cleanup (utils/db_cleanup.py:217), which deletes in
configurable batches and commits between them. The new path runs three plain
bulk DELETEs in a single session. Long-running bulk DELETE means:
- Row locks held for the duration (writers calling task_state.set() upserts on
matching rows block — they queue behind the cleanup transaction).
- On Postgres: massive WAL churn, autovacuum can't keep up, table bloat.
- On MySQL/InnoDB at REPEATABLE READ (Airflow's default): next-key/gap locks
make conflicts even more likely.
3. All three DELETEs share one transaction
with create_session() as session: opens one session; each session.execute()
runs inside it; commit happens at exit. If pass 1 takes 90s, the locks from
pass 1 are held while pass 2 and pass 3 run. A failure in pass 3 rolls back
passes 1 and 2 (cleanup makes no forward progress at all).
4. Scheduler main loop is blocked
_cleanup_expired_task_state is registered via call_regular_interval, which is
synchronous in the scheduler loop. Same pattern as
_remove_unreferenced_triggers and _update_asset_orphanage — but those have
small cardinality. task_state is user-driven and unbounded (the AIP encourages
users to write a lot of it). With a multi-minute cleanup the scheduler is not
scheduling for those minutes.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3206627881
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +259,84 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows.
+
+ Reads ``[state_store] default_retention_days`` and ``[state_store] state_cleanup_batch_size``
+ from config. Each pass runs in its own transaction so partial progress is committed even if a
+ later pass fails. Each pass is batched to avoid long-running locks on the table.
+
+ Two passes:
+ a. Rows where updated_at < now() - default_retention_days (global retention)
+ b. Rows where expires_at < now() (per-key early expiry set by the operator)
+ """
+ retention_days = conf.getint("state_store", "default_retention_days")
+ batch_size = conf.getint("state_store", "state_cleanup_batch_size")
+ now = timezone.utcnow()
+ older_than = now - timedelta(days=retention_days) if retention_days > 0 else None
+
+ pk_cols = (
+ TaskStateModel.dag_run_id,
+ TaskStateModel.task_id,
+ TaskStateModel.map_index,
+ TaskStateModel.key,
+ )
Review Comment:
Yeah I realised it too. Will use `INTEGER AUTOINCREMENT` here, it will work just fine and no overhead to generate UUID
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh closed pull request #66463: AIP-103: Adding periodic task state garbage collection and retention support
URL: https://github.com/apache/airflow/pull/66463
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3201411935
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +253,41 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows and orphaned asset state rows.
+
+ Reads ``[state_store] default_retention_days`` from config to determine the time-based
+ cutoff. Set to 0 to disable time-based cleanup (expires_at and orphan cleanup still run).
+
+ Two passes for task_state:
+ a. Rows where updated_at < now() - default_retention_days (global retention)
+ b. Rows where expires_at < now() (per-key early expiry set by the operator)
+
+ Asset state orphan cleanup: asset_state rows whose FK'd asset has no asset_active
+ entry are stale (asset removed from all DAGs) and are deleted here.
+ """
+ from datetime import timedelta
+
+ from airflow.configuration import conf
+
+ retention_days = conf.getint("state_store", "default_retention_days")
+ now = timezone.utcnow()
+ older_than = now - timedelta(days=retention_days) if retention_days > 0 else None
+ with create_session() as session:
+ if older_than:
+ session.execute(delete(TaskStateModel).where(TaskStateModel.updated_at < older_than))
+ session.execute(
+ delete(TaskStateModel).where(
+ TaskStateModel.expires_at.isnot(None),
+ TaskStateModel.expires_at < now,
+ )
+ )
+ active_asset_ids = select(AssetModel.id).join(
+ AssetActive, (AssetActive.name == AssetModel.name) & (AssetActive.uri == AssetModel.uri)
+ )
+ session.execute(delete(AssetStateModel).where(AssetStateModel.asset_id.not_in(active_asset_ids)))
+
Review Comment:
Thanks, good catches.
I will address all of them except last now cos its invalid from scheduler perspective, its a cli command now
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3225456247
##########
airflow-core/src/airflow/cli/commands/state_store_command.py:
##########
@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+
+from airflow.state import get_state_backend
+from airflow.state.metastore import MetastoreStateBackend
+
+log = logging.getLogger(__name__)
+
+# Other state operations (list, get, delete per key) will be added here in the future.
+
+
+def cleanup(args) -> None:
+ """Remove expired task state rows via the configured state backend."""
+ backend = get_state_backend()
+
+ if args.dry_run:
+ if isinstance(backend, MetastoreStateBackend):
+ summary = backend._summary_dry_run_()
+ expired = summary["expired"]
+ if not expired:
+ print("Nothing to delete.")
+ return
+ print(f"Would delete {len(expired)} task state row(s):\n")
+ for dag_id, run_id, task_id, map_index, key in expired:
+ print(
+ f" Dag {dag_id!r}, run {run_id!r}, task {task_id!r}, map_index {map_index!r}, key {key!r}"
+ )
+ else:
+ print("Custom backend configured — cannot preview rows.")
Review Comment:
or should we make `_summary_dry_run_` part of the base state backend? if it's not implented then we should this message
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3231985467
##########
airflow-core/tests/unit/state/test_metastore.py:
##########
Review Comment:
Handled in https://github.com/apache/airflow/commit/6b6968e1893f5198edef93e6c3cc6e1978935e8e
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3233320001
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -3025,6 +3025,24 @@ state_store:
type: string
example: "mypackage.state.CustomStateBackend"
default: "airflow.state.metastore.MetastoreStateBackend"
+ default_retention_days:
+ description: |
+ Number of days to retain task state after their last update.
+ Rows older than this are removed by the scheduler's periodic cleanup.
Review Comment:
this is outdated descirpiton
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3217772729
##########
airflow-core/src/airflow/cli/commands/state_store_command.py:
##########
@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+
+log = logging.getLogger(__name__)
+
+# Other state operations (list, get, delete per key) will be added here in the future.
+
+
+def cleanup(args) -> None:
+ """Remove expired task state rows via the configured state backend."""
+ from airflow.state import get_state_backend
+ from airflow.state.metastore import MetastoreStateBackend
+
+ backend = get_state_backend()
+
+ if args.dry_run:
+ if isinstance(backend, MetastoreStateBackend):
+ summary = backend._dry_run_summary()
+ expired = summary["expired"]
+ if not expired:
+ print("Nothing to delete.")
+ return
+ print(f"Would delete {len(expired)} task state row(s):\n")
+ for dag_id, run_id, task_id, map_index, key in expired:
+ print(
+ f" DAG {dag_id!r}, run {run_id!r}, task {task_id!r}, map_index {map_index!r}, key {key!r}"
Review Comment:
Handled in 7427d04ed2
##########
airflow-core/src/airflow/cli/commands/state_store_command.py:
##########
@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+
+log = logging.getLogger(__name__)
+
+# Other state operations (list, get, delete per key) will be added here in the future.
+
+
+def cleanup(args) -> None:
+ """Remove expired task state rows via the configured state backend."""
+ from airflow.state import get_state_backend
+ from airflow.state.metastore import MetastoreStateBackend
Review Comment:
Handled in 7427d04ed2
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3224009523
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +272,51 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows.
+
+ ``expires_at`` is set at write time on every ``set()`` call, so cleanup is a single
+ ``WHERE expires_at < now()`` pass. Rows with ``expires_at=NULL`` (default_retention_days=0)
+ are never deleted. Batching is configurable via ``[state_store] state_cleanup_batch_size``.
+ """
+ batch_size = conf.getint("state_store", "state_cleanup_batch_size")
+ now = timezone.utcnow()
+
+ def _delete_batched(where_clause) -> int:
+ total = 0
+ with create_session() as session:
+ while True:
+ id_query = select(TaskStateModel.id).where(where_clause)
+ if batch_size > 0:
+ id_query = id_query.limit(batch_size)
+ ids = session.scalars(id_query).all()
+ if not ids:
+ break
+ session.execute(delete(TaskStateModel).where(TaskStateModel.id.in_(ids)))
+ session.commit()
+ total += len(ids)
+ if batch_size <= 0 or len(ids) < batch_size:
+ break
+ return total
+
+ deleted = _delete_batched(TaskStateModel.expires_at < now)
Review Comment:
A second pass would be something like `WHERE expires_at IS NULL AND updated_at < now - default_retention_days`, to catch rows that were written when `default_retention_days=0` but the config was later raised.
Something like:
```python
# Pass 1: code right now
deleted_expired = _delete_batched(TaskStateModel.expires_at < now)
# Pass 2: rows with NULL expires_at that are stale under the current global default
if default_retention_days > 0:
cutoff = now - timedelta(days=default_retention_days)
deleted_stale = _delete_batched(
TaskStateModel.expires_at.is_(None) & (TaskStateModel.updated_at < cutoff)
)
```
It would run in the same `airflow state-store cleanup` command
But on thinking more, I do not think that it is needed. `expires_at=NULL` is an explicit signal — either `default_retention_days=0` was set, or `retention_days=0` was passed at write time. Both mean "keep this row forever." Retroactively deleting them on a config change would violate what was promised at write time.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on PR #66463:
URL: https://github.com/apache/airflow/pull/66463#issuecomment-4448211593
Thanks for review folks, merging this one, we can always revise it as need comes.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3217400057
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -3025,6 +3025,24 @@ state_store:
type: string
example: "mypackage.state.CustomStateBackend"
default: "airflow.state.metastore.MetastoreStateBackend"
+ default_retention_days:
+ description: |
+ Number of days to retain task_state rows after their last update.
+ Rows older than this are removed by the scheduler's periodic cleanup.
+ This config does not affect asset_state rows.
Review Comment:
if it does not affect asset state, should we have different config for each?
##########
airflow-core/src/airflow/cli/commands/state_store_command.py:
##########
@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+
+log = logging.getLogger(__name__)
+
+# Other state operations (list, get, delete per key) will be added here in the future.
+
+
+def cleanup(args) -> None:
+ """Remove expired task state rows via the configured state backend."""
+ from airflow.state import get_state_backend
+ from airflow.state.metastore import MetastoreStateBackend
Review Comment:
Why do we do inline imports here?
##########
airflow-core/src/airflow/cli/commands/state_store_command.py:
##########
@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+
+log = logging.getLogger(__name__)
+
+# Other state operations (list, get, delete per key) will be added here in the future.
+
+
+def cleanup(args) -> None:
+ """Remove expired task state rows via the configured state backend."""
+ from airflow.state import get_state_backend
+ from airflow.state.metastore import MetastoreStateBackend
+
+ backend = get_state_backend()
+
+ if args.dry_run:
+ if isinstance(backend, MetastoreStateBackend):
+ summary = backend._dry_run_summary()
+ expired = summary["expired"]
+ if not expired:
+ print("Nothing to delete.")
+ return
+ print(f"Would delete {len(expired)} task state row(s):\n")
+ for dag_id, run_id, task_id, map_index, key in expired:
+ print(
+ f" DAG {dag_id!r}, run {run_id!r}, task {task_id!r}, map_index {map_index!r}, key {key!r}"
Review Comment:
```suggestion
f" Dag {dag_id!r}, run {run_id!r}, task {task_id!r}, map_index {map_index!r}, key {key!r}"
```
##########
airflow-core/src/airflow/state/metastore.py:
##########
@@ -252,6 +272,51 @@ def _clear_asset_state(self, scope: AssetScope, *, session: Session) -> None:
)
)
+ def cleanup(self) -> None:
+ """
+ Remove expired task state rows.
+
+ ``expires_at`` is set at write time on every ``set()`` call, so cleanup is a single
+ ``WHERE expires_at < now()`` pass. Rows with ``expires_at=NULL`` (default_retention_days=0)
+ are never deleted. Batching is configurable via ``[state_store] state_cleanup_batch_size``.
+ """
+ batch_size = conf.getint("state_store", "state_cleanup_batch_size")
+ now = timezone.utcnow()
+
+ def _delete_batched(where_clause) -> int:
+ total = 0
+ with create_session() as session:
+ while True:
+ id_query = select(TaskStateModel.id).where(where_clause)
+ if batch_size > 0:
+ id_query = id_query.limit(batch_size)
+ ids = session.scalars(id_query).all()
+ if not ids:
+ break
+ session.execute(delete(TaskStateModel).where(TaskStateModel.id.in_(ids)))
+ session.commit()
+ total += len(ids)
+ if batch_size <= 0 or len(ids) < batch_size:
+ break
+ return total
+
+ deleted = _delete_batched(TaskStateModel.expires_at < now)
+ log.info("Deleted expired task_state rows", rows_deleted=deleted)
+
+ def _dry_run_summary(self) -> dict[str, list]:
Review Comment:
```suggestion
def _summary_dry_run_(self) -> dict[str, list]:
```
since it's not a property and probably shoudn't be one
##########
airflow-core/tests/unit/state/test_metastore.py:
##########
@@ -234,6 +239,112 @@ def test_clear_with_all_map_indices_flag_wipes_wide(
assert backend.get(scope0, "job_id", session=session) is None
assert backend.get(scope1, "job_id", session=session) is None
+ def test_set_populates_expires_at(
+ self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+ ):
+ """set() always populates expires_at so cleanup has a single pass."""
+ scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+ backend.set(scope, "job_id", "app_1234", session=session)
+ session.flush()
+
+ row = session.scalar(select(TaskStateModel).where(TaskStateModel.key == "job_id"))
+ assert row is not None
Review Comment:
it is possible for us to check the attributes instead of just existence instead
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3200,6 +3215,21 @@ def _activate_assets_generate_warnings() -> Iterator[tuple[str, str]]:
session.add(warning)
existing_warned_dag_ids.add(warning.dag_id)
+ @staticmethod
+ def _cleanup_orphaned_asset_state(*, session: Session) -> None:
+ """
+ Delete asset_state rows for assets no longer active in any DAG.
Review Comment:
```suggestion
Delete asset_state rows for assets no longer active in any Dag.
```
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -3025,6 +3025,24 @@ state_store:
type: string
example: "mypackage.state.CustomStateBackend"
default: "airflow.state.metastore.MetastoreStateBackend"
+ default_retention_days:
+ description: |
+ Number of days to retain task_state rows after their last update.
Review Comment:
```suggestion
Number of days to retain task state after their last update.
```
Since it's user-facing, let's not emphasize the DB aspect.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-103: Adding periodic task state garbage collection and retention support [airflow]
Posted by "amoghrajesh (via GitHub)" <gi...@apache.org>.
amoghrajesh commented on code in PR #66463:
URL: https://github.com/apache/airflow/pull/66463#discussion_r3217726360
##########
airflow-core/src/airflow/cli/commands/state_store_command.py:
##########
@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+
+log = logging.getLogger(__name__)
+
+# Other state operations (list, get, delete per key) will be added here in the future.
+
+
+def cleanup(args) -> None:
+ """Remove expired task state rows via the configured state backend."""
+ from airflow.state import get_state_backend
+ from airflow.state.metastore import MetastoreStateBackend
+
+ backend = get_state_backend()
+
+ if args.dry_run:
+ if isinstance(backend, MetastoreStateBackend):
+ summary = backend._dry_run_summary()
+ expired = summary["expired"]
+ if not expired:
+ print("Nothing to delete.")
+ return
+ print(f"Would delete {len(expired)} task state row(s):\n")
+ for dag_id, run_id, task_id, map_index, key in expired:
+ print(
+ f" DAG {dag_id!r}, run {run_id!r}, task {task_id!r}, map_index {map_index!r}, key {key!r}"
Review Comment:
Handled
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -3025,6 +3025,24 @@ state_store:
type: string
example: "mypackage.state.CustomStateBackend"
default: "airflow.state.metastore.MetastoreStateBackend"
+ default_retention_days:
+ description: |
+ Number of days to retain task_state rows after their last update.
Review Comment:
Handled, good call
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org