API ReferencePython APINemo LensLens

lens.distributed

View as Markdown

Cross-rank trace context broadcast and linked spans.

Provides helpers for sharing a trace_id across all ranks (so that distributed training appears as a single trace) and for creating linked spans that reference remote contexts.

Module Contents

Functions

NameDescription
broadcast_trace_contextBroadcast W3C trace context from src_rank to all other ranks.
create_linked_spanCreate a span with a link to a remote span context.

API

lens.distributed.broadcast_trace_context(
    rank: int,
    src_rank: int = 0
) -> dict | None

Broadcast W3C trace context from src_rank to all other ranks.

Must be called after torch.distributed.init_process_group(). Returns a carrier dict with the broadcasted trace context.

Parameters:

rank
int

Current process rank.

src_rank
int" default="0

Rank that holds the source context.

Returns: dict | None

A carrier dict with the broadcasted trace context, or None if

lens.distributed.create_linked_span(
    tracer: opentelemetry.trace.Tracer,
    name: str,
    remote_context: opentelemetry.context.Context | None = None,
    remote_carrier: dict | None = None,
    attributes = {}
) -> opentelemetry.trace.Span

Create a span with a link to a remote span context.

Use this for cross-rank or cross-service span correlation where parent-child relationships don’t make sense (e.g., pipeline parallel stages).

Parameters:

tracer
trace.Tracer

OTel tracer.

name
str

Span name.

remote_context
context.Context | None" default="None

An OTel Context from extract_context().

remote_carrier
dict | None" default="None

A W3C carrier dict (alternative to remote_context).

**attributes
Defaults to {}

Span attributes.

Returns: trace.Span

The created Span (already started, caller must end it).