potiuk commented on code in PR #66151:
URL: https://github.com/apache/airflow/pull/66151#discussion_r3173543594
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -738,6 +743,13 @@ def handle_requests(self, log: FilteringBoundLogger) ->
Generator[None, _Request
log.exception("Unable to decode message", body=request.body)
continue
+ # Restore the task runner's trace context so that any outbound
HTTP calls made while
+ # handling this request are linked to the correct task span, not
the supervisor's own span.
+ token = None
+ if request.context_carrier:
+ ctx = _trace_propagator.extract(request.context_carrier)
+ token = otel_context.attach(ctx)
Review Comment:
Minor correctness: `attach()` lives outside the `try`. In the unlikely event
something fires between line 751 and line 753 (e.g. extract returns a context
that triggers some custom processor on attach, or an interpreter-level
interrupt during dispatch), the `token` would leak un-detached. Cheap to harden
by moving the attach inside the try:
```python
token = None
try:
if request.context_carrier:
ctx = _trace_propagator.extract(request.context_carrier)
token = otel_context.attach(ctx)
self._handle_request(msg, log, request.id)
except ServerResponseError as e:
...
finally:
if token is not None:
otel_context.detach(token)
```
In practice `extract()` is robust against malformed carriers (returns an
invalid-context rather than raising), so this is mostly defense-in-depth. Worth
tightening anyway since the `finally` is already there.
---
Drafted-by: Claude Opus 4.7; reviewed by @potiuk before posting
##########
task-sdk/src/airflow/sdk/api/client.py:
##########
@@ -206,6 +210,24 @@ def add_correlation_id(request: httpx.Request):
request.headers["correlation-id"] = str(uuid7())
+def inject_trace_context(request: httpx.Request) -> None:
+ _trace_propagator.inject(request.headers)
+
+
+def _log_and_trace_retry(retry_state) -> None:
+ before_log(log, logging.WARNING)(retry_state)
Review Comment:
nit: `before_log(log, logging.WARNING)` constructs a fresh callable on every
retry. Hoisting it once at module scope is cleaner and avoids the (negligible)
per-retry allocation:
```python
_before_log_warning = before_log(log, logging.WARNING)
def _log_and_trace_retry(retry_state) -> None:
_before_log_warning(retry_state)
...
```
Non-blocking.
---
Drafted-by: Claude Opus 4.7; reviewed by @potiuk before posting
##########
task-sdk/src/airflow/sdk/execution_time/comms.py:
##########
@@ -193,10 +206,14 @@ class CommsDecoder(Generic[ReceiveMsgType, SendMsgType]):
# Async lock for async operations
_async_lock: asyncio.Lock = attrs.field(factory=asyncio.Lock, repr=False)
+ def _make_frame(self, msg: SendMsgType) -> _RequestFrame:
Review Comment:
Worth adding direct unit tests for the new propagation chain — manual
end-to-end (xBis7's screenshot) is great validation but doesn't catch silent
regressions if someone refactors this later. Concretely:
- A `test_comms.py` test that asserts `_make_frame()` populates
`context_carrier` when a span is active, and leaves it `None` (not an empty
dict) otherwise — the `carrier or None` here is the contract.
- A `test_supervisor.py` test that feeds a `_RequestFrame` with a known
carrier into `handle_requests()` and asserts the supervisor's downstream
`_handle_request` sees the extracted context (e.g. via a stub that records
`trace.get_current_span()`).
- A test for `inject_trace_context` confirming the httpx request gets a
`traceparent` header when a span is recording.
Grepping `task-sdk/tests/` for `traceparent` / `context_carrier` only turns
up the pre-existing TI/DagRun-level tests in `test_task_runner.py:507` —
nothing on the new IPC/HTTP propagation. Happy to leave for a follow-up if
you'd rather, but flagging since it's a protocol-shape change.
---
Drafted-by: Claude Opus 4.7; reviewed by @potiuk before posting
##########
task-sdk/src/airflow/sdk/execution_time/comms.py:
##########
@@ -159,7 +154,25 @@ def as_bytes(self) -> bytearray:
return buffer
-class _ResponseFrame(_RequestFrame, frozen=True): # type: ignore[call-arg]
+class _RequestFrame(_FrameMixin, msgspec.Struct, array_like=True, frozen=True,
omit_defaults=True): # type: ignore[call-arg]
+ id: int
+ """
+ The request id, set by the sender.
+
+ This is used to allow "pipeling" of requests and to be able to tie
response to requests, which is
+ particularly useful in the Triggerer where multiple async tasks can send a
requests concurrently.
+ """
+ body: dict[str, Any] | None
+ context_carrier: dict[str, str] | None = None
+ """W3C trace context carrier (traceparent + tracestate) of the task
runner's active span.
+
Review Comment:
nit: the field name `context_carrier` collides with the pre-existing
`context_carrier` on `TaskInstance` / `DagRun` (used for scheduler→task
propagation). Same W3C-carrier semantic but a different propagation hop. Future
readers landing on one might assume it's the other. A one-line cross-reference
in the docstring ("Distinct from `TaskInstance.context_carrier`, which carries
the scheduler's context to the task; this one carries the task runner's context
to the supervisor") would disambiguate.
Non-blocking.
---
Drafted-by: Claude Opus 4.7; reviewed by @potiuk before posting
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -2244,6 +2259,9 @@ def supervise_task(
finally:
if log_path and log_file_descriptor:
log_file_descriptor.close()
+ provider = trace.get_tracer_provider()
+ if hasattr(provider, "force_flush"):
+ provider.force_flush(timeout_millis=5000)
Review Comment:
nit: worth a one-line comment that the 5000ms is an *upper bound*, not a
fixed wait — `force_flush` returns as soon as the BSP has drained, so for the
common case this is sub-ms. Without that note, a future reader benchmarking
task teardown might see this and assume every supervisor exit costs 5s.
Non-blocking.
---
Drafted-by: Claude Opus 4.7; reviewed by @potiuk before posting
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
