amoghrajesh opened a new pull request, #75729:
URL: https://github.com/apache/airflow/pull/75729
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull
request.
Write a good git commit message following this guide:
https://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the
documentation.
For user-facing UI changes, please attach before/after screenshots (or a
short
screen recording) so reviewers can assess the visual impact.
Feel free to ping (in general) for the review if you do not see reaction for
a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we
sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR,
please
change below checkbox to `[X]` followed by the name of the tool, uncomment
the "Generated-by".
-->
- [ ] Yes (please specify the tool below)
<!--
Generated-by: [Tool Name] following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
### What
There is no way measure what / who is using the state store. Debugging /
Adoption questions like "which Dags touch it", "is it read heavy or write
heavy", "is anyone using `delete`" cannot be answered from Airflow today at all.
The `resumable_job.*` counters do not fill this gap today. They are emitted
by `ResumableJobMixin`, so they only see operators that use the mixin. The two
other ways to get durable execution are invisible to them:
`KubernetesPodOperator`, which reads and writes `pod_identifier` through a
class attribute rather than the mixin, and user code calling `task_state_store`
directly from a `@task` or `PythonOperator`.
### Current behaviour
The state store accessors emit nothing. Every call through
`task_state_store` and `asset_state_store` is unmeasured, whatever/whoever
issued it.
### Proposed change
Nine counters at the SDK accessors, which every caller passes through
regardless of operator:
| Task store | Asset store |
| -- | -- |
| `task_state_store.read_found` | `asset_state_store.read_found` |
| `task_state_store.read_absent` | `asset_state_store.read_absent` |
| `task_state_store.write` | `asset_state_store.write` |
| `task_state_store.delete` | `asset_state_store.delete` |
| `task_state_store.clear` | |
Tagged with `dag_id`, `task_id`, `run_type`, `operator`, and `team_name` on
multi-team deployments. There is no scope-level clear on the asset side, hence
eight plus one.
### Changes of Note
**These count calls, not saves.** A read that finds a value is not proof
that work was avoided. The accessor never learns whether the operator managed
to reconnect to the external job, because that lives outside Airflow and only
the operator can
ask. So these counters answer "who uses the store"; `resumable_job.*`
continues to answer "how often did it save work".
**Task store and asset store get separate names rather than one counter with
a `scope` tag.**
`TaskScope` carries `run_id`, so a task store read that finds something was
written by an earlier attempt of the same Dag run. `AssetScope` does not, so an
asset read that finds something is an ordinary watermark read across runs.
**The granularity is per call, not per attempt.** One `AgentOperator` task
emits roughly 21 counts where a `KubernetesPodOperator` task emits 1, so these
cannot be compared across operators by raw volume without normalising.
**No `state_key` tag.** `AgentOperator` writes one key per agent step, so
key cardinality is unbounded. Key-level detail belongs in logs, not metrics.
**Where the emits sit.** Reads go in `_extract_get_response`, which `get`
and `aget` already share on both accessors and which already distinguishes
found from absent, so one emit point covers sync and async. Writes, deletes and
clears emit at their call sites. Tags are built once in
`RuntimeTaskInstance.get_template_context` as `ti.stats_tags` plus `operator`,
and
passed into the accessor constructors, so nothing has to be looked up at
call time.
### User implications / backcompat
No breaking change. The tags only materialise on backends that carry them:
OTel and dogstatsd natively, StatsD only with `[metrics]
statsd_influxdb_enabled`.
`stats_tags` defaults to empty when an accessor is constructed directly, so
existing callers and tests keep working.
New series are added, bounded by `dag_id` x `task_id` x `operator`.
### Testing
I validated certain cases manually with breeze and otel integration -
`breeze start-airflow --integration otel`, reading the collector at
`localhost:28889/metrics`.
These cases matter to me:
**1. Plain `@task`, both read outcomes.**
Reads a missing key, writes it, reads it back.
```python
from __future__ import annotations
from datetime import datetime
from airflow.sdk import DAG, task
with DAG(
dag_id="state_store_metrics_01",
schedule=None,
start_date=datetime(2026, 1, 1),
catchup=False,
):
@task
def touch_store(**context):
store = context["task_state_store"]
print("first read:", store.get("counter"))
store.set("counter", 1)
print("second read:", store.get("counter"))
touch_store()
```
<img width="1650" height="809" alt="image"
src="https://github.com/user-attachments/assets/ab900908-7115-40df-9d70-215e6694f7c3";
/>
```result
airflow_task_state_store_read_absent_total{dag_id="state_store_metrics_01",instance="c5c9dc67-88f7-411f-850c-60a0e6c82c3a",job="test",operator="_PythonDecoratedOperator",otel_scope_name="airflow.sdk._shared.observability.metrics.otel_logger",otel_scope_schema_url="",otel_scope_version="",run_type="manual",task_id="touch_store"}
1
airflow_task_state_store_read_found_total{dag_id="state_store_metrics_01",instance="c5c9dc67-88f7-411f-850c-60a0e6c82c3a",job="test",operator="_PythonDecoratedOperator",otel_scope_name="airflow.sdk._shared.observability.metrics.otel_logger",otel_scope_schema_url="",otel_scope_version="",run_type="manual",task_id="touch_store"}
1
airflow_task_state_store_write_total{dag_id="state_store_metrics_01",instance="c5c9dc67-88f7-411f-850c-60a0e6c82c3a",job="test",operator="_PythonDecoratedOperator",otel_scope_name="airflow.sdk._shared.observability.metrics.otel_logger",otel_scope_schema_url="",otel_scope_version="",run_type="manual",task_id="touch_store"}
1
```
Both read outcomes distinguished, tags present on a task with no operator or
mixin involved.
**2. Asset store alongside task store.**
One task touching both.
```python
from __future__ import annotations
from datetime import datetime
from airflow.sdk import DAG, Asset, task
METRICS_ASSET = Asset(name="state_store_metrics/asset",
uri="s3://bucket/metrics")
with DAG(
dag_id="state_store_metrics_02",
schedule=None,
start_date=datetime(2026, 1, 1),
catchup=False,
):
@task(inlets=[METRICS_ASSET], outlets=[METRICS_ASSET])
def touch_asset_store(**context):
task_store = context["task_state_store"]
asset_store = context["asset_state_store"]
print("asset read:", asset_store.get("watermark"))
asset_store.set("watermark", "2026-01-01")
print("task read:", task_store.get("counter"))
touch_asset_store()
```
Ran it two times to show the difference:
<img width="1650" height="809" alt="image"
src="https://github.com/user-attachments/assets/d3809975-e75f-40b9-8662-2d65795dcb5f";
/>
<img width="1650" height="809" alt="image"
src="https://github.com/user-attachments/assets/8d030b8a-b9ca-46c3-a9d9-bbdafb20fa30";
/>
```result
amoghrajesh.desai …/airflow task-state-asset-state-metrics $!?
orbstack ♥ 13:09 curl -s localhost:28889/metrics \
| grep -E
'^airflow_(task|asset)_state_store.*dag_id="state_store_metrics_02"' \
| sed -E
's/(instance|job|otel_scope_name|otel_scope_schema_url|otel_scope_version)="[^"]*",?//g'
airflow_asset_state_store_read_absent_total{dag_id="state_store_metrics_02",operator="_PythonDecoratedOperator",run_type="manual",task_id="touch_asset_store"}
1
airflow_asset_state_store_write_total{dag_id="state_store_metrics_02",operator="_PythonDecoratedOperator",run_type="manual",task_id="touch_asset_store"}
1
airflow_task_state_store_read_absent_total{dag_id="state_store_metrics_02",operator="_PythonDecoratedOperator",run_type="manual",task_id="touch_asset_store"}
1
amoghrajesh.desai …/airflow task-state-asset-state-metrics $!?
orbstack ♥ 13:09 curl -s localhost:28889/metrics \
| grep -E
'^airflow_(task|asset)_state_store.*dag_id="state_store_metrics_02"' \
| sed -E
's/(instance|job|otel_scope_name|otel_scope_schema_url|otel_scope_version)="[^"]*",?//g'
airflow_asset_state_store_read_absent_total{dag_id="state_store_metrics_02",operator="_PythonDecoratedOperator",run_type="manual",task_id="touch_asset_store"}
1
airflow_asset_state_store_read_found_total{dag_id="state_store_metrics_02",operator="_PythonDecoratedOperator",run_type="manual",task_id="touch_asset_store"}
1
airflow_asset_state_store_write_total{dag_id="state_store_metrics_02",operator="_PythonDecoratedOperator",run_type="manual",task_id="touch_asset_store"}
1
airflow_asset_state_store_write_total{dag_id="state_store_metrics_02",operator="_PythonDecoratedOperator",run_type="manual",task_id="touch_asset_store"}
1
airflow_task_state_store_read_absent_total{dag_id="state_store_metrics_02",operator="_PythonDecoratedOperator",run_type="manual",task_id="touch_asset_store"}
1
airflow_task_state_store_read_absent_total{dag_id="state_store_metrics_02",operator="_PythonDecoratedOperator",run_type="manual",task_id="touch_asset_store"}
1
```
**3. Delete and clear.**
```python
from __future__ import annotations
from datetime import datetime
from airflow.sdk import DAG, Asset, task
METRICS_ASSET = Asset(name="state_store_metrics/asset3",
uri="s3://bucket/metrics3")
with DAG(
dag_id="state_store_metrics_03",
schedule=None,
start_date=datetime(2026, 1, 1),
catchup=False,
):
@task(inlets=[METRICS_ASSET], outlets=[METRICS_ASSET])
def delete_and_clear(**context):
task_store = context["task_state_store"]
asset_store = context["asset_state_store"]
task_store.set("a", 1)
task_store.set("b", 2)
task_store.delete("a")
task_store.clear()
asset_store.set("w", "x")
asset_store.delete("w")
delete_and_clear()
```
<img width="1650" height="809" alt="image"
src="https://github.com/user-attachments/assets/e2b8028f-03d7-4c93-8047-90263026ce72";
/>
```result
amoghrajesh.desai …/airflow task-state-asset-state-metrics $!?
orbstack ♥ 13:09 curl -s localhost:28889/metrics \
| grep -E
'^airflow_(task|asset)_state_store.*dag_id="state_store_metrics_03"' \
| sed -E
's/(instance|job|otel_scope_name|otel_scope_schema_url|otel_scope_version)="[^"]*",?//g'
airflow_asset_state_store_delete_total{dag_id="state_store_metrics_03",operator="_PythonDecoratedOperator",run_type="manual",task_id="delete_and_clear"}
1
airflow_asset_state_store_write_total{dag_id="state_store_metrics_03",operator="_PythonDecoratedOperator",run_type="manual",task_id="delete_and_clear"}
1
airflow_task_state_store_clear_total{dag_id="state_store_metrics_03",operator="_PythonDecoratedOperator",run_type="manual",task_id="delete_and_clear"}
1
airflow_task_state_store_delete_total{dag_id="state_store_metrics_03",operator="_PythonDecoratedOperator",run_type="manual",task_id="delete_and_clear"}
1
airflow_task_state_store_write_total{dag_id="state_store_metrics_03",operator="_PythonDecoratedOperator",run_type="manual",task_id="delete_and_clear"}
2
```
`write` at 2 for two `set` calls confirms the per-call grain. `delete` and
`clear` stay separate, which matters because `AgentOperator` deletes every key
it wrote on success.
**4. `ResumableJobMixin` operator, both families side by side.**
A mixin operator submitting a mock job, then cleared and re-run.
```python
from __future__ import annotations
import random
import string
import time
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any
from airflow.sdk import DAG, BaseOperator
from airflow.sdk.bases.resumablejobmixin import ResumableJobMixin
if TYPE_CHECKING:
from pydantic import JsonValue
from airflow.sdk.definitions.context import Context
_JOBS: dict[str, str] = {}
def _generate_job_id() -> str:
return "batch-" + "".join(random.choices(string.ascii_lowercase +
string.digits, k=8))
def _submit(job_id: str) -> None:
_JOBS[job_id] = "RUNNING"
def _get_status(job_id: str) -> str:
return _JOBS.get(job_id, "UNKNOWN")
def _poll(job_id: str) -> None:
time.sleep(1)
_JOBS[job_id] = "SUCCEEDED"
# ---------------------------------------------------------------------------
# Operator
# ---------------------------------------------------------------------------
class MockBatchOperator(ResumableJobMixin, BaseOperator):
"""
Submits a mock batch job and polls until completion.
Demonstrates how ``ResumableJobMixin`` handles a worker crash between
submission and the start of polling: on retry the job ID is read from
``task_store`` and the operator reconnects rather than resubmitting.
"""
external_id_key = "job_id"
def execute(self, context: Context) -> Any:
return self.execute_resumable(context)
def submit_job(self, context: Context) -> JsonValue:
job_id = _generate_job_id()
_submit(job_id)
self.log.info("Submitted batch job", job_id=job_id)
return job_id
def get_job_status(self, external_id: JsonValue, context: Context) ->
str:
simulate = context["params"].get("simulate_final_status", "")
if simulate:
return simulate
# Default to RUNNING if the ID is not in the local cache — this
simulates
# a real external system that keeps running after a worker crash.
return _JOBS.get(str(external_id), "RUNNING")
def is_job_active(self, status: str) -> bool:
return status == "RUNNING"
def is_job_succeeded(self, status: str) -> bool:
return status == "SUCCEEDED"
def poll_until_complete(self, external_id: JsonValue, context: Context)
-> None:
self.log.info("Polling job — kill the worker now to trigger a
retry", job_id=external_id)
time.sleep(10)
self.log.info("Polling job until complete", job_id=external_id)
_poll(str(external_id))
def get_job_result(self, external_id: JsonValue, context: Context) ->
str:
return f"Job {external_id} finished with status
{_get_status(str(external_id))}"
with DAG(
dag_id="example_resumable_job",
schedule=None,
start_date=datetime(2026, 1, 1),
catchup=False,
tags=["example", "resumable", "task-store"],
doc_md=__doc__,
params={"simulate_final_status": ""},
):
MockBatchOperator(
task_id="run_batch_job",
retries=2,
retry_delay=timedelta(seconds=5),
)
```
```
airflow_task_state_store_read_absent_total{dag_id="example_resumable_job",task_id="run_batch_job",operator="MockBatchOperator",run_type="manual"}
1
airflow_task_state_store_write_total{...} 1
airflow_resumable_job_fresh_submit_total{operator="MockBatchOperator"} 1
```
Two separate counter families at play here, no shared names, nothing double
counted. The mixin keeps reporting the resume decision; the accessor
independently reports the calls underneath it.
---
* Read the **[Pull Request
Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)**
for more information. Note: commit author/co-author name and email in commits
become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal
([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals))
is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License
Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment:
`{pr_number}.significant.rst`, in
[airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
You can add this file in a follow-up commit after the PR is created so you
know the PR number.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
