This is an automated email from the ASF dual-hosted git repository.
ferruzzi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 1bb880df42f Add team_name tags to Amazon executor metrics (#69072)
1bb880df42f is described below
commit 1bb880df42f8c185a27eca1e3d4ee7fee9c8d225
Author: Justin Pakzad <[email protected]>
AuthorDate: Mon Jun 29 18:37:19 2026 -0400
Add team_name tags to Amazon executor metrics (#69072)
---
.../aws/executors/aws_lambda/lambda_executor.py | 10 ++++++--
.../amazon/aws/executors/batch/batch_executor.py | 11 ++++++---
.../amazon/aws/executors/ecs/ecs_executor.py | 10 ++++++--
.../executors/aws_lambda/test_lambda_executor.py | 27 ++++++++++++++++++++++
.../aws/executors/batch/test_batch_executor.py | 25 ++++++++++++++++++++
.../amazon/aws/executors/ecs/test_ecs_executor.py | 27 +++++++++++++++++++++-
6 files changed, 102 insertions(+), 8 deletions(-)
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py
b/providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py
index 72455165e83..5d30d5c84a2 100644
---
a/providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py
+++
b/providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py
@@ -44,6 +44,7 @@ from airflow.providers.amazon.aws.hooks.lambda_function
import LambdaHook
from airflow.providers.amazon.aws.hooks.sqs import SqsHook
from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_3_PLUS
from airflow.providers.common.compat.sdk import AirflowException, Stats,
timezone
+from airflow.utils.helpers import prune_dict
if TYPE_CHECKING:
from sqlalchemy.orm import Session
@@ -94,7 +95,10 @@ class AwsLambdaExecutor(BaseExecutor):
from airflow.providers.common.compat.sdk import conf
self.conf = conf
-
+ # TODO: Remove this fallback once the min Airflow version for
providers is >= 3.1
+ # (BaseExecutor.team_name is always defined from Airflow 3.1 onward).
+ if not hasattr(self, "team_name"):
+ self.team_name = None
self.lambda_function_name = self.conf.get(CONFIG_GROUP_NAME,
AllLambdaConfigKeys.FUNCTION_NAME)
self.sqs_queue_url = self.conf.get(CONFIG_GROUP_NAME,
AllLambdaConfigKeys.QUEUE_URL)
self.dlq_url = self.conf.get(CONFIG_GROUP_NAME,
AllLambdaConfigKeys.DLQ_URL)
@@ -558,7 +562,9 @@ class AwsLambdaExecutor(BaseExecutor):
:param tis: The task instances to adopt.
"""
- with Stats.timer("lambda_executor.adopt_task_instances.duration"):
+ with Stats.timer(
+ "lambda_executor.adopt_task_instances.duration",
tags=prune_dict({"team_name": self.team_name})
+ ):
adopted_tis: list[TaskInstance] = []
if serialized_workload_keys := [
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py
b/providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py
index 74614961f89..00c7dcb7881 100644
---
a/providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py
+++
b/providers/amazon/src/airflow/providers/amazon/aws/executors/batch/batch_executor.py
@@ -35,7 +35,7 @@ from
airflow.providers.amazon.aws.executors.utils.exponential_backoff_retry impo
from airflow.providers.amazon.aws.hooks.batch_client import BatchClientHook
from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_3_PLUS
from airflow.providers.common.compat.sdk import AirflowException, Stats,
timezone
-from airflow.utils.helpers import merge_dicts
+from airflow.utils.helpers import merge_dicts, prune_dict
if TYPE_CHECKING:
from sqlalchemy.orm import Session
@@ -116,7 +116,10 @@ class AwsBatchExecutor(BaseExecutor):
from airflow.providers.common.compat.sdk import conf
self.conf = conf
-
+ # TODO: Remove this fallback once the min Airflow version for
providers is >= 3.1
+ # (BaseExecutor.team_name is always defined from Airflow 3.1 onward).
+ if not hasattr(self, "team_name"):
+ self.team_name = None
self.attempts_since_last_successful_connection = 0
self.load_batch_connection(check_connection=False)
self.IS_BOTO_CONNECTION_HEALTHY = False
@@ -540,7 +543,9 @@ class AwsBatchExecutor(BaseExecutor):
Anything that is not adopted will be cleared by the scheduler and
becomes eligible for re-scheduling.
"""
- with Stats.timer("batch_executor.adopt_task_instances.duration"):
+ with Stats.timer(
+ "batch_executor.adopt_task_instances.duration",
tags=prune_dict({"team_name": self.team_name})
+ ):
adopted_tis: list[TaskInstance] = []
if job_ids := [ti.external_executor_id for ti in tis if
ti.external_executor_id]:
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py
b/providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py
index 644fbb29ed6..cba439ddc62 100644
---
a/providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py
+++
b/providers/amazon/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py
@@ -50,7 +50,7 @@ from
airflow.providers.amazon.aws.executors.utils.exponential_backoff_retry impo
from airflow.providers.amazon.aws.hooks.ecs import EcsHook
from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_3_PLUS
from airflow.providers.common.compat.sdk import AirflowException, Stats,
timezone
-from airflow.utils.helpers import merge_dicts
+from airflow.utils.helpers import merge_dicts, prune_dict
from airflow.utils.state import State
if TYPE_CHECKING:
@@ -126,6 +126,10 @@ class AwsEcsExecutor(BaseExecutor):
from airflow.providers.common.compat.sdk import conf
self.conf = conf
+ # TODO: Remove this fallback once the min Airflow version for
providers is >= 3.1
+ # (BaseExecutor.team_name is always defined from Airflow 3.1 onward).
+ if not hasattr(self, "team_name"):
+ self.team_name = None
self.cluster = self.conf.get(CONFIG_GROUP_NAME,
AllEcsConfigKeys.CLUSTER)
self.container_name = self.conf.get(CONFIG_GROUP_NAME,
AllEcsConfigKeys.CONTAINER_NAME)
@@ -641,7 +645,9 @@ class AwsEcsExecutor(BaseExecutor):
Anything that is not adopted will be cleared by the scheduler and
becomes eligible for re-scheduling.
"""
- with Stats.timer("ecs_executor.adopt_task_instances.duration"):
+ with Stats.timer(
+ "ecs_executor.adopt_task_instances.duration",
tags=prune_dict({"team_name": self.team_name})
+ ):
adopted_tis: list[TaskInstance] = []
if task_arns := [ti.external_executor_id for ti in tis if
ti.external_executor_id]:
diff --git
a/providers/amazon/tests/unit/amazon/aws/executors/aws_lambda/test_lambda_executor.py
b/providers/amazon/tests/unit/amazon/aws/executors/aws_lambda/test_lambda_executor.py
index 284865285f1..6613b8e39d7 100644
---
a/providers/amazon/tests/unit/amazon/aws/executors/aws_lambda/test_lambda_executor.py
+++
b/providers/amazon/tests/unit/amazon/aws/executors/aws_lambda/test_lambda_executor.py
@@ -1082,6 +1082,33 @@ class TestAwsLambdaExecutor:
assert len(not_adopted) == 0
+ @pytest.mark.parametrize(
+ ("team_name", "expected_tags"),
+ [
+ pytest.param(None, {}, id="without_team"),
+ pytest.param(
+ "team_a",
+ {"team_name": "team_a"},
+ id="with_team",
+ marks=pytest.mark.skipif(
+ not AIRFLOW_V_3_1_PLUS, reason="Multi-team support
requires Airflow 3.1+"
+ ),
+ ),
+ ],
+ )
+ @mock.patch.object(lambda_executor.Stats, "timer")
+ def test_try_adopt_task_instances_emits_team_name_tag(
+ self, mock_timer, mock_executor, team_name, expected_tags
+ ):
+ """Test that the adopt task instances duration metric is tagged with
the team name."""
+ mock_executor.team_name = team_name
+
+ mock_executor.try_adopt_task_instances([])
+
+ mock_timer.assert_called_once_with(
+ "lambda_executor.adopt_task_instances.duration", tags=expected_tags
+ )
+
@mock.patch("airflow.providers.amazon.aws.executors.aws_lambda.lambda_executor.timezone")
def test_end_timeout(self, mock_timezone, mock_executor, mock_airflow_key):
"""Test that executor can end successfully; waiting for all workloads
to naturally exit."""
diff --git
a/providers/amazon/tests/unit/amazon/aws/executors/batch/test_batch_executor.py
b/providers/amazon/tests/unit/amazon/aws/executors/batch/test_batch_executor.py
index 211ed64acf6..f5800232610 100644
---
a/providers/amazon/tests/unit/amazon/aws/executors/batch/test_batch_executor.py
+++
b/providers/amazon/tests/unit/amazon/aws/executors/batch/test_batch_executor.py
@@ -931,6 +931,31 @@ class TestAwsBatchExecutor:
assert submit_kwargs["jobDefinition"] == "some-job-def"
assert submit_kwargs["jobName"] == "some-job-name"
+ @pytest.mark.parametrize(
+ ("team_name", "expected_tags"),
+ [
+ pytest.param(None, {}, id="without_team"),
+ pytest.param(
+ "team_a",
+ {"team_name": "team_a"},
+ id="with_team",
+ marks=pytest.mark.skipif(
+ not AIRFLOW_V_3_1_PLUS, reason="Multi-team support
requires Airflow 3.1+"
+ ),
+ ),
+ ],
+ )
+ @mock.patch.object(batch_executor.Stats, "timer")
+ def test_try_adopt_task_instances_emits_team_name_tag(
+ self, mock_timer, mock_executor, team_name, expected_tags
+ ):
+ """Test that the adopt task instances duration metric is tagged with
the team name."""
+ mock_executor.team_name = team_name
+
+ mock_executor.try_adopt_task_instances([])
+
+
mock_timer.assert_called_once_with("batch_executor.adopt_task_instances.duration",
tags=expected_tags)
+
class TestBatchExecutorConfig:
@staticmethod
diff --git
a/providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py
b/providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py
index de5a17acf96..b7c25d51277 100644
--- a/providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py
+++ b/providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py
@@ -60,7 +60,7 @@ from airflow.version import version as airflow_version_str
from tests_common import RUNNING_TESTS_AGAINST_AIRFLOW_PACKAGES
from tests_common.test_utils.config import conf_vars
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_3_PLUS
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_1_PLUS, AIRFLOW_V_3_3_PLUS
airflow_version = VersionInfo(*map(int, airflow_version_str.split(".")[:3]))
@@ -1334,6 +1334,31 @@ class TestAwsEcsExecutor:
# The remaining one task is unable to be adopted.
assert len(not_adopted_tasks) == 1
+ @pytest.mark.parametrize(
+ ("team_name", "expected_tags"),
+ [
+ pytest.param(None, {}, id="without_team"),
+ pytest.param(
+ "team_a",
+ {"team_name": "team_a"},
+ id="with_team",
+ marks=pytest.mark.skipif(
+ not AIRFLOW_V_3_1_PLUS, reason="Multi-team support
requires Airflow 3.1+"
+ ),
+ ),
+ ],
+ )
+ @mock.patch.object(ecs_executor.Stats, "timer")
+ def test_try_adopt_task_instances_emits_team_name_tag(
+ self, mock_timer, mock_executor, team_name, expected_tags
+ ):
+ """Test that the adopt task instances duration metric is tagged with
the team name."""
+ mock_executor.team_name = team_name
+
+ mock_executor.try_adopt_task_instances([])
+
+
mock_timer.assert_called_once_with("ecs_executor.adopt_task_instances.duration",
tags=expected_tags)
+
class TestEcsExecutorConfig:
@pytest.fixture
