From 6c8f5b6f5ad07a420ad81f33e70603d6f8bb3bb9 Mon Sep 17 00:00:00 2001 From: Shivam Rastogi <6463385+shivaam@users.noreply.github.com> Date: Thu, 25 Jun 2026 16:27:53 -0700 Subject: [PATCH 1/2] Add team name tag to Edge executor sync metric --- .../edge3/executors/edge_executor.py | 3 +- .../edge3/executors/test_edge_executor.py | 32 +++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py b/providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py index 9ac25874c3c92..5c582781a45f0 100644 --- a/providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py +++ b/providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py @@ -35,6 +35,7 @@ from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel, EdgeWorkerState, reset_metrics from airflow.providers.edge3.models.types import is_callback_execute from airflow.utils.db import DBLocks, create_global_lock +from airflow.utils.helpers import prune_dict from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.state import TaskInstanceState @@ -340,7 +341,7 @@ def _purge_jobs(self, session: Session) -> bool: @provide_session def sync(self, *, session: Session = NEW_SESSION) -> None: """Sync will get called periodically by the heartbeat method.""" - with Stats.timer("edge_executor.sync.duration"): + with Stats.timer("edge_executor.sync.duration", tags=prune_dict({"team_name": self.team_name})): orphaned = self._update_orphaned_jobs(session) purged = self._purge_jobs(session) liveness = self._check_worker_liveness(session) diff --git a/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py b/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py index 3554689b794d2..f2b4ff2d89e22 100644 --- a/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py +++ b/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py @@ -120,6 +120,38 @@ def test_sync_orphaned_tasks(self, mock_stats_incr): assert jobs[0].task_id == "started_running_orphaned" assert jobs[0].state == TaskInstanceState.REMOVED + @patch(f"{Stats.__module__}.Stats.timer") + def test_sync_duration_metric_includes_team_name(self, mock_stats_timer): + executor = EdgeExecutor(team_name="team_a") + + with ( + mock.patch.object(executor, "_update_orphaned_jobs", return_value=False), + mock.patch.object(executor, "_purge_jobs", return_value=False), + mock.patch.object(executor, "_check_worker_liveness", return_value=False), + ): + executor.sync(session=MagicMock()) + + mock_stats_timer.assert_called_once_with( + "edge_executor.sync.duration", + tags={"team_name": "team_a"}, + ) + + @patch(f"{Stats.__module__}.Stats.timer") + def test_sync_duration_metric_omits_team_name_for_global_executor(self, mock_stats_timer): + executor = EdgeExecutor() + + with ( + mock.patch.object(executor, "_update_orphaned_jobs", return_value=False), + mock.patch.object(executor, "_purge_jobs", return_value=False), + mock.patch.object(executor, "_check_worker_liveness", return_value=False), + ): + executor.sync(session=MagicMock()) + + mock_stats_timer.assert_called_once_with( + "edge_executor.sync.duration", + tags={}, + ) + @patch("airflow.providers.edge3.executors.edge_executor.EdgeExecutor.running_state") @patch("airflow.providers.edge3.executors.edge_executor.EdgeExecutor.success") @patch("airflow.providers.edge3.executors.edge_executor.EdgeExecutor.fail") From cb853135e8607f85e30c6bd5b76513bfc1aec062 Mon Sep 17 00:00:00 2001 From: Shivam Rastogi <6463385+shivaam@users.noreply.github.com> Date: Fri, 26 Jun 2026 18:47:55 -0700 Subject: [PATCH 2/2] Address Edge executor metric test review --- .../edge3/executors/test_edge_executor.py | 30 +++++++------------ 1 file changed, 11 insertions(+), 19 deletions(-) diff --git a/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py b/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py index f2b4ff2d89e22..8975eb7ad9443 100644 --- a/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py +++ b/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py @@ -120,25 +120,17 @@ def test_sync_orphaned_tasks(self, mock_stats_incr): assert jobs[0].task_id == "started_running_orphaned" assert jobs[0].state == TaskInstanceState.REMOVED + @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="team_name is only available in Airflow 3.2+") + @pytest.mark.parametrize( + ("team_name", "expected_tags"), + [ + ("team_a", {"team_name": "team_a"}), + (None, {}), + ], + ) @patch(f"{Stats.__module__}.Stats.timer") - def test_sync_duration_metric_includes_team_name(self, mock_stats_timer): - executor = EdgeExecutor(team_name="team_a") - - with ( - mock.patch.object(executor, "_update_orphaned_jobs", return_value=False), - mock.patch.object(executor, "_purge_jobs", return_value=False), - mock.patch.object(executor, "_check_worker_liveness", return_value=False), - ): - executor.sync(session=MagicMock()) - - mock_stats_timer.assert_called_once_with( - "edge_executor.sync.duration", - tags={"team_name": "team_a"}, - ) - - @patch(f"{Stats.__module__}.Stats.timer") - def test_sync_duration_metric_omits_team_name_for_global_executor(self, mock_stats_timer): - executor = EdgeExecutor() + def test_sync_duration_metric_tags_team_name(self, mock_stats_timer, team_name, expected_tags): + executor = EdgeExecutor(team_name=team_name) with ( mock.patch.object(executor, "_update_orphaned_jobs", return_value=False), @@ -149,7 +141,7 @@ def test_sync_duration_metric_omits_team_name_for_global_executor(self, mock_sta mock_stats_timer.assert_called_once_with( "edge_executor.sync.duration", - tags={}, + tags=expected_tags, ) @patch("airflow.providers.edge3.executors.edge_executor.EdgeExecutor.running_state")