Contrib Helpers

nemo.lens.contrib contains framework-specific integration helpers. Each is optional and isolated, so installing NeMo Lens without the corresponding extra raises a clear ImportError (with an install hint) when you call the helper, not when you import it.

FastAPI with contrib.fastapi

from nemo.lens.contrib.fastapi import instrument_fastapi
app = FastAPI()
instrument_fastapi(app)

Wraps opentelemetry-instrumentation-fastapi. After this call, every incoming HTTP request gets a span covering its lifetime, with W3C trace context automatically extracted from request headers (so upstream traces flow through).

Install: pip install 'nemo-lens[fastapi]'

Typical Integration

Gate on a span group so FastAPI spans respect your telemetry toggles:

from nemo.lens.state import is_span_group_enabled
if is_span_group_enabled('server'):
    instrument_fastapi(app)

aiohttp Client with contrib.aiohttp

from nemo.lens.contrib.aiohttp import instrument_aiohttp_client
instrument_aiohttp_client()
# After this, every aiohttp ClientSession request has W3C context injected automatically.

Wraps opentelemetry-instrumentation-aiohttp-client. Eliminates the need to manually call inject_context(kwargs['headers']) on every outbound HTTP call.

Install: pip install 'nemo-lens[aiohttp]'

Choose When to Call the Helper

Call this helper once at startup, after setup_telemetry returns and confirms you are actively exporting:

handle = setup_telemetry(config)
if handle.is_exporting:
    instrument_aiohttp_client()

Avoid running this helper on non-exporting ranks to prevent unnecessary overhead.

Ray with contrib.ray

Ray remote calls do not carry HTTP headers, so trace context must be passed explicitly. NeMo Lens exposes helpers that add a conventional _otel_carrier kwarg to remote calls.

Instrument Driver Side

from nemo.lens.contrib.ray import inject_ray_context, ray_dispatch_with_context
# Method 1: manual carrier
carrier = inject_ray_context()
future = my_actor.method.remote(arg1, arg2, _otel_carrier=carrier)
# Method 2: dispatch helper
future = ray_dispatch_with_context(my_actor.method, arg1, arg2)

Instrument Worker Side

Wrap remote methods with traced_remote_call to auto-extract context:

from nemo.lens.contrib.ray import traced_remote_call
@ray.remote
class MyActor:
    @traced_remote_call
    def method(self, arg1, arg2):
        # _otel_carrier kwarg is consumed by the decorator and used to set
        # the current span context before method body runs
        ...

Spans created inside method now appear as children of the driver’s span.

No extra install is needed, as this uses opentelemetry-api only.

NCCL with contrib.nccl

NCCL transfers use raw bytes and have no native header concept. For pipeline-parallel correlation, use the following pattern to piggyback trace context on a tensor transfer:

from nemo.lens.contrib.nccl import serialize_context, extract_nccl_context
# Sender
data = serialize_context()    # JSON-encoded W3C carrier as bytes
# ... send `data` alongside your tensor via NCCL ...
# Receiver
ctx = extract_nccl_context(data)
# attach ctx as parent context for new spans

If you need the intermediate carrier dict rather than a ready-to-use OTel Context, call deserialize_context(data: bytes) -> dict | None, which is the mid-layer that extract_nccl_context wraps. It returns the decoded carrier dict, or None if the bytes are malformed (it swallows JSONDecodeError and UnicodeDecodeError):

from nemo.lens.contrib.nccl import deserialize_context
carrier = deserialize_context(data)    # dict, or None on bad input

In practice, most pipeline-parallel users do not need this; broadcast_trace_context is simpler and more idiomatic (see Distributed Tracing). NCCL helpers exist for advanced cases where you are already passing metadata alongside tensors, and trace context can piggyback for free.

No extra install is needed, as this uses opentelemetry-api only.

Design Notes

Contrib modules are thin; each wraps an existing OTel instrumentation package or provides a couple of helper functions. They do not implement tracing logic themselves.

If you need to add a contrib module:

1. Check if an opentelemetry-instrumentation-<framework> package exists upstream. If yes, your module should be a single function that imports and calls it.
2. Add the package as an optional extra in pyproject.toml (nemo-lens[<framework>]).
3. Raise an ImportError with an actionable install hint if the instrumentation package is not present.

This keeps the contrib surface small and maintenance burden low.