This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch v3-4-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-4-test by this push:
new 602730282aa [v3-4-test] Add metrics for task and asset state store
access (#75729) (#75735)
602730282aa is described below
commit 602730282aadc666ccd46434d2edcd4c2fb5a138
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Tue Oct 27 10:22:08 2026 +0000
[v3-4-test] Add metrics for task and asset state store access (#75729)
(#75735)
(cherry picked from commit d340be05a4e2abc90a0b2014afc69135d48304ff)
Co-authored-by: Amogh Desai <[email protected]>
---
.../observability/metrics/metrics_template.yaml | 72 ++++++++++++++++
task-sdk/src/airflow/sdk/execution_time/context.py | 50 +++++++++--
.../src/airflow/sdk/execution_time/task_runner.py | 4 +-
.../tests/task_sdk/execution_time/test_context.py | 97 ++++++++++++++++++++++
.../task_sdk/execution_time/test_task_runner.py | 7 ++
5 files changed, 221 insertions(+), 9 deletions(-)
diff --git
a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
index 1b9bf737377..0150a550efe 100644
---
a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
+++
b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
@@ -88,6 +88,78 @@ metrics:
legacy_name: "-"
name_variables: []
+ - name: "task_state_store.read_found"
+ description: "Number of times a task read a key from the task state store
and found a value.
+ Counts each call, so one task may contribute several.
+ Metric with operator, dag_id and task_id tagging."
+ type: "counter"
+ legacy_name: "-"
+ name_variables: []
+
+ - name: "task_state_store.read_absent"
+ description: "Number of times a task read a key from the task state store
and found nothing.
+ Counts each call, so one task may contribute several.
+ Metric with operator, dag_id and task_id tagging."
+ type: "counter"
+ legacy_name: "-"
+ name_variables: []
+
+ - name: "task_state_store.write"
+ description: "Number of times a task wrote a key to the task state store.
+ Counts each call, so one task may contribute several.
+ Metric with operator, dag_id and task_id tagging."
+ type: "counter"
+ legacy_name: "-"
+ name_variables: []
+
+ - name: "task_state_store.delete"
+ description: "Number of times a task deleted a single key from the task
state store.
+ Counts each call, so one task may contribute several.
+ Metric with operator, dag_id and task_id tagging."
+ type: "counter"
+ legacy_name: "-"
+ name_variables: []
+
+ - name: "task_state_store.clear"
+ description: "Number of times a task cleared every key in its task state
store scope.
+ Counts each call, so one task may contribute several.
+ Metric with operator, dag_id and task_id tagging."
+ type: "counter"
+ legacy_name: "-"
+ name_variables: []
+
+ - name: "asset_state_store.read_found"
+ description: "Number of times a task read a key from the asset state store
and found a value.
+ Counts each call, so one task may contribute several.
+ Metric with operator, dag_id and task_id tagging."
+ type: "counter"
+ legacy_name: "-"
+ name_variables: []
+
+ - name: "asset_state_store.read_absent"
+ description: "Number of times a task read a key from the asset state store
and found nothing.
+ Counts each call, so one task may contribute several.
+ Metric with operator, dag_id and task_id tagging."
+ type: "counter"
+ legacy_name: "-"
+ name_variables: []
+
+ - name: "asset_state_store.write"
+ description: "Number of times a task wrote a key to the asset state store.
+ Counts each call, so one task may contribute several.
+ Metric with operator, dag_id and task_id tagging."
+ type: "counter"
+ legacy_name: "-"
+ name_variables: []
+
+ - name: "asset_state_store.delete"
+ description: "Number of times a task deleted a single key from the asset
state store.
+ Counts each call, so one task may contribute several.
+ Metric with operator, dag_id and task_id tagging."
+ type: "counter"
+ legacy_name: "-"
+ name_variables: []
+
- name: "ti_failures"
description: "Overall task instances failures. Metric with dag_id and
task_id tagging."
type: "counter"
diff --git a/task-sdk/src/airflow/sdk/execution_time/context.py
b/task-sdk/src/airflow/sdk/execution_time/context.py
index 02be651e9e2..b7e69b41a4f 100644
--- a/task-sdk/src/airflow/sdk/execution_time/context.py
+++ b/task-sdk/src/airflow/sdk/execution_time/context.py
@@ -32,6 +32,7 @@ from uuid import UUID
import attrs
import structlog
+from airflow.sdk._shared.observability.metrics import stats
from airflow.sdk._shared.state import AssetScope
from airflow.sdk.configuration import conf
from airflow.sdk.definitions._internal.contextmanager import _CURRENT_CONTEXT
@@ -708,9 +709,13 @@ def _get_worker_state_store_backend() -> BaseStoreBackend
| None:
class TaskStateStoreAccessor:
"""Accessor for task store scoped to the current task instance. Available
as ``context['task_state_store']`` at task execution time."""
- def __init__(self, ti_id: UUID, scope: TaskScope) -> None:
+ def __init__(self, ti_id: UUID, scope: TaskScope, stats_tags: dict[str,
str] | None = None) -> None:
self._ti_id = ti_id
self._scope = scope
+ self._stats_tags = stats_tags or {}
+
+ def _incr(self, event: str) -> None:
+ stats.incr(f"task_state_store.{event}", tags=self._stats_tags)
def __eq__(self, other: object) -> bool:
if not isinstance(other, TaskStateStoreAccessor):
@@ -750,6 +755,7 @@ class TaskStateStoreAccessor:
def _extract_get_response(self, resp: Any, key: str, default: JsonValue)
-> JsonValue:
if isinstance(resp, ErrorResponse) and resp.error !=
ErrorType.TASK_STORE_NOT_FOUND:
raise AirflowRuntimeError(resp)
+ self._incr("read_found" if isinstance(resp, TaskStateStoreResult) else
"read_absent")
if isinstance(resp, TaskStateStoreResult):
stored = resp.value
backend = _get_worker_state_store_backend()
@@ -781,12 +787,14 @@ class TaskStateStoreAccessor:
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
SUPERVISOR_COMMS.send(self._build_set_message(key, value, retention))
+ self._incr("write")
async def aset(self, key: str, value: JsonValue, *, retention: timedelta |
None = None) -> None:
"""Async version of :meth:`set` that awaits instead of blocking the
event loop."""
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
await SUPERVISOR_COMMS.asend(self._build_set_message(key, value,
retention))
+ self._incr("write")
def _build_set_message(
self, key: str, value: JsonValue, retention: timedelta | None
@@ -840,6 +848,7 @@ class TaskStateStoreAccessor:
# cleanup the DB ref first, if backend cleanup fails after this, the
ref is gone and
# deterministic keys are recoverable on next set().
SUPERVISOR_COMMS.send(DeleteTaskStateStore(ti_id=self._ti_id, key=key))
+ self._incr("delete")
backend = _get_worker_state_store_backend()
if backend is not None:
backend.delete(self._scope, key)
@@ -849,6 +858,7 @@ class TaskStateStoreAccessor:
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
await SUPERVISOR_COMMS.asend(DeleteTaskStateStore(ti_id=self._ti_id,
key=key))
+ self._incr("delete")
backend = _get_worker_state_store_backend()
if backend is not None:
await backend.adelete(self._scope, key)
@@ -860,6 +870,7 @@ class TaskStateStoreAccessor:
# cleanup the DB ref first, if backend cleanup fails after this, the
ref is gone and
# deterministic keys are recoverable on next set().
SUPERVISOR_COMMS.send(ClearTaskStateStore(ti_id=self._ti_id))
+ self._incr("clear")
backend = _get_worker_state_store_backend()
if backend is not None:
backend.clear(self._scope)
@@ -869,6 +880,7 @@ class TaskStateStoreAccessor:
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
await SUPERVISOR_COMMS.asend(ClearTaskStateStore(ti_id=self._ti_id))
+ self._incr("clear")
backend = _get_worker_state_store_backend()
if backend is not None:
await backend.aclear(self._scope)
@@ -893,11 +905,17 @@ class AssetStateStoreAccessor:
tasks, directly as ``context['asset_state_store']``.
"""
- def __init__(self, *, name: str | None = None, uri: str | None = None) ->
None:
+ def __init__(
+ self, *, name: str | None = None, uri: str | None = None, stats_tags:
dict[str, str] | None = None
+ ) -> None:
if not name and not uri:
raise ValueError("Either `name` or `uri` must be provided")
self._name = name
self._uri = uri
+ self._stats_tags = stats_tags or {}
+
+ def _incr(self, event: str) -> None:
+ stats.incr(f"asset_state_store.{event}", tags=self._stats_tags)
def __eq__(self, other: object) -> bool:
if not isinstance(other, AssetStateStoreAccessor):
@@ -937,6 +955,7 @@ class AssetStateStoreAccessor:
def _extract_get_response(self, resp: Any, key: str, default: JsonValue)
-> JsonValue:
if isinstance(resp, ErrorResponse) and resp.error !=
ErrorType.ASSET_STORE_NOT_FOUND:
raise AirflowRuntimeError(resp)
+ self._incr("read_found" if isinstance(resp, AssetStateStoreResult)
else "read_absent")
if isinstance(resp, AssetStateStoreResult):
stored = resp.value
backend = _get_worker_state_store_backend()
@@ -959,12 +978,14 @@ class AssetStateStoreAccessor:
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
SUPERVISOR_COMMS.send(self._build_set_message(key, value))
+ self._incr("write")
async def aset(self, key: str, value: JsonValue) -> None:
"""Async version of `set` that awaits instead of blocking the event
loop."""
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
await SUPERVISOR_COMMS.asend(self._build_set_message(key, value))
+ self._incr("write")
def _build_set_message(self, key: str, value: JsonValue) -> ToSupervisor:
if value is None:
@@ -1004,6 +1025,7 @@ class AssetStateStoreAccessor:
# DB ref first: if backend cleanup fails after this, the ref is gone
and
# deterministic keys are recoverable on next set().
SUPERVISOR_COMMS.send(self._build_delete_message(key))
+ self._incr("delete")
backend = _get_worker_state_store_backend()
if backend is not None:
backend.delete(AssetScope(name=self._name, uri=self._uri), key)
@@ -1013,6 +1035,7 @@ class AssetStateStoreAccessor:
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
await SUPERVISOR_COMMS.asend(self._build_delete_message(key))
+ self._incr("delete")
backend = _get_worker_state_store_backend()
if backend is not None:
await backend.adelete(AssetScope(name=self._name, uri=self._uri),
key)
@@ -1065,29 +1088,40 @@ class AssetStateStoreAccessors:
``adelete``, ``aclear``) can be called directly without subscripting.
"""
- def __init__(self, inlets: list, outlets: list | None = None) -> None:
+ def __init__(
+ self, inlets: list, outlets: list | None = None, stats_tags: dict[str,
str] | None = None
+ ) -> None:
self._by_name: dict[str, AssetStateStoreAccessor] = {}
self._by_uri: dict[str, AssetStateStoreAccessor] = {}
+ self._stats_tags = stats_tags
for inlet in inlets:
if isinstance(inlet, (Asset, AssetNameRef)):
- self._by_name[inlet.name] =
AssetStateStoreAccessor(name=inlet.name)
+ self._by_name[inlet.name] = AssetStateStoreAccessor(
+ name=inlet.name, stats_tags=self._stats_tags
+ )
elif isinstance(inlet, AssetUriRef):
- self._by_uri[inlet.uri] =
AssetStateStoreAccessor(uri=inlet.uri)
+ self._by_uri[inlet.uri] =
AssetStateStoreAccessor(uri=inlet.uri, stats_tags=self._stats_tags)
elif isinstance(inlet, AssetAlias):
from airflow.sdk.execution_time.task_runner import
SUPERVISOR_COMMS
resp =
SUPERVISOR_COMMS.send(GetAssetsByAlias(alias_name=inlet.name))
if isinstance(resp, AssetsByAliasResult):
for asset in resp.assets:
- self._by_name[asset.name] =
AssetStateStoreAccessor(name=asset.name)
+ self._by_name[asset.name] = AssetStateStoreAccessor(
+ name=asset.name, stats_tags=self._stats_tags
+ )
for outlet in outlets or []:
# AssetAlias outlets are for dynamic event emission, not state
access, so skip them
if isinstance(outlet, (Asset, AssetNameRef)) and outlet.name not
in self._by_name:
- self._by_name[outlet.name] =
AssetStateStoreAccessor(name=outlet.name)
+ self._by_name[outlet.name] = AssetStateStoreAccessor(
+ name=outlet.name, stats_tags=self._stats_tags
+ )
elif isinstance(outlet, AssetUriRef) and outlet.uri not in
self._by_uri:
- self._by_uri[outlet.uri] =
AssetStateStoreAccessor(uri=outlet.uri)
+ self._by_uri[outlet.uri] = AssetStateStoreAccessor(
+ uri=outlet.uri, stats_tags=self._stats_tags
+ )
self._total = len(self._by_name) + len(self._by_uri)
diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
index 811ae4542d0..71607ce353e 100644
--- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -339,6 +339,7 @@ class RuntimeTaskInstance(TaskInstance):
# Cache the context object, which ensures that all calls to
get_template_context
# are operating on the same context object.
if self._cached_template_context is None:
+ state_store_stats_tags = {**self.stats_tags, "operator":
self.task.task_type}
self._cached_template_context = {
# From the Task Execution interface
"dag": self.task.dag,
@@ -362,6 +363,7 @@ class RuntimeTaskInstance(TaskInstance):
"conn": ConnectionAccessor(),
"task_state_store": TaskStateStoreAccessor(
ti_id=self.id,
+ stats_tags=state_store_stats_tags,
scope=TaskScope(
dag_id=self.dag_id,
run_id=self.run_id,
@@ -374,7 +376,7 @@ class RuntimeTaskInstance(TaskInstance):
_asset_types = (Asset, AssetNameRef, AssetUriRef, AssetAlias)
if any(isinstance(i, _asset_types) for i in self.task.inlets +
self.task.outlets):
self._cached_template_context["asset_state_store"] =
AssetStateStoreAccessors(
- self.task.inlets, self.task.outlets
+ self.task.inlets, self.task.outlets,
stats_tags=state_store_stats_tags
)
# AssetAlias inlets are resolved to their concrete assets at
context build time via
# GetAssetsByAlias comms. If an alias maps to no active
assets, it doesn't contribute to
diff --git a/task-sdk/tests/task_sdk/execution_time/test_context.py
b/task-sdk/tests/task_sdk/execution_time/test_context.py
index 2d96f03e948..0f0cc999173 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_context.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_context.py
@@ -2441,6 +2441,103 @@ class TestAssetStateStoreAccessor:
backend.clear.assert_not_called()
+class TestStateStoreMetrics:
+ PATCH = "airflow.sdk._shared.observability.metrics.stats.incr"
+ TI_ID = UUID("01900000-0000-0000-0000-000000000001")
+ SCOPE = TaskScope(dag_id="dag", run_id="run", task_id="task")
+ TAGS = {"dag_id": "dag", "task_id": "task", "operator": "PythonOperator"}
+ ASSET_NAME = "my_asset"
+
+ def task_store(self):
+ return TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE,
stats_tags=self.TAGS)
+
+ def asset_store(self):
+ return AssetStateStoreAccessor(name=self.ASSET_NAME,
stats_tags=self.TAGS)
+
+ @pytest.mark.parametrize(
+ ("response", "expected"),
+ [
+ pytest.param(TaskStateStoreResult(value="job-001"), "read_found",
id="found"),
+ pytest.param(
+ ErrorResponse(error=ErrorType.TASK_STORE_NOT_FOUND,
detail={"key": "job_id"}),
+ "read_absent",
+ id="absent",
+ ),
+ ],
+ )
+ def test_task_store_read(self, mock_supervisor_comms, response, expected):
+ mock_supervisor_comms.send.return_value = response
+ with patch(self.PATCH) as incr:
+ self.task_store().get("job_id")
+ incr.assert_called_once_with(f"task_state_store.{expected}",
tags=self.TAGS)
+
+ @pytest.mark.parametrize(
+ ("response", "expected"),
+ [
+ pytest.param(AssetStateStoreResult(value="2026-01-01"),
"read_found", id="found"),
+ pytest.param(
+ ErrorResponse(error=ErrorType.ASSET_STORE_NOT_FOUND,
detail={"key": "watermark"}),
+ "read_absent",
+ id="absent",
+ ),
+ ],
+ )
+ def test_asset_store_read(self, mock_supervisor_comms, response, expected):
+ mock_supervisor_comms.send.return_value = response
+ with patch(self.PATCH) as incr:
+ self.asset_store().get("watermark")
+ incr.assert_called_once_with(f"asset_state_store.{expected}",
tags=self.TAGS)
+
+ @pytest.mark.parametrize(
+ ("call", "expected"),
+ [
+ pytest.param(lambda s: s.set("job_id", "v"),
"task_state_store.write", id="write"),
+ pytest.param(lambda s: s.delete("job_id"),
"task_state_store.delete", id="delete"),
+ pytest.param(lambda s: s.clear(), "task_state_store.clear",
id="clear"),
+ ],
+ )
+ def test_task_store_writes(self, mock_supervisor_comms, call, expected):
+ mock_supervisor_comms.send.return_value = OKResponse(ok=True)
+ with patch(self.PATCH) as incr:
+ call(self.task_store())
+ incr.assert_called_once_with(expected, tags=self.TAGS)
+
+ @pytest.mark.parametrize(
+ ("call", "expected"),
+ [
+ pytest.param(lambda s: s.set("watermark", "v"),
"asset_state_store.write", id="write"),
+ pytest.param(lambda s: s.delete("watermark"),
"asset_state_store.delete", id="delete"),
+ ],
+ )
+ def test_asset_store_writes(self, mock_supervisor_comms, call, expected):
+ mock_supervisor_comms.send.return_value = OKResponse(ok=True)
+ with patch(self.PATCH) as incr:
+ call(self.asset_store())
+ incr.assert_called_once_with(expected, tags=self.TAGS)
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ ("call", "expected"),
+ [
+ pytest.param(lambda s: s.aget("job_id"),
"task_state_store.read_found", id="read"),
+ pytest.param(lambda s: s.aset("job_id", "v"),
"task_state_store.write", id="write"),
+ pytest.param(lambda s: s.adelete("job_id"),
"task_state_store.delete", id="delete"),
+ pytest.param(lambda s: s.aclear(), "task_state_store.clear",
id="clear"),
+ ],
+ )
+ async def test_task_store_async(self, mock_supervisor_comms, call,
expected):
+ mock_supervisor_comms.asend.return_value =
TaskStateStoreResult(value="job-001")
+ with patch(self.PATCH) as incr:
+ await call(self.task_store())
+ incr.assert_called_once_with(expected, tags=self.TAGS)
+
+ def test_no_tags_when_not_supplied(self, mock_supervisor_comms):
+ mock_supervisor_comms.send.return_value =
TaskStateStoreResult(value="job-001")
+ with patch(self.PATCH) as incr:
+ TaskStateStoreAccessor(ti_id=self.TI_ID,
scope=self.SCOPE).get("job_id")
+ incr.assert_called_once_with("task_state_store.read_found", tags={})
+
+
class TestAssetStateStoreAccessors:
ASSET_NAME = "my_asset"
ASSET_URI = "s3://bucket/key"
diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
index b8d9a555376..0c44da0635a 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
@@ -2943,6 +2943,13 @@ class TestRuntimeTaskInstance:
"ti": runtime_ti,
}
+ def test_state_store_accessor_carries_stats_tags(self, create_runtime_ti):
+ runtime_ti = create_runtime_ti(task=BaseOperator(task_id="hello"))
+
+ accessor = runtime_ti.get_template_context()["task_state_store"]
+
+ assert accessor._stats_tags == {**runtime_ti.stats_tags, "operator":
"BaseOperator"}
+
def test_task_state_store_scope_carries_region_coordinates(self,
create_runtime_ti):
region_id = uuid7()
runtime_ti = create_runtime_ti(task=BaseOperator(task_id="hello"))
