Custom Export Strategies | NVIDIA NeMo Lens

NeMo Lens picks which ranks export by name: NemoLensConfig.export_strategy selects a callable from a process-wide registry. The four built-in strategies cover the common cases (see Sampling); the registry is the extension point for everything else, including fleet-aware sampling, locality-aware selection, and anything that needs more context than the built-in strategies expose.

Built-In Strategies

These strategies are registered automatically at import time and exposed as the frozenset nemo.lens.strategies.BUILTIN_STRATEGIES (this symbol lives on the submodule and is not a top-level nemo.lens export, so import it as from nemo.lens.strategies import BUILTIN_STRATEGIES):

single_rank (default): a single rank exports (config.export_rank, where -1 means the last rank).
all_ranks: every rank exports.
sampled: a deterministic hash-based fraction of ranks, controlled by config.export_sample_rate.
first_rank_per_node: the rank with LOCAL_RANK=0 on each node. This reads LOCAL_RANK from the environment (set by torchrun, deepspeed, etc.); a missing value is treated as "0".

These built-in strategies cannot be unregistered, and replacing them requires allow_override=True.

Select a Strategy

Select a strategy by name, through the config or env var:

1 from nemo.lens import NemoLensConfig, setup_telemetry
2 
3 cfg = NemoLensConfig(enabled=True, export_strategy="first_rank_per_node")
4 handle = setup_telemetry(cfg, rank=rank, world_size=world_size)

$ NEMO_LENS_EXPORT_STRATEGY=first_rank_per_node

Alternatively, pass the callable directly to setup_telemetry as an ad-hoc, one-off strategy:

1 handle = setup_telemetry(
2     cfg,
3     rank=rank,
4     world_size=world_size,
5     export_strategy=lambda config, rank, world_size: rank in {0, 7, 42},
6 )

The export_strategy= argument bypasses the registry, which is useful when you want a one-line strategy without polluting a process-wide name.

Register a Custom Strategy

1 from nemo.lens import register_export_strategy, NemoLensConfig, setup_telemetry
2 
3 def first_two_ranks(config, rank, world_size):
4     return rank < 2
5 
6 register_export_strategy("first_two_ranks", first_two_ranks)
7 
8 cfg = NemoLensConfig(enabled=True, export_strategy="first_two_ranks")
9 handle = setup_telemetry(cfg, rank=rank, world_size=world_size)

Register the strategy before calling setup_telemetry. Validation is lazy, so unknown names raise a ValueError at setup_telemetry time, not during NemoLensConfig construction.
The registry is process-wide. Call register_export_strategy exactly once per process, typically near the import-time setup.
allow_override is keyword-only. Pass allow_override=True only when you intentionally replace an existing entry. Without it, re-registering an already-registered name raises a ValueError. An empty name also raises a ValueError. Built-in strategies can only be replaced with allow_override=True, and they can never be unregistered.

Strategy Callable Signature

1 ExportStrategy = Callable[[NemoLensConfig, int, int], bool]

A strategy receives the resolved NemoLensConfig, the global rank, and the world_size. It returns True if this rank should export, False otherwise. The decision runs once per process at setup_telemetry time.

Strategies can read environment variables to obtain context omitted from function arguments, such as LOCAL_RANK, NODE_RANK, SLURM_PROCID, or other launcher-specific variables. The built-in first_rank_per_node strategy demonstrates this approach.

Discovery and Introspection

1 from nemo.lens import registered_strategies, unregister_export_strategy
2 
3 registered_strategies()
4 # ['all_ranks', 'first_rank_per_node', 'first_two_ranks', 'sampled', 'single_rank']
5 
6 unregister_export_strategy("first_two_ranks")

registered_strategies() returns a sorted list of every name currently in the registry. unregister_export_strategy(name) removes a custom entry; it raises a ValueError if name refers to a built-in strategy or is not registered.

register_export_strategy, registered_strategies, unregister_export_strategy, and ExportStrategy are top-level exports (from nemo.lens import ...). The lower-level get_export_strategy(name) lookup, like BUILTIN_STRATEGIES, is not a top-level export and must be imported from nemo.lens.strategies.

Why This API

Name-based dispatch lets the same env var (NEMO_LENS_EXPORT_STRATEGY) route to user code without touching NeMo Lens internals, allowing production runs and CI to pick a strategy from outside the program. The setup_telemetry(export_strategy=) argument is the lower-level escape hatch: pass any callable, skip the registry, and never name it. This parallels the span_exporter= argument documented in Custom Exporters, both of which keep the core surface small while staying open at the bottom for callers that need it.