> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/lens/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/lens/_mcp/server.

# Architecture

NeMo Lens is intentionally small. Its architecture prioritizes three properties:

1. **Low overhead when disabled**: instrumented code should pay a trivial gating cost (a `frozenset` lookup) when a span group is off, and never construct span objects.
2. **Optional dependency**: consumer libraries (such as Megatron-LM, NeMo RL, and NeMo Gym) can ship without bundling NeMo Lens.
3. **Clean layering**: each module has one responsibility and stable boundaries.

## Layer Diagram

```
┌─────────────────────────────────────────────────────────────┐
│ Public API (__init__.py)                                    │
│ — re-exports everything consumers use                       │
└─────────────────────────────────────────────────────────────┘
           │
           ▼
┌─────────────────────────────────────────────────────────────┐
│ Instrumentation layer                                        │
│ — helpers.py:  managed_span, trace_fn, span_cm              │
│ — distributed.py: broadcast_trace_context, create_linked_span│
│ — propagation.py: inject_context, extract_context           │
└─────────────────────────────────────────────────────────────┘
           │
           ▼
┌─────────────────────────────────────────────────────────────┐
│ State + config layer                                         │
│ — config.py:  NemoLensConfig                                │
│ — state.py:   enabled span groups, PP trace carrier         │
│ — groups.py:  SpanGroup + presets                           │
└─────────────────────────────────────────────────────────────┘
           │
           ▼
┌─────────────────────────────────────────────────────────────┐
│ Lifecycle + providers                                        │
│ — handle.py:     setup_telemetry, TelemetryHandle           │
│ — providers.py:  build_providers, build_noop_providers      │
│ — sampling.py:   RankAwareSampler                           │
└─────────────────────────────────────────────────────────────┘
           │
           ▼
┌─────────────────────────────────────────────────────────────┐
│ Extras                                                        │
│ — instruments/  inference, rl, gym metrics                  │
│ — resources/    SLURM, K8s, local detection                 │
│ — contrib/      FastAPI, aiohttp, Ray, NCCL                 │
│ — logging_bridge.py                                          │
│ — fallbacks.py  canonical no-ops for consumers              │
│ — semconv.py    attribute name constants                    │
└─────────────────────────────────────────────────────────────┘
```

Each layer depends only on layers above it. No circular imports.

## Module Responsibilities

| Module              | Responsibility                                                     |
| ------------------- | ------------------------------------------------------------------ |
| `__init__.py`       | Public API re-exports                                              |
| `config.py`         | `NemoLensConfig` dataclass, env var parsing, validation            |
| `state.py`          | Thread-safe global state: enabled span groups and PP trace carrier |
| `groups.py`         | `SpanGroup` base class, preset resolution                          |
| `helpers.py`        | `managed_span`, `trace_fn`, `span_cm`, attribute safety            |
| `handle.py`         | `setup_telemetry` entry point, `TelemetryHandle`, export strategy  |
| `providers.py`      | SDK provider construction (lazy imports)                           |
| `propagation.py`    | W3C context inject/extract                                         |
| `distributed.py`    | `broadcast_trace_context`, `create_linked_span`                    |
| `sampling.py`       | `RankAwareSampler` (OTel `Sampler` impl)                           |
| `semconv.py`        | Attribute name constants and version tracking                      |
| `fallbacks.py`      | Canonical no-op implementations for consumers                      |
| `logging_bridge.py` | Python logging → OTel LoggerProvider bridge                        |

## Call Flow for `setup_telemetry`

```
setup_telemetry(config, rank, world_size, resource_attributes)
  │
  ├─ if _INITIALIZED and config.enabled and not _allow_reinit:
  │     raise RuntimeError          # [double-init guard]
  │
  ├─ if not config.run_id:
  │     config.run_id = SLURM_JOB_ID or uuid4()
  │
  ├─ is_export_rank = _should_export(config, rank, world_size)
  │     └─ branches on config.export_strategy
  │
  ├─ if not config.enabled:
  │     build_noop_providers()
  │     set_enabled_span_groups(frozenset())
  │
  ├─ elif is_export_rank:
  │     build_providers(config, rank, world_size, resource_attributes, ...)
  │         ├─ imports SDK (primary SDK construction site)
  │         ├─ builds Resource with auto-detected + passed attributes
  │         ├─ builds TracerProvider (with optional RankAwareSampler)
  │         ├─ builds MeterProvider (with PeriodicExportingMetricReader)
  │         ├─ builds LoggerProvider (if config.logs_enabled)
  │         └─ sets W3C CompositePropagator
  │     set_enabled_span_groups(config.resolved_span_groups)
  │
  └─ else:    # non-exporting rank
        build_noop_providers()
        set_enabled_span_groups(frozenset())

  returns TelemetryHandle(tracer, meter, is_exporting)
```

## Call Flow for `managed_span`

```
with managed_span('step', 'train.step', iteration=42) as span:
  │
  ├─ is_span_group_enabled('step')   # frozenset lookup
  │
  ├─ if not enabled:
  │     yield None
  │     return                        # [disabled path: no span created]
  │
  ├─ tracer = tracer or trace.get_tracer(__name__)   # default tracer name "nemo.lens.helpers"
  ├─ span = tracer.start_span('train.step')
  ├─ safe_set_span_attributes(span, {'iteration': 42})
  ├─ token = context.attach(set_span_in_context(span))
  │
  ├─ try:
  │     yield span
  │   except Exception as exc:
  │     span.record_exception(exc)
  │     span.set_status(StatusCode.ERROR, str(exc))
  │     raise
  │   finally:
  │     context.detach(token)
  │     span.end()
```

The hot path when disabled is three Python statements: lookup, compare, yield `None`.

## Why This Factoring

**State separated from logic**: `state.py` is just a frozenset and a lock. `managed_span` in `helpers.py` imports `is_span_group_enabled`, but does not know how groups are stored. Swapping the state implementation (e.g., to add per-thread overrides) does not touch `helpers.py`.

**Config separated from providers**: `config.py` has no dependency on OTel SDK. A consumer can construct and validate `NemoLensConfig` in a process that does not have the SDK installed, then decide whether to initialize telemetry.

**Lazy SDK imports**: `providers.py` is the primary home for `opentelemetry.sdk.*` construction (a few other modules, such as `sampling.py` and `logging_bridge.py`, import the SDK lazily inside function bodies). Non-exporting ranks never execute `build_providers`, so they never incur the SDK import cost. On a large-rank job where most ranks do not export, avoiding that import on every non-exporting rank adds up.

**Public API minimal**: `__init__.py` exports only the documented public surface (the entries in `__all__`). Everything else is internal. Consumers can rely on the `__all__` list; internals can be refactored freely.

## Thread Safety

The NeMo Lens global state is protected where it matters:

* `state._ENABLED_GROUPS` is written under a lock, read lock-free (a frozenset is immutable so readers see a consistent snapshot).
* `state._PP_TRACE_CARRIER` is written during `setup_telemetry` (single-threaded) and read from multiple threads in the pipeline schedule.
* Instrument caches use `WeakKeyDictionary`, which is thread-safe under CPython's GIL for the operations NeMo Lens performs.

`setup_telemetry` itself is not thread-safe; it is documented as "call once per process." The [double-init guard](/nemo/lens/design/double-init-guard) surfaces violations.

## What Is Out of Scope

NeMo Lens is deliberately narrow. Log shipping beyond the OTel logging bridge remains with your existing log stack. Profiling CPU, memory, or GPU activity is the responsibility of the PyTorch profiler or `nsys`, not a tracing library. APM-style features such as service maps and error tracking are handled by whichever backend consumes the OTLP stream (such as Datadog, Sentry, or other compliant platforms).

The scope is also restricted to the NeMo ecosystem's training and inference workloads. The core primitives are domain-agnostic, but the opinionated components (span groups, metric instruments, and resource attributes) lean toward machine learning workloads. NeMo Lens does not serve as a general-purpose instrumentation library for unrelated domains.