feat: expand vLLM observability/telemetry by binaryaaron · Pull Request #526 · NVIDIA-NeMo/Safe-Synthesizer

binaryaaron · 2026-05-27T15:06:22Z

Summary

Adds a self-contained vllm_observability module that lifts four measurement primitives + the structured event schema out of any specific consumer (production VllmBackend or benchmark harness) and into a shared, well-tested surface. Wires the primitives into VllmBackend.generate() so every production generation invocation now emits a vllm.cell.complete structured event carrying peak device VRAM, host loadavg pre/post, vLLM's KV-cache fraction / prefix-cache hit rate / speculative-decoding acceptance, and the engine's effective runtime config.

The primitives are deliberately generic — they're observability infrastructure, not benchmark-specific. The companion PR (binaryaaron/vllm-benchmark-harness-work) imports them and composes the schema rather than re-defining its fields, which fixes the architectural smell where observability fields had been locked inside CandidateMetrics.

What this changes for production

Every VllmBackend.generate() call now writes:

Surface	What lands
Structured log	`logger.runtime.info("vllm.cell.complete", extra={"ctx": event.model_dump()})`
Wandb (when active run)	`wandb.log(event.to_wandb_payload())` — keys namespaced under `vllm_cell/`

Operators get flag-engagement drift detection (cross-checks intended config against probed vllm_config), KV-cache pressure visibility, out-of-process peak VRAM via NVML (sidesteps the VLLM_ENABLE_V1_MULTIPROCESSING=1 blind spot where torch.cuda.max_memory_allocated() in the harness reads 0), and host load bracketing per cell.

Primitives + schema

Primitive	Surface
`NvmlPeakSampler`	Context-manager + daemon thread polling `pynvml` at 250ms cadence
`read_loadavg()`	`/proc/loadavg` snapshot as `(1m, 5m, 15m)` tuple
`probe_engine_runtime_config(llm)`	Best-effort introspection of `vllm_config.{scheduler,cache,speculative}_config`
`read_vllm_runtime_metrics(llm)`	`LLM.get_metrics()` reader → `{kv_cache_usage_perc, prefix_cache_hit_rate, spec_accept_rate}`
`flag_engagement_mismatches(intended, actual)`	Dict-vs-dict cross-check; only checks explicitly-set fields
`CellObservability` pydantic model	The event schema, `extra="forbid"` so producers must update on schema change
`log_cell_observability(event)`	wandb_setup.py helper; no-op when no active run; best-effort emission

Every primitive is degraded-mode by design: missing pynvml / non-Linux / unavailable metric / failed call → None or empty dict, never raises.

Architectural choices worth flagging

phase=GENERATE + job_type="benchmark" instead of a new WandbPhase.BENCHMARK — a benchmark cell is structurally a GENERATE phase invocation that's measured rather than consumed; job_type is the right discriminator. Keeps production-generate and benchmark cells in the same wandb workspace section for cross-comparison.
Composition over duplication — the schema is one model. Consumers (production + benchmark) embed it via observability: CellObservability rather than re-declaring fields. Adding a new measurement primitive in the future requires touching this module only.
Soft dependency throughout — wandb missing / WANDB_MODE=disabled / wandb.init raises / wandb.log raises → warning, generation continues. Production reliability isn't gated on observability working.
No PR-1 dependency — primitives use pynvml + /proc/loadavg + vLLM's public llm.get_metrics() / llm.llm_engine.vllm_config surface. No imports from vllm_engine_factory / vllm_trace etc. This PR lands independently of PR-1's review timing.

Test coverage

tests/generation/test_vllm_observability.py — 28 contract tests, ~7s wall:

Schema: defaults all optional, JSON round-trip lossless, extra='forbid' enforced.
to_wandb_payload: namespacing, None-dropping, tuple unpacking, dict flattening.
flag_engagement_mismatches: parametrized matrix (clean match / disagreement / unset-intended / missing-actual / both-empty).
read_loadavg: shape on Linux, None on read failure.
probe_engine_runtime_config: empty dict on any failure (parametrized), extracts known fields when vllm_config is real-shaped.
read_vllm_runtime_metrics: stable dict keys, correct derivation, exception → degraded mode, zero-denominator → None.
NvmlPeakSampler: context-manager protocol, peak_gb typed float | None, pynvml ImportError → None.
log_cell_observability: no-op when no active run, calls wandb.log with the flattened payload when active, swallows wandb.log exceptions.

Tests deliberately focus on contracts, NOT implementation details (field counts, log wording, etc.) — that's documented inline so reviewers know the scope choice was intentional.

Test plan

Lint / typecheck passes
Unit tests (tests/generation/test_vllm_observability.py) pass — 28 tests
Manual: run safe-synthesizer run generate ... against a small dataset, confirm a vllm.cell.complete event lands in the structured log
Manual: run with WANDB_MODE=online, confirm the wandb run page shows the vllm_cell/* metrics
Manual: run with no GPU / pynvml unavailable, confirm peak_vram_gb: None and no exception

What's NOT in this PR

compilation_config + kv_cache_metrics fields on a future BenchmarkEngineConfig — those are benchmark-side knobs, deferred to the companion PR.
Within-call sampling cadence (peak-VRAM progression during long generates) — v1 emits one event per generate() call; future improvement via the existing NvmlPeakSampler thread.
VllmBackend.generate() end-to-end integration test — requires actual vLLM spin-up; covered by production usage + the companion PR's benchmark integration.

🤖 Generated with Claude Code

Summary by CodeRabbit

New Features
- End-of-generation observability: peak GPU memory sampling, vLLM runtime metrics, engine runtime config capture, and pre/post host load snapshots; formatted payloads mirrored to structured logs and WandB when available.
- Public helpers to produce WandB-friendly payloads.
Bug Fixes
- Observability is best-effort: telemetry/wandb failures are swallowed and won’t interrupt generation.
Tests
- Expanded tests covering observability schema, metric extraction, NVML sampler, payload formatting, and WandB/logging resilience.

coderabbitai · 2026-05-27T15:06:31Z

PR changed again? Review this PR in Change Stack to compare snapshots and stay oriented.

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

Walkthrough

Adds GenerationObservability schema and telemetry probes (NVML, loadavg, vLLM metrics, engine runtime config), a WandB logging Protocol/helper, integrates emission into VllmBackend.generate() as a finally-bracketed event, and adds unit tests for contracts and integration.

Changes

vLLM generation observability

Layer / File(s)	Summary
GenerationObservability schema and WandB helper `src/nemo_safe_synthesizer/generation/vllm_observability.py`, `src/nemo_safe_synthesizer/cli/wandb_setup.py`	Adds `GenerationObservability` Pydantic model with `to_wandb_payload()` and a `WandbLoggable` `Protocol` plus `log_observability_event()` helper.
NVML sampler, loadavg, metrics, engine probes `src/nemo_safe_synthesizer/generation/vllm_observability.py`, `src/nemo_safe_synthesizer/observability.py`	Implements `probe_engine_runtime_config()`, `flag_engagement_mismatches()`, `read_vllm_runtime_metrics()`, re-exports `NvmlPeakSampler`, `read_loadavg`, and adds degraded-mode NVML sampler and loadavg reader.
VllmBackend observability integration `src/nemo_safe_synthesizer/generation/vllm_backend.py`	Pre-declares and caches `_engine_runtime_config` in `initialize()`, refactors `generate()` to capture `loadavg_pre`, run generation inside `NvmlPeakSampler` and `_run_generation()`, and emit `GenerationObservability` in `_emit_generation_observability()` from a `finally` block while suppressing observability errors.
Observability module contract tests `tests/generation/test_vllm_observability.py`	Adds tests covering `GenerationObservability` schema, `to_wandb_payload()` flattening, `flag_engagement_mismatches`, `probe_engine_runtime_config` behaviors, `read_vllm_runtime_metrics`, `NvmlPeakSampler` contracts, and `log_observability_event()` semantics.
VllmBackend observability integration tests `tests/generation/test_vllm_backend.py`	Adds tests asserting `initialize()` caches engine runtime config and that `generate()` invokes `_emit_generation_observability()` (including on generation failure) and swallows emission errors.
General observability tests `tests/test_observability.py`	Adds tests for `read_loadavg()` shape, `_default_nvml_device_index()` parsing, and `NvmlPeakSampler` degraded and shutdown behaviors.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~75 minutes

Suggested labels

feature, test

Suggested reviewers

mckornfield
kendrickb-nvidia
nina-xu

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 78.57% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'feat: expand vLLM observability/telemetry' accurately summarizes the main change: adding comprehensive observability and telemetry infrastructure for vLLM generation, including new modules, primitives, and integration into VllmBackend.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests
Commit unit tests in branch binaryaaron/vllm-observability-improvements

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

codecov · 2026-05-27T15:43:45Z

Codecov Report

❌ Patch coverage is 93.96378% with 30 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/nemo_safe_synthesizer/observability.py	69.73%	23 Missing ⚠️
..._safe_synthesizer/generation/vllm_observability.py	97.63%	3 Missing ⚠️
...c/nemo_safe_synthesizer/generation/vllm_backend.py	91.66%	2 Missing ⚠️
tests/generation/test_vllm_observability.py	98.84%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

greptile-apps · 2026-06-03T23:04:45Z

Greptile Summary

Adds a self-contained vllm_observability module with four degraded-mode measurement primitives (NvmlPeakSampler, read_loadavg, probe_engine_runtime_config, read_vllm_runtime_metrics) and wires them into VllmBackend.generate() so every production invocation emits a GenerationObservability event to structured logs and wandb. Previous review issues (NVML resource leak on handle-lookup failure, device-index defaulting, logger call in except clause) are correctly resolved.

generate() is refactored into a thin observability bracket over _run_generation(): a finally-guarded _emit_generation_observability() emits the event regardless of success/failure, with a contextlib.suppress-protected inner logger so observability failures can never mask a successful generation result.
NvmlPeakSampler and read_loadavg land in the top-level observability.py as shared hardware primitives; vllm_observability.py re-exports them alongside the schema and vLLM-specific probes.
Test suite covers 28 contracts including degraded-mode paths, schema round-trip, wandb payload flattening, and the NVML shutdown-on-handle-failure path introduced to fix the previous resource-leak finding.

Confidence Score: 5/5

Safe to merge — observability failures are fully insulated from generation results and all previous review findings are resolved.

All three previously-flagged issues (NVML resource leak, wrong device index, logger-in-except propagation in vllm_backend.py) are correctly fixed. The generate() refactor is structurally sound: _run_generation always succeeds or raises, _emit_generation_observability is always called via finally and swallows every failure including a contextlib.suppress guard on the inner logger. Degraded-mode coverage is thorough and the 28 contract tests include the shutdown-on-handle-failure path. The two remaining notes are non-blocking style observations.

No files require special attention — the only open item is the unprotected logger.warning in wandb_setup.py's except branch, which is guarded by the generation path's outer try/except today and is a preventative concern for a future training wiring.

Important Files Changed

Filename	Overview
src/nemo_safe_synthesizer/generation/vllm_observability.py	New module introducing GenerationObservability schema, probe_engine_runtime_config, read_vllm_runtime_metrics, and flag_engagement_mismatches — all degraded-mode by design; minor all inconsistency with the private _default_nvml_device_index export.
src/nemo_safe_synthesizer/observability.py	Adds NvmlPeakSampler (daemon-thread NVML peak sampler), _default_nvml_device_index, and read_loadavg as shared hardware-sampling primitives; NVML resource-leak fix (shutdown called before nulling self._pynvml) is correctly implemented.
src/nemo_safe_synthesizer/generation/vllm_backend.py	generate() refactored into a thin observability bracket (NvmlPeakSampler + finally-based emit) over _run_generation(); _emit_generation_observability swallows all failures including a contextlib.suppress guard on the inner logger — generation reliability contract satisfied.
src/nemo_safe_synthesizer/cli/wandb_setup.py	Adds WandbLoggable Protocol and log_observability_event generic sink; the logger.warning call in the except branch lacks the contextlib.suppress guard applied to the equivalent site in vllm_backend.py — future training callers may not provide the outer protection the generation path currently has.
tests/generation/test_vllm_observability.py	28 contract tests covering schema round-trip, wandb payload flattening, degraded-mode paths, and the NVML sampler protocol; well-scoped and avoids over-testing implementation details.
tests/generation/test_vllm_backend.py	Adds TestGenerationObservabilityEmission covering the finalizer-always-runs contract, correct event assembly, and failure-swallowing; also adds test_initialize_caches_engine_runtime_config for the init-time probe wiring.
tests/test_observability.py	New test file for read_loadavg, _default_nvml_device_index (CUDA_VISIBLE_DEVICES parsing), and NvmlPeakSampler degraded-mode including the shutdown-on-handle-failure path.

Sequence Diagram

sequenceDiagram
    participant Caller
    participant Backend as VllmBackend
    participant Sampler as NvmlPeakSampler
    participant RunGen as _run_generation
    participant Emit as _emit_generation_observability
    participant Log as StructuredLog
    participant WB as wandb

    Caller->>Backend: generate()
    Backend->>Backend: "loadavg_pre = read_loadavg()"
    Backend->>Sampler: __enter__() start daemon poll thread
    Backend->>RunGen: _run_generation(data_actions_fn)
    Note over RunGen: batch loop sets self.gen_results
    RunGen-->>Backend: returns or raises
    Backend->>Sampler: __exit__() stop and join thread
    Note over Backend: finally block always runs
    Backend->>Emit: _emit_generation_observability(sampler, loadavg_pre)
    Emit->>Emit: read_vllm_runtime_metrics(self.llm)
    Emit->>Emit: read_loadavg() post
    Emit->>Emit: build GenerationObservability
    Emit->>Log: logger.runtime.info vLLM generation complete
    Emit->>WB: "log_observability_event(event, prefix=vllm_gen)"
    Note over Emit: all failures swallowed by except Exception
    Backend->>Caller: return self.gen_results

_{Reviews (6): Last reviewed commit: "fix(observability): plain log messages, ..." | Re-trigger Greptile}

coderabbitai

Actionable comments posted: 3

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: d2e5e7a9-cf3f-4117-ae86-b280f679957c

📥 Commits

Reviewing files that changed from the base of the PR and between 940844b and d96a6d0.

📒 Files selected for processing (5)

src/nemo_safe_synthesizer/cli/wandb_setup.py
src/nemo_safe_synthesizer/generation/vllm_backend.py
src/nemo_safe_synthesizer/generation/vllm_observability.py
tests/generation/test_vllm_backend.py
tests/generation/test_vllm_observability.py

📜 Review details

🧰 Additional context used

📓 Path-based instructions (9)

**/*.{md,markdown,py}

📄 CodeRabbit inference engine (.cursor/rules/agent-markdown-style.mdc)

**/*.{md,markdown,py}: Avoid decorative bold (**text**) in list items, body text, and docstrings; use structural cues (headers, list markers, colons, backticks) for emphasis instead
Use backticks for code identifiers, paths, and CLI commands in markdown and docstrings

Files:

src/nemo_safe_synthesizer/cli/wandb_setup.py
tests/generation/test_vllm_backend.py
src/nemo_safe_synthesizer/generation/vllm_observability.py
tests/generation/test_vllm_observability.py
src/nemo_safe_synthesizer/generation/vllm_backend.py

**/*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

**/*.py: Use American English spelling: 'initialize' not 'initialise', 'recognize' not 'recognise', 'color' not 'colour'.
Use observability.get_logger(__name__) for logging, never logging.getLogger() or structlog.get_logger() directly.
Use category loggers: .runtime for internals, .user for progress/results, .system for system events.
Never use print() for operational output. Use click.echo() for CLI output or sys.stdout.write() for raw output in tools.
Use extra={} in logging for structured data that downstream tools should query or aggregate; use f-strings for human-readable context.
Raise from the custom error hierarchy with dual inheritance: SafeSynthesizerError (base), UserError, DataError, ParameterError, GenerationError, InternalError.
Use NSSBaseModel for config/parameter models in config/ which define user-facing configuration. Use raw BaseModel or module-specific bases for data transfer objects and internal structures.
Use BaseSettings for env/CLI settings. Prefer AliasChoices on individual fields when a field needs to respond to both its Python name and an env var name.
Include Field(description=...) for Pydantic model fields as the canonical field docstring for API documentation and CLI help text.
Use assignment-style type = Field(default=..., description="...") as the default for Pydantic model fields because type checkers understand default, default_factory, and alias in assignment style.
Use Annotated only when the field carries additional metadata beyond Field() -- ValueValidator, AutoParam, DependsOnValidator, reusable constrained type aliases, nested-type constraints, or discriminated unions.
Put defaults as bare assignment (= value), not inside Field(default=...), when using Annotated. Exception: use assignment-style Field(default_factory=...) for defaults that cannot be expressed as bare assignments.
Use @dataclass(frozen=True) for immutable value objects and validators; mu...

Files:

src/nemo_safe_synthesizer/cli/wandb_setup.py
tests/generation/test_vllm_backend.py
src/nemo_safe_synthesizer/generation/vllm_observability.py
tests/generation/test_vllm_observability.py
src/nemo_safe_synthesizer/generation/vllm_backend.py

src/**/*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

src/**/*.py: Use relative imports in src/ (e.g., from ..observability import get_logger).
Do not use print() statements in library code. Use get_logger(__name__) from observability.py or click.echo() for CLI.
Do not use assert for validation in library code. Use if/raise for input validation. assert statements can be stripped by -O and must never guard correctness.

Files:

src/nemo_safe_synthesizer/cli/wandb_setup.py
src/nemo_safe_synthesizer/generation/vllm_observability.py
src/nemo_safe_synthesizer/generation/vllm_backend.py

⚙️ CodeRabbit configuration file

Review library code against STYLE_GUIDE.md. Focus on behavior, API contracts, error handling, resource cleanup, typing, logging, and user-facing failures. Public APIs and nontrivial functions need Google-style docstrings.

Files:

src/nemo_safe_synthesizer/cli/wandb_setup.py
src/nemo_safe_synthesizer/generation/vllm_observability.py
src/nemo_safe_synthesizer/generation/vllm_backend.py

**/*.{py,sh,yaml,yml}

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

Include SPDX copyright header at the top: # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. and # SPDX-License-Identifier: Apache-2.0. The make format command handles this automatically.

Files:

src/nemo_safe_synthesizer/cli/wandb_setup.py
tests/generation/test_vllm_backend.py
src/nemo_safe_synthesizer/generation/vllm_observability.py
tests/generation/test_vllm_observability.py
src/nemo_safe_synthesizer/generation/vllm_backend.py

**/*

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

**/*: Include a newline at the end of all files, never trailing whitespace. This is enforced by pre-commit.
Use line length of 120 characters for code, comments, and docstrings (configured in ruff.toml).

Files:

src/nemo_safe_synthesizer/cli/wandb_setup.py
tests/generation/test_vllm_backend.py
src/nemo_safe_synthesizer/generation/vllm_observability.py
tests/generation/test_vllm_observability.py
src/nemo_safe_synthesizer/generation/vllm_backend.py

⚙️ CodeRabbit configuration file

**/*: Review as a senior maintainer for NeMo Safe Synthesizer. Prioritize issues that can change behavior, break user workflows, weaken privacy guarantees, hide failures, make tests unreliable, or create maintenance risk. Avoid generic style commentary unless it points to a concrete project convention that automated tools will not catch.
Comment only when the finding is actionable and tied to changed code. For each finding, state the impact, the condition that triggers it, and the smallest practical fix. Prefer one precise comment over broad advice. Do not ask for refactors outside the PR scope unless the changed code creates the problem.
Review type guidance: - Potential issue: use for correctness bugs, data loss, privacy leaks,
security risks, broken public APIs, invalid config behavior, missing
validation, hidden failures, nondeterministic tests, or CI breakage.

Refactor suggestion: use for local maintainability problems introduced
by the diff when they have clear future cost, such as duplicated setup,
unclear boundaries, over-mocking, avoidable complexity, or opaque test
helpers.

Nitpick: avoid in chill mode. Do not emit formatting, import-order,
wording, or style-only comments unless automated tools cannot catch the
issue and it affects maintainability.

Severity guidance: - Critical: security/privacy leaks, data loss, training/test/holdout
contamination, or broken release/package/core pipeline execution.

Major: incorrect generation/training/evaluation behavior, broken
CLI/SDK public API, invalid config defaults or validators, or GPU/vLLM
cleanup and process-isolation bugs likely to fail CI or production
runs.

Minor: localized bugs, missing focused tests for changed behavior, or
bad test patterns that weaken regression coverage.

Trivial: small cleanup with no behavior impact. Usually suppress in
chill mode.

Info: context only. Avoid unless it helps reviewers understand risk.
Safe-Synthesizer-specific review focus: - Data ...

Files:

src/nemo_safe_synthesizer/cli/wandb_setup.py
tests/generation/test_vllm_backend.py
src/nemo_safe_synthesizer/generation/vllm_observability.py
tests/generation/test_vllm_observability.py
src/nemo_safe_synthesizer/generation/vllm_backend.py

tests/**

📄 CodeRabbit inference engine (.cursor/rules/repo-navigation.mdc)

Tests should mirror the src/ directory structure in tests/

Files:

tests/generation/test_vllm_backend.py
tests/generation/test_vllm_observability.py

tests/**/*.py

📄 CodeRabbit inference engine (.cursor/rules/repo-navigation.mdc)

Tests in tests/e2e/ should be auto-marked with e2e marker, tests in tests/smoke/ with smoke marker, others with unit marker

tests/**/*.py: Use absolute imports in tests/ (e.g., from nemo_safe_synthesizer.observability import get_logger).
Use fixture_ prefix convention for fixtures for grep-ability and to separate fixtures from test functions. Add a one-line docstring describing the fixture's purpose and data.
Use function-scoped fixtures by default. Session scope only when empirically justified by test runtime.
Use bare assert as the primary assertion style; pytest.raises() with match= for exceptions; pytest.approx() for floating-point comparisons.
Mark CUDA-dependent tests with @pytest.mark.e2e, @pytest.mark.smoke, or @pytest.mark.requires_gpu.
Mock only external boundaries, not internal implementation details.
Ensure test isolation: no shared mutable state or execution-order dependencies between tests. If something must be run first before executing a test, include it in the test or a fixture.
Use @pytest.mark.parametrize for testing multiple input combinations rather than copy-pasting similar tests.

Organize tests using pytest following the structure in tests/TESTING.md with support for unit tests, smoke tests, and end-to-end tests

Files:

tests/generation/test_vllm_backend.py
tests/generation/test_vllm_observability.py

⚙️ CodeRabbit configuration file

Review tests against tests/TESTING.md. Check marker usage, fixture naming, tmp_path usage, determinism, and GPU/vLLM process-isolation requirements. Flag slop tests that only check that code runs, assert result is not None when stronger invariants exist, over-mock internal implementation details, patch around the bug instead of reproducing it, or add broad snapshot/golden churn without a clear contract. Flag change detector tests that fail on harmless refactors, formatting, record ordering, incidental wording, or private implementation details without demonstrating a behavior regression. Prefer existing fixtures or focused new fixtures for repeated setup; keep tests DRY when reasonable without making the behavior under test opaque. print() is allowed in tests.

Files:

tests/generation/test_vllm_backend.py
tests/generation/test_vllm_observability.py

**/*test*.py

📄 CodeRabbit inference engine (AGENTS.md)

The unit_test marker is deprecated; use unit instead

Files:

tests/generation/test_vllm_backend.py
tests/generation/test_vllm_observability.py

src/nemo_safe_synthesizer/generation/**/*.py

⚙️ CodeRabbit configuration file

Review generation changes for retry loops, stopping conditions, invalid record handling, regex/structured output contracts, backend teardown, memory cleanup, and vLLM assumptions.

Files:

src/nemo_safe_synthesizer/generation/vllm_observability.py
src/nemo_safe_synthesizer/generation/vllm_backend.py

🧠 Learnings (4)

📚 Learning: 2026-05-14T21:47:20.140Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: tests/TESTING.md:0-0
Timestamp: 2026-05-14T21:47:20.140Z
Learning: Applies to tests/**/tests/**/test_*.py : Tests using vLLM generation backend must use the `vllm` marker and each file should have a dedicated `test-smoke-gpu-*` Make target to ensure separate process execution for GPU memory isolation.

Applied to files:

tests/generation/test_vllm_backend.py
tests/generation/test_vllm_observability.py
src/nemo_safe_synthesizer/generation/vllm_backend.py

📚 Learning: 2026-05-27T22:20:37.354Z

Learnt from: kendrickb-nvidia
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 520
File: tests/generation/test_vllm_backend.py:556-587
Timestamp: 2026-05-27T22:20:37.354Z
Learning: In NVIDIA-NeMo/Safe-Synthesizer, `tests/conftest.py`’s `pytest_collection_modifyitems` hook applies pytest category markers automatically based on each test file’s path: tests under `/e2e/` get `pytest.mark.e2e`, tests under `/smoke/` get `pytest.mark.smoke`, and all other tests get `pytest.mark.unit`. Therefore, when reviewing pytest tests outside `tests/e2e/` and `tests/smoke/`, do not flag missing explicit `pytest.mark.unit` decorators on test classes/functions as an issue (the hook will add them during collection). If a new test directory/category is introduced, ensure the hook is updated so it’s categorized correctly.

Applied to files:

tests/generation/test_vllm_backend.py
tests/generation/test_vllm_observability.py

📚 Learning: 2026-05-14T17:03:10.291Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: STYLE_GUIDE.md:0-0
Timestamp: 2026-05-14T17:03:10.291Z
Learning: Applies to **/*.py : Use bare `except Exception: pass` only in `__del__` methods where suppression is intentional.

Applied to files:

src/nemo_safe_synthesizer/generation/vllm_observability.py

📚 Learning: 2026-05-14T17:03:10.291Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: STYLE_GUIDE.md:0-0
Timestamp: 2026-05-14T17:03:10.291Z
Learning: Applies to tests/**/*.py : Use absolute imports in `tests/` (e.g., `from nemo_safe_synthesizer.observability import get_logger`).

Applied to files:

tests/generation/test_vllm_observability.py

🪛 Ruff (0.15.15)

src/nemo_safe_synthesizer/generation/vllm_observability.py

[error] 423-424: try-except-continue detected, consider logging the exception

(S112)

🔇 Additional comments (2)

src/nemo_safe_synthesizer/cli/wandb_setup.py (1)

24-25: LGTM!

Also applies to: 271-289

src/nemo_safe_synthesizer/generation/vllm_backend.py (1)

27-46: LGTM!

Also applies to: 227-230, 307-313, 669-680, 683-803

mckornfield

main q about adding this observability other than my stupid questions: do we know what the effect is on runtime?

mckornfield

would appreciate comments addressed, but we can move this through so I can review the other one a bit more cleanly

Respond to the PR #526 review round on the vLLM observability series. Rename the shared event primitive cell -> generation: GenerationObservability, vllm.generation.complete (+ .observability.emit_failed), _emit_generation_observability, log_generation_observability, and the vllm_gen wandb prefix. The benchmark harness's grid "cell" vocabulary is intentionally left unchanged. Review fixes: - guard the warning call in _emit_generation_observability so a faulty logger handler cannot turn a swallowed observability failure into a generation failure - resolve the NvmlPeakSampler device index from CUDA_VISIBLE_DEVICES instead of hardcoding physical GPU 0 on multi-GPU hosts - log degraded paths at debug (exc_info=True) in probe_engine_runtime_config and NvmlPeakSampler.__exit__ rather than swallowing silently - move _LOADAVG_HORIZON_LABELS above the model so it precedes its first reference - drop None values from engine_runtime_config in to_wandb_payload, symmetric with the scalar-field handling - add design reference URLs to the module docstring - assert the vllm.generation.complete structured-log emission in the finalizer test Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

coderabbitai

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tests/generation/test_vllm_backend.py (1)

498-512: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Add a test for logger-handler failure suppression in the emission fallback.

This test checks probe-failure swallowing, but not the new guard where logger.runtime.warning(...) itself raises. A regression there would reintroduce masking risk without failing this test.

Smallest practical test extension

     def test_emit_swallows_failures(self, base_params, mock_model_metadata, mock_schema, mock_workdir):
         """A failure inside emission must not propagate — observability is best-effort."""
         backend = create_backend(base_params, mock_model_metadata, mock_schema, mock_workdir)
         backend.llm = None

         sampler = MagicMock()
         sampler.peak_gb = 1.0

-        with patch(
-            "nemo_safe_synthesizer.generation.vllm_backend.read_vllm_runtime_metrics",
-            side_effect=RuntimeError("probe blew up"),
-        ):
+        with (
+            patch(
+                "nemo_safe_synthesizer.generation.vllm_backend.read_vllm_runtime_metrics",
+                side_effect=RuntimeError("probe blew up"),
+            ),
+            patch(
+                "nemo_safe_synthesizer.generation.vllm_backend.logger.runtime.warning",
+                side_effect=RuntimeError("logger blew up"),
+            ),
+        ):
             # Must not raise.
             backend._emit_generation_observability(sampler, None)

As per coding guidelines: "New behavior needs focused tests near the affected subsystem."

🧹 Nitpick comments (2)

src/nemo_safe_synthesizer/generation/vllm_observability.py (1)
306-307: ⚡ Quick win

Use runtime-category logging for degraded observability paths.

These are runtime-internal telemetry failures; uncategorized debug logs reduce consistency for downstream routing/filtering.
Suggested patch
-                logger.debug("nvml-sampler: nvmlShutdown failed", exc_info=True)
+                logger.runtime.debug("nvml-sampler: nvmlShutdown failed", exc_info=True)
@@
-        logger.debug("engine-probe: vllm_config unreachable; returning empty probe", exc_info=True)
+        logger.runtime.debug("engine-probe: vllm_config unreachable; returning empty probe", exc_info=True)
@@
-            logger.debug("engine-probe: field %r failed; skipping", spec.out_key, exc_info=True)
+            logger.runtime.debug("engine-probe: field %r failed; skipping", spec.out_key, exc_info=True)
As per coding guidelines, Use category loggers: .runtime for internals, .user for progress/results, .system for system events.

Also applies to: 445-460
src/nemo_safe_synthesizer/cli/wandb_setup.py (1)
286-290: ⚡ Quick win

Log swallowed WandB failures with traceback on the runtime logger.

This is a non-fatal observability boundary; keeping stack traces here improves debuggability without changing behavior.
Suggested patch
-    except Exception as exc:  # noqa: BLE001 — degraded mode
-        logger.warning(f"failed to log generation observability to wandb: {exc}")
+    except Exception:  # noqa: BLE001 — degraded mode
+        logger.runtime.debug("failed to log generation observability to wandb", exc_info=True)
As per coding guidelines, Use except Exception: + logger.debug(..., exc_info=True) for non-fatal cleanup at teardown boundaries and Use category loggers: .runtime for internals.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 987443c6-8dbf-411b-9a2a-cd88fc4f7529

📥 Commits

Reviewing files that changed from the base of the PR and between fd7f572 and 5fc70a0.

📒 Files selected for processing (5)

src/nemo_safe_synthesizer/cli/wandb_setup.py
src/nemo_safe_synthesizer/generation/vllm_backend.py
src/nemo_safe_synthesizer/generation/vllm_observability.py
tests/generation/test_vllm_backend.py
tests/generation/test_vllm_observability.py

📜 Review details

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (6)

GitHub Check: Unit Tests (3.12)
GitHub Check: Unit Tests (3.13)
GitHub Check: Unit Tests (3.11)
GitHub Check: Smoke Tests
GitHub Check: Analyze (Python)
GitHub Check: Analyze (python)

🧰 Additional context used

📓 Path-based instructions (9)

**/*.{md,markdown,py}

📄 CodeRabbit inference engine (.cursor/rules/agent-markdown-style.mdc)

**/*.{md,markdown,py}: Avoid decorative bold (**text**) in list items, body text, and docstrings; use structural cues (headers, list markers, colons, backticks) for emphasis instead
Use backticks for code identifiers, paths, and CLI commands in markdown and docstrings

Files:

src/nemo_safe_synthesizer/cli/wandb_setup.py
tests/generation/test_vllm_backend.py
src/nemo_safe_synthesizer/generation/vllm_backend.py
tests/generation/test_vllm_observability.py
src/nemo_safe_synthesizer/generation/vllm_observability.py

**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

**/*.py: Place durable implementation guidance in function and class docstrings for public contracts and source comments for local invariants
Target Python 3.11–3.13 with modern syntax (X | Y, list[str], Self). Python 3.14+ is not supported

**/*.py: Use American English spelling: 'initialize' not 'initialise', 'recognize' not 'recognise', 'color' not 'colour'.
For Pydantic models in config/, use NSSBaseModel for config/parameter models. Use raw BaseModel or module-specific bases (e.g., ReportBaseModel) for data transfer objects and internal structures.
Use BaseSettings for env/CLI settings. Prefer AliasChoices on individual fields when you need a field to respond to both its Python name and an env var name (e.g., validation_alias=AliasChoices("config_path", "NSS_CONFIG")). env_prefix is acceptable for simple settings classes.
Use Field(description=...) as the canonical field docstring for Pydantic models. Always include it.
Use assignment style (type = Field(default=..., description="...")) as the default for Pydantic model fields. Prefer it because type checkers understand default, default_factory, and alias in assignment-style Field() and synthesize correct __init__ signatures.
Use Annotated only when the field carries additional metadata beyond Field() -- ValueValidator, AutoParam, DependsOnValidator, reusable constrained type aliases, nested-type constraints, or discriminated unions.
For defaults with Annotated: put the default as a bare assignment (= value), not inside Field(default=...). Exception: default_factory has no bare-assignment equivalent, so use assignment-style Field(default_factory=...) even when the type is Annotated[...].
Use @dataclass(frozen=True) preferred for immutable value objects and validators. Mutable @dataclass acceptable for builders, accumulators, and pipeline state.
Use field(default_factory=list) for mutable defaults in dataclasses, never = [].
Use StrEnum for ...

Files:

src/nemo_safe_synthesizer/cli/wandb_setup.py
tests/generation/test_vllm_backend.py
src/nemo_safe_synthesizer/generation/vllm_backend.py
tests/generation/test_vllm_observability.py
src/nemo_safe_synthesizer/generation/vllm_observability.py

src/**/*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

src/**/*.py: Use relative imports in src/ (e.g., from ..observability import get_logger).
Never use print() statements in library code. Use get_logger(__name__) from observability.py or click.echo() for CLI.
Never use assert for validation in library code. Use if/raise for input validation. assert is fine in tests.
Every directory under src/ that contains Python files must include an __init__.py file, even if empty.

Write Google-style docstrings in source code for API reference auto-generation via mkdocstrings

Files:

src/nemo_safe_synthesizer/cli/wandb_setup.py
src/nemo_safe_synthesizer/generation/vllm_backend.py
src/nemo_safe_synthesizer/generation/vllm_observability.py

⚙️ CodeRabbit configuration file

Review library code against STYLE_GUIDE.md. Focus on behavior, API contracts, error handling, resource cleanup, typing, logging, and user-facing failures. Public APIs and nontrivial functions need Google-style docstrings.

Files:

src/nemo_safe_synthesizer/cli/wandb_setup.py
src/nemo_safe_synthesizer/generation/vllm_backend.py
src/nemo_safe_synthesizer/generation/vllm_observability.py

**/*

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

Include a newline at end of file, with no trailing whitespace.

Files:

src/nemo_safe_synthesizer/cli/wandb_setup.py
tests/generation/test_vllm_backend.py
src/nemo_safe_synthesizer/generation/vllm_backend.py
tests/generation/test_vllm_observability.py
src/nemo_safe_synthesizer/generation/vllm_observability.py

⚙️ CodeRabbit configuration file

**/*: Review as a senior maintainer for NeMo Safe Synthesizer. Prioritize issues that can change behavior, break user workflows, weaken privacy guarantees, hide failures, make tests unreliable, or create maintenance risk. Avoid generic style commentary unless it points to a concrete project convention that automated tools will not catch.
Comment only when the finding is actionable and tied to changed code. For each finding, state the impact, the condition that triggers it, and the smallest practical fix. Prefer one precise comment over broad advice. Do not ask for refactors outside the PR scope unless the changed code creates the problem.
Review type guidance: - Potential issue: use for correctness bugs, data loss, privacy leaks,
security risks, broken public APIs, invalid config behavior, missing
validation, hidden failures, nondeterministic tests, or CI breakage.

Refactor suggestion: use for local maintainability problems introduced
by the diff when they have clear future cost, such as duplicated setup,
unclear boundaries, over-mocking, avoidable complexity, or opaque test
helpers.

Nitpick: avoid in chill mode. Do not emit formatting, import-order,
wording, or style-only comments unless automated tools cannot catch the
issue and it affects maintainability.

Severity guidance: - Critical: security/privacy leaks, data loss, training/test/holdout
contamination, or broken release/package/core pipeline execution.

Major: incorrect generation/training/evaluation behavior, broken
CLI/SDK public API, invalid config defaults or validators, or GPU/vLLM
cleanup and process-isolation bugs likely to fail CI or production
runs.

Minor: localized bugs, missing focused tests for changed behavior, or
bad test patterns that weaken regression coverage.

Trivial: small cleanup with no behavior impact. Usually suppress in
chill mode.

Info: context only. Avoid unless it helps reviewers understand risk.
Safe-Synthesizer-specific review focus: - Data ...

Files:

src/nemo_safe_synthesizer/cli/wandb_setup.py
tests/generation/test_vllm_backend.py
src/nemo_safe_synthesizer/generation/vllm_backend.py
tests/generation/test_vllm_observability.py
src/nemo_safe_synthesizer/generation/vllm_observability.py

**/*.{py,sh,yaml,yml,md}

📄 CodeRabbit inference engine (CONTRIBUTING.md)

All source files (.py, .sh, .yaml, .yml, .md) must include SPDX copyright headers

Files:

src/nemo_safe_synthesizer/cli/wandb_setup.py
tests/generation/test_vllm_backend.py
src/nemo_safe_synthesizer/generation/vllm_backend.py
tests/generation/test_vllm_observability.py
src/nemo_safe_synthesizer/generation/vllm_observability.py

**/test_*.py

📄 CodeRabbit inference engine (AGENTS.md)

Use the unit marker instead of the deprecated unit_test marker for test identification

Files:

tests/generation/test_vllm_backend.py
tests/generation/test_vllm_observability.py

tests/**/*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

tests/**/*.py: Use absolute imports in tests/ (e.g., from nemo_safe_synthesizer.observability import get_logger).
Use fixtures with fixture_ prefix convention for grep-ability. Add a one-line docstring describing the fixture's purpose and data.
Use function-scoped fixtures by default. Session scope only when empirically justified by test runtime.
Use bare assert as the primary assertion style; pytest.raises() with match= for exceptions; pytest.approx() for floating-point comparisons.
Markers are auto-assigned by path via pytest_collection_modifyitems (/e2e/ -> e2e, /smoke/ -> smoke, default -> unit). Explicit markers: @pytest.mark.slow, @pytest.mark.requires_gpu, @pytest.mark.timeout().
Use tmp_path fixture for file operations, never write to the repo tree.
Mark CUDA-dependent tests with @pytest.mark.e2e, @pytest.mark.smoke, or @pytest.mark.requires_gpu.
Mock only external boundaries, not internal implementation details.
Ensure test isolation: no shared mutable state or execution-order dependencies between tests. Include necessary setup in the test or a fixture.
Use @pytest.mark.parametrize for testing multiple input combinations rather than copy-pasting similar tests.

tests/**/*.py: Run mise run test to execute unit tests (excludes slow unit tests, smoke and e2e)
New features must include tests; bug fixes must include regression tests

tests/**/*.py: Define pytest markers in pytest.ini with --strict-markers enabled. Use exactly one category marker (unit, smoke, e2e) per test. Category modifiers include slow for long-running tests, requires_gpu for CUDA-dependent tests, vllm for vLLM backend tests, smollm2 for SmolLM2 Hub download tests, and noautouse to skip autouse fixtures
For ParsedResponse mocking, use valid_records=[...], invalid_records=[...], errors=[...], and prompt_number=int. Use fixture_mock_processor or fixture_mock_processor_without_valid_records helpers
Use `pytest.impo...

Files:

tests/generation/test_vllm_backend.py
tests/generation/test_vllm_observability.py

⚙️ CodeRabbit configuration file

Review tests against tests/TESTING.md. Check marker usage, fixture naming, tmp_path usage, determinism, and GPU/vLLM process-isolation requirements. Flag slop tests that only check that code runs, assert result is not None when stronger invariants exist, over-mock internal implementation details, patch around the bug instead of reproducing it, or add broad snapshot/golden churn without a clear contract. Flag change detector tests that fail on harmless refactors, formatting, record ordering, incidental wording, or private implementation details without demonstrating a behavior regression. Prefer existing fixtures or focused new fixtures for repeated setup; keep tests DRY when reasonable without making the behavior under test opaque. print() is allowed in tests.

Files:

tests/generation/test_vllm_backend.py
tests/generation/test_vllm_observability.py

tests/**

📄 CodeRabbit inference engine (.cursor/rules/repo-navigation.mdc)

tests/**: Mirror src/ directory structure in tests/ directory for test organization
Auto-mark tests by directory: tests/e2e/ → e2e, tests/smoke/ → smoke, otherwise default to unit

Files:

tests/generation/test_vllm_backend.py
tests/generation/test_vllm_observability.py

src/nemo_safe_synthesizer/generation/**/*.py

⚙️ CodeRabbit configuration file

Review generation changes for retry loops, stopping conditions, invalid record handling, regex/structured output contracts, backend teardown, memory cleanup, and vLLM assumptions.

Files:

src/nemo_safe_synthesizer/generation/vllm_backend.py
src/nemo_safe_synthesizer/generation/vllm_observability.py

🧠 Learnings (14)

📚 Learning: 2026-06-04T16:14:09.868Z

Learnt from: binaryaaron
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 526
File: tests/generation/test_vllm_backend.py:399-509
Timestamp: 2026-06-04T16:14:09.868Z
Learning: In this repo, only apply `pytest.mark.vllm` to smoke tests under `tests/smoke/` that actually run real vLLM GPU generation and therefore require per-file process isolation (e.g., `test-smoke-gpu-*` Makefile targets). Do not apply `pytest.mark.vllm` to unit-style tests under `tests/generation/` that merely import `vllm_backend` but never instantiate a real vLLM engine and never call `.generate()` (GPU not required). Note that `tests/conftest.py` auto-marks these as `unit` via `pytest_collection_modifyitems`, and `vllm` is not among the auto-mark categories—so if a test in `tests/generation/` has `vllm`, it should be treated as a review issue unless it meets the real GPU generation criteria above.

Applied to files:

tests/generation/test_vllm_backend.py
tests/generation/test_vllm_observability.py

📚 Learning: 2026-06-03T23:08:50.142Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: tests/TESTING.md:0-0
Timestamp: 2026-06-03T23:08:50.142Z
Learning: Applies to tests/**/*.py : For vLLM tests that call `.generate()`, mark with `pytest.mark.vllm`, use per-file process isolation (`-n 0`), and create dedicated `test:smoke:gpu:*` mise tasks. vLLM pre-allocates all GPU memory and never releases it within a process, causing OOM in later tests if not isolated

Applied to files:

tests/generation/test_vllm_backend.py

📚 Learning: 2026-06-03T23:08:50.142Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: tests/TESTING.md:0-0
Timestamp: 2026-06-03T23:08:50.142Z
Learning: Applies to tests/**/*.py : Define pytest markers in `pytest.ini` with `--strict-markers` enabled. Use exactly one category marker (`unit`, `smoke`, `e2e`) per test. Category modifiers include `slow` for long-running tests, `requires_gpu` for CUDA-dependent tests, `vllm` for vLLM backend tests, `smollm2` for SmolLM2 Hub download tests, and `noautouse` to skip autouse fixtures

Applied to files:

tests/generation/test_vllm_backend.py

📚 Learning: 2026-05-27T22:20:37.354Z

Learnt from: kendrickb-nvidia
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 520
File: tests/generation/test_vllm_backend.py:556-587
Timestamp: 2026-05-27T22:20:37.354Z
Learning: In NVIDIA-NeMo/Safe-Synthesizer, `tests/conftest.py`’s `pytest_collection_modifyitems` hook applies pytest category markers automatically based on each test file’s path: tests under `/e2e/` get `pytest.mark.e2e`, tests under `/smoke/` get `pytest.mark.smoke`, and all other tests get `pytest.mark.unit`. Therefore, when reviewing pytest tests outside `tests/e2e/` and `tests/smoke/`, do not flag missing explicit `pytest.mark.unit` decorators on test classes/functions as an issue (the hook will add them during collection). If a new test directory/category is introduced, ensure the hook is updated so it’s categorized correctly.

Applied to files:

tests/generation/test_vllm_backend.py
tests/generation/test_vllm_observability.py

📚 Learning: 2026-06-03T23:08:14.151Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: STYLE_GUIDE.md:0-0
Timestamp: 2026-06-03T23:08:14.151Z
Learning: Applies to tests/**/*.py : Mark CUDA-dependent tests with `pytest.mark.e2e`, `pytest.mark.smoke`, or `pytest.mark.requires_gpu`.

Applied to files:

tests/generation/test_vllm_backend.py

📚 Learning: 2026-06-03T23:08:14.151Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: STYLE_GUIDE.md:0-0
Timestamp: 2026-06-03T23:08:14.151Z
Learning: Applies to tests/**/*.py : Markers are auto-assigned by path via `pytest_collection_modifyitems` (`/e2e/` -> `e2e`, `/smoke/` -> `smoke`, default -> `unit`). Explicit markers: `pytest.mark.slow`, `pytest.mark.requires_gpu`, `pytest.mark.timeout()`.

Applied to files:

tests/generation/test_vllm_backend.py

📚 Learning: 2026-06-03T23:06:56.798Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: AGENTS.md:0-0
Timestamp: 2026-06-03T23:06:56.798Z
Learning: Applies to **/test_*.py : Use the `unit` marker instead of the deprecated `unit_test` marker for test identification

Applied to files:

tests/generation/test_vllm_backend.py

📚 Learning: 2026-06-03T23:08:14.151Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: STYLE_GUIDE.md:0-0
Timestamp: 2026-06-03T23:08:14.151Z
Learning: Applies to tests/**/*.py : Use absolute imports in `tests/` (e.g., `from nemo_safe_synthesizer.observability import get_logger`).

Applied to files:

tests/generation/test_vllm_backend.py

📚 Learning: 2026-06-03T23:08:50.142Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: tests/TESTING.md:0-0
Timestamp: 2026-06-03T23:08:50.142Z
Learning: Read and understand key conftest.py files in order: `tests/conftest.py` (auto-marking, test helpers), `pytest.ini` (markers, asyncio, timeout), `tests/evaluation/conftest.py` (Faker-based data generation), and `tests/generation/conftest.py` (JSONL/schema fixtures) before writing tests

Applied to files:

tests/generation/test_vllm_backend.py

📚 Learning: 2026-06-03T23:09:02.641Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: .cursor/rules/repo-navigation.mdc:0-0
Timestamp: 2026-06-03T23:09:02.641Z
Learning: Applies to tests/** : Auto-mark tests by directory: `tests/e2e/` → `e2e`, `tests/smoke/` → `smoke`, otherwise default to `unit`

Applied to files:

tests/generation/test_vllm_backend.py

📚 Learning: 2026-06-03T23:09:02.641Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: .cursor/rules/repo-navigation.mdc:0-0
Timestamp: 2026-06-03T23:09:02.641Z
Learning: Applies to pytest.ini : Define test markers in `pytest.ini`: `unit`, `slow`, `smoke`, `e2e`, `requires_gpu`, `noautouse`

Applied to files:

tests/generation/test_vllm_backend.py

📚 Learning: 2026-06-04T16:14:09.868Z

Learnt from: binaryaaron
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 526
File: tests/generation/test_vllm_backend.py:399-509
Timestamp: 2026-06-04T16:14:09.868Z
Learning: In NVIDIA-NeMo/Safe-Synthesizer, the `pytest.mark.vllm` marker is scoped exclusively to tests under `tests/smoke/` that invoke real vLLM GPU generation and need per-file process isolation (via `test-smoke-gpu-*` Makefile targets). Mocked unit tests in `tests/generation/` that import from `vllm_backend` but never instantiate a real engine (no GPU required, no `.generate()` call) should NOT carry the `vllm` marker. `tests/conftest.py` auto-marks these files as `unit` via `pytest_collection_modifyitems`, and `vllm` is not part of the auto-mark categories. Sibling files like `test_vllm_shutdown.py` and `test_timeseries_backend.py` follow this convention.

Applied to files:

src/nemo_safe_synthesizer/generation/vllm_backend.py
src/nemo_safe_synthesizer/generation/vllm_observability.py

📚 Learning: 2026-06-03T23:08:14.151Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: STYLE_GUIDE.md:0-0
Timestamp: 2026-06-03T23:08:14.151Z
Learning: Applies to **/*.py : Add `from __future__ import annotations` to every module. This makes all annotations strings consistently, preventing accidental runtime evaluation and aligning with type-checker expectations.

Applied to files:

src/nemo_safe_synthesizer/generation/vllm_backend.py

📚 Learning: 2026-06-03T23:08:14.151Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: STYLE_GUIDE.md:0-0
Timestamp: 2026-06-03T23:08:14.151Z
Learning: Applies to **/*.py : Use `except Exception:` + `logger.debug(..., exc_info=True)` for non-fatal cleanup at teardown boundaries. Bare `except Exception: pass` only in `__del__` methods.

Applied to files:

src/nemo_safe_synthesizer/generation/vllm_observability.py

🔇 Additional comments (2)

src/nemo_safe_synthesizer/generation/vllm_backend.py (1)

8-680: LGTM!

Also applies to: 766-798

tests/generation/test_vllm_backend.py (1)

22-22: LGTM!

Also applies to: 399-497

Adds ``src/nemo_safe_synthesizer/generation/vllm_observability.py`` — a self-contained module housing four reusable measurement primitives plus the schema for the ``vllm.cell.complete`` structured event that ``VllmBackend.generate()`` will emit (next commit). This is a **pure-additive** commit: no existing code paths call the new module yet. Subsequent commits wire production (``VllmBackend.generate``) and benchmark (``tools/vllm_benchmark.py``) consumers. ## Primitives - ``NvmlPeakSampler`` — context-manager + daemon thread polling ``pynvml.nvmlDeviceGetMemoryInfo`` at 250 ms cadence, tracking peak device-wide VRAM. Reads at the driver layer so it sees vLLM worker-subprocess allocations regardless of which process holds the torch handle — sidesteps the ``VLLM_ENABLE_V1_MULTIPROCESSING=1`` blind spot where ``torch.cuda.max_memory_allocated()`` in the parent process reads 0. - ``read_loadavg() -> tuple[float, float, float] | None`` — host ``/proc/loadavg`` snapshot (1m, 5m, 15m). Linux-only; ``None`` on read failure. - ``probe_engine_runtime_config(llm) -> dict[str, Any]`` — best-effort introspection of ``llm.llm_engine.vllm_config`` for the scheduler/cache/speculative settings. Spans vLLM v0/v1 attribute naming (``llm_engine`` vs ``engine``). Empty dict on any failure. - ``read_vllm_runtime_metrics(llm) -> dict[str, float | None]`` — one-shot snapshot of ``llm.get_metrics()`` for ``vllm:kv_cache_usage_perc`` (gauge), ``prefix_cache_hit_rate`` (derived from ``hits/queries`` counters), and ``spec_accept_rate`` (derived from ``spec_decode_num_{accepted,draft}_tokens`` counters). ``None`` for any metric the engine didn't expose — distinguishes "not measured" from "measured zero". Plus ``flag_engagement_mismatches(intended, actual) -> list[str]`` — a helper that cross-checks an intended-config dict against a probed runtime-config dict, used by the caller to set the ``flag_did_not_engage`` bit. Dict-vs-dict shape (not pydantic) so it works whether the caller has a typed config model or just raw vLLM kwargs. ## Schema ``CellObservability`` pydantic model defines the ``vllm.cell.complete`` event payload: peak_vram_gb, kv_cache_usage_perc, prefix_cache_hit_rate, spec_accept_rate, loadavg_pre/post, engine_runtime_config (dict), flag_did_not_engage (bool). Every measurement field is optional with a sensible default — producers populate what they can capture and consumers (logs, wandb, benchmark aggregator) silently drop ``None`` values. ``model_config = ConfigDict(extra='forbid')`` so producers are forced to update the schema when they add new fields, preventing silent drift. ## Dependencies Only ``pydantic``, ``pynvml`` (already in the project's ``cu129`` extra via ``nvidia-ml-py``), and the existing ``observability.get_logger``. No PR-1-introduced modules (``vllm_engine_factory``, ``vllm_trace``, etc.) are imported — this module stands alone on top of ``main``. Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

…lmBackend.generate Wires the ``vllm_observability`` primitives into ``VllmBackend`` so every production generation invocation emits a ``vllm.cell.complete`` structured event carrying peak device VRAM, host loadavg pre/post, vLLM's kv_cache_usage_perc / prefix_cache_hit_rate / spec_accept_rate, and the engine's effective runtime config. ## What changes - ``VllmBackend.initialize`` (after ``self.llm = vLLM(...)``) caches ``self._engine_runtime_config = probe_engine_runtime_config(self.llm)``. Probed once at engine-init time; consumed by every subsequent ``generate()`` call. - ``VllmBackend.generate``: - Captures ``loadavg_pre`` at the top. - Enters an :class:`NvmlPeakSampler` context wrapping the entire body via manual ``__enter__`` / ``__exit__`` calls inside a try/finally, so the sampler thread always shuts down and the observability event always emits — even when generation raises mid-batch. - At end of body (finally branch): reads ``read_vllm_runtime_metrics(self.llm)``, captures ``loadavg_post``, builds a ``CellObservability`` event, emits via ``logger.runtime.info("vllm.cell.complete", extra={"ctx": ...})``. ## Degraded-mode behavior Every primitive returns ``None`` / empty-dict when its data source is unavailable, and the emission still fires. Specifically: - No GPU / pynvml missing → ``peak_vram_gb=None``. - Non-Linux → ``loadavg_pre/post=None``. - vLLM doesn't expose a given metric → that field is ``None``. - ``llm.get_metrics()`` raises → log a warning and emit with all metric fields ``None``. - Engine config probe fails → ``engine_runtime_config={}``. The event itself is never optional — every ``generate()`` call emits exactly one ``vllm.cell.complete`` to the structured-log surface that PR-1's telemetry work established (``logger.runtime.*`` aliases on the ``CategoryLogger`` from ``observability.py``). ## What ``flag_did_not_engage`` does (and doesn't) do here The bit is set to ``False`` unconditionally in production because ``VllmBackend`` doesn't carry an intended-overrides dict to compare against the probed runtime config — production passes whatever ``SafeSynthesizerParameters`` specifies, and the engine accepts or defaults each field independently. Benchmark callers (next PR's harness) carry their own intended-overrides dict per candidate and compute the mismatch themselves; they construct the ``CellObservability`` event with the bit populated correctly. The schema field is here so benchmark consumers don't have to extend it. ## Wandb wiring is a separate commit This commit only emits the structured-log event. The next commit adds the wandb-side emission via ``wandb_setup.py`` so production runs with ``WANDB_MODE!=disabled`` automatically log the same payload to the active wandb run. Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

``generate()`` reads ``self._engine_runtime_config`` when building the ``CellObservability`` event, but the attribute was only set inside ``initialize()`` (right after ``self.llm = vLLM(...)``). Existing test_vllm_backend.py tests construct VllmBackend without calling ``initialize()`` (they mock the engine separately), so ``generate()`` raised ``AttributeError`` on those tests. Fix: declare ``self._engine_runtime_config: dict[str, Any] = {}`` in ``__init__`` so the attribute always exists. ``initialize()`` overwrites it with the actual probe result post-engine-build; tests that skip ``initialize()`` see an empty dict, which flows through to ``CellObservability.engine_runtime_config`` as the documented 'probe unavailable' value. Closes 5 failures in test_vllm_backend.py: - test_native_eos_stopping_for_grouped_processor - test_no_stop_kwargs_for_tabular_processor - test_large_context_grouped_generation_has_eos_stop - test_uses_generation_max_tokens_for_with_cached_prompt_len - test_passes_cached_prompt_token_count_when_engine_initialized Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

…ntracts Review-driven cleanup of the vllm.cell.complete observability path: - generate(): wrap NvmlPeakSampler with `with`, extract _run_generation and _emit_cell_observability so generate() reads as a thin observability bracket; drop the manual __enter__/__exit__ calls. - read_vllm_runtime_metrics: widen the degraded-mode guard over the whole body and return a VllmRuntimeMetrics TypedDict (stable keys, float|None values); remove the caller's redundant try/except and duplicated literal. - probe_engine_runtime_config: replace nested if-ladders with a declarative _ProbeField table plus a single loop; degrade per-field instead of whole-probe; derive ENGINE_CONFIG_CHECKED_FIELDS from the table so the two cannot drift. - Narrow NvmlPeakSampler guards to pynvml.NVMLError; type the probe input as object and log_cell_observability's event as CellObservability. - Add integration tests for the generate() emission path and engine-config caching, plus probe-table and checked-fields coverage. Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

Collapse the per-metric if/elif dispatch into a set-driven collector (_collect_raw_metrics) feeding a functional VllmRuntimeMetrics construction, replacing the mutable out-dict with named, intention- revealing helpers: - _COLLECTED_METRICS: the raw counters pulled from llm.get_metrics() - _safe_ratio: factors out the duplicated numerator/denominator guard (prefix-cache hit rate, spec-accept rate) - _empty_runtime_metrics: the single degraded-mode snapshot Narrows the exception surface to wrap only the engine call + numeric coercion; the ratio derivation is pure dict arithmetic and stays outside the guard. Behavior and the stable three-key contract are unchanged (covered by existing test_vllm_observability tests). Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

Respond to the PR #526 review round on the vLLM observability series. Rename the shared event primitive cell -> generation: GenerationObservability, vllm.generation.complete (+ .observability.emit_failed), _emit_generation_observability, log_generation_observability, and the vllm_gen wandb prefix. The benchmark harness's grid "cell" vocabulary is intentionally left unchanged. Review fixes: - guard the warning call in _emit_generation_observability so a faulty logger handler cannot turn a swallowed observability failure into a generation failure - resolve the NvmlPeakSampler device index from CUDA_VISIBLE_DEVICES instead of hardcoding physical GPU 0 on multi-GPU hosts - log degraded paths at debug (exc_info=True) in probe_engine_runtime_config and NvmlPeakSampler.__exit__ rather than swallowing silently - move _LOADAVG_HORIZON_LABELS above the model so it precedes its first reference - drop None values from engine_runtime_config in to_wandb_payload, symmetric with the scalar-field handling - add design reference URLs to the module docstring - assert the vllm.generation.complete structured-log emission in the finalizer test Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

coderabbitai

Actionable comments posted: 1

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 5d01eadb-1d2e-4044-9003-13a7b1a698ee

📥 Commits

Reviewing files that changed from the base of the PR and between 5fc70a0 and 23c69c5.

📒 Files selected for processing (7)

src/nemo_safe_synthesizer/cli/wandb_setup.py
src/nemo_safe_synthesizer/generation/vllm_backend.py
src/nemo_safe_synthesizer/generation/vllm_observability.py
src/nemo_safe_synthesizer/observability.py
tests/generation/test_vllm_backend.py
tests/generation/test_vllm_observability.py
tests/test_observability.py

🚧 Files skipped from review as they are similar to previous changes (3)

src/nemo_safe_synthesizer/generation/vllm_backend.py
tests/generation/test_vllm_observability.py
tests/generation/test_vllm_backend.py

📜 Review details

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)

GitHub Check: Unit Tests (3.12)
GitHub Check: Unit Tests (3.11)
GitHub Check: Unit Tests (3.13)
GitHub Check: Smoke Tests

🧰 Additional context used

📓 Path-based instructions (13)

**/*.{md,markdown,py}

📄 CodeRabbit inference engine (.cursor/rules/agent-markdown-style.mdc)

**/*.{md,markdown,py}: Avoid decorative bold (**text**) in list items, body text, and docstrings; use structural cues (headers, list markers, colons, backticks) for emphasis instead
Use backticks for code identifiers, paths, and CLI commands in markdown and docstrings

Files:

src/nemo_safe_synthesizer/observability.py
tests/test_observability.py
src/nemo_safe_synthesizer/cli/wandb_setup.py
src/nemo_safe_synthesizer/generation/vllm_observability.py

**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

**/*.py: Place durable implementation guidance in function and class docstrings for public contracts and source comments for local invariants
Target Python 3.11–3.13 with modern syntax (X | Y, list[str], Self). Python 3.14+ is not supported

**/*.py: Use American English spelling: 'initialize' not 'initialise', 'recognize' not 'recognise', 'color' not 'colour'.
For Pydantic models in config/, use NSSBaseModel for config/parameter models. Use raw BaseModel or module-specific bases (e.g., ReportBaseModel) for data transfer objects and internal structures.
Use BaseSettings for env/CLI settings. Prefer AliasChoices on individual fields when you need a field to respond to both its Python name and an env var name (e.g., validation_alias=AliasChoices("config_path", "NSS_CONFIG")). env_prefix is acceptable for simple settings classes.
Use Field(description=...) as the canonical field docstring for Pydantic models. Always include it.
Use assignment style (type = Field(default=..., description="...")) as the default for Pydantic model fields. Prefer it because type checkers understand default, default_factory, and alias in assignment-style Field() and synthesize correct __init__ signatures.
Use Annotated only when the field carries additional metadata beyond Field() -- ValueValidator, AutoParam, DependsOnValidator, reusable constrained type aliases, nested-type constraints, or discriminated unions.
For defaults with Annotated: put the default as a bare assignment (= value), not inside Field(default=...). Exception: default_factory has no bare-assignment equivalent, so use assignment-style Field(default_factory=...) even when the type is Annotated[...].
Use @dataclass(frozen=True) preferred for immutable value objects and validators. Mutable @dataclass acceptable for builders, accumulators, and pipeline state.
Use field(default_factory=list) for mutable defaults in dataclasses, never = [].
Use StrEnum for ...

Files:

src/nemo_safe_synthesizer/observability.py
tests/test_observability.py
src/nemo_safe_synthesizer/cli/wandb_setup.py
src/nemo_safe_synthesizer/generation/vllm_observability.py

src/**/*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

src/**/*.py: Use relative imports in src/ (e.g., from ..observability import get_logger).
Never use print() statements in library code. Use get_logger(__name__) from observability.py or click.echo() for CLI.
Never use assert for validation in library code. Use if/raise for input validation. assert is fine in tests.
Every directory under src/ that contains Python files must include an __init__.py file, even if empty.

Files:

src/nemo_safe_synthesizer/observability.py
src/nemo_safe_synthesizer/cli/wandb_setup.py
src/nemo_safe_synthesizer/generation/vllm_observability.py

⚙️ CodeRabbit configuration file

Review library code against STYLE_GUIDE.md. Focus on behavior, API contracts, error handling, resource cleanup, typing, logging, and user-facing failures. Public APIs and nontrivial functions need Google-style docstrings.

Files:

src/nemo_safe_synthesizer/observability.py
src/nemo_safe_synthesizer/cli/wandb_setup.py
src/nemo_safe_synthesizer/generation/vllm_observability.py

**/*

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

Include a newline at end of file, with no trailing whitespace.

Files:

src/nemo_safe_synthesizer/observability.py
tests/test_observability.py
src/nemo_safe_synthesizer/cli/wandb_setup.py
src/nemo_safe_synthesizer/generation/vllm_observability.py

⚙️ CodeRabbit configuration file

**/*: Review as a senior maintainer for NeMo Safe Synthesizer. Prioritize issues that can change behavior, break user workflows, weaken privacy guarantees, hide failures, make tests unreliable, or create maintenance risk. Avoid generic style commentary unless it points to a concrete project convention that automated tools will not catch.
Comment only when the finding is actionable and tied to changed code. For each finding, state the impact, the condition that triggers it, and the smallest practical fix. Prefer one precise comment over broad advice. Do not ask for refactors outside the PR scope unless the changed code creates the problem.
Review type guidance: - Potential issue: use for correctness bugs, data loss, privacy leaks,
security risks, broken public APIs, invalid config behavior, missing
validation, hidden failures, nondeterministic tests, or CI breakage.

Refactor suggestion: use for local maintainability problems introduced
by the diff when they have clear future cost, such as duplicated setup,
unclear boundaries, over-mocking, avoidable complexity, or opaque test
helpers.

Nitpick: avoid in chill mode. Do not emit formatting, import-order,
wording, or style-only comments unless automated tools cannot catch the
issue and it affects maintainability.

Severity guidance: - Critical: security/privacy leaks, data loss, training/test/holdout
contamination, or broken release/package/core pipeline execution.

Major: incorrect generation/training/evaluation behavior, broken
CLI/SDK public API, invalid config defaults or validators, or GPU/vLLM
cleanup and process-isolation bugs likely to fail CI or production
runs.

Minor: localized bugs, missing focused tests for changed behavior, or
bad test patterns that weaken regression coverage.

Trivial: small cleanup with no behavior impact. Usually suppress in
chill mode.

Info: context only. Avoid unless it helps reviewers understand risk.
Safe-Synthesizer-specific review focus: - Data ...

Files:

src/nemo_safe_synthesizer/observability.py
tests/test_observability.py
src/nemo_safe_synthesizer/cli/wandb_setup.py
src/nemo_safe_synthesizer/generation/vllm_observability.py

**/*.{py,sh,yaml,yml,md}

📄 CodeRabbit inference engine (CONTRIBUTING.md)

All source files (.py, .sh, .yaml, .yml, .md) require SPDX copyright headers. Use mise run format to add them automatically

Files:

src/nemo_safe_synthesizer/observability.py
tests/test_observability.py
src/nemo_safe_synthesizer/cli/wandb_setup.py
src/nemo_safe_synthesizer/generation/vllm_observability.py

src/nemo_safe_synthesizer/**/*.py

📄 CodeRabbit inference engine (CONTRIBUTING.md)

All Python source code must follow Google-style docstring format for auto-generation in API reference

Files:

src/nemo_safe_synthesizer/observability.py
src/nemo_safe_synthesizer/cli/wandb_setup.py
src/nemo_safe_synthesizer/generation/vllm_observability.py

{src/**/*,test/**/*}

📄 CodeRabbit inference engine (CONTRIBUTING.md)

All src and test files are owned by @NVIDIA-NeMo/safe-synthesizer-reviewers and require their code review

Files:

src/nemo_safe_synthesizer/observability.py
src/nemo_safe_synthesizer/cli/wandb_setup.py
src/nemo_safe_synthesizer/generation/vllm_observability.py

**

⚙️ CodeRabbit configuration file

**:

AGENTS.md

Guide for AI agents (Cursor, Windsurf, Claude Code, etc.) working in the Safe-Synthesizer repo.

This project loads local developer preferences from @AGENTS.local.md. You MUST read this file if it exists and give its instructions top priority.

Skills

Repo-specific skills live in .agents/skills/; see .agents/README.md for the catalog. Read a skill when the task matches its scope instead of copying workflow details into this file.

Durable implementation guidance belongs with the code it describes: function and class docstrings for public contracts and source comments for local invariants. Test-suite guidance belongs in tests/TESTING.md.

Repo Conventions

See STYLE_GUIDE.md for detailed code style conventions (Python, markdown, Dockerfiles, shell scripts, testing, config files, docstrings).

Use uv for everything -- never pip or raw python. Python 3.11–3.13 with modern syntax (X | Y, list[str], Self). Python 3.14+ is not supported.

Common commands: mise run test (unit tests), mise run format (auto-fix formatting + lint + copyright), mise run check (read-only local quality checks), mise run validate (pre-PR quality, lock, and CI unit checks), mise run typecheck (ty only). Always use mise tasks or the wrapper scripts in tools/ instead of running ruff or ty directly. Use uv run for Python execution. When in doubt, inspect mise tasks and pytest --markers.

The canonical uv sync command for a full GPU/dev environment is:
uv sync --frozen --extra cu129 --extra engine --group dev
Bare uv sync --frozen (without extras) installs an incomplete environment -- ty, import checks, and GPU tests will fail.

Feature branches off main. Branch names often include an issue number prefix (e.g., <author>/123-short-name).

Do ...

Files:

src/nemo_safe_synthesizer/observability.py
tests/test_observability.py
src/nemo_safe_synthesizer/cli/wandb_setup.py
src/nemo_safe_synthesizer/generation/vllm_observability.py

**/test_*.py

📄 CodeRabbit inference engine (AGENTS.md)

Use the unit marker instead of the deprecated unit_test marker for test identification

Files:

tests/test_observability.py

tests/**/*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

tests/**/*.py: Use absolute imports in tests/ (e.g., from nemo_safe_synthesizer.observability import get_logger).
Use fixtures with fixture_ prefix convention for grep-ability. Add a one-line docstring describing the fixture's purpose and data.
Use function-scoped fixtures by default. Session scope only when empirically justified by test runtime.
Use bare assert as the primary assertion style; pytest.raises() with match= for exceptions; pytest.approx() for floating-point comparisons.
Markers are auto-assigned by path via pytest_collection_modifyitems (/e2e/ -> e2e, /smoke/ -> smoke, default -> unit). Explicit markers: @pytest.mark.slow, @pytest.mark.requires_gpu, @pytest.mark.timeout().
Use tmp_path fixture for file operations, never write to the repo tree.
Mark CUDA-dependent tests with @pytest.mark.e2e, @pytest.mark.smoke, or @pytest.mark.requires_gpu.
Mock only external boundaries, not internal implementation details.
Ensure test isolation: no shared mutable state or execution-order dependencies between tests. Include necessary setup in the test or a fixture.
Use @pytest.mark.parametrize for testing multiple input combinations rather than copy-pasting similar tests.

tests/**/*.py: Define pytest markers in pytest.ini with --strict-markers enabled. Use exactly one category marker (unit, smoke, e2e) per test. Category modifiers include slow for long-running tests, requires_gpu for CUDA-dependent tests, vllm for vLLM backend tests, smollm2 for SmolLM2 Hub download tests, and noautouse to skip autouse fixtures
For ParsedResponse mocking, use valid_records=[...], invalid_records=[...], errors=[...], and prompt_number=int. Use fixture_mock_processor or fixture_mock_processor_without_valid_records helpers
Use pytest.importorskip to gate tests on optional dependencies that require specific extras (e.g., sentence_transformers and vllm require cu129 extra)
For vLLM tests that call `.gen...

Files:

tests/test_observability.py

⚙️ CodeRabbit configuration file

Review tests against tests/TESTING.md. Check marker usage, fixture naming, tmp_path usage, determinism, and GPU/vLLM process-isolation requirements. Flag slop tests that only check that code runs, assert result is not None when stronger invariants exist, over-mock internal implementation details, patch around the bug instead of reproducing it, or add broad snapshot/golden churn without a clear contract. Flag change detector tests that fail on harmless refactors, formatting, record ordering, incidental wording, or private implementation details without demonstrating a behavior regression. Prefer existing fixtures or focused new fixtures for repeated setup; keep tests DRY when reasonable without making the behavior under test opaque. print() is allowed in tests.

Files:

tests/test_observability.py

tests/test_*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

File naming: test_*.py; class naming: Test*; function naming: test_<module>_<expected_behavior>.

Files:

tests/test_observability.py

tests/**

📄 CodeRabbit inference engine (.cursor/rules/repo-navigation.mdc)

tests/**: Mirror src/ directory structure in tests/ directory for test organization
Auto-mark tests by directory: tests/e2e/ → e2e, tests/smoke/ → smoke, otherwise default to unit

Files:

tests/test_observability.py

⚙️ CodeRabbit configuration file

tests/**:

Testing Guide

Comprehensive testing reference for Safe-Synthesizer developers. Covers commands, markers, test data, fixtures, and gotchas.

Read First

tests/conftest.py -- auto-marking, load_test_dataset/load_test_dataframe, fixture_mock_processor pattern
pytest.ini -- markers, asyncio, timeout
tests/evaluation/conftest.py -- most complex: Faker-based make_df, nullable dtype conversion
tests/generation/conftest.py -- JSONL/schema fixtures, fixture_valid_iris_dataset_jsonl_and_schema

Running Tests

All mise test tasks, grouped by scope:

mise run test                              # Unit (excludes slow, e2e, and smoke)
mise run test:unit-slow                    # Unit tests including slow (excludes e2e and smoke)
mise run test:smoke                        # CPU smoke tests (~few min, no GPU required)
mise run test:smoke:gpu                    # All staged GPU smoke tests (requires CUDA)
mise run test:smoke:gpu:train-only
mise run test:smoke:gpu:generation
mise run test:smoke:gpu:resume
mise run test:smoke:gpu:structured-generation
mise run test:smoke:gpu:timeseries
mise run test:smoke:gpu:smollm2
mise run test:e2e                          # All e2e (requires CUDA) -- runs default + dp
mise run test:e2e:default                  # e2e default (no-DP) tests only
mise run test:e2e:dp                       # e2e DP tests only
mise run test:ci                           # CI unit tests with coverage (excludes slow, e2e, gpu, smoke)
mise run test:ci-slow                      # CI slow tests with coverage
mise run test:ci-container                 # CI tests in a Linux container (Docker/Podman)

Run a single test:

uv run --frozen pytest tests/path/test_file.py::test_name -vvs -n0

Test runner: uv run --frozen pytest -n auto --dist loadscope -vv...

Files:

tests/test_observability.py

src/nemo_safe_synthesizer/generation/**/*.py

⚙️ CodeRabbit configuration file

Review generation changes for retry loops, stopping conditions, invalid record handling, regex/structured output contracts, backend teardown, memory cleanup, and vLLM assumptions.

Files:

src/nemo_safe_synthesizer/generation/vllm_observability.py

🧠 Learnings (6)

📚 Learning: 2026-06-04T16:14:16.006Z

Learnt from: binaryaaron
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 526
File: tests/generation/test_vllm_backend.py:399-509
Timestamp: 2026-06-04T16:14:16.006Z
Learning: In NVIDIA-NeMo/Safe-Synthesizer, the `pytest.mark.vllm` marker is scoped exclusively to tests under `tests/smoke/` that invoke real vLLM GPU generation and need per-file process isolation (via `test-smoke-gpu-*` Makefile targets). Mocked unit tests in `tests/generation/` that import from `vllm_backend` but never instantiate a real engine (no GPU required, no `.generate()` call) should NOT carry the `vllm` marker. `tests/conftest.py` auto-marks these files as `unit` via `pytest_collection_modifyitems`, and `vllm` is not part of the auto-mark categories. Sibling files like `test_vllm_shutdown.py` and `test_timeseries_backend.py` follow this convention.

Applied to files:

tests/test_observability.py
src/nemo_safe_synthesizer/generation/vllm_observability.py

📚 Learning: 2026-06-03T23:08:50.142Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: tests/TESTING.md:0-0
Timestamp: 2026-06-03T23:08:50.142Z
Learning: Applies to tests/**/*.py : For vLLM tests that call `.generate()`, mark with `pytest.mark.vllm`, use per-file process isolation (`-n 0`), and create dedicated `test:smoke:gpu:*` mise tasks. vLLM pre-allocates all GPU memory and never releases it within a process, causing OOM in later tests if not isolated

Applied to files:

tests/test_observability.py

📚 Learning: 2026-05-27T22:20:37.354Z

Learnt from: kendrickb-nvidia
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 520
File: tests/generation/test_vllm_backend.py:556-587
Timestamp: 2026-05-27T22:20:37.354Z
Learning: In NVIDIA-NeMo/Safe-Synthesizer, `tests/conftest.py`’s `pytest_collection_modifyitems` hook applies pytest category markers automatically based on each test file’s path: tests under `/e2e/` get `pytest.mark.e2e`, tests under `/smoke/` get `pytest.mark.smoke`, and all other tests get `pytest.mark.unit`. Therefore, when reviewing pytest tests outside `tests/e2e/` and `tests/smoke/`, do not flag missing explicit `pytest.mark.unit` decorators on test classes/functions as an issue (the hook will add them during collection). If a new test directory/category is introduced, ensure the hook is updated so it’s categorized correctly.

Applied to files:

tests/test_observability.py

📚 Learning: 2026-06-03T23:08:14.151Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: STYLE_GUIDE.md:0-0
Timestamp: 2026-06-03T23:08:14.151Z
Learning: Applies to **/*.py : Use `Protocol` for structural subtyping when you need duck-typing boundaries.

Applied to files:

src/nemo_safe_synthesizer/cli/wandb_setup.py

📚 Learning: 2026-06-03T23:08:14.151Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: STYLE_GUIDE.md:0-0
Timestamp: 2026-06-03T23:08:14.151Z
Learning: Applies to **/*.py : Use `except Exception:` + `logger.debug(..., exc_info=True)` for non-fatal cleanup at teardown boundaries. Bare `except Exception: pass` only in `__del__` methods.

Applied to files:

src/nemo_safe_synthesizer/generation/vllm_observability.py

📚 Learning: 2026-06-03T23:08:14.151Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: STYLE_GUIDE.md:0-0
Timestamp: 2026-06-03T23:08:14.151Z
Learning: Applies to **/*.py : Avoid defensive `try/except Exception` on trusted internal paths where exceptions shouldn't occur.

Applied to files:

src/nemo_safe_synthesizer/generation/vllm_observability.py

🔇 Additional comments (14)

src/nemo_safe_synthesizer/generation/vllm_observability.py (5)

1-91: LGTM!

98-239: LGTM!

246-349: LGTM!

352-379: LGTM!

425-488: LGTM!

src/nemo_safe_synthesizer/cli/wandb_setup.py (2)

14-14: LGTM!

268-301: LGTM!

src/nemo_safe_synthesizer/observability.py (2)

84-86: LGTM!

1012-1145: LGTM!

tests/test_observability.py (5)

1-13: LGTM!

15-21: LGTM!

24-43: LGTM!

53-77: LGTM!

46-50: Ensure NvmlPeakSampler.__enter__ handles the same failure mode your test injects for pynvml.
With monkeypatch.setitem(sys.modules, "pynvml", None), import pynvml will succeed and the first pynvml.* access (e.g., pynvml.nvmlInit()) will raise AttributeError; if NvmlPeakSampler.__enter__ only catches ImportError, the test won’t exercise the intended “no pynvml” fallback—either catch AttributeError in NvmlPeakSampler.__enter__ or mock the import to raise ImportError (e.g., remove the module from sys.modules so import pynvml fails).

- wandb_setup: use concrete "" default for WandbLoggable.to_wandb_payload prefix - vllm_observability: log read_vllm_runtime_metrics failure with exc_info=True for consistency with sibling degraded paths Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

Replace the dot-namespaced generation-observability event names ("vllm.generation.complete", "vllm.generation.observability.emit_failed") with plain human-readable log messages that match the repo's existing logger.runtime + extra={"ctx": ...} convention; the dotted form read like a module path. Give WandbLoggable.to_wandb_payload a docstring body instead of "..." so it is no longer flagged as a no-effect statement while staying idiomatic for a Protocol. Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

binaryaaron mentioned this pull request May 27, 2026

feat: vLLM benchmark harness #527

Draft

5 tasks

binaryaaron force-pushed the binaryaaron/vllm-observability-improvements branch from f7f177e to aa91b42 Compare June 3, 2026 21:44

github-code-quality Bot found potential problems Jun 3, 2026

View reviewed changes

Comment thread src/nemo_safe_synthesizer/generation/vllm_observability.py Fixed

Comment thread src/nemo_safe_synthesizer/generation/vllm_observability.py Fixed

binaryaaron marked this pull request as ready for review June 3, 2026 22:58

binaryaaron requested a review from a team as a code owner June 3, 2026 22:58

coderabbitai Bot added feature New feature or request test Test-only addition or change labels Jun 3, 2026

greptile-apps Bot reviewed Jun 3, 2026

View reviewed changes

Comment thread src/nemo_safe_synthesizer/generation/vllm_backend.py Outdated

Comment thread src/nemo_safe_synthesizer/generation/vllm_backend.py

Comment thread src/nemo_safe_synthesizer/generation/vllm_observability.py Outdated

coderabbitai Bot reviewed Jun 3, 2026

View reviewed changes

Comment thread src/nemo_safe_synthesizer/generation/vllm_observability.py

Comment thread tests/generation/test_vllm_backend.py

Comment thread tests/generation/test_vllm_backend.py

binaryaaron requested a review from mckornfield June 3, 2026 23:09

binaryaaron added area:generation area:observability labels Jun 3, 2026

mckornfield previously approved these changes Jun 4, 2026

View reviewed changes

binaryaaron dismissed mckornfield’s stale review via fd7f572 June 4, 2026 04:40

coderabbitai Bot removed feature New feature or request test Test-only addition or change labels Jun 4, 2026

mckornfield previously approved these changes Jun 4, 2026

View reviewed changes

binaryaaron dismissed mckornfield’s stale review via 5fc70a0 June 4, 2026 17:39

github-code-quality Bot found potential problems Jun 4, 2026

View reviewed changes

Comment thread src/nemo_safe_synthesizer/generation/vllm_observability.py Fixed

coderabbitai Bot added feature New feature or request test Test-only addition or change labels Jun 4, 2026

greptile-apps Bot reviewed Jun 4, 2026

View reviewed changes

Comment thread src/nemo_safe_synthesizer/generation/vllm_observability.py Outdated

mckornfield previously approved these changes Jun 4, 2026

View reviewed changes

coderabbitai Bot reviewed Jun 4, 2026

View reviewed changes

Comment thread src/nemo_safe_synthesizer/generation/vllm_backend.py

Comment thread tests/generation/test_vllm_observability.py

binaryaaron added 2 commits June 5, 2026 22:36

binaryaaron added 4 commits June 5, 2026 22:36

binaryaaron mentioned this pull request Jun 5, 2026

feat: add DP memory controls and physical batch sizing #518

Open

4 tasks

fix(observability): promote shared sampling and wandb sink

23c69c5

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

binaryaaron dismissed mckornfield’s stale review via 23c69c5 June 5, 2026 23:32

binaryaaron force-pushed the binaryaaron/vllm-observability-improvements branch from 5fc70a0 to 23c69c5 Compare June 5, 2026 23:32

github-actions Bot added area:sdk-cli area:tests labels Jun 5, 2026

github-code-quality Bot found potential problems Jun 5, 2026

View reviewed changes

Comment thread src/nemo_safe_synthesizer/generation/vllm_observability.py

Comment thread src/nemo_safe_synthesizer/cli/wandb_setup.py Fixed

coderabbitai Bot removed feature New feature or request test Test-only addition or change labels Jun 5, 2026

binaryaaron requested a review from mckornfield June 5, 2026 23:35

fix: formatting

a663d8a

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

coderabbitai Bot reviewed Jun 5, 2026

View reviewed changes

Comment thread src/nemo_safe_synthesizer/generation/vllm_observability.py

coderabbitai Bot added feature New feature or request test Test-only addition or change labels Jun 5, 2026

mckornfield previously approved these changes Jun 8, 2026

View reviewed changes

Comment thread src/nemo_safe_synthesizer/generation/vllm_backend.py Outdated

Comment thread src/nemo_safe_synthesizer/generation/vllm_backend.py

binaryaaron mentioned this pull request Jun 8, 2026

feat(generation): remote vLLM/NIM generation backend (GPU-free) #575

Draft

3 tasks

binaryaaron dismissed mckornfield’s stale review via 9c09160 June 8, 2026 22:54

github-code-quality Bot found potential problems Jun 8, 2026

View reviewed changes

Comment thread src/nemo_safe_synthesizer/cli/wandb_setup.py Fixed

binaryaaron requested a review from mckornfield June 9, 2026 04:08

mckornfield approved these changes Jun 9, 2026

View reviewed changes

binaryaaron added this pull request to the merge queue Jun 9, 2026

Merged via the queue into main with commit db28c48 Jun 9, 2026
19 checks passed

binaryaaron deleted the binaryaaron/vllm-observability-improvements branch June 9, 2026 18:08

Uh oh!

Conversation

binaryaaron commented May 27, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What this changes for production

Primitives + schema

Architectural choices worth flagging

Test coverage

Test plan

What's NOT in this PR

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented May 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reviews paused

Walkthrough

Changes

Estimated code review effort

Suggested labels

Suggested reviewers

❌ Failed checks (1 warning)

Uh oh!

codecov Bot commented May 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Uh oh!

greptile-apps Bot commented Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Sequence Diagram

Uh oh!

Uh oh!

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mckornfield left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mckornfield left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

AGENTS.md

Skills

Repo Conventions

Testing Guide

Read First

Running Tests

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

binaryaaron commented May 27, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented May 27, 2026 •

edited

Loading

codecov Bot commented May 27, 2026 •

edited

Loading

greptile-apps Bot commented Jun 3, 2026 •

edited

Loading