[v3-3-test] Add team_name tag to dag processor metrics for multi-team deployments

airflow-core/src/airflow/dag_processing/manager.py

@@ -69,6 +69,7 @@
 from airflow.sdk.log import init_log_file, logging_processors
 from airflow.typing_compat import assert_never
 from airflow.utils.file import list_py_file_paths, might_contain_dag
+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.process_utils import (
@@ -719,7 +720,12 @@ def _add_callback_to_queue(self, request: CallbackRequest) -> None:
         )
         self._callback_to_execute[file_info].append(request)
         self._add_files_to_queue([file_info], mode="front")
-        stats.incr("dag_processing.other_callback_count")
+        team_name = (
+            DagBundleModel.get_team_name(file_info.bundle_name)
+            if conf.getboolean("core", "multi_team")
+            else None
+        )
+        stats.incr("dag_processing.other_callback_count", tags=prune_dict({"team_name": team_name}))
 
     @provide_session
     def get_bundle_state(self, bundle_name: str, *, session: Session = NEW_SESSION) -> BundleState | None:
@@ -1020,6 +1026,14 @@ def _log_file_processing_stats(self, known_files: dict[str, set[DagFileInfo]]):
         utcnow = timezone.utcnow()
         now = time.monotonic()
 
+        if conf.getboolean("core", "multi_team"):
+            bundle_names = {bundle_name for bundle_name in known_files}
+            bundle_to_team = {
+                bundle_name: DagBundleModel.get_team_name(bundle_name) for bundle_name in bundle_names
+            }
+        else:
+            bundle_to_team = {}
+
         for files in known_files.values():
             for file in files:
                 stat = self._file_stats[file]
@@ -1040,11 +1054,14 @@ def _log_file_processing_stats(self, known_files: dict[str, set[DagFileInfo]]):
                     stats.gauge(
                         "dag_processing.last_run.seconds_ago",
                         seconds_ago,
-                        tags={
-                            "file_path": file.normalized_file_path_for_stats,
-                            "bundle_name": normalize_name_for_stats(file.bundle_name),
-                            "file_name": file_name,
-                        },
+                        tags=prune_dict(
+                            {
+                                "file_path": file.normalized_file_path_for_stats,
+                                "bundle_name": normalize_name_for_stats(file.bundle_name),
+                                "file_name": file_name,
+                                "team_name": bundle_to_team.get(file.bundle_name),
+                            }
+                        ),
                     )
 
                 rows.append(
@@ -1136,6 +1153,15 @@ def remove_orphaned_file_stats(self, present: set[DagFileInfo]):
     def terminate_orphan_processes(self, present: set[DagFileInfo]):
         """Stop processors that are working on deleted files."""
         present_keys = {file.presence_key for file in present}
+
+        if conf.getboolean("core", "multi_team"):
+            bundle_names = {file.bundle_name for file in self._processors}
+            bundle_to_team = {
+                bundle_name: DagBundleModel.get_team_name(bundle_name) for bundle_name in bundle_names
+            }
+        else:
+            bundle_to_team = {}
+
         for file in list(self._processors.keys()):
             if file.presence_key not in present_keys:
                 processor = self._processors.pop(file, None)
@@ -1145,7 +1171,13 @@ def terminate_orphan_processes(self, present: set[DagFileInfo]):
                 self.log.warning("Stopping processor for %s", file_name)
                 stats.decr(
                     "dag_processing.processes",
-                    tags={"file_path": file.normalized_file_path_for_stats, "action": "stop"},
+                    tags=prune_dict(
+                        {
+                            "file_path": file.normalized_file_path_for_stats,
+                            "action": "stop",
+                            "team_name": bundle_to_team.get(file.bundle_name),
+                        }
+                    ),
                 )
                 processor.kill(signal.SIGKILL)
                 processor.logger_filehandle.close()
@@ -1183,6 +1215,9 @@ def handle_parsing_result(
 
         run_duration = time.monotonic() - proc.start_time
         finish_time = timezone.utcnow()
+        team_name = (
+            DagBundleModel.get_team_name(file.bundle_name) if conf.getboolean("core", "multi_team") else None
+        )
         next_stat = process_parse_results(
             run_duration=run_duration,
             finish_time=finish_time,
@@ -1191,6 +1226,7 @@ def handle_parsing_result(
             parsing_result=proc.parsing_result,
             is_callback_only=is_callback_only,
             relative_fileloc=str(file.rel_path),
+            team_name=team_name,
         )
 
         if proc.parsing_result is not None:
@@ -1358,6 +1394,14 @@ def _create_process(self, dag_file: DagFileInfo) -> DagFileProcessorProcess:
 
     def _start_new_processes(self):
         """Start more processors if we have enough slots and files to process."""
+        if conf.getboolean("core", "multi_team"):
+            bundle_names = {file.bundle_name for file in self._file_queue}
+            bundle_to_team = {
+                bundle_name: DagBundleModel.get_team_name(bundle_name) for bundle_name in bundle_names
+            }
+        else:
+            bundle_to_team = {}
+
         while self._parallelism > len(self._processors) and self._file_queue:
             file, _ = self._file_queue.popitem(last=False)
             # Stop creating duplicate processor i.e. processor with the same filepath
@@ -1367,7 +1411,13 @@ def _start_new_processes(self):
             processor = self._create_process(file)
             stats.incr(
                 "dag_processing.processes",
-                tags={"file_path": file.normalized_file_path_for_stats, "action": "start"},
+                tags=prune_dict(
+                    {
+                        "file_path": file.normalized_file_path_for_stats,
+                        "action": "start",
+                        "team_name": bundle_to_team.get(file.bundle_name),
+                    }
+                ),
             )
 
             self._processors[file] = processor
@@ -1533,6 +1583,15 @@ def _kill_timed_out_processors(self):
         """Kill any file processors that timeout to defend against process hangs."""
         now = time.monotonic()
         processors_to_remove = []
+
+        if conf.getboolean("core", "multi_team"):
+            bundle_names = {file.bundle_name for file in self._processors}
+            bundle_to_team = {
+                bundle_name: DagBundleModel.get_team_name(bundle_name) for bundle_name in bundle_names
+            }
+        else:
+            bundle_to_team = {}
+
         for file, processor in self._processors.items():
             duration = now - processor.start_time
             if duration > self.processor_timeout:
@@ -1544,8 +1603,17 @@ def _kill_timed_out_processors(self):
                     self.processor_timeout,
                 )
                 file_path_tag = file.normalized_file_path_for_stats
-                stats.decr("dag_processing.processes", tags={"file_path": file_path_tag, "action": "timeout"})
-                stats.incr("dag_processing.processor_timeouts", tags={"file_path": file_path_tag})
+                team_name = bundle_to_team.get(file.bundle_name)
+                stats.decr(
+                    "dag_processing.processes",
+                    tags=prune_dict(
+                        {"file_path": file_path_tag, "action": "timeout", "team_name": team_name}
+                    ),
+                )
+                stats.incr(
+                    "dag_processing.processor_timeouts",
+                    tags=prune_dict({"file_path": file_path_tag, "team_name": team_name}),
+                )
                 processor.kill(signal.SIGKILL)
 
                 processors_to_remove.append(file)
@@ -1606,10 +1674,24 @@ def max_runs_reached(self):
 
     def terminate(self):
         """Stop all running processors."""
+        if conf.getboolean("core", "multi_team"):
+            bundle_names = {file.bundle_name for file in self._processors}
+            bundle_to_team = {
+                bundle_name: DagBundleModel.get_team_name(bundle_name) for bundle_name in bundle_names
+            }
+        else:
+            bundle_to_team = {}
+
         for file, processor in self._processors.items():
             stats.decr(
                 "dag_processing.processes",
-                tags={"file_path": file.normalized_file_path_for_stats, "action": "terminate"},
+                tags=prune_dict(
+                    {
+                        "file_path": file.normalized_file_path_for_stats,
+                        "action": "terminate",
+                        "team_name": bundle_to_team.get(file.bundle_name),
+                    }
+                ),
             )
             # SIGTERM, wait 5s, SIGKILL if still alive
             processor.kill(signal.SIGTERM, escalation_delay=5.0)
@@ -1642,6 +1724,7 @@ def process_parse_results(
     *,
     is_callback_only: bool = False,
     relative_fileloc: str | None = None,
+    team_name: str | None = None,
 ) -> DagFileStat:
     """
     Create a DagFileStat from parsing results and emit metrics.
@@ -1655,7 +1738,7 @@ def process_parse_results(
             last_duration=run_duration,
             run_count=run_count,  # Don't increment for callback-only processing
         )
-        stats.incr("dag_processing.callback_only_count")
+        stats.incr("dag_processing.callback_only_count", tags=prune_dict({"team_name": team_name}))
     else:
         # Actual DAG parsing or import error
         stat = DagFileStat(
@@ -1673,7 +1756,9 @@ def process_parse_results(
         stats.timing(
             "dag_processing.last_duration",
             stat.last_duration,
-            tags={"bundle_name": normalized_bundle, "file_name": file_name},
+            tags=prune_dict(
+                {"bundle_name": normalized_bundle, "file_name": file_name, "team_name": team_name}
+            ),
         )
 
     if parsing_result is None:

airflow-core/src/airflow/dag_processing/processor.py

@@ -38,6 +38,7 @@
 from airflow.configuration import conf
 from airflow.dag_processing.bundles.base import BundleVersionLock
 from airflow.dag_processing.dagbag import BundleDagBag, DagBag
+from airflow.models.dag import DagModel
 from airflow.sdk.exceptions import TaskNotFound
 from airflow.sdk.execution_time.comms import (
     ConnectionResult,
@@ -90,6 +91,7 @@
 from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG
 from airflow.utils.dag_version_inflation_checker import check_dag_file_stability
 from airflow.utils.file import iter_airflow_imports
+from airflow.utils.helpers import prune_dict
 from airflow.utils.state import TaskInstanceState
 
 if TYPE_CHECKING:
@@ -386,7 +388,19 @@ def _execute_dag_callbacks(dagbag: DagBag, request: DagCallbackRequest, log: Fil
             callback(context)
         except Exception:
             log.exception("Callback failed", dag_id=request.dag_id)
-            stats.incr("dag.callback_exceptions", tags={"dag_id": request.dag_id})
+            stats.incr(
+                "dag.callback_exceptions",
+                tags=prune_dict(
+                    {
+                        "dag_id": request.dag_id,
+                        "team_name": (
+                            DagModel.get_team_name(request.dag_id)
+                            if conf.getboolean("core", "multi_team")
+                            else None
+                        ),
+                    }
+                ),
+            )
 
 
 def _execute_task_callbacks(dagbag: DagBag, request: TaskCallbackRequest, log: FilteringBoundLogger) -> None:

airflow-core/src/airflow/models/dagrun.py

@@ -1420,6 +1420,7 @@ def execute_dag_callbacks(
             TaskInstance as TIDataModel,
             TIRunContext,
         )
+        from airflow.models.dag import DagModel
         from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
 
         if relevant_ti:
@@ -1463,7 +1464,19 @@ def execute_dag_callbacks(
                 callback(context)
             except Exception:
                 self.log.exception("Callback failed for %s", dag.dag_id)
-                stats.incr("dag.callback_exceptions", tags={"dag_id": dag.dag_id})
+                stats.incr(
+                    "dag.callback_exceptions",
+                    tags=prune_dict(
+                        {
+                            "dag_id": dag.dag_id,
+                            "team_name": (
+                                DagModel.get_team_name(dag.dag_id)
+                                if airflow_conf.getboolean("core", "multi_team")
+                                else None
+                            ),
+                        }
+                    ),
+                )
 
     def _get_ready_tis(
         self,