lens.contrib.ray | NVIDIA NeMo Lens

Ray context propagation helpers.

Provides utilities for propagating OTel trace context across Ray remote calls, enabling distributed traces that span driver -> worker boundaries.

Module Contents

Functions

Name	Description
`extract_ray_context`	Extract trace context from a Ray carrier.
`inject_ray_context`	Create a carrier dict with the current trace context for Ray.
`ray_dispatch_with_context`	Dispatch a Ray remote call with trace context injection.
`traced_remote_call`	Wrapper for Ray remote methods that auto-injects/extracts trace context.

API

lens.contrib.ray.extract_ray_context(
    carrier: dict | None = None
)

Extract trace context from a Ray carrier.

Call this on the worker side to resume the trace.

Parameters:

carrier

dict | None" default="None

Dict with W3C trace context headers. If None, returns an empty context.

Returns:

An OTel Context.

lens.contrib.ray.inject_ray_context() -> dict

Create a carrier dict with the current trace context for Ray.

Call this on the driver side before dispatching a remote call.

Returns: dict

A dict containing W3C trace context headers.

lens.contrib.ray.ray_dispatch_with_context(
    remote_fn,
    args: typing.Any = (),
    _carrier: dict | None = None,
    kwargs: typing.Any = {}
)

Dispatch a Ray remote call with trace context injection.

Parameters:

remote_fn

The Ray remote function/method reference.

*args

Any" default="()

Positional arguments for the remote call.

_carrier

dict | None" default="None

Pre-built carrier (if None, injects current context).

**kwargs

Any" default="{}

Keyword arguments for the remote call.

Returns:

The Ray ObjectRef from the remote call.

lens.contrib.ray.traced_remote_call(
    method: collections.abc.Callable,
    tracer: opentelemetry.trace.Tracer | None = None
) -> collections.abc.Callable

Wrapper for Ray remote methods that auto-injects/extracts trace context.

Apply to worker methods to automatically propagate traces across Ray remote boundaries.

Parameters:

method

Callable

The function to wrap.

tracer

trace.Tracer | None" default="None

OTel tracer. Defaults to the global tracer.

Returns: Callable

Wrapped function that accepts _otel_carrier kwarg.