diff --git a/airflow-core/docs/administration-and-deployment/logging-monitoring/traces.rst b/airflow-core/docs/administration-and-deployment/logging-monitoring/traces.rst index fae22f7c8873e..1c9258c6be24e 100644 --- a/airflow-core/docs/administration-and-deployment/logging-monitoring/traces.rst +++ b/airflow-core/docs/administration-and-deployment/logging-monitoring/traces.rst @@ -87,6 +87,36 @@ Custom spans created this way are automatically nested as children of the Airflo task span when tracing is enabled. When tracing is disabled, the no-op tracer provided by the OpenTelemetry API is used, so tasks run without any overhead. +Per-run trace controls (run conf) +--------------------------------- + +A few reserved keys in a Dag run's ``conf`` let you control tracing for that individual run. +They are read when the run is created (and when it is cleared), so set them when you trigger the +run — from the UI/API "Trigger with config" dialog, the ``airflow dags trigger --conf`` CLI, or +``TriggerDagRunOperator(conf=...)``. + +``airflow/trace_sampled`` + Force the head-sampling decision for this run. ``true`` always traces the run and ``false`` + never does, regardless of the configured sampler; omit the key to let the sampler decide. + Only an explicit boolean is honored. + +``airflow/task_span_detail_level`` + To enable detailed spans, set detail level greater than 1. This is intended as a debug-supporting feature as + such it is subject to change or removal at any time. + +``airflow/dagrun_parent_trace_context`` + Embed this run in an **external** trace instead of it being a root trace. Supply a W3C + ``traceparent`` string (optionally a mapping with ``traceparent`` and ``tracestate``) captured + from the system that triggered the run — an upstream orchestrator, event pipeline, CI job, or + another Airflow deployment. The whole run (the ``dag_run`` span and all task/worker spans) then + lives inside that trace, and parent-based samplers inherit the external sampling decision. When + the key is absent (default), the run is a root trace. A missing or malformed value is + ignored (the run stays a root) rather than failing run creation. + + .. code-block:: json + + {"airflow/dagrun_parent_trace_context": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"} + Enable Https ----------------- diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index edc4d903aa26c..e45a990a3fb3c 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/airflow-core/src/airflow/models/dagrun.py @@ -64,6 +64,7 @@ from airflow._shared.observability.metrics import stats from airflow._shared.observability.traces import ( + DAGRUN_PARENT_TRACE_CONTEXT_KEY, TASK_SPAN_DETAIL_LEVEL_KEY, TRACE_SAMPLED_KEY, new_dagrun_trace_carrier, @@ -194,6 +195,39 @@ def trace_sampled_override(conf) -> bool | None: return None +def parent_trace_context(conf) -> context.Context | None: + """ + External parent trace context from the ``airflow/dagrun_parent_trace_context`` run conf key. + + Lets a run be embedded in a trace owned by an external system (an upstream + orchestrator, event pipeline, CI job, another Airflow) instead of being a root + trace. The value is a W3C ``traceparent`` string, or a carrier dict carrying + ``traceparent`` (and optionally ``tracestate``); anything else -- including a + malformed traceparent -- is ignored so it can neither silently mis-parent the + run nor fail run creation. Returns an OpenTelemetry ``Context`` when a valid + parent is present, else None (root trace, the default). + """ + if not conf: + return None + match conf.get(DAGRUN_PARENT_TRACE_CONTEXT_KEY): + case str() as traceparent: + carrier = {"traceparent": traceparent} + case {"traceparent": str()} as raw: + # Keep only str members: a non-str tracestate reaches TraceState.from_header + # unvalidated and raises TypeError, which would drop the otherwise-valid parent. + carrier = {k: raw[k] for k in ("traceparent", "tracestate") if isinstance(raw.get(k), str)} + case _: + return None + try: + ctx = TraceContextTextMapPropagator().extract(carrier) + except Exception: + # Never let a malformed conf value fail run creation; fall back to a root trace. + return None + if not trace.get_current_span(ctx).get_span_context().is_valid: + return None + return ctx + + class DagRun(Base, LoggingMixin): """ Invocation instance of a DAG. @@ -425,6 +459,7 @@ def __init__( task_span_detail_level=self.conf.get(TASK_SPAN_DETAIL_LEVEL_KEY, None), attributes=dagrun_trace_attributes(self), # these are for potential use by head sampler force_sampled=trace_sampled_override(self.conf), + parent_context=parent_trace_context(self.conf), ) if not isinstance(partition_key, str | None): @@ -1116,11 +1151,11 @@ def _emit_dagrun_span(self, state: DagRunState): span_context = span.get_span_context() # Skip if the run was head-sampled out. Unlike the task spans, this guard is - # required (not just an optimization): the span below is forced to be a root span - # (context=context.Context()), so the configured sampler never sees the carrier's - # flag. A valid-but-unsampled carrier means head-sampled out; an invalid/empty - # carrier (legacy DagRun) recorded no decision, so it falls through and still - # emits — prior behavior we may want to reconsider. + # required (not just an optimization): the span below is started as a root span + # (or under an external parent) rather than under the carrier, so the configured + # sampler is not driven by the carrier's flag. A valid-but-unsampled carrier means + # head-sampled out; an invalid/empty carrier (legacy DagRun) recorded no decision, + # so it falls through and still emits — prior behavior we may want to reconsider. if span_context.is_valid and not span_context.trace_flags.sampled: return @@ -1143,18 +1178,21 @@ def _emit_dagrun_span(self, state: DagRunState): if self.partition_date: attributes["airflow.dag_run.partition_date"] = self.partition_date.isoformat() - # TODO: make the empty parent context optional. Default should be to - # nest the dag run span under the currently active parent span (by - # omitting `context` here); only use the empty `context.Context()` to - # force a root span when Airflow itself initiates the run (e.g. dag - # triggered via API, scheduler, or backfill). Today this forces a - # root span unconditionally. - # Tracked at https://github.com/apache/airflow/issues/67210 + # Root by default: an Airflow-initiated run owns its own trace, so we do not + # attach to whatever span is incidentally active in the scheduler. When the run + # opted into an external trace via airflow/dagrun_parent_trace_context, start under that + # parent so the run nests inside it. The carrier was already built riding the same + # external trace_id (both read the same conf), so the override_ids trace_id is a + # no-op here — only the span_id override takes effect and the parent link comes + # from parent_ctx. The emitted span's sampling then follows the external parent, so + # an unsampled external trace drops this dag_run span even when airflow/trace_sampled + # forced the head decision that the task/worker spans ride. + parent_ctx = parent_trace_context(self.conf) span = tracer.start_span( name=f"dag_run.{self.dag_id}", start_time=int((self.queued_at or self.start_date or timezone.utcnow()).timestamp() * 1e9), attributes=attributes, - context=context.Context(), + context=parent_ctx or context.Context(), ) status_code = StatusCode.OK if state == DagRunState.SUCCESS else StatusCode.ERROR span.set_status(status_code) diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index 8d1276999a295..9f8eb0805c4f4 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/airflow-core/src/airflow/models/taskinstance.py @@ -409,6 +409,7 @@ def clear_task_instances( from airflow.models.dagrun import ( # Avoid circular import DagRun, dagrun_trace_attributes, + parent_trace_context, trace_sampled_override, ) @@ -435,6 +436,7 @@ def clear_task_instances( task_span_detail_level=dr.conf.get(TASK_SPAN_DETAIL_LEVEL_KEY) if dr.conf else None, attributes=dagrun_trace_attributes(dr), force_sampled=trace_sampled_override(dr.conf), + parent_context=parent_trace_context(dr.conf), ) _recalculate_dagrun_queued_at_deadlines(dr, dr.queued_at, session) diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py index 3955ac572ef06..64b4947ec4c6f 100644 --- a/airflow-core/tests/unit/models/test_dagrun.py +++ b/airflow-core/tests/unit/models/test_dagrun.py @@ -45,7 +45,10 @@ from airflow import settings from airflow._shared.observability.metrics.base_stats_logger import StatsLogger -from airflow._shared.observability.traces import OverrideableRandomIdGenerator +from airflow._shared.observability.traces import ( + DAGRUN_PARENT_TRACE_CONTEXT_KEY, + OverrideableRandomIdGenerator, +) from airflow._shared.timezones import timezone from airflow.callbacks.callback_requests import DagCallbackRequest, DagRunContext from airflow.models.dag import DagModel, infer_automated_data_interval @@ -4158,6 +4161,100 @@ def failing_callback(context): mock_incr.assert_any_call("dag.callback_exceptions", tags=expected_tags) +_EXTERNAL_TRACE_ID = "11111111111111111111111111111111" +_EXTERNAL_SPAN_ID = "2222222222222222" + + +@pytest.mark.parametrize( + ("conf", "expected_trace_id"), + [ + pytest.param( + {DAGRUN_PARENT_TRACE_CONTEXT_KEY: f"00-{_EXTERNAL_TRACE_ID}-{_EXTERNAL_SPAN_ID}-01"}, + _EXTERNAL_TRACE_ID, + id="traceparent-string", + ), + pytest.param( + { + DAGRUN_PARENT_TRACE_CONTEXT_KEY: { + "traceparent": f"00-{_EXTERNAL_TRACE_ID}-{_EXTERNAL_SPAN_ID}-01" + } + }, + _EXTERNAL_TRACE_ID, + id="carrier-dict", + ), + pytest.param({DAGRUN_PARENT_TRACE_CONTEXT_KEY: "not-a-traceparent"}, None, id="malformed"), + # An unknown version prefix is accepted by the propagator's forward-compat parsing, + # so a 99- traceparent still rides the external trace rather than being rejected. + pytest.param( + {DAGRUN_PARENT_TRACE_CONTEXT_KEY: f"99-{_EXTERNAL_TRACE_ID}-{_EXTERNAL_SPAN_ID}-01"}, + _EXTERNAL_TRACE_ID, + id="future-version-prefix", + ), + pytest.param( + {DAGRUN_PARENT_TRACE_CONTEXT_KEY: f"00-{_EXTERNAL_TRACE_ID}-{_EXTERNAL_SPAN_ID}-zz"}, + None, + id="almost-valid-bad-flag-string", + ), + pytest.param( + { + DAGRUN_PARENT_TRACE_CONTEXT_KEY: { + "traceparent": f"00-{_EXTERNAL_TRACE_ID}-{_EXTERNAL_SPAN_ID}-zz" + } + }, + None, + id="almost-valid-bad-flag-dict", + ), + pytest.param({DAGRUN_PARENT_TRACE_CONTEXT_KEY: 123}, None, id="non-str"), + pytest.param({DAGRUN_PARENT_TRACE_CONTEXT_KEY: {"nope": "x"}}, None, id="dict-without-traceparent"), + pytest.param({}, None, id="empty-conf"), + pytest.param(None, None, id="no-conf"), + pytest.param({"other": "x"}, None, id="unrelated-key"), + ], +) +def test_parent_trace_context(conf, expected_trace_id): + """Only a valid W3C traceparent (string or carrier dict) yields a parent context; else None.""" + from airflow.models.dagrun import parent_trace_context + + ctx = parent_trace_context(conf) + if expected_trace_id is None: + assert ctx is None + else: + span_ctx = otel_trace.get_current_span(ctx).get_span_context() + assert span_ctx.is_valid + assert format(span_ctx.trace_id, "032x") == expected_trace_id + + +@pytest.mark.parametrize( + ("tracestate", "expected"), + [ + pytest.param("foo=bar", "bar", id="str-preserved"), + pytest.param(123, None, id="non-str-dropped"), + ], +) +def test_parent_trace_context_tracestate(tracestate, expected): + """A str tracestate rides alongside the traceparent; a non-str one is dropped, not raised.""" + from airflow.models.dagrun import parent_trace_context + + conf = { + DAGRUN_PARENT_TRACE_CONTEXT_KEY: { + "traceparent": f"00-{_EXTERNAL_TRACE_ID}-{_EXTERNAL_SPAN_ID}-01", + "tracestate": tracestate, + } + } + span_ctx = otel_trace.get_current_span(parent_trace_context(conf)).get_span_context() + assert format(span_ctx.trace_id, "032x") == _EXTERNAL_TRACE_ID + assert span_ctx.trace_state.get("foo") == expected + + +@mock.patch("airflow.models.dagrun.TraceContextTextMapPropagator.extract", side_effect=ValueError("boom")) +def test_parent_trace_context_swallows_propagator_error(mock_extract): + """A propagator failure degrades to a root trace instead of failing run creation.""" + from airflow.models.dagrun import parent_trace_context + + conf = {DAGRUN_PARENT_TRACE_CONTEXT_KEY: f"00-{_EXTERNAL_TRACE_ID}-{_EXTERNAL_SPAN_ID}-01"} + assert parent_trace_context(conf) is None + + class TestDagRunTracing: """Tests for DagRun OpenTelemetry span behavior.""" @@ -4441,6 +4538,30 @@ def test_context_carrier_honors_trace_sampled_conf(self, dag_maker, flag): span_ctx = trace.get_current_span(ctx).get_span_context() assert span_ctx.trace_flags.sampled is flag + @pytest.mark.parametrize( + ("conf", "embeds"), + [ + pytest.param( + {DAGRUN_PARENT_TRACE_CONTEXT_KEY: f"00-{_EXTERNAL_TRACE_ID}-{_EXTERNAL_SPAN_ID}-01"}, + True, + id="with-parent-conf-embeds", + ), + pytest.param(None, False, id="without-parent-conf-is-root"), + ], + ) + def test_context_carrier_parent_conf(self, dag_maker, conf, embeds): + """The carrier rides the external trace only when the parent conf key is set; else a root trace.""" + with dag_maker("test_tracing_parent_conf"): + EmptyOperator(task_id="t1") + dr = dag_maker.create_dagrun(conf=conf) + + ctx = TraceContextTextMapPropagator().extract(dr.context_carrier) + trace_id = format(otel_trace.get_current_span(ctx).get_span_context().trace_id, "032x") + if embeds: + assert trace_id == _EXTERNAL_TRACE_ID + else: + assert trace_id != _EXTERNAL_TRACE_ID + class TestDagRunStatsTagsTeamName: def test_stats_tags_without_team_name(self, dag_maker): diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py index 5f8e51c5b47c0..f2f5675ae566c 100644 --- a/airflow-core/tests/unit/models/test_taskinstance.py +++ b/airflow-core/tests/unit/models/test_taskinstance.py @@ -4151,6 +4151,26 @@ def test_clear_task_instances_honors_trace_sampled_conf(dag_maker, session, flag assert span_ctx.trace_flags.sampled is flag +@pytest.mark.db_test +def test_clear_task_instances_keeps_external_parent_trace(dag_maker, session): + """The regenerated carrier keeps riding the external trace from airflow/dagrun_parent_trace_context.""" + external_trace_id = "11111111111111111111111111111111" + with dag_maker("test_clear_parent_trace"): + EmptyOperator(task_id="t1") + dag_run = dag_maker.create_dagrun( + conf={"airflow/dagrun_parent_trace_context": f"00-{external_trace_id}-2222222222222222-01"} + ) + ti = dag_run.get_task_instance("t1", session=session) + ti.state = TaskInstanceState.SUCCESS + session.flush() + + clear_task_instances([ti], session) + + new_ctx = TraceContextTextMapPropagator().extract(dag_run.context_carrier) + span_ctx = trace.get_current_span(new_ctx).get_span_context() + assert format(span_ctx.trace_id, "032x") == external_trace_id + + @pytest.mark.db_test def test_task_instance_repr_does_not_raise_for_deferred_columns(dag_maker, session): """``TaskInstance.__repr__`` must survive *any* deferred column it reads. diff --git a/shared/observability/src/airflow_shared/observability/traces/__init__.py b/shared/observability/src/airflow_shared/observability/traces/__init__.py index f0ef435252786..ea866997db902 100644 --- a/shared/observability/src/airflow_shared/observability/traces/__init__.py +++ b/shared/observability/src/airflow_shared/observability/traces/__init__.py @@ -60,10 +60,11 @@ def generate_trace_id(self): TASK_SPAN_DETAIL_LEVEL_KEY = "airflow/task_span_detail_level" DEFAULT_TASK_SPAN_DETAIL_LEVEL = 1 TRACE_SAMPLED_KEY = "airflow/trace_sampled" +DAGRUN_PARENT_TRACE_CONTEXT_KEY = "airflow/dagrun_parent_trace_context" def new_dagrun_trace_carrier( - task_span_detail_level=None, attributes=None, force_sampled=None + task_span_detail_level=None, attributes=None, force_sampled=None, parent_context=None ) -> dict[str, str]: """ Generate a fresh W3C traceparent carrier without creating a recordable span. @@ -83,9 +84,29 @@ def new_dagrun_trace_carrier( SAMPLED flag directly (True = always trace this run, False = never) and the sampler is not consulted. Airflow wires this from the ``airflow/trace_sampled`` run conf key; when None the configured sampler makes the decision. + + ``parent_context`` optionally embeds the run in an *external* trace: when it + carries a valid span the carrier reuses that trace_id (so the whole run -- + dag_run, task_run, worker spans -- lives inside the external trace) and the + sampling decision is made with that parent, so parent-based samplers inherit + the external SAMPLED flag. Airflow wires this from the + ``airflow/dagrun_parent_trace_context`` run conf key; when None the run is a root + trace as before. """ + parent_span_context = ( + trace.get_current_span(context=parent_context).get_span_context() if parent_context else None + ) gen = RandomIdGenerator() - trace_id = gen.generate_trace_id() + if parent_span_context is not None and parent_span_context.is_valid: + # Embed in the external trace: ride its trace_id, inherit its tracestate, and + # let the sampler decide with this parent (so parent-based samplers follow it). + trace_id = parent_span_context.trace_id + parent_trace_state: TraceState | None = parent_span_context.trace_state + sampler_parent_context = parent_context + else: + trace_id = gen.generate_trace_id() + parent_trace_state = None + sampler_parent_context = None if force_sampled is not None: sampled = force_sampled @@ -95,7 +116,7 @@ def new_dagrun_trace_carrier( sampler = getattr(provider, "sampler", None) if sampler is not None: result = sampler.should_sample( - parent_context=None, # root decision + parent_context=sampler_parent_context, # None => root decision; else inherit parent trace_id=trace_id, name="dag_run", attributes=attributes or {}, @@ -109,9 +130,13 @@ def new_dagrun_trace_carrier( sampled = False sampler_trace_state = None - # Preserve the detail-level tracestate by merging it onto whatever the - # sampler returned. TraceState is immutable, so update() returns a new one. - trace_state = sampler_trace_state or TraceState() + # Preserve the detail-level tracestate by merging it onto whatever the sampler + # returned, falling back to the external parent's tracestate when embedding + # without a sampler decision. TraceState is immutable, so update() returns a new one. + trace_state = sampler_trace_state + if trace_state is None and parent_trace_state is not None: + trace_state = parent_trace_state + trace_state = trace_state or TraceState() for key, value in build_trace_state_entries(task_span_detail_level): trace_state = trace_state.update(key, value) diff --git a/shared/observability/tests/observability/test_traces.py b/shared/observability/tests/observability/test_traces.py index e2b43da92372f..eb0f499a250e1 100644 --- a/shared/observability/tests/observability/test_traces.py +++ b/shared/observability/tests/observability/test_traces.py @@ -18,7 +18,7 @@ from __future__ import annotations import pytest -from opentelemetry import trace +from opentelemetry import context, trace from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.sampling import ( ALWAYS_OFF, @@ -43,6 +43,23 @@ def _carrier_is_sampled(carrier: dict[str, str]) -> bool: return trace.get_current_span(ctx).get_span_context().trace_flags.sampled +def _carrier_span_context(carrier: dict[str, str]): + ctx = TraceContextTextMapPropagator().extract(carrier) + return trace.get_current_span(ctx).get_span_context() + + +def _parent_context(trace_id, span_id=0x1122334455667788, sampled=True, trace_state=None): + """A remote parent context standing in for an external trace.""" + span_ctx = SpanContext( + trace_id=trace_id, + span_id=span_id, + is_remote=True, + trace_flags=TraceFlags(TraceFlags.SAMPLED if sampled else 0), + trace_state=trace_state or TraceState(), + ) + return trace.set_span_in_context(NonRecordingSpan(span_ctx)) + + class TestBuildTraceStateEntries: def test_with_integer_level(self): entries = build_trace_state_entries(2) @@ -250,6 +267,65 @@ def test_force_sampled_preserves_detail_level(self, with_sampler): assert _carrier_is_sampled(carrier) is True +class TestNewDagrunTraceCarrierParentContext: + """parent_context embeds the run in an external trace instead of a fresh root.""" + + @pytest.fixture + def with_sampler(self, monkeypatch): + def _install(sampler): + provider = TracerProvider(sampler=sampler) + monkeypatch.setattr( + "airflow_shared.observability.traces.trace.get_tracer_provider", + lambda: provider, + ) + + return _install + + def test_embeds_in_parent_trace_with_own_span(self): + parent = _parent_context(trace_id=0xABC123, span_id=0xDEF456) + span_ctx = _carrier_span_context(new_dagrun_trace_carrier(parent_context=parent)) + assert span_ctx.trace_id == 0xABC123 # rides the external trace + assert span_ctx.span_id != 0xDEF456 # but is its own child span + + def test_without_parent_context_mints_fresh_root(self): + parent = _parent_context(trace_id=0xABC123) + embedded = _carrier_span_context(new_dagrun_trace_carrier(parent_context=parent)) + root = _carrier_span_context(new_dagrun_trace_carrier()) + assert root.trace_id != embedded.trace_id + + def test_empty_context_treated_as_no_parent(self): + # An empty Context carries an invalid span -> root trace, not an embed. + span_ctx = _carrier_span_context(new_dagrun_trace_carrier(parent_context=context.Context())) + assert span_ctx.is_valid + + @pytest.mark.parametrize( + ("sampler", "parent_sampled", "force_sampled", "expected_sampled"), + [ + # Root decision is OFF, but a sampled parent flips it to sampled. + pytest.param(ParentBased(root=ALWAYS_OFF), True, None, True, id="inherits-sampled-parent"), + # Root decision is ON, but an unsampled parent flips it to not-sampled. + pytest.param(ParentBased(root=ALWAYS_ON), False, None, False, id="inherits-unsampled-parent"), + # force_sampled bypasses the parent-based decision entirely. + pytest.param( + ParentBased(root=ALWAYS_ON), True, False, False, id="force-sampled-overrides-parent" + ), + ], + ) + def test_parent_based_sampler_decision( + self, with_sampler, sampler, parent_sampled, force_sampled, expected_sampled + ): + with_sampler(sampler) + parent = _parent_context(trace_id=0xAAA, sampled=parent_sampled) + carrier = new_dagrun_trace_carrier(parent_context=parent, force_sampled=force_sampled) + assert _carrier_is_sampled(carrier) is expected_sampled + + def test_embedding_preserves_parent_trace_state_when_forced(self): + # With force_sampled the sampler is skipped, so the external tracestate is the source. + parent = _parent_context(trace_id=0xAAA, trace_state=TraceState([("foo", "bar")])) + carrier = new_dagrun_trace_carrier(parent_context=parent, force_sampled=True) + assert _carrier_span_context(carrier).trace_state.get("foo") == "bar" + + class TestGetTaskSpanDetailLevel: def _make_span_with_trace_state(self, entries: list[tuple[str, str]]) -> NonRecordingSpan: from opentelemetry.sdk.trace.id_generator import RandomIdGenerator