`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#

`_campaign_stage`	Lifecycle stage for algorithm; :data:`_DEFAULT_CAMPAIGN_STAGE` if unknown.
`_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).
`_rank_attributes`	Resource attributes identifying this process within its group.
`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).
`instrument_aiohttp_client`	Make outgoing `aiohttp` requests carry the caller’s trace context.
`traced_worker_init`	Decorate a worker `__init__` so its model load lands in a span.
`get_telemetry_handle`	Return the process-global telemetry handle (`None` if uninitialised).
`shutdown_telemetry`	Flush and shut down telemetry providers, returning within timeout_ms.

Data#

`logger`
`_F`
`_TELEMETRY_HANDLE`
`_TELEMETRY_INITIALISED`
`_OTEL_PREFIX`
`_OTEL_FALLBACK_PREFIX`
`_RUN_ID_ENV`
`_WORKER_GROUP_ENV`
`_CAMPAIGN_STAGE_ATTR`
`_DEFAULT_CAMPAIGN_STAGE`
`_CAMPAIGN_STAGE_BY_ALGORITHM`
`_CAMPAIGN_STAGE_ENV`
`_ENV_FIELD_MAP`
`_SERVICE_NAME_ENV`
`_WORKER_GROUP_ATTR`
`_AIOHTTP_INSTRUMENTED`

API#

nemo_rl.telemetry.setup.logger#: ‘getLogger(…)’

nemo_rl.telemetry.setup._F#: ‘TypeVar(…)’

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._CAMPAIGN_STAGE_ATTR#: ‘nv.dl.campaign.stage’

nemo_rl.telemetry.setup._DEFAULT_CAMPAIGN_STAGE#: ‘RL’

nemo_rl.telemetry.setup._CAMPAIGN_STAGE_BY_ALGORITHM#: None

nemo_rl.telemetry.setup._CAMPAIGN_STAGE_ENV#: None

nemo_rl.telemetry.setup._campaign_stage(algorithm: str) → str#: Lifecycle stage for algorithm; :data:_DEFAULT_CAMPAIGN_STAGE if unknown.

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. Rank identity comes from :func:_rank_attributes, which the callers merge in.

nemo_rl.telemetry.setup._rank_attributes(rank: int, world_size: int) → dict[str, Any]#

Resource attributes identifying this process within its group.

Lens has no notion of rank: it neither filters nor samples on one, so every process that sets up telemetry exports. Rank is recorded as a resource attribute instead, which is what makes a single rank’s spans selectable afterwards – filter on nv.dl.rank in the collector. There is no per-rank switch to reach for instead: telemetry.enabled is one job-wide value every worker inherits through the Ray runtime_env.

Passing rank at all is what keeps that filter available; lens warns when nv.dl.rank is missing, because without it a process cannot be told apart from its peers downstream.

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.

nemo_rl.telemetry.setup._WORKER_GROUP_ATTR#: ‘rl.worker_group’

nemo_rl.telemetry.setup._worker_resource_attributes( worker_group: Optional[str], 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 nv.dl.rank alone cannot tell their spans apart. :data:_WORKER_GROUP_ATTR carries the group’s name_prefix (lm_policy, vllm_policy, …), which RayWorkerGroup exports as NRL_WORKER_GROUP. An explicit worker_group wins over the environment, for the actors built outside RayWorkerGroup, which nothing sets it for. Explicit extra attributes win over both.

nemo_rl.telemetry.setup.init_telemetry_worker( rank: Optional[int] = None, world_size: Optional[int] = None, worker_group: Optional[str] = None, resource_attributes: Optional[dict[str, Any]] = None, ) → 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, and are recorded as nv.dl.rank / nv.dl.world_size resource attributes.

Every worker that gets here exports, and telemetry.enabled is job-wide, so it cannot keep some ranks quiet while others export. Narrowing down is a downstream decision: filter on nv.dl.rank in the collector.

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.

Never raises: a worker must not fail a training run over optional observability. The driver resolves the same config before any worker starts, so a misconfiguration is already reported once, from the one process whose log the user is reading — rather than once per worker.

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

nemo_rl.telemetry.setup._AIOHTTP_INSTRUMENTED#: False

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

Make outgoing aiohttp requests carry the caller’s trace context.

The NeMo-Gym rollout path leaves this process over HTTP from inside Gym’s own client, so there is no request NeMo-RL can add a header to. Patching the client library instead means every request Gym makes is stamped with a W3C traceparent taken from the ambient context, which is what lets a Gym server running :func:nemo.lens.contrib.fastapi.instrument_fastapi continue the trace rather than start one.

Call after worker telemetry is up and before any ClientSession exists – the instrumentor patches the class, so sessions built earlier keep the unpatched behaviour.

Returns:: Whether the client is instrumented. False when telemetry is off, or when the optional dependency is missing – neither is fatal, because losing the Gym half of a trace is not a reason to fail a rollout.

nemo_rl.telemetry.setup.traced_worker_init(

span_name: str,

**attributes: Any,

) → Callable[[nemo_rl.telemetry.setup._F], nemo_rl.telemetry.setup._F]#

Decorate a worker __init__ so its model load lands in a span.

Training workers build and shard the model inside __init__, which on a large checkpoint is the longest single phase of a run’s startup. Without this the time is real but unattributed: the driver’s rl.setup.workers span covers it as one block, and nothing says what the worker was doing.

Initialises worker telemetry before the body runs. A worker normally calls

Func:: init_telemetry_worker itself a few lines into __init__, but a span cannot open before its provider exists, so the decorator hoists that call ahead of the constructor; the body’s own call then returns the handle already made, being idempotent per process.

A None handle – telemetry off, or its setup failed – runs the constructor untouched. That is also what keeps this safe when nemo.lens is absent: the span imports below are reached only once a handle exists, which cannot happen without lens, so an install without it neither imports nor pays for any of this.

The span belongs to the model_init umbrella and so carries no rl.bucket: it is startup rather than a phase of a training step, and the goodput rollup measures steps.

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, returning within timeout_ms.

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.

The budget is enforced here rather than passed down, because nothing below honours it. TelemetryHandle.shutdown forwards timeout_ms to force_flush(timeout_millis=...), and the SDK’s BatchProcessor accepts that argument and then drains the whole queue regardless; the provider.shutdown() calls that follow take no timeout at all and fall back to the SDK’s own 30s. Four unbounded calls in total. Against an unreachable collector each buffered batch is retried with exponential backoff, so the cost scales with the queue: 3k spans took 36s under a 5s budget when measured.

That mattered because callers size their own timeouts on this one – _flush_collector_telemetry in async GRPO allows 15s for a 3s quiesce plus “its 5s export”, then ray.kills the collector when the budget blows, dropping the last rollout spans the flush exists to save.

So: run the flush on a daemon thread and stop waiting at the deadline. A flush that overruns is abandoned rather than joined, which loses the same spans an unreachable collector was going to lose anyway, and being a daemon it never holds up process exit.

nemo_rl.telemetry.setup#

Module Contents#

Functions#

Data#

API#

`nemo_rl.telemetry.setup`#