Skip to navigation
Developer Guide

Register a Span Group

View as Markdown

Span groups are a runtime knob. You register one, wire instrumentation to it, and users get a new opt-in slice of telemetry without any code changes on their side. This page is for library authors.

See Span Groups for how groups work at runtime.

NeMo Lens Defines No Groups

There is no base set to extend and no SpanGroup class to subclass. Every group name belongs to the library that emits spans for it. Adding one means adding it to your own SpanRegistry.register call.

Register at Import Time

Put the registration where your telemetry module is first imported, so it runs before any instrumented code:

# Megatron-LM/megatron/core/telemetry/span_groups.py
from nemo.lens import SpanRegistry
NAMESPACE = "megatron"
GROUPS = {
    "step",
    "microbatch",
    "layer",
    "communication",
    "activation_offload",
    "data_loading",
    "inference",
    "my_new_group",          # <- the addition
}
SpanRegistry.register(
    NAMESPACE,
    groups=GROUPS,
    presets={
        "default": {"inference"},
        "per_step": {"step", "microbatch", "my_new_group"},
    },
)

Checklist:

  • Name added to your groups set.
  • Name added to the presets you intend it to appear in, if any.
  • Instrumentation added at the relevant call site(s) using managed_span("my_new_group", ...) or @trace_fn("my_new_group", ...).
  • Entry added to your library’s observability docs.
  • Test covering that the group resolves and that a wrapped site emits when enabled.

Register before setup_telemetry runs. Importing your library is what registers its groups, so put the call where the package’s own import reaches it — not lazily at the instrumentation site, which would be too late for the user’s spec to be validated against it.

Registering late still takes effect, but logs a warning. Any name a user puts in NEMO_LENS_SPAN_GROUPS that you have not registered by then is warned about and ignored — never raised, because in a multi-process job that name may belong to a library another process imports.

Register your groups under the conventional preset names (default, per_step) rather than only default. Presets union by name, so a spec of per_step leaves your spans dark in any process where you registered no per_step.

Choose a Namespace

One namespace per library — "megatron", "nemo_rl", "nemo_gym". It is the key for SpanRegistry.unregister, it appears in the collision and late-registration warnings, and it is what lets two libraries coexist in one process.

Group names are flat and shared across namespaces. If two libraries claim the same name, they share it: enabling that name turns on both libraries’ spans, and both get a warning. allow_override=True silences the warning. Pick names that are unlikely to clash, but do not rely on it — your library has no way to know what else is loaded.

Reference a Dependency’s Groups

If your library sits on top of another and wants that library’s groups in your presets, import it first and name them directly:

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

Import order does not matter. If nothing has registered step yet, the member is kept but ignored, and it starts working once Megatron registers. This means you can name a group from an optional dependency: your library still imports when that dependency is missing. It does log a warning, because a typo produces exactly the same situation.

Naming Conventions

  • Lowercase, snake_case: pipeline_parallel, not PipelineParallel or pipeline-parallel.
  • Concept-oriented, not verb-oriented: checkpoint, not saving_checkpoint.
  • Short: evaluate, not evaluation_phase.
  • Skim your existing set before inventing; if a nearby concept already has a group, extend instrumentation there rather than adding a new one.

Preset Inclusion

Do not add to default casually. The convention for default is that it stays quiet enough for production (structural job-level spans only). If your group fires more than a handful of times per iteration, it belongs in per_step or in no preset.

Presets union across namespaces, so a chatty group in your default raises the floor for every user who enables default — including users of a different library sharing the process. That is a wider blast radius than the subclass scheme this replaced, where each subclass owned its presets outright.

A reasonable default: add to per_step only, and wait for a concrete reason before promoting.

Changes Needed in fallbacks.py

None for a new group — group names are plain strings, and the no-op is_span_group_enabled returns False regardless of the name.

nemo.lens.fallbacks does provide a no-op SpanRegistry, because a consuming library calls SpanRegistry.register at import time and that has to keep working when NeMo Lens is absent:

try:
    from nemo.lens import SpanRegistry
except ImportError:
    from nemo.lens.fallbacks import SpanRegistry

If you change the signature of a SpanRegistry method in NeMo Lens, the no-op must match, including default values and not just parameter names. tests/test_fallbacks.py enforces this.

Testing

See Test for fixture conventions. Because NeMo Lens ships no groups, a test that needs one must register it:

def test_my_new_group_in_per_step_preset():
    SpanRegistry.register("mylib", {"my_new_group"}, {"per_step": {"my_new_group"}})
    cfg = NemoLensConfig(enabled=True, span_groups="per_step", exporter="console")
    setup_telemetry(cfg)
    assert is_span_group_enabled("my_new_group")

The reset_span_groups autouse fixture in conftest.py clears the registry around every test, so nothing carries over.

At minimum, cover:

  • The group resolves from its bare name.
  • The group is included in every preset you added it to, and excluded from the ones you didn’t.
  • An instrumentation site wrapped in managed_span("my_new_group", ...) emits a span when the group is enabled, and emits nothing when it is not.