Context Propagation

View as Markdown

When a traced request crosses a process boundary (such as an HTTP call, gRPC message, message queue, or Ray remote call), its trace context (trace ID, parent span ID, and baggage) must travel with it. Otherwise, downstream spans end up in a different trace, and the cross-service waterfall is lost.

NeMo Lens exposes two primitives for this, matching the OTel W3C TraceContext and Baggage propagators.

inject_context for Outbound Calls

1from nemo.lens import inject_context
2
3headers = {}
4inject_context(headers)
5# headers == {'traceparent': '00-<trace_id>-<span_id>-01'} # plus 'tracestate'/'baggage' only if present
6
7await http_client.post(url, headers=headers, json=body)

inject_context(carrier) writes the current OTel context into the carrier dict. The carrier is whatever your transport uses, such as HTTP headers, gRPC metadata, or message attributes.

extract_context for Inbound Calls

1from nemo.lens import extract_context
2from opentelemetry import trace, context
3
4# On the receiving side
5ctx = extract_context(request.headers)
6token = context.attach(ctx)
7try:
8 with trace.get_tracer(__name__).start_as_current_span("handle.request"):
9 ...
10finally:
11 context.detach(token)

extract_context(carrier) parses W3C headers and returns an OTel Context. Attach it before starting child spans, detach when done.

If the carrier has no valid trace context, the returned Context is empty, so new spans start a fresh trace, which is the correct behavior.

Auto-Instrumentation for Common Transports

Writing inject_context or extract_context by hand on every call is error-prone. For common frameworks, NeMo Lens ships auto-instrumentation helpers:

  • FastAPI: nemo.lens.contrib.fastapi.instrument_fastapi(app) extracts context from every incoming request and makes it the parent of the current span.
  • aiohttp client: nemo.lens.contrib.aiohttp.instrument_aiohttp_client() injects context into every outgoing request.
  • Ray: nemo.lens.contrib.ray.inject_ray_context(), extract_ray_context(), and traced_remote_call(method) provide helpers for Ray’s kwargs-based propagation.
  • NCCL: nemo.lens.contrib.nccl.serialize_context() and extract_nccl_context(data) provide helpers for piggy-backing context on NCCL byte transfers.

See Contrib Helpers for details.

When auto-instrumentation is available, use it instead of manual injection. This approach covers every call site, handles error paths, and prevents code rot as the codebase grows.

Cross-Rank Propagation in Distributed Training

HTTP-style propagation does not apply to torch.distributed, as those are not carrier-based transports. For distributed training, NeMo Lens provides:

  • broadcast_trace_context(rank, src_rank=0), which uses torch.distributed.broadcast to share trace context across ranks.
  • create_linked_span(tracer, name, remote_carrier=carrier), which creates a span with an OTel Link (not a parent-child relationship) to a remote span, which is useful for pipeline-parallel correlation.

See Distributed Tracing for the full pattern.

Baggage

Baggage is a W3C standard for propagating small key/value context alongside trace context (such as “which customer is this request for?” or “is this a canary deployment?”).

NeMo Lens’s propagator is a CompositePropagator of TraceContextTextMapPropagator and W3CBaggagePropagator, so both flow through inject_context and extract_context automatically. Set baggage with:

1from opentelemetry import baggage, context
2
3ctx = baggage.set_baggage("customer.id", "42")
4token = context.attach(ctx)
5try:
6 # All downstream inject_context calls will include customer.id in the baggage header
7 ...
8finally:
9 context.detach(token)

Baggage values are propagated across service boundaries but are not automatically added to spans. Use them for filtering/routing decisions, not for span attributes.