lens.contrib.nccl | NVIDIA NeMo Lens

NCCL context propagation helpers.

Provides utilities for serializing/deserializing OTel trace context alongside activation tensor transfers in pipeline parallel training.

Module Contents

Name	Description
`deserialize_context`	Deserialize trace context from bytes received via NCCL.
`extract_nccl_context`	Extract OTel context from NCCL-received bytes.
`serialize_context`	Serialize current trace context to bytes for piggy-backing on NCCL sends.

lens.contrib.nccl.deserialize_context(
    data: bytes
) -> dict | None

Deserialize trace context from bytes received via NCCL.

Parameters:

data

bytes

JSON-encoded W3C trace context bytes.

Returns: dict | None

A carrier dict, or None if deserialization fails.

lens.contrib.nccl.extract_nccl_context(
    data: bytes
)

Extract OTel context from NCCL-received bytes.

Convenience function combining deserialize + extract.

Parameters:

data

bytes

JSON-encoded W3C trace context bytes.

Returns:

An OTel Context (empty if deserialization fails).

lens.contrib.nccl.serialize_context() -> bytes

Serialize current trace context to bytes for piggy-backing on NCCL sends.

Returns: bytes

JSON-encoded W3C trace context as bytes.