diff --git a/airflow-core/src/airflow/executors/base_executor.py b/airflow-core/src/airflow/executors/base_executor.py index 31c307b163f88..ef5772331108f 100644 --- a/airflow-core/src/airflow/executors/base_executor.py +++ b/airflow-core/src/airflow/executors/base_executor.py @@ -40,6 +40,7 @@ from airflow.models import Log from airflow.models.taskinstancekey import TaskInstanceKey from airflow.observability.metrics import stats_utils +from airflow.utils.helpers import prune_dict from airflow.utils.log.logging_mixin import LoggingMixin PARALLELISM: int = conf.getint("core", "PARALLELISM") @@ -411,17 +412,17 @@ def _emit_metrics(self, open_slots, num_running_tasks, num_queued_tasks): stats.gauge( open_slots_metric_name, value=open_slots, - tags={"status": "open", "executor_class_name": name}, + tags=prune_dict({"status": "open", "executor_class_name": name, "team_name": self.team_name}), ) stats.gauge( queued_tasks_metric_name, value=num_queued_tasks, - tags={"status": "queued", "executor_class_name": name}, + tags=prune_dict({"status": "queued", "executor_class_name": name, "team_name": self.team_name}), ) stats.gauge( running_tasks_metric_name, value=num_running_tasks, - tags={"status": "running", "executor_class_name": name}, + tags=prune_dict({"status": "running", "executor_class_name": name, "team_name": self.team_name}), ) def order_queued_tasks_by_priority(self) -> list[tuple[TaskInstanceKey, workloads.ExecuteTask]]: diff --git a/airflow-core/tests/unit/executors/test_base_executor.py b/airflow-core/tests/unit/executors/test_base_executor.py index f74d1f1e08464..cb646f7ce8d55 100644 --- a/airflow-core/tests/unit/executors/test_base_executor.py +++ b/airflow-core/tests/unit/executors/test_base_executor.py @@ -174,30 +174,32 @@ def test_fail_and_success(): assert len(executor.get_event_buffer()) == 3 +@pytest.mark.parametrize( + ("team_name", "expected_tags"), + [ + pytest.param(None, {"status": "open", "executor_class_name": "BaseExecutor"}, id="without_team"), + pytest.param( + "team_a", + {"status": "open", "executor_class_name": "BaseExecutor", "team_name": "team_a"}, + id="with_team", + ), + ], +) @mock.patch("airflow.executors.base_executor.BaseExecutor.sync") @mock.patch("airflow.executors.base_executor.BaseExecutor.trigger_tasks") @mock.patch("airflow.executors.base_executor.stats.gauge") -def test_gauge_executor_metrics_single_executor(mock_stats_gauge, mock_trigger_tasks, mock_sync): - executor = BaseExecutor() +def test_gauge_executor_metrics_single_executor( + mock_stats_gauge, mock_trigger_tasks, mock_sync, team_name, expected_tags +): + executor = BaseExecutor(team_name=team_name) executor.heartbeat() - calls = [ - mock.call( - "executor.open_slots", - value=mock.ANY, - tags={"status": "open", "executor_class_name": "BaseExecutor"}, - ), - mock.call( - "executor.queued_tasks", - value=mock.ANY, - tags={"status": "queued", "executor_class_name": "BaseExecutor"}, - ), - mock.call( - "executor.running_tasks", - value=mock.ANY, - tags={"status": "running", "executor_class_name": "BaseExecutor"}, - ), - ] - mock_stats_gauge.assert_has_calls(calls) + # Verify all three gauges use the expected tag structure + for metric, status in [ + ("executor.open_slots", "open"), + ("executor.queued_tasks", "queued"), + ("executor.running_tasks", "running"), + ]: + mock_stats_gauge.assert_any_call(metric, value=mock.ANY, tags={**expected_tags, "status": status}) @pytest.mark.parametrize(