nemo_rl.telemetry.instrumentation#

Instrumentation helpers that attach efficiency tags.

Algorithms should import managed_span / trace_fn from here (not raw nemo-lens) so every leaf span gets rl.bucket when applicable.

Shared bucket tokens are productive | overhead | idle | wasted. Umbrella groups (job, step, rollout, …) are timed but not tagged; open those through umbrella_span / umbrella_trace_fn with the group’s U_ alias, so the call site says which of the two it is.

Module Contents#

Classes#

Bucket

Shared goodput buckets.

Functions#

safe_set_span_attributes

Like the lens helper, but a no-op when span is None.

bucket_for_span_group

Return the goodput bucket for a span group, or None if umbrella / unknown.

bucket_for_efficiency_category

Return the bucket for an async efficiency category label, if known.

bucket_scope

Reclassify every leaf span opened inside this block as bucket.

iteration_scope

Stamp :data:RL_ITERATION_ATTR on every span opened inside this block.

per_prompt_scope

Mark this block as per-prompt work.

in_per_prompt_scope

Whether the caller is running inside :func:per_prompt_scope.

goodput_span_attributes

Attributes to merge into managed_span for group.

current_trace_carrier

W3C traceparent carrier for the active span, to hand to another process.

remote_trace_context

Parent every span opened in this block to the span in carrier.

remote_trace_parent

The OTel context carrier names, or None when there is nothing to attach.

attached_context

Attach context for the duration of the block; a no-op for None.

trace_context_kwargs

Carrier kwarg for a Ray call, or {} when there is nothing to send.

dispatch_with_trace_context

remote_method.remote(...), plus the caller’s trace context.

_dispatch_method_name

Name of the method behind a dispatch handle, for the warning key.

_advertise_carrier

Declare the reserved kwarg on the signature Ray actually reads.

accepts_trace_context

Let a Ray method be parented to its caller’s span.

managed_span

Like lens managed_span, but injects rl.bucket for leaf groups.

_warn_leaf_group_at_umbrella_call

Report umbrella_span being handed a leaf group, once per span name.

umbrella_span

Span for an umbrella group: timed and nested, never bucketed.

streaming_umbrella_span

An umbrella span over an async generator, made current one step at a time.

_no_span_activation

Activation for a disabled group: nothing to attach, nothing to detach.

_efficiency_span_name_and_attrs

(span name, attributes) for one efficiency category.

efficiency_span

Span for one efficiency category, tagged with that category’s bucket.

start_efficiency_span

Start an efficiency span the caller ends by hand, or None if disabled.

startup_span

Umbrella over everything between process start and the first step.

setup_span

One startup phase, named rl.setup.<phase>.

evaluate_span

One validation pass, named rl.<algorithm>.evaluate.

checkpoint_finalize_span

The asynchronous tail of a checkpoint, named rl.checkpoint.finalize.

trace_fn

Decorator that wraps a function in a bucket-tagged managed_span.

umbrella_trace_fn

Decorator form of :func:umbrella_span, for whole-function umbrellas.

Data#

API#

nemo_rl.telemetry.instrumentation.RL_ALGORITHM#

‘rl.algorithm’

nemo_rl.telemetry.instrumentation.logger#

‘getLogger(…)’

nemo_rl.telemetry.instrumentation.RL_BUCKET_ATTR#

‘rl.bucket’

nemo_rl.telemetry.instrumentation.RL_EFFICIENCY_CATEGORY_ATTR#

‘rl.efficiency.category’

nemo_rl.telemetry.instrumentation.RL_ITERATION_ATTR#

‘rl.iteration’

nemo_rl.telemetry.instrumentation.RL_IDLE_POLLS_ATTR#

‘rl.idle.polls’

nemo_rl.telemetry.instrumentation.__all__#

[‘NO_SPAN’, ‘managed_span’, ‘umbrella_span’, ‘streaming_umbrella_span’, ‘trace_fn’, ‘umbrella_trace_…

nemo_rl.telemetry.instrumentation.safe_set_span_attributes(
span: Optional[Any],
attributes: dict[str, Any],
) → None#

Like the lens helper, but a no-op when span is None.

Every span helper in this module yields None for a disabled group, and lens’s version dereferences the span, so the raw one turns “telemetry is off” into an AttributeError at the call site. Absorbing that here rather than asking each caller to guard: the guard is invisible when it is missing (the disabled path is the one nobody exercises locally) and the failure lands in production code that has nothing to do with telemetry.

class nemo_rl.telemetry.instrumentation.Bucket#

Bases: str, enum.Enum

Shared goodput buckets.

Initialization

Initialize self. See help(type(self)) for accurate signature.

PRODUCTIVE#

‘productive’

OVERHEAD#

‘overhead’

IDLE#

‘idle’

WASTED#

‘wasted’

nemo_rl.telemetry.instrumentation.UMBRELLA_GROUPS: frozenset[str]#

None

nemo_rl.telemetry.instrumentation._DEFAULT_GROUP_BUCKET: Mapping[str, nemo_rl.telemetry.instrumentation.Bucket]#

None

nemo_rl.telemetry.instrumentation.EFFICIENCY_CATEGORY_BUCKET: Mapping[str, nemo_rl.telemetry.instrumentation.Bucket]#

None

nemo_rl.telemetry.instrumentation.bucket_for_span_group(
group: str,
) → Optional[nemo_rl.telemetry.instrumentation.Bucket]#

Return the goodput bucket for a span group, or None if umbrella / unknown.

Unknown non-umbrella groups default to overhead so new leaves are not silently dropped from the denominator.

nemo_rl.telemetry.instrumentation.bucket_for_efficiency_category(
category: str,
) → Optional[nemo_rl.telemetry.instrumentation.Bucket]#

Return the bucket for an async efficiency category label, if known.

nemo_rl.telemetry.instrumentation.COLLECTOR_LOOP_CATEGORIES: frozenset[str]#

‘frozenset(…)’

nemo_rl.telemetry.instrumentation.UNBUCKETED_SPAN_CATEGORIES: frozenset[str]#

None

nemo_rl.telemetry.instrumentation._BUCKET_OVERRIDE: contextvars.ContextVar[Optional[nemo_rl.telemetry.instrumentation.Bucket]]#

‘ContextVar(…)’

nemo_rl.telemetry.instrumentation.bucket_scope(
bucket: nemo_rl.telemetry.instrumentation.Bucket,
) → Iterator[None]#

Reclassify every leaf span opened inside this block as bucket.

For phases whose goodput meaning is set by the caller, not by the callee’s span group. Validation is the motivating case: it generates through the same

Data:

RLSpanGroup.GENERATION path as training rollouts, but the tokens are scored and discarded, so counting them as productive overstates goodput.

Applies to the group-derived bucket only. Umbrella groups stay unbucketed, a span that passes rl.bucket explicitly keeps it, and an

Func:

efficiency_span keeps its category’s bucket — that one names the phase it measures, so a caller cannot make idle/refit_bubble productive. So wrapping a region cannot start double-counting an interval that its children already account for.

Propagates like any :class:~contextvars.ContextVar: to nested calls, and to coroutines started inside the block (asyncio.run copies the current context), but not to raw threads or other processes.

nemo_rl.telemetry.instrumentation.iteration_scope(iteration: int) → Iterator[None]#

Stamp :data:RL_ITERATION_ATTR on every span opened inside this block.

Wrap a training step once and its whole subtree – leaves, checkpointing, rollout and generation spans – becomes filterable by step. Threading the counter into each call site instead leaves the answer to “what else was happening in step 412” dependent on which sites remembered to pass it, and silently excludes every span added afterwards.

A span that passes the attribute itself keeps its own value: lens starts the span before applying the caller’s attributes, so those overwrite whatever the scope stamped on. evaluate_span relies on that, since validation also runs once before training and reports iteration 0 there.

Only spans from the provider lens builds are stamped, which is every span NeMo-RL opens. Propagates like any :class:~contextvars.ContextVar: to nested calls and to coroutines started inside the block, but not to raw threads or other processes – a worker’s rl.policy.load_model is a root span in another process and is outside any step besides.

nemo_rl.telemetry.instrumentation._PER_PROMPT_SCOPE: contextvars.ContextVar[bool]#

‘ContextVar(…)’

nemo_rl.telemetry.instrumentation.per_prompt_scope() → Iterator[None]#

Mark this block as per-prompt work.

Spans opened inside it that consult :func:in_per_prompt_scope move to

Data:

RLSpanGroup.PER_PROMPT, so one group switch turns off every per-prompt span at once rather than each site needing its own flag.

Propagates like any :class:~contextvars.ContextVar: to nested calls and to coroutines started inside the block, but not to raw threads or other processes. That suits the motivating case – the rollout’s put happens in the same asyncio task that entered the scope – but a data-plane call handed to a thread pool would read as per-batch.

nemo_rl.telemetry.instrumentation.in_per_prompt_scope() → bool#

Whether the caller is running inside :func:per_prompt_scope.

nemo_rl.telemetry.instrumentation.goodput_span_attributes(group: str) → dict[str, str]#

Attributes to merge into managed_span for group.

Empty when the group is an umbrella (no rl.bucket). An enclosing

Func:

bucket_scope replaces the group’s default bucket.

nemo_rl.telemetry.instrumentation.current_trace_carrier() → dict[str, str]#

W3C traceparent carrier for the active span, to hand to another process.

Ray does not propagate OTel context, so a worker’s spans start their own trace unless the parent is passed explicitly. Capture this on the driver inside the span that should be the root, hand it to the actor, and reopen it there with :func:remote_trace_context.

Returns an empty dict when there is no active recording span — which is the case whenever the enclosing span’s group is disabled — so the caller needs no telemetry-specific branch.

nemo_rl.telemetry.instrumentation.remote_trace_context(
carrier: Optional[Mapping[str, str]],
) → Iterator[None]#

Parent every span opened in this block to the span in carrier.

A no-op for an empty carrier, so an uninstrumented or job-span-disabled run keeps emitting root spans instead of failing.

Attach per thread, not once per process: OTel context is a

Class:

~contextvars.ContextVar, and threading.Thread does not inherit them — a fire-and-forget worker thread starts with an empty context.

nemo_rl.telemetry.instrumentation.remote_trace_parent(
carrier: Optional[Mapping[str, str]],
) → Optional[Any]#

The OTel context carrier names, or None when there is nothing to attach.

nemo_rl.telemetry.instrumentation.attached_context(
context: Optional[Any],
) → Iterator[None]#

Attach context for the duration of the block; a no-op for None.

Attach/detach come from opentelemetry because lens wraps the extraction but not the activation.

nemo_rl.telemetry.instrumentation.TRACE_CARRIER_KWARG#

‘_otel_carrier’

nemo_rl.telemetry.instrumentation._F#

‘TypeVar(…)’

nemo_rl.telemetry.instrumentation.trace_context_kwargs() → dict[str, Any]#

Carrier kwarg for a Ray call, or {} when there is nothing to send.

For dispatch that does not go through remote() directly – notably RayWorkerGroup.run_all_workers_*, which takes the method by name and forwards **kwargs – spread this at the call site::

self.worker_group.run_all_workers_single_data(
    "begin_train_step_presharded", loss_fn=loss_fn, **trace_context_kwargs()
)

Every argument goes by keyword: run_all_workers_single_data asserts it was given none positionally.

Empty rather than a None carrier so a run with no recording span calls the method with exactly its original signature, which keeps an undecorated method working until it is wired.

nemo_rl.telemetry.instrumentation._CARRIER_REFUSED: set[str]#

‘set(…)’

nemo_rl.telemetry.instrumentation.dispatch_with_trace_context(
remote_method: Any,
/,
*args: Any,
**kwargs: Any,
) → Any#

remote_method.remote(...), plus the caller’s trace context.

The receiving method should be wrapped in :func:accepts_trace_context. When it is not, Ray rejects the extra kwarg on the caller – it validates against the callee’s recorded signature before submitting the task – and this retries without it, warning once. Losing the parent link degrades a trace; raising here would abort a rollout over an observability kwarg, which this module does not do anywhere else. test_every_context_accepting_method_is_dispatched_with_a_carrier is what actually keeps the two halves wired, at build time rather than at runtime.

Equivalent to lens’s ray_dispatch_with_context and uses the same kwarg, but skips the kwarg entirely when there is no recording span, so a run with the job group disabled calls the method with its original signature.

Returns:

normally an ObjectRef, or an ObjectRefGenerator when the handle was built with .options(num_returns="streaming") – which is how both NeMo-Gym call sites use it.

Return type:

Whatever .remote() returned

nemo_rl.telemetry.instrumentation._dispatch_method_name(remote_method: Any) → str#

Name of the method behind a dispatch handle, for the warning key.

.options(...) hands back _ActorMethodOptionsWrapper (Ray 2.56.1), which keeps the real ActorMethod in _actor_method and carries no _method_name itself. A repr would do for a key, except it embeds a fresh address per call, so the warning would repeat on every dispatch – and both NeMo-Gym call sites go through .options.

nemo_rl.telemetry.instrumentation._advertise_carrier(method: Any) → None#

Declare the reserved kwarg on the signature Ray actually reads.

Ray records an actor method’s signature from inspect.unwrap(method) (_ActorClassMethodMetadata.create) and validates every .remote() call against it on the caller, before the task is submitted. A functools.wraps wrapper sets __wrapped__, so unwrapping walks straight past our **kwargs to the original signature and the carrier comes back as TypeError: got an unexpected keyword argument. Ray hits exactly this with its own _ray_trace_ctx and fixes it the same way, on the same object, for the same reason – see ray/util/tracing.

A no-op when the unwrapped target already accepts **kwargs (everything stacked under wrap_with_nvtx_name, whose wrapper does not use functools.wraps and so absorbs the carrier on its own) and when the parameter is already declared, so decorating twice cannot duplicate it.

nemo_rl.telemetry.instrumentation.accepts_trace_context(
method: nemo_rl.telemetry.instrumentation._F,
) → nemo_rl.telemetry.instrumentation._F#

Let a Ray method be parented to its caller’s span.

Strips the reserved carrier kwarg and runs the body with that context attached, so every span the method opens – and every span opened by anything it calls, including an aiohttp request auto-instrumented by lens – nests under the caller instead of starting a new trace.

Unlike lens’s traced_remote_call this dispatches on the wrapped callable’s kind. That helper wraps everything in a def, which returns a coroutine or async generator without awaiting it, so the context detaches before the body runs – the async paths are precisely the ones that need this (NemoGym.run_rollouts is a streaming generator). It also does not open a span of its own: the methods this decorates already open theirs, and a second automatic span per call would double the trace’s depth.

The return is cast back to the decorated method’s own type: the wrappers take **kwargs, so without it every method this decorates would widen to an untyped callable, and several of them live in files the pyrefly whitelist expects to be fully typed.

nemo_rl.telemetry.instrumentation.managed_span(
group: str,
name: str,
tracer: Optional[opentelemetry.trace.Tracer] = None,
**attributes: Any,
) → Iterator[Any]#

Like lens managed_span, but injects rl.bucket for leaf groups.

Callers may override by passing rl.bucket=... explicitly. Umbrella groups (job / step / rollout / …) receive no bucket attribute.

nemo_rl.telemetry.instrumentation._LEAF_GROUP_AT_UMBRELLA_CALL: set[str]#

‘set(…)’

nemo_rl.telemetry.instrumentation._warn_leaf_group_at_umbrella_call(group: str, name: str) → None#

Report umbrella_span being handed a leaf group, once per span name.

A warning rather than an exception. The mistake is static, and a drift test rejects it at review time, which is both earlier and safer than a raise: umbrella spans open at very different points in a run — rl.<algo>.job at startup, rl.<algo>.evaluate not until the first validation — so raising would turn a mistyped group into a job that dies hours in. Nothing else in this package fails a run over telemetry either, and this check sits ahead of the group gate, so a raise here could kill a run that has tracing switched off entirely.

stack_info because the useful thing is the offending call site, and there is no exception here whose traceback would point at it.

nemo_rl.telemetry.instrumentation.umbrella_span(
group: str,
name: str,
tracer: Optional[opentelemetry.trace.Tracer] = None,
**attributes: Any,
) → Iterator[Any]#

Span for an umbrella group: timed and nested, never bucketed.

Behaves like :func:managed_span on an umbrella group — the distinction is at the call site, not in the output. Spell the group with its U_ alias (RLSpanGroup.U_ROLLOUT) so a reader can tell without a lookup that this span is absent from a goodput rollup.

Reach for it whenever a span can overlap another instance of itself: concurrent spans sum past the wall clock they happened in, so a bucket on them multiplies rather than measures. The work underneath is still counted, by the leaf spans nested inside.

Handed a leaf group, this warns and emits the span as a leaf instead — see

Func:

_warn_leaf_group_at_umbrella_call for why that rather than raising.

nemo_rl.telemetry.instrumentation.streaming_umbrella_span(
group: str,
name: str,
tracer: Optional[opentelemetry.trace.Tracer] = None,
**attributes: Any,
) → Iterator[Callable[[], ContextManager[Any]]]#

An umbrella span over an async generator, made current one step at a time.

Func:

umbrella_span cannot wrap a generator’s yield: Ray abandons a cancelled streaming call without closing the generator, so a context still attached at the yield is detached later by the garbage collector, in a different context, and OTel logs Failed to detach context.

The span itself still spans the whole stream – it is started here and ended when the block leaves. Only the attachment is per step, which is all that has to be, since nothing between two items belongs to this span.

Yields:

A callable returning a context manager to wrap each __anext__. It is a no-op when the group is disabled, so the caller needs no branch.

nemo_rl.telemetry.instrumentation.NO_SPAN: Final[ContextManager[None]]#

‘nullcontext(…)’

nemo_rl.telemetry.instrumentation._no_span_activation() → ContextManager[Any]#

Activation for a disabled group: nothing to attach, nothing to detach.

nemo_rl.telemetry.instrumentation._efficiency_span_name_and_attrs(
category: str,
attributes: Mapping[str, Any],
) → tuple[str, dict[str, Any]]#

(span name, attributes) for one efficiency category.

Shared by the two ways to open one of these spans so the bucket rules — including which categories are deliberately unbucketed — cannot drift between them.

nemo_rl.telemetry.instrumentation.efficiency_span(
category: str,
tracer: Optional[opentelemetry.trace.Tracer] = None,
**attributes: Any,
) → Iterator[Any]#

Span for one efficiency category, tagged with that category’s bucket.

category is the same label the Timer uses ("idle/refit_bubble", …), which keeps the span and the efficiency/* metric describing the identical phase. The bucket comes from

Data:

EFFICIENCY_CATEGORY_BUCKET, so idle/* lands in idle rather than defaulting to overhead the way an unknown leaf group would.

Categories in :data:UNBUCKETED_SPAN_CATEGORIES are emitted without a bucket — visible in a trace, invisible to a rollup. For the rest, two conditions have to hold at the call site. The phase must be measured on a single thread against wall time, since categories summed across concurrent threads are thread-seconds and would overcount (see

Data:

EFFICIENCY_CATEGORY_BUCKET). And the wrapped block must emit no bucketed child spans, because this span carries rl.bucket and a rollup that sums durations by bucket has no notion of nesting: a bucketed parent covering the same interval as its children is counted twice. Wrap a wait, not a phase that does instrumented work.

nemo_rl.telemetry.instrumentation.start_efficiency_span(
category: str,
tracer: Optional[opentelemetry.trace.Tracer] = None,
**attributes: Any,
) → Optional[Any]#

Start an efficiency span the caller ends by hand, or None if disabled.

For a wait that is implemented as a poll loop, where the span has to outlive a single iteration. :func:efficiency_span cannot express that twice over: it is a context manager, so it cannot span iterations at all, and it attaches the span to the context, so anything started while it was open would nest underneath it. In an asyncio pump that is not hypothetical — a task created during the wait copies the context it was created in, so unrelated rollout work would reparent under an idle span. This starts the span without attaching it: it adopts no children, and it does not have to, since a wait has nothing to nest.

The caller owns the whole lifecycle, which is the cost of the coalescing:

  • end it on the path that ends the wait, or the span runs to whenever the loop happens to stop;

  • end it on every other way out of the loop that lets the run continue – break, or the loop condition going false mid-wait – since an unended span stays open over whatever follows and is never exported. Exits that end the run can be left alone: dropping the span costs nothing there;

  • pass :data:RL_IDLE_POLLS_ATTR so the duration can be read against the number of retries it covers.

Returns None when the efficiency group is off, which doubles as the telemetry-off no-op: the is not None guards the caller already needs for a hand-managed span are the same ones that skip the disabled path.

Prefer :func:efficiency_span for anything that fits in a with.

nemo_rl.telemetry.instrumentation.startup_span(
tracer: Optional[opentelemetry.trace.Tracer] = None,
**attributes: Any,
) → Iterator[Any]#

Umbrella over everything between process start and the first step.

Open this in the entrypoint, around both init_ray() and the algorithm’s setup(). Those are separate top-level calls, so without something spanning them the startup phases arrive as unrelated root traces rather than one waterfall.

nemo_rl.telemetry.instrumentation.setup_span(
phase: str,
tracer: Optional[opentelemetry.trace.Tracer] = None,
**attributes: Any,
) → Iterator[Any]#

One startup phase, named rl.setup.<phase>.

Unbucketed and freely nestable — see :data:UMBRELLA_GROUPS for why the startup group carries no bucket at all.

Use for a phase that does work (building workers, opening collectives). For a phase that only waits, prefer :func:efficiency_span with the matching init/* category, so the span and the efficiency/* scalar describe the same interval.

nemo_rl.telemetry.instrumentation.evaluate_span(
algorithm: str,
tracer: Optional[opentelemetry.trace.Tracer] = None,
**attributes: Any,
) → Iterator[Any]#

One validation pass, named rl.<algorithm>.evaluate.

Validation generates through the same path as training rollouts, but its tokens are scored and thrown away – no weights advance – so the enclosed generate spans are overhead rather than goodput. The bucket scope is part of the helper because every algorithm’s validation has that property, and three of the six sites that predate it had already drifted without one.

The scope only bites on the sync rollout path, which is where those spans exist; async validation goes through generate_async, which carries no span yet. It is set regardless, so it applies as soon as that is instrumented.

nemo_rl.telemetry.instrumentation.checkpoint_finalize_span(
iteration: Optional[int],
*,
tracer: Optional[opentelemetry.trace.Tracer] = None,
**attributes: Any,
) → Iterator[Any]#

The asynchronous tail of a checkpoint, named rl.checkpoint.finalize.

Closes when the checkpoint is durable – every rank’s writer has flushed and tmp_step_N has been renamed to step_N – which is the moment a checkpoint can actually be resumed from. rl.<algo>.checkpointing cannot answer that: it ends when the training loop is released, which with async save is after D2H staging only, with the bytes still in flight.

iteration is passed rather than inherited, and that is the point of the helper. The work runs on a daemon thread spawned by begin_finalization and routinely outlives the step that triggered it, so the enclosing

Func:

iteration_scope is both unreachable (a raw thread starts with an empty context) and wrong by the time the span closes (the loop has moved on). Stamping it explicitly keeps the span filed under the step whose weights it holds, not whichever step happened to be running at the rename. Set directly rather than left to the scope because lens applies the caller’s attributes after the span starts, so an explicit value overwrites an inherited one.

Unbucketed, like the collector’s waits and for the same reason: it overlaps the next step’s productive spans, so an overhead tag here would charge one stretch of wall clock to two buckets at once. Reached through _managed_span rather than :func:managed_span to get that – the group stays checkpoint so one group switch still turns every checkpoint span on or off together.

Carries no async-save flag. Whether the bytes were still in flight is not knowable from here: every caller hands begin_finalization a wait function regardless of the setting, because those functions no-op when async save is off. The real setting lives per backend (Megatron’s mcore_state.cfg.checkpoint.async_save, Automodel’s is_async), and threading it through twelve call sites to label a span costs more than the label is worth. Read the span’s duration instead.

Parameters:

iteration – Step the checkpoint was triggered at, or None when it cannot be recovered from the directory name.

nemo_rl.telemetry.instrumentation.trace_fn(
group: str,
name: str,
tracer: Optional[opentelemetry.trace.Tracer] = None,
) → Callable[[nemo_rl.telemetry.instrumentation._F], nemo_rl.telemetry.instrumentation._F]#

Decorator that wraps a function in a bucket-tagged managed_span.

nemo_rl.telemetry.instrumentation.umbrella_trace_fn(
group: str,
name: str,
tracer: Optional[opentelemetry.trace.Tracer] = None,
) → Callable[[nemo_rl.telemetry.instrumentation._F], nemo_rl.telemetry.instrumentation._F]#

Decorator form of :func:umbrella_span, for whole-function umbrellas.

The rl.<algo>.job spans are all of this shape: one span over one call, covering everything the run does.