Move the traces and metrics code under a common observability package (apache#56187)

- Create shared/observability with base metrics and traces implementations
- Add task-sdk observability module with own Stats and Trace wrappers
- task SDK no longer imports from airflow-core observability
- Each component reads its own configuration
- Export only Trace in public API (Stats is internal)

---------

Co-authored-by: Jarek Potiuk <jarek@potiuk.com>
Co-authored-by: Amogh Desai <amoghrajesh1999@gmail.com>

Author: 3 people

airflow-core/pyproject.toml

@@ -227,6 +227,7 @@ exclude = [
 "../shared/secrets_masker/src/airflow_shared/secrets_masker" = "src/airflow/_shared/secrets_masker"
 "../shared/timezones/src/airflow_shared/timezones" = "src/airflow/_shared/timezones"
 "../shared/secrets_backend/src/airflow_shared/secrets_backend" = "src/airflow/_shared/secrets_backend"
+"../shared/observability/src/airflow_shared/observability" = "src/airflow/_shared/observability"
 
 [tool.hatch.build.targets.custom]
 path = "./hatch_build.py"
@@ -299,4 +300,5 @@ shared_distributions = [
     "apache-airflow-shared-secrets-backend",
     "apache-airflow-shared-secrets-masker",
     "apache-airflow-shared-timezones",
+    "apache-airflow-shared-observability",
 ]

airflow-core/src/airflow/__init__.py

@@ -89,6 +89,10 @@
     # Deprecated lazy imports
     "AirflowException": (".exceptions", "AirflowException", True),
     "Dataset": (".sdk", "Asset", True),
+    "Stats": (".observability.stats", "Stats", True),
+    "Trace": (".observability.trace", "Trace", True),
+    "metrics": (".observability.metrics", "", True),
+    "traces": (".observability.traces", "", True),
 }
 if TYPE_CHECKING:
     # These objects are imported by PEP-562, however, static analyzers and IDE's

airflow-core/src/airflow/_shared/observability

@@ -0,0 +1 @@
+../../../../shared/observability/src/airflow_shared/observability

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

@@ -38,7 +38,7 @@
     DagScheduleAssetUriReference,
     PartitionedAssetKeyLog,
 )
-from airflow.stats import Stats
+from airflow.observability.stats import Stats
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.sqlalchemy import get_dialect_name
 

airflow-core/src/airflow/cli/commands/daemon_utils.py

@@ -75,7 +75,7 @@ def run_command_with_daemon_option(
 
             with ctx:
                 # in daemon context stats client needs to be reinitialized.
-                from airflow.stats import Stats
+                from airflow.observability.stats import Stats
 
                 Stats.instance = None
                 callback()

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

@@ -58,10 +58,10 @@
 from airflow.models.dagwarning import DagWarning
 from airflow.models.db_callback_request import DbCallbackRequest
 from airflow.models.errors import ParseImportError
+from airflow.observability.stats import Stats
+from airflow.observability.trace import DebugTrace
 from airflow.sdk import SecretCache
 from airflow.sdk.log import init_log_file, logging_processors
-from airflow.stats import Stats
-from airflow.traces.tracer import DebugTrace
 from airflow.utils.file import list_py_file_paths, might_contain_dag
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.net import get_hostname

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

@@ -36,6 +36,7 @@
 )
 from airflow.configuration import conf
 from airflow.dag_processing.dagbag import DagBag
+from airflow.observability.stats import Stats
 from airflow.sdk.exceptions import TaskNotFound
 from airflow.sdk.execution_time.comms import (
     ConnectionResult,
@@ -66,7 +67,6 @@
 from airflow.sdk.execution_time.supervisor import WatchedSubprocess
 from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance, _send_error_email_notification
 from airflow.serialization.serialized_objects import LazyDeserializedDAG, SerializedDAG
-from airflow.stats import Stats
 from airflow.utils.file import iter_airflow_imports
 from airflow.utils.state import TaskInstanceState
 

airflow-core/src/airflow/executors/base_executor.py

@@ -27,15 +27,14 @@
 
 import pendulum
 
+from airflow._shared.observability.traces import NO_TRACE_ID
 from airflow.cli.cli_config import DefaultHelpParser
 from airflow.configuration import conf
 from airflow.executors import workloads
 from airflow.executors.executor_loader import ExecutorLoader
 from airflow.models import Log
-from airflow.stats import Stats
-from airflow.traces import NO_TRACE_ID
-from airflow.traces.tracer import DebugTrace, Trace, add_debug_span, gen_context
-from airflow.traces.utils import gen_span_id_from_ti_key
+from airflow.observability.stats import Stats
+from airflow.observability.trace import DebugTrace, Trace, add_debug_span
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.state import TaskInstanceState
 from airflow.utils.thread_safe_dict import ThreadSafeDict
@@ -419,11 +418,9 @@ def fail(self, key: TaskInstanceKey, info=None) -> None:
         """
         trace_id = Trace.get_current_span().get_span_context().trace_id
         if trace_id != NO_TRACE_ID:
-            span_id = int(gen_span_id_from_ti_key(key, as_int=True))
-            with DebugTrace.start_span(
+            with DebugTrace.start_child_span(
                 span_name="fail",
                 component="BaseExecutor",
-                parent_sc=gen_context(trace_id=trace_id, span_id=span_id),
             ) as span:
                 span.set_attributes(
                     {
@@ -446,11 +443,9 @@ def success(self, key: TaskInstanceKey, info=None) -> None:
         """
         trace_id = Trace.get_current_span().get_span_context().trace_id
         if trace_id != NO_TRACE_ID:
-            span_id = int(gen_span_id_from_ti_key(key, as_int=True))
-            with DebugTrace.start_span(
+            with DebugTrace.start_child_span(
                 span_name="success",
                 component="BaseExecutor",
-                parent_sc=gen_context(trace_id=trace_id, span_id=span_id),
             ) as span:
                 span.set_attributes(
                     {

airflow-core/src/airflow/jobs/dag_processor_job_runner.py

@@ -21,7 +21,7 @@
 
 from airflow.jobs.base_job_runner import BaseJobRunner
 from airflow.jobs.job import Job, perform_heartbeat
-from airflow.stats import Stats
+from airflow.observability.stats import Stats
 from airflow.utils.log.logging_mixin import LoggingMixin
 
 if TYPE_CHECKING:

airflow-core/src/airflow/jobs/job.py

@@ -35,8 +35,8 @@
 from airflow.executors.executor_loader import ExecutorLoader
 from airflow.listeners.listener import get_listener_manager
 from airflow.models.base import ID_LEN, Base
-from airflow.stats import Stats
-from airflow.traces.tracer import DebugTrace, add_debug_span
+from airflow.observability.stats import Stats
+from airflow.observability.trace import DebugTrace, add_debug_span
 from airflow.utils.helpers import convert_camel_to_snake
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.net import get_hostname