diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index af11f9fe701d1..50c05084f3c96 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -1280,6 +1280,17 @@ metrics: type: boolean example: ~ default: "False" + dag_tags_in_metrics: + description: | + Set to ``True`` to include Dag tags as metric tags on all Dag-run and task-instance metrics. + Tags that contain a colon (e.g. ``env:prod``) are split into a key/value pair. Plain tags + (e.g. ``production``) are emitted as standalone DogStatsd tags, or as ``tag=true`` in InfluxDB + line-protocol format. Disabled by default to avoid unexpected cardinality increases, since Dag + tags are free-form, user-defined strings. + version_added: 3.4.0 + type: boolean + example: ~ + default: "False" otel_on: description: | Enables sending metrics to OpenTelemetry. diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index f6d6a4787f629..cf7c952075ac6 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -334,6 +334,7 @@ def __init__( self._scheduler_use_job_schedule = conf.getboolean("scheduler", "use_job_schedule", fallback=True) self._parallelism = conf.getint("core", "parallelism") self._multi_team = conf.getboolean("core", "multi_team") + self._dag_tags_in_metrics = conf.getboolean("metrics", "dag_tags_in_metrics", fallback=False) self._max_partition_dag_runs_per_loop = MAX_PARTITION_DAG_RUNS_PER_LOOP self._dag_id_to_team_name: dict[str, str | None] = {} @@ -736,10 +737,10 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) - list(unique_dag_ids), ) for ti in task_instances_to_examine: - # Set team as a transient attribute; team lives on the Bundle, not - # on the TI/DagRun schema, so we resolve it at scheduling time. + # Team lives on the Bundle, not the TI/DagRun schema, so resolve it at scheduling + # time and stash it on the dag run, where stats_tags reads it for metric tagging. if team := dag_id_to_team_name.get(ti.dag_id): - ti._team_name = team + ti.dag_run._team_name = team executor_slots_available: dict[ExecutorName, int] = {} # First get a mapping of executor names to slots they have available @@ -1260,6 +1261,7 @@ def _process_executor_events(self, executor: BaseExecutor, session: Session) -> job_id=self.job.id, scheduler_dag_bag=self.scheduler_dag_bag, session=session, + eagerly_load_dag_tags=self._dag_tags_in_metrics, ) except Exception as exc: stats.incr("scheduler.executor_events.failed", tags={"exception_class": type(exc).__name__}) @@ -1272,7 +1274,12 @@ def _emit_executor_events_batch_metrics(num_events: int) -> None: @classmethod def process_executor_events( - cls, executor: BaseExecutor, job_id: int | None, scheduler_dag_bag: DBDagBag, session: Session + cls, + executor: BaseExecutor, + job_id: int | None, + scheduler_dag_bag: DBDagBag, + session: Session, + eagerly_load_dag_tags: bool = False, ) -> int: """ Process task completion events from the executor and update task instance states. @@ -1295,6 +1302,9 @@ def process_executor_events( :param job_id: The scheduler job ID, used to detect task requeuing by other schedulers :param scheduler_dag_bag: Serialized DAG bag for retrieving task definitions :param session: Database session for task instance updates + :param eagerly_load_dag_tags: When True, eager-load dag_model.tags so the per-finished-task + metrics carry Dag tags without a per-TI lazy load. The scheduler passes its cached flag so + the hot path never reads conf; other callers (e.g. ``dag.test()``) leave it at the default. :return: Number of events processed from the executor event buffer @@ -1386,6 +1396,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. + # TI already joins DagModel by dag_id, so warm tags off that relationship directly rather than + # via the dag_run hop; the DagModel is shared in the identity map, so dag_run.dag_model.tags is free. + if eagerly_load_dag_tags: + query = query.options(selectinload(TI.dag_model).selectinload(DagModel.tags)) # row lock this entire set of taskinstances to make sure the scheduler doesn't fail when we have # multi-schedulers locked_query = with_row_locks(query, of=TI, session=session, skip_locked=True) @@ -1922,7 +1938,11 @@ def _do_scheduling(self, session: Session) -> int: # examining, rather than making one query per DagRun. # Materialize into a list because the multi-team block below iterates # the result and ScalarResult is a one-pass iterator. - dag_runs = list(DagRun.get_running_dag_runs_to_examine(session=session)) + dag_runs = list( + DagRun.get_running_dag_runs_to_examine( + session=session, eagerly_load_dag_tags=self._dag_tags_in_metrics + ) + ) if self._multi_team and dag_runs: unique_dag_ids = {dr.dag_id for dr in dag_runs} @@ -3553,9 +3573,6 @@ def _purge_task_instances_without_heartbeats( if self._multi_team: unique_dag_ids = {ti.dag_id for ti in task_instances_without_heartbeats} dag_id_to_team_name = self._get_team_names_for_dag_ids(unique_dag_ids, session) - for ti in task_instances_without_heartbeats: - if team := dag_id_to_team_name.get(ti.dag_id): - ti._team_name = team else: dag_id_to_team_name = {} diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index a51f16e51f603..d6610b3663e7e 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/airflow-core/src/airflow/models/dagrun.py @@ -53,16 +53,25 @@ update, ) from sqlalchemy.dialects import postgresql -from sqlalchemy.exc import IntegrityError +from sqlalchemy.exc import IntegrityError, SQLAlchemyError from sqlalchemy.ext.associationproxy import association_proxy from sqlalchemy.ext.hybrid import hybrid_property from sqlalchemy.ext.mutable import MutableDict -from sqlalchemy.orm import Mapped, declared_attr, joinedload, mapped_column, relationship, synonym, validates +from sqlalchemy.orm import ( + Mapped, + declared_attr, + joinedload, + mapped_column, + relationship, + synonym, + validates, +) from sqlalchemy.orm.exc import StaleDataError from sqlalchemy.sql.expression import false, select from sqlalchemy.sql.functions import coalesce from airflow._shared.observability.metrics import stats +from airflow._shared.observability.metrics.stats import build_dag_metric_tags from airflow._shared.observability.traces import ( DAGRUN_PARENT_TRACE_CONTEXT_KEY, TASK_SPAN_DETAIL_LEVEL_KEY, @@ -576,11 +585,36 @@ def check_version_id_exists_in_dr(self, dag_version_id: UUID, *, session: Sessio ) return session.scalar(select_stmt) + def dag_tags_for_stats(self) -> dict[str, str]: + """Convert dag tags to metric tags. Tags with ':' become key:value; others are standalone (empty value).""" + if not airflow_conf.getboolean("metrics", "dag_tags_in_metrics", fallback=False): + return {} + try: + # Lazy-loads dag_model.tags if not already loaded. The scheduler hot loop + # (get_running_dag_runs_to_examine) eager-loads these to avoid a per-DagRun query; other, + # low-frequency emission paths fall back to this lazy load. On a detached/expired DagRun + # the load raises — swallow it so metric tagging never breaks the caller. + if not self.dag_model or not self.dag_model.tags: + return {} + return build_dag_metric_tags(tag.name for tag in self.dag_model.tags) + except SQLAlchemyError: + return {} + @property def stats_tags(self) -> dict[str, str]: - return prune_dict( - {"dag_id": self.dag_id, "run_type": self.run_type, "team_name": getattr(self, "_team_name", None)} + # prune_dict strips falsy values, so merge dag tags after it runs so standalone + # tags (empty value) are preserved for DogStatsD emission. + base = prune_dict( + { + "dag_id": self.dag_id, + # bare value so it serializes as e.g. "scheduled", not "dagruntype.scheduled" + "run_type": getattr(self.run_type, "value", self.run_type), + "team_name": getattr(self, "_team_name", None), + } ) + dag_tags = self.dag_tags_for_stats() + # Built-in keys win on collision; dag tags fill in everything else. + return {**dag_tags, **base} def get_state(self): return self._state @@ -702,7 +736,9 @@ def active_runs_of_dags( @classmethod @retry_db_transaction - def get_running_dag_runs_to_examine(cls, session: Session) -> ScalarResult[DagRun]: + def get_running_dag_runs_to_examine( + cls, *, session: Session, eagerly_load_dag_tags: bool + ) -> ScalarResult[DagRun]: """ Return the next DagRuns that the scheduler should attempt to schedule. @@ -733,6 +769,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. The caller owns the feature decision; + # the scheduler passes its cached flag so the loop never reads conf. + if eagerly_load_dag_tags: + query = query.options(joinedload(cls.dag_model).selectinload(DagModel.tags)) + query = query.where(DagRun.run_after <= func.now()) result = session.scalars(with_row_locks(query, of=cls, session=session, skip_locked=True)).unique() diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index 1c7de2185e607..0dd10507fc863 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/airflow-core/src/airflow/models/taskinstance.py @@ -100,7 +100,6 @@ from airflow.ti_deps.dep_context import DepContext from airflow.ti_deps.dependencies_deps import REQUEUEABLE_DEPS, RUNNING_DEPS from airflow.ti_deps.deps.ready_to_reschedule import ReadyToRescheduleDep -from airflow.utils.helpers import prune_dict from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.net import get_hostname from airflow.utils.platform import getuser @@ -757,9 +756,8 @@ def __hash__(self): @property def stats_tags(self) -> dict[str, str]: """Returns task instance tags.""" - return prune_dict( - {"dag_id": self.dag_id, "task_id": self.task_id, "team_name": getattr(self, "_team_name", None)} - ) + # Reuse the dag run's tags and add the task-level ones. + return {**self.dag_run.stats_tags, "task_id": self.task_id} @staticmethod def insert_mapping( diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 1f3b6992c3e1c..3c0664db66aa6 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -518,8 +518,13 @@ def test_process_executor_events(self, mock_get_backend, mock_task_callback, dag "scheduler.tasks.killed_externally", tags={"dag_id": dag_id, "task_id": ti1.task_id}, ), - mock.call("operator_failures_EmptyOperator", tags={"dag_id": dag_id, "task_id": ti1.task_id}), - mock.call("ti_failures", tags={"dag_id": dag_id, "task_id": ti1.task_id}), + mock.call( + "operator_failures_EmptyOperator", + tags={"dag_id": dag_id, "task_id": ti1.task_id, "run_type": "manual"}, + ), + mock.call( + "ti_failures", tags={"dag_id": dag_id, "task_id": ti1.task_id, "run_type": "manual"} + ), ], any_order=True, ) @@ -645,8 +650,11 @@ def test_process_executor_events_with_no_callback(self, mock_get_backend, mock_t "scheduler.tasks.killed_externally", tags={"dag_id": dag_id, "task_id": task_id}, ), - mock.call("operator_failures_EmptyOperator", tags={"dag_id": dag_id, "task_id": task_id}), - mock.call("ti_failures", tags={"dag_id": dag_id, "task_id": task_id}), + mock.call( + "operator_failures_EmptyOperator", + tags={"dag_id": dag_id, "task_id": task_id, "run_type": "manual"}, + ), + mock.call("ti_failures", tags={"dag_id": dag_id, "task_id": task_id, "run_type": "manual"}), ], any_order=True, ) @@ -9793,8 +9801,10 @@ def test_multi_team_config_disabled_uses_legacy_behavior(self, dag_maker, mock_e assert result2 == mock_executors[1] # Matched by executor name @conf_vars({("core", "multi_team"): "true"}) - def test_multi_team_sets_team_name_on_task_instances(self, dag_maker, mock_executors, session): - """Test that _team_name is set on TaskInstance objects during the scheduling loop.""" + @mock.patch("airflow._shared.observability.metrics.stats.timing") + def test_multi_team_sets_team_name_on_task_instances(self, mock_timing, dag_maker, session): + """The scheduling loop resolves the bundle's team onto the dag run, so the QUEUED + state-change metric (emitted via TaskInstance.stats_tags) carries team_name.""" clear_db_teams() clear_db_dag_bundles() @@ -9813,19 +9823,21 @@ def test_multi_team_sets_team_name_on_task_instances(self, dag_maker, mock_execu dr = dag_maker.create_dagrun() ti = dr.get_task_instance("task_a", session=session) ti.state = State.SCHEDULED + ti.scheduled_dttm = timezone.utcnow() + session.merge(ti) session.flush() scheduler_job = Job() self.job_runner = SchedulerJobRunner(job=scheduler_job) - self.job_runner._multi_team = True - # Simulate what _executable_task_instances_to_queued does - dag_id_to_team_name = self.job_runner._get_team_names_for_dag_ids(["dag_a"], session) - if team_name := dag_id_to_team_name.get(ti.dag_id): - ti._team_name = team_name + queued_tis = self.job_runner._executable_task_instances_to_queued(max_tis=32, session=session) - assert ti._team_name == "team_a" - assert ti.stats_tags == {"dag_id": "dag_a", "task_id": "task_a", "team_name": "team_a"} + assert {t.key for t in queued_tis} == {ti.key} + scheduled_calls = [ + c for c in mock_timing.call_args_list if c.args and c.args[0] == "task.scheduled_duration" + ] + assert scheduled_calls, "expected a task.scheduled_duration metric on QUEUED transition" + assert scheduled_calls[0].kwargs["tags"]["team_name"] == "team_a" @conf_vars({("core", "multi_team"): "true"}) def test_do_scheduling_multi_team_schedules_task_instances(self, dag_maker, session): diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py index 64b4947ec4c6f..c46c30d1194e6 100644 --- a/airflow-core/tests/unit/models/test_dagrun.py +++ b/airflow-core/tests/unit/models/test_dagrun.py @@ -21,7 +21,7 @@ from collections import defaultdict from collections.abc import Mapping from contextlib import contextmanager -from functools import reduce +from functools import partial, reduce from typing import TYPE_CHECKING from unittest import mock from unittest.mock import ANY, call @@ -1086,10 +1086,12 @@ def test_next_dagruns_to_examine_only_unpaused(self, session, state, testing_dag ) if state == DagRunState.RUNNING: - func = DagRun.get_running_dag_runs_to_examine + fetch = partial( + DagRun.get_running_dag_runs_to_examine, session=session, eagerly_load_dag_tags=False + ) else: - func = DagRun.get_queued_dag_runs_to_set_running - runs = func(session).all() + fetch = partial(DagRun.get_queued_dag_runs_to_set_running, session) + runs = fetch().all() assert runs == [dr] @@ -1097,7 +1099,7 @@ def test_next_dagruns_to_examine_only_unpaused(self, session, state, testing_dag session.merge(orm_dag) session.commit() - runs = func(session).all() + runs = fetch().all() assert runs == [] @mock.patch("airflow._shared.observability.metrics.stats.timing") @@ -4563,34 +4565,115 @@ def test_context_carrier_parent_conf(self, dag_maker, conf, embeds): assert trace_id != _EXTERNAL_TRACE_ID -class TestDagRunStatsTagsTeamName: - def test_stats_tags_without_team_name(self, dag_maker): - """stats_tags should not include team_name when _team_name is not set.""" - with dag_maker("test_dag"): - EmptyOperator(task_id="t1") - dr = dag_maker.create_dagrun() - tags = dr.stats_tags - assert "team_name" not in tags - assert tags == {"dag_id": "test_dag", "run_type": "manual"} +def test_stats_tags_without_team_name(dag_maker): + """stats_tags omits team_name when _team_name is not set.""" + with dag_maker("test_dag"): + EmptyOperator(task_id="t1") + dr = dag_maker.create_dagrun() + assert dr.stats_tags == {"dag_id": "test_dag", "run_type": dr.run_type} - def test_stats_tags_with_team_name(self, dag_maker): - """stats_tags should include team_name when _team_name is set.""" - with dag_maker("test_dag"): - EmptyOperator(task_id="t1") - dr = dag_maker.create_dagrun() - dr._team_name = "my_team" - tags = dr.stats_tags - assert tags["team_name"] == "my_team" - assert tags == {"dag_id": "test_dag", "run_type": "manual", "team_name": "my_team"} - - def test_stats_tags_with_none_team_name(self, dag_maker): - """stats_tags should not include team_name when _team_name is None.""" - with dag_maker("test_dag"): - EmptyOperator(task_id="t1") - dr = dag_maker.create_dagrun() - dr._team_name = None - tags = dr.stats_tags - assert "team_name" not in tags + +def test_stats_tags_with_team_name(dag_maker): + """stats_tags includes team_name when _team_name is set.""" + with dag_maker("test_dag"): + EmptyOperator(task_id="t1") + dr = dag_maker.create_dagrun() + dr._team_name = "my_team" + assert dr.stats_tags == {"dag_id": "test_dag", "run_type": dr.run_type, "team_name": "my_team"} + + +def test_stats_tags_with_none_team_name(dag_maker): + """stats_tags omits team_name when _team_name is None.""" + with dag_maker("test_dag"): + EmptyOperator(task_id="t1") + dr = dag_maker.create_dagrun() + dr._team_name = None + assert dr.stats_tags == {"dag_id": "test_dag", "run_type": dr.run_type} + + +def test_stats_tags_dag_tags_disabled_by_default(dag_maker, session): + """With the flag off (the default), dag tags must not leak into metrics.""" + with dag_maker("disabled_tag_dag", tags=["production", "env:prod"], session=session): + pass + dr = dag_maker.create_dagrun() + _ = dr.dag_model.tags + assert dr.stats_tags == {"dag_id": "disabled_tag_dag", "run_type": dr.run_type} + + +@conf_vars({("metrics", "dag_tags_in_metrics"): "True"}) +def test_stats_tags_without_dag_tags(dag_maker, session): + with dag_maker("no_tags_dag", session=session): + pass + dr = dag_maker.create_dagrun() + _ = dr.dag_model.tags + assert dr.stats_tags == {"dag_id": "no_tags_dag", "run_type": dr.run_type} + + +@conf_vars({("metrics", "dag_tags_in_metrics"): "True"}) +def test_stats_tags_with_standalone_dag_tag(dag_maker, session): + with dag_maker("standalone_tag_dag", tags=["production"], session=session): + pass + dr = dag_maker.create_dagrun() + _ = dr.dag_model.tags # eager-load so _dag_tags_for_stats sees the tags + tags = dr.stats_tags + assert tags == {"dag_id": "standalone_tag_dag", "run_type": "manual", "production": ""} + # run_type is the bare value, not a DagRunType enum (serializes as "manual", not "dagruntype.manual") + assert type(tags["run_type"]) is str + + +@conf_vars({("metrics", "dag_tags_in_metrics"): "True"}) +def test_stats_tags_with_key_value_dag_tag(dag_maker, session): + with dag_maker("kv_tag_dag", tags=["env:staging"], session=session): + pass + dr = dag_maker.create_dagrun() + _ = dr.dag_model.tags + assert dr.stats_tags == {"dag_id": "kv_tag_dag", "run_type": dr.run_type, "env": "staging"} + + +@conf_vars({("metrics", "dag_tags_in_metrics"): "True"}) +def test_stats_tags_builtin_keys_win_on_collision(dag_maker, session): + with dag_maker("collision_dag", tags=["dag_id:sneaky"], session=session): + pass + dr = dag_maker.create_dagrun() + _ = dr.dag_model.tags + # built-in dag_id wins over the colliding "dag_id:sneaky" tag + assert dr.stats_tags == {"dag_id": "collision_dag", "run_type": dr.run_type} + + +@conf_vars({("metrics", "dag_tags_in_metrics"): "True"}) +def test_stats_tags_lazy_loads_dag_tags_when_not_eager_loaded(dag_maker, session): + """When dag_model is not eager-loaded, stats_tags lazy-loads it in-session so tags still appear.""" + from sqlalchemy import inspect as sa_inspect + + with dag_maker("lazy_tag_dag", tags=["env:prod"], session=session): + pass + dr = dag_maker.create_dagrun() + session.expire(dr, ["dag_model"]) # not eager-loaded + assert "dag_model" in sa_inspect(dr).unloaded + + # lazy fallback loads dag_model.tags in-session, so the tags are still emitted + assert dr.stats_tags == {"dag_id": "lazy_tag_dag", "run_type": dr.run_type, "env": "prod"} + + +@conf_vars({("metrics", "dag_tags_in_metrics"): "True"}) +def test_get_running_dag_runs_to_examine_eager_loads_dag_tags(dag_maker, session): + """With the flag on, the scheduler query eager-loads dag_model.tags so stats_tags fires no lazy load.""" + from sqlalchemy import inspect as sa_inspect + + with dag_maker("eager_tag_dag", tags=["env:prod"], session=session): + pass + dag_maker.create_dagrun(state=DagRunState.RUNNING) + session.commit() + + dr = next( + r + for r in DagRun.get_running_dag_runs_to_examine(session=session, eagerly_load_dag_tags=True) + if r.dag_id == "eager_tag_dag" + ) + # dag_model and its tags are already populated — no lazy load needed at metric-emission time. + assert "dag_model" not in sa_inspect(dr).unloaded + assert "tags" not in sa_inspect(dr.dag_model).unloaded + assert dr.stats_tags == {"dag_id": "eager_tag_dag", "run_type": dr.run_type, "env": "prod"} class TestClearPartitionRuns: diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py index 7f1a193d96380..e0d4a1063e935 100644 --- a/airflow-core/tests/unit/models/test_taskinstance.py +++ b/airflow-core/tests/unit/models/test_taskinstance.py @@ -2431,7 +2431,7 @@ def test_handle_failure_no_task(self, mock_get_backend, dag_maker): ti.task = None ti.state = State.QUEUED session.flush() - expected_stats_tags = {"dag_id": ti.dag_id, "task_id": ti.task_id} + expected_stats_tags = {"dag_id": ti.dag_id, "task_id": ti.task_id, "run_type": dr.run_type} assert ti.task is None, "Check critical pre-condition" @@ -4374,47 +4374,82 @@ def test_task_instance_repr_does_not_raise_for_deferred_columns(dag_maker, sessi class TestTaskInstanceStatsTagsTeamName: def test_stats_tags_without_team_name(self, dag_maker, session): - """stats_tags should not include team_name when _team_name is not set.""" + """stats_tags should not include team_name when the dag run has no team_name.""" with dag_maker("test_dag"): EmptyOperator(task_id="my_task") dr = dag_maker.create_dagrun() ti = dr.get_task_instance("my_task", session=session) tags = ti.stats_tags assert "team_name" not in tags - assert tags == {"dag_id": "test_dag", "task_id": "my_task"} + assert tags == {"dag_id": "test_dag", "task_id": "my_task", "run_type": dr.run_type} def test_stats_tags_with_team_name(self, dag_maker, session): - """stats_tags should include team_name when _team_name is set.""" + """stats_tags takes team_name from the dag run's stats_tags.""" with dag_maker("test_dag"): EmptyOperator(task_id="my_task") dr = dag_maker.create_dagrun() + dr._team_name = "my_team" ti = dr.get_task_instance("my_task", session=session) - ti._team_name = "my_team" tags = ti.stats_tags assert tags["team_name"] == "my_team" - assert tags == {"dag_id": "test_dag", "task_id": "my_task", "team_name": "my_team"} + assert tags == { + "dag_id": "test_dag", + "task_id": "my_task", + "team_name": "my_team", + "run_type": dr.run_type, + } def test_stats_tags_with_none_team_name(self, dag_maker, session): - """stats_tags should not include team_name when _team_name is None.""" + """stats_tags should not include team_name when the dag run's team_name is None.""" with dag_maker("test_dag"): EmptyOperator(task_id="my_task") dr = dag_maker.create_dagrun() + dr._team_name = None ti = dr.get_task_instance("my_task", session=session) - ti._team_name = None tags = ti.stats_tags assert "team_name" not in tags + @conf_vars({("metrics", "dag_tags_in_metrics"): "True"}) + def test_stats_tags_match_worker_tag_set(self, dag_maker, session): + """ti_failures (and other ti.* metrics) are emitted from both the worker + (RuntimeTaskInstance.stats_tags) and the scheduler (TaskInstance.stats_tags); both must produce + the same tag set. The worker side is asserted in test_task_runner.py + (test_stats_tags_with_standalone_and_key_value_tags); this guards the scheduler side against drift: + dag tags + dag_id + task_id + run_type. + """ + with dag_maker("parity_dag", tags=["env:prod", "production"], session=session): + EmptyOperator(task_id="t1") + dr = dag_maker.create_dagrun() + ti = dr.get_task_instance("t1", session=session) + tags = ti.stats_tags + assert tags == { + "dag_id": "parity_dag", + "task_id": "t1", + "run_type": "manual", + "env": "prod", + "production": "", + } + # run_type must be the bare value (not a DagRunType enum) so it serializes identically to + # the worker side, e.g. "manual" not "dagruntype.manual". + assert type(tags["run_type"]) is str + @pytest.mark.parametrize( ("team_name", "expected_tags"), [ pytest.param( "my_team", - {"dag_id": "test_dag", "task_id": "my_task", "team_name": "my_team", "queue": "default"}, + { + "dag_id": "test_dag", + "task_id": "my_task", + "team_name": "my_team", + "queue": "default", + "run_type": "manual", + }, id="with_team", ), pytest.param( None, - {"dag_id": "test_dag", "task_id": "my_task", "queue": "default"}, + {"dag_id": "test_dag", "task_id": "my_task", "queue": "default", "run_type": "manual"}, id="without_team", ), ], @@ -4430,7 +4465,7 @@ def test_emit_state_change_metric_includes_team_name( ti.state = TaskInstanceState.SCHEDULED ti.scheduled_dttm = timezone.utcnow() if team_name: - ti._team_name = team_name + dr._team_name = team_name session.merge(ti) session.flush() diff --git a/shared/observability/src/airflow_shared/observability/metrics/datadog_logger.py b/shared/observability/src/airflow_shared/observability/metrics/datadog_logger.py index 129354c3a4381..b1fda8208aa5b 100644 --- a/shared/observability/src/airflow_shared/observability/metrics/datadog_logger.py +++ b/shared/observability/src/airflow_shared/observability/metrics/datadog_logger.py @@ -58,6 +58,15 @@ def __init__( self.stat_name_handler = stat_name_handler self.statsd_influxdb_enabled = statsd_influxdb_enabled + def _build_tags_list(self, tags: dict[str, str] | None) -> list[str]: + if not (self.metrics_tags and isinstance(tags, dict)): + return [] + return [ + (f"{key}:{value}" if value != "" else key) + for key, value in tags.items() + if self.metric_tags_validator.test(key) + ] + @validate_stat def incr( self, @@ -68,12 +77,7 @@ 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) if self.metrics_validator.test(stat): return self.dogstatsd.increment(metric=stat, value=count, tags=tags_list, sample_rate=rate) return None @@ -88,12 +92,7 @@ def decr( tags: dict[str, str] | None = None, ) -> None: """Decrement 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) if self.metrics_validator.test(stat): return self.dogstatsd.decrement(metric=stat, value=count, tags=tags_list, sample_rate=rate) return None @@ -109,12 +108,7 @@ def gauge( tags: dict[str, str] | None = None, ) -> None: """Gauge 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) if self.metrics_validator.test(stat): return self.dogstatsd.gauge(metric=stat, value=value, tags=tags_list, sample_rate=rate) return None @@ -128,12 +122,7 @@ def timing( tags: dict[str, str] | None = None, ) -> None: """Stats timing.""" - 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) if self.metrics_validator.test(stat): if isinstance(dt, datetime.timedelta): dt = dt.total_seconds() * 1000.0 @@ -148,12 +137,7 @@ def timer( **kwargs, ) -> Timer: """Timer metric that can be cancelled.""" - 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) if stat and self.metrics_validator.test(stat): return Timer(self.dogstatsd.timed(stat, tags=tags_list, **kwargs)) return Timer() diff --git a/shared/observability/src/airflow_shared/observability/metrics/stats.py b/shared/observability/src/airflow_shared/observability/metrics/stats.py index 5140314922a67..7b51e580cd036 100644 --- a/shared/observability/src/airflow_shared/observability/metrics/stats.py +++ b/shared/observability/src/airflow_shared/observability/metrics/stats.py @@ -20,7 +20,7 @@ import os import re import socket -from collections.abc import Callable +from collections.abc import Callable, Iterable from typing import TYPE_CHECKING, Any from .base_stats_logger import NoStatsLogger @@ -35,6 +35,23 @@ _VALID_STAT_NAME_CHARS_RE = re.compile(r"^[a-zA-Z0-9_.-]+$") _INVALID_STAT_NAME_CHARS_RE = re.compile(r"[^a-zA-Z0-9_.-]") + +def build_dag_metric_tags(tag_names: Iterable[str]) -> dict[str, str]: + """ + Convert Dag tag strings into metric tags. + + Tags with a non-empty value after a ``:`` (e.g. ``env:prod``) split into a + ``key: value`` pair. Plain tags (e.g. ``production``) and tags with no value + after the colon (e.g. ``env:``) map to an empty string, emitted as a standalone + DogStatsd tag or as ``tag=true`` in InfluxDB line protocol. + """ + result: dict[str, str] = {} + for name in tag_names: + key, _, value = name.partition(":") + result[key] = value + return result + + # Module-level singleton state. _factory: Callable[[], StatsLogger | NoStatsLogger] | None = None _backend: StatsLogger | NoStatsLogger | None = None diff --git a/shared/observability/src/airflow_shared/observability/metrics/statsd_logger.py b/shared/observability/src/airflow_shared/observability/metrics/statsd_logger.py index 3500a04dc7be3..547f680db50f0 100644 --- a/shared/observability/src/airflow_shared/observability/metrics/statsd_logger.py +++ b/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 + if all(c not in [",", "="] for c in f"{v_str}{k}"): + stat += f",{k}={v_str}" else: log.error("Dropping invalid tag: %s=%s.", k, v) return fn(self, stat, *args, tags=tags, **kwargs) diff --git a/shared/observability/tests/observability/metrics/test_stats.py b/shared/observability/tests/observability/metrics/test_stats.py index 4121ffe9eee85..25deeb114e03a 100644 --- a/shared/observability/tests/observability/metrics/test_stats.py +++ b/shared/observability/tests/observability/metrics/test_stats.py @@ -34,6 +34,7 @@ from airflow_shared.observability.metrics import datadog_logger, statsd_logger from airflow_shared.observability.metrics.base_stats_logger import StatsLogger from airflow_shared.observability.metrics.datadog_logger import SafeDogStatsdLogger +from airflow_shared.observability.metrics.stats import build_dag_metric_tags from airflow_shared.observability.metrics.statsd_logger import SafeStatsdLogger from airflow_shared.observability.metrics.validators import ( PatternAllowListValidator, @@ -259,6 +260,20 @@ def test_decr(self): metric="empty", sample_rate=1, value=1, tags=[] ) + @pytest.mark.parametrize( + ("tags", "expected"), + [ + ({"env": "prod"}, {"env:prod"}), + ({"production": ""}, {"production"}), + ({"production": "", "env": "staging"}, {"production", "env:staging"}), + ], + ) + def test_key_value_and_standalone_tags(self, tags, expected): + dogstatsd = SafeDogStatsdLogger(self.dogstatsd_client, metrics_tags=True) + dogstatsd.incr("my_metric", tags=tags) + call_kwargs = self.dogstatsd_client.increment.call_args + assert set(call_kwargs.kwargs["tags"]) == expected + class TestStatsAllowAndBlockLists: @pytest.mark.parametrize( @@ -457,23 +472,35 @@ def test_increment_counter(self): ) self.statsd_client.incr.assert_called_once_with("test_stats_run.delay", 1, 1) - def test_increment_counter_with_tags(self): - self.stats.incr( - "test_stats_run.delay", - tags={"key0": 0, "key1": "val1", "key2": "val2"}, - ) - self.statsd_client.incr.assert_called_once_with("test_stats_run.delay,key0=0,key1=val1", 1, 1) - - def test_increment_counter_with_tags_and_forward_slash(self): - self.stats.incr("test_stats_run.dag", tags={"path": "/some/path/dag.py"}) - self.statsd_client.incr.assert_called_once_with("test_stats_run.dag,path=/some/path/dag.py", 1, 1) - - def test_does_not_increment_counter_drops_invalid_tags(self): - self.stats.incr( - "test_stats_run.delay", - tags={"key0,": "val0", "key1": "val1", "key2": "val2", "key3": "val3"}, - ) - self.statsd_client.incr.assert_called_once_with("test_stats_run.delay,key1=val1", 1, 1) + @pytest.mark.parametrize( + ("stat", "tags", "expected"), + [ + ( + "test_stats_run.delay", + {"key0": 0, "key1": "val1", "key2": "val2"}, + "test_stats_run.delay,key0=0,key1=val1", + ), + ( + "test_stats_run.dag", + {"path": "/some/path/dag.py"}, + "test_stats_run.dag,path=/some/path/dag.py", + ), + ( + "test_stats_run.delay", + {"key0,": "val0", "key1": "val1", "key2": "val2", "key3": "val3"}, + "test_stats_run.delay,key1=val1", + ), + # Empty value renders as `=true` in influxdb line protocol. + ( + "test_stats_run.delay", + {"production": "", "key1": "val1"}, + "test_stats_run.delay,production=true,key1=val1", + ), + ], + ) + def test_increment_counter_with_tags(self, stat, tags, expected): + self.stats.incr(stat, tags=tags) + self.statsd_client.incr.assert_called_once_with(expected, 1, 1) def always_invalid(stat_name): @@ -769,3 +796,26 @@ def test_does_send_stats_using_dogstatsd_when_the_name_is_valid(self): def teardown_method(self) -> None: # To avoid side-effect importlib.reload(airflow_shared.observability.metrics.stats) + + +@pytest.mark.parametrize( + ("tag_names", "expected"), + [ + pytest.param([], {}, id="empty"), + pytest.param(["production"], {"production": ""}, id="standalone"), + pytest.param(["env:prod"], {"env": "prod"}, id="key-value"), + pytest.param( + ["production", "env:prod", "team:data"], + {"production": "", "env": "prod", "team": "data"}, + id="mixed", + ), + pytest.param(["a:b:c"], {"a": "b:c"}, id="value-with-colon"), + pytest.param(["env:"], {"env": ""}, id="trailing-colon-is-standalone"), + ], +) +def test_build_dag_metric_tags(tag_names: list[str], expected: dict[str, str]) -> None: + assert build_dag_metric_tags(tag_names) == expected + + +def test_build_dag_metric_tags_accepts_generator() -> None: + assert build_dag_metric_tags(name for name in ["env:prod"]) == {"env": "prod"} 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 94363b984b7f6..7fb29e37a2728 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -45,6 +45,7 @@ from airflow.dag_processing.bundles.base import BaseDagBundle, BundleVersionLock from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.sdk._shared.observability.metrics import stats +from airflow.sdk._shared.observability.metrics.stats import build_dag_metric_tags from airflow.sdk._shared.observability.traces import get_task_span_detail_level from airflow.sdk._shared.template_rendering import truncate_rendered_value from airflow.sdk.api.client import get_hostname, getuser @@ -254,10 +255,20 @@ 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.update(dag_id=self.dag_id, task_id=self.task_id) + if self._ti_context_from_server: + # run_type keeps the tag set consistent with the scheduler-side TaskInstance.stats_tags. + # Coerce the DagRunType enum to its bare value so it serializes as e.g. "scheduled" rather + # than "dagruntype.scheduled" (matching the scheduler, which emits the plain string). + run_type = self._ti_context_from_server.dag_run.run_type + tags["run_type"] = getattr(run_type, "value", run_type) + if self._ti_context_from_server.dag_run.team_name: + tags["team_name"] = self._ti_context_from_server.dag_run.team_name return tags def __rich_repr__(self): 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 9d15ac40c6967..ce70760f83396 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 @@ -5434,11 +5434,11 @@ def test_ti_start_metric_emitted(self, create_runtime_ti, mock_supervisor_comms) run(ti, context=ti.get_template_context(), log=mock.MagicMock()) # verify ti.start was called in legacy format - backend.incr.assert_any_call(f"ti.start.{ti.dag_id}.{ti.task_id}") + backend.incr.assert_any_call(f"ti.start.{ti.dag_id}.{ti.task_id}", tags={"run_type": "manual"}) # verify ti.start was called in tagged format backend.incr.assert_any_call( "ti.start", - tags={"dag_id": ti.dag_id, "task_id": ti.task_id}, + tags={"dag_id": ti.dag_id, "task_id": ti.task_id, "run_type": "manual"}, ) @pytest.mark.parametrize( @@ -5463,11 +5463,18 @@ def test_ti_finish_metric_emitted_for_terminal_states( run(ti, context=ti.get_template_context(), log=mock.MagicMock()) # verify ti.finish was called in legacy format - backend.incr.assert_any_call(f"ti.finish.{ti.dag_id}.{ti.task_id}.{expected_state}") + backend.incr.assert_any_call( + f"ti.finish.{ti.dag_id}.{ti.task_id}.{expected_state}", tags={"run_type": "manual"} + ) # verify ti.finish was called in tagged format backend.incr.assert_any_call( "ti.finish", - tags={"dag_id": ti.dag_id, "task_id": ti.task_id, "state": expected_state}, + tags={ + "dag_id": ti.dag_id, + "task_id": ti.task_id, + "run_type": "manual", + "state": expected_state, + }, ) def test_operator_successes_metrics_emitted(self, create_runtime_ti, mock_supervisor_comms): @@ -5480,7 +5487,7 @@ def test_operator_successes_metrics_emitted(self, create_runtime_ti, mock_superv mock_get_backend.return_value = backend run(ti, context=ti.get_template_context(), log=mock.MagicMock()) - stats_tags = {"dag_id": ti.dag_id, "task_id": ti.task_id} + stats_tags = {"dag_id": ti.dag_id, "task_id": ti.task_id, "run_type": "manual"} # verify operator_successes in legacy format backend.incr.assert_any_call("operator_successes_PythonOperator", tags=stats_tags) @@ -5501,7 +5508,7 @@ def test_operator_failures_metrics_emitted(self, create_runtime_ti, mock_supervi mock_get_backend.return_value = backend run(ti, context=ti.get_template_context(), log=mock.MagicMock()) - stats_tags = {"dag_id": ti.dag_id, "task_id": ti.task_id} + stats_tags = {"dag_id": ti.dag_id, "task_id": ti.task_id, "run_type": "manual"} # verify operator_failures in legacy format backend.incr.assert_any_call("operator_failures_PythonOperator", tags=stats_tags) @@ -5532,7 +5539,12 @@ def test_ti_start_metric_respects_team_name( mock_get_backend.return_value = backend run(ti, context=ti.get_template_context(), log=mock.MagicMock()) - expected = {"dag_id": ti.dag_id, "task_id": ti.task_id, **expected_tags_extra} + expected = { + "dag_id": ti.dag_id, + "task_id": ti.task_id, + "run_type": "manual", + **expected_tags_extra, + } backend.incr.assert_any_call("ti.start", tags=expected) @pytest.mark.parametrize( @@ -5554,7 +5566,12 @@ def test_operator_metrics_respect_team_name( mock_get_backend.return_value = backend run(ti, context=ti.get_template_context(), log=mock.MagicMock()) - stats_tags = {"dag_id": ti.dag_id, "task_id": ti.task_id, "team_name": "team_a"} + stats_tags = { + "dag_id": ti.dag_id, + "task_id": ti.task_id, + "run_type": "manual", + "team_name": "team_a", + } backend.incr.assert_any_call( operator_metric, tags={**stats_tags, "operator_name": "PythonOperator"}, @@ -6343,3 +6360,40 @@ 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 + + with DAG("tagged_dag", tags=tags): + task = BaseOperator(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): + tags = _make_dag_tagged_ti(create_runtime_ti, []).stats_tags + assert tags == {"dag_id": "tagged_dag", "task_id": "t", "run_type": "manual"} + # run_type must be a plain str, not a DagRunType enum member: DagRunType is a str-enum, so the + # dict equality above passes either way, but the enum serializes as "dagruntype.manual" on the wire. + assert type(tags["run_type"]) is str + + +@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", + }