Span Groups

Span groups are the NeMo Lens mechanism for controlling trace granularity at runtime without code changes. Every instrumentation site tags itself with a group name; at startup, only the enabled groups actually emit spans.

NeMo Lens defines no group names of its own. Each consuming library declares what it emits by registering with SpanRegistry, and users select from the union of everything registered.

Why Use Span Groups

A training job can use different levels of tracing detail for different environments:

• Minimize production overhead. Enable a coarse preset to emit only job-level spans for the lowest performance cost.
• Add iteration boundaries. Enable a per-step preset in staging to add per-iteration boundaries for a moderate performance cost.
• Diagnose performance hangs. Enable all during active debugging to instrument every registered site, which carries the highest performance cost.

The primary goal is to toggle granularity through a single env var without code changes, relying on a fast gating check when disabled: is_span_group_enabled is one frozenset membership test.

Register the Groups Your Library Emits

from nemo.lens import SpanRegistry
SpanRegistry.register(
    "megatron",
    groups={"step", "microbatch", "layer", "communication", "inference"},
    presets={
        "default": {"inference"},
        "per_step": {"step", "microbatch"},
    },
)

Group names share one flat namespace across libraries, which keeps instrumentation call sites terse. If two libraries claim the same name, they share it: the name maps to a single gate, and enabling it turns on both libraries’ spans. This logs a warning rather than raising, because it happens inside library code during import and neither library knows it is second. Pass allow_override=True when you mean it.

Presets Union Across Libraries

An RL job driving Megatron has both libraries registered in one process. Presets merge rather than compete:

SpanRegistry.register("megatron", {"step"}, {"default": {"step"}})
SpanRegistry.register("nemo_rl", {"rollout"}, {"default": {"rollout"}})
SpanRegistry.resolve("default")   # -> ({"step", "rollout"}, frozenset())

This is why registration replaced the SpanGroup subclass hook it succeeds: a subclass could only ever describe one library, and a subclass that overrode _PRESETS silently discarded the base class’s.

all is built in and always means “every group currently registered”. It is not a fixed list, and it cannot be used as a preset name.

Select Groups at Runtime

export NEMO_LENS_SPAN_GROUPS=default          # a preset
export NEMO_LENS_SPAN_GROUPS=per_step,layer   # a preset plus one extra group
export NEMO_LENS_SPAN_GROUPS=all              # everything registered

The spec is a comma-separated, case-insensitive mix of preset names and bare group names.

Register Before setup_telemetry

Each library owns its own telemetry, so importing a library is what registers its groups. Do that before initializing telemetry — by then everything the process uses is known, and the spec can be validated against it:

import megatron.core.telemetry     # registers "megatron"
import nemo_rl.telemetry           # registers "nemo_rl"
setup_telemetry(NemoLensConfig.from_env())   # validates NEMO_LENS_SPAN_GROUPS

A spec entry that names nothing registered in this process is reported as a warning naming what is available, and everything else in the spec still applies:

WARNING Span group spec 'per_stpe' names ['per_stpe'], which no library
registered in this process provides; nothing will be emitted for those.
Registered groups: ['layer', 'microbatch', 'step']. Presets: ['all', 'default',
'per_step']. Namespaces: ['megatron']. ...

It is a warning and not an error on purpose. A registry is per process; a spec is usually job-wide. A launcher agent and a spawned checkpoint worker inherit one NEMO_LENS_SPAN_GROUPS from the trainer but import a different set of libraries, so a value that is perfectly valid in the trainer names nothing in them. That is not a typo, and telemetry setup must not take those processes down over it — nor leave them exporting with no handle to shut down.

If a process needs its own vocabulary, give it its own env prefix rather than trying to find one spec that suits every process:

NemoLensConfig.from_env(prefix="NVRX_OTEL", fallback_prefix="NEMO_LENS")

If your library’s telemetry module is imported lazily at the instrumentation site rather than eagerly with the package, it will not have registered in time. Register it from the package’s __init__.

Registering after setup_telemetry still takes effect — NeMo Lens retains the spec and re-resolves it — but it logs a warning, because the group was not selectable when the user’s spec was validated. It is a consumer import-order bug, treated as loud rather than fatal: refusing the registration would drop spans silently, and raising would let a telemetry misconfiguration take down a training job.

Reference Another Library’s Groups

A preset may name any group already in the registry, not only the ones the same call declares. That is how a library layers on one it depends on:

SpanRegistry.register(
    "nemo_rl",
    groups={"rollout"},
    presets={"default": {"rollout", "step"}},   # "step" belongs to megatron
)

Import order does not matter, and Megatron does not have to be present at all. Each read compares a preset against the groups currently registered, so a member no one owns is ignored, and it starts working if its owner is imported later. A library can therefore name a group from an optional dependency without breaking its own import when that dependency is missing. It logs a warning, because a typo produces exactly the same situation.

Referencing a group does not transfer ownership: if Megatron unregisters, step leaves the registry, and any preset that borrowed it drops it too — a preset can never name a group outside all. If that leaves the preset your spec named empty, NeMo Lens warns that the spec resolved to nothing. pending_span_groups() is empty in this case, so without that warning there would be nothing to go on.

Inspect What Resolved

from nemo.lens import enabled_span_groups
from nemo.lens.groups import SpanRegistry
enabled_span_groups()        # frozenset({'step', 'rollout'})
SpanRegistry.namespaces()    # ['megatron', 'nemo_rl']
SpanRegistry.groups()        # every group any library registered

Use Groups at Instrumentation Sites

Group names are plain strings; nothing needs importing from NeMo Lens:

from nemo.lens import managed_span, trace_fn
with managed_span("step", "train.step") as span:
    if span:
        span.set_attribute("dl.iteration", i)
@trace_fn("forward_backward", "dl.forward_backward")
def forward_backward(...):
    ...

Keep attribute construction inside the if span: guard. Work done before the group check is paid by every user on every call, whether or not they enabled telemetry.

Naming Conventions

• Lowercase, snake_case: pipeline_parallel, not PipelineParallel.
• Concept-oriented, not verb-oriented: checkpoint, not saving_checkpoint.
• Short: evaluate, not evaluation_phase.
• Because the namespace is flat, prefer a name unlikely to collide with another NeMo library’s.

Choose Preset Membership Carefully

The convention for default is that it stays quiet enough for production. If a group fires more than a handful of times per iteration, it belongs in per_step or in no preset at all. Since presets union across libraries, adding a chatty group to your default raises the floor for every user who enables default — including users of a different library that happens to share the process.