diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py b/airflow-core/src/airflow/jobs/triggerer_job_runner.py index 97fd669204456..3612feb011fc3 100644 --- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py +++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py @@ -110,7 +110,7 @@ from airflow.serialization.serialized_objects import DagSerialization from airflow.triggers.base import BaseEventTrigger, BaseTrigger, DiscrimatedTriggerEvent, TriggerEvent from airflow.triggers.shared_stream import SharedStreamManager -from airflow.utils.helpers import log_filename_template_renderer +from airflow.utils.helpers import log_filename_template_renderer, prune_dict from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.session import create_session, provide_session @@ -679,7 +679,7 @@ def heartbeat(self): perform_heartbeat(self.job, heartbeat_callback=self.heartbeat_callback, only_if_necessary=True) def heartbeat_callback(self, session: Session | None = None) -> None: - stats.incr("triggerer_heartbeat", 1, 1) + stats.incr("triggerer_heartbeat", 1, 1, tags=prune_dict({"team_name": self.team_name})) def load_triggers(self) -> None: """Assign triggers to this triggerer and update the runner with the IDs it should run.""" @@ -710,7 +710,7 @@ def handle_events(self): if entry.persist_seq is not None: self.persisted_event_seqs.append(entry.persist_seq) # Emit stat event - stats.incr("triggers.succeeded") + stats.incr("triggers.succeeded", tags=prune_dict({"team_name": self.team_name})) def on_trigger_event(self, trigger_id: int, event: TriggerEvent) -> None: """Record that a trigger fired an event.""" @@ -731,7 +731,7 @@ def handle_failed_triggers(self): trigger_id, exc = self.failed_triggers.popleft() self.on_trigger_failure(trigger_id=trigger_id, exc=exc) # Emit stat event - stats.incr("triggers.failed") + stats.incr("triggers.failed", tags=prune_dict({"team_name": self.team_name})) def on_trigger_failure(self, trigger_id: int, exc: list[str] | None) -> None: """Record that a trigger failed.""" @@ -753,7 +753,7 @@ def metric_tags(self) -> dict[str, str]: "TriggerRunnerSupervisor.metric_tags() requires a Job with a hostname; " "subclasses without a metadata-DB Job must override this method." ) - return {"hostname": hostname} + return prune_dict({"hostname": hostname, "team_name": self.team_name}) def emit_metrics(self): tags = self.metric_tags() diff --git a/airflow-core/src/airflow/models/callback.py b/airflow-core/src/airflow/models/callback.py index 78247be9ec018..1e8e758450ded 100644 --- a/airflow-core/src/airflow/models/callback.py +++ b/airflow-core/src/airflow/models/callback.py @@ -39,6 +39,7 @@ from airflow.models import Base from airflow.models.base import StringID from airflow.models.dagbundle import DagBundleModel +from airflow.utils.helpers import prune_dict from airflow.utils.sqlalchemy import ExtendedJSON, UtcDateTime from airflow.utils.state import CallbackState @@ -156,7 +157,7 @@ def __init__(self, priority_weight: int = 1, prefix: str = "", **kwargs): def queue(self, *, session: Session) -> None: self.state = CallbackState.QUEUED - def get_metric_info(self, status: CallbackState, result: Any) -> dict: + def get_metric_info(self, status: CallbackState, result: Any, team_name: str | None = None) -> dict: tags = {"result": result, **self.data} tags.pop("prefix", None) @@ -177,6 +178,10 @@ def get_metric_info(self, status: CallbackState, result: Any) -> dict: for k, v in tags.items() } + # team_name is omitted entirely when not in a multi-team deployment or when the + # callback's bundle is not mapped to a team. + tags = prune_dict({**tags, "team_name": team_name}) + prefix = self.data.get("prefix", "") name = f"{prefix}.callback_{status}" if prefix else f"callback_{status}" @@ -250,9 +255,12 @@ def handle_event(self, event: TriggerEvent, session: Session): if (status := event.payload.get(PAYLOAD_STATUS_KEY)) and status in (ACTIVE_STATES | TERMINAL_STATES): self.state = status if status in TERMINAL_STATES: + team_name: str | None = None + if self.bundle_name and conf.getboolean("core", "multi_team"): + team_name = DagBundleModel.get_team_name(self.bundle_name, session=session) self.trigger = None self.output = event.payload.get(PAYLOAD_BODY_KEY) - stats.incr(**self.get_metric_info(status, self.output)) + stats.incr(**self.get_metric_info(status, self.output, team_name=team_name)) session.add(self) else: diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index 740596f9d69b6..bbf4b0db92377 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/airflow-core/src/airflow/models/taskinstance.py @@ -1505,7 +1505,7 @@ def emit_state_change_metric(self, new_state: TaskInstanceState) -> None: stats.timing( f"task.{metric_name}", timing, - tags={"task_id": self.task_id, "dag_id": self.dag_id, "queue": self.queue}, + tags={**self.stats_tags, "queue": self.queue}, ) def clear_next_method_args(self) -> None: diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py b/airflow-core/tests/unit/jobs/test_triggerer_job.py index 923779df0034c..2a875fbfbf6c1 100644 --- a/airflow-core/tests/unit/jobs/test_triggerer_job.py +++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py @@ -2460,6 +2460,83 @@ def test_handle_events_does_not_confirm_seq_when_persist_fails(jobless_superviso assert list(jobless_supervisor.persisted_event_seqs) == [] +@pytest.mark.parametrize( + ("team_name", "expected_tags"), + [ + pytest.param("team_alpha", {"team_name": "team_alpha"}, id="with_team"), + pytest.param(None, {}, id="without_team"), + ], +) +def test_handle_events_emits_team_name(jobless_supervisor, team_name, expected_tags): + """triggers.succeeded carries the triggerer's team_name (omitted when the triggerer has none).""" + jobless_supervisor.team_name = team_name + jobless_supervisor.events.append(TriggerEventEntry(1, TriggerEvent(True), 7)) + + with ( + mock.patch.object(TriggerRunnerSupervisor, "on_trigger_event", autospec=True), + mock.patch("airflow.jobs.triggerer_job_runner.stats.incr") as mock_incr, + ): + jobless_supervisor.handle_events() + + mock_incr.assert_called_once_with("triggers.succeeded", tags=expected_tags) + + +@pytest.mark.parametrize( + ("team_name", "expected_tags"), + [ + pytest.param("team_alpha", {"team_name": "team_alpha"}, id="with_team"), + pytest.param(None, {}, id="without_team"), + ], +) +def test_handle_failed_triggers_emits_team_name(jobless_supervisor, team_name, expected_tags): + """triggers.failed carries the triggerer's team_name (omitted when the triggerer has none).""" + jobless_supervisor.team_name = team_name + jobless_supervisor.failed_triggers.append((1, None)) + + with ( + mock.patch.object(TriggerRunnerSupervisor, "on_trigger_failure", autospec=True), + mock.patch("airflow.jobs.triggerer_job_runner.stats.incr") as mock_incr, + ): + jobless_supervisor.handle_failed_triggers() + + mock_incr.assert_called_once_with("triggers.failed", tags=expected_tags) + + +@pytest.mark.parametrize( + ("team_name", "expected_tags"), + [ + pytest.param("team_alpha", {"team_name": "team_alpha"}, id="with_team"), + pytest.param(None, {}, id="without_team"), + ], +) +def test_heartbeat_callback_emits_team_name(jobless_supervisor, team_name, expected_tags): + jobless_supervisor.team_name = team_name + + with mock.patch("airflow.jobs.triggerer_job_runner.stats.incr") as mock_incr: + jobless_supervisor.heartbeat_callback() + + mock_incr.assert_called_once_with("triggerer_heartbeat", 1, 1, tags=expected_tags) + + +@pytest.mark.parametrize( + ("team_name", "expected_extra"), + [ + pytest.param("team_alpha", {"team_name": "team_alpha"}, id="with_team"), + pytest.param(None, {}, id="without_team"), + ], +) +def test_emit_metrics_includes_team_name(supervisor_builder, mocker, team_name, expected_extra): + supervisor = supervisor_builder() + supervisor.team_name = team_name + gauge = mocker.patch("airflow.jobs.triggerer_job_runner.stats.gauge") + + supervisor.emit_metrics() + + expected_tags = {"hostname": supervisor.job.hostname, **expected_extra} + gauge.assert_any_call("triggers.running", mock.ANY, tags=expected_tags) + gauge.assert_any_call("triggerer.capacity_left", mock.ANY, tags=expected_tags) + + def test_state_sync_carries_and_drains_persist_confirmations(jobless_supervisor): """The state-sync response carries pending confirmations once, then None when there are none.""" jobless_supervisor.persisted_event_seqs.extend([3, 9]) diff --git a/airflow-core/tests/unit/models/test_callback.py b/airflow-core/tests/unit/models/test_callback.py index b5296979ed417..b2b13582710fc 100644 --- a/airflow-core/tests/unit/models/test_callback.py +++ b/airflow-core/tests/unit/models/test_callback.py @@ -143,6 +143,21 @@ def test_get_metric_info_dict_values_are_stringified(self): # insertion order collapse to one metric series (no needless cardinality split). assert metric_info["tags"]["result"] == '{"code": 0, "output": [1, 2]}' + @pytest.mark.parametrize( + ("team_name", "expect_tag"), + [ + pytest.param("team_alpha", True, id="with_team"), + pytest.param(None, False, id="without_team"), + ], + ) + def test_get_metric_info_includes_team_name(self, team_name, expect_tag): + callback = TriggererCallback(TEST_ASYNC_CALLBACK, prefix="deadline_alerts", dag_id=TEST_DAG_ID) + metric_info = callback.get_metric_info(CallbackState.SUCCESS, "0", team_name=team_name) + if expect_tag: + assert metric_info["tags"]["team_name"] == team_name + else: + assert "team_name" not in metric_info["tags"] + class TestTriggererCallback: def test_polymorphic_serde(self, session): @@ -249,6 +264,36 @@ def test_handle_event(self, session, event, terminal_state): assert callback.trigger is None assert callback.output == event.payload[PAYLOAD_BODY_KEY] + @pytest.mark.parametrize( + ("multi_team", "team_name", "expect_tag"), + [ + pytest.param("true", "team_alpha", True, id="with_team"), + pytest.param("false", None, False, id="without_team"), + ], + ) + @patch("airflow.models.callback.stats.incr") + def test_handle_event_emits_team_name(self, mock_incr, multi_team, team_name, expect_tag, session): + """On a terminal event, callback_{status} carries team_name resolved from the bundle.""" + callback = TriggererCallback(TEST_ASYNC_CALLBACK, dag_id=TEST_DAG_ID) + callback.bundle_name = "test_bundle" + event = TriggerEvent({PAYLOAD_STATUS_KEY: CallbackState.SUCCESS, PAYLOAD_BODY_KEY: "0"}) + + with ( + conf_vars({("core", "multi_team"): multi_team}), + patch( + "airflow.models.callback.DagBundleModel.get_team_name", return_value=team_name + ) as mock_get_team_name, + ): + callback.handle_event(event, session) + + mock_incr.assert_called_once() + _, kwargs = mock_incr.call_args + if expect_tag: + assert kwargs["tags"]["team_name"] == team_name + else: + mock_get_team_name.assert_not_called() + assert "team_name" not in kwargs["tags"] + class TestExecutorCallback: def test_polymorphic_serde(self, session): diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py index b9d8c85f61321..4e3fc7090284a 100644 --- a/airflow-core/tests/unit/models/test_taskinstance.py +++ b/airflow-core/tests/unit/models/test_taskinstance.py @@ -4126,3 +4126,41 @@ def test_stats_tags_with_none_team_name(self, dag_maker, session): ti._team_name = None tags = ti.stats_tags assert "team_name" not in tags + + @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"}, + id="with_team", + ), + pytest.param( + None, + {"dag_id": "test_dag", "task_id": "my_task", "queue": "default"}, + id="without_team", + ), + ], + ) + @mock.patch("airflow._shared.observability.metrics.stats.timing") + def test_emit_state_change_metric_includes_team_name( + self, mock_timing, team_name, expected_tags, dag_maker, session + ): + with dag_maker("test_dag"): + EmptyOperator(task_id="my_task") + dr = dag_maker.create_dagrun() + ti = dr.get_task_instance("my_task", session=session) + ti.state = TaskInstanceState.SCHEDULED + ti.scheduled_dttm = timezone.utcnow() + if team_name: + ti._team_name = team_name + session.merge(ti) + session.flush() + + ti.emit_state_change_metric(TaskInstanceState.QUEUED) + + mock_timing.assert_called_once_with( + "task.scheduled_duration", + mock.ANY, + tags=expected_tags, + ) diff --git a/task-sdk/src/airflow/sdk/bases/resumablejobmixin.py b/task-sdk/src/airflow/sdk/bases/resumablejobmixin.py index 8bfda28a66442..7066d10cced4e 100644 --- a/task-sdk/src/airflow/sdk/bases/resumablejobmixin.py +++ b/task-sdk/src/airflow/sdk/bases/resumablejobmixin.py @@ -107,7 +107,13 @@ def execute_resumable(self, context: Context) -> Any: Closing this window would require atomic "submit + persist", which is not possible across an external system boundary. """ - operator_tag = {"operator": type(self).__name__} + stats_tags = {"operator": type(self).__name__} + # The task is team-scoped in multi-team deployments; surface team_name on the + # resumable_job metrics via the running task instance's stats tags (omitted when + # not multi-team or the task has no team). + ti = context.get("ti") + if ti is not None and (team_name := ti.stats_tags.get("team_name")): + stats_tags["team_name"] = team_name reconnect_to: Any = None already_succeeded_id: Any = None @@ -125,7 +131,7 @@ def execute_resumable(self, context: Context) -> Any: else: external_id = task_state_store.get(self.external_id_key) if external_id: - stats.incr("resumable_job.reconnect_attempt", tags=operator_tag) + stats.incr("resumable_job.reconnect_attempt", tags=stats_tags) status = self.get_job_status(external_id, context) @@ -135,7 +141,7 @@ def execute_resumable(self, context: Context) -> Any: if self.is_job_active(status): # Job is still running, skip submission and reconnect to it. span.set_attribute("resumable.decision", "reconnect") - stats.incr("resumable_job.reconnect_success", tags=operator_tag) + stats.incr("resumable_job.reconnect_success", tags=stats_tags) self.log.info( "Reconnecting to existing job", external_id_key=self.external_id_key, @@ -146,7 +152,7 @@ def execute_resumable(self, context: Context) -> Any: elif self.is_job_succeeded(status): # Job already finished successfully, skip polling and return result directly. span.set_attribute("resumable.decision", "already_succeeded") - stats.incr("resumable_job.already_succeeded", tags=operator_tag) + stats.incr("resumable_job.already_succeeded", tags=stats_tags) self.log.info( "Job already completed successfully, skipping resubmission", external_id_key=self.external_id_key, @@ -156,7 +162,7 @@ def execute_resumable(self, context: Context) -> Any: else: # Job is in a terminal failed state, fall through and submit a new job. span.set_attribute("resumable.decision", "terminal_resubmit") - stats.incr("resumable_job.terminal_resubmit", tags=operator_tag) + stats.incr("resumable_job.terminal_resubmit", tags=stats_tags) self.log.warning( "Prior job in terminal state, resubmitting fresh", external_id_key=self.external_id_key, @@ -165,7 +171,7 @@ def execute_resumable(self, context: Context) -> Any: ) else: span.set_attribute("resumable.decision", "fresh_submit") - stats.incr("resumable_job.fresh_submit", tags=operator_tag) + stats.incr("resumable_job.fresh_submit", tags=stats_tags) self.log.debug( "No stored external ID found; submitting fresh job", external_id_key=self.external_id_key, diff --git a/task-sdk/src/airflow/sdk/types.py b/task-sdk/src/airflow/sdk/types.py index 711dbe71ca46c..029af34496f07 100644 --- a/task-sdk/src/airflow/sdk/types.py +++ b/task-sdk/src/airflow/sdk/types.py @@ -151,6 +151,9 @@ def log_url(self) -> str: ... @property def mark_success_url(self) -> str: ... + @property + def stats_tags(self) -> dict[str, str]: ... + def xcom_pull( self, task_ids: str | Iterable[str] | None = None, diff --git a/task-sdk/tests/task_sdk/bases/test_resumablejobmixin.py b/task-sdk/tests/task_sdk/bases/test_resumablejobmixin.py index 11f97616cdf78..e186cc19a367c 100644 --- a/task-sdk/tests/task_sdk/bases/test_resumablejobmixin.py +++ b/task-sdk/tests/task_sdk/bases/test_resumablejobmixin.py @@ -85,6 +85,20 @@ def make_context(task_store: FakeTaskState | None = None) -> dict: return ctx +class FakeTI: + """Minimal stand-in for RuntimeTaskInstance exposing stats_tags with an optional team_name.""" + + def __init__(self, team_name: str | None = None): + self._team_name = team_name + + @property + def stats_tags(self) -> dict[str, str]: + tags = {"dag_id": "d", "task_id": "t"} + if self._team_name: + tags["team_name"] = self._team_name + return tags + + class TestFirstSubmission: def test_submits_and_polls_when_no_prior_state(self): op = ConcreteResumableOperator(task_id="test_task") @@ -253,6 +267,26 @@ def test_terminal_resubmit_fires_when_job_failed(self): assert "resumable_job.reconnect_success" not in called_names assert "resumable_job.fresh_submit" not in called_names + @pytest.mark.parametrize( + ("team_name", "expected_tag"), + [ + pytest.param( + "team_alpha", + {"operator": "ConcreteResumableOperator", "team_name": "team_alpha"}, + id="with_team", + ), + pytest.param(None, {"operator": "ConcreteResumableOperator"}, id="without_team"), + ], + ) + def test_team_name_added_to_metric_tags(self, team_name, expected_tag): + op = ConcreteResumableOperator(task_id="test_task") + ctx = make_context(FakeTaskState()) + ctx["ti"] = FakeTI(team_name) + mock_incr = MagicMock() + with patch(self._PATCH, mock_incr): + op.execute_resumable(ctx) + mock_incr.assert_called_once_with("resumable_job.fresh_submit", tags=expected_tag) + class TestTracing: _MODULE_TRACER = "airflow.sdk.bases.resumablejobmixin.tracer"