lens.strategies | NVIDIA NeMo Lens

Export strategy registry: maps strategy names to per-rank decision callables.

A strategy is a callable taking (config, rank, world_size) and returning True if this rank should export telemetry data. Lens ships four built-ins and lets users register their own under a chosen name; the chosen name is selected via NemoLensConfig.export_strategy or the NEMO_LENS_EXPORT_STRATEGY env var.

Module Contents

Functions

Name	Description
`_all_ranks`	Built-in: every rank exports.
`_first_rank_per_node`	Built-in: the first local rank on each node exports.
`_sampled`	Built-in: deterministic hash-based sampling. Same rank + same rate → same outcome.
`_single_rank`	Built-in: only `config.export_rank` exports (`-1` means last rank).
`get_export_strategy`	Look up a registered strategy by name.
`register_export_strategy`	Register a custom export strategy under `name`.
`registered_strategies`	Return a sorted list of registered strategy names (built-ins + custom).
`unregister_export_strategy`	Remove a strategy from the registry.

Data

BUILTIN_STRATEGIES

_REGISTRY

_REGISTRY_LOCK

API

lens.strategies._all_ranks(
    config: nemo.lens.config.NemoLensConfig,
    rank: int,
    world_size: int
) -> bool

Built-in: every rank exports.

lens.strategies._first_rank_per_node(
    config: nemo.lens.config.NemoLensConfig,
    rank: int,
    world_size: int
) -> bool

Built-in: the first local rank on each node exports.

Reads LOCAL_RANK (set by torchrun, deepspeed, etc.). Treats a missing LOCAL_RANK as "0", which gives the expected single-node behaviour.

lens.strategies._sampled(
    config: nemo.lens.config.NemoLensConfig,
    rank: int,
    world_size: int
) -> bool

Built-in: deterministic hash-based sampling. Same rank + same rate → same outcome.

lens.strategies._single_rank(
    config: nemo.lens.config.NemoLensConfig,
    rank: int,
    world_size: int
) -> bool

Built-in: only config.export_rank exports (-1 means last rank).

lens.strategies.get_export_strategy(
    name: str
) -> lens.strategies.ExportStrategy

Look up a registered strategy by name.

Raises:

ValueError: If name is not registered. Error message lists the currently registered names.

lens.strategies.register_export_strategy(
    name: str,
    strategy: lens.strategies.ExportStrategy,
    allow_override: bool = False
) -> None

Once registered, the strategy is selectable via NemoLensConfig(export_strategy=name) or NEMO_LENS_EXPORT_STRATEGY=<name>.

Parameters:

name

str

Strategy identifier. Must be a non-empty string.

strategy

ExportStrategy

Callable (config, rank, world_size) -> bool.

allow_override

boolDefaults to False

If False (default), raises ValueError when name is already registered. Built-in names are protected and additionally require allow_override=True to replace.

Raises:

ValueError: If name is empty, or already registered without allow_override=True.

lens.strategies.registered_strategies() -> list[str]

Return a sorted list of registered strategy names (built-ins + custom).

lens.strategies.unregister_export_strategy(
    name: str
) -> None

Remove a strategy from the registry.

Raises:

ValueError: If name refers to a built-in strategy, or is unknown.

lens.strategies.BUILTIN_STRATEGIES: frozenset[str] = frozenset({'all_ranks', 'single_rank', 'sampled', 'first_rank_per_node'})

lens.strategies._REGISTRY: dict[str, ExportStrategy] = {}

lens.strategies._REGISTRY_LOCK = threading.Lock()