nemo_rl.telemetry.setup#

Process-global nemo-lens telemetry lifecycle for NeMo-RL.

Two entry points, mirroring Megatron’s global_vars._set_telemetry / get_telemetry_handle pattern but adapted to NeMo-RL’s Ray driver + worker process model:

  • func:

    init_telemetry_driver — called once on the driver, before init_ray(). It reads the telemetry: config block, exports the settings as NEMO_RL_OTEL_* env vars so every Ray worker inherits them, and sets up the driver’s own telemetry (the training loop and the metrics logger run here, so the driver always exports).

  • func:

    init_telemetry_worker — called once inside each Ray actor process, from the worker’s __init__ (policy, value and vLLM generation workers). It reads the propagated env and sets up that worker’s telemetry, then

    func:

    shutdown_telemetry flushes it from the worker’s shutdown. __init__ rather than post_init, because some post_init fan-outs run on only one rank per parallel group, while OTel providers have to be set up in every actor process.

nemo-lens is a base dependency, so telemetry.enabled is the only switch: when it is false the init functions return None and every instrumentation site is a ~0-cost no-op. Lens imports stay function-local to keep the cost off the import path of modules that never emit anything.

Module Contents#

Functions#

_is_env_truthy

Return True if env var name is set to a truthy value.

telemetry_enabled_in_env

Whether telemetry.enabled reached this process as truthy.

vllm_native_tracing_requested

Whether telemetry.vllm_native_tracing reached this process as truthy.

_config_to_env

Translate a TelemetryConfig into NEMO_RL_OTEL_* env vars.

_dig

Best-effort nested lookup that works for both dicts and objects.

_build_resource_attributes

Build process-lifetime resource attributes (Jaeger “Process” tags).

_always_export

Export-strategy override for singleton processes: always export.

_unrank

Disable rank-based span filtering for a process that has no real rank.

init_telemetry_driver

Initialise driver-side telemetry (call once, before init_ray()).

_worker_resource_attributes

Build resource attributes identifying this worker process.

init_telemetry_worker

Initialise telemetry inside a Ray actor (call once per worker process).

get_telemetry_handle

Return the process-global telemetry handle (None if uninitialised).

shutdown_telemetry

Flush and shut down telemetry providers.

Data#

logger

_TELEMETRY_HANDLE

_TELEMETRY_INITIALISED

_OTEL_PREFIX

_OTEL_FALLBACK_PREFIX

_RUN_ID_ENV

_WORKER_GROUP_ENV

_ENV_FIELD_MAP

_SERVICE_NAME_ENV

API#

nemo_rl.telemetry.setup.logger#

‘getLogger(…)’

nemo_rl.telemetry.setup._TELEMETRY_HANDLE: Optional[nemo.lens.TelemetryHandle]#

None

nemo_rl.telemetry.setup._TELEMETRY_INITIALISED#

False

nemo_rl.telemetry.setup._OTEL_PREFIX#

‘NEMO_RL_OTEL’

nemo_rl.telemetry.setup._OTEL_FALLBACK_PREFIX#

‘NEMO_LENS’

nemo_rl.telemetry.setup._RUN_ID_ENV#

None

nemo_rl.telemetry.setup._WORKER_GROUP_ENV#

‘NRL_WORKER_GROUP’

nemo_rl.telemetry.setup._ENV_FIELD_MAP#

None

nemo_rl.telemetry.setup._SERVICE_NAME_ENV#

‘OTEL_SERVICE_NAME’

nemo_rl.telemetry.setup._is_env_truthy(name: str) → bool#

Return True if env var name is set to a truthy value.

nemo_rl.telemetry.setup.telemetry_enabled_in_env() → bool#

Whether telemetry.enabled reached this process as truthy.

Exists for instrumentation that never asks for a handle: vLLM’s native tracing runs on vLLM’s own exporter, so :func:get_telemetry_handle would not gate it and the master switch would not reach it. Reads the environment rather than a config object because that is the only channel a worker process has.

nemo_rl.telemetry.setup.vllm_native_tracing_requested() → bool#

Whether telemetry.vllm_native_tracing reached this process as truthy.

Callers must also honour :func:telemetry_enabled_in_env; this reports only the one field.

nemo_rl.telemetry.setup._config_to_env(tel: Any) → None#

Translate a TelemetryConfig into NEMO_RL_OTEL_* env vars.

Uses os.environ.setdefault so raw env vars always win over YAML. Runs on the driver before init_ray(), so the resulting environment is snapshotted into the Ray runtime_env and inherited by every worker process.

nemo_rl.telemetry.setup._dig(obj: Any, *path: str) → Any#

Best-effort nested lookup that works for both dicts and objects.

Returns None as soon as any level is missing. Used to pull resource attributes out of a MasterConfig whose nested nodes may be pydantic models (attribute access) or TypedDict-derived dicts (key access).

nemo_rl.telemetry.setup._build_resource_attributes(
master_config: Any,
algorithm: str,
) → dict#

Build process-lifetime resource attributes (Jaeger “Process” tags).

Only stable-for-the-run values belong here (algorithm, model, precision, parallelism). Per-step values are span tags; time-series values are metrics. Best-effort: a missing key simply omits that attribute — never raises. dl.rank / dl.world_size are set by lens from the setup call, not here.

nemo_rl.telemetry.setup._always_export(
config: nemo.lens.NemoLensConfig,
rank: int,
world_size: int,
) → bool#

Export-strategy override for singleton processes: always export.

nemo_rl.telemetry.setup._unrank(config: nemo.lens.NemoLensConfig) → Any#

Disable rank-based span filtering for a process that has no real rank.

The driver and singleton actors such as AsyncTrajectoryCollector are not members of a distributed group, so they pass a synthetic rank=0 / world_size=1. Both of lens’s rank filters then misfire on that made-up rank, and the two are independent, so both have to be neutralised:

  • export_strategy decides whether this process exports at all. A strategy that selects among the ranks of a group has no meaning for a singleton, and mutes it outright for any export_rank >= 1.

  • sampler_enabled installs a RankAwareSampler on the tracer provider, which drops every span on ranks whose md5(rank) bucket lands above export_sample_rate. Rank 0’s bucket is 0.785, so any sample rate at or below that discards the process’s spans before the export decision is ever consulted.

The driver hosts the training loop and the metrics logger, and the collector generates every async rollout, so either filter silently drops the telemetry that matters most.

Mutating config here is process-local: it is rebuilt from the environment in each process, and the propagated NEMO_RL_OTEL_* vars are untouched, so ranked workers still honour what the user configured.

Returns the export-strategy override to hand to setup_telemetry.

nemo_rl.telemetry.setup.init_telemetry_driver(
master_config: Any,
algorithm: str,
) → Optional[nemo.lens.TelemetryHandle]#

Initialise driver-side telemetry (call once, before init_ray()).

Reads master_config.telemetry, exports the resolved settings as NEMO_RL_OTEL_* env vars (so workers inherit them), and sets up the driver’s OTel providers. The driver always exports (it hosts the training loop and the metrics logger).

Returns the :class:TelemetryHandle, or None if telemetry is disabled. Idempotent.

Raises:

ValueError – if telemetry.export_strategy names a strategy lens does not have. Deliberately fatal on the driver, where the user sees it.

nemo_rl.telemetry.setup._worker_resource_attributes(
extra: Optional[dict[str, Any]],
) → dict[str, Any]#

Build resource attributes identifying this worker process.

RANK is group-local — the policy group and the generation group each number their workers from zero — so dl.rank alone cannot tell their spans apart. rl.worker_group carries the group’s name_prefix (lm_policy, vllm_policy, …), which RayWorkerGroup exports as NRL_WORKER_GROUP. Explicit extra attributes win.

nemo_rl.telemetry.setup.init_telemetry_worker(
rank: Optional[int] = None,
world_size: Optional[int] = None,
resource_attributes: Optional[dict[str, Any]] = None,
always_export: bool = False,
) → Optional[nemo.lens.TelemetryHandle]#

Initialise telemetry inside a Ray actor (call once per worker process).

Reads the NEMO_RL_OTEL_* env propagated from the driver via the Ray runtime_env. rank / world_size default to the RANK / WORLD_SIZE env vars the worker was launched with, which — together with the export strategy — decide whether this worker exports.

Parameters:

  • rank – This process’s rank. Defaults to the RANK env var.

  • world_size – Size of this process’s group. Defaults to WORLD_SIZE.

  • resource_attributes – Extra resource attributes for this process.

  • always_export – Bypass the configured rank filters for this process, the way the driver does (see :func:_unrank). Set it for a singleton actor passing a synthetic rank / world_size: the filters select among the ranks of a distributed group, so applying them to a made-up rank has no meaning and silently mutes the actor — e.g. export_rank: 3 never matches a synthetic rank 0. Ranked members of a real worker group must leave this false.

Never raises: unlike the driver, a worker must not fail a training run over optional observability, and the driver has already validated the same config before any worker starts — so a genuine misconfiguration surfaces there, loudly, rather than here.

Returns the :class:TelemetryHandle, or None if telemetry is disabled or setup failed. Idempotent per process.

nemo_rl.telemetry.setup.get_telemetry_handle() → Optional[nemo.lens.TelemetryHandle]#

Return the process-global telemetry handle (None if uninitialised).

Named for the handle rather than the telemetry because callers reach through it – .tracer, .meter, .is_exporting – rather than using the return value as a value.

nemo_rl.telemetry.setup.shutdown_telemetry(timeout_ms: int = 5000) → None#

Flush and shut down telemetry providers.

Call on the driver at job end, and in each Ray actor’s shutdown: span and metric processors buffer in the background, so an actor that exits without flushing silently drops whatever it had not exported yet. A no-op when this process never initialised telemetry.