lens.state | NVIDIA NeMo Lens

Module-level span group state — importable anywhere without circular deps.

Holds a frozenset of enabled span groups so that any module can call :func:is_span_group_enabled without importing the full nemo.lens package.

Two ways in, and the last one called wins:

:func:set_span_group_spec stores the raw user spec (e.g. "default,step") and resolves it against :class:~nemo.lens.groups.SpanRegistry. It never raises: an entry naming nothing is warned about and kept pending, because this process’s registry is not the authority on a job-wide spec. Several processes in one job routinely share a spec while importing different libraries, so a name that resolves nowhere here may be perfectly valid next door.

The spec is retained, so a library registering late still takes effect (:func:refresh_enabled_span_groups) instead of silently emitting nothing. That path warns — see SpanRegistry._warn_if_late.
:func:set_enabled_span_groups pins an explicit set and drops the spec, so a later registration cannot reopen it. This is how a disabled process stays disabled, and how a test enables exactly the groups it means to.

Before either is called every :func:is_span_group_enabled query returns False.

Module Contents

Functions

Name	Description
`_report`	Warn about a spec that did not get the caller what they asked for.
`_warn_key`	The parts of a resolution the warning text depends on.
`enabled_span_groups`	The groups currently enabled. Diagnostic; not for the hot path.
`get_pp_trace_carrier`	Return the current PP trace carrier, or None.
`is_span_group_enabled`	Return `True` if the named span group is currently enabled.
`pending_span_groups`	Spec entries that resolved to nothing. Diagnostic; not for the hot path.
`refresh_enabled_span_groups`	Re-resolve the stored spec. Called whenever the registry changes.
`set_enabled_span_groups`	Pin the active span groups explicitly, discarding any stored spec.
`set_pp_trace_carrier`	Store the pipeline-parallel trace carrier for cross-stage linking.
`set_span_group_spec`	Store the user’s span-group spec and resolve it against the registry.

Data

API

lens.state._report(
    spec: str,
    resolution: nemo.lens.groups._Resolution
) -> None

Warn about a spec that did not get the caller what they asked for.

Deliberately a warning and not an exception: see the module docstring. Carries what is registered, because that is what turns “I set SPAN_GROUPS and got nothing” into a diagnosis — either a typo, or a library this process never imported.

Everything reported comes from resolution, which was taken in a single registry hold. Building the message from a fresh query let a registration land in between, so a warning could call a group unresolved and also list it as registered.

lens.state._warn_key(
    resolution: nemo.lens.groups._Resolution
) -> tuple[frozenset, bool, bool]

The parts of a resolution the warning text depends on.

lens.state.enabled_span_groups() -> frozenset

The groups currently enabled. Diagnostic; not for the hot path.

lens.state.get_pp_trace_carrier() -> dict | None

Return the current PP trace carrier, or None.

lens.state.is_span_group_enabled(
    group: str
) -> bool

Return True if the named span group is currently enabled.

This is the primary check at every instrumentation site (~2ns overhead). Returns False before any group has been enabled.

lens.state.pending_span_groups() -> frozenset

Spec entries that resolved to nothing. Diagnostic; not for the hot path.

Non-empty here is the usual explanation for “I set SPAN_GROUPS and got nothing”: either a typo, or a name belonging to a library this process never imported.

lens.state.refresh_enabled_span_groups() -> None

Re-resolve the stored spec. Called whenever the registry changes.

A no-op when the groups were pinned explicitly rather than from a spec. Registration happens at import time, so the cost never lands on the hot path — :func:is_span_group_enabled stays one frozenset membership test.

lens.state.set_enabled_span_groups(
    groups: frozenset
) -> None

Pin the active span groups explicitly, discarding any stored spec.

Subsequent registrations will not change the set — use :func:set_span_group_spec if you want it to track the registry.

lens.state.set_pp_trace_carrier(
    carrier: dict | None
) -> None

Store the pipeline-parallel trace carrier for cross-stage linking.

Called from the training loop after :func:broadcast_trace_context.

lens.state.set_span_group_spec(
    spec: str
) -> None

Store the user’s span-group spec and resolve it against the registry.

Called from :func:~nemo.lens.handle.setup_telemetry. Never raises: an entry that resolves to nothing is warned about and kept pending, because this process’s registry is not the authority on a job-wide spec.

lens.state._ENABLED_GROUPS: frozenset = frozenset()

lens.state._LOCK = threading.Lock()

lens.state._PENDING: frozenset = frozenset()

lens.state._PP_TRACE_CARRIER: dict | None = None

lens.state._SPEC: str = ''

lens.state._WARNED: tuple[frozenset, bool, bool] | None = None