Provide an alternative OpenTelemetry implementation for traces that follows standard otel practices

airflow-core/docs/img/airflow_erd.sha256

@@ -1 +1 @@
-040d02de3ade319dccca8e797c88a2d118b3f231e293fbac9b7cd22c40dee94b
+6d34a69fdd22fb001604f5fdbc5d456cf7c25d5657536a74df8d0ac8ac63bd9a

airflow-core/docs/migrations-ref.rst

@@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | Revision ID             | Revises ID       | Airflow Version   | Description                                                  |
 +=========================+==================+===================+==============================================================+
-| ``be2cc2f742cf`` (head) | ``d469d27e2a64`` | ``3.0.0``         | Support bundles in DagPriorityParsingRequest.                |
+| ``ec62e120484d`` (head) | ``be2cc2f742cf`` | ``3.0.0``         | Add new otel span fields.                                    |
++-------------------------+------------------+-------------------+--------------------------------------------------------------+
+| ``be2cc2f742cf``        | ``d469d27e2a64`` | ``3.0.0``         | Support bundles in DagPriorityParsingRequest.                |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | ``d469d27e2a64``        | ``16f7f5ee874e`` | ``3.0.0``         | Use ti_id as FK to TaskReschedule.                           |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+

airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py

@@ -241,6 +241,7 @@ class TaskInstance(StrictBaseModel):
     try_number: int
     map_index: int = -1
     hostname: str | None = None
+    context_carrier: dict | None = None
 
 
 class DagRun(StrictBaseModel):

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

@@ -37,6 +37,7 @@
 from airflow.traces.utils import gen_span_id_from_ti_key, gen_trace_id
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.state import TaskInstanceState
+from airflow.utils.thread_safe_dict import ThreadSafeDict
 
 PARALLELISM: int = conf.getint("core", "PARALLELISM")
 
@@ -116,6 +117,8 @@ class BaseExecutor(LoggingMixin):
     :param parallelism: how many jobs should run at one time.
     """
 
+    active_spans = ThreadSafeDict()
+
     supports_ad_hoc_ti_run: bool = False
     supports_sentry: bool = False
 
@@ -178,6 +181,10 @@ def __init__(self, parallelism: int = PARALLELISM, team_id: str | None = None):
     def __repr__(self):
         return f"{self.__class__.__name__}(parallelism={self.parallelism})"
 
+    @classmethod
+    def set_active_spans(cls, active_spans: ThreadSafeDict):
+        cls.active_spans = active_spans
+
     def start(self):  # pragma: no cover
         """Executors may need to get things started."""
 
@@ -375,7 +382,7 @@ def trigger_tasks(self, open_slots: int) -> None:
         """
         sorted_queue = self.order_queued_tasks_by_priority()
         task_tuples = []
-        workloads = []
+        workload_list = []
 
         for _ in range(min((open_slots, len(self.queued_tasks)))):
             key, item = sorted_queue.pop(0)
@@ -419,16 +426,46 @@ def trigger_tasks(self, open_slots: int) -> None:
                 # TODO: TaskSDK: Compat, remove when KubeExecutor is fully moved over to TaskSDK too.
                 # TODO: TaskSDK: We need to minimum version requirements on executors with Airflow 3.
                 # How/where do we do that? Executor loader?
+                from airflow.executors import workloads
+
+                if isinstance(item, workloads.ExecuteTask) and hasattr(item, "ti"):
+                    ti = item.ti
+
+                    # If it's None, then the span for the current TaskInstanceKey hasn't been started.
+                    if self.active_spans is not None and self.active_spans.get(key) is None:
+                        from airflow.models.taskinstance import SimpleTaskInstance
+
+                        if isinstance(ti, SimpleTaskInstance):
+                            parent_context = Trace.extract(ti.parent_context_carrier)
+                        elif isinstance(ti, workloads.TaskInstance):
+                            parent_context = Trace.extract(ti.parent_context_carrier)
+                        else:
+                            parent_context = Trace.extract(ti.dag_run.context_carrier)
+                        # Start a new span using the context from the parent.
+                        # Attributes will be set once the task has finished so that all
+                        # values will be available (end_time, duration, etc.).
+
+                        span = Trace.start_child_span(
+                            span_name=f"{ti.task_id}",
+                            parent_context=parent_context,
+                            component="task",
+                            start_as_current=False,
+                        )
+                        self.active_spans.set(key, span)
+                        # Inject the current context into the carrier.
+                        carrier = Trace.inject()
+                        ti.context_carrier = carrier
+
                 if hasattr(self, "_process_workloads"):
-                    workloads.append(item)
+                    workload_list.append(item)
                 else:
                     (command, _, queue, ti) = item
                     task_tuples.append((key, command, queue, getattr(ti, "executor_config", None)))
 
         if task_tuples:
             self._process_tasks(task_tuples)
-        elif workloads:
-            self._process_workloads(workloads)  # type: ignore[attr-defined]
+        elif workload_list:
+            self._process_workloads(workload_list)  # type: ignore[attr-defined]
 
     @add_span
     def _process_tasks(self, task_tuples: list[TaskTuple]) -> None:

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

@@ -67,6 +67,10 @@ class TaskInstance(BaseModel):
     priority_weight: int
     executor_config: dict | None = Field(default=None, exclude=True)
 
+    parent_context_carrier: dict | None = None
+    context_carrier: dict | None = None
+    queued_dttm: datetime | None = None
+
     # TODO: Task-SDK: Can we replace TastInstanceKey with just the uuid across the codebase?
     @property
     def key(self) -> TaskInstanceKey:
@@ -105,6 +109,7 @@ def make(
         from airflow.utils.helpers import log_filename_template_renderer
 
         ser_ti = TaskInstance.model_validate(ti, from_attributes=True)
+        ser_ti.parent_context_carrier = ti.dag_run.context_carrier
         bundle_info = BundleInfo(
             name=ti.dag_model.bundle_name,
             version=ti.dag_run.bundle_version,