diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 5ad3618048f13..bcec1d541d3c3 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -45,6 +45,7 @@ from tabulate import tabulate from uuid6 import uuid7 +from airflow._shared.observability.metrics.dual_stats_manager import DualStatsManager from airflow._shared.observability.metrics.stats import Stats, normalize_name_for_stats from airflow._shared.timezones import timezone from airflow.api_fastapi.execution_api.app import InProcessExecutionAPI @@ -1288,11 +1289,11 @@ def process_parse_results( file_name = normalize_name_for_stats(Path(relative_fileloc).stem) # bundle_name is included to distinguish files with the same name across different bundles normalized_bundle = normalize_name_for_stats(bundle_name) - Stats.timing(f"dag_processing.last_duration.{normalized_bundle}.{file_name}", stat.last_duration) - Stats.timing( + DualStatsManager.timing( "dag_processing.last_duration", stat.last_duration, - tags={"file_name": file_name, "bundle_name": normalized_bundle}, + tags={}, + legacy_name_tags={"bundle_name": normalized_bundle, "file_name": file_name}, ) if parsing_result is None: diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index c667631756edc..aafe0691e5147 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -886,7 +886,7 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) - "pool.starving_tasks", num_starving_tasks, tags={}, - extra_tags={"pool_name": pool_name}, + legacy_name_tags={"pool_name": pool_name}, ) Stats.gauge("scheduler.tasks.starving", num_starving_tasks_total) @@ -2145,7 +2145,7 @@ def _update_state(dag: SerializedDAG, dag_run: DagRun): "dagrun.schedule_delay", schedule_delay, tags={}, - extra_tags={"dag_id": dag.dag_id}, + legacy_name_tags={"dag_id": dag.dag_id}, ) # cache saves time during scheduling of many dag_runs for same dag @@ -2289,7 +2289,7 @@ def _schedule_dag_run( "dagrun.duration.failed", duration, tags={}, - extra_tags={"dag_id": dag_run.dag_id}, + legacy_name_tags={"dag_id": dag_run.dag_id}, ) return callback_to_execute @@ -2584,7 +2584,7 @@ def _emit_ti_metrics(self, session: Session = NEW_SESSION) -> None: f"ti.{state}", float(count), tags={}, - extra_tags={"queue": queue, "dag_id": dag_id, "task_id": task_id}, + legacy_name_tags={"queue": queue, "dag_id": dag_id, "task_id": task_id}, ) for prev_key in self.previous_ti_metrics[state]: @@ -2595,7 +2595,7 @@ def _emit_ti_metrics(self, session: Session = NEW_SESSION) -> None: f"ti.{state}", 0, tags={}, - extra_tags={"queue": queue, "dag_id": dag_id, "task_id": task_id}, + legacy_name_tags={"queue": queue, "dag_id": dag_id, "task_id": task_id}, ) self.previous_ti_metrics[state] = ti_metrics @@ -2617,31 +2617,31 @@ def _emit_pool_metrics(self, session: Session = NEW_SESSION) -> None: "pool.open_slots", slot_stats["open"], tags={}, - extra_tags={"pool_name": normalized_pool_name}, + legacy_name_tags={"pool_name": normalized_pool_name}, ) DualStatsManager.gauge( "pool.queued_slots", slot_stats["queued"], tags={}, - extra_tags={"pool_name": normalized_pool_name}, + legacy_name_tags={"pool_name": normalized_pool_name}, ) DualStatsManager.gauge( "pool.running_slots", slot_stats["running"], tags={}, - extra_tags={"pool_name": normalized_pool_name}, + legacy_name_tags={"pool_name": normalized_pool_name}, ) DualStatsManager.gauge( "pool.deferred_slots", slot_stats["deferred"], tags={}, - extra_tags={"pool_name": normalized_pool_name}, + legacy_name_tags={"pool_name": normalized_pool_name}, ) DualStatsManager.gauge( "pool.scheduled_slots", slot_stats["scheduled"], tags={}, - extra_tags={"pool_name": normalized_pool_name}, + legacy_name_tags={"pool_name": normalized_pool_name}, ) @provide_session diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py b/airflow-core/src/airflow/jobs/triggerer_job_runner.py index 1406283c05cb3..894748ea9bdca 100644 --- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py +++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py @@ -615,7 +615,7 @@ def emit_metrics(self): "triggers.running", len(self.running_triggers), tags={}, - extra_tags={"hostname": self.job.hostname}, + legacy_name_tags={"hostname": self.job.hostname}, ) capacity_left = self.capacity - len(self.running_triggers) @@ -623,7 +623,7 @@ def emit_metrics(self): "triggerer.capacity_left", capacity_left, tags={}, - extra_tags={"hostname": self.job.hostname}, + legacy_name_tags={"hostname": self.job.hostname}, ) def update_triggers(self, requested_trigger_ids: set[int]): diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index 61242e45390d6..b4ec0cf3fa6a6 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/airflow-core/src/airflow/models/dagrun.py @@ -1083,7 +1083,7 @@ def recalculate(self) -> _UnfinishedStates: with DualStatsManager.timer( "dagrun.dependency-check", tags={}, - extra_tags=self.stats_tags, + legacy_name_tags=self.stats_tags, ): dag = self.get_dag() info = self.task_instance_scheduling_decisions(session) @@ -1562,10 +1562,12 @@ def _emit_true_scheduling_delay_stats_for_finished_state(self, finished_tis: lis else: true_delay = first_start_date - self.run_after if true_delay.total_seconds() > 0: - Stats.timing( - f"dagrun.{dag.dag_id}.first_task_scheduling_delay", true_delay, tags=self.stats_tags + DualStatsManager.timing( + "dagrun.first_task_scheduling_delay", + true_delay, + tags=self.stats_tags, + legacy_name_tags={"dag_id": dag.dag_id}, ) - Stats.timing("dagrun.first_task_scheduling_delay", true_delay, tags=self.stats_tags) except Exception: self.log.warning("Failed to record first_task_scheduling_delay metric:", exc_info=True) @@ -1584,7 +1586,7 @@ def _emit_duration_stats_for_finished_state(self): f"dagrun.duration.{self.state}", dt=duration, tags=self.stats_tags, - extra_tags={"dag_id": self.dag_id}, + legacy_name_tags={"dag_id": self.dag_id}, ) @provide_session @@ -1664,7 +1666,7 @@ def _check_for_removed_or_restored_tasks( DualStatsManager.incr( "task_restored_to_dag", tags=self.stats_tags, - extra_tags={"dag_id": dag.dag_id}, + legacy_name_tags={"dag_id": dag.dag_id}, ) ti.state = None except AirflowException: @@ -1675,7 +1677,7 @@ def _check_for_removed_or_restored_tasks( DualStatsManager.incr( "task_removed_from_dag", tags=self.stats_tags, - extra_tags={"dag_id": dag.dag_id}, + legacy_name_tags={"dag_id": dag.dag_id}, ) ti.state = TaskInstanceState.REMOVED continue @@ -1843,7 +1845,7 @@ def _create_task_instances( "task_instance_created", count, tags=self.stats_tags, - extra_tags={"task_type": task_type}, + legacy_name_tags={"task_type": task_type}, ) session.flush() except IntegrityError: diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index 32daf41aec3e7..b0ce2db7963b1 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/airflow-core/src/airflow/models/taskinstance.py @@ -1314,7 +1314,7 @@ def emit_state_change_metric(self, new_state: TaskInstanceState) -> None: f"task.{metric_name}", timing, tags={}, - extra_tags={"task_id": self.task_id, "dag_id": self.dag_id, "queue": self.queue}, + legacy_name_tags={"task_id": self.task_id, "dag_id": self.dag_id, "queue": self.queue}, ) def clear_next_method_args(self) -> None: @@ -1534,7 +1534,7 @@ def fetch_handle_failure_context( DualStatsManager.incr( "operator_failures", tags=ti.stats_tags, - extra_tags={"operator_name": ti.operator}, + legacy_name_tags={"operator_name": ti.operator}, ) Stats.incr("ti_failures", tags=ti.stats_tags) diff --git a/airflow-core/tests/unit/core/test_dual_stats_manager.py b/airflow-core/tests/unit/core/test_dual_stats_manager.py index f7024c72cd8d5..f90bb0f66adbb 100644 --- a/airflow-core/tests/unit/core/test_dual_stats_manager.py +++ b/airflow-core/tests/unit/core/test_dual_stats_manager.py @@ -119,79 +119,81 @@ def test_get_dict_with_defined_args( assert sorted(args_dict) == sorted(expected_args_dict) @pytest.mark.parametrize( - ("args_dict", "tags", "extra_tags", "expected_args_dict"), + ("args_dict", "tags", "legacy_name_tags", "expected_args_dict"), [ pytest.param( {"count": 1}, {"test": True}, {}, {"count": 1, "tags": {"test": True}}, - id="no_extra_tags", + id="no_legacy_name_tags", ), pytest.param( {}, {"test": True}, {}, {"tags": {"test": True}}, - id="no_args_no_extra_but_tags", + id="no_args_no_legacy_but_tags", ), pytest.param( {}, {"test": True}, - {"test_extra": True}, - {"tags": {"test": True, "test_extra": True}}, - id="no_args_but_tags_and_extra", + {"test_legacy": True}, + {"tags": {"test": True, "test_legacy": True}}, + id="no_args_but_tags_and_legacy", ), pytest.param( {"count": 1}, {"test": True}, {}, {"count": 1, "tags": {"test": True}}, - id="no_args_no_tags_but_extra_tags", + id="no_args_no_tags_but_legacy_name_tags", ), pytest.param( {"count": 1}, {"test": True}, - {"test_extra": True}, - {"count": 1, "tags": {"test": True, "test_extra": True}}, + {"test_legacy": True}, + {"count": 1, "tags": {"test": True, "test_legacy": True}}, id="all_params_provided", ), pytest.param( {"count": 1, "rate": 3}, {"test1": True, "test2": False}, - {"test_extra1": True, "test_extra2": False, "test_extra3": True}, + {"test_legacy1": True, "test_legacy2": False, "test_legacy3": True}, { "count": 1, "rate": 3, "tags": { "test1": True, "test2": False, - "test_extra1": True, - "test_extra2": False, - "test_extra3": True, + "test_legacy1": True, + "test_legacy2": False, + "test_legacy3": True, }, }, id="multiple_params", ), ], ) - def test_get_args_dict_with_extra_tags_if_set( + def test_get_args_dict_with_legacy_name_tags_if_set( self, args_dict: dict[str, Any] | None, tags: dict[str, Any] | None, - extra_tags: dict[str, Any] | None, + legacy_name_tags: dict[str, Any] | None, expected_args_dict: dict[str, Any], ): - dict_full = dual_stats_manager._get_args_dict_with_extra_tags_if_set(args_dict, tags, extra_tags) + dict_full = dual_stats_manager._get_args_dict_with_legacy_name_tags_if_set( + args_dict, tags, legacy_name_tags + ) assert sorted(dict_full) == sorted(expected_args_dict) @pytest.mark.parametrize( - ("tags", "extra_tags", "expected_tags_dict"), + ("tags", "legacy_name_tags", "expected_tags_dict"), [ pytest.param( {"test": True}, - {"test_extra": True}, - {"test": True, "test_extra": True}, + {"test_legacy": True}, + {"test": True, "test_legacy": True}, id="all_params_provided", ), pytest.param( @@ -208,31 +210,31 @@ def test_get_args_dict_with_extra_tags_if_set( ), pytest.param( {}, - {"test_extra": True}, - {"test_extra": True}, - id="only_extra", + {"test_legacy": True}, + {"test_legacy": True}, + id="only_legacy", ), pytest.param( {"test1": True, "test2": False}, - {"test_extra1": True, "test_extra2": False, "test_extra3": True}, + {"test_legacy1": True, "test_legacy2": False, "test_legacy3": True}, { "test1": True, "test2": False, - "test_extra1": True, - "test_extra2": False, - "test_extra3": True, + "test_legacy1": True, + "test_legacy2": False, + "test_legacy3": True, }, id="multiple_params", ), ], ) - def test_get_tags_with_extra( + def test_get_tags_with_legacy( self, tags: dict[str, Any] | None, - extra_tags: dict[str, Any] | None, + legacy_name_tags: dict[str, Any] | None, expected_tags_dict: dict[str, Any], ): - tags_full = dual_stats_manager._get_tags_with_extra(tags, extra_tags) + tags_full = dual_stats_manager._get_tags_with_legacy(tags, legacy_name_tags) assert sorted(tags_full) == sorted(expected_tags_dict) @pytest.mark.parametrize( 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 58f516ee1229e..dcf3184bd19f6 100644 --- a/providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py +++ b/providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py @@ -29,6 +29,7 @@ from airflow.executors.base_executor import BaseExecutor from airflow.models.taskinstance import TaskInstance from airflow.providers.common.compat.sdk import Stats, timezone +from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_2_PLUS from airflow.providers.edge3.models.db import EdgeDBManager, check_db_manager_config from airflow.providers.edge3.models.edge_job import EdgeJobModel from airflow.providers.edge3.models.edge_logs import EdgeLogsModel @@ -192,11 +193,16 @@ def _update_orphaned_jobs(self, session: Session) -> bool: "queue": job.queue, "state": str(TaskInstanceState.FAILED), } - Stats.incr( - f"edge_worker.ti.finish.{job.queue}.{TaskInstanceState.FAILED}.{job.dag_id}.{job.task_id}", - tags=tags, - ) - Stats.incr("edge_worker.ti.finish", tags=tags) + if AIRFLOW_V_3_2_PLUS: + from airflow.sdk.observability.stats import DualStatsManager + + DualStatsManager.incr("edge_worker.ti.finish", tags={}, legacy_name_tags=tags) + else: + Stats.incr( + f"edge_worker.ti.finish.{job.queue}.{TaskInstanceState.FAILED}.{job.dag_id}.{job.task_id}", + tags=tags, + ) + Stats.incr("edge_worker.ti.finish", tags=tags) return bool(lifeless_jobs) diff --git a/providers/edge3/src/airflow/providers/edge3/models/edge_worker.py b/providers/edge3/src/airflow/providers/edge3/models/edge_worker.py index 5b037f2903cb0..e46b510b54a36 100644 --- a/providers/edge3/src/airflow/providers/edge3/models/edge_worker.py +++ b/providers/edge3/src/airflow/providers/edge3/models/edge_worker.py @@ -27,13 +27,9 @@ from sqlalchemy.orm import Mapped from airflow.providers.common.compat.sdk import AirflowException, Stats, timezone -from airflow.providers.edge3.models.edge_base import Base - -try: - from airflow.sdk.observability.stats import DualStatsManager -except ImportError: - DualStatsManager = None # type: ignore[assignment,misc] # Airflow < 3.2 compat from airflow.providers.common.compat.sqlalchemy.orm import mapped_column +from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_2_PLUS +from airflow.providers.edge3.models.edge_base import Base from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.providers_configuration_loader import providers_configuration_loaded from airflow.utils.session import NEW_SESSION, provide_session @@ -179,47 +175,49 @@ def set_metrics( EdgeWorkerState.OFFLINE_MAINTENANCE, ) - if DualStatsManager is not None: + if AIRFLOW_V_3_2_PLUS: + from airflow.sdk.observability.stats import DualStatsManager + DualStatsManager.gauge( "edge_worker.connected", int(connected), tags={}, - extra_tags={"worker_name": worker_name}, + legacy_name_tags={"worker_name": worker_name}, ) DualStatsManager.gauge( "edge_worker.maintenance", int(maintenance), tags={}, - extra_tags={"worker_name": worker_name}, + legacy_name_tags={"worker_name": worker_name}, ) DualStatsManager.gauge( "edge_worker.jobs_active", jobs_active, tags={}, - extra_tags={"worker_name": worker_name}, + legacy_name_tags={"worker_name": worker_name}, ) DualStatsManager.gauge( "edge_worker.concurrency", concurrency, tags={}, - extra_tags={"worker_name": worker_name}, + legacy_name_tags={"worker_name": worker_name}, ) DualStatsManager.gauge( "edge_worker.free_concurrency", free_concurrency, tags={}, - extra_tags={"worker_name": worker_name}, + legacy_name_tags={"worker_name": worker_name}, ) DualStatsManager.gauge( "edge_worker.num_queues", len(queues), tags={}, - extra_tags={"worker_name": worker_name, "queues": ",".join(queues)}, + legacy_name_tags={"worker_name": worker_name, "queues": ",".join(queues)}, ) else: Stats.gauge(f"edge_worker.connected.{worker_name}", int(connected)) diff --git a/providers/edge3/src/airflow/providers/edge3/worker_api/routes/jobs.py b/providers/edge3/src/airflow/providers/edge3/worker_api/routes/jobs.py index 1191d068014a0..539c53a30ee8b 100644 --- a/providers/edge3/src/airflow/providers/edge3/worker_api/routes/jobs.py +++ b/providers/edge3/src/airflow/providers/edge3/worker_api/routes/jobs.py @@ -27,11 +27,7 @@ from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc from airflow.executors.workloads import ExecuteTask from airflow.providers.common.compat.sdk import Stats, timezone - -try: - from airflow.sdk.observability.stats import DualStatsManager -except ImportError: - DualStatsManager = None # type: ignore[assignment,misc] # Airflow < 3.2 compat +from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_2_PLUS from airflow.providers.edge3.models.edge_job import EdgeJobModel from airflow.providers.edge3.worker_api.auth import jwt_token_authorization_rest from airflow.providers.edge3.worker_api.datamodels import ( @@ -41,6 +37,9 @@ ) from airflow.utils.state import TaskInstanceState +if AIRFLOW_V_3_2_PLUS: + from airflow.sdk.observability.stats import DualStatsManager + jobs_router = AirflowRouter(tags=["Jobs"], prefix="/jobs") @@ -91,8 +90,8 @@ def fetch( session.commit() # Edge worker does not backport emitted Airflow metrics, so export some metrics tags = {"dag_id": job.dag_id, "task_id": job.task_id, "queue": job.queue} - if DualStatsManager is not None: - DualStatsManager.incr("edge_worker.ti.start", tags=tags) + if AIRFLOW_V_3_2_PLUS: + DualStatsManager.incr("edge_worker.ti.start", tags={}, legacy_name_tags=tags) else: Stats.incr(f"edge_worker.ti.start.{job.queue}.{job.dag_id}.{job.task_id}", tags=tags) Stats.incr("edge_worker.ti.start", tags=tags) @@ -146,13 +145,10 @@ def state( "dag_id": job.dag_id, "task_id": job.task_id, "queue": job.queue, - "state": str(state), + "state": state, } - if DualStatsManager is not None: - DualStatsManager.incr( - "edge_worker.ti.finish", - tags=tags, - ) + if AIRFLOW_V_3_2_PLUS: + DualStatsManager.incr("edge_worker.ti.finish", tags={}, legacy_name_tags=tags) else: Stats.incr( f"edge_worker.ti.finish.{job.queue}.{state}.{job.dag_id}.{job.task_id}", diff --git a/providers/edge3/src/airflow/providers/edge3/worker_api/routes/worker.py b/providers/edge3/src/airflow/providers/edge3/worker_api/routes/worker.py index 34368c98c4ae2..d61d7888a45c8 100644 --- a/providers/edge3/src/airflow/providers/edge3/worker_api/routes/worker.py +++ b/providers/edge3/src/airflow/providers/edge3/worker_api/routes/worker.py @@ -27,11 +27,7 @@ from airflow.api_fastapi.common.router import AirflowRouter from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc from airflow.providers.common.compat.sdk import Stats, timezone - -try: - from airflow.sdk.observability.stats import DualStatsManager -except ImportError: - DualStatsManager = None # type: ignore[assignment,misc] # Airflow < 3.2 compat +from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_2_PLUS from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel, EdgeWorkerState, set_metrics from airflow.providers.edge3.worker_api.auth import jwt_token_authorization_rest from airflow.providers.edge3.worker_api.datamodels import ( @@ -217,13 +213,15 @@ def set_state( worker.sysinfo = json.dumps(body.sysinfo) worker.last_update = timezone.utcnow() session.commit() - if DualStatsManager is not None: + if AIRFLOW_V_3_2_PLUS: + from airflow.sdk.observability.stats import DualStatsManager + DualStatsManager.incr( "edge_worker.heartbeat_count", 1, 1, tags={}, - extra_tags={"worker_name": worker_name}, + legacy_name_tags={"worker_name": worker_name}, ) else: Stats.incr(f"edge_worker.heartbeat_count.{worker_name}", 1, 1) 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 c38dffed3e918..11ea9d93e4db1 100644 --- a/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py +++ b/providers/edge3/tests/unit/edge3/executors/test_edge_executor.py @@ -25,7 +25,7 @@ from airflow.configuration import conf from airflow.models.taskinstancekey import TaskInstanceKey -from airflow.providers.common.compat.sdk import Stats, timezone +from airflow.providers.common.compat.sdk import timezone from airflow.providers.edge3.executors.edge_executor import EdgeExecutor from airflow.providers.edge3.models.edge_job import EdgeJobModel from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel, EdgeWorkerState @@ -33,6 +33,35 @@ from airflow.utils.state import TaskInstanceState from tests_common.test_utils.config import conf_vars +from tests_common.test_utils.version_compat import AIRFLOW_V_3_2_PLUS + +if AIRFLOW_V_3_2_PLUS: + from airflow.sdk._shared.observability.metrics.dual_stats_manager import DualStatsManager # noqa: F401 + + stats_reference = "airflow.sdk._shared.observability.metrics.dual_stats_manager.DualStatsManager" + expected_incr_kwargs = { + "tags": {}, + "legacy_name_tags": { + "dag_id": "test_dag", + "queue": "default", + "state": "failed", + "task_id": "started_running_orphaned", + }, + } + expected_call_count = 1 +else: + from airflow.providers.common.compat.sdk import Stats + + stats_reference = f"{Stats.__module__}.Stats" + expected_incr_kwargs = { + "tags": { + "dag_id": "test_dag", + "queue": "default", + "state": "failed", + "task_id": "started_running_orphaned", + }, + } + expected_call_count = 2 pytestmark = pytest.mark.db_test @@ -57,7 +86,7 @@ def get_test_executor(self, pool_slots=1): return (executor, key) - @patch(f"{Stats.__module__}.Stats.incr") + @patch(f"{stats_reference}.incr") def test_sync_orphaned_tasks(self, mock_stats_incr): executor = EdgeExecutor() @@ -93,16 +122,8 @@ def test_sync_orphaned_tasks(self, mock_stats_incr): executor.sync() - mock_stats_incr.assert_called_with( - "edge_worker.ti.finish", - tags={ - "dag_id": "test_dag", - "queue": "default", - "state": "failed", - "task_id": "started_running_orphaned", - }, - ) - mock_stats_incr.call_count == 2 + mock_stats_incr.assert_called_with("edge_worker.ti.finish", **expected_incr_kwargs) + assert mock_stats_incr.call_count == expected_call_count with create_session() as session: jobs = session.scalars(select(EdgeJobModel)).all() diff --git a/providers/edge3/tests/unit/edge3/worker_api/routes/test_jobs.py b/providers/edge3/tests/unit/edge3/worker_api/routes/test_jobs.py index 667ec184c0aa5..86140ae0dc98d 100644 --- a/providers/edge3/tests/unit/edge3/worker_api/routes/test_jobs.py +++ b/providers/edge3/tests/unit/edge3/worker_api/routes/test_jobs.py @@ -27,29 +27,47 @@ from airflow.utils.session import create_session from airflow.utils.state import TaskInstanceState +from tests_common.test_utils.version_compat import AIRFLOW_V_3_2_PLUS + if TYPE_CHECKING: from sqlalchemy.orm import Session -try: +DAG_ID = "my_dag" +TASK_ID = "my_task" +RUN_ID = "manual__2024-11-24T21:03:01+01:00" +QUEUE = "test" + +if AIRFLOW_V_3_2_PLUS: from airflow.sdk._shared.observability.metrics.dual_stats_manager import DualStatsManager # noqa: F401 stats_reference = "airflow.sdk._shared.observability.metrics.dual_stats_manager.DualStatsManager" + expected_incr_kwargs = { + "tags": {}, + "legacy_name_tags": { + "dag_id": DAG_ID, + "queue": QUEUE, + "state": TaskInstanceState.SUCCESS, + "task_id": TASK_ID, + }, + } expected_call_count = 1 -except ImportError: +else: from airflow.providers.common.compat.sdk import Stats stats_reference = f"{Stats.__module__}.Stats" + expected_incr_kwargs = { + "tags": { + "dag_id": DAG_ID, + "queue": QUEUE, + "state": TaskInstanceState.SUCCESS, + "task_id": TASK_ID, + }, + } expected_call_count = 2 pytestmark = pytest.mark.db_test -DAG_ID = "my_dag" -TASK_ID = "my_task" -RUN_ID = "manual__2024-11-24T21:03:01+01:00" -QUEUE = "test" - - class TestJobsApiRoutes: @pytest.fixture(autouse=True) def setup_test_cases(self, dag_maker, session: Session): @@ -94,15 +112,7 @@ def test_state(self, mock_stats_incr, session: Session): session=session, ) - mock_stats_incr.assert_called_with( - "edge_worker.ti.finish", - tags={ - "dag_id": DAG_ID, - "queue": QUEUE, - "state": TaskInstanceState.SUCCESS, - "task_id": TASK_ID, - }, - ) + mock_stats_incr.assert_called_with("edge_worker.ti.finish", **expected_incr_kwargs) assert mock_stats_incr.call_count == expected_call_count db_job: EdgeJobModel | None = session.scalar(select(EdgeJobModel)) diff --git a/providers/openlineage/src/airflow/providers/openlineage/plugins/adapter.py b/providers/openlineage/src/airflow/providers/openlineage/plugins/adapter.py index 076730e9bf6d1..8047e6a0abb3a 100644 --- a/providers/openlineage/src/airflow/providers/openlineage/plugins/adapter.py +++ b/providers/openlineage/src/airflow/providers/openlineage/plugins/adapter.py @@ -36,11 +36,7 @@ ) from airflow.providers.common.compat.sdk import Stats, conf as airflow_conf - -try: - from airflow.sdk.observability.stats import DualStatsManager -except ImportError: - DualStatsManager = None # type: ignore[assignment,misc] # Airflow < 3.2 compat +from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_2_PLUS from airflow.providers.openlineage import __version__ as OPENLINEAGE_PROVIDER_VERSION, conf from airflow.providers.openlineage.utils.utils import ( OpenLineageRedactor, @@ -163,11 +159,13 @@ def emit(self, event: RunEvent): try: with ExitStack() as stack: - if DualStatsManager is not None: + if AIRFLOW_V_3_2_PLUS: + from airflow.sdk.observability.stats import DualStatsManager + stack.enter_context( DualStatsManager.timer( "ol.emit.attempts", - extra_tags={"event_type": event_type, "transport_type": transport_type}, + legacy_name_tags={"event_type": event_type, "transport_type": transport_type}, ) ) else: diff --git a/providers/openlineage/tests/unit/openlineage/plugins/test_adapter.py b/providers/openlineage/tests/unit/openlineage/plugins/test_adapter.py index 4c3e5ba67238e..368ab810c1d4a 100644 --- a/providers/openlineage/tests/unit/openlineage/plugins/test_adapter.py +++ b/providers/openlineage/tests/unit/openlineage/plugins/test_adapter.py @@ -60,11 +60,11 @@ from tests_common.test_utils.taskinstance import create_task_instance from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_2_PLUS -try: +if AIRFLOW_V_3_2_PLUS: from airflow.sdk._shared.observability.metrics.dual_stats_manager import DualStatsManager # noqa: F401 stats_reference = "airflow.sdk._shared.observability.metrics.dual_stats_manager.DualStatsManager" -except ImportError: +else: stats_reference = "airflow.providers.openlineage.plugins.adapter.Stats" @@ -207,7 +207,7 @@ def test_emit_start_event(mock_stats_incr, mock_stats_timer): mock_stats_incr.assert_not_called() if AIRFLOW_V_3_2_PLUS: mock_stats_timer.assert_called_with( - "ol.emit.attempts", extra_tags={"event_type": ANY, "transport_type": ANY} + "ol.emit.attempts", legacy_name_tags={"event_type": ANY, "transport_type": ANY} ) else: mock_stats_timer.assert_called_with("ol.emit.attempts") @@ -325,7 +325,7 @@ def test_emit_start_event_with_additional_information(mock_stats_incr, mock_stat mock_stats_incr.assert_not_called() if AIRFLOW_V_3_2_PLUS: mock_stats_timer.assert_called_with( - "ol.emit.attempts", extra_tags={"event_type": ANY, "transport_type": ANY} + "ol.emit.attempts", legacy_name_tags={"event_type": ANY, "transport_type": ANY} ) else: mock_stats_timer.assert_called_with("ol.emit.attempts") @@ -388,7 +388,7 @@ def test_emit_complete_event(mock_stats_incr, mock_stats_timer): mock_stats_incr.assert_not_called() if AIRFLOW_V_3_2_PLUS: mock_stats_timer.assert_called_with( - "ol.emit.attempts", extra_tags={"event_type": ANY, "transport_type": ANY} + "ol.emit.attempts", legacy_name_tags={"event_type": ANY, "transport_type": ANY} ) else: mock_stats_timer.assert_called_with("ol.emit.attempts") @@ -508,7 +508,7 @@ def test_emit_complete_event_with_additional_information(mock_stats_incr, mock_s mock_stats_incr.assert_not_called() if AIRFLOW_V_3_2_PLUS: mock_stats_timer.assert_called_with( - "ol.emit.attempts", extra_tags={"event_type": ANY, "transport_type": ANY} + "ol.emit.attempts", legacy_name_tags={"event_type": ANY, "transport_type": ANY} ) else: mock_stats_timer.assert_called_with("ol.emit.attempts") @@ -571,7 +571,7 @@ def test_emit_failed_event(mock_stats_incr, mock_stats_timer): mock_stats_incr.assert_not_called() if AIRFLOW_V_3_2_PLUS: mock_stats_timer.assert_called_with( - "ol.emit.attempts", extra_tags={"event_type": ANY, "transport_type": ANY} + "ol.emit.attempts", legacy_name_tags={"event_type": ANY, "transport_type": ANY} ) else: mock_stats_timer.assert_called_with("ol.emit.attempts") @@ -692,7 +692,7 @@ def test_emit_failed_event_with_additional_information(mock_stats_incr, mock_sta mock_stats_incr.assert_not_called() if AIRFLOW_V_3_2_PLUS: mock_stats_timer.assert_called_with( - "ol.emit.attempts", extra_tags={"event_type": ANY, "transport_type": ANY} + "ol.emit.attempts", legacy_name_tags={"event_type": ANY, "transport_type": ANY} ) else: mock_stats_timer.assert_called_with("ol.emit.attempts") @@ -870,7 +870,7 @@ def test_emit_dag_started_event(mock_stats_incr, mock_stats_timer, build_ol_id, mock_stats_incr.assert_not_called() if AIRFLOW_V_3_2_PLUS: mock_stats_timer.assert_called_with( - "ol.emit.attempts", extra_tags={"event_type": ANY, "transport_type": ANY} + "ol.emit.attempts", legacy_name_tags={"event_type": ANY, "transport_type": ANY} ) else: mock_stats_timer.assert_called_with("ol.emit.attempts") @@ -1042,7 +1042,7 @@ def test_emit_dag_complete_event( mock_stats_incr.assert_not_called() if AIRFLOW_V_3_2_PLUS: mock_stats_timer.assert_called_with( - "ol.emit.attempts", extra_tags={"event_type": ANY, "transport_type": ANY} + "ol.emit.attempts", legacy_name_tags={"event_type": ANY, "transport_type": ANY} ) else: mock_stats_timer.assert_called_with("ol.emit.attempts") @@ -1218,7 +1218,7 @@ def test_emit_dag_failed_event( mock_stats_incr.assert_not_called() if AIRFLOW_V_3_2_PLUS: mock_stats_timer.assert_called_with( - "ol.emit.attempts", extra_tags={"event_type": ANY, "transport_type": ANY} + "ol.emit.attempts", legacy_name_tags={"event_type": ANY, "transport_type": ANY} ) else: mock_stats_timer.assert_called_with("ol.emit.attempts") @@ -1238,7 +1238,7 @@ def test_openlineage_adapter_stats_emit_failed( if AIRFLOW_V_3_2_PLUS: mock_stats_timer.assert_called_with( - "ol.emit.attempts", extra_tags={"event_type": ANY, "transport_type": ANY} + "ol.emit.attempts", legacy_name_tags={"event_type": ANY, "transport_type": ANY} ) else: mock_stats_timer.assert_called_with("ol.emit.attempts") diff --git a/shared/observability/src/airflow_shared/observability/metrics/dual_stats_manager.py b/shared/observability/src/airflow_shared/observability/metrics/dual_stats_manager.py index af787596d7ffb..c40891e900c49 100644 --- a/shared/observability/src/airflow_shared/observability/metrics/dual_stats_manager.py +++ b/shared/observability/src/airflow_shared/observability/metrics/dual_stats_manager.py @@ -60,22 +60,22 @@ def _get_dict_with_defined_args( return defined_args_dict -def _get_args_dict_with_extra_tags_if_set( +def _get_args_dict_with_legacy_name_tags_if_set( args_dict: dict[str, Any] | None = None, prov_tags: dict[str, Any] | None = None, - prov_tags_extra: dict[str, Any] | None = None, + prov_tags_legacy: dict[str, Any] | None = None, ) -> dict[str, Any]: """ - Create a new merged tags dict if there are extra tags. + Create a new merged tags dict if there are legacy name tags. The new merged tags dict will replace the existing one, in the args dict. """ # The args_dict already has the base tags. - # If there are no `extra_tags`, this method is basically + # If there are no `legacy_name_tags`, this method is basically # returning the `args_dict` unchanged. args_dict_full = dict(args_dict) if args_dict is not None else {} - tags_full = _get_tags_with_extra(prov_tags, prov_tags_extra) + tags_full = _get_tags_with_legacy(prov_tags, prov_tags_legacy) # Set `tags` only if there's something in `tags_full`. # If it's empty, remove any inherited key. @@ -87,19 +87,19 @@ def _get_args_dict_with_extra_tags_if_set( return args_dict_full -def _get_tags_with_extra( +def _get_tags_with_legacy( prov_tags: dict[str, Any] | None = None, - prov_tags_extra: dict[str, Any] | None = None, + prov_tags_legacy: dict[str, Any] | None = None, ) -> dict[str, Any]: - """Return a new dict with all tags if extra have been provided.""" - # If there are no extra tags then return the original tags. + """Return a new dict with all tags if legacy name tags have been provided.""" + # If there are no legacy name tags then return the original tags. tags_full: dict[str, Any] = {} if prov_tags: tags_full.update(prov_tags) - # If there are `extra_tags`, then add them to the dict. - if prov_tags_extra is not None: - tags_full.update(prov_tags_extra) + # If there are `legacy_name_tags`, then add them to the dict. + if prov_tags_legacy is not None: + tags_full.update(prov_tags_legacy) return tags_full @@ -158,12 +158,21 @@ def incr( rate: int | float | None = None, *, tags: dict[str, Any] | None = None, - extra_tags: dict[str, Any] | None = None, + legacy_name_tags: dict[str, Any] | None = None, ) -> None: + """ + Increases the stat counter. + + Legacy name tags are cross-referenced with the registry yaml file. + The 'tags' and 'legacy_name_tags' dictionaries are merged. + + :param tags: Extra tags that are set to both new and legacy metrics + :param legacy_name_tags: Tags that will be used for the legacy name metric + """ kw = _get_dict_with_defined_args(count, rate, None, tags) - if cls.export_legacy_names and extra_tags is not None: - legacy_stat = cls.get_legacy_stat(stat=stat, variables=extra_tags) + if cls.export_legacy_names and legacy_name_tags is not None: + legacy_stat = cls.get_legacy_stat(stat=stat, variables=legacy_name_tags) if legacy_stat is not None: # Emit legacy metric @@ -171,8 +180,10 @@ def incr( else: raise ValueError(f"Stat '{stat}' doesn't have a legacy name registered in the YAML file.") - kw_with_extra_tags_if_set = _get_args_dict_with_extra_tags_if_set(kw, tags, extra_tags) - Stats.incr(stat, **kw_with_extra_tags_if_set) + kw_with_legacy_name_tags_if_set = _get_args_dict_with_legacy_name_tags_if_set( + kw, tags, legacy_name_tags + ) + Stats.incr(stat, **kw_with_legacy_name_tags_if_set) @classmethod def decr( @@ -182,12 +193,21 @@ def decr( rate: int | float | None = None, *, tags: dict[str, Any] | None = None, - extra_tags: dict[str, Any] | None = None, + legacy_name_tags: dict[str, Any] | None = None, ) -> None: + """ + Decreases the stat counter. + + Legacy name tags are cross-referenced with the registry yaml file. + The 'tags' and 'legacy_name_tags' dictionaries are merged. + + :param tags: Extra tags that are set to both new and legacy metrics + :param legacy_name_tags: Tags that will be used for the legacy name metric + """ kw = _get_dict_with_defined_args(count, rate, None, tags) - if cls.export_legacy_names and extra_tags is not None: - legacy_stat = cls.get_legacy_stat(stat=stat, variables=extra_tags) + if cls.export_legacy_names and legacy_name_tags is not None: + legacy_stat = cls.get_legacy_stat(stat=stat, variables=legacy_name_tags) if legacy_stat is not None: # Emit legacy metric @@ -195,8 +215,10 @@ def decr( else: raise ValueError(f"Stat '{stat}' doesn't have a legacy name registered in the YAML file.") - kw_with_extra_tags_if_set = _get_args_dict_with_extra_tags_if_set(kw, tags, extra_tags) - Stats.decr(stat, **kw_with_extra_tags_if_set) + kw_with_legacy_name_tags_if_set = _get_args_dict_with_legacy_name_tags_if_set( + kw, tags, legacy_name_tags + ) + Stats.decr(stat, **kw_with_legacy_name_tags_if_set) @classmethod def gauge( @@ -207,12 +229,21 @@ def gauge( delta: bool | None = None, *, tags: dict[str, Any] | None = None, - extra_tags: dict[str, Any] | None = None, + legacy_name_tags: dict[str, Any] | None = None, ) -> None: + """ + Gauges the stat. + + Legacy name tags are cross-referenced with the registry yaml file. + The 'tags' and 'legacy_name_tags' dictionaries are merged. + + :param tags: Extra tags that are set to both new and legacy metrics + :param legacy_name_tags: Tags that will be used for the legacy name metric + """ kw = _get_dict_with_defined_args(None, rate, delta, tags) - if cls.export_legacy_names and extra_tags is not None: - legacy_stat = cls.get_legacy_stat(stat=stat, variables=extra_tags) + if cls.export_legacy_names and legacy_name_tags is not None: + legacy_stat = cls.get_legacy_stat(stat=stat, variables=legacy_name_tags) if legacy_stat is not None: # Emit legacy metric @@ -220,8 +251,10 @@ def gauge( else: raise ValueError(f"Stat '{stat}' doesn't have a legacy name registered in the YAML file.") - kw_with_extra_tags_if_set = _get_args_dict_with_extra_tags_if_set(kw, tags, extra_tags) - Stats.gauge(stat, value, **kw_with_extra_tags_if_set) + kw_with_legacy_name_tags_if_set = _get_args_dict_with_legacy_name_tags_if_set( + kw, tags, legacy_name_tags + ) + Stats.gauge(stat, value, **kw_with_legacy_name_tags_if_set) @classmethod def timing( @@ -230,10 +263,19 @@ def timing( dt: DeltaType, *, tags: dict[str, Any] | None = None, - extra_tags: dict[str, Any] | None = None, + legacy_name_tags: dict[str, Any] | None = None, ) -> None: - if cls.export_legacy_names and extra_tags is not None: - legacy_stat = cls.get_legacy_stat(stat=stat, variables=extra_tags) + """ + Stat timing. + + Legacy name tags are cross-referenced with the registry yaml file. + The 'tags' and 'legacy_name_tags' dictionaries are merged. + + :param tags: Extra tags that are set to both new and legacy metrics + :param legacy_name_tags: Tags that will be used for the legacy name metric + """ + if cls.export_legacy_names and legacy_name_tags is not None: + legacy_stat = cls.get_legacy_stat(stat=stat, variables=legacy_name_tags) if legacy_stat is not None: if tags: @@ -243,10 +285,10 @@ def timing( else: raise ValueError(f"Stat '{stat}' doesn't have a legacy name registered in the YAML file.") - tags_with_extra = _get_tags_with_extra(tags, extra_tags) + tags_with_legacy = _get_tags_with_legacy(tags, legacy_name_tags) - if tags_with_extra: - Stats.timing(stat, dt, tags=tags_with_extra) + if tags_with_legacy: + Stats.timing(stat, dt, tags=tags_with_legacy) else: Stats.timing(stat, dt) @@ -255,9 +297,18 @@ def timer( cls, stat: str, tags: dict[str, Any] | None = None, - extra_tags: dict[str, Any] | None = None, + legacy_name_tags: dict[str, Any] | None = None, **kwargs, ): + """ + Stat timer with a context manager. + + Legacy name tags are cross-referenced with the registry yaml file. + The 'tags' and 'legacy_name_tags' dictionaries are merged. + + :param tags: Extra tags that are set to both new and legacy metrics + :param legacy_name_tags: Tags that will be used for the legacy name metric + """ kw = dict(kwargs) if tags is not None: kw["tags"] = tags @@ -265,8 +316,8 @@ def timer( # Used with a context manager. stack = ExitStack() - if cls.export_legacy_names and extra_tags is not None: - legacy_stat = cls.get_legacy_stat(stat=stat, variables=extra_tags) + if cls.export_legacy_names and legacy_name_tags is not None: + legacy_stat = cls.get_legacy_stat(stat=stat, variables=legacy_name_tags) ctx_mg1: AbstractContextManager[Any] = cast( "AbstractContextManager[Any]", Stats.timer(legacy_stat, **kw) @@ -276,7 +327,9 @@ def timer( stack.enter_context(ctx_mg1) - kw_with_extra_tags_if_set = _get_args_dict_with_extra_tags_if_set(kw, tags, extra_tags) - stack.enter_context(Stats.timer(stat, **kw_with_extra_tags_if_set)) + kw_with_legacy_name_tags_if_set = _get_args_dict_with_legacy_name_tags_if_set( + kw, tags, legacy_name_tags + ) + stack.enter_context(Stats.timer(stat, **kw_with_legacy_name_tags_if_set)) return stack diff --git a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml index e171980c41d71..188a740bd6366 100644 --- a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml +++ b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml @@ -262,6 +262,19 @@ metrics: legacy_name: "-" name_variables: [] + - name: "edge_worker.ti.start" + description: "Number of task instances started in an edge worker" + type: "counter" + legacy_name: "edge_worker.ti.start.{queue}.{dag_id}.{task_id}" + name_variables: ["queue", "dag_id", "task_id"] + + - name: "edge_worker.ti.finish" + description: "Number of task instances finished in an edge worker" + type: "counter" + legacy_name: "edge_worker.ti.finish.{queue}.{state}.{dag_id}.{task_id}" + name_variables: ["queue", "state", "dag_id", "task_id"] + + # ========== # Gauges # ========== @@ -486,8 +499,8 @@ metrics: - name: "dag_processing.last_duration" description: "Milliseconds taken to load the given Dag file" type: "timer" - legacy_name: "dag_processing.last_duration.{dag_file}" - name_variables: ["dag_file"] + legacy_name: "dag_processing.last_duration.{bundle_name}.{file_name}" + name_variables: ["bundle_name", "file_name"] - name: "dagrun.duration.success" description: "Milliseconds taken for a DagRun to reach success state" diff --git a/shared/observability/tests/observability/metrics/test_dual_stats_manager.py b/shared/observability/tests/observability/metrics/test_dual_stats_manager.py index 46f2aca515567..f9af2184d3544 100644 --- a/shared/observability/tests/observability/metrics/test_dual_stats_manager.py +++ b/shared/observability/tests/observability/metrics/test_dual_stats_manager.py @@ -119,79 +119,81 @@ def test_get_dict_with_defined_args( assert sorted(args_dict) == sorted(expected_args_dict) @pytest.mark.parametrize( - ("args_dict", "tags", "extra_tags", "expected_args_dict"), + ("args_dict", "tags", "legacy_name_tags", "expected_args_dict"), [ pytest.param( {"count": 1}, {"test": True}, {}, {"count": 1, "tags": {"test": True}}, - id="no_extra_tags", + id="no_legacy_name_tags", ), pytest.param( {}, {"test": True}, {}, {"tags": {"test": True}}, - id="no_args_no_extra_but_tags", + id="no_args_no_legacy_but_tags", ), pytest.param( {}, {"test": True}, - {"test_extra": True}, - {"tags": {"test": True, "test_extra": True}}, - id="no_args_but_tags_and_extra", + {"test_legacy": True}, + {"tags": {"test": True, "test_legacy": True}}, + id="no_args_but_tags_and_legacy", ), pytest.param( {"count": 1}, {"test": True}, {}, {"count": 1, "tags": {"test": True}}, - id="no_args_no_tags_but_extra_tags", + id="no_args_no_tags_but_legacy_name_tags", ), pytest.param( {"count": 1}, {"test": True}, - {"test_extra": True}, - {"count": 1, "tags": {"test": True, "test_extra": True}}, + {"test_legacy": True}, + {"count": 1, "tags": {"test": True, "test_legacy": True}}, id="all_params_provided", ), pytest.param( {"count": 1, "rate": 3}, {"test1": True, "test2": False}, - {"test_extra1": True, "test_extra2": False, "test_extra3": True}, + {"test_legacy1": True, "test_legacy2": False, "test_legacy3": True}, { "count": 1, "rate": 3, "tags": { "test1": True, "test2": False, - "test_extra1": True, - "test_extra2": False, - "test_extra3": True, + "test_legacy1": True, + "test_legacy2": False, + "test_legacy3": True, }, }, id="multiple_params", ), ], ) - def test_get_args_dict_with_extra_tags_if_set( + def test_get_args_dict_with_legacy_name_tags_if_set( self, args_dict: dict[str, Any] | None, tags: dict[str, Any] | None, - extra_tags: dict[str, Any] | None, + legacy_name_tags: dict[str, Any] | None, expected_args_dict: dict[str, Any], ): - dict_full = dual_stats_manager._get_args_dict_with_extra_tags_if_set(args_dict, tags, extra_tags) + dict_full = dual_stats_manager._get_args_dict_with_legacy_name_tags_if_set( + args_dict, tags, legacy_name_tags + ) assert sorted(dict_full) == sorted(expected_args_dict) @pytest.mark.parametrize( - ("tags", "extra_tags", "expected_tags_dict"), + ("tags", "legacy_name_tags", "expected_tags_dict"), [ pytest.param( {"test": True}, - {"test_extra": True}, - {"test": True, "test_extra": True}, + {"test_legacy": True}, + {"test": True, "test_legacy": True}, id="all_params_provided", ), pytest.param( @@ -208,31 +210,31 @@ def test_get_args_dict_with_extra_tags_if_set( ), pytest.param( {}, - {"test_extra": True}, - {"test_extra": True}, - id="only_extra", + {"test_legacy": True}, + {"test_legacy": True}, + id="only_legacy", ), pytest.param( {"test1": True, "test2": False}, - {"test_extra1": True, "test_extra2": False, "test_extra3": True}, + {"test_legacy1": True, "test_legacy2": False, "test_legacy3": True}, { "test1": True, "test2": False, - "test_extra1": True, - "test_extra2": False, - "test_extra3": True, + "test_legacy1": True, + "test_legacy2": False, + "test_legacy3": True, }, id="multiple_params", ), ], ) - def test_get_tags_with_extra( + def test_get_tags_with_legacy( self, tags: dict[str, Any] | None, - extra_tags: dict[str, Any] | None, + legacy_name_tags: dict[str, Any] | None, expected_tags_dict: dict[str, Any], ): - tags_full = dual_stats_manager._get_tags_with_extra(tags, extra_tags) + tags_full = dual_stats_manager._get_tags_with_legacy(tags, legacy_name_tags) assert sorted(tags_full) == sorted(expected_tags_dict) @pytest.mark.parametrize( diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index 674935f5eaecb..f17507f79fc8e 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -41,6 +41,7 @@ from airflow.dag_processing.bundles.base import BaseDagBundle, BundleVersionLock from airflow.dag_processing.bundles.manager import DagBundlesManager +from airflow.sdk._shared.observability.metrics.dual_stats_manager import DualStatsManager from airflow.sdk._shared.observability.metrics.stats import Stats from airflow.sdk.api.client import get_hostname, getuser from airflow.sdk.api.datamodels._generated import ( @@ -1746,8 +1747,12 @@ def finalize( duration_ms = (ti.end_date - ti.start_date).total_seconds() * 1000 stats_tags = {"dag_id": ti.dag_id, "task_id": ti.task_id} - Stats.timing(f"dag.{ti.dag_id}.{ti.task_id}.duration", duration_ms) - Stats.timing("task.duration", duration_ms, tags=stats_tags) + DualStatsManager.timing( + "task.duration", + duration_ms, + tags=stats_tags, + legacy_name_tags={"dag_id": ti.dag_id, "task_id": ti.task_id}, + ) task = ti.task # Pushing xcom for each operator extra links defined on the operator only.