amoghrajesh opened a new pull request, #75721:
URL: https://github.com/apache/airflow/pull/75721
<!-- 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".
-->
- [x] Yes - claude sonnet
<!--
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
The five `resumable_job.*` counters ship with a single `operator` tag. That
answers "how often did durable execution avoid redoing work, for which
operator", but not "for which Dag". Per-Dag is the question people ask when
they want to know whether the feature is earning its keep on a specific
pipeline, and today it cannot be answered from these metrics at all.
### Current behaviour
Every `resumable_job.*` counter carries `operator`, plus `team_name` on
multi-team deployments. Nothing identifies the Dag or task that produced the
count, so all executions of every Dag using a given operator collapse into one
series.
### Proposed change
Take the whole of `RuntimeTaskInstance.stats_tags` instead of pulling one
key out of it. That property already assembles the label set every other task
metric uses, so the counters now carry `dag_id`, `task_id` and `run_type`
alongside `operator`, with `team_name` continuing to appear where it applies.
The metric descriptions in `metrics_template.yaml` are updated to match.
### Changes of Note
`stats_tags` also includes author-chosen Dag tags when `[metrics]
dag_tags_in_metrics` is enabled. These counters now inherit that, same as
`ti.start` and `ti.finish`.
The alternative was filtering to a fixed key list (`dag_id`, `task_id`,
`run_type`, `team_name`), which bounds cardinality by construction. I kept the
consistent behaviour instead, because the config is off by default and a
deployment that turns it on has opted into that blast radius everywhere else
already. Happy to switch if reviewers would rather these five counters be
special-cased.
### User implications / backcompat
No breaking change. The labels only materialise on backends that carry tags:
OTel and dogstatsd natively, StatsD only with `[metrics]
statsd_influxdb_enabled`. That is the same constraint the existing `operator`
tag already has, not a new one, so plain StatsD deployments see no difference.
Series cardinality on these five counters increases for anyone already
receiving the tags, bounded by `dag_id` x `task_id` x `operator`.
### Testing
Using this dag:
```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),
)
```
#### Case 1: Running it normally without interruption
<img width="1650" height="809" alt="image"
src="https://github.com/user-attachments/assets/b3a01f81-a909-4e27-9680-28141f15f810";
/>
Metrics have the added tags now:
```shell
amoghrajesh.desai …/airflow
add-more-ti-stats-to-resumablemixin-counters $!? orbstack ♥ 12:21
curl -s localhost:28889/metrics | grep '^airflow_resumable_job'
airflow_resumable_job_fresh_submit_total{dag_id="example_resumable_job",instance="36b43c68-344c-4222-ad4a-3c93128716d2",job="test",operator="MockBatchOperator",otel_scope_name="airflow.sdk._shared.observability.metrics.otel_logger",otel_scope_schema_url="",otel_scope_version="",run_type="manual",task_id="run_batch_job"}
1
```
Before this change the same line read `{operator="MockBatchOperator"}`.
#### Case 2: Running with `simulate_final_status` as `SUCCEEDED`
Crashed the worker when the task waits for it
<img width="1650" height="809" alt="image"
src="https://github.com/user-attachments/assets/15f5a1a9-e1f5-4ca6-9f34-c7f21b2b469e";
/>
<img width="1650" height="809" alt="image"
src="https://github.com/user-attachments/assets/273070ad-484b-41ab-a859-2151fa3c217d";
/>
Now the metrics:
```shell
airflow_resumable_job_already_succeeded_total{dag_id="example_resumable_job",instance="36ec9a2f-9faf-475e-93b8-6590f4fc4a73",job="test",operator="MockBatchOperator",otel_scope_name="airflow.sdk._shared.observability.metrics.otel_logger",otel_scope_schema_url="",otel_scope_version="",run_type="manual",task_id="run_batch_job"}
1
airflow_resumable_job_fresh_submit_total{dag_id="example_resumable_job",instance="36b43c68-344c-4222-ad4a-3c93128716d2",job="test",operator="MockBatchOperator",otel_scope_name="airflow.sdk._shared.observability.metrics.otel_logger",otel_scope_schema_url="",otel_scope_version="",run_type="manual",task_id="run_batch_job"}
1
airflow_resumable_job_reconnect_attempt_total{dag_id="example_resumable_job",instance="36ec9a2f-9faf-475e-93b8-6590f4fc4a73",job="test",operator="MockBatchOperator",otel_scope_name="airflow.sdk._shared.observability.metrics.otel_logger",otel_scope_schema_url="",otel_scope_version="",run_type="manual",task_id="run_batch_job"}
1
```
---
* 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]
