nemo_rl.telemetry.span_groups

NeMo-RL span groups, declared into lens’s SpanRegistry.

Lens ships no span-group names of its own: a consuming library registers what it emits under its own namespace and users select from that with telemetry.span_groups. So the groups that used to arrive from lens’s base class – job, step, checkpoint, evaluate, model_init – are declared here alongside the RL-specific ones. They keep their original names, because the names are the user-facing vocabulary and Megatron registers some of the same ones for the same phases; a group two libraries both register is shared, which is the intended behaviour for a job that drives both.

Only names NeMo-RL actually emits are registered. Lens’s remaining base names (load_checkpoint, forward_backward, optimizer) are absent because no call site here opens them, and registering a group that emits nothing makes a preset advertise spans it cannot deliver. They stay selectable in a job that also drives Megatron: the driver hands workers the raw spec string, so a worker that imports Megatron resolves the name against Megatron’s own registration.

RLSpanGroup is a bag of str constants rather than an enum or a lens subclass, because managed_span and trace_fn take the group as a plain string. Call sites keep reading RLSpanGroup.GENERATION for discoverability and to keep the spelling in one place.

Registration happens at import, which is what makes these groups selectable: lens resolves the span_groups spec against whatever is registered at the time, so this module must be imported before setup_telemetry. nemo_rl.telemetry.setup imports it directly for that reason.

Module Contents

Classes

RLSpanGroup

Span group names for NeMo-RL instrumentation.

Functions

_group_values	Group names declared on :class:RLSpanGroup, by spelling.
register_span_groups	Declare NeMo-RL’s groups and presets to lens.

Data

NAMESPACE
UMBRELLA_GROUP_VALUES

API

nemo_rl.telemetry.span_groups.NAMESPACE

‘nemo_rl’

class nemo_rl.telemetry.span_groups.RLSpanGroup

Span group names for NeMo-RL instrumentation.

JOB

‘job’

The whole run: one span wrapping everything from setup to teardown.

CHECKPOINT

‘checkpoint’

Checkpoint save spans.

EVALUATE

‘evaluate’

Validation / evaluation spans.

MODEL_INIT

‘model_init’

Model construction and sharding inside a worker.

STEP

‘step’

One training step.

SETUP

‘setup’

Driver-side startup phases: Ray init, worker builds, collective init.

Distinct from MODEL_INIT, which covers the model load itself inside a worker. These are the driver’s view of the same startup — the phases between process start and the first training step.

ROLLOUT

‘rollout’

Rollout collection spans.

GENERATION

‘generation’

Text generation spans.

LOGPROB

‘logprob’

Log-probability computation spans.

REWARD

‘reward’

Reward computation spans.

ADVANTAGE

‘advantage’

Advantage computation spans.

POLICY_UPDATE

‘policy_update’

Policy gradient update spans.

DATA_PROCESSING

‘data_processing’

Data processing / batching spans.

DATA_PLANE

‘data_plane’

Transfer-queue / data-plane operations (put, claim, get, clear).

Its own group rather than part of DATA_PROCESSING because it is far finer-grained: one span per data-plane RPC, several per step, so a user who finds the volume noisy can drop these without losing the coarse phases.

PER_PROMPT

‘per_prompt’

Spans emitted once per prompt rather than once per step or batch.

A cardinality axis, unlike every other group here, which names a phase. The two axes are independent: rl.sc.generate_and_push and the rollout path’s rl.data_plane.put are a rollout span and a data-plane span respectively, but what governs whether a user wants them is neither of those things — it is that their count scales with the prompt count, so a 10k-prompt rollout emits ~20k spans where the phase groups emit a fixed handful per step.

Kept out of the per_step preset for that reason: per_step is meant to be the detailed-but-usable choice, and its span count should scale with steps, not with dataset size. Reach for these with span_groups: "per_step,per_prompt" when debugging an individual rollout, or all.

An umbrella group, so its spans carry no rl.bucket — see instrumentation.UMBRELLA_GROUPS for why concurrency forces that.

EFFICIENCY

‘efficiency’

Async efficiency phases (idle / wasted accounting).

Unlike the other leaf groups these do not have one fixed bucket — the rl.bucket comes from the category, so emit them via instrumentation.efficiency_span.

U_JOB

None

U_STEP

None

U_MODEL_INIT

None

U_EVALUATE

None

An eval pass; an umbrella unless its inside is timed as idle.

U_ROLLOUT

None

The collect_rollouts umbrella.

U_SETUP

None

Startup is real overhead, but no subset of these spans is summable.

The phases nest (rl.startup over rl.setup.workers over rl.vllm.load_model) and the worker builds run concurrently under parallel init, so a rollup adding them by rl.bucket would multiply startup rather than measure it. The flat number lives in the rl.setup.duration metric at phase=total_setup, which is the one value that cannot double-count. These spans are for shape only.

U_PER_PROMPT

None

Per-prompt work on the single-controller path overlaps itself.

Many rollouts are in flight at once (max_inflight_prompts), so any bucket these carried would sum to a large multiple of the wall clock they happened in. That applies to the data-plane put inside a rollout as much as to the rollout span itself, which is why this group overrides DATA_PLANE’s overhead bucket there – see per_prompt_scope.

ALL_GROUPS: ClassVar[frozenset[str]]

None

_PRESETS: ClassVar[dict]

None

classmethod resolve(spec: str) → frozenset

Resolve a span_groups spec to the group names it selects.

Thin wrapper over :meth:SpanRegistry.resolve that drops the pending half of the return. Unlike the fixed-membership scheme this replaces, an unrecognised entry is not an error — the library that owns it may simply not have been imported — so callers that want to report typos should use

Meth:

resolve_with_pending.

classmethod resolve_with_pending(spec: str) → tuple[frozenset, frozenset]

(enabled, pending) for spec, straight from the registry.

pending holds spec entries that matched nothing registered. For NeMo-RL that is very likely a typo rather than a not-yet-imported library, which is why the driver reports it.

nemo_rl.telemetry.span_groups._group_values(*, umbrella: bool) → frozenset[str]

Group names declared on :class:RLSpanGroup, by spelling.

Read off the class rather than listed by hand: a group added above but forgotten in a list is never registered with lens, so span_groups: all leaves it out and selecting it yields silence rather than an error.

nemo_rl.telemetry.span_groups.UMBRELLA_GROUP_VALUES: Final[frozenset[str]]

‘_group_values(…)’

nemo_rl.telemetry.span_groups.register_span_groups(*, allow_override: bool = True) → None

Declare NeMo-RL’s groups and presets to lens.

Called at import. Idempotent by default so a re-import, or a test that cleared the registry, can safely call it again – re-registering replaces the namespace wholesale, which is exactly the intent here since everything is passed every time.