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.

Span groups are registered once via :func:set_enabled_span_groups (called from :func:~nemo.lens.handle.setup_telemetry). Before that call every :func:is_span_group_enabled query returns False.

Module Contents

Functions

Name	Description
`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.
`set_enabled_span_groups`	Register the active span groups.
`set_pp_trace_carrier`	Store the pipeline-parallel trace carrier for cross-stage linking.

Data

_ENABLED_GROUPS

_LOCK

_PP_TRACE_CARRIER

API

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 :func:set_enabled_span_groups is called.

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

Called once from :func:~nemo.lens.handle.setup_telemetry. Subsequent calls override the previous value (useful for testing).

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._ENABLED_GROUPS: frozenset = frozenset()

lens.state._LOCK = threading.Lock()

lens.state._PP_TRACE_CARRIER: dict | None = None