From b188b398c341ea40af31813c16cdff5f580472c7 Mon Sep 17 00:00:00 2001 From: Raj Singh Date: Wed, 7 Oct 2026 14:13:42 -0700 Subject: [PATCH 01/17] chore(deps): bump nemo-lens to the tagged v0.3.0 release The pin was b0f977d, an untagged commit four behind the v0.3.0 tag, so the pin now names a release instead of a point on main. Picks up shared attribute encoding and typed reconstruction, the training resource semconv fields, optional group gating on explicit span emission, and scoped span attributes. The >=0.2.0 floor in [dependencies] stays as is: package_info.py still declares 0.2.0 at the v0.3.0 tag, so raising it to >=0.3.0 would fail to resolve. Lens' pyproject.toml is byte-identical across the two revs, so the lockfile change is confined to the rev and version strings. Signed-off-by: Raj Singh --- pyproject.toml | 4 ++-- uv.lock | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 1e3f0fac33a..ef7001a250b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -338,7 +338,7 @@ nemo_run = { git = "https://github.com/NVIDIA-NeMo/Run", rev = "414f0077c648fde2 # actually reconciles the two revs; the floor in [dependencies] stays above the # newest PyPI release so dropping either one fails to resolve instead of # falling back to 0.1.0. -nemo-lens = { git = "https://github.com/NVIDIA-NeMo/Lens.git", rev = "b0f977d414b2f89938604a0b7eaa78ee08bc8700" } +nemo-lens = { git = "https://github.com/NVIDIA-NeMo/Lens.git", rev = "7bcc03ce3565d2be6858b819b049bcc4c3ec05b4" } sglang = { git = "https://github.com/sgl-project/sglang.git", rev = "3003d70f680d41c59d1b7acbf65cb47795dfd19e", subdirectory = "python" } # torch/torchvision/triton all come from the torch index in order to pick up aarch64 wheels torch = [ @@ -511,7 +511,7 @@ override-dependencies = [ # shipped `SpanGroup` class) and the consumer-driven metric registry # (`MetricSpec`/`register_metric_group`), both of which landed after v0.2.0. Drop # this once Megatron-LM bumps to the same or newer rev. - "nemo-lens[sdk,aiohttp] @ git+https://github.com/NVIDIA-NeMo/Lens.git@b0f977d414b2f89938604a0b7eaa78ee08bc8700", + "nemo-lens[sdk,aiohttp] @ git+https://github.com/NVIDIA-NeMo/Lens.git@7bcc03ce3565d2be6858b819b049bcc4c3ec05b4", ] # CVE fixes diff --git a/uv.lock b/uv.lock index 481be455796..3522fa1dc79 100644 --- a/uv.lock +++ b/uv.lock @@ -97,7 +97,7 @@ overrides = [ { name = "llguidance", specifier = ">=1.7.0,<1.8.0" }, { name = "mamba-ssm", git = "https://github.com/state-spaces/mamba.git?rev=a14b1dff0454a3bc27d9eb31355dc01e4b2490ec" }, { name = "mlflow", specifier = ">=3.13.0" }, - { name = "nemo-lens", extras = ["sdk", "aiohttp"], git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=b0f977d414b2f89938604a0b7eaa78ee08bc8700" }, + { name = "nemo-lens", extras = ["sdk", "aiohttp"], git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=7bcc03ce3565d2be6858b819b049bcc4c3ec05b4" }, { name = "numpy", specifier = ">=2.1.0,<2.5" }, { name = "nvidia-cublas", marker = "sys_platform != 'darwin'", specifier = "==13.5.1.27" }, { name = "nvidia-cudnn-cu13", marker = "sys_platform != 'darwin'", specifier = "==9.26.0.51" }, @@ -4133,8 +4133,8 @@ codec = [ [[package]] name = "nemo-lens" -version = "0.2.0+b0f977d" -source = { git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=b0f977d414b2f89938604a0b7eaa78ee08bc8700#b0f977d414b2f89938604a0b7eaa78ee08bc8700" } +version = "0.2.0+7bcc03c" +source = { git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=7bcc03ce3565d2be6858b819b049bcc4c3ec05b4#7bcc03ce3565d2be6858b819b049bcc4c3ec05b4" } dependencies = [ { name = "opentelemetry-api" }, ] @@ -4387,7 +4387,7 @@ requires-dist = [ { name = "nccl4py", marker = "sys_platform != 'darwin'" }, { name = "nemo-automodel", extras = ["moe"], marker = "extra == 'automodel'", editable = "3rdparty/Automodel-workspace/Automodel" }, { name = "nemo-gym", marker = "extra == 'nemo-gym'", editable = "3rdparty/Gym-workspace/Gym" }, - { name = "nemo-lens", extras = ["aiohttp", "sdk"], git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=b0f977d414b2f89938604a0b7eaa78ee08bc8700" }, + { name = "nemo-lens", extras = ["aiohttp", "sdk"], git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=7bcc03ce3565d2be6858b819b049bcc4c3ec05b4" }, { name = "ninja" }, { name = "nixl", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')", specifier = "==1.3.0" }, { name = "num2words", specifier = ">=0.5.14" }, From 8fea3826eee6c733d7d5e72d183cff72c6c1a455 Mon Sep 17 00:00:00 2001 From: Raj Singh Date: Wed, 7 Oct 2026 17:14:34 -0700 Subject: [PATCH 02/17] feat(telemetry)!: carry rl.iteration on every span in a step rl.step and rl.iteration were the same counter under two names, split by which span emitted it rather than by meaning: the training loops read total_steps and tagged rl.iteration, while validate() takes a parameter named step and tagged rl.step. Both resolve to the same value, so "everything in training step 412" needed two filters and knowledge of which spans used which. Collapse onto rl.iteration and set it once per step with iteration_scope, built on lens' span_attributes. Leaves, checkpoint spans and the rollout and generation spans under a step now all carry it without the counter being threaded to each call site -- which is why checkpoint spans had no attributes at all, and why a span added later would have missed it again. Spans left outside the scope are the ones with no single step: rl..job wraps the run, and rl.sc.generate_and_push runs on the rollout pump's clock and keeps rl.target_step. Worker-side model init is in another process, so the ContextVar does not reach it. Also tags the ppo and rm evaluate spans, which had step in scope but never attached it, and scopes sft_v2 at the loop body since its step span is a decorator and it checkpoints from the caller. BREAKING CHANGE: spans no longer carry rl.step. Dashboards and queries filtering on it should use rl.iteration, which holds the same value. Signed-off-by: Raj Singh --- docs/observability/extending.md | 6 +- docs/observability/span-groups.md | 27 ++++++- nemo_rl/algorithms/distillation.py | 8 +- nemo_rl/algorithms/dpo.py | 8 +- nemo_rl/algorithms/grpo.py | 11 +-- nemo_rl/algorithms/ppo.py | 9 ++- nemo_rl/algorithms/rm.py | 6 +- nemo_rl/algorithms/sft.py | 8 +- nemo_rl/algorithms/sft_v2.py | 29 +++++--- nemo_rl/algorithms/single_controller.py | 12 +-- nemo_rl/telemetry/instrumentation.py | 35 +++++++++ tests/unit/telemetry/test_instrumentation.py | 78 ++++++++++++++++++++ 12 files changed, 197 insertions(+), 40 deletions(-) diff --git a/docs/observability/extending.md b/docs/observability/extending.md index 6a42c115f2a..339150b8c07 100644 --- a/docs/observability/extending.md +++ b/docs/observability/extending.md @@ -161,7 +161,7 @@ For a hot path where you want minimal cost when the group is disabled: ```python with umbrella_span(RLSpanGroup.U_ROLLOUT, "rl.grpo.generation", - **{"rl.iteration": iteration}) as span: + **{"rl.num_prompt_groups": n_groups}) as span: result = collect() if span is not None: span.set_attribute("rl.num_generations_per_prompt", n) @@ -169,6 +169,10 @@ with umbrella_span(RLSpanGroup.U_ROLLOUT, "rl.grpo.generation", Both yield `None` when the group is disabled; the body still runs, so guard attribute-setting with `if span is not None`. +Do not pass `rl.iteration` here. The training loops set it once per step with +`iteration_scope`, so any span opened inside a step already carries it — see +[`rl.iteration` is set once per step](span-groups.md#rliteration-is-set-once-per-step). + ### Always-on block — `span_cm` `span_cm` always creates a span when telemetry is active (no group gate) — for cold, top-level paths only: diff --git a/docs/observability/span-groups.md b/docs/observability/span-groups.md index c9755d756d8..5ade636dbdc 100644 --- a/docs/observability/span-groups.md +++ b/docs/observability/span-groups.md @@ -120,9 +120,8 @@ These are set on spans for filtering — they answer "which one?" / "what kind?" | Tag | Meaning | |---|---| -| `rl.iteration` | training iteration index | +| `rl.iteration` | training step index, 1-based and counted across the whole run rather than reset per epoch. Set once per step, so every span inside that step carries it — see [below](#rliteration-is-set-once-per-step) | | `rl.epoch` | epoch index (omitted on `rl.sc.step` — see below) | -| `rl.step` | step index | | `rl.num_generations_per_prompt` | GRPO group size | | `rl.weight_version` / `rl.target_weight_version` | async rollout batch: the weights it generated from, and the training step it targets | | `rl.num_prompt_groups` | async rollout batch width, so a gap-filling batch is not read as an unexplained speed-up | @@ -144,6 +143,30 @@ the epoch counter at the time a train step runs describes how far *generation* has read into the dataset — not the epoch this step's batch came from. Use the rollout-side spans for that. +### `rl.iteration` is set once per step + +Each algorithm opens its step inside +`nemo_rl.telemetry.instrumentation.iteration_scope`, which stamps +`rl.iteration` on every span started within it. So the step umbrella, the leaves +under it, the checkpoint span and the rollout and generation spans all carry the +same value, and `rl.iteration = 412` selects everything that happened in that +step rather than the step span alone. + +A span that passes `rl.iteration` itself keeps its own value — the scope is +applied by a span processor that skips keys already set. `rl..evaluate` +does exactly that, because validation also runs once before training starts, +where it reports `rl.iteration = 0`: a sentinel for "no step has run yet". + +Two kinds of span are deliberately outside the scope. `rl..job` wraps the +whole run and belongs to no single step. `rl.sc.generate_and_push` runs on the +rollout pump's own clock rather than the training loop's, so it carries +`rl.target_step` — the step it is generating *for* — instead. + +The scope travels on a `ContextVar`, which reaches nested +calls and coroutines but not raw threads or other processes. Worker-side spans +such as `rl.policy.load_model` are therefore unstamped; they are model init, not +step work, so they sit outside any step regardless. + ### Span group → `rl.bucket` Leaf groups are tagged automatically when using diff --git a/nemo_rl/algorithms/distillation.py b/nemo_rl/algorithms/distillation.py index ff9a06e58fc..e3de12a4ab0 100644 --- a/nemo_rl/algorithms/distillation.py +++ b/nemo_rl/algorithms/distillation.py @@ -74,7 +74,9 @@ from nemo_rl.models.policy.lm_policy import Policy from nemo_rl.telemetry.config import TelemetryConfig from nemo_rl.telemetry.instrumentation import ( + RL_ITERATION_ATTR, evaluate_span, + iteration_scope, managed_span, umbrella_span, umbrella_trace_fn, @@ -792,12 +794,13 @@ def _distillation_train_impl( val_metrics, validation_timings = None, None with ( + iteration_scope(total_steps + 1), timer.time("total_step_time"), umbrella_span( RLSpanGroup.U_STEP, "rl.distillation.step", tracer=_tracer, - **{"rl.iteration": total_steps + 1, "rl.epoch": current_epoch + 1}, + **{"rl.epoch": current_epoch + 1}, ), ): # Prepare batch @@ -973,7 +976,6 @@ def _distillation_train_impl( RLSpanGroup.POLICY_UPDATE, "rl.distillation.policy_training", tracer=_tracer, - **{"rl.iteration": total_steps + 1}, ), ): train_results = student_policy.train( @@ -1294,7 +1296,7 @@ def validate( timer = Timer() with ( timer.time("total_validation_time"), - evaluate_span("distillation", **{"rl.step": step}), + evaluate_span("distillation", **{RL_ITERATION_ATTR: step}), ): print(f"▶ Starting validation at step {step}...", flush=True) diff --git a/nemo_rl/algorithms/dpo.py b/nemo_rl/algorithms/dpo.py index 2502cde0180..230158193d5 100644 --- a/nemo_rl/algorithms/dpo.py +++ b/nemo_rl/algorithms/dpo.py @@ -40,7 +40,9 @@ from nemo_rl.models.policy.lm_policy import Policy from nemo_rl.telemetry.config import TelemetryConfig from nemo_rl.telemetry.instrumentation import ( + RL_ITERATION_ATTR, evaluate_span, + iteration_scope, managed_span, umbrella_span, umbrella_trace_fn, @@ -468,7 +470,7 @@ def validate_one_dataset( timer = Timer() with ( timer.time("total_validation_time"), - evaluate_span("dpo", **{"rl.step": step}), + evaluate_span("dpo", **{RL_ITERATION_ATTR: step}), ): print(f"▶ Starting validation at step {step} for `{dataset_name}` set..") @@ -658,12 +660,13 @@ def dpo_train( val_metrics, validation_timings = None, None with ( + iteration_scope(total_steps + 1), timer.time("total_step_time"), umbrella_span( RLSpanGroup.U_STEP, "rl.dpo.step", tracer=_tracer, - **{"rl.iteration": total_steps + 1, "rl.epoch": current_epoch + 1}, + **{"rl.epoch": current_epoch + 1}, ), ): print("▶ Taking a training step...") @@ -673,7 +676,6 @@ def dpo_train( RLSpanGroup.POLICY_UPDATE, "rl.dpo.policy_training", tracer=_tracer, - **{"rl.iteration": total_steps + 1}, ), ): train_results = policy.train( diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index fd36d8aa4c1..ff54db6191d 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -154,9 +154,11 @@ from nemo_rl.models.policy.lm_policy import Policy from nemo_rl.telemetry.config import TelemetryConfig from nemo_rl.telemetry.instrumentation import ( + RL_ITERATION_ATTR, current_trace_carrier, efficiency_span, evaluate_span, + iteration_scope, managed_span, umbrella_span, umbrella_trace_fn, @@ -3301,12 +3303,13 @@ def _grpo_train_impl( val_metrics, validation_timings = None, None with ( + iteration_scope(total_steps + 1), timer.time("total_step_time"), umbrella_span( RLSpanGroup.U_STEP, "rl.grpo.step", tracer=_tracer, - **{"rl.iteration": total_steps + 1, "rl.epoch": current_epoch + 1}, + **{"rl.epoch": current_epoch + 1}, ), ): # Prepare batch @@ -3907,7 +3910,6 @@ def _grpo_train_impl( RLSpanGroup.POLICY_UPDATE, "rl.grpo.policy_training", tracer=_tracer, - **{"rl.iteration": total_steps + 1}, ), ): train_results = policy.train( @@ -4457,7 +4459,7 @@ def validate( timer = Timer(context={"worker": "validator"}) with ( timer.time("total_validation_time"), - evaluate_span("grpo", **{"rl.step": step}), + evaluate_span("grpo", **{RL_ITERATION_ATTR: step}), ): print(f"▶ Starting validation at step {step}...", flush=True) # >= 1 is validated in setup(). @@ -5301,12 +5303,12 @@ def _flush_collector_telemetry() -> None: maybe_gpu_profile_step(policy_generation, step + 1) with ( + iteration_scope(step + 1), timer.time("total_step_time"), umbrella_span( RLSpanGroup.U_STEP, "rl.grpo.step", tracer=_tracer, - **{"rl.iteration": step + 1}, ), ): num_mask_sample_filtered = 0 @@ -5733,7 +5735,6 @@ def _flush_collector_telemetry() -> None: RLSpanGroup.POLICY_UPDATE, "rl.grpo.policy_training", tracer=_tracer, - **{"rl.iteration": step + 1}, ), ): train_results = policy.train( diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index 4cd51bc0110..f474509138a 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -104,7 +104,9 @@ from nemo_rl.models.value.interfaces import ValueInterface from nemo_rl.telemetry.config import TelemetryConfig from nemo_rl.telemetry.instrumentation import ( + RL_ITERATION_ATTR, evaluate_span, + iteration_scope, managed_span, umbrella_span, umbrella_trace_fn, @@ -1429,12 +1431,13 @@ def ppo_train( val_metrics, validation_timings = None, None with ( + iteration_scope(total_steps + 1), timer.time("total_step_time"), umbrella_span( RLSpanGroup.U_STEP, "rl.ppo.step", tracer=_tracer, - **{"rl.iteration": total_steps + 1, "rl.epoch": current_epoch + 1}, + **{"rl.epoch": current_epoch + 1}, ), ): # Prepare batch @@ -1790,7 +1793,6 @@ def ppo_train( RLSpanGroup.POLICY_UPDATE, "rl.ppo.value_training", tracer=_tracer, - **{"rl.iteration": total_steps + 1}, ), ): value_results = value_model.train( @@ -1829,7 +1831,6 @@ def ppo_train( RLSpanGroup.POLICY_UPDATE, "rl.ppo.policy_training", tracer=_tracer, - **{"rl.iteration": total_steps + 1}, ), ): train_results = policy.train( @@ -3242,7 +3243,7 @@ def validate( timer = Timer() with ( timer.time("total_validation_time"), - evaluate_span("ppo"), + evaluate_span("ppo", **{RL_ITERATION_ATTR: step}), ): print(f"▶ Starting validation at step {step}...", flush=True) diff --git a/nemo_rl/algorithms/rm.py b/nemo_rl/algorithms/rm.py index 11fc863b66d..9fd11c4a94c 100644 --- a/nemo_rl/algorithms/rm.py +++ b/nemo_rl/algorithms/rm.py @@ -41,7 +41,9 @@ from nemo_rl.models.policy.lm_policy import Policy from nemo_rl.telemetry.config import TelemetryConfig from nemo_rl.telemetry.instrumentation import ( + RL_ITERATION_ATTR, evaluate_span, + iteration_scope, managed_span, umbrella_span, umbrella_trace_fn, @@ -375,7 +377,7 @@ def validate_one_dataset( timer = Timer() with ( timer.time("total_validation_time"), - evaluate_span("rm"), + evaluate_span("rm", **{RL_ITERATION_ATTR: step}), ): print(f"▶ Starting validation at step {step} for `{dataset_name}` set..") @@ -548,12 +550,12 @@ def rm_train( val_metrics, validation_timings = None, None with ( + iteration_scope(total_steps + 1), timer.time("total_step_time"), umbrella_span( RLSpanGroup.U_STEP, "rl.rm.step", tracer=_tracer, - **{"rl.iteration": total_steps + 1}, ), ): # Prepare batch and generate responses diff --git a/nemo_rl/algorithms/sft.py b/nemo_rl/algorithms/sft.py index 0f308a94d66..1527e60030f 100644 --- a/nemo_rl/algorithms/sft.py +++ b/nemo_rl/algorithms/sft.py @@ -50,7 +50,9 @@ from nemo_rl.models.policy.lm_policy import Policy from nemo_rl.telemetry.config import TelemetryConfig from nemo_rl.telemetry.instrumentation import ( + RL_ITERATION_ATTR, evaluate_span, + iteration_scope, managed_span, umbrella_span, umbrella_trace_fn, @@ -549,7 +551,7 @@ def validate( timer = Timer() with ( timer.time("total_validation_time"), - evaluate_span("sft", **{"rl.step": step}), + evaluate_span("sft", **{RL_ITERATION_ATTR: step}), ): print(f"▶ Starting validation at step {step}...") @@ -699,12 +701,13 @@ def sft_train( val_metrics, validation_timings = None, None with ( + iteration_scope(total_steps + 1), timer.time("total_step_time"), umbrella_span( RLSpanGroup.U_STEP, "rl.sft.step", tracer=_tracer, - **{"rl.iteration": total_steps + 1, "rl.epoch": current_epoch + 1}, + **{"rl.epoch": current_epoch + 1}, ), ): # Prepare batch and generate responses @@ -732,7 +735,6 @@ def sft_train( RLSpanGroup.POLICY_UPDATE, "rl.sft.policy_training", tracer=_tracer, - **{"rl.iteration": total_steps + 1}, ), ): train_results = policy.train( diff --git a/nemo_rl/algorithms/sft_v2.py b/nemo_rl/algorithms/sft_v2.py index f713511a91c..b355ea60af4 100644 --- a/nemo_rl/algorithms/sft_v2.py +++ b/nemo_rl/algorithms/sft_v2.py @@ -57,6 +57,7 @@ from nemo_rl.telemetry.config import TelemetryConfig from nemo_rl.telemetry.instrumentation import ( accepts_trace_context, + iteration_scope, managed_span, trace_context_kwargs, umbrella_span, @@ -222,18 +223,22 @@ def run(self) -> dict[str, Any]: with umbrella_span(RLSpanGroup.U_JOB, "rl.sft_v2.job", tracer=self._tracer): self._trainer.prepare_for_training() while self._save_state.total_steps < self._max_steps: - metrics = self._run_train_step() - self._logger.log_metrics(metrics, self._save_state.total_steps) - metric = self._checkpoint_metric(metrics) - self._timeout.mark_iteration() - save_by_timeout = self._timeout.check_save() - if self._should_save(save_by_timeout=save_by_timeout): - self._save_checkpoint(metric) - if save_by_timeout: - # check_save fires once and then latches, so continuing - # would train unsaved until the walltime kill. - print("Timeout has been reached, stopping training early") - break + # Read before _run_train_step increments it, so the whole + # pass is tagged with the step that is running rather than + # the one that just finished. + with iteration_scope(self._save_state.total_steps + 1): + metrics = self._run_train_step() + self._logger.log_metrics(metrics, self._save_state.total_steps) + metric = self._checkpoint_metric(metrics) + self._timeout.mark_iteration() + save_by_timeout = self._timeout.check_save() + if self._should_save(save_by_timeout=save_by_timeout): + self._save_checkpoint(metric) + if save_by_timeout: + # check_save fires once and then latches, so continuing + # would train unsaved until the walltime kill. + print("Timeout has been reached, stopping training early") + break return vars(self._save_state).copy() finally: for cleanup, name in ( diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 6a4fc883f7c..75fd5d069c7 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -196,6 +196,7 @@ RL_IDLE_POLLS_ATTR, efficiency_span, is_span_group_enabled, + iteration_scope, managed_span, per_prompt_scope, safe_set_span_attributes, @@ -2956,6 +2957,11 @@ async def _train_pump(self) -> None: step_finalizer_group_ids: set[str] = set() with ( + # Pinned here rather than read per span: _train_steps is + # incremented part-way through the step, so a later read would + # tag the checkpoint spans with a different value than the + # step they belong to. + iteration_scope(self._train_steps + 1), self._timer.time("total_step_time"), umbrella_span( RLSpanGroup.U_STEP, @@ -2964,10 +2970,7 @@ async def _train_pump(self) -> None: # No rl.epoch: the rollout pump advances the epoch on its own # clock, so its value here would describe whichever epoch that # pump had reached, not the one this step's data came from. - **{ - "rl.iteration": self._train_steps + 1, - "rl.weight_version": version_during_step, - }, + **{"rl.weight_version": version_during_step}, ), ): # One span per starvation episode, not per 5ms poll. @@ -3775,7 +3778,6 @@ async def _train_pump(self) -> None: RLSpanGroup.CHECKPOINT, "rl.sc.checkpointing", tracer=self._tracer, - **{"rl.step": self._train_steps}, ), ): await self._save_checkpoint( diff --git a/nemo_rl/telemetry/instrumentation.py b/nemo_rl/telemetry/instrumentation.py index 3589d07d156..f9b38098456 100644 --- a/nemo_rl/telemetry/instrumentation.py +++ b/nemo_rl/telemetry/instrumentation.py @@ -54,6 +54,9 @@ from nemo.lens import ( safe_set_span_attributes as _safe_set_span_attributes, ) +from nemo.lens import ( + span_attributes as _span_attributes, +) from nemo_rl.telemetry.span_groups import UMBRELLA_GROUP_VALUES, RLSpanGroup from nemo_rl.telemetry.vocabulary import ( @@ -73,6 +76,11 @@ # without parsing it back out of the span name. RL_EFFICIENCY_CATEGORY_ATTR = "rl.efficiency.category" +# Training step the span belongs to, 1-based and counted across the whole run +# rather than reset per epoch. Set once per step by :func:`iteration_scope`, so +# a leaf does not have to be threaded the counter to be filterable by step. +RL_ITERATION_ATTR = "rl.iteration" + # Retry count for a wait that one span covers rather than one span per poll. # Without it the coalesced span's duration is unreadable: the same ten seconds # could be two thousand clean 5ms polls or two hundred polls whose selection @@ -102,6 +110,7 @@ "accepts_trace_context", "TRACE_CARRIER_KWARG", "bucket_scope", + "iteration_scope", "per_prompt_scope", "in_per_prompt_scope", "efficiency_span", @@ -111,6 +120,7 @@ "evaluate_span", "RL_EFFICIENCY_CATEGORY_ATTR", "RL_IDLE_POLLS_ATTR", + "RL_ITERATION_ATTR", ] @@ -311,6 +321,31 @@ def bucket_scope(bucket: Bucket) -> Iterator[None]: _BUCKET_OVERRIDE.reset(token) +@contextmanager +def iteration_scope(iteration: int) -> Iterator[None]: + """Stamp :data:`RL_ITERATION_ATTR` on every span opened inside this block. + + Wrap a training step once and its whole subtree -- leaves, checkpointing, + rollout and generation spans -- becomes filterable by step. Threading the + counter into each call site instead leaves the answer to "what else was + happening in step 412" dependent on which sites remembered to pass it, and + silently excludes every span added afterwards. + + A span that passes the attribute itself keeps its own value: lens applies + the scope in a span processor that skips keys already set. ``evaluate_span`` + relies on that, since validation also runs once before training and reports + iteration 0 there. + + Only spans from the provider lens builds are stamped, which is every span + NeMo-RL opens. Propagates like any :class:`~contextvars.ContextVar`: to + nested calls and to coroutines started inside the block, but not to raw + threads or other processes -- a worker's ``rl.policy.load_model`` is a root + span in another process and is outside any step besides. + """ + with _span_attributes({RL_ITERATION_ATTR: iteration}): + yield + + # Marks a region as per-prompt work, for spans opened below the caller. # # Needed because cardinality is a property of the call site, not of the diff --git a/tests/unit/telemetry/test_instrumentation.py b/tests/unit/telemetry/test_instrumentation.py index e9fa7e3b3f8..ce77e98b260 100644 --- a/tests/unit/telemetry/test_instrumentation.py +++ b/tests/unit/telemetry/test_instrumentation.py @@ -25,6 +25,7 @@ RL_BUCKET_ATTR, RL_EFFICIENCY_CATEGORY_ATTR, RL_IDLE_POLLS_ATTR, + RL_ITERATION_ATTR, TRACE_CARRIER_KWARG, UMBRELLA_GROUPS, Bucket, @@ -37,12 +38,14 @@ efficiency_span, goodput_span_attributes, in_per_prompt_scope, + iteration_scope, managed_span, per_prompt_scope, remote_trace_context, start_efficiency_span, trace_context_kwargs, trace_fn, + umbrella_span, ) from nemo_rl.telemetry.span_groups import RLSpanGroup @@ -1131,3 +1134,78 @@ def test_step_nests_under_job(): step, job = spans["rl.grpo.step"], spans["rl.grpo.job"] assert step.parent is not None assert step.parent.span_id == job.context.span_id + + +@requires_lens +def test_iteration_scope_reaches_leaf_and_checkpoint_spans(): + """The point of the scope: one filter selects a whole step. + + Threading the counter per call site is what left the checkpoint spans and + most leaves untagged, so the assertion is deliberately about the spans that + never passed it themselves rather than about the step umbrella. + """ + handle, exporter = _setup("all") + with iteration_scope(412): + with umbrella_span(RLSpanGroup.U_STEP, "rl.grpo.step", tracer=handle.tracer): + with managed_span( + RLSpanGroup.POLICY_UPDATE, + "rl.grpo.policy_training", + tracer=handle.tracer, + ): + pass + with managed_span( + RLSpanGroup.CHECKPOINT, + "rl.grpo.checkpointing", + tracer=handle.tracer, + ): + pass + handle.shutdown() + + emitted = {s.name: s for s in exporter.get_finished_spans()} + assert set(emitted) == { + "rl.grpo.step", + "rl.grpo.policy_training", + "rl.grpo.checkpointing", + } + for span in emitted.values(): + assert span.attributes[RL_ITERATION_ATTR] == 412 + + +@requires_lens +def test_iteration_scope_does_not_override_an_explicit_value(): + """``rl..evaluate`` depends on this. + + Validation also runs once before training, where it reports iteration 0. If + the scope overwrote what a span passed, a baseline validation that happened + to be called from inside a step would be renumbered into that step. + """ + handle, exporter = _setup("all") + with iteration_scope(412): + with managed_span( + RLSpanGroup.CHECKPOINT, + "rl.grpo.checkpointing", + tracer=handle.tracer, + **{RL_ITERATION_ATTR: 0}, + ): + pass + handle.shutdown() + + (emitted,) = exporter.get_finished_spans() + assert emitted.attributes[RL_ITERATION_ATTR] == 0 + + +@requires_lens +def test_iteration_scope_does_not_leak_past_the_step(): + """``rl..job`` outlives every step, so it must stay unstamped.""" + handle, exporter = _setup("all") + with umbrella_span(RLSpanGroup.U_JOB, "rl.grpo.job", tracer=handle.tracer): + with iteration_scope(1): + pass + with managed_span( + RLSpanGroup.CHECKPOINT, "rl.grpo.checkpointing", tracer=handle.tracer + ): + pass + handle.shutdown() + + for span in exporter.get_finished_spans(): + assert RL_ITERATION_ATTR not in span.attributes From c16ab3496e6fce67fd4e6f41d8861f930f8ae895 Mon Sep 17 00:00:00 2001 From: Raj Singh Date: Thu, 8 Oct 2026 09:24:05 -0700 Subject: [PATCH 03/17] feat(telemetry): report the shared nv.dl.* resource attributes An RL run previously set two attributes from lens's shared distributed-learning namespace, nv.dl.rank and nv.dl.world_size, so a backend collecting NeMo-RL alongside Megatron or NeMo could not compare the two on topology, batch shape, step target or library versions -- the facts a reader reaches for first when one run is slower than another. Populate the rest of that surface from the config RL already resolves: provider name, run uuid, local rank, tp/pp/dp sizes, global and micro batch size, sequence length, optimizer, recompute granularity, the configured step target, and the torch/CUDA/NCCL/transformer-engine versions as actually loaded. Three values that were written as string literals now come from nemo.lens.semconv, so a typo fails as an ImportError rather than becoming silent data loss. Two of them change spelling in the process: dl.tensor_parallel.size and dl.pipeline_parallel.size become nv.dl.topology.size.tp and .pp. Both spell the same quantity, but only the topology one is in lens's DL_RESOURCE_TYPES registry, which is what normalize_resource_attributes types a resource attribute against -- the dl.* family is the span and metric spelling. A dashboard filtering on the old keys has to be repointed. Metric names are deliberately left in the rl.* scope. Lens registers a dl.* series for whichever trainer owns the step, so publishing dl.grad_norm from an RL job would write into a co-resident Megatron trainer's series; rl.grad_norm and rl.learning_rate are spelling matches for exactly that reason. The dp size is derived rather than read, since RL configures none, and is omitted when generation runs on dedicated GPUs -- the cluster then holds ranks the policy does not have. Signed-off-by: Raj Singh --- docs/observability/configuration.md | 19 +- docs/observability/extending.md | 8 +- docs/observability/metrics.md | 2 +- docs/observability/span-groups.md | 2 +- nemo_rl/algorithms/single_controller.py | 3 +- nemo_rl/telemetry/instrumentation.py | 6 + nemo_rl/telemetry/setup.py | 229 ++++++++++++++++++++++-- tests/unit/telemetry/test_setup.py | 156 ++++++++++++++-- 8 files changed, 393 insertions(+), 32 deletions(-) diff --git a/docs/observability/configuration.md b/docs/observability/configuration.md index dad3e76c3b4..012c846f8d9 100644 --- a/docs/observability/configuration.md +++ b/docs/observability/configuration.md @@ -90,12 +90,25 @@ Each process sets stable-for-the-run values on its OTel `Resource` — `init_tel | `rl.algorithm` | the `algorithm=""` passed to `init_telemetry_driver` | | `rl.model` | `policy.model_name` | | `nemo.precision` | `policy.precision` | -| `dl.tensor_parallel.size` | `policy.megatron_cfg` / `automodel_cfg` TP size | -| `dl.pipeline_parallel.size` | `policy.megatron_cfg` PP size | +| `nv.dl.provider.name` | always the literal `nemo-rl`: which training framework produced the telemetry, for a backend collecting several. Separate from the service name, which a user may rename | +| `nv.dl.run.uuid` | the same run id as `nemo.run.id`. Duplicated on purpose — a backend joining RL runs to Megatron ones groups on the `nv.dl.*` spelling | +| `nv.dl.topology.size.tp` | `policy.megatron_cfg` / `automodel_cfg` TP size | +| `nv.dl.topology.size.pp` | `policy.megatron_cfg` PP size | +| `nv.dl.topology.size.dp` | derived: `cluster.num_nodes × cluster.gpus_per_node ÷ (TP × PP × CP)`. Omitted when `policy.generation.colocated.enabled` is false, because generation then holds ranks the policy does not have, and when the product does not divide the world size | +| `nv.dl.training.config.global_batch_size` | `policy.train_global_batch_size` | +| `nv.dl.training.config.micro_batch_size` | `policy.train_micro_batch_size` | +| `nv.dl.training.config.sequence_length` | `policy.max_total_sequence_length` | +| `nv.dl.training.config.optimizer` | `policy.megatron_cfg.optimizer.optimizer` (`adam`), else `policy.optimizer.name` (`torch.optim.AdamW`) — the two backends spell it differently | +| `nv.dl.training.config.recompute_granularity` | `policy.megatron_cfg.recompute_granularity`, only when `activation_checkpointing` is on. Megatron ignores the granularity otherwise, so reporting it would describe recompute that is not happening | +| `nv.dl.training.target.train_iters` | `.max_num_steps`, when positive. `rm` spells "train for one epoch" as `-1`, which is omitted rather than reported as a target | +| `nv.dl.software.torch`, `.cuda`, `.nccl`, `.transformer_engine` | versions as loaded in this process, not as pinned in the lockfile. Each is omitted when unavailable — CUDA and NCCL on a CPU-only build, transformer-engine when it is not installed | | `nv.dl.rank`, `nv.dl.world_size` | this process's rank and group size (`RANK` / `WORLD_SIZE`, or `0` / `1` for the driver and singleton actors) | +| `nv.dl.local_rank` | worker processes only: `LOCAL_RANK`, which `RayWorkerGroup` exports beside `RANK`. Identifies a device on one node, which is what you have when the complaint came from `nvidia-smi` rather than from the job | | `rl.worker_group` | worker processes only: the worker group's `name_prefix` (`lm_policy`, `vllm_policy`, ...), from `NRL_WORKER_GROUP`. The `NemoGym` actor reports the literal `nemo_gym` — it is built from the environment registry rather than by `RayWorkerGroup`, so nothing sets the env var for it | -Attribute construction is best-effort: a missing config key simply omits that attribute; it never raises. Plus auto-detected host / GPU / SLURM / Kubernetes attributes from lens's resource detection. +The config-derived rows are driver-only: a worker is handed the propagated environment, not the `MasterConfig`. Attribute construction is best-effort: a missing config key simply omits that attribute; it never raises. Plus auto-detected host / GPU / SLURM / Kubernetes attributes from lens's resource detection. + +Every key above that lens declares is taken from `nemo.lens.semconv` rather than written as a string, so a typo is an `ImportError` instead of silent data loss. Note the two spellings of the parallelism sizes: `nv.dl.topology.size.*` is the resource-scoped one lens types through its `DL_RESOURCE_TYPES` registry, while the `dl.tensor_parallel.*` family is for span and metric attributes. ## Typical configurations diff --git a/docs/observability/extending.md b/docs/observability/extending.md index 339150b8c07..5e82a2a2b98 100644 --- a/docs/observability/extending.md +++ b/docs/observability/extending.md @@ -190,10 +190,14 @@ if telemetry is not None: |---|---|---| | Span name | `rl..`, matching the block's `Timer` key (the two umbrella spans excepted — see [span groups](span-groups.md#per-algorithm-span-names)) | `rl.grpo.generation` | | Span tag | `rl.` categorical | `rl.iteration`, `rl.backend` | -| Resource attribute | `rl.` / shared `dl.` | `rl.model`, `dl.tensor_parallel.size` | +| Resource attribute | `rl.` for RL's own, shared `nv.dl.` for anything a distributed training job has | `rl.model`, `nv.dl.topology.size.tp` | | Metric name | `rl..` (application scope) | `rl.efficiency.seconds` | -Metric names use the **application scope** (`rl.*`) — never `dl.*`. NeMo-RL names its own attributes, in `nemo_rl/telemetry/`, rather than importing them from `nemo.lens.semconv`: they describe RL's properties, so RL is where a rename has to happen. The lens constants are for the few attributes lens itself sets, such as `nv.dl.rank`. +**Metric names use the application scope (`rl.*`) — never `dl.*`.** This is not a style preference: lens registers a `dl.*` series for the trainer that owns the step, so an RL job publishing `dl.grad_norm` would be writing into a co-resident Megatron trainer's series. `rl.grad_norm` and `rl.learning_rate` are deliberate spelling matches for exactly that reason, so query the two by their own names rather than by swapping the prefix. Lens's own semconv file says the same: RL's metric series are declared through `register_metric_group` in this tree and are deliberately not defined there. + +**Attributes split by who owns the name.** If lens declares it, import the constant from `nemo.lens.semconv` — a typo then fails as an `ImportError` instead of becoming silent data loss. That covers the shared `nv.dl.*` resource surface and the handful of `rl.*` span attributes lens itself declares (`rl.algorithm`, `rl.reward`, `rl.generation.backend`, `rl.num_rollouts`; `RL_ALGORITHM` is re-exported from `nemo_rl/telemetry/instrumentation.py` so call sites take every `rl.*` key from one module). Everything else — `rl.iteration`, `rl.bucket`, `rl.worker_group` — describes RL's own properties, so it is named in `nemo_rl/telemetry/` where a rename can land. + +Mind the scope when picking a shared name: `nv.dl.topology.size.tp` is the resource-scoped parallelism size, typed by lens's `DL_RESOURCE_TYPES` registry, while `dl.tensor_parallel.size` is the span and metric spelling of the same quantity. ## Choosing a span group diff --git a/docs/observability/metrics.md b/docs/observability/metrics.md index aaff030f8e7..5fd98556b63 100644 --- a/docs/observability/metrics.md +++ b/docs/observability/metrics.md @@ -211,7 +211,7 @@ The one rule that trips people up. Classify each value before you emit it: |---|---|---| | **Metric** | numerical value that changes over time | per-category efficiency seconds → `rl.efficiency.seconds` | | **Span tag** | categorical per-span context for filtering | `rl.iteration`, `rl.bucket`, `rl.num_generations_per_prompt`, `rl.weight_version` | -| **Resource attribute** | stable for the whole run | `rl.algorithm`, `rl.model`, `dl.tensor_parallel.size` | +| **Resource attribute** | stable for the whole run | `rl.algorithm`, `rl.model`, `nv.dl.topology.size.tp` | Do **not** put a time-series number (loss, reward) on a span attribute — it produces no useful series in your backend and wastes storage. Do **not** put a per-step categorical (iteration number) on a metric label — that is unbounded cardinality. See [lens: metrics — metric vs span attribute vs resource attribute](https://github.com/NVIDIA-NeMo/Lens). diff --git a/docs/observability/span-groups.md b/docs/observability/span-groups.md index 5ade636dbdc..bb703efdd45 100644 --- a/docs/observability/span-groups.md +++ b/docs/observability/span-groups.md @@ -599,7 +599,7 @@ blanks today, so an empty trace is not read as a broken exporter: ## Resource attributes (process tags) -Stable-for-the-run values, set once at init and attached to every span/metric: `nv.dl.campaign.stage` (always `RL`), `rl.algorithm`, `rl.model`, `nemo.precision`, `dl.tensor_parallel.size`, `dl.pipeline_parallel.size`, plus `nv.dl.rank` / `nv.dl.world_size`. See [Configuration — Resource attributes](configuration.md#resource-attributes). +Stable-for-the-run values, set once at init and attached to every span/metric: `nv.dl.campaign.stage` (always `RL`), `rl.algorithm`, `rl.model`, `nemo.precision`, `nv.dl.provider.name`, the `nv.dl.topology.size.*` parallelism sizes, the `nv.dl.training.config.*` / `nv.dl.training.target.*` settings, the `nv.dl.software.*` versions, plus `nv.dl.rank` / `nv.dl.world_size` / `nv.dl.local_rank`. See [Configuration — Resource attributes](configuration.md#resource-attributes) for the full table and where each value comes from. ## Granularity guidance diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 75fd5d069c7..1950ef8de55 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -193,6 +193,7 @@ from nemo_rl.models.value.tq_value import TQValue from nemo_rl.telemetry.instrumentation import ( NO_SPAN, + RL_ALGORITHM, RL_IDLE_POLLS_ATTR, efficiency_span, is_span_group_enabled, @@ -784,7 +785,7 @@ async def run(self) -> dict[str, Any]: RLSpanGroup.U_JOB, "rl.sc.job", tracer=self._tracer, - **{"rl.algorithm": "ppo" if self._is_ppo else "grpo"}, + **{RL_ALGORITHM: "ppo" if self._is_ppo else "grpo"}, ): result = await self._run_pumps() finally: diff --git a/nemo_rl/telemetry/instrumentation.py b/nemo_rl/telemetry/instrumentation.py index f9b38098456..7c5e4e6f8d8 100644 --- a/nemo_rl/telemetry/instrumentation.py +++ b/nemo_rl/telemetry/instrumentation.py @@ -58,6 +58,11 @@ span_attributes as _span_attributes, ) +# Span-attribute key lens declares itself. Re-exported so call sites take every +# ``rl.*`` key from this one module, and so a rename lands in lens's semconv +# instead of in a string literal scattered across the algorithms. +from nemo.lens.semconv import RL_ALGORITHM + from nemo_rl.telemetry.span_groups import UMBRELLA_GROUP_VALUES, RLSpanGroup from nemo_rl.telemetry.vocabulary import ( INIT_TOTAL_CATEGORY, @@ -118,6 +123,7 @@ "startup_span", "setup_span", "evaluate_span", + "RL_ALGORITHM", "RL_EFFICIENCY_CATEGORY_ATTR", "RL_IDLE_POLLS_ATTR", "RL_ITERATION_ATTR", diff --git a/nemo_rl/telemetry/setup.py b/nemo_rl/telemetry/setup.py index e8bea9c8cf0..a8b8c446beb 100644 --- a/nemo_rl/telemetry/setup.py +++ b/nemo_rl/telemetry/setup.py @@ -67,6 +67,11 @@ # Set per worker by ``RayWorkerGroup`` from the group's ``name_prefix``. _WORKER_GROUP_ENV = "NRL_WORKER_GROUP" +# Reported as ``nv.dl.provider.name``: which training framework produced this +# telemetry, for a backend collecting several of them. Spelled like the default +# service name, but separate from it -- a user is free to rename the service. +_PROVIDER_NAME = "nemo-rl" + # Which stage of a model's lifecycle (pretrain -> SFT -> RL) produced this # telemetry. A backend collecting several NeMo products selects a stage on this # rather than on service names, which differ per launcher. Seeded into both the @@ -186,13 +191,20 @@ def _build_resource_attributes( """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. + parallelism, batch and optimizer settings). Per-step values are span tags; + time-series values are metrics. + + Driver-only, all of it: a worker is handed the propagated environment rather + than the ``MasterConfig``, so nothing here is reachable from one. Best-effort: a missing key simply omits that attribute — never raises. - Rank identity comes from :func:`_rank_attributes`, which the callers merge in. + Process identity comes from :func:`_process_attributes`, which the callers + merge in. """ + from nemo.lens.semconv import RL_ALGORITHM + attrs: dict[str, Any] = { _CAMPAIGN_STAGE_ATTR: _campaign_stage(algorithm), - "rl.algorithm": algorithm, + RL_ALGORITHM: algorithm, } model = _dig(master_config, "policy", "model_name") @@ -203,21 +215,179 @@ def _build_resource_attributes( if precision: attrs["nemo.precision"] = precision - # Parallelism lives under the active policy backend (megatron vs dtensor). + attrs.update(_topology_attributes(master_config)) + attrs.update(_training_config_attributes(master_config, algorithm)) + return attrs + + +def _topology_attributes(master_config: Any) -> dict[str, Any]: + """Parallelism sizes, under lens's resource-scoped topology keys. + + ``nv.dl.topology.size.*`` rather than the ``dl.tensor_parallel.size`` family: + only the former is in lens's ``DL_RESOURCE_TYPES``, which is the registry + ``normalize_resource_attributes`` types a resource attribute against. The + ``dl.*`` spelling is the span and metric one. + + Each size lives under whichever policy backend is active, so each is looked + up under both megatron and automodel. Only the current spelling of the + latter: ``check_outdated_config`` rejects a config still saying + ``dtensor_cfg`` before a run reaches here. + """ + from nemo.lens.semconv import ( + NV_DL_TOPOLOGY_SIZE_DP, + NV_DL_TOPOLOGY_SIZE_PP, + NV_DL_TOPOLOGY_SIZE_TP, + ) + tp = _dig( master_config, "policy", "megatron_cfg", "tensor_model_parallel_size" ) or _dig(master_config, "policy", "automodel_cfg", "tensor_parallel_size") - if tp: - attrs["dl.tensor_parallel.size"] = tp pp = _dig(master_config, "policy", "megatron_cfg", "pipeline_model_parallel_size") + cp = _dig(master_config, "policy", "megatron_cfg", "context_parallel_size") or _dig( + master_config, "policy", "automodel_cfg", "context_parallel_size" + ) + + attrs: dict[str, Any] = {} + if tp: + attrs[NV_DL_TOPOLOGY_SIZE_TP] = tp if pp: - attrs["dl.pipeline_parallel.size"] = pp + attrs[NV_DL_TOPOLOGY_SIZE_PP] = pp + dp = _data_parallel_size(master_config, tp=tp or 1, pp=pp or 1, cp=cp or 1) + if dp: + attrs[NV_DL_TOPOLOGY_SIZE_DP] = dp + return attrs + + +def _data_parallel_size( + master_config: Any, *, tp: int, pp: int, cp: int +) -> Optional[int]: + """The ranks left for data parallelism once tp, pp and cp have claimed theirs. + + Derived rather than read: RL configures no dp size. Training spans the + cluster's GPUs and dp is whatever the other dimensions leave over. + + That identity breaks when generation runs on dedicated GPUs, because the + cluster then holds ranks the policy does not have. Returns ``None`` there, + rather than reporting a dp size for a group that is smaller than the + arithmetic says. Also ``None`` when the product does not divide the world + size, which the policy worker rejects with a far better message than a + resource attribute could carry. + """ + if _dig(master_config, "policy", "generation", "colocated", "enabled") is False: + return None + nodes = _dig(master_config, "cluster", "num_nodes") + gpus_per_node = _dig(master_config, "cluster", "gpus_per_node") + if not nodes or not gpus_per_node: + return None + + world_size = nodes * gpus_per_node + model_parallel_size = tp * pp * cp + if world_size % model_parallel_size: + return None + return world_size // model_parallel_size + + +def _training_config_attributes(master_config: Any, algorithm: str) -> dict[str, Any]: + """Batch, sequence and optimizer settings, plus the configured step target. + + These are the knobs a reader compares two runs on, and lens types them so + that a backend can select on them numerically. All best-effort. + """ + from nemo.lens.semconv import ( + NV_DL_TRAINING_CONFIG_GLOBAL_BATCH_SIZE, + NV_DL_TRAINING_CONFIG_MICRO_BATCH_SIZE, + NV_DL_TRAINING_CONFIG_OPTIMIZER, + NV_DL_TRAINING_CONFIG_RECOMPUTE_GRANULARITY, + NV_DL_TRAINING_CONFIG_SEQUENCE_LENGTH, + NV_DL_TRAINING_TARGET_TRAIN_ITERS, + ) + + attrs: dict[str, Any] = {} + for key, path in ( + (NV_DL_TRAINING_CONFIG_GLOBAL_BATCH_SIZE, ("train_global_batch_size",)), + (NV_DL_TRAINING_CONFIG_MICRO_BATCH_SIZE, ("train_micro_batch_size",)), + (NV_DL_TRAINING_CONFIG_SEQUENCE_LENGTH, ("max_total_sequence_length",)), + ): + value = _dig(master_config, "policy", *path) + if value: + attrs[key] = value + + # Two spellings, one per backend: megatron names the algorithm ("adam"), + # dtensor names the class ("torch.optim.AdamW"). + optimizer = _dig( + master_config, "policy", "megatron_cfg", "optimizer", "optimizer" + ) or _dig(master_config, "policy", "optimizer", "name") + if optimizer: + attrs[NV_DL_TRAINING_CONFIG_OPTIMIZER] = optimizer + + # Gated on activation checkpointing because megatron ignores the + # granularity without it, so reporting it would describe recompute that is + # not happening. + if _dig(master_config, "policy", "megatron_cfg", "activation_checkpointing"): + granularity = _dig( + master_config, "policy", "megatron_cfg", "recompute_granularity" + ) + if granularity: + attrs[NV_DL_TRAINING_CONFIG_RECOMPUTE_GRANULARITY] = granularity + + # The algorithm's own config section is named after it (``grpo:``, ``sft:``). + # Non-positive values are dropped rather than reported: ``rm`` spells "train + # for one epoch" as ``max_num_steps: -1``, which downstream would otherwise + # read as a target of -1 iterations. + max_num_steps = _dig(master_config, algorithm, "max_num_steps") + if isinstance(max_num_steps, int) and max_num_steps > 0: + attrs[NV_DL_TRAINING_TARGET_TRAIN_ITERS] = max_num_steps + return attrs + + +def _software_versions() -> dict[str, Any]: + """Library versions as actually loaded in this process. + + The lockfile says what was meant to be installed; these say what the run + imported, which is the pair worth comparing when a run regresses after an + image rebuild. Every lookup is optional: CUDA and NCCL are absent from a + CPU-only build, and transformer-engine is not installed for every backend. + """ + from nemo.lens.semconv import ( + NV_DL_SOFTWARE_CUDA, + NV_DL_SOFTWARE_NCCL, + NV_DL_SOFTWARE_TORCH, + NV_DL_SOFTWARE_TRANSFORMER_ENGINE, + ) + + attrs: dict[str, Any] = {} + try: + import torch + except ImportError: + return attrs + + attrs[NV_DL_SOFTWARE_TORCH] = str(torch.__version__) + if torch.version.cuda: + attrs[NV_DL_SOFTWARE_CUDA] = str(torch.version.cuda) + # Deliberately broad: the NCCL version is read out of the loaded extension, + # which raises differently depending on how torch was built, and no spelling + # of that failure is worth risking the driver's setup over. + try: + nccl_version = torch.cuda.nccl.version() + except Exception: + pass + else: + attrs[NV_DL_SOFTWARE_NCCL] = ".".join(str(part) for part in nccl_version) + + try: + import transformer_engine + except ImportError: + pass + else: + te_version = getattr(transformer_engine, "__version__", "") + if te_version: + attrs[NV_DL_SOFTWARE_TRANSFORMER_ENGINE] = str(te_version) return attrs -def _rank_attributes(rank: int, world_size: int) -> dict[str, Any]: - """Resource attributes identifying this process within its group. +def _process_attributes(rank: int, world_size: int, run_id: str = "") -> dict[str, Any]: + """Resource attributes identifying this process and the stack it loaded. 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 @@ -228,11 +398,38 @@ def _rank_attributes(rank: int, world_size: int) -> dict[str, Any]: 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. + apart from its peers downstream. ``nv.dl.local_rank`` narrows that to a + device on one node, which is what a reader has when the complaint came from + ``nvidia-smi`` rather than from the job. + + Called by both entry points, since a resource is per process and neither + path sees the other's. """ - from nemo.lens.semconv import NV_DL_RANK, NV_DL_WORLD_SIZE + from nemo.lens.semconv import ( + NV_DL_LOCAL_RANK, + NV_DL_PROVIDER_NAME, + NV_DL_RANK, + NV_DL_RUN_UUID, + NV_DL_WORLD_SIZE, + ) - return {NV_DL_RANK: rank, NV_DL_WORLD_SIZE: world_size} + attrs: dict[str, Any] = { + NV_DL_RANK: rank, + NV_DL_WORLD_SIZE: world_size, + NV_DL_PROVIDER_NAME: _PROVIDER_NAME, + } + # Exported next to RANK by ``RayWorkerGroup``; absent on the driver, which + # owns no device. + local_rank = os.environ.get("LOCAL_RANK", "").strip() + if local_rank.isdigit(): + attrs[NV_DL_LOCAL_RANK] = int(local_rank) + # Duplicates lens's own ``nemo.run.id`` by design: the two namespaces are + # read by different consumers, and a backend joining RL runs to Megatron + # ones groups on the ``nv.dl.*`` spelling. + if run_id: + attrs[NV_DL_RUN_UUID] = run_id + attrs.update(_software_versions()) + return attrs def init_telemetry_driver( @@ -296,7 +493,9 @@ def init_telemetry_driver( # The driver is a singleton, not a member of a distributed group. Rank 0 of # 1 is the honest description, and stating it silences lens's warning about # a process that cannot be identified by rank downstream. - resource_attrs.update(_rank_attributes(rank=0, world_size=1)) + resource_attrs.update( + _process_attributes(rank=0, world_size=1, run_id=config.run_id) + ) handle = setup_telemetry(config, resource_attributes=resource_attrs) # Only now, past everything that can raise: setting the guard earlier would @@ -425,7 +624,9 @@ def init_telemetry_worker( return None attrs = _worker_resource_attributes(worker_group, resource_attributes) - attrs.update(_rank_attributes(rank=rank, world_size=world_size)) + attrs.update( + _process_attributes(rank=rank, world_size=world_size, run_id=config.run_id) + ) handle = setup_telemetry(config, resource_attributes=attrs) logger.info( diff --git a/tests/unit/telemetry/test_setup.py b/tests/unit/telemetry/test_setup.py index d1aa6d83f80..40aa32cd132 100644 --- a/tests/unit/telemetry/test_setup.py +++ b/tests/unit/telemetry/test_setup.py @@ -16,7 +16,10 @@ from nemo_rl.telemetry.config import TelemetryConfig from nemo_rl.telemetry.setup import ( _build_resource_attributes, + _data_parallel_size, _dig, + _process_attributes, + _training_config_attributes, _worker_resource_attributes, get_telemetry_handle, init_telemetry_driver, @@ -58,8 +61,8 @@ def test_build_resource_attributes(): assert attrs["rl.algorithm"] == "grpo" assert attrs["rl.model"] == "org/Model-1B" assert attrs["nemo.precision"] == "bfloat16" - assert attrs["dl.tensor_parallel.size"] == 2 - assert attrs["dl.pipeline_parallel.size"] == 1 + assert attrs["nv.dl.topology.size.tp"] == 2 + assert attrs["nv.dl.topology.size.pp"] == 1 def test_campaign_stage_is_tagged_on_driver_and_worker(): @@ -86,8 +89,135 @@ def test_build_resource_attributes_automodel_tp(): } ) attrs = _build_resource_attributes(cfg, "grpo") - assert attrs["dl.tensor_parallel.size"] == 4 + assert attrs["nv.dl.topology.size.tp"] == 4 + assert "nv.dl.topology.size.pp" not in attrs + + +def test_parallelism_sizes_use_the_resource_scoped_spelling(): + """``nv.dl.topology.size.*``, not the ``dl.tensor_parallel.*`` family. + + Both spell the same quantity, but only the topology one is in lens's + ``DL_RESOURCE_TYPES`` -- the registry ``normalize_resource_attributes`` + types a resource attribute against. The ``dl.*`` spelling is the span and + metric one, so setting it here puts the value outside that typing. + """ + pytest.importorskip("nemo.lens") + from nemo.lens.semconv.resources import DL_RESOURCE_TYPES + + attrs = _build_resource_attributes(_FakeMasterConfig(), "grpo") + assert "dl.tensor_parallel.size" not in attrs assert "dl.pipeline_parallel.size" not in attrs + assert "nv.dl.topology.size.tp" in DL_RESOURCE_TYPES + assert "dl.tensor_parallel.size" not in DL_RESOURCE_TYPES + + +def test_data_parallel_size_is_derived_from_the_cluster(): + # RL configures no dp size: training spans the cluster's GPUs and dp is + # whatever tp/pp/cp leave over. + cfg = {"cluster": {"num_nodes": 2, "gpus_per_node": 8}} + assert _data_parallel_size(cfg, tp=2, pp=1, cp=1) == 8 + assert _data_parallel_size(cfg, tp=2, pp=2, cp=2) == 2 + + +def test_data_parallel_size_is_omitted_for_dedicated_generation_gpus(): + """The cluster then holds ranks the policy does not have. + + Reporting the cluster's arithmetic would describe a training group larger + than the one that exists, so this reports nothing instead. + """ + cfg = { + "cluster": {"num_nodes": 2, "gpus_per_node": 8}, + "policy": {"generation": {"colocated": {"enabled": False}}}, + } + assert _data_parallel_size(cfg, tp=1, pp=1, cp=1) is None + + +def test_data_parallel_size_is_omitted_when_it_does_not_divide_evenly(): + # A misconfiguration the policy worker rejects with a far better message + # than a resource attribute could carry. + cfg = {"cluster": {"num_nodes": 1, "gpus_per_node": 8}} + assert _data_parallel_size(cfg, tp=3, pp=1, cp=1) is None + + +def test_data_parallel_size_is_omitted_without_a_cluster_block(): + assert _data_parallel_size({}, tp=1, pp=1, cp=1) is None + + +def test_training_target_drops_a_non_positive_step_count(): + """``rm`` spells "train for one epoch" as ``max_num_steps: -1``. + + Emitting that verbatim would read downstream as a configured target of -1 + iterations, which is worse than having no target attribute at all. + """ + attrs = _training_config_attributes({"rm": {"max_num_steps": -1}}, "rm") + assert "nv.dl.training.target.train_iters" not in attrs + + attrs = _training_config_attributes({"rm": {"max_num_steps": 150}}, "rm") + assert attrs["nv.dl.training.target.train_iters"] == 150 + + +def test_training_config_reads_either_backends_optimizer_spelling(): + # megatron names the algorithm, dtensor names the class. + megatron = {"policy": {"megatron_cfg": {"optimizer": {"optimizer": "adam"}}}} + dtensor = {"policy": {"optimizer": {"name": "torch.optim.AdamW"}}} + key = "nv.dl.training.config.optimizer" + assert _training_config_attributes(megatron, "grpo")[key] == "adam" + assert _training_config_attributes(dtensor, "grpo")[key] == "torch.optim.AdamW" + + +def test_recompute_granularity_is_reported_only_when_recompute_is_on(): + # Megatron ignores the granularity without activation checkpointing, so + # reporting it would describe recompute that is not happening. + key = "nv.dl.training.config.recompute_granularity" + off = { + "policy": { + "megatron_cfg": { + "activation_checkpointing": False, + "recompute_granularity": "full", + } + } + } + assert key not in _training_config_attributes(off, "grpo") + + on = { + "policy": { + "megatron_cfg": { + "activation_checkpointing": True, + "recompute_granularity": "full", + } + } + } + assert _training_config_attributes(on, "grpo")[key] == "full" + + +def test_process_attributes_report_the_local_rank_when_the_launcher_set_one( + monkeypatch, +): + """``RayWorkerGroup`` exports ``LOCAL_RANK`` beside ``RANK``. + + It identifies a device on one node, which is what a reader has when the + complaint came from ``nvidia-smi`` rather than from the job. The driver owns + no device, so a process without the variable reports no local rank rather + than defaulting to zero and claiming one. + """ + pytest.importorskip("nemo.lens") + monkeypatch.setenv("LOCAL_RANK", "3") + assert _process_attributes(rank=11, world_size=16)["nv.dl.local_rank"] == 3 + + monkeypatch.delenv("LOCAL_RANK") + assert "nv.dl.local_rank" not in _process_attributes(rank=0, world_size=1) + + +def test_process_attributes_name_the_provider_and_the_run(monkeypatch): + # nv.dl.run.uuid duplicates lens's own nemo.run.id by design: the two + # namespaces are read by different consumers, and a backend joining RL runs + # to Megatron ones groups on the nv.dl.* spelling. + pytest.importorskip("nemo.lens") + monkeypatch.delenv("LOCAL_RANK", raising=False) + attrs = _process_attributes(rank=0, world_size=1, run_id="job-7") + assert attrs["nv.dl.provider.name"] == "nemo-rl" + assert attrs["nv.dl.run.uuid"] == "job-7" + assert "nv.dl.run.uuid" not in _process_attributes(rank=0, world_size=1) def test_init_driver_returns_none_when_disabled(): @@ -289,13 +419,19 @@ def _fake_setup_telemetry(config, **kwargs): handle = init_telemetry_worker() assert handle is not None # Rank reaches lens as a resource attribute now, alongside the group name: - # RANK is group-local, so the two are only useful together. - assert captured["resource_attributes"] == { - "nv.dl.campaign.stage": "RL", - "rl.worker_group": "vllm_policy", - "nv.dl.rank": 3, - "nv.dl.world_size": 8, - } + # RANK is group-local, so the two are only useful together. Checked by + # subset rather than by whole-dict equality, because the rest of what a + # worker reports depends on its environment -- the nv.dl.software.* versions + # are absent without torch, and nv.dl.local_rank without a launcher. + assert ( + captured["resource_attributes"].items() + >= { + "nv.dl.campaign.stage": "RL", + "rl.worker_group": "vllm_policy", + "nv.dl.rank": 3, + "nv.dl.world_size": 8, + }.items() + ) def test_init_worker_explicit_rank_overrides_env(monkeypatch): From 8d017b45ba23f343bcef5494f8a3b78981a38cbf Mon Sep 17 00:00:00 2001 From: Raj Singh Date: Thu, 8 Oct 2026 13:45:31 -0700 Subject: [PATCH 04/17] feat(telemetry): trace checkpoint finalization to durability rl..checkpointing ends when the training loop is released, which with async save is after D2H staging only -- the writer processes are still going and tmp_step_N has not been renamed. Nothing answered "when did step 7 become resumable", which is the question asked after a crash. Add rl.checkpoint.finalize, opened on the thread begin_finalization spawns and closed when the rename lands. It carries rl.checkpoint.async_save, because both save shapes emit the same span and a near-zero duration is otherwise ambiguous between a sync save and an async one that had already finished. Two properties need stating. Its rl.iteration is the step that triggered the save rather than the step running at the rename: a bare thread starts with an empty context, so the enclosing iteration_scope is both unreachable and, by the time the span closes, wrong. The step is read back from the tmp_step_N name, which is also now the only parser of that name -- _rename_checkpoint used to read it a second way, and on an unexpected name that one published the checkpoint under whatever sat between the second and third underscore. And it carries no rl.bucket. It overlaps the next step's productive spans, so an overhead tag would charge one stretch of wall clock to two buckets and push the goodput denominator past the wall clock. It is a child of the staging span for navigability, which does mean step 7's span subtree stays open into step 9; read step cost off rl..step rather than off the subtree. Opening the span is guarded. The work runs inside the try whose exception finalize_pending re-raises, so an unguarded failure here would abort the run and skip the rename -- reporting a lens incompatibility as a failed save. Failures from the wait and the rename still reach both the span and the caller. Signed-off-by: Raj Singh --- docs/observability/span-groups.md | 64 +++++- nemo_rl/telemetry/instrumentation.py | 61 +++++ nemo_rl/utils/checkpoint.py | 101 ++++++++- tests/unit/telemetry/test_checkpoint_spans.py | 212 ++++++++++++++++++ tests/unit/telemetry/test_instrumentation.py | 136 +++++++++++ tests/unit/telemetry/test_setup.py | 6 +- tests/unit/telemetry/test_source_drift.py | 6 + tests/unit/utils/test_checkpoint.py | 30 +++ 8 files changed, 600 insertions(+), 16 deletions(-) create mode 100644 tests/unit/telemetry/test_checkpoint_spans.py diff --git a/docs/observability/span-groups.md b/docs/observability/span-groups.md index bb703efdd45..1884084afea 100644 --- a/docs/observability/span-groups.md +++ b/docs/observability/span-groups.md @@ -50,7 +50,7 @@ One consequence worth knowing: registration is an *import side effect*, so a gro |---|---|---| | `job` | base | the whole-run root span (`rl..job`) | | `setup` | RL | startup, before the first step: `rl.startup` and the `rl.setup.*` phases under it | -| `checkpoint` | base | `rl..checkpointing` | +| `checkpoint` | base | `rl..checkpointing` — staging, which is what blocks the loop; `rl.checkpoint.finalize` — the asynchronous tail, open until the checkpoint is durable. See [Checkpoint spans](#a-checkpoint-takes-two-spans) | | `evaluate` | base | `rl..evaluate` | | `model_init` | base | `rl.vllm.load_model` (generation worker), `rl.policy.load_model` / `rl.value.load_model` (training workers) | | `step` | base | `rl..step` (one per training step) | @@ -127,6 +127,7 @@ These are set on spans for filtering — they answer "which one?" / "what kind?" | `rl.num_prompt_groups` | async rollout batch width, so a gap-filling batch is not read as an unexplained speed-up | | `rl.gym.batch_size` | how many examples one `rl.gym.run_rollouts` span covers — the NeMo-Gym counterpart to `rl.num_prompt_groups` | | `rl.idle.polls` | how many retries one `rl.idle.buffer_starvation` span covers on the single-controller path, where the span is coalesced over a poll loop. Read it against the duration: the same ten seconds is two thousand clean 5 ms polls or two hundred polls whose selection ran long, which are opposite diagnoses | +| `rl.checkpoint.async_save` | whether bytes were still being written when the loop resumed. `true` means `rl.checkpoint.finalize` spans a real wait on the writer processes; `false` means the write finished inside the staging span and the finalize span is just a rename. Both shapes emit the same span, so without this a near-zero duration is ambiguous between the two | | `rl.rollout.attempt` | SingleController dispatch attempt: `0` is a first try, `> 0` a substitution after a skipped group, whose tokens were discarded | | `rl.target_step` | the training step an `rl.sc.generate_and_push` dispatch is aimed at; omitted when the dispatch is unstamped | | `rl.critic_epochs` | critic epochs covered by one `rl.sc.value_training` span | @@ -167,6 +168,65 @@ calls and coroutines but not raw threads or other processes. Worker-side spans such as `rl.policy.load_model` are therefore unstamped; they are model init, not step work, so they sit outside any step regardless. +`rl.checkpoint.finalize` is the one place the scope would have been both +unreachable *and* wrong, so it sets `rl.iteration` itself — see below. + +### A checkpoint takes two spans + +An async save returns as soon as the weights are staged off the GPU; the writer +processes keep going in the background and the checkpoint is not resumable +until they finish and `tmp_step_N` is renamed to `step_N`. One span cannot +describe both facts, so there are two. + +| Span | Opens | Closes | Answers | +|---|---|---|---| +| `rl..checkpointing` | when the step decides to save | when the loop is released | how long checkpointing **cost** the step | +| `rl.checkpoint.finalize` | when `begin_finalization` hands off | when the rename lands | when the checkpoint became **durable** | + +`rl.checkpoint.finalize` is a **child** of the staging span, not a sibling: the +trace context is captured inside `rl..checkpointing` and handed to the +thread, which keeps the pair navigable as one unit. A child that ends after its +parent is unusual but legal, and it is the honest shape — so a trace viewer +shows step 7's subtree still open during step 9: + +``` +step 7 step 8 step 9 +└─ rl.grpo.checkpointing ──┤ + (stage, blocking) │ + └─ rl.checkpoint.finalize ──────────────────────────┤ + rl.iteration = 7 │ + (writers drain, then rename) ┘ +``` + +Three consequences follow from that overlap: + +- **Step 7's span subtree extends past step 7.** A rollup that sums span + durations under a step span will attribute the finalize wait to the step that + triggered it, which is the right answer for "when did step 7 become durable" + and the wrong one for "how long did step 7 take". Read step cost off + `rl..step` and the buckets below, not off the subtree. + +- **`rl.checkpoint.finalize` carries no `rl.bucket`.** It runs concurrently with + the next step's productive spans, so tagging it `overhead` would charge one + stretch of wall clock to two buckets at once and push the goodput denominator + past the wall clock. Same reasoning as the collector-side `efficiency` + categories. The staging span it follows *is* `overhead`, because that one + genuinely blocks. +- **Its `rl.iteration` is the step that triggered the save**, not the step + running when the rename lands. The work happens on a daemon thread that the + enclosing `iteration_scope` cannot reach, so the step is read back from the + `tmp_step_N` directory name and set on the span explicitly. `rl.iteration = 7` + therefore selects the save of step 7's weights however late it completed. + +A finalization that fails — a writer rank dying, a rename onto a full disk — +ends the span with an error status, and the same error is re-raised from the +next `finalize_pending()`. Both spans are in the `checkpoint` group, so one +group switch turns the pair on or off together. + +The reverse never happens: failing to *open* the finalize span is warned about +and the finalization proceeds untraced. Instrumentation does not get to decide +whether a checkpoint becomes resumable. + ### Span group → `rl.bucket` Leaf groups are tagged automatically when using @@ -180,7 +240,7 @@ so the call site shows which of the two it is — see |---|---| | `job`, `step`, `rollout`, `model_init`, `evaluate`, `setup`, `per_prompt` (aliased `U_JOB`, `U_STEP`, …) | *(none — umbrella)* | | `generation`, `reward`, `policy_update` | `productive` | -| `data_processing`, `data_plane`, `checkpoint`, `logprob`, `advantage` | `overhead` | +| `data_processing`, `data_plane`, `checkpoint`, `logprob`, `advantage` | `overhead` — except `rl.checkpoint.finalize`, which is unbucketed because it overlaps the next step (see below) | | `efficiency` | `idle` for the two driver-side phases; *none* for the two collector-side ones (see below) | Rolled-up `rl.goodput` is **monitor-derived**, not emitted by NeMo-RL. diff --git a/nemo_rl/telemetry/instrumentation.py b/nemo_rl/telemetry/instrumentation.py index 7c5e4e6f8d8..70cadcaebae 100644 --- a/nemo_rl/telemetry/instrumentation.py +++ b/nemo_rl/telemetry/instrumentation.py @@ -92,6 +92,14 @@ # work ran long, which are opposite diagnoses. RL_IDLE_POLLS_ATTR = "rl.idle.polls" +# Whether a checkpoint's bytes were still being written when the training loop +# moved on. Distinguishes the two shapes ``rl.checkpoint.finalize`` can take: +# with async save it spans a real wait on the writer processes, with sync save +# the write already finished inside the staging span and this is just a rename. +# Same span either way, so without this a near-zero duration is ambiguous +# between "sync save" and "async save that had already landed". +RL_CHECKPOINT_ASYNC_ATTR = "rl.checkpoint.async_save" + __all__ = [ "NO_SPAN", "managed_span", @@ -123,7 +131,9 @@ "startup_span", "setup_span", "evaluate_span", + "checkpoint_finalize_span", "RL_ALGORITHM", + "RL_CHECKPOINT_ASYNC_ATTR", "RL_EFFICIENCY_CATEGORY_ATTR", "RL_IDLE_POLLS_ATTR", "RL_ITERATION_ATTR", @@ -956,6 +966,57 @@ def evaluate_span( yield span +@contextmanager +def checkpoint_finalize_span( + iteration: Optional[int], + *, + async_save: bool, + tracer: Optional[Tracer] = None, + **attributes: Any, +) -> Iterator[Any]: + """The asynchronous tail of a checkpoint, named ``rl.checkpoint.finalize``. + + Closes when the checkpoint is durable -- every rank's writer has flushed and + ``tmp_step_N`` has been renamed to ``step_N`` -- which is the moment a + checkpoint can actually be resumed from. ``rl..checkpointing`` cannot + answer that: it ends when the training loop is released, which with async + save is after D2H staging only, with the bytes still in flight. + + *iteration* is passed rather than inherited, and that is the point of the + helper. The work runs on a daemon thread spawned by ``begin_finalization`` + and routinely outlives the step that triggered it, so the enclosing + :func:`iteration_scope` is both unreachable (a raw thread starts with an + empty context) and wrong by the time the span closes (the loop has moved + on). Stamping it explicitly keeps the span filed under the step whose + weights it holds, not whichever step happened to be running at the rename. + Set directly rather than left to the scope because lens's span processor + skips keys already present, so an explicit value wins over an inherited one. + + Unbucketed, like the collector's waits and for the same reason: it overlaps + the next step's productive spans, so an ``overhead`` tag here would charge + one stretch of wall clock to two buckets at once. Reached through + ``_managed_span`` rather than :func:`managed_span` to get that -- the group + stays ``checkpoint`` so one group switch still turns every checkpoint span + on or off together. + + Args: + iteration: Step the checkpoint was triggered at, or None when it cannot + be recovered from the directory name. + async_save: Whether the caller handed a wait function, i.e. whether + bytes were still being written when the training loop resumed. + """ + attrs: dict[str, Any] = {RL_CHECKPOINT_ASYNC_ATTR: async_save, **attributes} + if iteration is not None: + attrs[RL_ITERATION_ATTR] = iteration + with _managed_span( + RLSpanGroup.CHECKPOINT, + "rl.checkpoint.finalize", + tracer=tracer, + **attrs, + ) as span: + yield span + + def trace_fn( group: str, name: str, tracer: Optional[Tracer] = None ) -> Callable[[_F], _F]: diff --git a/nemo_rl/utils/checkpoint.py b/nemo_rl/utils/checkpoint.py index 5f6839cc84d..18010d92331 100644 --- a/nemo_rl/utils/checkpoint.py +++ b/nemo_rl/utils/checkpoint.py @@ -26,10 +26,12 @@ import time import warnings from concurrent.futures import Future, ThreadPoolExecutor +from contextlib import ExitStack, contextmanager from pathlib import Path from typing import ( Any, Callable, + Iterator, Literal, Mapping, NotRequired, @@ -43,8 +45,66 @@ import yaml from pydantic import BaseModel +from nemo_rl.telemetry.instrumentation import ( + checkpoint_finalize_span, + current_trace_carrier, + remote_trace_context, +) + PathLike = Union[str, "os.PathLike[Any]"] +# ``init_tmp_checkpoint`` names the staging directory ``tmp_step_{step}``, and +# that name is the only place the step survives into ``begin_finalization``, +# which takes the path alone. Parsed rather than added as an argument, so every +# algorithm's call site stays untouched. +_TMP_STEP_PATTERN = re.compile(r"^tmp_step_(\d+)$") + + +def _step_from_tmp_path(checkpoint_path: PathLike) -> Optional[int]: + """Step a ``tmp_step_N`` directory belongs to, or None if unrecognised. + + None rather than a raise: one caller only labels a span with this, and a + checkpoint staged under some other name should still be finalized. + """ + match = _TMP_STEP_PATTERN.match(Path(checkpoint_path).name) + return int(match.group(1)) if match else None + + +@contextmanager +def _finalize_telemetry( + carrier: Mapping[str, str], step: Optional[int], *, async_save: bool +) -> Iterator[None]: + """Open ``rl.checkpoint.finalize``, or carry on without it if that fails. + + Guarded because the block it wraps performs the rename that makes a + checkpoint resumable, and nothing in this package lets observability decide + whether that happens -- the same stance :meth:`_warn_on_delete_failure` + takes for pruning. Unguarded, a lens incompatibility would surface as a + failed save: the exception reaches ``finalize_pending`` as + ``_finalize_error``, so the run dies reporting a checkpoint failure *and* + skips the rename. + + Only entering is guarded. A failure inside the block is the checkpoint's + own and must still reach both the span and the caller. + """ + stack = ExitStack() + try: + stack.enter_context(remote_trace_context(carrier)) + stack.enter_context(checkpoint_finalize_span(step, async_save=async_save)) + # Deliberately broad: the point is that no failure from the telemetry stack, + # whatever it is, reaches the caller from here. + except Exception as e: + stack.close() + warnings.warn( + f"Could not open the checkpoint finalization span: {e!r}. " + "Finalization continues untraced; the checkpoint is unaffected.", + stacklevel=2, + ) + yield + else: + with stack: + yield + def _load_megatron_common_state_dict(iteration_dir: Path) -> dict[str, Any]: """Load common state from either legacy or current MCore checkpoints.""" @@ -374,9 +434,18 @@ def _rename_checkpoint(self, checkpoint_path: PathLike) -> None: If step_N already exists (defensive guard for edge cases, e.g. resuming training), performs a pseudo-atomic swap via an intermediate old_step_N directory. + + Raises: + ValueError: *checkpoint_path* is not a ``tmp_step_N`` directory, so + there is no step to publish it under. """ checkpoint_path = Path(checkpoint_path) - step = checkpoint_path.name.split("_")[2] + step = _step_from_tmp_path(checkpoint_path) + if step is None: + raise ValueError( + f"Expected a tmp_step_ directory from init_tmp_checkpoint(), " + f"got {checkpoint_path.name!r}" + ) to_checkpoint_path = checkpoint_path.parent / f"step_{step}" if to_checkpoint_path.exists(): old_checkpoint_path = checkpoint_path.parent / f"old_step_{step}" @@ -422,19 +491,29 @@ def begin_finalization( self.finalize_pending() self._pending_checkpoint_path = Path(checkpoint_path) self._finalize_error = None + step = _step_from_tmp_path(checkpoint_path) + # Captured here because a bare Thread starts with an empty context: read + # it on the worker instead and the finalize span re-roots into a trace of + # its own, detached from the run that produced the checkpoint. + carrier = current_trace_carrier() def _finalize(): + # The work is inside the span so a failed wait or rename is recorded + # on it, and the span is entered through a guard so the reverse + # cannot happen: telemetry must not be what skips the rename. try: - if wait_fn is not None: - wait_fn() - self._rename_checkpoint(checkpoint_path) - # Prune old checkpoints off the critical path. Surface any - # failure via a done-callback so a broken delete is not silently - # swallowed (the discarded Future would otherwise hide it). - delete_future = self._delete_executor.submit( - self.remove_old_checkpoints - ) - delete_future.add_done_callback(self._warn_on_delete_failure) + with _finalize_telemetry(carrier, step, async_save=wait_fn is not None): + if wait_fn is not None: + wait_fn() + self._rename_checkpoint(checkpoint_path) + # Prune old checkpoints off the critical path. Surface any + # failure via a done-callback so a broken delete is not + # silently swallowed (the discarded Future would otherwise + # hide it). + delete_future = self._delete_executor.submit( + self.remove_old_checkpoints + ) + delete_future.add_done_callback(self._warn_on_delete_failure) except Exception as e: self._finalize_error = e finally: diff --git a/tests/unit/telemetry/test_checkpoint_spans.py b/tests/unit/telemetry/test_checkpoint_spans.py new file mode 100644 index 00000000000..78961729d6f --- /dev/null +++ b/tests/unit/telemetry/test_checkpoint_spans.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""``rl.checkpoint.finalize`` against the real ``CheckpointManager``. + +Here rather than in ``tests/unit/utils/test_checkpoint.py`` for the autouse +fixture in this package's conftest, which resets the process-global lens and +OTel providers between tests. The subject is the telemetry rather than the +rename, so the span assertions and that fixture belong together. +""" + +import threading + +import pytest + +from nemo_rl.utils.checkpoint import CheckpointManager + +try: + from nemo.lens import NemoLensConfig, setup_telemetry + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, + ) + + from nemo_rl.telemetry.instrumentation import ( + RL_CHECKPOINT_ASYNC_ATTR, + RL_ITERATION_ATTR, + iteration_scope, + managed_span, + ) + from nemo_rl.telemetry.span_groups import RLSpanGroup + + _HAS_LENS = True +except ImportError: + _HAS_LENS = False + +pytestmark = pytest.mark.skipif( + not _HAS_LENS, reason="nemo-lens (+ opentelemetry sdk) not installed" +) + + +@pytest.fixture +def checkpoint_dir(tmp_path): + return tmp_path.resolve() / "checkpoints" + + +@pytest.fixture +def checkpointer(checkpoint_dir): + return CheckpointManager( + { + "enabled": True, + "checkpoint_dir": checkpoint_dir, + "metric_name": "loss", + "higher_is_better": False, + "save_period": 1, + "keep_top_k": 3, + "save_optimizer": True, + } + ) + + +@pytest.fixture +def exporter(): + recorder = InMemorySpanExporter() + handle = setup_telemetry( + NemoLensConfig(enabled=True, span_groups="all"), span_exporter=recorder + ) + yield handle, recorder + handle.shutdown() + + +def _named(recorder, name): + (span,) = [s for s in recorder.get_finished_spans() if s.name == name] + return span + + +def test_the_span_stays_open_until_the_checkpoint_is_durable( + checkpointer, checkpoint_dir, exporter +): + """The gap this closes: ``rl..checkpointing`` ends at staging. + + Reproduces the async-save shape. ``wait_fn`` blocks past the step that + triggered the save, so the only span that lines up with "resumable from + step 7" is one that is still open while the loop runs step 8. + """ + handle, recorder = exporter + writers_done = threading.Event() + tmp_dir = checkpointer.init_tmp_checkpoint(7, {"loss": 0.5}) + + with iteration_scope(7): + with managed_span( + RLSpanGroup.CHECKPOINT, "rl.grpo.checkpointing", tracer=handle.tracer + ): + checkpointer.begin_finalization( + tmp_dir, wait_fn=lambda: writers_done.wait(5) + ) + # Staging is over and the loop is free, but nothing is resumable yet. + assert not (checkpoint_dir / "step_7").exists() + + with iteration_scope(8): + writers_done.set() + checkpointer.finalize_pending() + handle.shutdown() + + assert (checkpoint_dir / "step_7").exists() + staging = _named(recorder, "rl.grpo.checkpointing") + durable = _named(recorder, "rl.checkpoint.finalize") + assert durable.end_time > staging.end_time + assert durable.parent.span_id == staging.get_span_context().span_id + + +def test_the_span_is_filed_under_the_step_that_triggered_the_save( + checkpointer, exporter +): + """Not step 8, which is what the loop had moved on to at the rename. + + ``begin_finalization`` runs its work on a bare thread, so the enclosing + ``iteration_scope`` does not reach it; the step comes from the staging + directory name instead. + """ + handle, recorder = exporter + writers_done = threading.Event() + tmp_dir = checkpointer.init_tmp_checkpoint(7, {"loss": 0.5}) + + with iteration_scope(7): + checkpointer.begin_finalization(tmp_dir, wait_fn=lambda: writers_done.wait(5)) + with iteration_scope(8): + writers_done.set() + checkpointer.finalize_pending() + handle.shutdown() + + durable = _named(recorder, "rl.checkpoint.finalize") + assert durable.attributes[RL_ITERATION_ATTR] == 7 + assert durable.attributes[RL_CHECKPOINT_ASYNC_ATTR] is True + + +def test_a_checkpoint_that_never_landed_is_not_reported_as_a_clean_one( + checkpointer, exporter +): + """The error reaches ``finalize_pending`` and the span at the same time.""" + handle, recorder = exporter + tmp_dir = checkpointer.init_tmp_checkpoint(7, {"loss": 0.5}) + + def wait_fn(): + raise RuntimeError("writer rank 3 died") + + checkpointer.begin_finalization(tmp_dir, wait_fn=wait_fn) + with pytest.raises(RuntimeError, match="Background checkpoint finalization failed"): + checkpointer.finalize_pending() + handle.shutdown() + + durable = _named(recorder, "rl.checkpoint.finalize") + assert not durable.status.is_ok + assert durable.attributes[RL_ITERATION_ATTR] == 7 + + +def test_a_span_that_cannot_be_opened_does_not_cost_the_checkpoint( + checkpointer, checkpoint_dir, exporter, monkeypatch +): + """Instrumentation does not get to decide whether a checkpoint lands. + + The span is opened on the finalization thread, inside the try whose + exception ``finalize_pending`` re-raises -- so an unguarded failure here + would both abort the run and skip the rename, reporting a lens + incompatibility as a failed save. ``NemoLensConfig.from_env`` dropping a + parameter is the kind of change that reaches this path. + """ + + def _broken_span(*args, **kwargs): + raise RuntimeError("lens dropped a parameter") + + monkeypatch.setattr( + "nemo_rl.utils.checkpoint.checkpoint_finalize_span", _broken_span + ) + tmp_dir = checkpointer.init_tmp_checkpoint(7, {"loss": 0.5}) + + # Both calls inside the block: the warning is emitted on the finalization + # thread, which begin_finalization has already started. + with pytest.warns(UserWarning, match="Could not open the checkpoint"): + checkpointer.begin_finalization(tmp_dir) + checkpointer.finalize_pending() + + assert (checkpoint_dir / "step_7").exists() + + +def test_a_sync_save_is_distinguishable_from_an_async_one(checkpointer, exporter): + """No ``wait_fn`` means the bytes already landed and this is just a rename. + + Both shapes emit the same span, so without the attribute a near-zero + duration could be either one. + """ + handle, recorder = exporter + tmp_dir = checkpointer.init_tmp_checkpoint(7, {"loss": 0.5}) + + checkpointer.begin_finalization(tmp_dir) + checkpointer.finalize_pending() + handle.shutdown() + + assert ( + _named(recorder, "rl.checkpoint.finalize").attributes[RL_CHECKPOINT_ASYNC_ATTR] + is False + ) diff --git a/tests/unit/telemetry/test_instrumentation.py b/tests/unit/telemetry/test_instrumentation.py index ce77e98b260..d1bbf92e793 100644 --- a/tests/unit/telemetry/test_instrumentation.py +++ b/tests/unit/telemetry/test_instrumentation.py @@ -17,12 +17,14 @@ import asyncio import inspect import logging +import threading import pytest from nemo_rl.telemetry.instrumentation import ( EFFICIENCY_CATEGORY_BUCKET, RL_BUCKET_ATTR, + RL_CHECKPOINT_ASYNC_ATTR, RL_EFFICIENCY_CATEGORY_ATTR, RL_IDLE_POLLS_ATTR, RL_ITERATION_ATTR, @@ -33,6 +35,7 @@ bucket_for_efficiency_category, bucket_for_span_group, bucket_scope, + checkpoint_finalize_span, current_trace_carrier, dispatch_with_trace_context, efficiency_span, @@ -1209,3 +1212,136 @@ def test_iteration_scope_does_not_leak_past_the_step(): for span in exporter.get_finished_spans(): assert RL_ITERATION_ATTR not in span.attributes + + +# --------------------------------------------------------------------------- # +# Durable-checkpoint span # +# --------------------------------------------------------------------------- # +@requires_lens +def test_checkpoint_finalize_span_outlives_the_step_that_triggered_it(): + """The reason the span exists: ``rl..checkpointing`` ends too early. + + Mirrors ``begin_finalization``: the loop is released after staging and runs + another step while the writers drain, so the finalize span must still be + open once the next step's span has closed. + """ + handle, exporter = _setup("all") + writers_done = threading.Event() + + def finalize(carrier): + with ( + remote_trace_context(carrier), + checkpoint_finalize_span(7, async_save=True, tracer=handle.tracer), + ): + writers_done.wait(5) + + with iteration_scope(7): + with managed_span( + RLSpanGroup.CHECKPOINT, "rl.grpo.checkpointing", tracer=handle.tracer + ): + # begin_finalization is called from inside the staging span. + worker = threading.Thread(target=finalize, args=(current_trace_carrier(),)) + worker.start() + + with iteration_scope(8): + with managed_span( + RLSpanGroup.POLICY_UPDATE, "rl.grpo.policy_update", tracer=handle.tracer + ): + pass + writers_done.set() + worker.join(5) + handle.shutdown() + + staging = _named(exporter, "rl.grpo.checkpointing") + next_step = _named(exporter, "rl.grpo.policy_update") + durable = _named(exporter, "rl.checkpoint.finalize") + assert durable.end_time > staging.end_time + assert durable.end_time > next_step.end_time + + +@requires_lens +def test_checkpoint_finalize_span_is_filed_under_the_triggering_step(): + """Not the step that happens to be running when the rename lands. + + The thread is started under step 7 but the span closes while the driver is + inside ``iteration_scope(8)``, which is the case that makes an inherited + iteration wrong. + """ + handle, exporter = _setup("all") + started = threading.Event() + release = threading.Event() + + def finalize(carrier): + with ( + remote_trace_context(carrier), + checkpoint_finalize_span(7, async_save=True, tracer=handle.tracer), + ): + started.set() + release.wait(5) + + with iteration_scope(7): + worker = threading.Thread(target=finalize, args=(current_trace_carrier(),)) + worker.start() + started.wait(5) + + with iteration_scope(8): + release.set() + worker.join(5) + handle.shutdown() + + durable = _named(exporter, "rl.checkpoint.finalize") + assert durable.attributes[RL_ITERATION_ATTR] == 7 + assert durable.attributes[RL_CHECKPOINT_ASYNC_ATTR] is True + + +@requires_lens +def test_checkpoint_finalize_span_is_parented_to_the_run(): + """A bare ``Thread`` starts with an empty context and would re-root here.""" + handle, exporter = _setup("all") + + def finalize(carrier): + with ( + remote_trace_context(carrier), + checkpoint_finalize_span(7, async_save=False, tracer=handle.tracer), + ): + pass + + with umbrella_span(RLSpanGroup.U_JOB, "rl.grpo.job", tracer=handle.tracer) as job: + job_id = job.get_span_context().span_id + worker = threading.Thread(target=finalize, args=(current_trace_carrier(),)) + worker.start() + worker.join(5) + handle.shutdown() + + assert _named(exporter, "rl.checkpoint.finalize").parent.span_id == job_id + + +@requires_lens +def test_checkpoint_finalize_span_carries_no_bucket(): + """It overlaps the next step, so a bucket would double-count wall clock. + + The staging span it follows is bucketed, because that one does block. + """ + handle, exporter = _setup("all") + with managed_span( + RLSpanGroup.CHECKPOINT, "rl.grpo.checkpointing", tracer=handle.tracer + ): + pass + with checkpoint_finalize_span(7, async_save=True, tracer=handle.tracer): + pass + handle.shutdown() + + assert RL_BUCKET_ATTR not in _named(exporter, "rl.checkpoint.finalize").attributes + staging = _named(exporter, "rl.grpo.checkpointing") + assert staging.attributes[RL_BUCKET_ATTR] == Bucket.OVERHEAD.value + + +@requires_lens +def test_checkpoint_finalize_span_follows_the_checkpoint_group(): + """One switch turns staging and finalization on or off together.""" + handle, exporter = _setup("generation") + with checkpoint_finalize_span(7, async_save=True, tracer=handle.tracer): + pass + handle.shutdown() + + assert exporter.get_finished_spans() == () diff --git a/tests/unit/telemetry/test_setup.py b/tests/unit/telemetry/test_setup.py index 40aa32cd132..f11a4f4f95f 100644 --- a/tests/unit/telemetry/test_setup.py +++ b/tests/unit/telemetry/test_setup.py @@ -157,12 +157,12 @@ def test_training_target_drops_a_non_positive_step_count(): def test_training_config_reads_either_backends_optimizer_spelling(): - # megatron names the algorithm, dtensor names the class. + # megatron names the algorithm, automodel names the class. megatron = {"policy": {"megatron_cfg": {"optimizer": {"optimizer": "adam"}}}} - dtensor = {"policy": {"optimizer": {"name": "torch.optim.AdamW"}}} + automodel = {"policy": {"optimizer": {"name": "torch.optim.AdamW"}}} key = "nv.dl.training.config.optimizer" assert _training_config_attributes(megatron, "grpo")[key] == "adam" - assert _training_config_attributes(dtensor, "grpo")[key] == "torch.optim.AdamW" + assert _training_config_attributes(automodel, "grpo")[key] == "torch.optim.AdamW" def test_recompute_granularity_is_reported_only_when_recompute_is_on(): diff --git a/tests/unit/telemetry/test_source_drift.py b/tests/unit/telemetry/test_source_drift.py index e7bb4325885..cae819ea925 100644 --- a/tests/unit/telemetry/test_source_drift.py +++ b/tests/unit/telemetry/test_source_drift.py @@ -33,6 +33,7 @@ _REPO / "nemo_rl" / "distributed", _REPO / "nemo_rl" / "data_plane", _REPO / "nemo_rl" / "environments", + _REPO / "nemo_rl" / "utils", _REPO / "examples", ) @@ -310,6 +311,10 @@ def _emitted_span_names( if called == "evaluate_span": algorithm = _string_arg(node) return {f"rl.{algorithm}.evaluate"} if algorithm else set() + if called == "checkpoint_finalize_span": + # Fixed name: CheckpointManager is shared across algorithms, so there + # is no algorithm to interpolate. + return {"rl.checkpoint.finalize"} if called == "efficiency_span": category = _string_arg(node) return {f"rl.{category.replace('/', '.')}"} if category else set() @@ -401,6 +406,7 @@ def test_every_registered_group_has_an_emitter(): "startup_span": RLSpanGroup.SETUP, "setup_span": RLSpanGroup.SETUP, "evaluate_span": RLSpanGroup.U_EVALUATE, + "checkpoint_finalize_span": RLSpanGroup.CHECKPOINT, "efficiency_span": RLSpanGroup.EFFICIENCY, "traced_worker_init": RLSpanGroup.MODEL_INIT, } diff --git a/tests/unit/utils/test_checkpoint.py b/tests/unit/utils/test_checkpoint.py index 6d77dcb989b..6a6a565e081 100644 --- a/tests/unit/utils/test_checkpoint.py +++ b/tests/unit/utils/test_checkpoint.py @@ -1200,6 +1200,25 @@ def test_swaps_when_target_exists(self, checkpoint_manager, checkpoint_dir): with open(step_dir / "training_info.json") as f: assert json.load(f)["loss"] == 0.2 + def test_rejects_a_directory_it_cannot_read_a_step_from( + self, checkpoint_manager, checkpoint_dir + ): + """The step comes out of the name, so an unreadable name has to raise. + + The previous spelling of that read, ``name.split("_")[2]``, returned + whatever sat between the second and third underscore -- publishing + ``tmp_weights_12_a`` as ``step_12``, a directory claiming to be a + checkpoint of a step it has nothing to do with. + """ + staged = checkpoint_dir / "tmp_weights_12_a" + staged.mkdir(parents=True) + + with pytest.raises(ValueError, match="Expected a tmp_step_ directory"): + checkpoint_manager._rename_checkpoint(staged) + + assert staged.exists() + assert not (checkpoint_dir / "step_12").exists() + # --------------------------------------------------------------------------- # Fault tolerance (ft_keep_latest_k / ft_save_period) retention tests @@ -1428,3 +1447,14 @@ def test_metric_topk_with_ft_keeps_best_and_recent_disjoint(self, checkpoint_dir # top-2 by metric -> {1, 2}; ft latest-2 -> {5, 6}; union -> {1, 2, 5, 6}. assert self._remaining_steps(checkpoint_dir) == [1, 2, 5, 6] + + +def test_step_from_tmp_path_tolerates_an_unexpected_name(): + """Reading the step is only used to label a span, so it must never raise. + + A caller that stages under some other name still gets a finalized + checkpoint; it just loses the iteration attribute on the span. + """ + assert checkpoint_module._step_from_tmp_path("/ckpt/tmp_step_412") == 412 + assert checkpoint_module._step_from_tmp_path("/ckpt/step_412") is None + assert checkpoint_module._step_from_tmp_path("/ckpt/tmp_step_final") is None From 9e1a83a1ee58b10d6817afcd86e115007bb22d3f Mon Sep 17 00:00:00 2001 From: Raj Singh Date: Thu, 8 Oct 2026 15:11:25 -0700 Subject: [PATCH 05/17] fix(telemetry): read resource attributes from the enabled backend only Every example config carries both policy backend blocks fully populated and only the enabled one describes the run, so reading whichever value is non-empty picked the disabled block. dpo.yaml reported tp 2 against the Automodel tp 1 it really uses, and grpo_math_1B.yaml reported optimizer adam against AdamW -- wrong rather than missing, which is worse. Select the backend from megatron_cfg.enabled once, then read only that block. Drop nv.dl.topology.size.dp. It was the one attribute derived rather than read, and dividing the cluster's GPU count by the model-parallel sizes assumes the policy owns every GPU. A reward model with its own resources, distillation teachers, expert parallelism and non-colocated generation each leave it fewer, and the division came out too high rather than failing. Read the worker's device mesh if the real number is wanted. Also: - sft_v2 reports itself as sft_v2 but keeps its settings under the sft section, so it never reported a training target. Map the section name. - nv.dl.run.uuid carried the raw run id, which is already reported as nemo.run.id. A backend joining RL runs to Megatron ones groups on lens's derived uuid5, so the raw id there matched nothing. - Read transformer-engine's version from distribution metadata instead of importing it, which loads its PyTorch extension in every process that sets up telemetry and raises OSError, not ImportError, when the shared library will not load. - Correct the docs: rl.iteration only stamps driver-process spans, the campaign stage is not always RL, local_rank is not a device index, and an unregistered resource attribute is passed through untyped rather than rejected. Signed-off-by: Raj Singh --- docs/observability/configuration.md | 24 ++-- docs/observability/span-groups.md | 28 +++-- nemo_rl/telemetry/setup.py | 188 +++++++++++++++------------- pyproject.toml | 13 +- tests/unit/telemetry/test_setup.py | 133 +++++++++++++------- 5 files changed, 229 insertions(+), 157 deletions(-) diff --git a/docs/observability/configuration.md b/docs/observability/configuration.md index 012c846f8d9..f20ffab845d 100644 --- a/docs/observability/configuration.md +++ b/docs/observability/configuration.md @@ -86,29 +86,31 @@ Each process sets stable-for-the-run values on its OTel `Resource` — `init_tel | Attribute | Source | |---|---| | `nemo.run.id` | the run id resolved above (`NEMO_RL_OTEL_RUN_ID` / `SLURM_JOB_ID` / generated). Emitted by lens from the `run_id` config field — see [Run identification](#run-identification) | -| `nv.dl.campaign.stage` | always the literal `RL`. Set by every NeMo-RL process, driver and worker, so a backend collecting several stages of a model's lifecycle (pretrain → SFT → RL) can select this stage without matching on service names | +| `nv.dl.campaign.stage` | the lifecycle stage this run occupies: `SFT` for `sft` and `sft_v2`, `RM` for `rm`, `RL` for everything else including DPO. Set by every NeMo-RL process, driver and worker, so a backend collecting several stages of a model's lifecycle (pretrain → SFT → RL) can select one stage without matching on service names | | `rl.algorithm` | the `algorithm=""` passed to `init_telemetry_driver` | | `rl.model` | `policy.model_name` | | `nemo.precision` | `policy.precision` | | `nv.dl.provider.name` | always the literal `nemo-rl`: which training framework produced the telemetry, for a backend collecting several. Separate from the service name, which a user may rename | -| `nv.dl.run.uuid` | the same run id as `nemo.run.id`. Duplicated on purpose — a backend joining RL runs to Megatron ones groups on the `nv.dl.*` spelling | -| `nv.dl.topology.size.tp` | `policy.megatron_cfg` / `automodel_cfg` TP size | -| `nv.dl.topology.size.pp` | `policy.megatron_cfg` PP size | -| `nv.dl.topology.size.dp` | derived: `cluster.num_nodes × cluster.gpus_per_node ÷ (TP × PP × CP)`. Omitted when `policy.generation.colocated.enabled` is false, because generation then holds ranks the policy does not have, and when the product does not divide the world size | +| `nv.dl.run.uuid` | lens's `derive_nv_dl_run_uuid` over the run id — a uuid5, not the raw id, which is reported separately as `nemo.run.id`. A backend joining RL runs to Megatron ones groups on the derived value, so the raw id under this key would match nothing | +| `nv.dl.topology.size.tp` | TP size, read from whichever backend `policy.megatron_cfg.enabled` selects. Every config carries both blocks populated and only the live one is read, so a disabled `megatron_cfg` never contributes | +| `nv.dl.topology.size.pp` | PP size, Megatron runs only — Automodel has no pipeline dimension | | `nv.dl.training.config.global_batch_size` | `policy.train_global_batch_size` | | `nv.dl.training.config.micro_batch_size` | `policy.train_micro_batch_size` | | `nv.dl.training.config.sequence_length` | `policy.max_total_sequence_length` | -| `nv.dl.training.config.optimizer` | `policy.megatron_cfg.optimizer.optimizer` (`adam`), else `policy.optimizer.name` (`torch.optim.AdamW`) — the two backends spell it differently | -| `nv.dl.training.config.recompute_granularity` | `policy.megatron_cfg.recompute_granularity`, only when `activation_checkpointing` is on. Megatron ignores the granularity otherwise, so reporting it would describe recompute that is not happening | -| `nv.dl.training.target.train_iters` | `.max_num_steps`, when positive. `rm` spells "train for one epoch" as `-1`, which is omitted rather than reported as a target | -| `nv.dl.software.torch`, `.cuda`, `.nccl`, `.transformer_engine` | versions as loaded in this process, not as pinned in the lockfile. Each is omitted when unavailable — CUDA and NCCL on a CPU-only build, transformer-engine when it is not installed | +| `nv.dl.training.config.optimizer` | on Megatron `policy.megatron_cfg.optimizer.optimizer` (`adam`), otherwise `policy.optimizer.name` (`torch.optim.AdamW`) — the backends both spell it and store it differently, and only the enabled one is read | +| `nv.dl.training.config.recompute_granularity` | Megatron runs only: `policy.megatron_cfg.recompute_granularity`, and only when `activation_checkpointing` is on. Megatron ignores the granularity otherwise, so reporting it would describe recompute that is not happening | +| `nv.dl.training.target.train_iters` | `.max_num_steps`, when positive. `sft_v2` reads the `sft` section, which is where it keeps its settings. `rm` spells "train for one epoch" as `-1`, which is omitted rather than reported as a target | +| `nv.dl.software.torch`, `.cuda`, `.nccl` | versions as loaded in this process, not as pinned in the lockfile. CUDA and NCCL are omitted on a CPU-only build | +| `nv.dl.software.transformer_engine` | the installed distribution's version, read from package metadata rather than by importing it — importing transformer-engine loads its PyTorch extension in every process that sets up telemetry. Omitted when it is not installed | | `nv.dl.rank`, `nv.dl.world_size` | this process's rank and group size (`RANK` / `WORLD_SIZE`, or `0` / `1` for the driver and singleton actors) | -| `nv.dl.local_rank` | worker processes only: `LOCAL_RANK`, which `RayWorkerGroup` exports beside `RANK`. Identifies a device on one node, which is what you have when the complaint came from `nvidia-smi` rather than from the job | +| `nv.dl.local_rank` | worker processes only: `LOCAL_RANK`, which `RayWorkerGroup` exports beside `RANK`. Places a process within its own node, which is how you group a job's processes by the host they share. Not a physical device index — Ray gives each worker its own `CUDA_VISIBLE_DEVICES` | | `rl.worker_group` | worker processes only: the worker group's `name_prefix` (`lm_policy`, `vllm_policy`, ...), from `NRL_WORKER_GROUP`. The `NemoGym` actor reports the literal `nemo_gym` — it is built from the environment registry rather than by `RayWorkerGroup`, so nothing sets the env var for it | The config-derived rows are driver-only: a worker is handed the propagated environment, not the `MasterConfig`. Attribute construction is best-effort: a missing config key simply omits that attribute; it never raises. Plus auto-detected host / GPU / SLURM / Kubernetes attributes from lens's resource detection. -Every key above that lens declares is taken from `nemo.lens.semconv` rather than written as a string, so a typo is an `ImportError` instead of silent data loss. Note the two spellings of the parallelism sizes: `nv.dl.topology.size.*` is the resource-scoped one lens types through its `DL_RESOURCE_TYPES` registry, while the `dl.tensor_parallel.*` family is for span and metric attributes. +Every key above that lens declares is taken from `nemo.lens.semconv` rather than written as a string, so a typo is an `ImportError` instead of silent data loss. Note the two spellings of the parallelism sizes: `nv.dl.topology.size.*` is the resource-scoped one, registered in lens's `DL_RESOURCE_TYPES` and so typed by `normalize_resource_attributes`, while the `dl.tensor_parallel.*` family is for span and metric attributes. A resource attribute outside that registry is passed through rather than rejected, so the wrong spelling arrives as an untyped string instead of failing. + +Lens also declares `nv.dl.topology.size.dp`, which NeMo-RL does not report. RL configures no data-parallel size, so it could only be derived by dividing the cluster's GPU count by the model-parallel sizes — and that assumes the policy owns every GPU in the cluster. A reward model with its own `resources` block, distillation teachers, expert parallelism and non-colocated generation each leave the policy fewer, and the division comes out too high rather than failing. Read the worker's device mesh if you need the real number. ## Typical configurations diff --git a/docs/observability/span-groups.md b/docs/observability/span-groups.md index 1884084afea..be20c1e0382 100644 --- a/docs/observability/span-groups.md +++ b/docs/observability/span-groups.md @@ -148,10 +148,10 @@ rollout-side spans for that. Each algorithm opens its step inside `nemo_rl.telemetry.instrumentation.iteration_scope`, which stamps -`rl.iteration` on every span started within it. So the step umbrella, the leaves -under it, the checkpoint span and the rollout and generation spans all carry the -same value, and `rl.iteration = 412` selects everything that happened in that -step rather than the step span alone. +`rl.iteration` on every span started within it **in the driver process**. So +the step umbrella, the leaves under it and the checkpoint span all carry the +same value, and `rl.iteration = 412` selects the driver's whole step rather +than the step span alone. A span that passes `rl.iteration` itself keeps its own value — the scope is applied by a span processor that skips keys already set. `rl..evaluate` @@ -163,10 +163,22 @@ whole run and belongs to no single step. `rl.sc.generate_and_push` runs on the rollout pump's own clock rather than the training loop's, so it carries `rl.target_step` — the step it is generating *for* — instead. -The scope travels on a `ContextVar`, which reaches nested -calls and coroutines but not raw threads or other processes. Worker-side spans -such as `rl.policy.load_model` are therefore unstamped; they are model init, not -step work, so they sit outside any step regardless. +The scope travels on a `ContextVar`, which reaches nested calls and coroutines +but not raw threads or other processes. Only the trace parent crosses a Ray +boundary, so a span opened inside an actor is correctly placed in the trace but +carries no `rl.iteration` of its own. That covers more than worker internals: + +| Unstamped span | Runs in | +| --- | --- | +| `rl.policy.load_model` and the other worker-side leaves | the policy workers | +| `rl.gym.run_rollouts` | the `NemoGym` actor | +| `rl.sft_v2.read_batch`, `rl.sft_v2.prepare_batch` | the Energon SFT worker | +| the async-GRPO collector's rollout spans | the collector actor | + +To scope those to a step, walk up to the driver-side ancestor that is stamped, +or filter on the trace rather than the attribute. `rl.policy.load_model` is the +one case where nothing is lost: it is model init, not step work, so it sits +outside any step regardless. `rl.checkpoint.finalize` is the one place the scope would have been both unreachable *and* wrong, so it sets `rl.iteration` itself — see below. diff --git a/nemo_rl/telemetry/setup.py b/nemo_rl/telemetry/setup.py index a8b8c446beb..8188a67df70 100644 --- a/nemo_rl/telemetry/setup.py +++ b/nemo_rl/telemetry/setup.py @@ -40,6 +40,7 @@ from __future__ import annotations import functools +import importlib.metadata import logging import os import threading @@ -103,6 +104,19 @@ def _campaign_stage(algorithm: str) -> str: return _CAMPAIGN_STAGE_BY_ALGORITHM.get(algorithm, _DEFAULT_CAMPAIGN_STAGE) +# An algorithm whose settings live under a differently named config section. +# ``run_sft_v2.py`` reports itself as ``sft_v2`` while reading ``master_config +# .sft``, so looking its settings up by algorithm name finds nothing. Mapped +# rather than renamed because the name is also reported as ``rl.algorithm``, +# where the two implementations do need telling apart. +_CONFIG_SECTION_BY_ALGORITHM = {"sft_v2": "sft"} + + +def _config_section(algorithm: str) -> str: + """Config section holding *algorithm*'s settings; its own name by default.""" + return _CONFIG_SECTION_BY_ALGORITHM.get(algorithm, algorithm) + + # TelemetryConfig field -> NEMO_RL_OTEL_* env var. ``service_name`` maps to the # standard ``OTEL_SERVICE_NAME`` (lens reads it directly, unprefixed). _ENV_FIELD_MAP = { @@ -220,73 +234,59 @@ def _build_resource_attributes( return attrs +def _megatron_enabled(master_config: Any) -> bool: + """Whether the policy trains on Megatron rather than Automodel. + + Every example config carries both backend blocks fully populated, and only + the enabled one describes the run. So the other's values have to be ignored + rather than read as a fallback: ``dpo.yaml`` ships a disabled + ``megatron_cfg`` saying tp 2 beside the Automodel tp 1 it actually uses. + """ + return bool(_dig(master_config, "policy", "megatron_cfg", "enabled")) + + def _topology_attributes(master_config: Any) -> dict[str, Any]: """Parallelism sizes, under lens's resource-scoped topology keys. - ``nv.dl.topology.size.*`` rather than the ``dl.tensor_parallel.size`` family: - only the former is in lens's ``DL_RESOURCE_TYPES``, which is the registry - ``normalize_resource_attributes`` types a resource attribute against. The - ``dl.*`` spelling is the span and metric one. - - Each size lives under whichever policy backend is active, so each is looked - up under both megatron and automodel. Only the current spelling of the - latter: ``check_outdated_config`` rejects a config still saying - ``dtensor_cfg`` before a run reaches here. + ``nv.dl.topology.size.*`` rather than the ``dl.tensor_parallel.size`` + family: only the former is in lens's ``DL_RESOURCE_TYPES``, which feeds the + ``RESOURCE_TYPES`` table ``normalize_resource_attributes`` consults. A key + outside that table is passed through rather than rejected, so the ``dl.*`` + spelling would survive but arrive as an untyped string -- it is the span + and metric spelling, not the resource one. + + Read from the enabled backend alone, under that backend's own spelling: the + two disagree on names as well as values, and Automodel has no pipeline + dimension to report at all. + + No ``nv.dl.topology.size.dp``, though lens defines it. RL configures no + data-parallel size, so it could only be derived from the cluster's GPU + count, and that assumes the policy owns every one of them. A reward model, + distillation teachers, expert parallelism and non-colocated generation each + break the assumption in a way that reports a dp too high rather than + failing, so there is nothing to report instead of a plausible wrong number. """ - from nemo.lens.semconv import ( - NV_DL_TOPOLOGY_SIZE_DP, - NV_DL_TOPOLOGY_SIZE_PP, - NV_DL_TOPOLOGY_SIZE_TP, - ) + from nemo.lens.semconv import NV_DL_TOPOLOGY_SIZE_PP, NV_DL_TOPOLOGY_SIZE_TP - tp = _dig( - master_config, "policy", "megatron_cfg", "tensor_model_parallel_size" - ) or _dig(master_config, "policy", "automodel_cfg", "tensor_parallel_size") - pp = _dig(master_config, "policy", "megatron_cfg", "pipeline_model_parallel_size") - cp = _dig(master_config, "policy", "megatron_cfg", "context_parallel_size") or _dig( - master_config, "policy", "automodel_cfg", "context_parallel_size" - ) + if _megatron_enabled(master_config): + tp = _dig(master_config, "policy", "megatron_cfg", "tensor_model_parallel_size") + pp = _dig( + master_config, "policy", "megatron_cfg", "pipeline_model_parallel_size" + ) + else: + # Only the current spelling: ``check_outdated_config`` rejects a config + # still saying ``dtensor_cfg`` before a run reaches here. + tp = _dig(master_config, "policy", "automodel_cfg", "tensor_parallel_size") + pp = None attrs: dict[str, Any] = {} if tp: attrs[NV_DL_TOPOLOGY_SIZE_TP] = tp if pp: attrs[NV_DL_TOPOLOGY_SIZE_PP] = pp - dp = _data_parallel_size(master_config, tp=tp or 1, pp=pp or 1, cp=cp or 1) - if dp: - attrs[NV_DL_TOPOLOGY_SIZE_DP] = dp return attrs -def _data_parallel_size( - master_config: Any, *, tp: int, pp: int, cp: int -) -> Optional[int]: - """The ranks left for data parallelism once tp, pp and cp have claimed theirs. - - Derived rather than read: RL configures no dp size. Training spans the - cluster's GPUs and dp is whatever the other dimensions leave over. - - That identity breaks when generation runs on dedicated GPUs, because the - cluster then holds ranks the policy does not have. Returns ``None`` there, - rather than reporting a dp size for a group that is smaller than the - arithmetic says. Also ``None`` when the product does not divide the world - size, which the policy worker rejects with a far better message than a - resource attribute could carry. - """ - if _dig(master_config, "policy", "generation", "colocated", "enabled") is False: - return None - nodes = _dig(master_config, "cluster", "num_nodes") - gpus_per_node = _dig(master_config, "cluster", "gpus_per_node") - if not nodes or not gpus_per_node: - return None - - world_size = nodes * gpus_per_node - model_parallel_size = tp * pp * cp - if world_size % model_parallel_size: - return None - return world_size // model_parallel_size - - def _training_config_attributes(master_config: Any, algorithm: str) -> dict[str, Any]: """Batch, sequence and optimizer settings, plus the configured step target. @@ -312,29 +312,32 @@ def _training_config_attributes(master_config: Any, algorithm: str) -> dict[str, if value: attrs[key] = value - # Two spellings, one per backend: megatron names the algorithm ("adam"), - # dtensor names the class ("torch.optim.AdamW"). - optimizer = _dig( - master_config, "policy", "megatron_cfg", "optimizer", "optimizer" - ) or _dig(master_config, "policy", "optimizer", "name") + # Each backend keeps its optimizer somewhere else and spells it differently: + # megatron names the algorithm ("adam") inside its own block, Automodel + # names the class ("torch.optim.AdamW") in the policy's. Reading whichever + # is non-empty would report the disabled backend's answer. + if _megatron_enabled(master_config): + optimizer = _dig( + master_config, "policy", "megatron_cfg", "optimizer", "optimizer" + ) + # Gated on activation checkpointing because megatron ignores the + # granularity without it, so reporting it would describe recompute that + # is not happening. + if _dig(master_config, "policy", "megatron_cfg", "activation_checkpointing"): + granularity = _dig( + master_config, "policy", "megatron_cfg", "recompute_granularity" + ) + if granularity: + attrs[NV_DL_TRAINING_CONFIG_RECOMPUTE_GRANULARITY] = granularity + else: + optimizer = _dig(master_config, "policy", "optimizer", "name") if optimizer: attrs[NV_DL_TRAINING_CONFIG_OPTIMIZER] = optimizer - # Gated on activation checkpointing because megatron ignores the - # granularity without it, so reporting it would describe recompute that is - # not happening. - if _dig(master_config, "policy", "megatron_cfg", "activation_checkpointing"): - granularity = _dig( - master_config, "policy", "megatron_cfg", "recompute_granularity" - ) - if granularity: - attrs[NV_DL_TRAINING_CONFIG_RECOMPUTE_GRANULARITY] = granularity - - # The algorithm's own config section is named after it (``grpo:``, ``sft:``). # Non-positive values are dropped rather than reported: ``rm`` spells "train # for one epoch" as ``max_num_steps: -1``, which downstream would otherwise # read as a target of -1 iterations. - max_num_steps = _dig(master_config, algorithm, "max_num_steps") + max_num_steps = _dig(master_config, _config_section(algorithm), "max_num_steps") if isinstance(max_num_steps, int) and max_num_steps > 0: attrs[NV_DL_TRAINING_TARGET_TRAIN_ITERS] = max_num_steps return attrs @@ -344,9 +347,12 @@ def _software_versions() -> dict[str, Any]: """Library versions as actually loaded in this process. The lockfile says what was meant to be installed; these say what the run - imported, which is the pair worth comparing when a run regresses after an - image rebuild. Every lookup is optional: CUDA and NCCL are absent from a - CPU-only build, and transformer-engine is not installed for every backend. + actually has, which is the pair worth comparing when a run regresses after + an image rebuild. Torch and what it was built against are read from the + imported module, transformer-engine from its distribution metadata -- see + below for why that one is not imported. Every lookup is optional: CUDA and + NCCL are absent from a CPU-only build, and transformer-engine is not + installed for every backend. """ from nemo.lens.semconv import ( NV_DL_SOFTWARE_CUDA, @@ -375,14 +381,17 @@ def _software_versions() -> dict[str, Any]: else: attrs[NV_DL_SOFTWARE_NCCL] = ".".join(str(part) for part in nccl_version) + # Read from the installed distribution's metadata rather than by importing + # the package. Importing transformer-engine pulls in its PyTorch extension, + # which is seconds of work and a chunk of memory in every process that sets + # up telemetry, and which raises ``OSError`` rather than ``ImportError`` + # when the shared library will not load against the driver present. try: - import transformer_engine - except ImportError: + te_version = importlib.metadata.version("transformer-engine") + except importlib.metadata.PackageNotFoundError: pass else: - te_version = getattr(transformer_engine, "__version__", "") - if te_version: - attrs[NV_DL_SOFTWARE_TRANSFORMER_ENGINE] = str(te_version) + attrs[NV_DL_SOFTWARE_TRANSFORMER_ENGINE] = te_version return attrs @@ -398,13 +407,16 @@ def _process_attributes(rank: int, world_size: int, run_id: str = "") -> dict[st 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. ``nv.dl.local_rank`` narrows that to a - device on one node, which is what a reader has when the complaint came from - ``nvidia-smi`` rather than from the job. + apart from its peers downstream. ``nv.dl.local_rank`` adds its position + within its own node, which is what groups a job's processes by the host + they share. It is not a physical device index: Ray gives each worker its + own ``CUDA_VISIBLE_DEVICES``, so local rank 0 is whichever GPU that worker + was placed on. Called by both entry points, since a resource is per process and neither path sees the other's. """ + from nemo.lens.resources.slurm import derive_nv_dl_run_uuid from nemo.lens.semconv import ( NV_DL_LOCAL_RANK, NV_DL_PROVIDER_NAME, @@ -423,11 +435,15 @@ def _process_attributes(rank: int, world_size: int, run_id: str = "") -> dict[st local_rank = os.environ.get("LOCAL_RANK", "").strip() if local_rank.isdigit(): attrs[NV_DL_LOCAL_RANK] = int(local_rank) - # Duplicates lens's own ``nemo.run.id`` by design: the two namespaces are - # read by different consumers, and a backend joining RL runs to Megatron - # ones groups on the ``nv.dl.*`` spelling. - if run_id: - attrs[NV_DL_RUN_UUID] = run_id + # Derived rather than set to ``run_id`` itself, which is already reported + # as lens's ``nemo.run.id``. ``nv.dl.run.uuid`` is a uuid5 over the + # scheduler's identity for this attempt, and a backend grouping RL runs + # with Megatron ones joins on that; a raw run id in the same key would + # match nothing. Returns None off a scheduler with no run id to fall back + # on, which cannot happen here but is handled rather than asserted. + run_uuid = derive_nv_dl_run_uuid(run_id=run_id) if run_id else None + if run_uuid: + attrs[NV_DL_RUN_UUID] = run_uuid attrs.update(_software_versions()) return attrs diff --git a/pyproject.toml b/pyproject.toml index ef7001a250b..ffd4ff61696 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,7 +80,9 @@ dependencies = [ # venv; `telemetry.enabled` in the config is the on/off switch. Sourced from # git in [tool.uv.sources] — see the note there; the floor is above the # newest PyPI release so a dropped pin fails to resolve instead of silently - # falling back to 0.1.0. The aiohttp extra is what stamps a traceparent on + # falling back to 0.1.0. It reads 0.2.0 rather than 0.3.0 because the pinned + # rev is the v0.3.0 tag but still declares 0.2.0 in its package_info, so a + # 0.3.0 floor would not resolve. The aiohttp extra is what stamps a traceparent on # the HTTP calls NeMo-Gym's client makes, so a Gym rollout continues the # run's trace rather than starting its own. "nemo-lens[sdk,aiohttp]>=0.2.0", @@ -507,10 +509,11 @@ override-dependencies = [ "mamba-ssm==2.3.2.post1", # megatron-core still pins lens v0.2.0 in its own [tool.uv.sources]; without an # override the two git URLs are a hard conflict, since the workspace resolves a - # single nemo-lens. We need the newer rev for `SpanRegistry` (which replaced the - # shipped `SpanGroup` class) and the consumer-driven metric registry - # (`MetricSpec`/`register_metric_group`), both of which landed after v0.2.0. Drop - # this once Megatron-LM bumps to the same or newer rev. + # single nemo-lens. The rev here is the v0.3.0 tag, which we need for + # `SpanRegistry` (which replaced the shipped `SpanGroup` class), the + # consumer-driven metric registry (`MetricSpec`/`register_metric_group`) and + # the scoped span attributes behind `rl.iteration`. Drop this once Megatron-LM + # bumps to the same or newer rev. "nemo-lens[sdk,aiohttp] @ git+https://github.com/NVIDIA-NeMo/Lens.git@7bcc03ce3565d2be6858b819b049bcc4c3ec05b4", ] diff --git a/tests/unit/telemetry/test_setup.py b/tests/unit/telemetry/test_setup.py index f11a4f4f95f..33206ba86db 100644 --- a/tests/unit/telemetry/test_setup.py +++ b/tests/unit/telemetry/test_setup.py @@ -16,7 +16,6 @@ from nemo_rl.telemetry.config import TelemetryConfig from nemo_rl.telemetry.setup import ( _build_resource_attributes, - _data_parallel_size, _dig, _process_attributes, _training_config_attributes, @@ -39,6 +38,7 @@ def __init__(self, telemetry=None, policy=None): "model_name": "org/Model-1B", "precision": "bfloat16", "megatron_cfg": { + "enabled": True, "tensor_model_parallel_size": 2, "pipeline_model_parallel_size": 1, }, @@ -85,7 +85,7 @@ def test_build_resource_attributes_automodel_tp(): policy={ "model_name": "org/Model-1B", "precision": "bfloat16", - "automodel_cfg": {"tensor_parallel_size": 4}, + "automodel_cfg": {"enabled": True, "tensor_parallel_size": 4}, } ) attrs = _build_resource_attributes(cfg, "grpo") @@ -93,13 +93,60 @@ def test_build_resource_attributes_automodel_tp(): assert "nv.dl.topology.size.pp" not in attrs +def test_the_disabled_backends_settings_are_never_reported(): + """Every example config carries both backend blocks fully populated. + + ``dpo.yaml`` is the live case: a disabled ``megatron_cfg`` saying tp 2 and + optimizer adam, beside the Automodel tp 1 and AdamW the run really uses. + Reading whichever value is non-empty picks the disabled block, and the + result is wrong rather than missing -- a plausible topology nobody ran. + """ + cfg = _FakeMasterConfig( + policy={ + "automodel_cfg": {"enabled": True, "tensor_parallel_size": 1}, + "optimizer": {"name": "torch.optim.AdamW"}, + "megatron_cfg": { + "enabled": False, + "tensor_model_parallel_size": 2, + "pipeline_model_parallel_size": 4, + "activation_checkpointing": True, + "recompute_granularity": "full", + "optimizer": {"optimizer": "adam"}, + }, + } + ) + attrs = _build_resource_attributes(cfg, "dpo") + assert attrs["nv.dl.topology.size.tp"] == 1 + assert attrs["nv.dl.training.config.optimizer"] == "torch.optim.AdamW" + # Automodel has no pipeline dimension, and recompute granularity is a + # megatron setting, so neither has anything to report here. + assert "nv.dl.topology.size.pp" not in attrs + assert "nv.dl.training.config.recompute_granularity" not in attrs + + +def test_the_data_parallel_size_is_not_reported(): + """Though lens defines ``nv.dl.topology.size.dp``. + + RL configures no dp size, so it could only come from dividing the cluster's + GPU count by the model-parallel sizes, and that assumes the policy owns + every GPU in the cluster. A reward model with its own resources + (``grpo_rm_1B.yaml``), distillation teachers, expert parallelism and + non-colocated generation each leave the policy fewer, and the division is + too high rather than failing. No attribute beats a plausible wrong one. + """ + cfg = _FakeMasterConfig() + cfg.cluster = {"num_nodes": 2, "gpus_per_node": 8} + assert "nv.dl.topology.size.dp" not in _build_resource_attributes(cfg, "grpo") + + def test_parallelism_sizes_use_the_resource_scoped_spelling(): """``nv.dl.topology.size.*``, not the ``dl.tensor_parallel.*`` family. Both spell the same quantity, but only the topology one is in lens's - ``DL_RESOURCE_TYPES`` -- the registry ``normalize_resource_attributes`` - types a resource attribute against. The ``dl.*`` spelling is the span and - metric one, so setting it here puts the value outside that typing. + ``DL_RESOURCE_TYPES``, which feeds the table + ``normalize_resource_attributes`` consults. Keys outside it are passed + through untyped rather than rejected, so the ``dl.*`` spelling would reach + the backend as a string -- it is the span and metric spelling. """ pytest.importorskip("nemo.lens") from nemo.lens.semconv.resources import DL_RESOURCE_TYPES @@ -111,38 +158,6 @@ def test_parallelism_sizes_use_the_resource_scoped_spelling(): assert "dl.tensor_parallel.size" not in DL_RESOURCE_TYPES -def test_data_parallel_size_is_derived_from_the_cluster(): - # RL configures no dp size: training spans the cluster's GPUs and dp is - # whatever tp/pp/cp leave over. - cfg = {"cluster": {"num_nodes": 2, "gpus_per_node": 8}} - assert _data_parallel_size(cfg, tp=2, pp=1, cp=1) == 8 - assert _data_parallel_size(cfg, tp=2, pp=2, cp=2) == 2 - - -def test_data_parallel_size_is_omitted_for_dedicated_generation_gpus(): - """The cluster then holds ranks the policy does not have. - - Reporting the cluster's arithmetic would describe a training group larger - than the one that exists, so this reports nothing instead. - """ - cfg = { - "cluster": {"num_nodes": 2, "gpus_per_node": 8}, - "policy": {"generation": {"colocated": {"enabled": False}}}, - } - assert _data_parallel_size(cfg, tp=1, pp=1, cp=1) is None - - -def test_data_parallel_size_is_omitted_when_it_does_not_divide_evenly(): - # A misconfiguration the policy worker rejects with a far better message - # than a resource attribute could carry. - cfg = {"cluster": {"num_nodes": 1, "gpus_per_node": 8}} - assert _data_parallel_size(cfg, tp=3, pp=1, cp=1) is None - - -def test_data_parallel_size_is_omitted_without_a_cluster_block(): - assert _data_parallel_size({}, tp=1, pp=1, cp=1) is None - - def test_training_target_drops_a_non_positive_step_count(): """``rm`` spells "train for one epoch" as ``max_num_steps: -1``. @@ -156,9 +171,25 @@ def test_training_target_drops_a_non_positive_step_count(): assert attrs["nv.dl.training.target.train_iters"] == 150 +def test_sft_v2_finds_its_step_target_under_the_sft_section(): + """``run_sft_v2.py`` reports ``sft_v2`` but reads ``master_config.sft``. + + Looking the settings up by algorithm name finds no section at all, so the + whole of sft_v2 silently reports no training target. + """ + cfg = {"sft": {"max_num_steps": 60}} + attrs = _training_config_attributes(cfg, "sft_v2") + assert attrs["nv.dl.training.target.train_iters"] == 60 + + def test_training_config_reads_either_backends_optimizer_spelling(): - # megatron names the algorithm, automodel names the class. - megatron = {"policy": {"megatron_cfg": {"optimizer": {"optimizer": "adam"}}}} + # megatron names the algorithm and keeps it in its own block, automodel + # names the class and keeps it in the policy's. + megatron = { + "policy": { + "megatron_cfg": {"enabled": True, "optimizer": {"optimizer": "adam"}} + } + } automodel = {"policy": {"optimizer": {"name": "torch.optim.AdamW"}}} key = "nv.dl.training.config.optimizer" assert _training_config_attributes(megatron, "grpo")[key] == "adam" @@ -172,6 +203,7 @@ def test_recompute_granularity_is_reported_only_when_recompute_is_on(): off = { "policy": { "megatron_cfg": { + "enabled": True, "activation_checkpointing": False, "recompute_granularity": "full", } @@ -182,6 +214,7 @@ def test_recompute_granularity_is_reported_only_when_recompute_is_on(): on = { "policy": { "megatron_cfg": { + "enabled": True, "activation_checkpointing": True, "recompute_granularity": "full", } @@ -195,10 +228,10 @@ def test_process_attributes_report_the_local_rank_when_the_launcher_set_one( ): """``RayWorkerGroup`` exports ``LOCAL_RANK`` beside ``RANK``. - It identifies a device on one node, which is what a reader has when the - complaint came from ``nvidia-smi`` rather than from the job. The driver owns - no device, so a process without the variable reports no local rank rather - than defaulting to zero and claiming one. + It places a process within its own node, which is what groups a job's + processes by the host they share. The driver owns no device, so a process + without the variable reports no local rank rather than defaulting to zero + and claiming one. """ pytest.importorskip("nemo.lens") monkeypatch.setenv("LOCAL_RANK", "3") @@ -209,14 +242,20 @@ def test_process_attributes_report_the_local_rank_when_the_launcher_set_one( def test_process_attributes_name_the_provider_and_the_run(monkeypatch): - # nv.dl.run.uuid duplicates lens's own nemo.run.id by design: the two - # namespaces are read by different consumers, and a backend joining RL runs - # to Megatron ones groups on the nv.dl.* spelling. + """``nv.dl.run.uuid`` is lens's derived uuid5, not the raw run id. + + The raw id is already reported as ``nemo.run.id``. A backend grouping RL + runs with Megatron ones joins on the derived value, so publishing the run + id under this key would look populated and match nothing. + """ pytest.importorskip("nemo.lens") + from nemo.lens.resources.slurm import derive_nv_dl_run_uuid + monkeypatch.delenv("LOCAL_RANK", raising=False) attrs = _process_attributes(rank=0, world_size=1, run_id="job-7") assert attrs["nv.dl.provider.name"] == "nemo-rl" - assert attrs["nv.dl.run.uuid"] == "job-7" + assert attrs["nv.dl.run.uuid"] == derive_nv_dl_run_uuid(run_id="job-7") + assert attrs["nv.dl.run.uuid"] != "job-7" assert "nv.dl.run.uuid" not in _process_attributes(rank=0, world_size=1) From bfbaf9fe1dbcfdf4435b91a868760a3dd1bf8d3c Mon Sep 17 00:00:00 2001 From: Raj Singh Date: Thu, 8 Oct 2026 16:11:40 -0700 Subject: [PATCH 06/17] docs(telemetry): describe nv.dl.local_rank as a bundle index RayWorkerGroup sets LOCAL_RANK from the worker's bundle index in its placement group, not from a per-node counter. That is node-local only when each node has its own placement group, and a generation backend spanning nodes asks for a single unified one instead -- SGLang always, vLLM, TRT-LLM and Megatron when their parallelism crosses nodes. The unified group's bundles are summed across nodes, so the index runs across all of them and the previous "within its own node" wording would have had a reader grouping processes by a host that is not implied. Signed-off-by: Raj Singh --- docs/observability/configuration.md | 2 +- nemo_rl/telemetry/setup.py | 13 ++++++++----- tests/unit/telemetry/test_setup.py | 8 ++++---- 3 files changed, 13 insertions(+), 10 deletions(-) diff --git a/docs/observability/configuration.md b/docs/observability/configuration.md index f20ffab845d..c9723b7fcd4 100644 --- a/docs/observability/configuration.md +++ b/docs/observability/configuration.md @@ -103,7 +103,7 @@ Each process sets stable-for-the-run values on its OTel `Resource` — `init_tel | `nv.dl.software.torch`, `.cuda`, `.nccl` | versions as loaded in this process, not as pinned in the lockfile. CUDA and NCCL are omitted on a CPU-only build | | `nv.dl.software.transformer_engine` | the installed distribution's version, read from package metadata rather than by importing it — importing transformer-engine loads its PyTorch extension in every process that sets up telemetry. Omitted when it is not installed | | `nv.dl.rank`, `nv.dl.world_size` | this process's rank and group size (`RANK` / `WORLD_SIZE`, or `0` / `1` for the driver and singleton actors) | -| `nv.dl.local_rank` | worker processes only: `LOCAL_RANK`, which `RayWorkerGroup` exports beside `RANK`. Places a process within its own node, which is how you group a job's processes by the host they share. Not a physical device index — Ray gives each worker its own `CUDA_VISIBLE_DEVICES` | +| `nv.dl.local_rank` | worker processes only: `LOCAL_RANK`, which `RayWorkerGroup` exports beside `RANK` as the worker's bundle index in its placement group. Node-local only when each node has its own placement group — a generation backend spanning nodes requests a single unified one (SGLang always; vLLM, TRT-LLM and Megatron when their parallelism crosses nodes) and the index then runs across nodes. Not a physical device index either, since Ray gives each worker its own `CUDA_VISIBLE_DEVICES` | | `rl.worker_group` | worker processes only: the worker group's `name_prefix` (`lm_policy`, `vllm_policy`, ...), from `NRL_WORKER_GROUP`. The `NemoGym` actor reports the literal `nemo_gym` — it is built from the environment registry rather than by `RayWorkerGroup`, so nothing sets the env var for it | The config-derived rows are driver-only: a worker is handed the propagated environment, not the `MasterConfig`. Attribute construction is best-effort: a missing config key simply omits that attribute; it never raises. Plus auto-detected host / GPU / SLURM / Kubernetes attributes from lens's resource detection. diff --git a/nemo_rl/telemetry/setup.py b/nemo_rl/telemetry/setup.py index 8188a67df70..d4bb4d17dd8 100644 --- a/nemo_rl/telemetry/setup.py +++ b/nemo_rl/telemetry/setup.py @@ -407,11 +407,14 @@ def _process_attributes(rank: int, world_size: int, run_id: str = "") -> dict[st 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. ``nv.dl.local_rank`` adds its position - within its own node, which is what groups a job's processes by the host - they share. It is not a physical device index: Ray gives each worker its - own ``CUDA_VISIBLE_DEVICES``, so local rank 0 is whichever GPU that worker - was placed on. + apart from its peers downstream. ``nv.dl.local_rank`` adds the worker's + bundle index within its placement group, which ``RayWorkerGroup`` exports + as ``LOCAL_RANK``. That is node-local only when each node has its own + placement group: a generation backend that spans nodes asks for a single + unified one instead -- SGLang always, vLLM, TRT-LLM and Megatron when + their parallelism crosses nodes -- and the index then runs across every + node in the group. Nor is it a physical device index, since Ray gives each + worker its own ``CUDA_VISIBLE_DEVICES``. Called by both entry points, since a resource is per process and neither path sees the other's. diff --git a/tests/unit/telemetry/test_setup.py b/tests/unit/telemetry/test_setup.py index 33206ba86db..5b5ec3d10b7 100644 --- a/tests/unit/telemetry/test_setup.py +++ b/tests/unit/telemetry/test_setup.py @@ -228,10 +228,10 @@ def test_process_attributes_report_the_local_rank_when_the_launcher_set_one( ): """``RayWorkerGroup`` exports ``LOCAL_RANK`` beside ``RANK``. - It places a process within its own node, which is what groups a job's - processes by the host they share. The driver owns no device, so a process - without the variable reports no local rank rather than defaulting to zero - and claiming one. + It is the worker's bundle index in its placement group, which is node-local + only when each node has one of its own. The driver owns no device, so a + process without the variable reports no local rank rather than defaulting + to zero and claiming one. """ pytest.importorskip("nemo.lens") monkeypatch.setenv("LOCAL_RANK", "3") From eccaf46ea763fe610df791356a1c9bb80871bd61 Mon Sep 17 00:00:00 2001 From: Raj Singh Date: Thu, 8 Oct 2026 16:13:52 -0700 Subject: [PATCH 07/17] docs(telemetry): correct the second always-RL campaign stage claim The resource-attribute summary in span-groups.md carried its own copy of the claim, which the configuration.md fix did not reach. sft and sft_v2 report SFT and rm reports RM. Also take derive_nv_dl_run_uuid from nemo.lens rather than reaching into nemo.lens.resources.slurm for it. The function is re-exported at the top level and listed in __all__, so the public path is the one that survives lens moving it. Signed-off-by: Raj Singh --- docs/observability/span-groups.md | 2 +- nemo_rl/telemetry/setup.py | 2 +- tests/unit/telemetry/test_setup.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/observability/span-groups.md b/docs/observability/span-groups.md index be20c1e0382..c0191126929 100644 --- a/docs/observability/span-groups.md +++ b/docs/observability/span-groups.md @@ -671,7 +671,7 @@ blanks today, so an empty trace is not read as a broken exporter: ## Resource attributes (process tags) -Stable-for-the-run values, set once at init and attached to every span/metric: `nv.dl.campaign.stage` (always `RL`), `rl.algorithm`, `rl.model`, `nemo.precision`, `nv.dl.provider.name`, the `nv.dl.topology.size.*` parallelism sizes, the `nv.dl.training.config.*` / `nv.dl.training.target.*` settings, the `nv.dl.software.*` versions, plus `nv.dl.rank` / `nv.dl.world_size` / `nv.dl.local_rank`. See [Configuration — Resource attributes](configuration.md#resource-attributes) for the full table and where each value comes from. +Stable-for-the-run values, set once at init and attached to every span/metric: `nv.dl.campaign.stage` (`SFT`, `RM` or `RL`, depending on the algorithm), `rl.algorithm`, `rl.model`, `nemo.precision`, `nv.dl.provider.name`, the `nv.dl.topology.size.*` parallelism sizes, the `nv.dl.training.config.*` / `nv.dl.training.target.*` settings, the `nv.dl.software.*` versions, plus `nv.dl.rank` / `nv.dl.world_size` / `nv.dl.local_rank`. See [Configuration — Resource attributes](configuration.md#resource-attributes) for the full table and where each value comes from. ## Granularity guidance diff --git a/nemo_rl/telemetry/setup.py b/nemo_rl/telemetry/setup.py index d4bb4d17dd8..a8f2cf49dc2 100644 --- a/nemo_rl/telemetry/setup.py +++ b/nemo_rl/telemetry/setup.py @@ -419,7 +419,7 @@ def _process_attributes(rank: int, world_size: int, run_id: str = "") -> dict[st Called by both entry points, since a resource is per process and neither path sees the other's. """ - from nemo.lens.resources.slurm import derive_nv_dl_run_uuid + from nemo.lens import derive_nv_dl_run_uuid from nemo.lens.semconv import ( NV_DL_LOCAL_RANK, NV_DL_PROVIDER_NAME, diff --git a/tests/unit/telemetry/test_setup.py b/tests/unit/telemetry/test_setup.py index 5b5ec3d10b7..c923c8137ed 100644 --- a/tests/unit/telemetry/test_setup.py +++ b/tests/unit/telemetry/test_setup.py @@ -249,7 +249,7 @@ def test_process_attributes_name_the_provider_and_the_run(monkeypatch): id under this key would look populated and match nothing. """ pytest.importorskip("nemo.lens") - from nemo.lens.resources.slurm import derive_nv_dl_run_uuid + from nemo.lens import derive_nv_dl_run_uuid monkeypatch.delenv("LOCAL_RANK", raising=False) attrs = _process_attributes(rank=0, world_size=1, run_id="job-7") From 0a26cee611402e703ca10d865da1b1d351610670 Mon Sep 17 00:00:00 2001 From: Raj Singh Date: Thu, 8 Oct 2026 16:57:54 -0700 Subject: [PATCH 08/17] refactor(telemetry): keep checkpoint renaming exactly as it was Reading the step for the finalize span had grown into a second parser for tmp_step_N, and pointing _rename_checkpoint at it changed how a checkpoint gets published: a name the regex rejected raised instead of publishing under whatever sat between the second and third underscore. That is a behaviour change to the checkpoint path in a telemetry PR, which is not where a reviewer looks for one. Drop the new parser. _rename_checkpoint is byte-identical to main again, and begin_finalization reads the step the same way it does, guarded and dropped to None on anything unexpected because the value only labels a span. The pre-existing quirk -- tmp_step_final publishes as step_final -- is left alone for a checkpoint-scoped change to fix. Signed-off-by: Raj Singh --- nemo_rl/utils/checkpoint.py | 38 ++++++++------------------- tests/unit/utils/test_checkpoint.py | 40 ++++++++++++----------------- 2 files changed, 27 insertions(+), 51 deletions(-) diff --git a/nemo_rl/utils/checkpoint.py b/nemo_rl/utils/checkpoint.py index 18010d92331..912b8068506 100644 --- a/nemo_rl/utils/checkpoint.py +++ b/nemo_rl/utils/checkpoint.py @@ -53,22 +53,6 @@ PathLike = Union[str, "os.PathLike[Any]"] -# ``init_tmp_checkpoint`` names the staging directory ``tmp_step_{step}``, and -# that name is the only place the step survives into ``begin_finalization``, -# which takes the path alone. Parsed rather than added as an argument, so every -# algorithm's call site stays untouched. -_TMP_STEP_PATTERN = re.compile(r"^tmp_step_(\d+)$") - - -def _step_from_tmp_path(checkpoint_path: PathLike) -> Optional[int]: - """Step a ``tmp_step_N`` directory belongs to, or None if unrecognised. - - None rather than a raise: one caller only labels a span with this, and a - checkpoint staged under some other name should still be finalized. - """ - match = _TMP_STEP_PATTERN.match(Path(checkpoint_path).name) - return int(match.group(1)) if match else None - @contextmanager def _finalize_telemetry( @@ -434,18 +418,9 @@ def _rename_checkpoint(self, checkpoint_path: PathLike) -> None: If step_N already exists (defensive guard for edge cases, e.g. resuming training), performs a pseudo-atomic swap via an intermediate old_step_N directory. - - Raises: - ValueError: *checkpoint_path* is not a ``tmp_step_N`` directory, so - there is no step to publish it under. """ checkpoint_path = Path(checkpoint_path) - step = _step_from_tmp_path(checkpoint_path) - if step is None: - raise ValueError( - f"Expected a tmp_step_ directory from init_tmp_checkpoint(), " - f"got {checkpoint_path.name!r}" - ) + step = checkpoint_path.name.split("_")[2] to_checkpoint_path = checkpoint_path.parent / f"step_{step}" if to_checkpoint_path.exists(): old_checkpoint_path = checkpoint_path.parent / f"old_step_{step}" @@ -491,7 +466,16 @@ def begin_finalization( self.finalize_pending() self._pending_checkpoint_path = Path(checkpoint_path) self._finalize_error = None - step = _step_from_tmp_path(checkpoint_path) + # The step reaches here only in the staging directory's name, which + # ``init_tmp_checkpoint`` formats as ``tmp_step_{step}``; read back the + # same way :meth:`_rename_checkpoint` reads it, rather than adding a + # second spelling of the same parse to this file. Guarded and dropped + # to None on anything unexpected, since this only labels a span and a + # checkpoint staged under some other name must still be finalized. + try: + step = int(Path(checkpoint_path).name.split("_")[2]) + except (IndexError, ValueError): + step = None # Captured here because a bare Thread starts with an empty context: read # it on the worker instead and the finalize span re-roots into a trace of # its own, detached from the run that produced the checkpoint. diff --git a/tests/unit/utils/test_checkpoint.py b/tests/unit/utils/test_checkpoint.py index 6a6a565e081..e10de34c005 100644 --- a/tests/unit/utils/test_checkpoint.py +++ b/tests/unit/utils/test_checkpoint.py @@ -1200,23 +1200,6 @@ def test_swaps_when_target_exists(self, checkpoint_manager, checkpoint_dir): with open(step_dir / "training_info.json") as f: assert json.load(f)["loss"] == 0.2 - def test_rejects_a_directory_it_cannot_read_a_step_from( - self, checkpoint_manager, checkpoint_dir - ): - """The step comes out of the name, so an unreadable name has to raise. - - The previous spelling of that read, ``name.split("_")[2]``, returned - whatever sat between the second and third underscore -- publishing - ``tmp_weights_12_a`` as ``step_12``, a directory claiming to be a - checkpoint of a step it has nothing to do with. - """ - staged = checkpoint_dir / "tmp_weights_12_a" - staged.mkdir(parents=True) - - with pytest.raises(ValueError, match="Expected a tmp_step_ directory"): - checkpoint_manager._rename_checkpoint(staged) - - assert staged.exists() assert not (checkpoint_dir / "step_12").exists() @@ -1449,12 +1432,21 @@ def test_metric_topk_with_ft_keeps_best_and_recent_disjoint(self, checkpoint_dir assert self._remaining_steps(checkpoint_dir) == [1, 2, 5, 6] -def test_step_from_tmp_path_tolerates_an_unexpected_name(): - """Reading the step is only used to label a span, so it must never raise. +def test_a_non_numeric_staging_name_costs_only_the_span_label( + checkpoint_manager, checkpoint_dir +): + """The span's step is read from the directory name, so it must never raise. - A caller that stages under some other name still gets a finalized - checkpoint; it just loses the iteration attribute on the span. + ``_rename_checkpoint`` publishes whatever sits between the second and third + underscore, so ``tmp_step_final`` becomes ``step_final``. The span wants an + int and cannot have one here; it drops the label rather than failing the + finalization that carries it. """ - assert checkpoint_module._step_from_tmp_path("/ckpt/tmp_step_412") == 412 - assert checkpoint_module._step_from_tmp_path("/ckpt/step_412") is None - assert checkpoint_module._step_from_tmp_path("/ckpt/tmp_step_final") is None + staged = checkpoint_dir / "tmp_step_final" + staged.mkdir(parents=True) + + checkpoint_manager.begin_finalization(staged, wait_fn=None) + checkpoint_manager.finalize_pending() + + assert (checkpoint_dir / "step_final").exists() + assert not staged.exists() From c8aa4190bbbdc880608c27df2657a55d53cfa207 Mon Sep 17 00:00:00 2001 From: Raj Singh Date: Fri, 9 Oct 2026 17:30:46 -0700 Subject: [PATCH 09/17] feat(telemetry): export per-node host memory from Ray to OTel Every raylet publishes ray_node_mem_used / ray_node_mem_total on its own Prometheus endpoint, and RayGpuMonitorLogger already reads them for W&B. Nothing carried them into OTel, so a run's telemetry had GPU-side memory and no host-side memory at all. HostMemoryMonitor scrapes the same endpoints and records rl.host.memory.used / rl.host.memory.total (bytes, not the GB the W&B path reports), one point per live node. Scraped separately rather than teed out of the GPU monitor: that path flattens the node's identity into the metric name (node.3.mem_gb) and discards the hostname, and it is gated on logger.monitor_gpus, which would make an OTel series disappear when a W&B feature is turned off. host.name is a metric attribute here rather than the resource attribute of the same name, because the driver reports about hosts other than its own. rl.worker_groups names the groups resident on the host, resolved through placement_group_table as each RayWorkerGroup is built. Plural and comma-joined on purpose: it labels whole-host memory, so it describes the box rather than partitioning the value. The monitor is owned by init_telemetry_driver / shutdown_telemetry rather than by Logger. That placement is load-bearing: the group-to-node map is a module global, visible only to the process that wrote it, and RayWorkerGroup writes it from the driver. Single-controller GRPO and SFT v2 build their worker groups on the driver but their Logger inside the controller actor, so a Logger-owned monitor would read an empty map and label every host as holding no group -- right numbers, wrong label, on every point. Also drops a stale assertion left in test_swaps_when_target_exists by the checkpoint-rename revert in 5551a17c3, where it is vacuously true. Signed-off-by: Raj Singh --- docs/observability/metrics.md | 42 ++ nemo_rl/distributed/worker_groups.py | 15 + nemo_rl/telemetry/host_memory.py | 545 +++++++++++++++++++++++ nemo_rl/telemetry/metrics.py | 41 +- nemo_rl/telemetry/setup.py | 36 ++ nemo_rl/telemetry/vocabulary.py | 28 +- pyproject.toml | 14 +- pyrefly.toml | 1 + tests/unit/telemetry/test_host_memory.py | 517 +++++++++++++++++++++ tests/unit/utils/test_checkpoint.py | 2 - uv.lock | 8 +- 11 files changed, 1232 insertions(+), 17 deletions(-) create mode 100644 nemo_rl/telemetry/host_memory.py create mode 100644 tests/unit/telemetry/test_host_memory.py diff --git a/docs/observability/metrics.md b/docs/observability/metrics.md index 5fd98556b63..da98371bbcd 100644 --- a/docs/observability/metrics.md +++ b/docs/observability/metrics.md @@ -203,6 +203,48 @@ These overlap with `rl.vllm.*` on token counts but are not redundant: `gen_ai.*` These ride the normal `http/protobuf` OTLP path and reach the same backend as everything else. They are distinct from vLLM's **native** engine tracing (opt-in, and exported over vLLM's own gRPC-by-default exporter) — see [vLLM Tracing](vllm-tracing.md). +## Host memory (`rl.host.memory.*`) + +Every raylet publishes its node's memory on its own Prometheus endpoint. `HostMemoryMonitor` in `nemo_rl/telemetry/host_memory.py` scrapes all of them from the driver on a timer and records one point per node — the same source `RayGpuMonitorLogger` already reads for the W&B `ray/node.*.mem_gb` rows, but keeping the node's identity instead of flattening it into the metric name. + +| Metric | Type | Attributes | Source series | +|---|---|---|---| +| `rl.host.memory.used` | Observable gauge (`By`) | `host.name`, `rl.worker_groups` | `ray_node_mem_used` | +| `rl.host.memory.total` | Observable gauge (`By`) | `host.name`, `rl.worker_groups` | `ray_node_mem_total` | + +Bytes, not the gigabytes the W&B path reports. The interval follows `logger.gpu_monitoring.collection_interval` so the host-memory and GPU series line up, but collection is **not** gated on `logger.monitor_gpus`: that flag governs what reaches W&B, and an OTel series should not vanish when someone turns off an unrelated W&B feature. It is gated on telemetry exporting, and on nothing else. + +The monitor is started by `init_telemetry_driver` and stopped by `shutdown_telemetry`, which puts it in the **driver** process on every entry point. That placement is load-bearing rather than incidental — see `rl.worker_groups` below — which is why it does not live in `Logger` beside the GPU monitor it otherwise resembles. + +### Why these are observable rather than pushed + +A host's `rl.worker_groups` label changes during a run: it is `none` until the worker groups are built, which is minutes in, and changes again when they regroup or shut down. Pushed to an ordinary gauge that would be a permanent error, not a transient one — the SDK keeps a synchronous gauge's last value *per attribute set* and re-exports it every cycle, and nothing in the OTel API can retract an attribute set. Every host's startup-time `none` point would sit next to its correctly labelled one for the rest of the run, so `sum by (host.name)` would count each host twice and filtering on `none` would show every training host as idle. + +An observable gauge is rebuilt from its callback on each collection cycle, so a label the callback stops reporting simply ends. The monitor therefore scrapes on its own timer into a snapshot that it *replaces* wholesale, and the callbacks return whatever is in it. Three consequences worth knowing: + +- Scraping stays on the monitor's thread rather than inside the callback, so the exporter never blocks on HTTP to an unreachable node — a wait that would otherwise hold up every other metric in the process and grow with the cluster. +- The scrape and export intervals are independent, so a point can be up to one collection interval old, and scrapes taken between two export cycles are never seen. Both are normal for a gauge. +- A node that fails a scrape leaves a gap rather than a flat line at its last known value, because the snapshot is replaced and not merged. For "how much memory does this host have in use right now" a gap is the honest answer. + +The instruments are created by `create_rl_metric_instruments()` when the monitor starts. This is not optional bookkeeping: nothing ever records these series, so without that call no instrument exists and the callbacks are never polled. + +`host.name` is a metric attribute here, not the resource attribute of the same name that lens already sets. The resource one names the host the *emitting* process runs on, which for every one of these points is the head node; the driver is reporting about other hosts. It prefers Ray's `NodeManagerHostname` and falls back to `NodeManagerAddress` (the IP) when Ray resolved no hostname. + +### `rl.worker_groups` is a descriptor, not a partition key + +It lists the worker groups resident on the host, sorted and comma-joined (`lm_policy,vllm_policy`), or `none` for a host running no group — the head node, or any node whose groups had not been built when the sample was taken. `RayWorkerGroup` registers its placement as it is created, resolving bundles to nodes through `placement_group_table`, so no worker is asked anything. + +There is exactly one point per `host.name` per cycle. That matters where several raylets share a machine, as in a local multi-node cluster: each publishes the same whole-host memory on its own endpoint, so the monitor scrapes a host once and merges the worker groups of every node on it, rather than emitting a point per node that `sum by (host.name)` would add together. + +That registry is an ordinary module global, so it is visible only to the process that wrote it, and the writer is always the driver. This is why the monitor has to run there too: the single-controller entry points (`run_grpo_single_controller.py`, `run_sft_v2.py`) build their worker groups on the driver but their `Logger` inside the controller actor, so a monitor owned by the `Logger` would read an empty map and label every host `none` — right numbers, wrong label, on every point. + +**Do not aggregate on it.** The value it labels is the *whole host's* memory: + +- Under colocation `lm_policy` and `vllm_policy` share GPUs and therefore a host, so `sum by (rl.worker_groups)` would count that host's memory once per group. +- No worker group owns the raylet, the Ray **object store** (often tens of GB, and frequently the thing that actually ran a node out of memory) or — on the head node — the driver, yet all of them are in the number. + +Group on it to decide *which* hosts to look at; read `host.name` for the number itself. There is no breakdown inside `ray_node_mem_used`: getting one means per-process sampling (`psutil` RSS from inside each worker, where `rl.worker_group` and `nv.dl.rank` are already on the resource and nothing has to be inferred), and attributing a single process's RSS to tensors versus pinned buffers versus allocator arenas is a profiler's job, not a metric's. + ## Metric vs span tag vs resource attribute The one rule that trips people up. Classify each value before you emit it: diff --git a/nemo_rl/distributed/worker_groups.py b/nemo_rl/distributed/worker_groups.py index ffbf7f3ab2a..959d0554709 100644 --- a/nemo_rl/distributed/worker_groups.py +++ b/nemo_rl/distributed/worker_groups.py @@ -31,6 +31,10 @@ ) from nemo_rl.distributed.virtual_cluster import RayVirtualCluster from nemo_rl.distributed.worker_group_utils import recursive_merge_options +from nemo_rl.telemetry.host_memory import ( + forget_worker_group_nodes, + register_worker_group_placement, +) from nemo_rl.utils.venvs import ( add_hf_modules_cache_to_pythonpath, create_local_venv_on_each_node, @@ -766,6 +770,14 @@ def _create_workers_from_bundle_indices( } ) + # Here rather than anywhere in telemetry: the host-memory monitor + # scrapes raylets, which know about nodes and nothing about worker + # groups, so the only process that can relate the two is this one. + register_worker_group_placement( + self.name_prefix, + [(spec["pg"], spec["bundle_idx"]) for spec in self._worker_specs.values()], + ) + def log_worker_gpu_state( self, worker_idx: int, *, label: str, timeout_s: float = 30.0 ) -> None: @@ -1320,5 +1332,8 @@ def shutdown( # Clear worker lists self._workers = [] self._worker_metadata = [] + # So the host-memory monitor stops attributing nodes to a group that no + # longer occupies them -- a teacher group torn down mid-run, say. + forget_worker_group_nodes(self.name_prefix) return success diff --git a/nemo_rl/telemetry/host_memory.py b/nemo_rl/telemetry/host_memory.py new file mode 100644 index 00000000000..138f9610b56 --- /dev/null +++ b/nemo_rl/telemetry/host_memory.py @@ -0,0 +1,545 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Host memory per node, scraped from Ray and exported through lens. + +Same source as the W&B path: every raylet publishes ``ray_node_mem_used`` and +``ray_node_mem_total`` on its own Prometheus endpoint, and +``RayGpuMonitorLogger`` in :mod:`nemo_rl.utils.logger` already reads them. +Nothing else in NeMo-RL reports host memory -- ``utils/memory_tracker.py`` +samples the driver's own RSS, prints it, and is wired only into PPO -- so +without this a run's OTel stream has GPU-side numbers and no host-side ones. + +Scraped here rather than teed out of the GPU monitor because that path flattens +the node's identity into the metric *name* (``node.3.mem_gb``) and discards the +hostname, and because it is gated on ``logger.monitor_gpus``: reading it would +make an OTel series disappear when someone turns off a W&B feature. The cost is +a second HTTP GET per node per interval, against a local endpoint, which is not +measurable next to a training step. + +What the number is, and is not: ``ray_node_mem_used`` is whole-host memory. It +covers the raylet, the Ray object store (often tens of GB, and frequently the +thing that actually ran the node out of memory), every worker resident on the +host, and -- on the head node -- the driver. It is not a per-process or +per-worker-group figure, and there is no breakdown inside it; see +:data:`RL_WORKER_GROUPS_ATTR` for what the worker-group attribute can and +cannot be used for. +""" + +from __future__ import annotations + +import logging +import threading +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Iterable, Optional + +import requests +from prometheus_client.parser import text_string_to_metric_families + +from nemo_rl.telemetry.metrics import create_rl_metric_instruments, warn_once +from nemo_rl.telemetry.setup import get_telemetry_handle +from nemo_rl.telemetry.vocabulary import ( + RecordedMetric, + register_recorded_metrics, + registry_key, +) + +if TYPE_CHECKING: + from opentelemetry.metrics import CallbackOptions, Observation + +logger = logging.getLogger(__name__) + +#: Resident and installed memory on one host, in bytes. +RL_HOST_MEMORY_USED_METRIC = "rl.host.memory.used" +RL_HOST_MEMORY_TOTAL_METRIC = "rl.host.memory.total" + +#: Which host a point describes. OTel's ``host.name`` is normally a *resource* +#: attribute naming the host a process runs on, and lens already sets it that +#: way. Here it is a metric attribute instead, because the driver is reporting +#: about hosts other than its own -- the resource would say "the head node" for +#: every point. +HOST_NAME_ATTR = "host.name" + +#: Worker groups resident on the host, sorted and comma-joined. +#: +#: A descriptor of the host, not a partition key. The value it labels is the +#: whole host's memory, so ``sum by (rl.worker_groups)`` double-counts under +#: colocation (where ``lm_policy`` and ``vllm_policy`` share GPUs and therefore +#: a host) and over-counts always, since no worker group owns the raylet, the +#: object store or the driver. Plural, and joined rather than emitted as one +#: point per group, so that this reads as "what is on this box" rather than as +#: something summable. Group it to find *which* hosts to look at; read +#: ``host.name`` for the number itself. +RL_WORKER_GROUPS_ATTR = "rl.worker_groups" + +#: Value of :data:`RL_WORKER_GROUPS_ATTR` for a host running no worker group: +#: a head node that only hosts the driver, or any node whose groups have not +#: been built yet when the sample is taken. +NO_WORKER_GROUPS = "none" + +#: Prometheus series published by each raylet, and the key each is recorded +#: against. Bytes already, despite the GB the W&B path reports: that divides by +#: 1024**3 on the way out. +_SERIES_BY_RAY_NAME = { + "ray_node_mem_used": RL_HOST_MEMORY_USED_METRIC, + "ray_node_mem_total": RL_HOST_MEMORY_TOTAL_METRIC, +} + +_USED_KEY = registry_key(RL_HOST_MEMORY_USED_METRIC) +_TOTAL_KEY = registry_key(RL_HOST_MEMORY_TOTAL_METRIC) + +#: The most recent scrape, keyed by the attributes the points carry. +#: +#: Replaced wholesale each cycle rather than updated, which is what makes a +#: host's label stop being exported when it changes: the exporter's points are +#: exactly what the callbacks below return, so a key that is no longer in here +#: is a series that ends. Read under :data:`_SNAPSHOT_LOCK` from the exporter's +#: thread and written from the monitor's. +_SNAPSHOT: dict[tuple[str, str], dict[str, float]] = {} +_SNAPSHOT_LOCK = threading.Lock() + + +def _observations(key: str) -> list[Observation]: + """One observation per host in the latest scrape, for registry *key*.""" + # Deferred: this module is imported to declare its rows, which must not + # require OTel to be installed. By the time a callback runs, it is. + from opentelemetry.metrics import Observation + + with _SNAPSHOT_LOCK: + latest = list(_SNAPSHOT.items()) + return [ + Observation( + values[key], + {HOST_NAME_ATTR: host, RL_WORKER_GROUPS_ATTR: groups}, + ) + for (host, groups), values in latest + if key in values + ] + + +# Two callbacks rather than one parameterised closure: lens holds the function +# for the life of the process, so each has to be a stable module-level object. +def _observe_used(options: CallbackOptions) -> list[Observation]: + return _observations(_USED_KEY) + + +def _observe_total(options: CallbackOptions) -> list[Observation]: + return _observations(_TOTAL_KEY) + + +#: Observable rather than pushed, because the attributes move: a host is +#: labelled ``rl.worker_groups="none"`` until its groups are built, and the +#: label changes again when they shut down. A synchronous gauge keeps its last +#: value per attribute set and re-exports it every cycle for the rest of the +#: run, so the startup-time point would sit next to the correct one forever and +#: ``sum by (host.name)`` would count every host twice. An observable gauge is +#: rebuilt from its callback each cycle, so the old label simply ends. +HOST_MEMORY_RECORDED_METRICS = ( + RecordedMetric( + RL_HOST_MEMORY_USED_METRIC, + kind="observable_gauge", + unit="By", + description=( + "Memory in use on one host, across every process on it. Labelled by " + "host.name; see rl.worker_groups before aggregating." + ), + callback=_observe_used, + ), + RecordedMetric( + RL_HOST_MEMORY_TOTAL_METRIC, + kind="observable_gauge", + unit="By", + description="Memory installed on one host.", + callback=_observe_total, + ), +) + +register_recorded_metrics(HOST_MEMORY_RECORDED_METRICS) + +#: Matches ``GPUMonitoringConfig.collection_interval``, so the host-memory and +#: GPU series line up on a dashboard rather than interleaving. +DEFAULT_COLLECTION_INTERVAL = 10.0 + +#: Short, because the endpoint is a raylet on the cluster's own network and the +#: loop would rather miss a sample than hold the interval open behind a node +#: that is already wedged. +_SCRAPE_TIMEOUT_S = 5.0 + +#: How long :meth:`HostMemoryMonitor.stop` waits for the loop to notice. +#: +#: Twice the scrape budget because ``requests`` applies its timeout to the +#: connect and the read separately, plus a margin to parse the response it has +#: already fetched. Deliberately not a function of the collection interval: the +#: loop is either asleep on an :class:`threading.Event` that ``stop`` has +#: already set, or inside one node's scrape, and it checks between nodes. +_SHUTDOWN_GRACE_S = _SCRAPE_TIMEOUT_S * 2 + 1.0 + + +# --- which worker groups sit on which node -------------------------------- +# +# Filled by RayWorkerGroup as each group is built, because the mapping is not +# knowable when the monitor starts: the driver initialises telemetry before +# init_ray(), and the groups are placed later still. A node missing from here +# is reported as NO_WORKER_GROUPS rather than skipped -- its memory is just as +# real, and the head node never appears here at all. +_WORKER_GROUP_NODES: dict[str, frozenset[str]] = {} +_WORKER_GROUP_LOCK = threading.Lock() + + +def register_worker_group_nodes(name_prefix: str, node_ids: Iterable[str]) -> None: + """Record that worker group *name_prefix* occupies *node_ids*. + + Args: + name_prefix: The group's ``RayWorkerGroup.name_prefix`` -- the same + value its workers carry as the ``rl.worker_group`` resource + attribute, so a reader can join the two. + node_ids: Ray ``NodeID`` hex strings the group's workers are placed on. + + Replaces any previous entry for the group, so a group rebuilt onto different + nodes after a worker restart is not reported on both sets. + """ + with _WORKER_GROUP_LOCK: + _WORKER_GROUP_NODES[name_prefix] = frozenset(node_ids) + + +def register_worker_group_placement( + name_prefix: str, bundles: Iterable[tuple[Any, int]] +) -> None: + """Resolve *bundles* to nodes and register them. Never raises. + + The form :class:`~nemo_rl.distributed.worker_groups.RayWorkerGroup` can + supply directly: it knows each worker's ``(placement group, bundle index)`` + but not the node underneath, and ``placement_group_table`` answers that from + the GCS's own scheduling record -- driver-side, with no worker asked + anything. Each placement group is looked up once however many bundles it + contributes, since the table covers all of them. + + Args: + name_prefix: See :func:`register_worker_group_nodes`. + bundles: ``(placement_group, bundle_index)`` pairs the group occupies. + """ + # Deferred with the rest of Ray: see ray_nodes(). + from ray.util.placement_group import placement_group_table + + try: + node_ids = set() + tables: dict[int, dict[int, str]] = {} + for pg, bundle_index in bundles: + table = tables.get(id(pg)) + if table is None: + table = placement_group_table(pg)["bundles_to_node_id"] + tables[id(pg)] = table + node_id = table.get(bundle_index) + if node_id: + node_ids.add(node_id) + except Exception: + warn_once( + f"host_memory_placement:{name_prefix}", + f"could not resolve which nodes worker group {name_prefix!r} is on; " + "its host memory will be reported without a worker-group label", + ) + else: + register_worker_group_nodes(name_prefix, node_ids) + + +def forget_worker_group_nodes(name_prefix: str) -> None: + """Drop *name_prefix*'s placement, for a group that has been shut down.""" + with _WORKER_GROUP_LOCK: + _WORKER_GROUP_NODES.pop(name_prefix, None) + + +def worker_group_names_on_node(node_id: str) -> set[str]: + """Worker groups with at least one worker on *node_id*, unformatted.""" + with _WORKER_GROUP_LOCK: + return { + prefix + for prefix, node_ids in _WORKER_GROUP_NODES.items() + if node_id in node_ids + } + + +def format_worker_groups(names: Iterable[str]) -> str: + """Render *names* as :data:`RL_WORKER_GROUPS_ATTR` carries them.""" + return ",".join(sorted(names)) or NO_WORKER_GROUPS + + +def worker_groups_on_node(node_id: str) -> str: + """Value of :data:`RL_WORKER_GROUPS_ATTR` for *node_id*.""" + return format_worker_groups(worker_group_names_on_node(node_id)) + + +@dataclass(frozen=True) +class RayNode: + """One live Ray node: who it is, and where to scrape it.""" + + node_id: str + host: str + metrics_url: str + + +def ray_nodes() -> list[RayNode]: + """Live Ray nodes that publish a metrics endpoint. + + Returns an empty list when Ray is not up, which makes the collection loop + idle rather than fail: the driver starts telemetry before ``init_ray()``. + """ + # Deferred: this module is imported from nemo_rl.utils.logger at module + # scope, and must stay importable in a process that never initialises Ray. + import ray + + if not ray.is_initialized(): + return [] + + nodes = [] + for node in ray.nodes(): + # "Alive" because ray.nodes() keeps dead entries, whose endpoints would + # be scraped every interval until they time out. + if not node.get("Alive"): + continue + port = node.get("MetricsExportPort") + address = node.get("NodeManagerAddress") + if not port or not address: + continue + nodes.append( + RayNode( + node_id=node.get("NodeID", ""), + # NodeManagerAddress is the IP; the hostname is a sibling key. + # Preferred because it is what a user recognises and what the + # cluster's other tooling reports, and it falls back to the IP + # rather than going unlabelled when Ray has not resolved one. + host=node.get("NodeManagerHostname") or address, + metrics_url=f"http://{address}:{port}/metrics", + ) + ) + return nodes + + +def parse_memory_samples(metrics_text: str) -> dict[str, float]: + """Extract ``{registry key: bytes}`` from one node's Prometheus payload. + + Pure function (no OTel side effects) so it is trivially unit-testable. + Unknown series are ignored, which is nearly all of them: a raylet publishes + hundreds, and only the two in :data:`_SERIES_BY_RAY_NAME` are wanted. + """ + values: dict[str, float] = {} + for family in text_string_to_metric_families(metrics_text): + for sample in family.samples: + name = _SERIES_BY_RAY_NAME.get(sample.name) + if name is None: + continue + values[registry_key(name)] = float(sample.value) + return values + + +class HostMemoryMonitor: + """Samples every Ray node's host memory on a timer, for lens to export. + + Driver-side and read-only: it scrapes endpoints the cluster already serves + and touches no worker. Start it once the cluster is up; it tolerates being + started earlier and simply finds no nodes until there are some. + + It publishes into :data:`_SNAPSHOT` rather than recording: the series are + observable, so the exporter pulls the latest sample on its own schedule. + The two intervals are independent, which means a sample can be up to one + collection interval old when it is exported, and that samples taken + between two export cycles are never seen. Both are properties of a gauge + and neither is worth coupling the threads to avoid. + """ + + def __init__( + self, collection_interval: float = DEFAULT_COLLECTION_INTERVAL + ) -> None: + """Initialise the monitor. + + Args: + collection_interval: Seconds between samples of every node. + """ + self.collection_interval = collection_interval + self._stop = threading.Event() + self._thread: Optional[threading.Thread] = None + + def start(self) -> None: + """Begin sampling on a background thread. Idempotent.""" + if self._thread is not None: + return + self._stop.clear() + self._thread = threading.Thread( + target=self._collection_loop, + name="nemo-rl-host-memory", + # Daemon so a run that exits without calling stop() -- a crash, a + # KeyboardInterrupt -- is not held open by an observability thread. + daemon=True, + ) + self._thread.start() + logger.info( + "host memory monitoring started (interval=%ss)", self.collection_interval + ) + + def stop(self) -> None: + """Stop sampling and wait for the thread to finish. Idempotent.""" + thread = self._thread + self._stop.set() + self._thread = None + if thread is not None: + thread.join(timeout=_SHUTDOWN_GRACE_S) + + def collect_once(self) -> None: + """Sample every node and publish the result for export. Never raises. + + The guarantee lives here rather than in the loop so that the loop cannot + be killed by a single bad sample, and so a caller driving the monitor + by hand gets the same protection. + """ + try: + self._collect_once() + except Exception: + warn_once("host_memory", "failed to collect host memory from Ray") + + def _collect_once(self) -> None: + """Body of :meth:`collect_once`, inside its exception guard.""" + handle = get_telemetry_handle() + if handle is None or not handle.is_exporting: + return + # Scraped on this thread rather than inside the callbacks so that the + # exporter never blocks on HTTP: a callback that waits on an + # unreachable node would hold up every other metric in the process, and + # the wait scales with the size of the cluster. + # Accumulated per host rather than per node, because the number is the + # host's: several raylets on one box -- a local multi-node setup -- each + # publish the same whole-host memory on their own endpoint. Keyed by + # node that would be two points carrying one host.name, which double + # counts under `sum by (host.name)` and contradicts the one-point-per- + # host reading the attribute is documented to have. + values_by_host: dict[str, dict[str, float]] = {} + groups_by_host: dict[str, set[str]] = {} + for node in ray_nodes(): + # Between nodes, so stopping costs the scrape in flight rather than + # the rest of the cluster: this loop is one scrape timeout per node + # and runs them in sequence, which on a large cluster is minutes. + # Returning without publishing is the point -- a partial sweep would + # end the series of every host it had not reached yet. + if self._stop.is_set(): + return + groups_by_host.setdefault(node.host, set()).update( + worker_group_names_on_node(node.node_id) + ) + if node.host in values_by_host: + continue + values = self._scrape(node) + if values: + values_by_host[node.host] = values + sampled = { + (host, format_worker_groups(groups_by_host[host])): values + for host, values in values_by_host.items() + } + with _SNAPSHOT_LOCK: + # Replaced, not merged: a host that has dropped out of the cluster + # or changed its worker-group label must stop being reported under + # the old one, and leaving the old key in would be the stale series + # this design exists to avoid. The cost is that a node which fails + # one scrape has a gap until the next, rather than a flat line at + # its last known value -- which is the honest reading anyway. + _SNAPSHOT.clear() + _SNAPSHOT.update(sampled) + + def _scrape(self, node: RayNode) -> dict[str, float]: + """Fetch and parse one node's memory series, or ``{}`` if unreachable. + + A node that cannot be scraped is skipped rather than retried: the next + interval is the retry, and a node that is down or overloaded is exactly + the one that must not be allowed to stall the loop for the others. + """ + try: + response = requests.get(node.metrics_url, timeout=_SCRAPE_TIMEOUT_S) + response.raise_for_status() + except requests.RequestException: + warn_once( + f"host_memory_scrape:{node.host}", + f"could not scrape host memory from {node.metrics_url}", + ) + return {} + return parse_memory_samples(response.text) + + def _collection_loop(self) -> None: + """Sample until stopped, on :attr:`collection_interval`.""" + # Sampled before the first wait, so a run shorter than one interval + # still reports something. Event.wait rather than sleep so stop() + # returns promptly instead of blocking for up to a whole interval. + while True: + self.collect_once() + if self._stop.wait(self.collection_interval): + return + + +#: The one monitor a process runs, owned by the telemetry lifecycle below. +_MONITOR: Optional[HostMemoryMonitor] = None + + +def start_host_memory_monitoring( + collection_interval: float = DEFAULT_COLLECTION_INTERVAL, +) -> None: + """Start this process's monitor, if it has none. Never raises. + + Called from ``init_telemetry_driver``, which is what puts the monitor in + the **driver** process on every entry point. That is not a detail: the + placement map above is an ordinary module global, so it is only visible to + the process that wrote it, and ``RayWorkerGroup`` writes it from the driver. + Owning the monitor anywhere else reports the right memory under the wrong + label -- the single-controller entry points build their worker groups on + the driver (``examples/run_grpo_single_controller.py``, + ``examples/run_sft_v2.py``) but their ``Logger`` inside the controller + actor, whose copy of the map would never be anything but empty. + + Args: + collection_interval: Seconds between samples of every node. + """ + global _MONITOR + if _MONITOR is not None: + return + monitor = HostMemoryMonitor(collection_interval=collection_interval) + try: + # Before start(), so the instruments exist no matter when the first + # export cycle lands. Nothing records these series, so this call is the + # only thing that ever creates them -- without it the monitor would + # scrape happily into a snapshot no exporter ever reads. + create_rl_metric_instruments() + monitor.start() + except Exception: + warn_once("host_memory_start", "could not start host memory monitoring") + else: + _MONITOR = monitor + + +def stop_host_memory_monitoring() -> None: + """Stop this process's monitor. Never raises; a no-op if none is running. + + Called from ``shutdown_telemetry``, which every worker also calls -- hence + the no-op rather than an assertion that one is running. + """ + global _MONITOR + monitor = _MONITOR + _MONITOR = None + if monitor is None: + return + try: + monitor.stop() + except Exception: + warn_once("host_memory_stop", "could not stop host memory monitoring") + # Outside the guard: the callbacks stay registered for the life of the + # process -- OTel has no way to detach an observable instrument -- so a + # snapshot left behind here would keep being exported by every cycle after + # the monitor that produced it stopped, which is exactly the frozen series + # the observable kinds were chosen to avoid. + with _SNAPSHOT_LOCK: + _SNAPSHOT.clear() diff --git a/nemo_rl/telemetry/metrics.py b/nemo_rl/telemetry/metrics.py index a8eacf25103..cc3ae971fbe 100644 --- a/nemo_rl/telemetry/metrics.py +++ b/nemo_rl/telemetry/metrics.py @@ -221,7 +221,20 @@ def _metric_specs() -> list[MetricSpec]: unit=row.unit, description=row.description, ) - for row in (*teed_metrics(), *recorded_metrics()) + for row in teed_metrics() + ) + specs.extend( + MetricSpec( + row.key, + row.name, + row.kind, + unit=row.unit, + description=row.description, + # Only the recorded rows can carry one; a teed row is pushed by + # definition, since its value arrives through Logger.log_metrics. + callback=row.callback, + ) + for row in recorded_metrics() ) return specs @@ -300,6 +313,32 @@ def record_rl_metrics( _record(handle.meter, values, attributes) +def create_rl_metric_instruments() -> None: + """Build the ``rl`` group's instruments now. No-op unless exporting. + + Only the observable rows need this. A synchronous series gets its + instrument on the first :func:`record_rl_metrics` that names it, but an + observable one is never recorded to -- the exporter pulls it through the + callback -- so without this nothing ever creates it and it is simply never + exported. Call it once from the owner that declared the callback, after + telemetry is up. + + Because this registers the group, it also closes declaration (see + ``freeze_metrics``). That makes telemetry startup the deadline for + declaring a row, where it used to be the first metric recorded, so a module + that declares one has to be imported by then -- which every entry point + already does. + """ + handle = get_telemetry_handle() + if handle is None or not handle.is_exporting: + return + if not ensure_metric_group_registered(): + return + from nemo.lens.instruments import create_metric_instruments + + create_metric_instruments(handle.meter, RL_METRIC_GROUP) + + def efficiency_measurements() -> dict[str, str]: """Map each canonical efficiency category to its measurement kind. diff --git a/nemo_rl/telemetry/setup.py b/nemo_rl/telemetry/setup.py index a8f2cf49dc2..2dc901855d5 100644 --- a/nemo_rl/telemetry/setup.py +++ b/nemo_rl/telemetry/setup.py @@ -530,6 +530,32 @@ def init_telemetry_driver( setup_logging_bridge() + if handle.is_exporting: + # Deferred to break the cycle: host_memory reads get_telemetry_handle + # from this module. + from nemo_rl.telemetry.host_memory import ( + DEFAULT_COLLECTION_INTERVAL, + start_host_memory_monitoring, + ) + + # Here rather than in ``Logger`` beside the GPU monitor, even though the + # two sample the same endpoints: the host-memory series is labelled with + # the worker groups each node holds, and that map is a module global + # only the driver writes. The single-controller entry points run their + # ``Logger`` inside the controller actor, which would label every host + # as holding no group at all. + # + # Shares the GPU monitor's cadence so the two line up on a dashboard, + # read from the logger config rather than restated -- but not *gated* on + # ``monitor_gpus``, which decides what reaches W&B and has no business + # deciding whether an OTel series exists. + interval = _dig( + master_config, "logger", "gpu_monitoring", "collection_interval" + ) + start_host_memory_monitoring( + collection_interval=interval or DEFAULT_COLLECTION_INTERVAL + ) + # Every resolved field, not just the headline ones: the env projection uses # setdefault, so a stray NEMO_RL_OTEL_* in the shell silently overrides the # YAML. Logging what was actually resolved keeps "how was this run @@ -800,6 +826,16 @@ def shutdown_telemetry(timeout_ms: int = 5000) -> None: spans an unreachable collector was going to lose anyway, and being a daemon it never holds up process exit. """ + # Deferred to break the cycle: host_memory reads get_telemetry_handle from + # this module. + from nemo_rl.telemetry.host_memory import stop_host_memory_monitoring + + # Called before the early return below, and before the flush: the monitor + # records through the providers this is about to shut down, so leaving it + # sampling would have it writing into a closed meter for a whole interval. + + stop_host_memory_monitoring() + global _TELEMETRY_HANDLE handle = _TELEMETRY_HANDLE if handle is None: diff --git a/nemo_rl/telemetry/vocabulary.py b/nemo_rl/telemetry/vocabulary.py index d1a7d4e1c7b..24df7fba385 100644 --- a/nemo_rl/telemetry/vocabulary.py +++ b/nemo_rl/telemetry/vocabulary.py @@ -31,7 +31,12 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any, Iterable, Optional +from typing import TYPE_CHECKING, Any, Iterable, Optional + +if TYPE_CHECKING: + # Behind the guard to keep the no-runtime-imports promise above: an owner + # declaring a row must not have to have OTel installed to do it. + from opentelemetry.metrics import CallbackT #: Prefix every NeMo-RL series carries, stripped to form the registry key. METRIC_NAME_PREFIX = "rl." @@ -80,15 +85,26 @@ class RecordedMetric: Attributes: name: OTel series name emitted; :func:`registry_key` derives the key. - kind: One of lens's ``gauge`` / ``counter`` / ``histogram`` / ``up_down_counter``. + kind: One of lens's ``gauge`` / ``counter`` / ``histogram`` / + ``up_down_counter``, or one of its observable kinds + (``observable_gauge`` / ``observable_counter`` / + ``observable_up_down_counter``), which require *callback*. unit: UCUM-ish unit string, or empty when dimensionless. description: Help text carried to the backend. + callback: For the observable kinds only, what the exporter calls to + read the series. Each cycle's points are exactly what it returns, + so an attribute set it stops yielding stops being exported -- which + is the reason to choose an observable kind over pushing to a gauge, + whose last value per attribute set is otherwise re-exported for the + rest of the run. Must be a function that outlives the process's + first export, not a closure built per call. """ name: str kind: str = "gauge" unit: str = "" description: str = "" + callback: Optional[CallbackT] = None @property def key(self) -> str: @@ -142,8 +158,12 @@ def _check_not_frozen(name: str) -> None: raise RuntimeError( f"{name!r} was declared after the metric group was registered with " "lens, so it has no instrument and would be silently dropped. " - "Import the module that declares it before the first metric is " - "recorded -- at startup, not lazily from inside the training loop." + "Import the module that declares it before init_telemetry_driver " + "runs, which is where registration now happens -- the observable " + "host-memory series have to have their instruments built at " + "startup, because nothing records them into existence later. " + "Declare at import time from a module the entry point already " + "imports, not lazily from inside the training loop." ) diff --git a/pyproject.toml b/pyproject.toml index ffd4ff61696..88b4abe577c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -340,7 +340,7 @@ nemo_run = { git = "https://github.com/NVIDIA-NeMo/Run", rev = "414f0077c648fde2 # actually reconciles the two revs; the floor in [dependencies] stays above the # newest PyPI release so dropping either one fails to resolve instead of # falling back to 0.1.0. -nemo-lens = { git = "https://github.com/NVIDIA-NeMo/Lens.git", rev = "7bcc03ce3565d2be6858b819b049bcc4c3ec05b4" } +nemo-lens = { git = "https://github.com/NVIDIA-NeMo/Lens.git", rev = "4d045c03186b90afdc815db46b671f6a21bc7e98" } sglang = { git = "https://github.com/sgl-project/sglang.git", rev = "3003d70f680d41c59d1b7acbf65cb47795dfd19e", subdirectory = "python" } # torch/torchvision/triton all come from the torch index in order to pick up aarch64 wheels torch = [ @@ -509,12 +509,14 @@ override-dependencies = [ "mamba-ssm==2.3.2.post1", # megatron-core still pins lens v0.2.0 in its own [tool.uv.sources]; without an # override the two git URLs are a hard conflict, since the workspace resolves a - # single nemo-lens. The rev here is the v0.3.0 tag, which we need for + # single nemo-lens. The rev here is lens main past v0.3.0, which we need for # `SpanRegistry` (which replaced the shipped `SpanGroup` class), the - # consumer-driven metric registry (`MetricSpec`/`register_metric_group`) and - # the scoped span attributes behind `rl.iteration`. Drop this once Megatron-LM - # bumps to the same or newer rev. - "nemo-lens[sdk,aiohttp] @ git+https://github.com/NVIDIA-NeMo/Lens.git@7bcc03ce3565d2be6858b819b049bcc4c3ec05b4", + # consumer-driven metric registry (`MetricSpec`/`register_metric_group`), the + # scoped span attributes behind `rl.iteration`, and the observable metric + # kinds behind `rl.host.memory.*` -- a rev rather than a tag because there is + # no release carrying the last of those yet. Drop this once Megatron-LM bumps + # to the same or newer rev. + "nemo-lens[sdk,aiohttp] @ git+https://github.com/NVIDIA-NeMo/Lens.git@4d045c03186b90afdc815db46b671f6a21bc7e98", ] # CVE fixes diff --git a/pyrefly.toml b/pyrefly.toml index 4aa2faa91cd..cc0b44e9669 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -303,6 +303,7 @@ project-includes = [ "nemo_rl/models/value/workers/__init__.py", "nemo_rl/telemetry/__init__.py", "nemo_rl/telemetry/config.py", + "nemo_rl/telemetry/host_memory.py", "nemo_rl/telemetry/instrumentation.py", "nemo_rl/telemetry/metrics.py", "nemo_rl/telemetry/setup.py", diff --git a/tests/unit/telemetry/test_host_memory.py b/tests/unit/telemetry/test_host_memory.py new file mode 100644 index 00000000000..d3dd913d3fc --- /dev/null +++ b/tests/unit/telemetry/test_host_memory.py @@ -0,0 +1,517 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the host-memory monitor. + +Ray and its Prometheus endpoints are both faked: the point of these is the +attribution and the never-raises guarantee, neither of which needs a cluster. +""" + +import sys +import types +from types import SimpleNamespace + +import pytest + +import nemo_rl.telemetry.host_memory as host_memory +from nemo_rl.telemetry.host_memory import ( + HOST_NAME_ATTR, + NO_WORKER_GROUPS, + RL_HOST_MEMORY_TOTAL_METRIC, + RL_HOST_MEMORY_USED_METRIC, + RL_WORKER_GROUPS_ATTR, + HostMemoryMonitor, + forget_worker_group_nodes, + parse_memory_samples, + ray_nodes, + register_worker_group_nodes, + start_host_memory_monitoring, + stop_host_memory_monitoring, + worker_groups_on_node, +) +from nemo_rl.telemetry.metrics import create_rl_metric_instruments + +_USED_BYTES = 16_000_000_000.0 +_TOTAL_BYTES = 540_000_000_000.0 + +#: A raylet's payload, cut down to the two series that are read plus one that +#: is not, so "ignores everything else" is actually exercised. +_SCRAPE = f"""\ +# HELP ray_node_mem_used Memory in use on the node. +# TYPE ray_node_mem_used gauge +ray_node_mem_used{{SessionName="s"}} {_USED_BYTES} +# HELP ray_node_mem_total Memory installed on the node. +# TYPE ray_node_mem_total gauge +ray_node_mem_total{{SessionName="s"}} {_TOTAL_BYTES} +# HELP ray_node_cpu_utilization CPU percentage. +# TYPE ray_node_cpu_utilization gauge +ray_node_cpu_utilization{{SessionName="s"}} 42.0 +""" + + +@pytest.fixture(autouse=True) +def _reset_host_memory_state(): + """Placement map, snapshot and monitor are process-global; none leaks.""" + host_memory._WORKER_GROUP_NODES.clear() + host_memory._SNAPSHOT.clear() + stop_host_memory_monitoring() + yield + host_memory._WORKER_GROUP_NODES.clear() + host_memory._SNAPSHOT.clear() + stop_host_memory_monitoring() + + +def _fake_ray(monkeypatch, nodes): + """Install a ``ray`` module reporting *nodes* from ``ray.nodes()``.""" + module = types.ModuleType("ray") + module.is_initialized = lambda: True + module.nodes = lambda: nodes + monkeypatch.setitem(sys.modules, "ray", module) + return module + + +def _node(node_id="n0", *, hostname="worker-0", address="10.0.0.1", alive=True): + return { + "NodeID": node_id, + "NodeManagerHostname": hostname, + "NodeManagerAddress": address, + "MetricsExportPort": 8080, + "Alive": alive, + } + + +def _fake_scrape(monkeypatch, by_url): + """Serve *by_url* from ``requests.get``; a missing URL raises like a dead node.""" + import requests + + import nemo_rl.telemetry.host_memory as host_memory + + class _Response: + def __init__(self, text): + self.text = text + + def raise_for_status(self): + return None + + def _get(url, timeout=None): + if url not in by_url: + raise requests.ConnectionError(url) + return _Response(by_url[url]) + + monkeypatch.setattr(host_memory.requests, "get", _get) + + +def _gauge_points(data, name): + """Data points for gauge *name* in one already-collected metrics batch.""" + if data is None: + return [] + return [ + point + for rm in data.resource_metrics + for sm in rm.scope_metrics + for metric in sm.metrics + if metric.name == name + for point in metric.data.data_points + ] + + +def _start_exporting_telemetry(): + """Install an exporting telemetry handle; returns its metric reader.""" + from nemo.lens import NemoLensConfig, setup_telemetry + from opentelemetry.sdk.metrics.export import InMemoryMetricReader + + import nemo_rl.telemetry.setup as setup_mod + + instruments = pytest.importorskip("nemo.lens.instruments") + if not hasattr(instruments, "register_metric_group"): + pytest.skip("installed nemo-lens has no metric registry (lens PR #46)") + if not hasattr(instruments, "create_metric_instruments"): + pytest.skip("installed nemo-lens has no observable kinds (lens PR #84)") + + reader = InMemoryMetricReader() + setup_mod._TELEMETRY_HANDLE = setup_telemetry( + NemoLensConfig(enabled=True), metric_reader=reader + ) + # These series are observable, so nothing records them into existence -- + # the reader only sees them once their instruments exist. In a run this is + # start_host_memory_monitoring's job; here the tests drive collection by + # hand, so it is the handle's. + create_rl_metric_instruments() + return reader + + +def test_only_the_two_node_memory_series_are_read(): + """A raylet publishes hundreds of series; all but two are noise here.""" + assert parse_memory_samples(_SCRAPE) == { + "host_memory_used": _USED_BYTES, + "host_memory_total": _TOTAL_BYTES, + } + + +def test_bytes_are_reported_unscaled(): + """The W&B path divides to GB on the way out; OTel wants the base unit.""" + values = parse_memory_samples(_SCRAPE) + + assert values["host_memory_used"] == _USED_BYTES + assert values["host_memory_used"] > 1e9 + + +def test_a_node_is_named_by_hostname_rather_than_ip(monkeypatch): + """``NodeManagerAddress`` is the IP; the hostname is what a reader knows.""" + _fake_ray(monkeypatch, [_node(hostname="dgx-042", address="10.0.0.7")]) + + (node,) = ray_nodes() + + assert node.host == "dgx-042" + # The IP is still what the endpoint is reached on. + assert node.metrics_url == "http://10.0.0.7:8080/metrics" + + +def test_a_node_falls_back_to_its_ip_when_ray_resolved_no_hostname(monkeypatch): + """Better an address than an unlabelled point.""" + _fake_ray(monkeypatch, [_node(hostname="", address="10.0.0.7")]) + + (node,) = ray_nodes() + + assert node.host == "10.0.0.7" + + +def test_dead_nodes_are_not_scraped(monkeypatch): + """``ray.nodes()`` keeps dead entries, whose endpoints only ever time out.""" + _fake_ray( + monkeypatch, + [_node("n0", hostname="live"), _node("n1", hostname="dead", alive=False)], + ) + + assert [node.host for node in ray_nodes()] == ["live"] + + +def test_no_endpoint_means_no_node(monkeypatch): + """A node without a metrics port publishes nothing to scrape.""" + node = _node() + del node["MetricsExportPort"] + _fake_ray(monkeypatch, [node]) + + assert ray_nodes() == [] + + +def test_every_group_resident_on_a_node_is_named(): + """Colocation puts several groups on one host; the label says so.""" + register_worker_group_nodes("vllm_policy", ["n0", "n1"]) + register_worker_group_nodes("lm_policy", ["n0"]) + + # Sorted, so the attribute value is stable rather than insertion-ordered -- + # an unstable value would split one host's series in two. + assert worker_groups_on_node("n0") == "lm_policy,vllm_policy" + assert worker_groups_on_node("n1") == "vllm_policy" + + +def test_a_node_running_no_worker_group_is_still_labelled(): + """The head node hosts only the driver, and its memory is just as real.""" + register_worker_group_nodes("lm_policy", ["n1"]) + + assert worker_groups_on_node("head") == NO_WORKER_GROUPS + + +def test_a_regrouped_worker_group_is_not_reported_on_both_placements(): + """A restart can move a group; the old nodes must stop claiming it.""" + register_worker_group_nodes("lm_policy", ["n0"]) + register_worker_group_nodes("lm_policy", ["n1"]) + + assert worker_groups_on_node("n0") == NO_WORKER_GROUPS + assert worker_groups_on_node("n1") == "lm_policy" + + +def test_a_shut_down_group_stops_claiming_its_nodes(): + register_worker_group_nodes("teacher", ["n0"]) + forget_worker_group_nodes("teacher") + + assert worker_groups_on_node("n0") == NO_WORKER_GROUPS + + +def test_each_host_is_reported_with_its_name_and_its_worker_groups(monkeypatch): + pytest.importorskip("nemo.lens") + reader = _start_exporting_telemetry() + + _fake_ray( + monkeypatch, + [ + _node("n0", hostname="dgx-000", address="10.0.0.1"), + _node("n1", hostname="dgx-001", address="10.0.0.2"), + ], + ) + _fake_scrape( + monkeypatch, + { + "http://10.0.0.1:8080/metrics": _SCRAPE, + "http://10.0.0.2:8080/metrics": _SCRAPE, + }, + ) + register_worker_group_nodes("lm_policy", ["n0"]) + + HostMemoryMonitor().collect_once() + + data = reader.get_metrics_data() + used = { + point.attributes[HOST_NAME_ATTR]: point + for point in _gauge_points(data, RL_HOST_MEMORY_USED_METRIC) + } + assert set(used) == {"dgx-000", "dgx-001"} + assert used["dgx-000"].value == _USED_BYTES + assert used["dgx-000"].attributes[RL_WORKER_GROUPS_ATTR] == "lm_policy" + assert used["dgx-001"].attributes[RL_WORKER_GROUPS_ATTR] == NO_WORKER_GROUPS + + total = _gauge_points(data, RL_HOST_MEMORY_TOTAL_METRIC) + assert {point.value for point in total} == {_TOTAL_BYTES} + + +def test_a_label_the_scrape_stops_reporting_stops_being_exported(monkeypatch): + """Why these series are observable instead of pushed to a gauge. + + A host is labelled ``none`` until its worker groups are built, which is + minutes into a run. Pushed to a synchronous gauge, that startup point would + keep being re-exported next to the correct one for the rest of the run: the + SDK holds a gauge's last value per attribute set, and nothing can retract + one. ``sum by (host.name)`` would then count every host twice, and + filtering on ``none`` would show every training host as idle. + """ + pytest.importorskip("nemo.lens") + reader = _start_exporting_telemetry() + + _fake_ray(monkeypatch, [_node("n0", hostname="dgx-000", address="10.0.0.1")]) + _fake_scrape(monkeypatch, {"http://10.0.0.1:8080/metrics": _SCRAPE}) + monitor = HostMemoryMonitor() + + monitor.collect_once() + at_startup = _gauge_points(reader.get_metrics_data(), RL_HOST_MEMORY_USED_METRIC) + assert [point.attributes[RL_WORKER_GROUPS_ATTR] for point in at_startup] == [ + NO_WORKER_GROUPS + ] + + register_worker_group_nodes("lm_policy", ["n0"]) + monitor.collect_once() + + once_grouped = _gauge_points(reader.get_metrics_data(), RL_HOST_MEMORY_USED_METRIC) + assert [point.attributes[RL_WORKER_GROUPS_ATTR] for point in once_grouped] == [ + "lm_policy" + ] + + +def test_two_raylets_on_one_host_are_one_point_scraped_once(monkeypatch): + """A local multi-node cluster must not report its host twice. + + Each raylet has its own endpoint but publishes the same whole-host memory, + so a point per node would double count under ``sum by (host.name)``. The + worker groups of both are merged, since the label describes the host. + """ + pytest.importorskip("nemo.lens") + reader = _start_exporting_telemetry() + + _fake_ray( + monkeypatch, + [ + _node("n0", hostname="dev-box", address="10.0.0.1"), + _node("n1", hostname="dev-box", address="10.0.0.1"), + ], + ) + _fake_scrape(monkeypatch, {"http://10.0.0.1:8080/metrics": _SCRAPE}) + served = host_memory.requests.get + scraped: list[str] = [] + + def _counting(url, timeout=None): + scraped.append(url) + return served(url, timeout=timeout) + + monkeypatch.setattr(host_memory.requests, "get", _counting) + register_worker_group_nodes("lm_policy", ["n0"]) + register_worker_group_nodes("vllm_policy", ["n1"]) + + HostMemoryMonitor().collect_once() + + points = _gauge_points(reader.get_metrics_data(), RL_HOST_MEMORY_USED_METRIC) + assert len(points) == 1 + assert points[0].attributes[HOST_NAME_ATTR] == "dev-box" + assert points[0].attributes[RL_WORKER_GROUPS_ATTR] == "lm_policy,vllm_policy" + assert len(scraped) == 1 + + +def test_stopping_abandons_the_sweep_instead_of_finishing_it(monkeypatch): + """Shutdown costs the scrape in flight, not one per remaining node. + + The sweep is sequential and each node has its own timeout, so finishing it + on the way out would make shutdown scale with the size of the cluster. The + partial result is thrown away rather than published: it would otherwise end + the series of every host the sweep had not reached. + """ + pytest.importorskip("nemo.lens") + _start_exporting_telemetry() + + nodes = [ + _node(f"n{index}", hostname=f"dgx-{index}", address=f"10.0.0.{index}") + for index in range(5) + ] + _fake_ray(monkeypatch, nodes) + monitor = HostMemoryMonitor() + scraped: list[str] = [] + + class _Response: + text = _SCRAPE + + def raise_for_status(self): + return None + + def _get(url, timeout=None): + scraped.append(url) + # Shutdown lands while the first node's request is in flight. + monitor._stop.set() + return _Response() + + monkeypatch.setattr(host_memory.requests, "get", _get) + + monitor.collect_once() + + assert len(scraped) == 1 + assert host_memory._SNAPSHOT == {} + + +def test_a_stopped_monitor_leaves_nothing_behind_to_export(monkeypatch): + """The callbacks outlive the monitor: OTel cannot detach an instrument. + + So whatever the last scrape left behind would otherwise be exported, frozen, + by every cycle between shutdown and process exit. + """ + pytest.importorskip("nemo.lens") + reader = _start_exporting_telemetry() + + _fake_ray(monkeypatch, [_node("n0", hostname="dgx-000", address="10.0.0.1")]) + _fake_scrape(monkeypatch, {"http://10.0.0.1:8080/metrics": _SCRAPE}) + + # An interval long enough that the loop samples once and then sleeps, so + # the points under test are the ones collected below and not a race. + start_host_memory_monitoring(collection_interval=3600.0) + HostMemoryMonitor().collect_once() + assert _gauge_points(reader.get_metrics_data(), RL_HOST_MEMORY_USED_METRIC) + + stop_host_memory_monitoring() + + assert not _gauge_points(reader.get_metrics_data(), RL_HOST_MEMORY_USED_METRIC) + + +def test_an_unreachable_node_costs_only_its_own_point(monkeypatch): + """The node that is wedged is exactly the one whose host memory matters.""" + pytest.importorskip("nemo.lens") + reader = _start_exporting_telemetry() + + _fake_ray( + monkeypatch, + [ + _node("n0", hostname="dgx-000", address="10.0.0.1"), + _node("n1", hostname="wedged", address="10.0.0.2"), + ], + ) + _fake_scrape(monkeypatch, {"http://10.0.0.1:8080/metrics": _SCRAPE}) + + HostMemoryMonitor().collect_once() + + points = _gauge_points(reader.get_metrics_data(), RL_HOST_MEMORY_USED_METRIC) + assert [point.attributes[HOST_NAME_ATTR] for point in points] == ["dgx-000"] + + +def test_collection_never_raises(monkeypatch): + """It runs on a timer beside a training run and must not be able to end it.""" + import nemo_rl.telemetry.host_memory as host_memory + + def _explode(): + raise RuntimeError("GCS is unreachable") + + monkeypatch.setattr(host_memory, "ray_nodes", _explode) + _start_exporting_telemetry() + + HostMemoryMonitor().collect_once() + + +def test_nothing_is_scraped_when_telemetry_is_not_exporting(monkeypatch): + """No handle means no consumer, so the HTTP calls are wasted work.""" + import nemo_rl.telemetry.setup as setup_mod + + setup_mod._TELEMETRY_HANDLE = None + scraped = [] + _fake_ray(monkeypatch, [_node()]) + monkeypatch.setattr( + HostMemoryMonitor, + "_scrape", + lambda self, node: scraped.append(node) or {}, + ) + + HostMemoryMonitor().collect_once() + + assert scraped == [] + + +def test_stopping_a_monitor_that_was_never_started_is_harmless(): + HostMemoryMonitor().stop() + + +def test_a_process_runs_at_most_one_monitor(monkeypatch): + """``init_telemetry_driver`` is idempotent, so starting twice must be too.""" + _fake_ray(monkeypatch, []) + + start_host_memory_monitoring(collection_interval=60.0) + first = host_memory._MONITOR + start_host_memory_monitoring(collection_interval=60.0) + + assert host_memory._MONITOR is first + + +def test_stopping_releases_the_monitor_so_a_later_run_can_start_one(monkeypatch): + _fake_ray(monkeypatch, []) + + start_host_memory_monitoring(collection_interval=60.0) + stop_host_memory_monitoring() + + assert host_memory._MONITOR is None + + +def test_stopping_without_a_monitor_is_a_no_op(): + """Every worker calls ``shutdown_telemetry``, and none of them started one.""" + stop_host_memory_monitoring() + + assert host_memory._MONITOR is None + + +def test_the_driver_starts_the_monitor_that_the_worker_groups_register_with( + monkeypatch, tmp_path +): + """The whole point of owning it here: one process writes and reads the map. + + ``RayWorkerGroup`` registers placements from the driver into a module + global, so a monitor anywhere else -- the controller actor that owns the + ``Logger`` under single-controller GRPO and SFT v2 -- would read an empty + map and label every host ``none``. + """ + pytest.importorskip("nemo.lens") + from nemo_rl.telemetry.setup import init_telemetry_driver, shutdown_telemetry + + monkeypatch.setenv("NEMO_RL_OTEL_ENABLED", "true") + monkeypatch.setenv("NEMO_RL_OTEL_EXPORTER", "console") + _fake_ray(monkeypatch, []) + + config = SimpleNamespace( + telemetry=None, + logger=SimpleNamespace( + gpu_monitoring=SimpleNamespace(collection_interval=60.0) + ), + policy=None, + cluster=None, + ) + handle = init_telemetry_driver(config, algorithm="grpo") + if handle is None or not handle.is_exporting: + pytest.skip("telemetry did not come up exporting in this environment") + + assert host_memory._MONITOR is not None + # Taken from the logger's GPU cadence rather than restated, so the two + # families land on the same points of a dashboard. + assert host_memory._MONITOR.collection_interval == 60.0 + + shutdown_telemetry(timeout_ms=100) + assert host_memory._MONITOR is None diff --git a/tests/unit/utils/test_checkpoint.py b/tests/unit/utils/test_checkpoint.py index e10de34c005..13141b362c7 100644 --- a/tests/unit/utils/test_checkpoint.py +++ b/tests/unit/utils/test_checkpoint.py @@ -1200,8 +1200,6 @@ def test_swaps_when_target_exists(self, checkpoint_manager, checkpoint_dir): with open(step_dir / "training_info.json") as f: assert json.load(f)["loss"] == 0.2 - assert not (checkpoint_dir / "step_12").exists() - # --------------------------------------------------------------------------- # Fault tolerance (ft_keep_latest_k / ft_save_period) retention tests diff --git a/uv.lock b/uv.lock index 3522fa1dc79..cfcff447d0b 100644 --- a/uv.lock +++ b/uv.lock @@ -97,7 +97,7 @@ overrides = [ { name = "llguidance", specifier = ">=1.7.0,<1.8.0" }, { name = "mamba-ssm", git = "https://github.com/state-spaces/mamba.git?rev=a14b1dff0454a3bc27d9eb31355dc01e4b2490ec" }, { name = "mlflow", specifier = ">=3.13.0" }, - { name = "nemo-lens", extras = ["sdk", "aiohttp"], git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=7bcc03ce3565d2be6858b819b049bcc4c3ec05b4" }, + { name = "nemo-lens", extras = ["sdk", "aiohttp"], git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=4d045c03186b90afdc815db46b671f6a21bc7e98" }, { name = "numpy", specifier = ">=2.1.0,<2.5" }, { name = "nvidia-cublas", marker = "sys_platform != 'darwin'", specifier = "==13.5.1.27" }, { name = "nvidia-cudnn-cu13", marker = "sys_platform != 'darwin'", specifier = "==9.26.0.51" }, @@ -4133,8 +4133,8 @@ codec = [ [[package]] name = "nemo-lens" -version = "0.2.0+7bcc03c" -source = { git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=7bcc03ce3565d2be6858b819b049bcc4c3ec05b4#7bcc03ce3565d2be6858b819b049bcc4c3ec05b4" } +version = "0.2.0+4d045c0" +source = { git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=4d045c03186b90afdc815db46b671f6a21bc7e98#4d045c03186b90afdc815db46b671f6a21bc7e98" } dependencies = [ { name = "opentelemetry-api" }, ] @@ -4387,7 +4387,7 @@ requires-dist = [ { name = "nccl4py", marker = "sys_platform != 'darwin'" }, { name = "nemo-automodel", extras = ["moe"], marker = "extra == 'automodel'", editable = "3rdparty/Automodel-workspace/Automodel" }, { name = "nemo-gym", marker = "extra == 'nemo-gym'", editable = "3rdparty/Gym-workspace/Gym" }, - { name = "nemo-lens", extras = ["aiohttp", "sdk"], git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=7bcc03ce3565d2be6858b819b049bcc4c3ec05b4" }, + { name = "nemo-lens", extras = ["aiohttp", "sdk"], git = "https://github.com/NVIDIA-NeMo/Lens.git?rev=4d045c03186b90afdc815db46b671f6a21bc7e98" }, { name = "ninja" }, { name = "nixl", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')", specifier = "==1.3.0" }, { name = "num2words", specifier = ">=0.5.14" }, From 099fa4a90a512d529890745a32f16b3cf82a5916 Mon Sep 17 00:00:00 2001 From: Raj Singh Date: Tue, 13 Oct 2026 17:39:38 -0700 Subject: [PATCH 10/17] fix(telemetry): act on review of the host-memory series Thirteen review threads, grouped by what they change. Behaviour. parse_memory_samples ran outside the scrape's try, which only caught requests.RequestException; prometheus_client raises ValueError on a body that is not the exposition format, so one raylet serving an HTML error page aborted the whole sweep and left every host's point frozen at the previous snapshot. It is now guarded on its own, and a bad page costs only its own point. rl.checkpoint.async_save is gone: all twelve begin_finalization callers pass wait_fn whether or not async save is on, so the attribute read true on every checkpoint, sync saves included. Reasons that were wrong. The observable-gauge rationale claimed the SDK re-exports a synchronous gauge's last value forever; on the locked 1.44.0 it drops the value at each collect instead, so the cycles between two samples carry no point at all, and only from 1.45 does it hold the last value per attribute set. Both make an observable gauge the right choice, for different reasons, and the comment, metrics.md and the RecordedMetric docstring now say so. The nv.dl.topology.size.* rename was justified by lens typing resource_attributes, which it does not -- it merges the caller's as-is, so the name is a convention. iteration_scope was said to win against an explicit value because lens skips keys already present; lens starts the span before applying the caller's attributes, so they overwrite. The deferred import named nemo_rl.utils.logger as the importer; it is worker_groups.py, at module scope. Claims that were too broad. setup.py said only the driver builds worker groups, but RewardModelEnvironment builds its Policy inside a Ray actor, so its hosts report rl.worker_groups="none" -- documented rather than fixed, since completing the label means a named actor and a round trip per sample. span-groups.md said every algorithm opens iteration_scope; four loops do not. index.md said every entry point initialises telemetry; six do not. extending.md told contributors to import rl.* names from lens, reversing main's rule that RL names its own attributes, so rl.algorithm is RL's constant again and the rule is back. Tests. register_worker_group_placement had no coverage despite every RayWorkerGroup calling it: one table lookup per placement group, a bundle index the table does not place, and the never-raises guard. The lens import guards could never skip, since the telemetry conftest imports lens for every test in the directory, and a handle that comes up not exporting is the regression the driver test exists to catch rather than a reason to skip it. The fixture also unfreezes the metric vocabulary it froze. FALLBACK_COLLECTION_INTERVAL is no longer any parameter's default: it is reached only when the config carries no logger block at all, so a run that does configure a cadence cannot silently fall back to it. Read from GPUMonitoringConfig instead would drag mlflow and wandb into telemetry startup for one float. Signed-off-by: Raj Singh --- docs/observability/configuration.md | 2 +- docs/observability/extending.md | 4 +- docs/observability/index.md | 6 +- docs/observability/metrics.md | 12 +- docs/observability/span-groups.md | 18 ++- nemo_rl/telemetry/host_memory.py | 54 ++++--- nemo_rl/telemetry/instrumentation.py | 45 +++--- nemo_rl/telemetry/setup.py | 42 +++-- nemo_rl/telemetry/vocabulary.py | 9 +- nemo_rl/utils/checkpoint.py | 14 +- pyproject.toml | 4 +- tests/unit/telemetry/test_checkpoint_spans.py | 21 --- tests/unit/telemetry/test_host_memory.py | 148 +++++++++++++++--- tests/unit/telemetry/test_instrumentation.py | 12 +- 14 files changed, 255 insertions(+), 136 deletions(-) diff --git a/docs/observability/configuration.md b/docs/observability/configuration.md index c9723b7fcd4..b4b87202f73 100644 --- a/docs/observability/configuration.md +++ b/docs/observability/configuration.md @@ -108,7 +108,7 @@ Each process sets stable-for-the-run values on its OTel `Resource` — `init_tel The config-derived rows are driver-only: a worker is handed the propagated environment, not the `MasterConfig`. Attribute construction is best-effort: a missing config key simply omits that attribute; it never raises. Plus auto-detected host / GPU / SLURM / Kubernetes attributes from lens's resource detection. -Every key above that lens declares is taken from `nemo.lens.semconv` rather than written as a string, so a typo is an `ImportError` instead of silent data loss. Note the two spellings of the parallelism sizes: `nv.dl.topology.size.*` is the resource-scoped one, registered in lens's `DL_RESOURCE_TYPES` and so typed by `normalize_resource_attributes`, while the `dl.tensor_parallel.*` family is for span and metric attributes. A resource attribute outside that registry is passed through rather than rejected, so the wrong spelling arrives as an untyped string instead of failing. +Shared `nv.dl.*` keys are taken from `nemo.lens.semconv` rather than written as strings, so a typo is an `ImportError` instead of silent data loss. Note the two spellings of the parallelism sizes: `nv.dl.topology.size.*` is the resource-scoped one and the `dl.tensor_parallel.*` family is the span and metric spelling of the same quantity. That split is a convention this repo keeps, not something lens checks — `normalize_resource_attributes` runs only over the attributes lens derives from `OTEL_RESOURCE_ATTRIBUTES`, and the `resource_attributes` RL passes in are merged as they are. Either spelling would arrive intact, so the wrong one is a silently misnamed series rather than a failure. Lens also declares `nv.dl.topology.size.dp`, which NeMo-RL does not report. RL configures no data-parallel size, so it could only be derived by dividing the cluster's GPU count by the model-parallel sizes — and that assumes the policy owns every GPU in the cluster. A reward model with its own `resources` block, distillation teachers, expert parallelism and non-colocated generation each leave the policy fewer, and the division comes out too high rather than failing. Read the worker's device mesh if you need the real number. diff --git a/docs/observability/extending.md b/docs/observability/extending.md index 5e82a2a2b98..465cfd69068 100644 --- a/docs/observability/extending.md +++ b/docs/observability/extending.md @@ -195,9 +195,9 @@ if telemetry is not None: **Metric names use the application scope (`rl.*`) — never `dl.*`.** This is not a style preference: lens registers a `dl.*` series for the trainer that owns the step, so an RL job publishing `dl.grad_norm` would be writing into a co-resident Megatron trainer's series. `rl.grad_norm` and `rl.learning_rate` are deliberate spelling matches for exactly that reason, so query the two by their own names rather than by swapping the prefix. Lens's own semconv file says the same: RL's metric series are declared through `register_metric_group` in this tree and are deliberately not defined there. -**Attributes split by who owns the name.** If lens declares it, import the constant from `nemo.lens.semconv` — a typo then fails as an `ImportError` instead of becoming silent data loss. That covers the shared `nv.dl.*` resource surface and the handful of `rl.*` span attributes lens itself declares (`rl.algorithm`, `rl.reward`, `rl.generation.backend`, `rl.num_rollouts`; `RL_ALGORITHM` is re-exported from `nemo_rl/telemetry/instrumentation.py` so call sites take every `rl.*` key from one module). Everything else — `rl.iteration`, `rl.bucket`, `rl.worker_group` — describes RL's own properties, so it is named in `nemo_rl/telemetry/` where a rename can land. +**Attributes split by who owns the name.** RL names every `rl.*` attribute in `nemo_rl/telemetry/`, so a rename lands where the producers are. That holds even for the few `rl.*` keys lens also declares in its semconv (`rl.algorithm`, `rl.reward`, `rl.generation.backend`, `rl.num_rollouts`): lens declares those and reads none of them, so importing the constant would buy a dependency with no contract behind it while RL remains the only producer. Import from `nemo.lens.semconv` only for names lens actually acts on — the shared `nv.dl.*` resource surface — where a typo should fail as an `ImportError` rather than become silent data loss. -Mind the scope when picking a shared name: `nv.dl.topology.size.tp` is the resource-scoped parallelism size, typed by lens's `DL_RESOURCE_TYPES` registry, while `dl.tensor_parallel.size` is the span and metric spelling of the same quantity. +Mind the scope when picking a shared name: `nv.dl.topology.size.tp` is the resource-scoped parallelism size, while `dl.tensor_parallel.size` is the span and metric spelling of the same quantity. ## Choosing a span group diff --git a/docs/observability/index.md b/docs/observability/index.md index 3f82cf40122..227f807e59f 100644 --- a/docs/observability/index.md +++ b/docs/observability/index.md @@ -60,7 +60,9 @@ Keeping the settings in the config file is what makes a run's telemetry reproduc ## What gets instrumented -Each algorithm's `examples/run_.py` calls `init_telemetry_driver(config, algorithm="")` **before** `init_ray()` (so the resolved `NEMO_RL_OTEL_*` settings are snapshotted into the Ray `runtime_env` and inherited by every worker) and `shutdown_telemetry()` from a `finally` block wrapping the whole run, so buffered spans are flushed even when the run fails. +The instrumented entry points — those in the table below, plus `examples/run_sft_v2.py` — call `init_telemetry_driver(config, algorithm="")` **before** `init_ray()` (so the resolved `NEMO_RL_OTEL_*` settings are snapshotted into the Ray `runtime_env` and inherited by every worker) and `shutdown_telemetry()` from a `finally` block wrapping the whole run, so buffered spans are flushed even when the run fails. + +Not every entry point is wired up: `run_eval.py`, `run_grpo_sliding_puzzle.py`, `run_vlm_grpo.py`, `run_vlm_mpo.py`, `run_vlm_sft.py` and `run_xtoken_off_policy_distillation.py` emit nothing today. A run launched through one of those produces no spans and no metrics, whatever `telemetry.enabled` says. | Algorithm | Entry point | Representative spans | |---|---|---| @@ -80,7 +82,7 @@ Each span belongs to a **span group** that controls whether it is emitted at run ## What gets exported - **Traces**: any OTLP-compatible backend (Jaeger, Grafana Tempo, an OpenTelemetry Collector, ...) via OTLP. -- **Metrics**: the `rl.efficiency.*` async accounting teed from the driver's metrics logger, the training scalars (`rl.reward.mean`, `rl.policy.loss`, …), `rl.setup.duration` for the startup phases, the `rl.vllm.*` engine deltas, and the driver-side `gen_ai.*` generation series — see [Metrics](metrics.md). +- **Metrics**: the `rl.efficiency.*` async accounting teed from the driver's metrics logger, the training scalars (`rl.reward.mean`, `rl.policy.loss`, …), `rl.setup.duration` for the startup phases, the `rl.vllm.*` engine deltas, the driver-side `gen_ai.*` generation series, and `rl.host.memory.*` per node — see [Metrics](metrics.md). - **Logs** (optional): via the OTel log bridge when `telemetry.logs_enabled` is true — correlates Python `logging` records with the active span's trace ID. **Every** process that enables telemetry exports, each labelled with `nv.dl.rank` / `nv.dl.world_size`. Narrowing a large fleet down is a collector-side filter rather than a config setting — see [Configuration — Which ranks export](configuration.md#which-ranks-export). diff --git a/docs/observability/metrics.md b/docs/observability/metrics.md index da98371bbcd..ce58a5f7b31 100644 --- a/docs/observability/metrics.md +++ b/docs/observability/metrics.md @@ -1,8 +1,8 @@ # Metrics -NeMo-RL emits five namespaces of metrics: async efficiency metrics (`rl.efficiency.*`), startup phase durations (`rl.setup.duration`), training scalars (`rl.reward.*`, `rl.policy.*`, …), vLLM engine metrics (`rl.vllm.*`), and driver-side vLLM generation metrics (`gen_ai.*`, following the OTel GenAI semantic conventions). +NeMo-RL emits six namespaces of metrics: async efficiency metrics (`rl.efficiency.*`), startup phase durations (`rl.setup.duration`), training scalars (`rl.reward.*`, `rl.policy.*`, …), vLLM engine metrics (`rl.vllm.*`), driver-side vLLM generation metrics (`gen_ai.*`, following the OTel GenAI semantic conventions), and per-host memory (`rl.host.memory.*`). -Metrics are emitted **only when telemetry is exporting** — the driver always exports, so the `rl.*` series come from the driver's metrics logger. +Metrics are emitted **only when telemetry is exporting**. All of them are produced on the driver, which always exports: most pass through the driver's metrics logger, while `rl.host.memory.*` is scraped from every node by a driver-side monitor on its own timer. Every `rl.*` series is **declared by NeMo-RL**, not by lens. lens provides a consumer-driven metric registry: NeMo-RL calls `register_metric_group("rl", [...])` once per process with a `MetricSpec` per series, then records against it with `record_metrics(meter, "rl", ...)`. The declarations live in one table in `nemo_rl/telemetry/metrics.py`, so adding a series needs no lens release and no agreement on field names. `record_metrics` never raises into the caller and skips `None`. @@ -218,9 +218,9 @@ The monitor is started by `init_telemetry_driver` and stopped by `shutdown_telem ### Why these are observable rather than pushed -A host's `rl.worker_groups` label changes during a run: it is `none` until the worker groups are built, which is minutes in, and changes again when they regroup or shut down. Pushed to an ordinary gauge that would be a permanent error, not a transient one — the SDK keeps a synchronous gauge's last value *per attribute set* and re-exports it every cycle, and nothing in the OTel API can retract an attribute set. Every host's startup-time `none` point would sit next to its correctly labelled one for the rest of the run, so `sum by (host.name)` would count each host twice and filtering on `none` would show every training host as idle. +The monitor samples on its own timer, once per collection interval, which does not line up with the exporter's cycle — and a synchronous gauge is reported only in the cycles that follow a `set`. On the locked SDK (1.44.0) the cycles in between would carry no point at all; from 1.45 the SDK holds a synchronous gauge's last value *per attribute set* instead, which trades the gaps for a staler problem, since nothing in the OTel API can retract an attribute set. A host's `rl.worker_groups` label does change during a run — it is `none` until the worker groups are built, which is minutes in, and changes again when they regroup or shut down — so on 1.45 every host's startup-time `none` point would sit next to its correctly labelled one for the rest of the run, and `sum by (host.name)` would count each host twice. -An observable gauge is rebuilt from its callback on each collection cycle, so a label the callback stops reporting simply ends. The monitor therefore scrapes on its own timer into a snapshot that it *replaces* wholesale, and the callbacks return whatever is in it. Three consequences worth knowing: +An observable gauge is read on every export instead, so it always reports the latest snapshot and a label the callback stops reporting simply ends. The monitor therefore scrapes on its own timer into a snapshot that it *replaces* wholesale, and the callbacks return whatever is in it. Three consequences worth knowing: - Scraping stays on the monitor's thread rather than inside the callback, so the exporter never blocks on HTTP to an unreachable node — a wait that would otherwise hold up every other metric in the process and grow with the cluster. - The scrape and export intervals are independent, so a point can be up to one collection interval old, and scrapes taken between two export cycles are never seen. Both are normal for a gauge. @@ -236,7 +236,9 @@ It lists the worker groups resident on the host, sorted and comma-joined (`lm_po There is exactly one point per `host.name` per cycle. That matters where several raylets share a machine, as in a local multi-node cluster: each publishes the same whole-host memory on its own endpoint, so the monitor scrapes a host once and merges the worker groups of every node on it, rather than emitting a point per node that `sum by (host.name)` would add together. -That registry is an ordinary module global, so it is visible only to the process that wrote it, and the writer is always the driver. This is why the monitor has to run there too: the single-controller entry points (`run_grpo_single_controller.py`, `run_sft_v2.py`) build their worker groups on the driver but their `Logger` inside the controller actor, so a monitor owned by the `Logger` would read an empty map and label every host `none` — right numbers, wrong label, on every point. +That registry is an ordinary module global, so it is visible only to the process that wrote it. This is why the monitor has to run on the driver: the single-controller entry points (`run_grpo_single_controller.py`, `run_sft_v2.py`) build their worker groups on the driver but their `Logger` inside the controller actor, so a monitor owned by the `Logger` would read an empty map and label every host `none` — right numbers, wrong label, on every point. + +One known gap follows from the same property. `RewardModelEnvironment` is a Ray actor and builds its `Policy` inside itself, so that worker group registers in the actor's copy of the map and the driver never sees it: hosts running only the reward model report `none`. The memory figure is unaffected; only the label is incomplete. Closing it would mean sharing the map through a named actor and paying a round trip per sample, which is more than a descriptive label is worth. **Do not aggregate on it.** The value it labels is the *whole host's* memory: diff --git a/docs/observability/span-groups.md b/docs/observability/span-groups.md index c0191126929..6be9a64b153 100644 --- a/docs/observability/span-groups.md +++ b/docs/observability/span-groups.md @@ -127,7 +127,6 @@ These are set on spans for filtering — they answer "which one?" / "what kind?" | `rl.num_prompt_groups` | async rollout batch width, so a gap-filling batch is not read as an unexplained speed-up | | `rl.gym.batch_size` | how many examples one `rl.gym.run_rollouts` span covers — the NeMo-Gym counterpart to `rl.num_prompt_groups` | | `rl.idle.polls` | how many retries one `rl.idle.buffer_starvation` span covers on the single-controller path, where the span is coalesced over a poll loop. Read it against the duration: the same ten seconds is two thousand clean 5 ms polls or two hundred polls whose selection ran long, which are opposite diagnoses | -| `rl.checkpoint.async_save` | whether bytes were still being written when the loop resumed. `true` means `rl.checkpoint.finalize` spans a real wait on the writer processes; `false` means the write finished inside the staging span and the finalize span is just a rename. Both shapes emit the same span, so without this a near-zero duration is ambiguous between the two | | `rl.rollout.attempt` | SingleController dispatch attempt: `0` is a first try, `> 0` a substitution after a skipped group, whose tokens were discarded | | `rl.target_step` | the training step an `rl.sc.generate_and_push` dispatch is aimed at; omitted when the dispatch is unstamped | | `rl.critic_epochs` | critic epochs covered by one `rl.sc.value_training` span | @@ -146,17 +145,20 @@ rollout-side spans for that. ### `rl.iteration` is set once per step -Each algorithm opens its step inside +Most algorithms open their step inside `nemo_rl.telemetry.instrumentation.iteration_scope`, which stamps `rl.iteration` on every span started within it **in the driver process**. So the step umbrella, the leaves under it and the checkpoint span all carry the same value, and `rl.iteration = 412` selects the driver's whole step rather -than the step span alone. - -A span that passes `rl.iteration` itself keeps its own value — the scope is -applied by a span processor that skips keys already set. `rl..evaluate` -does exactly that, because validation also runs once before training starts, -where it reports `rl.iteration = 0`: a sentinel for "no step has run yet". +than the step span alone. `grpo_train_sync`, `async_ppo_train`, +`xtoken_off_policy_distillation_train` and `mpo_train` do not open the scope +yet, so their spans carry no `rl.iteration`. + +A span that passes `rl.iteration` itself keeps its own value — lens sets the +caller's attributes after the span starts, so they overwrite the scope's value. +`rl..evaluate` does exactly that, because validation also runs once +before training starts, where it reports `rl.iteration = 0`: a sentinel for +"no step has run yet". Two kinds of span are deliberately outside the scope. `rl..job` wraps the whole run and belongs to no single step. `rl.sc.generate_and_push` runs on the diff --git a/nemo_rl/telemetry/host_memory.py b/nemo_rl/telemetry/host_memory.py index 138f9610b56..352b4c1bfae 100644 --- a/nemo_rl/telemetry/host_memory.py +++ b/nemo_rl/telemetry/host_memory.py @@ -138,13 +138,15 @@ def _observe_total(options: CallbackOptions) -> list[Observation]: return _observations(_TOTAL_KEY) -#: Observable rather than pushed, because the attributes move: a host is -#: labelled ``rl.worker_groups="none"`` until its groups are built, and the -#: label changes again when they shut down. A synchronous gauge keeps its last -#: value per attribute set and re-exports it every cycle for the rest of the -#: run, so the startup-time point would sit next to the correct one forever and -#: ``sum by (host.name)`` would count every host twice. An observable gauge is -#: rebuilt from its callback each cycle, so the old label simply ends. +#: Observable rather than pushed, because this samples once per collection +#: interval rather than on the exporter's cycle. A synchronous gauge is wrong +#: for that either way the SDK leans: on the locked 1.44.0 it drops its value +#: at each collect, so the cycles between two samples carry no point at all, +#: and from 1.45 it holds the last value per attribute set instead, which +#: re-exports a host's startup-time ``rl.worker_groups="none"`` label long +#: after its groups are built. An observable gauge is read on every export, so +#: it always reports the latest snapshot and a label the callback stops +#: yielding simply ends. HOST_MEMORY_RECORDED_METRICS = ( RecordedMetric( RL_HOST_MEMORY_USED_METRIC, @@ -167,9 +169,12 @@ def _observe_total(options: CallbackOptions) -> list[Observation]: register_recorded_metrics(HOST_MEMORY_RECORDED_METRICS) -#: Matches ``GPUMonitoringConfig.collection_interval``, so the host-memory and +#: Used only when the config carries no ``logger.gpu_monitoring`` block, which +#: is why it is not the default of any parameter below -- a run that *does* +#: configure a cadence must not be able to silently fall back to this one. Set +#: to match ``GPUMonitoringConfig.collection_interval`` so the host-memory and #: GPU series line up on a dashboard rather than interleaving. -DEFAULT_COLLECTION_INTERVAL = 10.0 +FALLBACK_COLLECTION_INTERVAL = 10.0 #: Short, because the endpoint is a raylet on the cluster's own network and the #: loop would rather miss a sample than hold the interval open behind a node @@ -294,8 +299,9 @@ def ray_nodes() -> list[RayNode]: Returns an empty list when Ray is not up, which makes the collection loop idle rather than fail: the driver starts telemetry before ``init_ray()``. """ - # Deferred: this module is imported from nemo_rl.utils.logger at module - # scope, and must stay importable in a process that never initialises Ray. + # Deferred: this module is imported at module scope by + # nemo_rl.distributed.worker_groups, and must stay importable in a process + # that never initialises Ray. import ray if not ray.is_initialized(): @@ -357,9 +363,7 @@ class HostMemoryMonitor: and neither is worth coupling the threads to avoid. """ - def __init__( - self, collection_interval: float = DEFAULT_COLLECTION_INTERVAL - ) -> None: + def __init__(self, collection_interval: float) -> None: """Initialise the monitor. Args: @@ -454,7 +458,7 @@ def _collect_once(self) -> None: _SNAPSHOT.update(sampled) def _scrape(self, node: RayNode) -> dict[str, float]: - """Fetch and parse one node's memory series, or ``{}`` if unreachable. + """Fetch and parse one node's memory series, or ``{}`` if unusable. A node that cannot be scraped is skipped rather than retried: the next interval is the retry, and a node that is down or overloaded is exactly @@ -469,7 +473,21 @@ def _scrape(self, node: RayNode) -> dict[str, float]: f"could not scrape host memory from {node.metrics_url}", ) return {} - return parse_memory_samples(response.text) + try: + return parse_memory_samples(response.text) + except ValueError: + # Guarded separately from the fetch, and not left to the guard + # around the sweep: prometheus_client raises ValueError on a body + # that is not the exposition format -- an HTML error page from a + # proxy in front of the raylet, a truncated response -- and letting + # that reach the sweep would abandon every other node's sample and + # leave the previous snapshot exporting, stale labels included, for + # as long as the page stayed broken. + warn_once( + f"host_memory_parse:{node.host}", + f"could not parse host memory from {node.metrics_url}", + ) + return {} def _collection_loop(self) -> None: """Sample until stopped, on :attr:`collection_interval`.""" @@ -486,9 +504,7 @@ def _collection_loop(self) -> None: _MONITOR: Optional[HostMemoryMonitor] = None -def start_host_memory_monitoring( - collection_interval: float = DEFAULT_COLLECTION_INTERVAL, -) -> None: +def start_host_memory_monitoring(collection_interval: float) -> None: """Start this process's monitor, if it has none. Never raises. Called from ``init_telemetry_driver``, which is what puts the monitor in diff --git a/nemo_rl/telemetry/instrumentation.py b/nemo_rl/telemetry/instrumentation.py index 70cadcaebae..695a60177bf 100644 --- a/nemo_rl/telemetry/instrumentation.py +++ b/nemo_rl/telemetry/instrumentation.py @@ -58,10 +58,12 @@ span_attributes as _span_attributes, ) -# Span-attribute key lens declares itself. Re-exported so call sites take every -# ``rl.*`` key from this one module, and so a rename lands in lens's semconv -# instead of in a string literal scattered across the algorithms. -from nemo.lens.semconv import RL_ALGORITHM +# Which algorithm a run is: "grpo", "ppo", "dpo". Named here, not imported from +# lens's semconv, even though lens declares the same string: lens declares it +# and reads it nowhere, so importing it buys a dependency with no contract +# behind it, and RL is the only producer. Same rule as every other ``rl.*`` key +# in this module -- RL names the properties that are RL's. +RL_ALGORITHM = "rl.algorithm" from nemo_rl.telemetry.span_groups import UMBRELLA_GROUP_VALUES, RLSpanGroup from nemo_rl.telemetry.vocabulary import ( @@ -92,14 +94,6 @@ # work ran long, which are opposite diagnoses. RL_IDLE_POLLS_ATTR = "rl.idle.polls" -# Whether a checkpoint's bytes were still being written when the training loop -# moved on. Distinguishes the two shapes ``rl.checkpoint.finalize`` can take: -# with async save it spans a real wait on the writer processes, with sync save -# the write already finished inside the staging span and this is just a rename. -# Same span either way, so without this a near-zero duration is ambiguous -# between "sync save" and "async save that had already landed". -RL_CHECKPOINT_ASYNC_ATTR = "rl.checkpoint.async_save" - __all__ = [ "NO_SPAN", "managed_span", @@ -133,7 +127,6 @@ "evaluate_span", "checkpoint_finalize_span", "RL_ALGORITHM", - "RL_CHECKPOINT_ASYNC_ATTR", "RL_EFFICIENCY_CATEGORY_ATTR", "RL_IDLE_POLLS_ATTR", "RL_ITERATION_ATTR", @@ -347,10 +340,10 @@ def iteration_scope(iteration: int) -> Iterator[None]: happening in step 412" dependent on which sites remembered to pass it, and silently excludes every span added afterwards. - A span that passes the attribute itself keeps its own value: lens applies - the scope in a span processor that skips keys already set. ``evaluate_span`` - relies on that, since validation also runs once before training and reports - iteration 0 there. + A span that passes the attribute itself keeps its own value: lens starts the + span before applying the caller's attributes, so those overwrite whatever + the scope stamped on. ``evaluate_span`` relies on that, since validation + also runs once before training and reports iteration 0 there. Only spans from the provider lens builds are stamped, which is every span NeMo-RL opens. Propagates like any :class:`~contextvars.ContextVar`: to @@ -970,7 +963,6 @@ def evaluate_span( def checkpoint_finalize_span( iteration: Optional[int], *, - async_save: bool, tracer: Optional[Tracer] = None, **attributes: Any, ) -> Iterator[Any]: @@ -989,8 +981,9 @@ def checkpoint_finalize_span( empty context) and wrong by the time the span closes (the loop has moved on). Stamping it explicitly keeps the span filed under the step whose weights it holds, not whichever step happened to be running at the rename. - Set directly rather than left to the scope because lens's span processor - skips keys already present, so an explicit value wins over an inherited one. + Set directly rather than left to the scope because lens applies the caller's + attributes after the span starts, so an explicit value overwrites an + inherited one. Unbucketed, like the collector's waits and for the same reason: it overlaps the next step's productive spans, so an ``overhead`` tag here would charge @@ -999,13 +992,19 @@ def checkpoint_finalize_span( stays ``checkpoint`` so one group switch still turns every checkpoint span on or off together. + Carries no async-save flag. Whether the bytes were still in flight is not + knowable from here: every caller hands ``begin_finalization`` a wait + function regardless of the setting, because those functions no-op when + async save is off. The real setting lives per backend (Megatron's + ``mcore_state.cfg.checkpoint.async_save``, Automodel's ``is_async``), and + threading it through twelve call sites to label a span costs more than the + label is worth. Read the span's duration instead. + Args: iteration: Step the checkpoint was triggered at, or None when it cannot be recovered from the directory name. - async_save: Whether the caller handed a wait function, i.e. whether - bytes were still being written when the training loop resumed. """ - attrs: dict[str, Any] = {RL_CHECKPOINT_ASYNC_ATTR: async_save, **attributes} + attrs: dict[str, Any] = dict(attributes) if iteration is not None: attrs[RL_ITERATION_ATTR] = iteration with _managed_span( diff --git a/nemo_rl/telemetry/setup.py b/nemo_rl/telemetry/setup.py index 2dc901855d5..48d830a1ad8 100644 --- a/nemo_rl/telemetry/setup.py +++ b/nemo_rl/telemetry/setup.py @@ -214,7 +214,7 @@ def _build_resource_attributes( Process identity comes from :func:`_process_attributes`, which the callers merge in. """ - from nemo.lens.semconv import RL_ALGORITHM + from nemo_rl.telemetry.instrumentation import RL_ALGORITHM attrs: dict[str, Any] = { _CAMPAIGN_STAGE_ATTR: _campaign_stage(algorithm), @@ -249,11 +249,14 @@ def _topology_attributes(master_config: Any) -> dict[str, Any]: """Parallelism sizes, under lens's resource-scoped topology keys. ``nv.dl.topology.size.*`` rather than the ``dl.tensor_parallel.size`` - family: only the former is in lens's ``DL_RESOURCE_TYPES``, which feeds the - ``RESOURCE_TYPES`` table ``normalize_resource_attributes`` consults. A key - outside that table is passed through rather than rejected, so the ``dl.*`` - spelling would survive but arrive as an untyped string -- it is the span - and metric spelling, not the resource one. + family, as a naming convention and nothing more: the former is the + resource-scoped spelling of a parallelism size, the latter the span and + metric spelling of the same quantity. Nothing enforces it from lens's side. + ``normalize_resource_attributes`` runs only over the attributes lens + derives from ``OTEL_RESOURCE_ATTRIBUTES`` (providers.py), while the + ``resource_attributes`` handed in from here are merged as they are, so + either spelling would arrive intact and untyped. Getting it right is on + this function. Read from the enabled backend alone, under that backend's own spelling: the two disagree on names as well as values, and Automodel has no pipeline @@ -534,26 +537,43 @@ def init_telemetry_driver( # Deferred to break the cycle: host_memory reads get_telemetry_handle # from this module. from nemo_rl.telemetry.host_memory import ( - DEFAULT_COLLECTION_INTERVAL, + FALLBACK_COLLECTION_INTERVAL, start_host_memory_monitoring, ) # Here rather than in ``Logger`` beside the GPU monitor, even though the # two sample the same endpoints: the host-memory series is labelled with # the worker groups each node holds, and that map is a module global - # only the driver writes. The single-controller entry points run their - # ``Logger`` inside the controller actor, which would label every host - # as holding no group at all. + # visible only to the process that wrote it. The driver is where almost + # every worker group is built, so it is where the monitor sees the most. + # The single-controller entry points run their ``Logger`` inside the + # controller actor, which would label every host as holding no group. + # + # Not every group, though: ``RewardModelEnvironment`` is itself a Ray + # actor and builds its ``Policy`` there, so that group registers in the + # actor's copy of the map and the driver never learns about it. Its + # hosts report ``rl.worker_groups="none"``. The memory is still right -- + # the label is what is incomplete -- and making it complete means + # sharing the map across processes, which is a named actor and a round + # trip per sample for a descriptive label. Documented instead. # # Shares the GPU monitor's cadence so the two line up on a dashboard, # read from the logger config rather than restated -- but not *gated* on # ``monitor_gpus``, which decides what reaches W&B and has no business # deciding whether an OTel series exists. + # + # The fallback is for the entry points whose config carries no logger + # block at all (eval). Spelled here rather than read from + # ``GPUMonitoringConfig``: importing that means importing + # ``nemo_rl.utils.logger``, which pulls mlflow, wandb and the rest of + # the logging stack into telemetry startup for one float. interval = _dig( master_config, "logger", "gpu_monitoring", "collection_interval" ) start_host_memory_monitoring( - collection_interval=interval or DEFAULT_COLLECTION_INTERVAL + collection_interval=interval + if interval is not None + else FALLBACK_COLLECTION_INTERVAL ) # Every resolved field, not just the headline ones: the env projection uses diff --git a/nemo_rl/telemetry/vocabulary.py b/nemo_rl/telemetry/vocabulary.py index 24df7fba385..d8d3449a971 100644 --- a/nemo_rl/telemetry/vocabulary.py +++ b/nemo_rl/telemetry/vocabulary.py @@ -93,11 +93,10 @@ class RecordedMetric: description: Help text carried to the backend. callback: For the observable kinds only, what the exporter calls to read the series. Each cycle's points are exactly what it returns, - so an attribute set it stops yielding stops being exported -- which - is the reason to choose an observable kind over pushing to a gauge, - whose last value per attribute set is otherwise re-exported for the - rest of the run. Must be a function that outlives the process's - first export, not a closure built per call. + so an attribute set it stops yielding stops being exported, and it + is read on every export, where a pushed gauge reports only in the + cycles that follow a set. Must be a function that outlives the + process's first export, not a closure built per call. """ name: str diff --git a/nemo_rl/utils/checkpoint.py b/nemo_rl/utils/checkpoint.py index 912b8068506..fa26273610a 100644 --- a/nemo_rl/utils/checkpoint.py +++ b/nemo_rl/utils/checkpoint.py @@ -56,7 +56,7 @@ @contextmanager def _finalize_telemetry( - carrier: Mapping[str, str], step: Optional[int], *, async_save: bool + carrier: Mapping[str, str], step: Optional[int] ) -> Iterator[None]: """Open ``rl.checkpoint.finalize``, or carry on without it if that fails. @@ -74,7 +74,7 @@ def _finalize_telemetry( stack = ExitStack() try: stack.enter_context(remote_trace_context(carrier)) - stack.enter_context(checkpoint_finalize_span(step, async_save=async_save)) + stack.enter_context(checkpoint_finalize_span(step)) # Deliberately broad: the point is that no failure from the telemetry stack, # whatever it is, reaches the caller from here. except Exception as e: @@ -84,10 +84,12 @@ def _finalize_telemetry( "Finalization continues untraced; the checkpoint is unaffected.", stacklevel=2, ) + # Yielded outside the except block so that a rename which then also fails is + # not reported as having happened "during handling" of the telemetry error, + # which points the reader at observability instead of at the rename. A + # closed ExitStack re-enters as an empty one, so this covers both paths. + with stack: yield - else: - with stack: - yield def _load_megatron_common_state_dict(iteration_dir: Path) -> dict[str, Any]: @@ -486,7 +488,7 @@ def _finalize(): # on it, and the span is entered through a guard so the reverse # cannot happen: telemetry must not be what skips the rename. try: - with _finalize_telemetry(carrier, step, async_save=wait_fn is not None): + with _finalize_telemetry(carrier, step): if wait_fn is not None: wait_fn() self._rename_checkpoint(checkpoint_path) diff --git a/pyproject.toml b/pyproject.toml index 88b4abe577c..0ad6b3f9ebb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -81,8 +81,8 @@ dependencies = [ # git in [tool.uv.sources] — see the note there; the floor is above the # newest PyPI release so a dropped pin fails to resolve instead of silently # falling back to 0.1.0. It reads 0.2.0 rather than 0.3.0 because the pinned - # rev is the v0.3.0 tag but still declares 0.2.0 in its package_info, so a - # 0.3.0 floor would not resolve. The aiohttp extra is what stamps a traceparent on + # rev is lens main past v0.3.0, which still declares 0.2.0 in its + # package_info, so a 0.3.0 floor would not resolve. The aiohttp extra is what stamps a traceparent on # the HTTP calls NeMo-Gym's client makes, so a Gym rollout continues the # run's trace rather than starting its own. "nemo-lens[sdk,aiohttp]>=0.2.0", diff --git a/tests/unit/telemetry/test_checkpoint_spans.py b/tests/unit/telemetry/test_checkpoint_spans.py index 78961729d6f..fcaaa723288 100644 --- a/tests/unit/telemetry/test_checkpoint_spans.py +++ b/tests/unit/telemetry/test_checkpoint_spans.py @@ -33,7 +33,6 @@ ) from nemo_rl.telemetry.instrumentation import ( - RL_CHECKPOINT_ASYNC_ATTR, RL_ITERATION_ATTR, iteration_scope, managed_span, @@ -141,7 +140,6 @@ def test_the_span_is_filed_under_the_step_that_triggered_the_save( durable = _named(recorder, "rl.checkpoint.finalize") assert durable.attributes[RL_ITERATION_ATTR] == 7 - assert durable.attributes[RL_CHECKPOINT_ASYNC_ATTR] is True def test_a_checkpoint_that_never_landed_is_not_reported_as_a_clean_one( @@ -191,22 +189,3 @@ def _broken_span(*args, **kwargs): checkpointer.finalize_pending() assert (checkpoint_dir / "step_7").exists() - - -def test_a_sync_save_is_distinguishable_from_an_async_one(checkpointer, exporter): - """No ``wait_fn`` means the bytes already landed and this is just a rename. - - Both shapes emit the same span, so without the attribute a near-zero - duration could be either one. - """ - handle, recorder = exporter - tmp_dir = checkpointer.init_tmp_checkpoint(7, {"loss": 0.5}) - - checkpointer.begin_finalization(tmp_dir) - checkpointer.finalize_pending() - handle.shutdown() - - assert ( - _named(recorder, "rl.checkpoint.finalize").attributes[RL_CHECKPOINT_ASYNC_ATTR] - is False - ) diff --git a/tests/unit/telemetry/test_host_memory.py b/tests/unit/telemetry/test_host_memory.py index d3dd913d3fc..5b4890bc82a 100644 --- a/tests/unit/telemetry/test_host_memory.py +++ b/tests/unit/telemetry/test_host_memory.py @@ -14,6 +14,7 @@ import pytest import nemo_rl.telemetry.host_memory as host_memory +import nemo_rl.telemetry.vocabulary as vocabulary from nemo_rl.telemetry.host_memory import ( HOST_NAME_ATTR, NO_WORKER_GROUPS, @@ -25,6 +26,7 @@ parse_memory_samples, ray_nodes, register_worker_group_nodes, + register_worker_group_placement, start_host_memory_monitoring, stop_host_memory_monitoring, worker_groups_on_node, @@ -59,6 +61,10 @@ def _reset_host_memory_state(): host_memory._WORKER_GROUP_NODES.clear() host_memory._SNAPSHOT.clear() stop_host_memory_monitoring() + # The driver test below freezes the metric vocabulary for the rest of the + # process. Nothing in the suite registers a metric after collection today, + # but a module that did would fail here rather than where it went wrong. + vocabulary._FROZEN = False def _fake_ray(monkeypatch, nodes): @@ -122,12 +128,6 @@ def _start_exporting_telemetry(): import nemo_rl.telemetry.setup as setup_mod - instruments = pytest.importorskip("nemo.lens.instruments") - if not hasattr(instruments, "register_metric_group"): - pytest.skip("installed nemo-lens has no metric registry (lens PR #46)") - if not hasattr(instruments, "create_metric_instruments"): - pytest.skip("installed nemo-lens has no observable kinds (lens PR #84)") - reader = InMemoryMetricReader() setup_mod._TELEMETRY_HANDLE = setup_telemetry( NemoLensConfig(enabled=True), metric_reader=reader @@ -229,8 +229,80 @@ def test_a_shut_down_group_stops_claiming_its_nodes(): assert worker_groups_on_node("n0") == NO_WORKER_GROUPS +def _fake_placement_groups(monkeypatch, tables, *, raises=False): + """Install a ``placement_group_table`` answering from *tables*. + + Keyed by ``id(pg)``, like the code under test, because a real placement + group is not hashable either. + + Returns the list of placement groups it was asked about, in order, so a + test can tell one lookup from several. + """ + asked = [] + + def placement_group_table(pg): + asked.append(pg) + if raises: + raise RuntimeError("GCS is unreachable") + return {"bundles_to_node_id": tables[id(pg)]} + + ray = types.ModuleType("ray") + util = types.ModuleType("ray.util") + module = types.ModuleType("ray.util.placement_group") + module.placement_group_table = placement_group_table + util.placement_group = module + ray.util = util + # All three levels, since the import under test walks the package path. + monkeypatch.setitem(sys.modules, "ray", ray) + monkeypatch.setitem(sys.modules, "ray.util", util) + monkeypatch.setitem(sys.modules, "ray.util.placement_group", module) + return asked + + +def test_a_groups_bundles_name_the_nodes_it_is_on(monkeypatch): + """The form the worker groups actually call with: (pg, bundle index).""" + pg = SimpleNamespace(name="policy_pg") + _fake_placement_groups(monkeypatch, {id(pg): {0: "n0", 1: "n1"}}) + + register_worker_group_placement("lm_policy", [(pg, 0), (pg, 1)]) + + assert worker_groups_on_node("n0") == "lm_policy" + assert worker_groups_on_node("n1") == "lm_policy" + + +def test_a_placement_group_is_looked_up_once_however_many_bundles_it_holds( + monkeypatch, +): + """One table covers every bundle, and the lookup is a GCS round trip.""" + pg = SimpleNamespace(name="policy_pg") + asked = _fake_placement_groups(monkeypatch, {id(pg): {i: "n0" for i in range(4)}}) + + register_worker_group_placement("lm_policy", [(pg, i) for i in range(4)]) + + assert asked == [pg] + + +def test_a_bundle_the_table_does_not_place_is_left_out(monkeypatch): + """A pending bundle has no node yet; the ones that do are still labelled.""" + pg = SimpleNamespace(name="policy_pg") + _fake_placement_groups(monkeypatch, {id(pg): {0: "n0"}}) + + register_worker_group_placement("lm_policy", [(pg, 0), (pg, 1)]) + + assert worker_groups_on_node("n0") == "lm_policy" + + +def test_an_unresolvable_placement_costs_only_the_label(monkeypatch): + """Attribution is descriptive; failing to resolve it must not reach Ray.""" + pg = SimpleNamespace(name="policy_pg") + _fake_placement_groups(monkeypatch, {id(pg): {0: "n0"}}, raises=True) + + register_worker_group_placement("lm_policy", [(pg, 0)]) + + assert worker_groups_on_node("n0") == NO_WORKER_GROUPS + + def test_each_host_is_reported_with_its_name_and_its_worker_groups(monkeypatch): - pytest.importorskip("nemo.lens") reader = _start_exporting_telemetry() _fake_ray( @@ -249,7 +321,7 @@ def test_each_host_is_reported_with_its_name_and_its_worker_groups(monkeypatch): ) register_worker_group_nodes("lm_policy", ["n0"]) - HostMemoryMonitor().collect_once() + HostMemoryMonitor(collection_interval=10.0).collect_once() data = reader.get_metrics_data() used = { @@ -275,12 +347,11 @@ def test_a_label_the_scrape_stops_reporting_stops_being_exported(monkeypatch): one. ``sum by (host.name)`` would then count every host twice, and filtering on ``none`` would show every training host as idle. """ - pytest.importorskip("nemo.lens") reader = _start_exporting_telemetry() _fake_ray(monkeypatch, [_node("n0", hostname="dgx-000", address="10.0.0.1")]) _fake_scrape(monkeypatch, {"http://10.0.0.1:8080/metrics": _SCRAPE}) - monitor = HostMemoryMonitor() + monitor = HostMemoryMonitor(collection_interval=10.0) monitor.collect_once() at_startup = _gauge_points(reader.get_metrics_data(), RL_HOST_MEMORY_USED_METRIC) @@ -304,7 +375,6 @@ def test_two_raylets_on_one_host_are_one_point_scraped_once(monkeypatch): so a point per node would double count under ``sum by (host.name)``. The worker groups of both are merged, since the label describes the host. """ - pytest.importorskip("nemo.lens") reader = _start_exporting_telemetry() _fake_ray( @@ -326,7 +396,7 @@ def _counting(url, timeout=None): register_worker_group_nodes("lm_policy", ["n0"]) register_worker_group_nodes("vllm_policy", ["n1"]) - HostMemoryMonitor().collect_once() + HostMemoryMonitor(collection_interval=10.0).collect_once() points = _gauge_points(reader.get_metrics_data(), RL_HOST_MEMORY_USED_METRIC) assert len(points) == 1 @@ -343,7 +413,6 @@ def test_stopping_abandons_the_sweep_instead_of_finishing_it(monkeypatch): partial result is thrown away rather than published: it would otherwise end the series of every host the sweep had not reached. """ - pytest.importorskip("nemo.lens") _start_exporting_telemetry() nodes = [ @@ -351,7 +420,7 @@ def test_stopping_abandons_the_sweep_instead_of_finishing_it(monkeypatch): for index in range(5) ] _fake_ray(monkeypatch, nodes) - monitor = HostMemoryMonitor() + monitor = HostMemoryMonitor(collection_interval=10.0) scraped: list[str] = [] class _Response: @@ -380,7 +449,6 @@ def test_a_stopped_monitor_leaves_nothing_behind_to_export(monkeypatch): So whatever the last scrape left behind would otherwise be exported, frozen, by every cycle between shutdown and process exit. """ - pytest.importorskip("nemo.lens") reader = _start_exporting_telemetry() _fake_ray(monkeypatch, [_node("n0", hostname="dgx-000", address="10.0.0.1")]) @@ -389,7 +457,7 @@ def test_a_stopped_monitor_leaves_nothing_behind_to_export(monkeypatch): # An interval long enough that the loop samples once and then sleeps, so # the points under test are the ones collected below and not a race. start_host_memory_monitoring(collection_interval=3600.0) - HostMemoryMonitor().collect_once() + HostMemoryMonitor(collection_interval=10.0).collect_once() assert _gauge_points(reader.get_metrics_data(), RL_HOST_MEMORY_USED_METRIC) stop_host_memory_monitoring() @@ -397,9 +465,39 @@ def test_a_stopped_monitor_leaves_nothing_behind_to_export(monkeypatch): assert not _gauge_points(reader.get_metrics_data(), RL_HOST_MEMORY_USED_METRIC) +def test_an_unparsable_page_costs_only_its_own_point(monkeypatch): + """One broken page must not abandon the sweep and freeze the snapshot. + + A proxy in front of a raylet can answer 200 with an HTML error page. Left to + the guard around the sweep, the ValueError would discard every other node's + sample, and the exporter would keep publishing the previous snapshot -- + stale worker-group labels included -- for as long as the page stayed broken. + """ + reader = _start_exporting_telemetry() + + _fake_ray( + monkeypatch, + [ + _node("n0", hostname="dgx-000", address="10.0.0.1"), + _node("n1", hostname="proxied", address="10.0.0.2"), + ], + ) + _fake_scrape( + monkeypatch, + { + "http://10.0.0.1:8080/metrics": _SCRAPE, + "http://10.0.0.2:8080/metrics": "502 Bad Gateway", + }, + ) + + HostMemoryMonitor(collection_interval=10.0).collect_once() + + points = _gauge_points(reader.get_metrics_data(), RL_HOST_MEMORY_USED_METRIC) + assert [point.attributes[HOST_NAME_ATTR] for point in points] == ["dgx-000"] + + def test_an_unreachable_node_costs_only_its_own_point(monkeypatch): """The node that is wedged is exactly the one whose host memory matters.""" - pytest.importorskip("nemo.lens") reader = _start_exporting_telemetry() _fake_ray( @@ -411,7 +509,7 @@ def test_an_unreachable_node_costs_only_its_own_point(monkeypatch): ) _fake_scrape(monkeypatch, {"http://10.0.0.1:8080/metrics": _SCRAPE}) - HostMemoryMonitor().collect_once() + HostMemoryMonitor(collection_interval=10.0).collect_once() points = _gauge_points(reader.get_metrics_data(), RL_HOST_MEMORY_USED_METRIC) assert [point.attributes[HOST_NAME_ATTR] for point in points] == ["dgx-000"] @@ -427,7 +525,7 @@ def _explode(): monkeypatch.setattr(host_memory, "ray_nodes", _explode) _start_exporting_telemetry() - HostMemoryMonitor().collect_once() + HostMemoryMonitor(collection_interval=10.0).collect_once() def test_nothing_is_scraped_when_telemetry_is_not_exporting(monkeypatch): @@ -443,13 +541,13 @@ def test_nothing_is_scraped_when_telemetry_is_not_exporting(monkeypatch): lambda self, node: scraped.append(node) or {}, ) - HostMemoryMonitor().collect_once() + HostMemoryMonitor(collection_interval=10.0).collect_once() assert scraped == [] def test_stopping_a_monitor_that_was_never_started_is_harmless(): - HostMemoryMonitor().stop() + HostMemoryMonitor(collection_interval=10.0).stop() def test_a_process_runs_at_most_one_monitor(monkeypatch): @@ -489,7 +587,6 @@ def test_the_driver_starts_the_monitor_that_the_worker_groups_register_with( ``Logger`` under single-controller GRPO and SFT v2 -- would read an empty map and label every host ``none``. """ - pytest.importorskip("nemo.lens") from nemo_rl.telemetry.setup import init_telemetry_driver, shutdown_telemetry monkeypatch.setenv("NEMO_RL_OTEL_ENABLED", "true") @@ -505,8 +602,11 @@ def test_the_driver_starts_the_monitor_that_the_worker_groups_register_with( cluster=None, ) handle = init_telemetry_driver(config, algorithm="grpo") - if handle is None or not handle.is_exporting: - pytest.skip("telemetry did not come up exporting in this environment") + # Asserted rather than skipped: the lens guard at the top of this test + # already covers "no telemetry here", so a handle that comes up not + # exporting is the regression this test exists to catch, and skipping on it + # would hide exactly that. + assert handle is not None and handle.is_exporting assert host_memory._MONITOR is not None # Taken from the logger's GPU cadence rather than restated, so the two diff --git a/tests/unit/telemetry/test_instrumentation.py b/tests/unit/telemetry/test_instrumentation.py index d1bbf92e793..2e474515c25 100644 --- a/tests/unit/telemetry/test_instrumentation.py +++ b/tests/unit/telemetry/test_instrumentation.py @@ -24,7 +24,6 @@ from nemo_rl.telemetry.instrumentation import ( EFFICIENCY_CATEGORY_BUCKET, RL_BUCKET_ATTR, - RL_CHECKPOINT_ASYNC_ATTR, RL_EFFICIENCY_CATEGORY_ATTR, RL_IDLE_POLLS_ATTR, RL_ITERATION_ATTR, @@ -1231,7 +1230,7 @@ def test_checkpoint_finalize_span_outlives_the_step_that_triggered_it(): def finalize(carrier): with ( remote_trace_context(carrier), - checkpoint_finalize_span(7, async_save=True, tracer=handle.tracer), + checkpoint_finalize_span(7, tracer=handle.tracer), ): writers_done.wait(5) @@ -1274,7 +1273,7 @@ def test_checkpoint_finalize_span_is_filed_under_the_triggering_step(): def finalize(carrier): with ( remote_trace_context(carrier), - checkpoint_finalize_span(7, async_save=True, tracer=handle.tracer), + checkpoint_finalize_span(7, tracer=handle.tracer), ): started.set() release.wait(5) @@ -1291,7 +1290,6 @@ def finalize(carrier): durable = _named(exporter, "rl.checkpoint.finalize") assert durable.attributes[RL_ITERATION_ATTR] == 7 - assert durable.attributes[RL_CHECKPOINT_ASYNC_ATTR] is True @requires_lens @@ -1302,7 +1300,7 @@ def test_checkpoint_finalize_span_is_parented_to_the_run(): def finalize(carrier): with ( remote_trace_context(carrier), - checkpoint_finalize_span(7, async_save=False, tracer=handle.tracer), + checkpoint_finalize_span(7, tracer=handle.tracer), ): pass @@ -1327,7 +1325,7 @@ def test_checkpoint_finalize_span_carries_no_bucket(): RLSpanGroup.CHECKPOINT, "rl.grpo.checkpointing", tracer=handle.tracer ): pass - with checkpoint_finalize_span(7, async_save=True, tracer=handle.tracer): + with checkpoint_finalize_span(7, tracer=handle.tracer): pass handle.shutdown() @@ -1340,7 +1338,7 @@ def test_checkpoint_finalize_span_carries_no_bucket(): def test_checkpoint_finalize_span_follows_the_checkpoint_group(): """One switch turns staging and finalization on or off together.""" handle, exporter = _setup("generation") - with checkpoint_finalize_span(7, async_save=True, tracer=handle.tracer): + with checkpoint_finalize_span(7, tracer=handle.tracer): pass handle.shutdown() From 94fe5b8d23029317b6f7fa0189143805bc328c8a Mon Sep 17 00:00:00 2001 From: Raj Singh Date: Tue, 13 Oct 2026 17:51:30 -0700 Subject: [PATCH 11/17] refactor(telemetry): type each shared name once, trim the long comments Five more review threads, none of them behavioural except the parse. parse_memory_samples ran prometheus_client over the whole raylet page to keep two series. Dropping the lines that cannot match first is the same result for ~28x less work: 143 ms to 5 ms on a synthetic 48k-line page in this venv. The filter is a prefix test, so it keeps near misses like ray_node_mem_used_percent, which the name lookup then drops -- that series is now in the test fixture so the two steps stay honest together. Names typed twice. Ray's own series names were spelled here and again in RayGpuMonitorLogger, where renaming one side for a Ray bump would leave this module matching nothing and its series silently empty; they are RAY_NODE_MEM_USED / RAY_NODE_MEM_TOTAL now, imported by logger.py. The same for NRL_WORKER_GROUP, which worker_groups.py exports and telemetry reads back, and for the two efficiency keys algorithms/utils.py writes and metrics.py reads: all three are public in vocabulary.py, which imports nothing from nemo_rl and so cannot close a cycle. "init/total" inside WALL_CLOCK_EFFICIENCY_CATEGORIES stays a literal on purpose: that list is the canonical declaration, and the telemetry suite parses it as string literals to check the constants that restate it. Comments. Eight blocks of five or more lines inside function bodies, cut to the line that says what breaks: the design reasoning is in metrics.md and in the docstrings, and a call-site paragraph only drifts from it. configuration.md now says the software versions are per process, so transformer_engine missing from a driver, controller or vLLM span does not mean the run had none, and that SGLang, TRT-LLM and Dynamo workers report no versions at all. Signed-off-by: Raj Singh --- docs/observability/configuration.md | 4 +- nemo_rl/algorithms/grpo.py | 11 ++--- nemo_rl/algorithms/ppo.py | 12 +++-- nemo_rl/algorithms/utils.py | 13 ++++-- nemo_rl/distributed/worker_groups.py | 3 +- nemo_rl/telemetry/host_memory.py | 53 ++++++++++------------ nemo_rl/telemetry/metrics.py | 8 ++-- nemo_rl/telemetry/setup.py | 57 +++++++----------------- nemo_rl/telemetry/vocabulary.py | 10 +++++ nemo_rl/utils/checkpoint.py | 8 +--- nemo_rl/utils/logger.py | 5 ++- tests/unit/telemetry/test_host_memory.py | 9 +++- 12 files changed, 95 insertions(+), 98 deletions(-) diff --git a/docs/observability/configuration.md b/docs/observability/configuration.md index b4b87202f73..30767a543a2 100644 --- a/docs/observability/configuration.md +++ b/docs/observability/configuration.md @@ -100,8 +100,8 @@ Each process sets stable-for-the-run values on its OTel `Resource` — `init_tel | `nv.dl.training.config.optimizer` | on Megatron `policy.megatron_cfg.optimizer.optimizer` (`adam`), otherwise `policy.optimizer.name` (`torch.optim.AdamW`) — the backends both spell it and store it differently, and only the enabled one is read | | `nv.dl.training.config.recompute_granularity` | Megatron runs only: `policy.megatron_cfg.recompute_granularity`, and only when `activation_checkpointing` is on. Megatron ignores the granularity otherwise, so reporting it would describe recompute that is not happening | | `nv.dl.training.target.train_iters` | `.max_num_steps`, when positive. `sft_v2` reads the `sft` section, which is where it keeps its settings. `rm` spells "train for one epoch" as `-1`, which is omitted rather than reported as a target | -| `nv.dl.software.torch`, `.cuda`, `.nccl` | versions as loaded in this process, not as pinned in the lockfile. CUDA and NCCL are omitted on a CPU-only build | -| `nv.dl.software.transformer_engine` | the installed distribution's version, read from package metadata rather than by importing it — importing transformer-engine loads its PyTorch extension in every process that sets up telemetry. Omitted when it is not installed | +| `nv.dl.software.torch`, `.cuda`, `.nccl` | versions as loaded in the process that opened the span, not as pinned in the lockfile — they are per process, so filter on `rl.worker_group` to read one backend's versions. CUDA and NCCL are omitted on a CPU-only build. SGLang, TRT-LLM and Dynamo workers do not report versions yet | +| `nv.dl.software.transformer_engine` | the installed distribution's version, read from package metadata rather than by importing it — importing transformer-engine loads its PyTorch extension in every process that sets up telemetry. Omitted when it is not installed, which is every process except Megatron and Automodel workers: its absence on a driver, controller or vLLM span does not mean the run lacked TE | | `nv.dl.rank`, `nv.dl.world_size` | this process's rank and group size (`RANK` / `WORLD_SIZE`, or `0` / `1` for the driver and singleton actors) | | `nv.dl.local_rank` | worker processes only: `LOCAL_RANK`, which `RayWorkerGroup` exports beside `RANK` as the worker's bundle index in its placement group. Node-local only when each node has its own placement group — a generation backend spanning nodes requests a single unified one (SGLang always; vLLM, TRT-LLM and Megatron when their parallelism crosses nodes) and the index then runs across nodes. Not a physical device index either, since Ray gives each worker its own `CUDA_VISIBLE_DEVICES` | | `rl.worker_group` | worker processes only: the worker group's `name_prefix` (`lm_policy`, `vllm_policy`, ...), from `NRL_WORKER_GROUP`. The `NemoGym` actor reports the literal `nemo_gym` — it is built from the environment registry rather than by `RayWorkerGroup`, so nothing sets the env var for it | diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index ff54db6191d..56b632a1af2 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -165,6 +165,7 @@ ) from nemo_rl.telemetry.setup import get_telemetry_handle from nemo_rl.telemetry.span_groups import RLSpanGroup +from nemo_rl.telemetry.vocabulary import INIT_TOTAL_CATEGORY, WORKER_GROUP_ENV from nemo_rl.utils.checkpoint import ( CheckpointingConfig, CheckpointManager, @@ -5045,7 +5046,7 @@ def async_grpo_train( "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector", extra_env_vars={ # Names this actor's spans the way RayWorkerGroup names its groups'. - "NRL_WORKER_GROUP": "trajectory_collector", + WORKER_GROUP_ENV: "trajectory_collector", }, ) @@ -5221,8 +5222,8 @@ def _flush_collector_telemetry() -> None: # rather than only as a scalar. Unbucketed by construction (see # UNBUCKETED_SPAN_CATEGORIES): the generation fleet is busy for this whole # window, and its spans join this trace. - with efficiency_span("init/total", tracer=_tracer): - timer.start("init/total") + with efficiency_span(INIT_TOTAL_CATEGORY, tracer=_tracer): + timer.start(INIT_TOTAL_CATEGORY) wait_iterations = 0 while True: buffer_size_current = ray.get(replay_buffer.size.remote()) @@ -5286,7 +5287,7 @@ def _flush_collector_telemetry() -> None: # Retained because the per-step timer.reset() below discards it; the # efficiency snapshot re-supplies it every step. - init_total_s = timer.stop("init/total") + init_total_s = timer.stop(INIT_TOTAL_CATEGORY) print(f"✅ Buffer ready for step {step}! Starting training loop...") ft_save_period = master_config.checkpointing.get("ft_save_period") @@ -6246,7 +6247,7 @@ def _flush_collector_telemetry() -> None: # at the end of every step drops it -- so re-supply the captured # value, or the series reports the real startup cost at step 1 and # zero for the rest of the run. - driver_efficiency["init/total"] = init_total_s + driver_efficiency[INIT_TOTAL_CATEGORY] = init_total_s merged_efficiency = {**driver_efficiency} for cat, dur in collector_efficiency.items(): merged_efficiency[cat] = merged_efficiency.get(cat, 0.0) + dur diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index f474509138a..970c2465c03 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -113,7 +113,11 @@ ) from nemo_rl.telemetry.setup import get_telemetry_handle from nemo_rl.telemetry.span_groups import RLSpanGroup -from nemo_rl.telemetry.vocabulary import TeedMetric, register_teed_metrics +from nemo_rl.telemetry.vocabulary import ( + INIT_TOTAL_CATEGORY, + TeedMetric, + register_teed_metrics, +) from nemo_rl.utils.checkpoint import ( CheckpointingConfig, CheckpointManager, @@ -2518,7 +2522,7 @@ def _raise_if_collector_stopped(waiting_for: str) -> None: print("📦 Started continuous background trajectory collection") print(f"⏳ Waiting for replay buffer to be ready for step {step}...") - timer.start("init/total") + timer.start(INIT_TOTAL_CATEGORY) wait_iterations = 0 while True: current_step_ready = ray.get( @@ -2561,7 +2565,7 @@ def _raise_if_collector_stopped(waiting_for: str) -> None: _raise_if_collector_stopped("waiting for the initial replay batch") wait_iterations += 1 time.sleep(1.0) - timer.stop("init/total") + timer.stop(INIT_TOTAL_CATEGORY) print(f"✅ Buffer ready for step {step}! Starting async PPO training loop...") except Exception: _shutdown_workers(propagate_checkpoint_error=False) @@ -3175,7 +3179,7 @@ def _raise_if_collector_stopped(waiting_for: str) -> None: driver_efficiency = { cat: timer.reduce(cat, "sum") for cat in [ - "init/total", + INIT_TOTAL_CATEGORY, "idle/buffer_starvation", "idle/refit_bubble", "idle/validation", diff --git a/nemo_rl/algorithms/utils.py b/nemo_rl/algorithms/utils.py index 062e3e0e67c..4382d7f7637 100644 --- a/nemo_rl/algorithms/utils.py +++ b/nemo_rl/algorithms/utils.py @@ -32,7 +32,11 @@ should_use_deepseek_v4_chat_template, ) from nemo_rl.models.policy import TokenizerConfig -from nemo_rl.telemetry.vocabulary import RUN_WINDOW_WALL_CLOCK_CATEGORIES +from nemo_rl.telemetry.vocabulary import ( + EFFICIENCY_PCT_KEY, + EFFICIENCY_PCT_PER_STEP_KEY, + RUN_WINDOW_WALL_CLOCK_CATEGORIES, +) from nemo_rl.utils.fastokens import maybe_patch_fastokens from nemo_rl.utils.logger import Logger @@ -1014,6 +1018,9 @@ def log_generation_metrics( WALL_CLOCK_EFFICIENCY_CATEGORIES = [ + # Spelled out rather than INIT_TOTAL_CATEGORY: this list is the canonical + # declaration, and the telemetry suite parses it as string literals to keep + # the constants that restate it honest. "init/total", "idle/buffer_starvation", "idle/refit_bubble", @@ -1130,13 +1137,13 @@ def print_efficiency_summary( loggable["efficiency/thread_seconds_total_s"] = thread_seconds_total loggable["efficiency/total_waste_s"] = wall_waste loggable["efficiency/productive_time_s"] = productive - loggable["efficiency/efficiency_pct"] = efficiency_pct + loggable[EFFICIENCY_PCT_KEY] = efficiency_pct loggable["efficiency/total_wall_time_s"] = total_wall_time_s # Carries the window of the ratio above to the OTel tee, which tags it. A # caller that supplies no per-step denominator gets a run-cumulative ratio, # and publishing that as per-step would state the opposite of what it is. # A float, not the string it represents, because everything in this dict is # also logged to WandB/TensorBoard as a scalar. - loggable["efficiency/efficiency_pct_is_per_step"] = float(pct_is_per_step) + loggable[EFFICIENCY_PCT_PER_STEP_KEY] = float(pct_is_per_step) return loggable diff --git a/nemo_rl/distributed/worker_groups.py b/nemo_rl/distributed/worker_groups.py index 959d0554709..3f0e735d2d5 100644 --- a/nemo_rl/distributed/worker_groups.py +++ b/nemo_rl/distributed/worker_groups.py @@ -35,6 +35,7 @@ forget_worker_group_nodes, register_worker_group_placement, ) +from nemo_rl.telemetry.vocabulary import WORKER_GROUP_ENV from nemo_rl.utils.venvs import ( add_hf_modules_cache_to_pythonpath, create_local_venv_on_each_node, @@ -632,7 +633,7 @@ def _create_workers_from_bundle_indices( # RANK is group-local, so it alone cannot tell a policy # worker from a generation worker. Observability consumers # need the group to disambiguate them. - "NRL_WORKER_GROUP": self.name_prefix, + WORKER_GROUP_ENV: self.name_prefix, } ) # Remove Ray-specific environment variables, let the worker itself set them. diff --git a/nemo_rl/telemetry/host_memory.py b/nemo_rl/telemetry/host_memory.py index 352b4c1bfae..5b04c8dc5eb 100644 --- a/nemo_rl/telemetry/host_memory.py +++ b/nemo_rl/telemetry/host_memory.py @@ -88,12 +88,18 @@ #: been built yet when the sample is taken. NO_WORKER_GROUPS = "none" +#: Ray's own series names, public because ``RayGpuMonitorLogger`` reads the +#: same two for the W&B path. Typed once: renaming only one side would leave +#: this module matching nothing and its series silently empty. +RAY_NODE_MEM_USED = "ray_node_mem_used" +RAY_NODE_MEM_TOTAL = "ray_node_mem_total" + #: Prometheus series published by each raylet, and the key each is recorded #: against. Bytes already, despite the GB the W&B path reports: that divides by #: 1024**3 on the way out. _SERIES_BY_RAY_NAME = { - "ray_node_mem_used": RL_HOST_MEMORY_USED_METRIC, - "ray_node_mem_total": RL_HOST_MEMORY_TOTAL_METRIC, + RAY_NODE_MEM_USED: RL_HOST_MEMORY_USED_METRIC, + RAY_NODE_MEM_TOTAL: RL_HOST_MEMORY_TOTAL_METRIC, } _USED_KEY = registry_key(RL_HOST_MEMORY_USED_METRIC) @@ -339,7 +345,15 @@ def parse_memory_samples(metrics_text: str) -> dict[str, float]: hundreds, and only the two in :data:`_SERIES_BY_RAY_NAME` are wanted. """ values: dict[str, float] = {} - for family in text_string_to_metric_families(metrics_text): + # Parsing costs ~16 us a line and only two series are wanted, so drop the + # other lines first; each kept series keeps its HELP/TYPE lines. + wanted = tuple(_SERIES_BY_RAY_NAME) + kept = "\n".join( + line + for line in metrics_text.splitlines() + if line.removeprefix("# HELP ").removeprefix("# TYPE ").startswith(wanted) + ) + for family in text_string_to_metric_families(kept): for sample in family.samples: name = _SERIES_BY_RAY_NAME.get(sample.name) if name is None: @@ -415,24 +429,13 @@ def _collect_once(self) -> None: handle = get_telemetry_handle() if handle is None or not handle.is_exporting: return - # Scraped on this thread rather than inside the callbacks so that the - # exporter never blocks on HTTP: a callback that waits on an - # unreachable node would hold up every other metric in the process, and - # the wait scales with the size of the cluster. - # Accumulated per host rather than per node, because the number is the - # host's: several raylets on one box -- a local multi-node setup -- each - # publish the same whole-host memory on their own endpoint. Keyed by - # node that would be two points carrying one host.name, which double - # counts under `sum by (host.name)` and contradicts the one-point-per- - # host reading the attribute is documented to have. + # Scraped here, not in the callbacks, so the exporter never waits on HTTP. + # Keyed by host, not node: raylets on one box all report the same memory. values_by_host: dict[str, dict[str, float]] = {} groups_by_host: dict[str, set[str]] = {} for node in ray_nodes(): - # Between nodes, so stopping costs the scrape in flight rather than - # the rest of the cluster: this loop is one scrape timeout per node - # and runs them in sequence, which on a large cluster is minutes. - # Returning without publishing is the point -- a partial sweep would - # end the series of every host it had not reached yet. + # Checked between nodes; return without publishing, or unreached + # hosts' series end. if self._stop.is_set(): return groups_by_host.setdefault(node.host, set()).update( @@ -448,12 +451,8 @@ def _collect_once(self) -> None: for host, values in values_by_host.items() } with _SNAPSHOT_LOCK: - # Replaced, not merged: a host that has dropped out of the cluster - # or changed its worker-group label must stop being reported under - # the old one, and leaving the old key in would be the stale series - # this design exists to avoid. The cost is that a node which fails - # one scrape has a gap until the next, rather than a flat line at - # its last known value -- which is the honest reading anyway. + # Replaced, not merged, so a host that left or changed label stops + # reporting. _SNAPSHOT.clear() _SNAPSHOT.update(sampled) @@ -552,10 +551,6 @@ def stop_host_memory_monitoring() -> None: monitor.stop() except Exception: warn_once("host_memory_stop", "could not stop host memory monitoring") - # Outside the guard: the callbacks stay registered for the life of the - # process -- OTel has no way to detach an observable instrument -- so a - # snapshot left behind here would keep being exported by every cycle after - # the monitor that produced it stopped, which is exactly the frozen series - # the observable kinds were chosen to avoid. + # Outside the guard: the callbacks outlive the monitor and would export this. with _SNAPSHOT_LOCK: _SNAPSHOT.clear() diff --git a/nemo_rl/telemetry/metrics.py b/nemo_rl/telemetry/metrics.py index cc3ae971fbe..515f9ad66bb 100644 --- a/nemo_rl/telemetry/metrics.py +++ b/nemo_rl/telemetry/metrics.py @@ -51,6 +51,8 @@ ) from nemo_rl.telemetry.setup import get_telemetry_handle from nemo_rl.telemetry.vocabulary import ( + EFFICIENCY_PCT_KEY, + EFFICIENCY_PCT_PER_STEP_KEY, RUN_WINDOW_WALL_CLOCK_CATEGORIES, as_scalar, freeze_metrics, @@ -123,8 +125,6 @@ # Keys that ``print_efficiency_summary`` puts in the Logger dict. _EFFICIENCY_KEY_PREFIX = "efficiency/" _EFFICIENCY_SECONDS_SUFFIX = "_s" -_EFFICIENCY_PCT_KEY = "efficiency/efficiency_pct" -_EFFICIENCY_PCT_PER_STEP_KEY = "efficiency/efficiency_pct_is_per_step" # --- lens metric registry ------------------------------------------------- # @@ -448,7 +448,7 @@ def _tee_efficiency_metrics(meter: Meter, metrics: dict[str, Any]) -> None: """Emit the ``efficiency/*`` phase durations as ``rl.efficiency.*``.""" measurement_by_category = efficiency_measurements() seconds = map_efficiency_seconds(metrics, measurement_by_category) - pct = as_scalar(metrics.get(_EFFICIENCY_PCT_KEY)) + pct = as_scalar(metrics.get(EFFICIENCY_PCT_KEY)) if not seconds and pct is None: return @@ -472,7 +472,7 @@ def _tee_efficiency_metrics(meter: Meter, metrics: dict[str, Any]) -> None: # per-step one -- and defaulting to the run window when the flag is # absent, since mislabelling a run ratio as per-step is the harmful # direction. - is_per_step = as_scalar(metrics.get(_EFFICIENCY_PCT_PER_STEP_KEY)) + is_per_step = as_scalar(metrics.get(EFFICIENCY_PCT_PER_STEP_KEY)) _record( meter, {_EFFICIENCY_PCT_KEY_REGISTRY: pct}, diff --git a/nemo_rl/telemetry/setup.py b/nemo_rl/telemetry/setup.py index 48d830a1ad8..8a57e8e715c 100644 --- a/nemo_rl/telemetry/setup.py +++ b/nemo_rl/telemetry/setup.py @@ -47,6 +47,10 @@ import uuid from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, cast +# The one nemo_rl import at module scope: vocabulary imports nothing from +# nemo_rl, so it cannot close a cycle with the modules that import this one. +from nemo_rl.telemetry.vocabulary import WORKER_GROUP_ENV + if TYPE_CHECKING: from nemo.lens import TelemetryHandle @@ -65,9 +69,6 @@ _OTEL_FALLBACK_PREFIX = "NEMO_LENS" _RUN_ID_ENV = f"{_OTEL_PREFIX}_RUN_ID" -# Set per worker by ``RayWorkerGroup`` from the group's ``name_prefix``. -_WORKER_GROUP_ENV = "NRL_WORKER_GROUP" - # Reported as ``nv.dl.provider.name``: which training framework produced this # telemetry, for a backend collecting several of them. Spelled like the default # service name, but separate from it -- a user is free to rename the service. @@ -384,11 +385,8 @@ def _software_versions() -> dict[str, Any]: else: attrs[NV_DL_SOFTWARE_NCCL] = ".".join(str(part) for part in nccl_version) - # Read from the installed distribution's metadata rather than by importing - # the package. Importing transformer-engine pulls in its PyTorch extension, - # which is seconds of work and a chunk of memory in every process that sets - # up telemetry, and which raises ``OSError`` rather than ``ImportError`` - # when the shared library will not load against the driver present. + # Metadata, not import: importing transformer-engine is slow and can raise + # OSError. try: te_version = importlib.metadata.version("transformer-engine") except importlib.metadata.PackageNotFoundError: @@ -441,12 +439,8 @@ def _process_attributes(rank: int, world_size: int, run_id: str = "") -> dict[st local_rank = os.environ.get("LOCAL_RANK", "").strip() if local_rank.isdigit(): attrs[NV_DL_LOCAL_RANK] = int(local_rank) - # Derived rather than set to ``run_id`` itself, which is already reported - # as lens's ``nemo.run.id``. ``nv.dl.run.uuid`` is a uuid5 over the - # scheduler's identity for this attempt, and a backend grouping RL runs - # with Megatron ones joins on that; a raw run id in the same key would - # match nothing. Returns None off a scheduler with no run id to fall back - # on, which cannot happen here but is handled rather than asserted. + # uuid5 of the run id, not the raw id: backends join RL and Megatron runs + # on it. run_uuid = derive_nv_dl_run_uuid(run_id=run_id) if run_id else None if run_uuid: attrs[NV_DL_RUN_UUID] = run_uuid @@ -541,32 +535,13 @@ def init_telemetry_driver( start_host_memory_monitoring, ) - # Here rather than in ``Logger`` beside the GPU monitor, even though the - # two sample the same endpoints: the host-memory series is labelled with - # the worker groups each node holds, and that map is a module global - # visible only to the process that wrote it. The driver is where almost - # every worker group is built, so it is where the monitor sees the most. - # The single-controller entry points run their ``Logger`` inside the - # controller actor, which would label every host as holding no group. - # - # Not every group, though: ``RewardModelEnvironment`` is itself a Ray - # actor and builds its ``Policy`` there, so that group registers in the - # actor's copy of the map and the driver never learns about it. Its - # hosts report ``rl.worker_groups="none"``. The memory is still right -- - # the label is what is incomplete -- and making it complete means - # sharing the map across processes, which is a named actor and a round - # trip per sample for a descriptive label. Documented instead. - # - # Shares the GPU monitor's cadence so the two line up on a dashboard, - # read from the logger config rather than restated -- but not *gated* on - # ``monitor_gpus``, which decides what reaches W&B and has no business - # deciding whether an OTel series exists. + # Here, not in ``Logger``: the worker-group map is a module global, and + # the driver is the process that builds almost every group. Almost -- + # ``RewardModelEnvironment`` builds its ``Policy`` in its own actor, so + # its hosts report ``rl.worker_groups="none"``; see metrics.md. # - # The fallback is for the entry points whose config carries no logger - # block at all (eval). Spelled here rather than read from - # ``GPUMonitoringConfig``: importing that means importing - # ``nemo_rl.utils.logger``, which pulls mlflow, wandb and the rest of - # the logging stack into telemetry startup for one float. + # Not gated on monitor_gpus: that decides what reaches W&B, not OTel. + # The fallback covers the entry points with no logger block at all. interval = _dig( master_config, "logger", "gpu_monitoring", "collection_interval" ) @@ -620,7 +595,7 @@ def _worker_resource_attributes( _CAMPAIGN_STAGE_ATTR: os.environ.get(_CAMPAIGN_STAGE_ENV, "").strip() or _DEFAULT_CAMPAIGN_STAGE } - group = worker_group or os.environ.get(_WORKER_GROUP_ENV, "").strip() + group = worker_group or os.environ.get(WORKER_GROUP_ENV, "").strip() if group: attrs[_WORKER_GROUP_ATTR] = group if extra: @@ -696,7 +671,7 @@ def init_telemetry_worker( handle = setup_telemetry(config, resource_attributes=attrs) logger.info( "nemo-lens worker telemetry initialised (group=%s, rank=%s/%s, exporting=%s)", - os.environ.get(_WORKER_GROUP_ENV, "?"), + os.environ.get(WORKER_GROUP_ENV, "?"), rank, world_size, handle.is_exporting, diff --git a/nemo_rl/telemetry/vocabulary.py b/nemo_rl/telemetry/vocabulary.py index d8d3449a971..a9a5b186c64 100644 --- a/nemo_rl/telemetry/vocabulary.py +++ b/nemo_rl/telemetry/vocabulary.py @@ -44,6 +44,16 @@ #: Efficiency category covering startup, measured once before the training loop. INIT_TOTAL_CATEGORY = "init/total" +#: Env var ``RayWorkerGroup`` exports so a worker's spans carry +#: ``rl.worker_group``; telemetry reads it back when setting a worker up. +WORKER_GROUP_ENV = "NRL_WORKER_GROUP" + +#: Keys ``algorithms/utils.py`` writes the efficiency ratio under and telemetry +#: reads back. Named here, not at either end, because a rename on one side +#: alone would leave the metric reading nothing with no error. +EFFICIENCY_PCT_KEY = "efficiency/efficiency_pct" +EFFICIENCY_PCT_PER_STEP_KEY = "efficiency/efficiency_pct_is_per_step" + # Wall-clock categories whose value covers the whole run rather than one step. # The driver's Timer is reset every step, so its idle categories are per-step # deltas -- but init/total is measured once before the loop and republished diff --git a/nemo_rl/utils/checkpoint.py b/nemo_rl/utils/checkpoint.py index fa26273610a..0e128096c60 100644 --- a/nemo_rl/utils/checkpoint.py +++ b/nemo_rl/utils/checkpoint.py @@ -468,12 +468,8 @@ def begin_finalization( self.finalize_pending() self._pending_checkpoint_path = Path(checkpoint_path) self._finalize_error = None - # The step reaches here only in the staging directory's name, which - # ``init_tmp_checkpoint`` formats as ``tmp_step_{step}``; read back the - # same way :meth:`_rename_checkpoint` reads it, rather than adding a - # second spelling of the same parse to this file. Guarded and dropped - # to None on anything unexpected, since this only labels a span and a - # checkpoint staged under some other name must still be finalized. + # The step is only in the staging dir name (tmp_step_); None if it + # does not parse, since this only labels a span. try: step = int(Path(checkpoint_path).name.split("_")[2]) except (IndexError, ValueError): diff --git a/nemo_rl/utils/logger.py b/nemo_rl/utils/logger.py index 7d9e45a3e9f..faf27f0c395 100644 --- a/nemo_rl/utils/logger.py +++ b/nemo_rl/utils/logger.py @@ -45,6 +45,7 @@ from nemo_rl.data.interfaces import LLMMessageLogType from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.experience.metric_utils import is_histogram_metric +from nemo_rl.telemetry.host_memory import RAY_NODE_MEM_TOTAL, RAY_NODE_MEM_USED from nemo_rl.telemetry.metrics import tee_rl_metrics_to_otel from nemo_rl.utils.train_data_dump import json_safe_nonfinite @@ -905,10 +906,10 @@ def _parse_metric(self, sample: Sample, node_idx: int) -> dict[str, Any]: metric_name = f"node.{node_idx}.gpu.{index}.mem_gb" # NOTE: It appears their docs say bytes, but it appears to be MB value /= 1024 - elif metric_name == "ray_node_mem_used": + elif metric_name == RAY_NODE_MEM_USED: metric_name = f"node.{node_idx}.mem_gb" value /= 1024 * 1024 * 1024 - elif metric_name == "ray_node_mem_total": + elif metric_name == RAY_NODE_MEM_TOTAL: metric_name = f"node.{node_idx}.mem_total_gb" value /= 1024 * 1024 * 1024 else: diff --git a/tests/unit/telemetry/test_host_memory.py b/tests/unit/telemetry/test_host_memory.py index 5b4890bc82a..d4a214fe3dd 100644 --- a/tests/unit/telemetry/test_host_memory.py +++ b/tests/unit/telemetry/test_host_memory.py @@ -48,6 +48,9 @@ # HELP ray_node_cpu_utilization CPU percentage. # TYPE ray_node_cpu_utilization gauge ray_node_cpu_utilization{{SessionName="s"}} 42.0 +# HELP ray_node_mem_used_percent Share of memory in use. +# TYPE ray_node_mem_used_percent gauge +ray_node_mem_used_percent{{SessionName="s"}} 3.0 """ @@ -141,7 +144,11 @@ def _start_exporting_telemetry(): def test_only_the_two_node_memory_series_are_read(): - """A raylet publishes hundreds of series; all but two are noise here.""" + """A raylet publishes hundreds of series; all but two are noise here. + + ``ray_node_mem_used_percent`` is the near miss that matters: the parser + drops most lines on a prefix before parsing them, and that prefix keeps it. + """ assert parse_memory_samples(_SCRAPE) == { "host_memory_used": _USED_BYTES, "host_memory_total": _TOTAL_BYTES, From 882bbc4d064553e86c8f7088cadfff33b62da73d Mon Sep 17 00:00:00 2001 From: Raj Singh Date: Tue, 13 Oct 2026 17:58:41 -0700 Subject: [PATCH 12/17] refactor(telemetry): look the step target up by algorithm name _CONFIG_SECTION_BY_ALGORITHM existed for one entry: sft_v2 reports itself as sft_v2 while reading master_config.sft, so looking its settings up by algorithm name found nothing. One map entry is cheap; a table of which section each entry point lays its config out under is not, because every new entry point has to remember to extend it and nothing fails if it does not. Gone, along with _config_section and the test that pinned the mapping. The lookup is by algorithm name now, so an algorithm whose settings live under another section reports no training target: _dig returns None on the missing section and the isinstance check drops it. Checked that this is the only effect -- sft_v2 omits nv.dl.training.target.train_iters, and sft, grpo, dpo and rm are unchanged, rm still omitting its -1. Neither rl.algorithm nor the campaign stage reads the section. Not replaced with a config_section argument from run_sft_v2.py: that is API for an attribute that is a placeholder cap in several configs anyway. Signed-off-by: Raj Singh --- docs/observability/configuration.md | 2 +- nemo_rl/telemetry/setup.py | 23 ++++++----------------- tests/unit/telemetry/test_setup.py | 11 ----------- 3 files changed, 7 insertions(+), 29 deletions(-) diff --git a/docs/observability/configuration.md b/docs/observability/configuration.md index 30767a543a2..6391192236a 100644 --- a/docs/observability/configuration.md +++ b/docs/observability/configuration.md @@ -99,7 +99,7 @@ Each process sets stable-for-the-run values on its OTel `Resource` — `init_tel | `nv.dl.training.config.sequence_length` | `policy.max_total_sequence_length` | | `nv.dl.training.config.optimizer` | on Megatron `policy.megatron_cfg.optimizer.optimizer` (`adam`), otherwise `policy.optimizer.name` (`torch.optim.AdamW`) — the backends both spell it and store it differently, and only the enabled one is read | | `nv.dl.training.config.recompute_granularity` | Megatron runs only: `policy.megatron_cfg.recompute_granularity`, and only when `activation_checkpointing` is on. Megatron ignores the granularity otherwise, so reporting it would describe recompute that is not happening | -| `nv.dl.training.target.train_iters` | `.max_num_steps`, when positive. `sft_v2` reads the `sft` section, which is where it keeps its settings. `rm` spells "train for one epoch" as `-1`, which is omitted rather than reported as a target | +| `nv.dl.training.target.train_iters` | `.max_num_steps`, when positive. Looked up by algorithm name, so an entry point that keeps its settings under a differently named section reports no target at all — `sft_v2` reads `master_config.sft`, and omits this. `rm` spells "train for one epoch" as `-1`, which is omitted rather than reported as a target | | `nv.dl.software.torch`, `.cuda`, `.nccl` | versions as loaded in the process that opened the span, not as pinned in the lockfile — they are per process, so filter on `rl.worker_group` to read one backend's versions. CUDA and NCCL are omitted on a CPU-only build. SGLang, TRT-LLM and Dynamo workers do not report versions yet | | `nv.dl.software.transformer_engine` | the installed distribution's version, read from package metadata rather than by importing it — importing transformer-engine loads its PyTorch extension in every process that sets up telemetry. Omitted when it is not installed, which is every process except Megatron and Automodel workers: its absence on a driver, controller or vLLM span does not mean the run lacked TE | | `nv.dl.rank`, `nv.dl.world_size` | this process's rank and group size (`RANK` / `WORLD_SIZE`, or `0` / `1` for the driver and singleton actors) | diff --git a/nemo_rl/telemetry/setup.py b/nemo_rl/telemetry/setup.py index 8a57e8e715c..905b932b1f4 100644 --- a/nemo_rl/telemetry/setup.py +++ b/nemo_rl/telemetry/setup.py @@ -105,19 +105,6 @@ def _campaign_stage(algorithm: str) -> str: return _CAMPAIGN_STAGE_BY_ALGORITHM.get(algorithm, _DEFAULT_CAMPAIGN_STAGE) -# An algorithm whose settings live under a differently named config section. -# ``run_sft_v2.py`` reports itself as ``sft_v2`` while reading ``master_config -# .sft``, so looking its settings up by algorithm name finds nothing. Mapped -# rather than renamed because the name is also reported as ``rl.algorithm``, -# where the two implementations do need telling apart. -_CONFIG_SECTION_BY_ALGORITHM = {"sft_v2": "sft"} - - -def _config_section(algorithm: str) -> str: - """Config section holding *algorithm*'s settings; its own name by default.""" - return _CONFIG_SECTION_BY_ALGORITHM.get(algorithm, algorithm) - - # TelemetryConfig field -> NEMO_RL_OTEL_* env var. ``service_name`` maps to the # standard ``OTEL_SERVICE_NAME`` (lens reads it directly, unprefixed). _ENV_FIELD_MAP = { @@ -338,10 +325,12 @@ def _training_config_attributes(master_config: Any, algorithm: str) -> dict[str, if optimizer: attrs[NV_DL_TRAINING_CONFIG_OPTIMIZER] = optimizer - # Non-positive values are dropped rather than reported: ``rm`` spells "train - # for one epoch" as ``max_num_steps: -1``, which downstream would otherwise - # read as a target of -1 iterations. - max_num_steps = _dig(master_config, _config_section(algorithm), "max_num_steps") + # By algorithm name, with no table of which section each entry point keeps + # its settings in: one that does not match a section simply reports no + # target. Non-positive values are dropped too -- ``rm`` spells "train for + # one epoch" as ``max_num_steps: -1``, which downstream would otherwise read + # as a target of -1 iterations. + max_num_steps = _dig(master_config, algorithm, "max_num_steps") if isinstance(max_num_steps, int) and max_num_steps > 0: attrs[NV_DL_TRAINING_TARGET_TRAIN_ITERS] = max_num_steps return attrs diff --git a/tests/unit/telemetry/test_setup.py b/tests/unit/telemetry/test_setup.py index c923c8137ed..4f782fd3e8f 100644 --- a/tests/unit/telemetry/test_setup.py +++ b/tests/unit/telemetry/test_setup.py @@ -171,17 +171,6 @@ def test_training_target_drops_a_non_positive_step_count(): assert attrs["nv.dl.training.target.train_iters"] == 150 -def test_sft_v2_finds_its_step_target_under_the_sft_section(): - """``run_sft_v2.py`` reports ``sft_v2`` but reads ``master_config.sft``. - - Looking the settings up by algorithm name finds no section at all, so the - whole of sft_v2 silently reports no training target. - """ - cfg = {"sft": {"max_num_steps": 60}} - attrs = _training_config_attributes(cfg, "sft_v2") - assert attrs["nv.dl.training.target.train_iters"] == 60 - - def test_training_config_reads_either_backends_optimizer_spelling(): # megatron names the algorithm and keeps it in its own block, automodel # names the class and keeps it in the policy's. From 8432921208cc489dfe983e60b47cc0e9be65252a Mon Sep 17 00:00:00 2001 From: Raj Singh Date: Tue, 13 Oct 2026 18:15:37 -0700 Subject: [PATCH 13/17] fix(telemetry): refuse a non-positive collection interval Reading the cadence as `interval if interval is not None` let a configured `collection_interval: 0` through, where the previous `interval or ...` had made it fall back. At zero the loop scrapes every raylet as fast as HTTP allows: the new test takes 14s and spins without the guard. Refused rather than clamped. Clamping reports a cadence the config did not ask for, which is harder to notice than no series at all, and the check sits in start_host_memory_monitoring rather than the constructor so the documented never-raises contract holds -- the constructor runs before that function's guard. No config ships a zero today; the GPU monitor spins the same way on one, which is its own fix. The telemetry conftest also resets the metric vocabulary's freeze now, beside the other process globals it already clears. Any test that brings telemetry up closes the vocabulary for the rest of the process, and a module whose rows are declared by an import inside a later test then raises on registering them -- surfacing as a missing series in whichever test runs next rather than as leaked state. Moved there from the host-memory fixture, which was the narrower half of the same fix. Signed-off-by: Raj Singh --- nemo_rl/telemetry/host_memory.py | 15 ++++++++++++++- tests/unit/telemetry/conftest.py | 7 +++++++ tests/unit/telemetry/test_host_memory.py | 17 ++++++++++++----- 3 files changed, 33 insertions(+), 6 deletions(-) diff --git a/nemo_rl/telemetry/host_memory.py b/nemo_rl/telemetry/host_memory.py index 5b04c8dc5eb..64785cb61ee 100644 --- a/nemo_rl/telemetry/host_memory.py +++ b/nemo_rl/telemetry/host_memory.py @@ -381,7 +381,9 @@ def __init__(self, collection_interval: float) -> None: """Initialise the monitor. Args: - collection_interval: Seconds between samples of every node. + collection_interval: Seconds between samples of every node. Must be + positive; :func:`start_host_memory_monitoring` is what enforces + that, since it is the only caller a run goes through. """ self.collection_interval = collection_interval self._stop = threading.Event() @@ -522,6 +524,17 @@ def start_host_memory_monitoring(collection_interval: float) -> None: global _MONITOR if _MONITOR is not None: return + # Refused rather than clamped to something sane: at zero the loop would + # scrape every raylet as fast as HTTP allows, and a descriptive series is + # not worth hammering the cluster for. Clamping would report a cadence the + # config did not ask for, which is harder to notice than no series at all. + if collection_interval <= 0: + warn_once( + "host_memory_interval", + "host memory monitoring needs a positive collection interval, not " + f"{collection_interval}; not starting it", + ) + return monitor = HostMemoryMonitor(collection_interval=collection_interval) try: # Before start(), so the instruments exist no matter when the first diff --git a/tests/unit/telemetry/conftest.py b/tests/unit/telemetry/conftest.py index de918667bed..6b0deadd676 100644 --- a/tests/unit/telemetry/conftest.py +++ b/tests/unit/telemetry/conftest.py @@ -66,10 +66,17 @@ def _clear_telemetry_env() -> None: def _reset_rl_telemetry() -> None: import nemo_rl.telemetry.metrics as metrics_mod import nemo_rl.telemetry.setup as setup_mod + import nemo_rl.telemetry.vocabulary as vocabulary_mod setup_mod._TELEMETRY_HANDLE = None setup_mod._TELEMETRY_INITIALISED = False + # Any test that brings telemetry up closes the vocabulary for the rest of + # the process, and a module whose rows are declared by an import inside a + # later test then raises on registering them -- a failure that reads as a + # missing series in whichever test runs next, not as leaked state here. + vocabulary_mod._FROZEN = False + # The lens metric registry is process-global, so a group declared by one # test would otherwise satisfy the next test's registration guard and hide # a regression in it. diff --git a/tests/unit/telemetry/test_host_memory.py b/tests/unit/telemetry/test_host_memory.py index d4a214fe3dd..db39a172707 100644 --- a/tests/unit/telemetry/test_host_memory.py +++ b/tests/unit/telemetry/test_host_memory.py @@ -14,7 +14,6 @@ import pytest import nemo_rl.telemetry.host_memory as host_memory -import nemo_rl.telemetry.vocabulary as vocabulary from nemo_rl.telemetry.host_memory import ( HOST_NAME_ATTR, NO_WORKER_GROUPS, @@ -64,10 +63,6 @@ def _reset_host_memory_state(): host_memory._WORKER_GROUP_NODES.clear() host_memory._SNAPSHOT.clear() stop_host_memory_monitoring() - # The driver test below freezes the metric vocabulary for the rest of the - # process. Nothing in the suite registers a metric after collection today, - # but a module that did would fail here rather than where it went wrong. - vocabulary._FROZEN = False def _fake_ray(monkeypatch, nodes): @@ -568,6 +563,18 @@ def test_a_process_runs_at_most_one_monitor(monkeypatch): assert host_memory._MONITOR is first +@pytest.mark.parametrize("interval", [0.0, -5.0]) +def test_a_nonpositive_interval_starts_nothing_rather_than_spinning( + monkeypatch, interval +): + """At zero the loop would scrape every raylet as fast as HTTP allows.""" + _fake_ray(monkeypatch, []) + + start_host_memory_monitoring(collection_interval=interval) + + assert host_memory._MONITOR is None + + def test_stopping_releases_the_monitor_so_a_later_run_can_start_one(monkeypatch): _fake_ray(monkeypatch, []) From 77ec2030f58e8c35036841c6cc077510ecaa8299 Mon Sep 17 00:00:00 2001 From: Raj Singh Date: Wed, 14 Oct 2026 14:13:01 -0700 Subject: [PATCH 14/17] docs(telemetry): say the step target is a cap, trim two stale comments Three review follow-ups, all comment or prose. nv.dl.training.target.train_iters is the configured cap, not a prediction: grpo_math_1B ships max_num_steps: 1000000 with max_num_epochs: 1 and stops on epochs long before the cap, so the row says so rather than leaving 1,000,000 to be read as a run length. What to report instead is a follow-up. The parse guard added back a seven-line comment right after the eight long ones came out; the reasoning is already in the docstring of test_an_unparsable_page_costs_only_its_own_point, so two lines do here. The assert in the driver test justified itself by "the lens guard at the top of this test", which this branch removed two commits earlier. It asserts because lens is a core dependency. Signed-off-by: Raj Singh --- docs/observability/configuration.md | 2 +- nemo_rl/telemetry/host_memory.py | 9 ++------- tests/unit/telemetry/test_host_memory.py | 6 ++---- 3 files changed, 5 insertions(+), 12 deletions(-) diff --git a/docs/observability/configuration.md b/docs/observability/configuration.md index 6391192236a..9fab877b536 100644 --- a/docs/observability/configuration.md +++ b/docs/observability/configuration.md @@ -99,7 +99,7 @@ Each process sets stable-for-the-run values on its OTel `Resource` — `init_tel | `nv.dl.training.config.sequence_length` | `policy.max_total_sequence_length` | | `nv.dl.training.config.optimizer` | on Megatron `policy.megatron_cfg.optimizer.optimizer` (`adam`), otherwise `policy.optimizer.name` (`torch.optim.AdamW`) — the backends both spell it and store it differently, and only the enabled one is read | | `nv.dl.training.config.recompute_granularity` | Megatron runs only: `policy.megatron_cfg.recompute_granularity`, and only when `activation_checkpointing` is on. Megatron ignores the granularity otherwise, so reporting it would describe recompute that is not happening | -| `nv.dl.training.target.train_iters` | `.max_num_steps`, when positive. Looked up by algorithm name, so an entry point that keeps its settings under a differently named section reports no target at all — `sft_v2` reads `master_config.sft`, and omits this. `rm` spells "train for one epoch" as `-1`, which is omitted rather than reported as a target | +| `nv.dl.training.target.train_iters` | `.max_num_steps`, when positive. This is the configured cap, not a prediction: many configs set a placeholder (`grpo_math_1B` uses `1000000` with `max_num_epochs: 1`) and stop earlier on epochs. Looked up by algorithm name, so an entry point that keeps its settings under a differently named section reports no target at all — `sft_v2` reads `master_config.sft`, and omits this. `rm` spells "train for one epoch" as `-1`, which is omitted rather than reported as a target | | `nv.dl.software.torch`, `.cuda`, `.nccl` | versions as loaded in the process that opened the span, not as pinned in the lockfile — they are per process, so filter on `rl.worker_group` to read one backend's versions. CUDA and NCCL are omitted on a CPU-only build. SGLang, TRT-LLM and Dynamo workers do not report versions yet | | `nv.dl.software.transformer_engine` | the installed distribution's version, read from package metadata rather than by importing it — importing transformer-engine loads its PyTorch extension in every process that sets up telemetry. Omitted when it is not installed, which is every process except Megatron and Automodel workers: its absence on a driver, controller or vLLM span does not mean the run lacked TE | | `nv.dl.rank`, `nv.dl.world_size` | this process's rank and group size (`RANK` / `WORLD_SIZE`, or `0` / `1` for the driver and singleton actors) | diff --git a/nemo_rl/telemetry/host_memory.py b/nemo_rl/telemetry/host_memory.py index 64785cb61ee..5a81ab030d8 100644 --- a/nemo_rl/telemetry/host_memory.py +++ b/nemo_rl/telemetry/host_memory.py @@ -477,13 +477,8 @@ def _scrape(self, node: RayNode) -> dict[str, float]: try: return parse_memory_samples(response.text) except ValueError: - # Guarded separately from the fetch, and not left to the guard - # around the sweep: prometheus_client raises ValueError on a body - # that is not the exposition format -- an HTML error page from a - # proxy in front of the raylet, a truncated response -- and letting - # that reach the sweep would abandon every other node's sample and - # leave the previous snapshot exporting, stale labels included, for - # as long as the page stayed broken. + # Caught here so one bad page (an HTML error page, a truncated + # body) costs only this node's point, not the whole sweep. warn_once( f"host_memory_parse:{node.host}", f"could not parse host memory from {node.metrics_url}", diff --git a/tests/unit/telemetry/test_host_memory.py b/tests/unit/telemetry/test_host_memory.py index db39a172707..0197c4915fc 100644 --- a/tests/unit/telemetry/test_host_memory.py +++ b/tests/unit/telemetry/test_host_memory.py @@ -616,10 +616,8 @@ def test_the_driver_starts_the_monitor_that_the_worker_groups_register_with( cluster=None, ) handle = init_telemetry_driver(config, algorithm="grpo") - # Asserted rather than skipped: the lens guard at the top of this test - # already covers "no telemetry here", so a handle that comes up not - # exporting is the regression this test exists to catch, and skipping on it - # would hide exactly that. + # Asserted, not skipped: lens is a core dependency, so a handle that is not + # exporting is the regression this test exists to catch. assert handle is not None and handle.is_exporting assert host_memory._MONITOR is not None From a382370a4065c9c4a3022d82cd8d7bc2e1c8f852 Mon Sep 17 00:00:00 2001 From: Raj Singh Date: Wed, 14 Oct 2026 16:31:47 -0700 Subject: [PATCH 15/17] test(telemetry): drop unreachable lens guards and a false typing rationale The autouse fixture in this package's conftest imports nemo.lens unconditionally, so a missing lens errors in fixture setup before any pytest.importorskip in a test body can skip. Remove the three this branch added, and the module-level _HAS_LENS skipif in test_checkpoint_spans.py, importing lens at the top instead. test_parallelism_sizes_use_the_resource_scoped_spelling credited DL_RESOURCE_TYPES and normalize_resource_attributes for the choice of spelling. normalize_resource_attributes runs only over the attributes lens derives from OTEL_RESOURCE_ATTRIBUTES, so the resource_attributes handed in from setup.py are merged untouched and either spelling would arrive intact. setup.py was corrected already; this test was missed. Reword it to the convention reason and assert the key is reported rather than asserting against the converter registry. Signed-off-by: Raj Singh --- tests/unit/telemetry/test_checkpoint_spans.py | 31 ++++++------------- tests/unit/telemetry/test_setup.py | 23 ++++++-------- 2 files changed, 19 insertions(+), 35 deletions(-) diff --git a/tests/unit/telemetry/test_checkpoint_spans.py b/tests/unit/telemetry/test_checkpoint_spans.py index fcaaa723288..80134be1ea8 100644 --- a/tests/unit/telemetry/test_checkpoint_spans.py +++ b/tests/unit/telemetry/test_checkpoint_spans.py @@ -23,29 +23,18 @@ import threading import pytest +from nemo.lens import NemoLensConfig, setup_telemetry +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) -from nemo_rl.utils.checkpoint import CheckpointManager - -try: - from nemo.lens import NemoLensConfig, setup_telemetry - from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( - InMemorySpanExporter, - ) - - from nemo_rl.telemetry.instrumentation import ( - RL_ITERATION_ATTR, - iteration_scope, - managed_span, - ) - from nemo_rl.telemetry.span_groups import RLSpanGroup - - _HAS_LENS = True -except ImportError: - _HAS_LENS = False - -pytestmark = pytest.mark.skipif( - not _HAS_LENS, reason="nemo-lens (+ opentelemetry sdk) not installed" +from nemo_rl.telemetry.instrumentation import ( + RL_ITERATION_ATTR, + iteration_scope, + managed_span, ) +from nemo_rl.telemetry.span_groups import RLSpanGroup +from nemo_rl.utils.checkpoint import CheckpointManager @pytest.fixture diff --git a/tests/unit/telemetry/test_setup.py b/tests/unit/telemetry/test_setup.py index 4f782fd3e8f..ee2e02fc5e2 100644 --- a/tests/unit/telemetry/test_setup.py +++ b/tests/unit/telemetry/test_setup.py @@ -11,6 +11,7 @@ from types import SimpleNamespace import pytest +from nemo.lens import derive_nv_dl_run_uuid from pydantic import ValidationError from nemo_rl.telemetry.config import TelemetryConfig @@ -142,20 +143,18 @@ def test_the_data_parallel_size_is_not_reported(): def test_parallelism_sizes_use_the_resource_scoped_spelling(): """``nv.dl.topology.size.*``, not the ``dl.tensor_parallel.*`` family. - Both spell the same quantity, but only the topology one is in lens's - ``DL_RESOURCE_TYPES``, which feeds the table - ``normalize_resource_attributes`` consults. Keys outside it are passed - through untyped rather than rejected, so the ``dl.*`` spelling would reach - the backend as a string -- it is the span and metric spelling. + Both spell the same quantity, and nothing on lens's side picks between + them: ``normalize_resource_attributes`` runs only over the attributes lens + derives from ``OTEL_RESOURCE_ATTRIBUTES``, while the ``resource_attributes`` + handed in from here are merged as they are, so either spelling would arrive + intact. The convention is that the topology keys are the resource-scoped + spelling and ``dl.*`` the span and metric one, which makes getting it right + entirely this function's job -- hence the test. """ - pytest.importorskip("nemo.lens") - from nemo.lens.semconv.resources import DL_RESOURCE_TYPES - attrs = _build_resource_attributes(_FakeMasterConfig(), "grpo") assert "dl.tensor_parallel.size" not in attrs assert "dl.pipeline_parallel.size" not in attrs - assert "nv.dl.topology.size.tp" in DL_RESOURCE_TYPES - assert "dl.tensor_parallel.size" not in DL_RESOURCE_TYPES + assert attrs["nv.dl.topology.size.tp"] == 2 def test_training_target_drops_a_non_positive_step_count(): @@ -222,7 +221,6 @@ def test_process_attributes_report_the_local_rank_when_the_launcher_set_one( process without the variable reports no local rank rather than defaulting to zero and claiming one. """ - pytest.importorskip("nemo.lens") monkeypatch.setenv("LOCAL_RANK", "3") assert _process_attributes(rank=11, world_size=16)["nv.dl.local_rank"] == 3 @@ -237,9 +235,6 @@ def test_process_attributes_name_the_provider_and_the_run(monkeypatch): runs with Megatron ones joins on the derived value, so publishing the run id under this key would look populated and match nothing. """ - pytest.importorskip("nemo.lens") - from nemo.lens import derive_nv_dl_run_uuid - monkeypatch.delenv("LOCAL_RANK", raising=False) attrs = _process_attributes(rank=0, world_size=1, run_id="job-7") assert attrs["nv.dl.provider.name"] == "nemo-rl" From 58b5f0c9aef705ff9c039422749f62f5b3dc9d4f Mon Sep 17 00:00:00 2001 From: shuyixiong <219646547+shuyixiong@users.noreply.github.com> Date: Thu, 15 Oct 2026 09:38:37 -0700 Subject: [PATCH 16/17] fix(trtllm): stop stripping MPI/SLURM env vars; allow MPI to run as root ray.sub (head and workers) and init_ray() no longer unset the PMI/PMIX/MPI/OMPI/SLURM variables inherited from the srun launcher. The TRT-LLM generation worker instead sets OMPI_ALLOW_RUN_AS_ROOT(_CONFIRM) and PRTE_ALLOW_RUN_AS_ROOT(_CONFIRM) unconditionally, since the container runs as root and Open MPI / PRTE refuse to start as root otherwise. grpo-qwen2.5-0.5b-1n4g-megatron-trtllm-noncolocated-async runs its 30 steps with the launcher's SLURM_/OMPI_MCA_/PRTE_MCA_ variables left in place. Co-Authored-By: Claude Opus 5.5 Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com> (cherry picked from commit 546081fee8dc22c6de38f43a2670ce9be0d399d9) --- nemo_rl/distributed/virtual_cluster.py | 14 -------------- .../generation/trtllm/trtllm_worker_async.py | 7 +++++++ ray.sub | 14 -------------- .../generation/trtllm/test_trtllm_worker_async.py | 7 +++++++ 4 files changed, 14 insertions(+), 28 deletions(-) diff --git a/nemo_rl/distributed/virtual_cluster.py b/nemo_rl/distributed/virtual_cluster.py index 830ca473b94..e9e15050bb0 100644 --- a/nemo_rl/distributed/virtual_cluster.py +++ b/nemo_rl/distributed/virtual_cluster.py @@ -457,20 +457,6 @@ def _init_ray(log_dir: Optional[str] = None) -> str: Returns: How the cluster was obtained: one of the ``RAY_CLUSTER_*`` constants. """ - # Strip MPI/PMIx/SLURM launcher vars from the driver env before they get - # captured into runtime_env (both by `dict(os.environ)` below and by - # RayWorkerGroup, which re-reads os.environ). Otherwise they are forwarded - # into every actor, where they cause two failures in the TRT-LLM generation - # worker: (1) PMIX_/SLURM_STEP_ID make OMPI's MPI_Init abort with "OMPI was - # not built with SLURM's PMI support"; (2) any residual OMPI_/MPI_/SLURM_ - # var makes TRT-LLM think it runs under an MPI launcher and pick the MPI - # orchestrator (MPI_Comm_Spawn -> MPI_ERR_SPAWN) instead of the Ray - # orchestrator. init_ray() runs before any worker group is built, so - # popping here cleans the env for all downstream captures. - for _k in list(os.environ): - if _k.startswith(("PMIX_", "PMI_", "MPI_", "OMPI_", "SLURM_")): - os.environ.pop(_k, None) - # Ray actors deserialize constructor arguments before importing NeMo-RL. # Put Hugging Face's generated ``transformers_modules`` package on the # cluster-wide PYTHONPATH so trust_remote_code objects can be unpickled at diff --git a/nemo_rl/models/generation/trtllm/trtllm_worker_async.py b/nemo_rl/models/generation/trtllm/trtllm_worker_async.py index 87ad4bb7ce6..6126ee902f8 100644 --- a/nemo_rl/models/generation/trtllm/trtllm_worker_async.py +++ b/nemo_rl/models/generation/trtllm/trtllm_worker_async.py @@ -94,6 +94,13 @@ def configure_worker( "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1", "NCCL_CUMEM_ENABLE": "1", "PATH": worker_path, + # The container runs as root, and Open MPI / PRTE refuse to start + # as root unless explicitly allowed. Set unconditionally so the + # engine's MPI init works whatever the launch environment carries. + "OMPI_ALLOW_RUN_AS_ROOT": "1", + "OMPI_ALLOW_RUN_AS_ROOT_CONFIRM": "1", + "PRTE_ALLOW_RUN_AS_ROOT": "1", + "PRTE_ALLOW_RUN_AS_ROOT_CONFIRM": "1", } init_kwargs: dict[str, Any] = {} diff --git a/ray.sub b/ray.sub index 5c49e8f2812..37145c5c518 100644 --- a/ray.sub +++ b/ray.sub @@ -889,13 +889,6 @@ if [[ "$DEDICATED_RAY_HEAD" == "1" ]]; then echo "[INFO] Dedicated Ray head enabled: advertising zero GPUs and no worker_units" fi -# Clear MPI/PMIx/SLURM env vars inherited from the srun launcher so TRT-LLM's -# MPI_Init (in the ray daemon child process below) starts clean. -# NOTE: keep this after topology_probe.sh, which reads SLURM_TOPOLOGY_ADDR/SLURM_PROCID. -for v in \$(env | awk -F= '/^(PMI|PMIX|MPI|OMPI|SLURM)_/{print \$1}'); do - unset "\$v" -done - # Start the head without --block. 'ray start --head' then returns as soon as the # GCS is listening, so we touch STARTED_RAY_HEAD only after the head is actually up. # Workers gate on that file, so they only try to connect once the GCS is accepting @@ -1094,13 +1087,6 @@ sed -i 's/context\.py_executable = " "\.join(self\.nsight_cmd) + " python"/conte source $LOG_DIR/topology_probe.sh -# Clear MPI/PMIx/SLURM env vars inherited from the srun launcher so TRT-LLM's -# MPI_Init (in the ray daemon child process below) starts clean. -# NOTE: keep this after topology_probe.sh, which reads SLURM_TOPOLOGY_ADDR/SLURM_PROCID. -for v in \$(env | awk -F= '/^(PMI|PMIX|MPI|OMPI|SLURM)_/{print \$1}'); do - unset "\$v" -done - # Wait for the head to signal that its GCS is listening before connecting. echo "[INFO] Worker \$SLURM_PROCID: waiting for Ray head GCS to be ready..." while true; do diff --git a/tests/unit/models/generation/trtllm/test_trtllm_worker_async.py b/tests/unit/models/generation/trtllm/test_trtllm_worker_async.py index 625563c8ec6..4a7e7ebe1b6 100644 --- a/tests/unit/models/generation/trtllm/test_trtllm_worker_async.py +++ b/tests/unit/models/generation/trtllm/test_trtllm_worker_async.py @@ -77,6 +77,13 @@ def test_configure_worker_reserves_gpus_for_internal_trtllm_actors(monkeypatch): assert env_vars["RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES"] == "1" assert env_vars["NCCL_CUMEM_ENABLE"] == "1" assert env_vars["PATH"].endswith(":/usr/bin") + for name in ( + "OMPI_ALLOW_RUN_AS_ROOT", + "OMPI_ALLOW_RUN_AS_ROOT_CONFIRM", + "PRTE_ALLOW_RUN_AS_ROOT", + "PRTE_ALLOW_RUN_AS_ROOT_CONFIRM", + ): + assert env_vars[name] == "1" assert init_kwargs == { "bundle_indices": [2, 3], "bundle_pg_idx": 0, From bc2dbec16641d302475b099689d486008814b2ba Mon Sep 17 00:00:00 2001 From: Raj Singh Date: Thu, 15 Oct 2026 10:26:24 -0700 Subject: [PATCH 17/17] fix(ray): drop a comment referring to the removed env scrub The Alloy sidecar is sourced early so it inherits SLURM_* before the MPI/PMIx/SLURM scrub that used to follow. The preceding commit removed that scrub from both roles, so the ordering the comment justifies no longer has anything to order against. Signed-off-by: Raj Singh --- ray.sub | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/ray.sub b/ray.sub index 37145c5c518..eb5b03a2a5f 100644 --- a/ray.sub +++ b/ray.sub @@ -693,8 +693,7 @@ else fi # Alloy log export, sourced by both roles once RAY_NODE_ROLE is set. Kept in its -# own file so the two generated container commands don't each need a copy, and -# sourced early in them so it inherits SLURM_* before the scrub. +# own file so the two generated container commands don't each need a copy. ALLOY_SIDECAR_FILE="$LOG_DIR/alloy_sidecar.sh" cat > "$ALLOY_SIDECAR_FILE" <