ashb commented on code in PR #68568:
URL: https://github.com/apache/airflow/pull/68568#discussion_r3536647292
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -1374,6 +1374,12 @@ def process_executor_events(
.options(joinedload(TI.dag_run).selectinload(DagRun.created_dag_version))
.options(joinedload(TI.dag_version))
)
+ # When emitting Dag tags as metric tags, eager-load dag_model.tags so
the per-finished-task
+ # ti_failures / operator_failures / task.*_duration metrics carry them
without a per-TI lazy load.
+ if conf.getboolean("metrics", "dag_tags_in_metrics", fallback=False):
Review Comment:
Missed this first time around, but `conf.getboolean` is surprisingly
expensive (perhaps we should have added caching at the config layer, but for
this PR), and since this is called for every task instance that reaches
terminal state, please could we create a `self._dag_tags_in_metrics` which we
load in `__init__` to avoid calling conf.getboolean in scheduler hot path.
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -1374,6 +1374,12 @@ def process_executor_events(
.options(joinedload(TI.dag_run).selectinload(DagRun.created_dag_version))
.options(joinedload(TI.dag_version))
)
+ # When emitting Dag tags as metric tags, eager-load dag_model.tags so
the per-finished-task
+ # ti_failures / operator_failures / task.*_duration metrics carry them
without a per-TI lazy load.
+ if conf.getboolean("metrics", "dag_tags_in_metrics", fallback=False):
+ query = query.options(
+
joinedload(TI.dag_run).joinedload(DagRun.dag_model).selectinload(DagModel.tags)
Review Comment:
Do we actually need to join to dag table? We've got `dag_id` column on TI
already and _could_ join directly from TI.dag_id to DagTags table?
Not sure on this one, yes it's "more correct" to do it via dag, and it's a
pure pk index lookup to DagModel so maybe it's fine as is?
##########
shared/observability/tests/observability/metrics/test_stats.py:
##########
@@ -475,6 +496,12 @@ def
test_does_not_increment_counter_drops_invalid_tags(self):
)
self.statsd_client.incr.assert_called_once_with("test_stats_run.delay,key1=val1",
1, 1)
+ def test_standalone_tag_empty_value_emitted_as_true(self):
+ self.stats.incr("test_stats_run.delay", tags={"production": "",
"key1": "val1"})
+ self.statsd_client.incr.assert_called_once_with(
+ "test_stats_run.delay,production=true,key1=val1", 1, 1
+ )
Review Comment:
Could be added to previous parametrized case as this doesn't seem to be
related to Allow or Block lists?
##########
shared/observability/src/airflow_shared/observability/metrics/statsd_logger.py:
##########
@@ -52,8 +52,9 @@ def wrapper(
if stat is not None and tags is not None:
for k, v in tags.items():
if self.metric_tags_validator.test(k):
- if all(c not in [",", "="] for c in f"{v}{k}"):
- stat += f",{k}={v}"
+ v_str = "true" if v == "" else v
Review Comment:
Does it maybe make more sense to put this in build_dag_metric_tags?
##########
task-sdk/tests/task_sdk/execution_time/test_task_runner.py:
##########
@@ -6267,3 +6284,49 @@ def test_bad_declaration_is_skipped_not_fatal(self):
with patch("airflow.sdk.execution_time.task_runner.allow_class",
side_effect=ValueError("nope")):
# Must not raise -- the walk swallows per-class registration
errors.
_register_deserialization_allowed_classes(dag,
structlog.get_logger())
+
+
+def _make_dag_tagged_ti(create_runtime_ti, tags):
+ """Build a RuntimeTaskInstance whose in-memory Dag carries the given
tags."""
+ from airflow.sdk import DAG
+ from airflow.sdk.bases.operator import BaseOperator
+
+ class _NoopOperator(BaseOperator):
+ def execute(self, context):
+ return None
+
+ with DAG("tagged_dag", tags=tags):
+ task = _NoopOperator(task_id="t")
+ return create_runtime_ti(task=task)
+
+
+def test_stats_tags_dag_tags_disabled_by_default(create_runtime_ti):
+ """With the flag off (the default), dag tags must not leak into metrics."""
+ ti = _make_dag_tagged_ti(create_runtime_ti, ["env:prod", "validation"])
+ assert ti.stats_tags == {"dag_id": "tagged_dag", "task_id": "t",
"run_type": "manual"}
+
+
+@conf_vars({("metrics", "dag_tags_in_metrics"): "True"})
+def test_stats_tags_without_dag_tags(create_runtime_ti):
+ ti = _make_dag_tagged_ti(create_runtime_ti, [])
+ assert ti.stats_tags == {"dag_id": "tagged_dag", "task_id": "t",
"run_type": "manual"}
+
+
+@conf_vars({("metrics", "dag_tags_in_metrics"): "True"})
+def test_stats_tags_with_standalone_and_key_value_tags(create_runtime_ti):
+ ti = _make_dag_tagged_ti(create_runtime_ti, ["env:prod", "validation"])
+ assert ti.stats_tags == {
+ "env": "prod",
+ "validation": "",
+ "dag_id": "tagged_dag",
+ "task_id": "t",
+ "run_type": "manual",
+ }
+
+
+def test_stats_tags_run_type_is_bare_string(create_runtime_ti):
+ """run_type must be the bare DagRunType value (e.g. 'manual'), not the
enum — otherwise it
+ serializes as 'dagruntype.manual' on the wire and disagrees with the
scheduler-side tag."""
+ run_type = _make_dag_tagged_ti(create_runtime_ti,
[]).stats_tags["run_type"]
+ assert run_type == "manual"
+ assert type(run_type) is str # plain str, not a DagRunType enum member
Review Comment:
```suggestion
"run_type": "manual", # plain str, not a DagRunType enum member
}
```
Nit: No need for an extra test fn when we are already asserting it's a
string three times.
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -702,6 +736,12 @@ def get_running_dag_runs_to_examine(cls, session: Session)
-> ScalarResult[DagRu
.limit(cls.DEFAULT_DAGRUNS_TO_EXAMINE)
)
+ # When dag tags are emitted as metric tags, eagerly load
dag_model.tags so stats_tags
+ # does not fire a per-DagRun N+1 lazy load in the scheduler loop. Off
by default, so the
+ # extra load is only paid when the feature is enabled.
+ if airflow_conf.getboolean("metrics", "dag_tags_in_metrics",
fallback=False):
Review Comment:
This one too is called in scheduler hotpath, but this is a class method, so
I'm not sure where/how to avoid the per-loop conf.getboolean -- see what you
can come up with?
##########
shared/observability/tests/observability/metrics/test_stats.py:
##########
@@ -259,6 +260,26 @@ def test_decr(self):
metric="empty", sample_rate=1, value=1, tags=[]
)
+ def test_key_value_tag_emitted_with_colon(self):
+ dogstatsd = SafeDogStatsdLogger(self.dogstatsd_client,
metrics_tags=True)
+ dogstatsd.incr("my_metric", tags={"env": "prod"})
+ self.dogstatsd_client.increment.assert_called_once_with(
+ metric="my_metric", sample_rate=1, value=1, tags=["env:prod"]
+ )
+
+ def test_standalone_tag_empty_value_emitted_without_colon(self):
+ dogstatsd = SafeDogStatsdLogger(self.dogstatsd_client,
metrics_tags=True)
+ dogstatsd.incr("my_metric", tags={"production": ""})
+ self.dogstatsd_client.increment.assert_called_once_with(
+ metric="my_metric", sample_rate=1, value=1, tags=["production"]
+ )
+
+ def test_mixed_tags_standalone_and_key_value(self):
+ dogstatsd = SafeDogStatsdLogger(self.dogstatsd_client,
metrics_tags=True)
+ dogstatsd.incr("my_metric", tags={"production": "", "env": "staging"})
+ call_kwargs = self.dogstatsd_client.increment.call_args
+ assert set(call_kwargs.kwargs["tags"]) == {"production", "env:staging"}
Review Comment:
Better as a single parametrized test with three cases I think.
##########
task-sdk/tests/task_sdk/execution_time/test_task_runner.py:
##########
@@ -6267,3 +6284,49 @@ def test_bad_declaration_is_skipped_not_fatal(self):
with patch("airflow.sdk.execution_time.task_runner.allow_class",
side_effect=ValueError("nope")):
# Must not raise -- the walk swallows per-class registration
errors.
_register_deserialization_allowed_classes(dag,
structlog.get_logger())
+
+
+def _make_dag_tagged_ti(create_runtime_ti, tags):
+ """Build a RuntimeTaskInstance whose in-memory Dag carries the given
tags."""
+ from airflow.sdk import DAG
+ from airflow.sdk.bases.operator import BaseOperator
+
+ class _NoopOperator(BaseOperator):
+ def execute(self, context):
+ return None
+
+ with DAG("tagged_dag", tags=tags):
+ task = _NoopOperator(task_id="t")
Review Comment:
Nit: the operator isn't important, so I'd suggest using BaseOperator instead:
```suggestion
with DAG("tagged_dag", tags=tags):
task = task = BaseOperator(task_id="t")
```
##########
shared/observability/src/airflow_shared/observability/metrics/datadog_logger.py:
##########
@@ -68,12 +79,11 @@ def incr(
tags: dict[str, str] | None = None,
) -> None:
"""Increment stat."""
- if self.metrics_tags and isinstance(tags, dict):
- tags_list = [
- f"{key}:{value}" for key, value in tags.items() if
self.metric_tags_validator.test(key)
- ]
- else:
- tags_list = []
+ tags_list = (
+ self._build_tags_list(tags, self.metric_tags_validator)
+ if self.metrics_tags and isinstance(tags, dict)
+ else []
+ )
Review Comment:
This block appears 5 times. Worth DRYing it up? (Your call either way)
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -254,10 +255,21 @@ class RuntimeTaskInstance(TaskInstance):
@property
def stats_tags(self) -> dict[str, str]:
- """Metric tags for this task instance, including team_name when
available."""
- tags: dict[str, str] = {"dag_id": self.dag_id, "task_id": self.task_id}
- if self._ti_context_from_server and
self._ti_context_from_server.dag_run.team_name:
- tags["team_name"] = self._ti_context_from_server.dag_run.team_name
+ """Metric tags for this task instance, including dag tags and
team_name when available."""
+ tags: dict[str, str] = {}
+ if conf.getboolean("metrics", "dag_tags_in_metrics", fallback=False):
+ tags.update(build_dag_metric_tags(self.task.dag.tags))
+ # Built-in keys always win on collision.
+ tags["dag_id"] = self.dag_id
+ tags["task_id"] = self.task_id
Review Comment:
Nit:
```suggestion
tags.update(dag_id=self.dag_id, task_id=self.task_id)
```
--
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]
