`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.

Module Contents#

Classes#

Bucket

Shared goodput buckets.

Functions#

`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.
`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.
`managed_span`	Like lens `managed_span`, but injects `rl.bucket` for leaf groups.
`efficiency_span`	Span for one efficiency category, tagged with that category’s bucket.
`trace_fn`	Decorator that wraps a function in a bucket-tagged `managed_span`.

Data#

`RL_BUCKET_ATTR`
`RL_EFFICIENCY_CATEGORY_ATTR`
`__all__`
`UMBRELLA_GROUPS`
`_DEFAULT_GROUP_BUCKET`
`EFFICIENCY_CATEGORY_BUCKET`
`COLLECTOR_LOOP_CATEGORIES`
`_BUCKET_OVERRIDE`

API#

nemo_rl.telemetry.instrumentation.RL_BUCKET_ATTR#: ‘rl.bucket’

nemo_rl.telemetry.instrumentation.RL_EFFICIENCY_CATEGORY_ATTR#: ‘rl.efficiency.category’

nemo_rl.telemetry.instrumentation.__all__#: [‘managed_span’, ‘trace_fn’, ‘span_cm’, ‘is_span_group_enabled’, ‘RL_BUCKET_ATTR’, ‘Bucket’, ‘UMBREL…

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]#: ‘frozenset(…)’

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._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.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.managed_span(

group: str,

name: str,

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.efficiency_span(

category: str,

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:COLLECTOR_LOOP_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.trace_fn(group: str, name: str, tracer=None)#: Decorator that wraps a function in a bucket-tagged managed_span.

nemo_rl.telemetry.instrumentation#