fix(llm): size VRAM headroom from reclaimable memory on integrated GPUs

[marcusds]

Problem

_get_vram_allocations sizes available VRAM from torch.cuda.mem_get_info().free. On integrated GPUs (such as that on a DGX Spark), that call counts only unallocated pages and ignores the reclaimable page cache. On such a box with most memory held in buff/cache, free reads only a couple GiB; after the 2 GiB safety buffer the usable fraction collapses to 0, and the gpu.vram preflight check hard-fails every job:

vram_exceeds_capacity: Estimated required VRAM ~8.4 GiB total ... exceeds
available ~0.0 GiB (training.max_vram_fraction=0.8). Training is expected to OOM.

…even though tens of GiB are actually reclaimable and usable. Observed on a GB10 with 121 GiB unified memory (is_integrated == 1), 73 GiB in reclaimable cache, mem_get_info().free ≈ 2 GiB → check reports ~0.0 GiB and fails.

Fix

On integrated GPUs, use the kernel's MemAvailable (from /proc/meminfo, which accounts for reclaimable memory), capped at device total, instead of mem_get_info's free figure.

• Discrete-GPU behavior is unchanged — the new path is gated on torch.cuda.get_device_properties(i).is_integrated.
• Falls back to the CUDA free value when /proc/meminfo can't be read (e.g. non-Linux hosts).

free, total = torch.cuda.mem_get_info(device=i)
if getattr(torch.cuda.get_device_properties(i), "is_integrated", False):
    available = _reclaimable_available_bytes()   # MemAvailable, or None
    if available is not None:
        free = min(max(free, available), total)
safe_free = max(free - (2 * 1024**3), 0)

Tests

tests/llm/test_utils.py:

• discrete GPU path unchanged (existing test, now also mocks get_device_properties);
• integrated GPU uses reclaimable MemAvailable (2 GiB free + 60 GiB available → 58 GiB usable);
• integrated fallback to CUDA free when MemAvailable is unavailable;
• MemAvailable parsing + absent-file handling.

Validated against the modified source with mocked CUDA (discrete {0: 0.5} unchanged; integrated {0: 0.483} from reclaimable memory). ruff check / ruff format --check clean.

Notes

• Commit is DCO signed-off and SSH-signed.
• Complements a separate ergonomic guard on the NeMo Platform plugin side; this is the root-cause fix for unified-memory hosts (GB10 / DGX Spark).

Summary by CodeRabbit

• Improvements
  • Enhanced VRAM headroom and allocation calculations for integrated GPUs on Linux by using reclaimable system memory (via MemAvailable) rather than relying only on CUDA “free”.
  • When applicable, reclaimable estimates are bounded by the active container/cgroup memory limit, supporting both cgroup v1 and v2 and treating “unlimited” limits appropriately.
• Bug Fixes
  • Improved integrated-GPU VRAM reservation behavior by replacing CUDA free headroom with capped reclaimable-based headroom while preserving safety buffer and max fraction logic.
• Tests
  • Expanded unit coverage for integrated vs. discrete GPU paths, MemAvailable parsing/fallback behavior, and cgroup limit edge cases (finite vs. unlimited, unreadable configurations).

[coderabbitai]

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

Changes

The VRAM helpers now account for Linux MemAvailable and finite cgroup memory allowances when estimating integrated-GPU headroom. Tests cover cgroup versions, parsing failures, device-total caps, discrete GPUs, and CUDA fallback behavior.

Integrated GPU memory handling

Layer / File(s)	Summary
Memory source helpers and validation src/nemo_safe_synthesizer/llm/utils.py, tests/llm/test_utils.py	Adds MemAvailable parsing and cgroup allowance helpers, including cgroup v1/v2 handling, unlimited-limit detection, and failure cases.
Integrated-GPU allocation adjustment src/nemo_safe_synthesizer/llm/utils.py, tests/llm/test_utils.py	Bounds reclaimable memory by cgroup allowance and device total before applying existing VRAM allocation calculations.
GPU path and fallback coverage tests/llm/test_utils.py	Covers non-integrated GPUs and fallback to CUDA-reported free memory when MemAvailable is unavailable.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Suggested labels: bug, test

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 33.33% 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
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.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: sizing VRAM headroom from reclaimable memory on integrated GPUs.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 1

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: da1a70ab-c4fa-4444-9e2d-63d380bff6d7

📥 Commits

Reviewing files that changed from the base of the PR and between d73679a and cce90c1.

📒 Files selected for processing (2)

• src/nemo_safe_synthesizer/llm/utils.py
• tests/llm/test_utils.py

📜 Review details

⏰ Context from checks skipped due to timeout. (8)

• GitHub Check: Unit Tests (3.13)
• GitHub Check: Unit Tests (3.12)
• GitHub Check: Unit Tests (3.11)
• GitHub Check: Smoke Tests
• GitHub Check: End-user Wheel Install
• GitHub Check: Greptile Review
• GitHub Check: Typecheck
• GitHub Check: Analyze (Python)

🧰 Additional context used

📓 Path-based instructions (11)

**/*.{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/llm/utils.py
• tests/llm/test_utils.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 in Python code, documentation, and messages.
Use Field(description=...) for every Pydantic model field.
Use assignment-style Field() by default; use Annotated only for additional metadata such as validators, constrained aliases, or discriminated unions.
Use @dataclass(frozen=True) for immutable value objects and validators; use mutable dataclasses only for builders, accumulators, and pipeline state.
Use field(default_factory=list) instead of mutable list defaults.
Use StrEnum for string-valued configuration or serialization enums and plain Enum for internal constants.
Obtain loggers with observability.get_logger(__name__); do not call logging.getLogger() or structlog.get_logger() directly.
Use .runtime, .user, and .system category loggers appropriately.
Do not use print() for operational library output; use the approved logger, click.echo(), or sys.stdout.write() where appropriate.
Use extra={} for machine-queryable logging data and f-strings only for human-readable context.
Raise errors from the custom Safe Synthesizer error hierarchy, using the documented dual inheritance for user and internal errors.
Keep shared package code compatible with Python 3.11; do not use Python 3.12-only syntax such as PEP 695 type statements or bracketed generic parameters.
Prefer X | Y, built-in collection generics, and Self over Optional, Union, and legacy typing collections.
Use collection ABCs for function arguments and concrete collection types for return values.
Use Protocol for structural subtyping and avoid Any when object, generics, or protocols are suitable.
Use TYPE_CHECKING guards for heavy imports such as pandas, torch, and transformers.
...

Files:

• src/nemo_safe_synthesizer/llm/utils.py
• tests/llm/test_utils.py

src/**/*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

src/**/*.py: Use relative imports for package code under src/.
Do not use assert for validation in library code; raise an appropriate exception instead.

Write Google-style docstrings for public Python APIs because API reference pages are generated from source docstrings.

Files:

• src/nemo_safe_synthesizer/llm/utils.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/llm/utils.py

**/*

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

**/*: Every source file requires the SPDX copyright and license header appropriate to its file format.
End files with a newline, remove trailing whitespace, use one space between sentences, and keep code, comments, and docstrings within 120 characters.

**/*: All contributions must use verified Git commits and DCO sign-off; unsigned or unsigned-off commits cannot be merged.
Branches other than main must follow <author>/<description>, optionally including an issue ID or type; branch names must use lowercase alphanumeric characters and hyphens.
Commits merged to main must follow Conventional Commits, using a valid lowercase type and a description of at most 100 characters.

Files:

• src/nemo_safe_synthesizer/llm/utils.py
• tests/llm/test_utils.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/llm/utils.py
• tests/llm/test_utils.py

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

📄 CodeRabbit inference engine (CONTRIBUTING.md)

All Python, shell, YAML, YML, and Markdown source files require SPDX copyright headers.

Files:

• src/nemo_safe_synthesizer/llm/utils.py
• tests/llm/test_utils.py

**/*.{py,pyi}

📄 CodeRabbit inference engine (CONTRIBUTING.md)

Shared Python package code must remain compatible with Python 3.11 syntax; do not use Python 3.12-only syntax such as PEP 695 type statements or bracketed generic class/function parameters.

Files:

• src/nemo_safe_synthesizer/llm/utils.py
• tests/llm/test_utils.py

**/*.{py,md}

📄 CodeRabbit inference engine (CONTRIBUTING.md)

Use the repository's documented Python and Markdown style conventions and validate changes with the pinned mise formatting and checking tasks.

Files:

• src/nemo_safe_synthesizer/llm/utils.py
• tests/llm/test_utils.py

**/*.{py,sh}

📄 CodeRabbit inference engine (CONTRIBUTING.md)

Use the repository's pinned mise tasks for formatting, linting, type checking, and testing rather than invoking ruff or ty directly for project-wide checks.

Files:

• src/nemo_safe_synthesizer/llm/utils.py
• tests/llm/test_utils.py

**/test_*.py

📄 CodeRabbit inference engine (AGENTS.md)

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

Files:

• tests/llm/test_utils.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

Mirror source code directory structure in tests directory (e.g., tests/training/, tests/generation/ parallel to source structure)

Files:

• tests/llm/test_utils.py

tests/**/*.py

📄 CodeRabbit inference engine (tests/TESTING.md)

tests/**/*.py: Auto-mark tests based on file path: tests under /e2e/ get e2e marker, tests under /smoke/ get smoke marker, all others get unit marker (only if no category marker already present)
Every test should have exactly one category marker: unit, smoke, or e2e
Use pytest.mark.requires_gpu modifier on tests that need CUDA hardware
Use pytest.mark.vllm on tests using vLLM generation backend and ensure each vLLM test file runs in its own process for GPU memory isolation
Use pytest.mark.slow on long-running tests
Use pytest.mark.smollm2 for SmolLM2 Hub download tests to enable process isolation
Use pytest.mark.noautouse to skip autouse fixtures for specific tests
Use load_test_dataset(filename) helper to load test datasets from tests/stub_datasets/ as HuggingFace Dataset objects
Use load_test_dataframe(filename) helper to load test data files from tests/stub_datasets/ as pandas DataFrames
Convert pandas columns to nullable dtypes (pd.Int64Dtype(), pd.BooleanDtype()) before assigning np.nan values
Use fake.seed_instance(seed) and random.seed(seed) together for Faker-based test data reproducibility
When sharing methods across multiple test files, define them in conftest.py and import them using relative imports (e.g., from .conftest import train_with_sdk); note that importing from other test files like tests/cli/helpers.py does not work
Use fixture_mock_processor or fixture_mock_processor_without_valid_records for mocking ParsedResponse objects with valid_records, invalid_records, errors, and prompt_number fields
Use pytest.importorskip to gate tests on optional dependencies that require specific extras (e.g., sentence_transformers, vllm)
Run vLLM tests with separate pytest invocations (one per file) using -n 0 (single process) for GPU memory isolation, or use staged mise tasks for CI visibility
Print statements are allowed in tests (ruff T201 is suppressed for tests/ directory) and should...

Files:

• tests/llm/test_utils.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/llm/test_utils.py

🧠 Learnings (1)

📚 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/llm/test_utils.py

🔇 Additional comments (2)

src/nemo_safe_synthesizer/llm/utils.py (1)

439-454: LGTM!

Also applies to: 457-488

tests/llm/test_utils.py (1)

7-13: LGTM!

Also applies to: 77-86, 105-127

[codecov]

Codecov Report

❌ Patch coverage is 97.36842% with 1 line in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/nemo_safe_synthesizer/llm/utils.py	91.66%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[greptile-apps]

Greptile Summary

This PR fixes a false-negative preflight VRAM check on integrated GPUs (e.g. DGX Spark / GB10) where torch.cuda.mem_get_info().free only counts unallocated pages and ignores the reclaimable page cache, causing the headroom calculation to collapse to ~0 GiB even with tens of GiB available.

• Adds _reclaimable_available_bytes() to read MemAvailable from /proc/meminfo (returns None on non-Linux) and applies it in _get_vram_allocations for is_integrated devices — hoisted outside the loop so the file is opened once per call, not once per GPU.
• Discrete-GPU behavior is entirely unchanged; integrated devices get free = min(max(cuda_free, MemAvailable), device_total) before the 2 GiB safety buffer is subtracted.
• Test coverage added for integrated happy path, capped-reclaimable edge case, meminfo-unreadable fallback, and MemAvailable parsing.

Confidence Score: 5/5

Safe to merge for bare-metal integrated-GPU hosts; one note about cgroup-limited container environments worth considering before wide platform rollout.

The fix is correct and well-scoped: discrete-GPU behavior is unchanged, the reclaimable value is read once before the loop, the min(max(free, reclaimable), total) expression is mathematically sound, and the tests cover the main paths including the cap and the fallback. The only open question is cgroup-limited containers, which is a deployment-time concern that doesn't affect bare-metal or full-memory-access container scenarios.

No files require special attention; the cgroup note in utils.py is worth revisiting if this path runs inside containers with per-container memory limits on the NeMo Platform.

Important Files Changed

Filename	Overview
src/nemo_safe_synthesizer/llm/utils.py	Adds _reclaimable_available_bytes() helper and applies MemAvailable-based free estimate for integrated GPUs in _get_vram_allocations; reclaimable is read once before the loop, and the result is capped at device total before the safety buffer subtraction.
tests/llm/test_utils.py	Adds five new test cases covering discrete (patched is_integrated), integrated reclaimable, capped-reclaimable, meminfo-fallback, and MemAvailable parsing/absent-file paths; all assertions are mathematically consistent with the implementation.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[_get_vram_allocations called] --> B[torch.cuda.is_available?]
    B -- No --> Z[Return empty dict]
    B -- Yes --> C[_reclaimable_available_bytes
Read /proc/meminfo once]
    C --> D{MemAvailable
parsed OK?}
    D -- Yes --> E[reclaimable = MemAvailable bytes]
    D -- No/OSError/ValueError --> F[reclaimable = None]
    E --> G[For each GPU i]
    F --> G
    G --> H[mem_get_info device i: free, total]
    H --> I{is_integrated AND
reclaimable not None?}
    I -- Yes --> J[free = min max free reclaimable total]
    I -- No --> K[Keep CUDA free as-is]
    J --> L[safe_free = max free minus 2GiB 0]
    K --> L
    L --> M[utilization = min max_vram_fraction safe_free/total]
    M --> N[memory_bytes = utilization x total]
    N --> O[Append to allocations]
    O --> P{More GPUs?}
    P -- Yes --> G
    P -- No --> Q[Return allocations dict]

Loading

_{Reviews (6): Last reviewed commit: "fix(llm): hoist MemAvailable read and te..." | Re-trigger Greptile}

[coderabbitai]

Actionable comments posted: 1

Caution

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

⚠️ Outside diff range comments (1)

src/nemo_safe_synthesizer/llm/utils.py (1)

573-574: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Apply the cgroup-bounded value rather than taking the larger value.

When CUDA reports 16 GiB free but cgroup-bounded reclaimable memory is 1 GiB, max() retains 16 GiB and bypasses the new container limit. This can exceed the container allowance on integrated GPUs. Replace this with free = min(reclaimable, total); the existing reclaimable is None branch already preserves CUDA fallback.

Proposed fix

-                free = min(max(free, reclaimable), total)
+                free = min(reclaimable, total)

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 91e78d2b-1690-47d8-9655-c869128dcf02

📥 Commits

Reviewing files that changed from the base of the PR and between c69bbe8 and 98bd6d5.

📒 Files selected for processing (2)

• src/nemo_safe_synthesizer/llm/utils.py
• tests/llm/test_utils.py

📜 Review details

⏰ Context from checks skipped due to timeout. (8)

• GitHub Check: Unit Tests (3.12)
• GitHub Check: Unit Tests (3.13)
• GitHub Check: Unit Tests (3.11)
• GitHub Check: Smoke Tests
• GitHub Check: End-user Wheel Install
• GitHub Check: Greptile Review
• GitHub Check: Typecheck
• GitHub Check: Analyze (Python)

🧰 Additional context used

📓 Path-based instructions (11)

**/*.{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:

• tests/llm/test_utils.py
• src/nemo_safe_synthesizer/llm/utils.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 in Python code, documentation, and messages.
Use Field(description=...) for every Pydantic model field.
Use assignment-style Field() by default; use Annotated only for additional metadata such as validators, constrained aliases, or discriminated unions.
Use @dataclass(frozen=True) for immutable value objects and validators; use mutable dataclasses only for builders, accumulators, and pipeline state.
Use field(default_factory=list) instead of mutable list defaults.
Use StrEnum for string-valued configuration or serialization enums and plain Enum for internal constants.
Obtain loggers with observability.get_logger(__name__); do not call logging.getLogger() or structlog.get_logger() directly.
Use .runtime, .user, and .system category loggers appropriately.
Do not use print() for operational library output; use the approved logger, click.echo(), or sys.stdout.write() where appropriate.
Use extra={} for machine-queryable logging data and f-strings only for human-readable context.
Raise errors from the custom Safe Synthesizer error hierarchy, using the documented dual inheritance for user and internal errors.
Keep shared package code compatible with Python 3.11; do not use Python 3.12-only syntax such as PEP 695 type statements or bracketed generic parameters.
Prefer X | Y, built-in collection generics, and Self over Optional, Union, and legacy typing collections.
Use collection ABCs for function arguments and concrete collection types for return values.
Use Protocol for structural subtyping and avoid Any when object, generics, or protocols are suitable.
Use TYPE_CHECKING guards for heavy imports such as pandas, torch, and transformers.
...

Files:

• tests/llm/test_utils.py
• src/nemo_safe_synthesizer/llm/utils.py

**/test_*.py

📄 CodeRabbit inference engine (AGENTS.md)

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

Files:

• tests/llm/test_utils.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

Mirror source code directory structure in tests directory (e.g., tests/training/, tests/generation/ parallel to source structure)

Files:

• tests/llm/test_utils.py

tests/**/*.py

📄 CodeRabbit inference engine (tests/TESTING.md)

tests/**/*.py: Auto-mark tests based on file path: tests under /e2e/ get e2e marker, tests under /smoke/ get smoke marker, all others get unit marker (only if no category marker already present)
Every test should have exactly one category marker: unit, smoke, or e2e
Use pytest.mark.requires_gpu modifier on tests that need CUDA hardware
Use pytest.mark.vllm on tests using vLLM generation backend and ensure each vLLM test file runs in its own process for GPU memory isolation
Use pytest.mark.slow on long-running tests
Use pytest.mark.smollm2 for SmolLM2 Hub download tests to enable process isolation
Use pytest.mark.noautouse to skip autouse fixtures for specific tests
Use load_test_dataset(filename) helper to load test datasets from tests/stub_datasets/ as HuggingFace Dataset objects
Use load_test_dataframe(filename) helper to load test data files from tests/stub_datasets/ as pandas DataFrames
Convert pandas columns to nullable dtypes (pd.Int64Dtype(), pd.BooleanDtype()) before assigning np.nan values
Use fake.seed_instance(seed) and random.seed(seed) together for Faker-based test data reproducibility
When sharing methods across multiple test files, define them in conftest.py and import them using relative imports (e.g., from .conftest import train_with_sdk); note that importing from other test files like tests/cli/helpers.py does not work
Use fixture_mock_processor or fixture_mock_processor_without_valid_records for mocking ParsedResponse objects with valid_records, invalid_records, errors, and prompt_number fields
Use pytest.importorskip to gate tests on optional dependencies that require specific extras (e.g., sentence_transformers, vllm)
Run vLLM tests with separate pytest invocations (one per file) using -n 0 (single process) for GPU memory isolation, or use staged mise tasks for CI visibility
Print statements are allowed in tests (ruff T201 is suppressed for tests/ directory) and should...

Files:

• tests/llm/test_utils.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/llm/test_utils.py

**/*

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

**/*: Every source file requires the SPDX copyright and license header appropriate to its file format.
End files with a newline, remove trailing whitespace, use one space between sentences, and keep code, comments, and docstrings within 120 characters.

**/*: All contributions must use verified Git commits and DCO sign-off; unsigned or unsigned-off commits cannot be merged.
Branches other than main must follow <author>/<description>, optionally including an issue ID or type; branch names must use lowercase alphanumeric characters and hyphens.
Commits merged to main must follow Conventional Commits, using a valid lowercase type and a description of at most 100 characters.

Files:

• tests/llm/test_utils.py
• src/nemo_safe_synthesizer/llm/utils.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:

• tests/llm/test_utils.py
• src/nemo_safe_synthesizer/llm/utils.py

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

📄 CodeRabbit inference engine (CONTRIBUTING.md)

All Python, shell, YAML, YML, and Markdown source files require SPDX copyright headers.

Files:

• tests/llm/test_utils.py
• src/nemo_safe_synthesizer/llm/utils.py

**/*.{py,pyi}

📄 CodeRabbit inference engine (CONTRIBUTING.md)

Shared Python package code must remain compatible with Python 3.11 syntax; do not use Python 3.12-only syntax such as PEP 695 type statements or bracketed generic class/function parameters.

Files:

• tests/llm/test_utils.py
• src/nemo_safe_synthesizer/llm/utils.py

**/*.{py,md}

📄 CodeRabbit inference engine (CONTRIBUTING.md)

Use the repository's documented Python and Markdown style conventions and validate changes with the pinned mise formatting and checking tasks.

Files:

• tests/llm/test_utils.py
• src/nemo_safe_synthesizer/llm/utils.py

**/*.{py,sh}

📄 CodeRabbit inference engine (CONTRIBUTING.md)

Use the repository's pinned mise tasks for formatting, linting, type checking, and testing rather than invoking ruff or ty directly for project-wide checks.

Files:

• tests/llm/test_utils.py
• src/nemo_safe_synthesizer/llm/utils.py

src/**/*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

src/**/*.py: Use relative imports for package code under src/.
Do not use assert for validation in library code; raise an appropriate exception instead.

Write Google-style docstrings for public Python APIs because API reference pages are generated from source docstrings.

Files:

• src/nemo_safe_synthesizer/llm/utils.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/llm/utils.py

🧠 Learnings (1)

📚 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/llm/test_utils.py

🪛 ast-grep (0.44.1)

src/nemo_safe_synthesizer/llm/utils.py

[warning] 464-464: File path is request-/variable-derived; validate and normalize to prevent path traversal.
Context: open(path, encoding="utf-8")
Note: [CWE-22] Improper Limitation of a Pathname to a Restricted Directory ('Path Traversal').

(open-filename-from-request)

🔇 Additional comments (3)

tests/llm/test_utils.py (3)

13-13: LGTM!

---

100-100: LGTM!

Also applies to: 117-138

---

165-225: LGTM!

[coderabbitai]

♻️ Duplicate comments (1)

src/nemo_safe_synthesizer/llm/utils.py (1)

471-497: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Unreadable/unparsable cgroup usage silently treated as zero, overestimating headroom.

At lines 490-494, when raw_usage is None (unreadable) or fails int() parsing, usage defaults to 0, so headroom = max(limit - usage, 0) becomes the full limit. In a memory-limited container where the usage file is transiently unreadable, this reports the entire limit as "remaining," which _get_vram_allocations then uses to size an allocation near the full limit — risking OOM. This is the same issue flagged in a prior review of this function ("Do not treat unreadable cgroup usage as zero ... Propagate an unavailable cgroup-data state").

Skip the depth (mirroring the existing raw_limit is None or raw_limit == "max" pattern) instead of assuming zero usage, so an unknown usage doesn't silently produce an overly generous headroom estimate.

🛡️ Proposed fix

         raw_usage = _read_cgroup_file(f"{directory}/{usage_file}")
-        try:
-            usage = int(raw_usage) if raw_usage is not None else 0
-        except ValueError:
-            usage = 0
+        if raw_usage is None:
+            continue
+        try:
+            usage = int(raw_usage)
+        except ValueError:
+            continue
         headroom = max(limit - usage, 0)
         remaining = headroom if remaining is None else min(remaining, headroom)

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c6d74e54-23b3-441b-b35e-756e1acf880a

📥 Commits

Reviewing files that changed from the base of the PR and between 98bd6d5 and 0448df1.

📒 Files selected for processing (2)

• src/nemo_safe_synthesizer/llm/utils.py
• tests/llm/test_utils.py

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

• tests/llm/test_utils.py

📜 Review details

⏰ Context from checks skipped due to timeout. (2)

• GitHub Check: Greptile Review
• GitHub Check: Analyze (Python)

🧰 Additional context used

📓 Path-based instructions (8)

**/*.{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/llm/utils.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 in Python code, documentation, and messages.
Use Field(description=...) for every Pydantic model field.
Use assignment-style Field() by default; use Annotated only for additional metadata such as validators, constrained aliases, or discriminated unions.
Use @dataclass(frozen=True) for immutable value objects and validators; use mutable dataclasses only for builders, accumulators, and pipeline state.
Use field(default_factory=list) instead of mutable list defaults.
Use StrEnum for string-valued configuration or serialization enums and plain Enum for internal constants.
Obtain loggers with observability.get_logger(__name__); do not call logging.getLogger() or structlog.get_logger() directly.
Use .runtime, .user, and .system category loggers appropriately.
Do not use print() for operational library output; use the approved logger, click.echo(), or sys.stdout.write() where appropriate.
Use extra={} for machine-queryable logging data and f-strings only for human-readable context.
Raise errors from the custom Safe Synthesizer error hierarchy, using the documented dual inheritance for user and internal errors.
Keep shared package code compatible with Python 3.11; do not use Python 3.12-only syntax such as PEP 695 type statements or bracketed generic parameters.
Prefer X | Y, built-in collection generics, and Self over Optional, Union, and legacy typing collections.
Use collection ABCs for function arguments and concrete collection types for return values.
Use Protocol for structural subtyping and avoid Any when object, generics, or protocols are suitable.
Use TYPE_CHECKING guards for heavy imports such as pandas, torch, and transformers.
...

Files:

• src/nemo_safe_synthesizer/llm/utils.py

src/**/*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

src/**/*.py: Use relative imports for package code under src/.
Do not use assert for validation in library code; raise an appropriate exception instead.

Write Google-style docstrings for public Python APIs because API reference pages are generated from source docstrings.

Files:

• src/nemo_safe_synthesizer/llm/utils.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/llm/utils.py

**/*

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

**/*: Every source file requires the SPDX copyright and license header appropriate to its file format.
End files with a newline, remove trailing whitespace, use one space between sentences, and keep code, comments, and docstrings within 120 characters.

**/*: All contributions must use verified Git commits and DCO sign-off; unsigned or unsigned-off commits cannot be merged.
Branches other than main must follow <author>/<description>, optionally including an issue ID or type; branch names must use lowercase alphanumeric characters and hyphens.
Commits merged to main must follow Conventional Commits, using a valid lowercase type and a description of at most 100 characters.

Files:

• src/nemo_safe_synthesizer/llm/utils.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/llm/utils.py

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

📄 CodeRabbit inference engine (CONTRIBUTING.md)

All Python, shell, YAML, YML, and Markdown source files require SPDX copyright headers.

Files:

• src/nemo_safe_synthesizer/llm/utils.py

**/*.{py,pyi}

📄 CodeRabbit inference engine (CONTRIBUTING.md)

Shared Python package code must remain compatible with Python 3.11 syntax; do not use Python 3.12-only syntax such as PEP 695 type statements or bracketed generic class/function parameters.

Files:

• src/nemo_safe_synthesizer/llm/utils.py

**/*.{py,md}

📄 CodeRabbit inference engine (CONTRIBUTING.md)

Use the repository's documented Python and Markdown style conventions and validate changes with the pinned mise formatting and checking tasks.

Files:

• src/nemo_safe_synthesizer/llm/utils.py

**/*.{py,sh}

📄 CodeRabbit inference engine (CONTRIBUTING.md)

Use the repository's pinned mise tasks for formatting, linting, type checking, and testing rather than invoking ruff or ty directly for project-wide checks.

Files:

• src/nemo_safe_synthesizer/llm/utils.py

🪛 ast-grep (0.44.1)

src/nemo_safe_synthesizer/llm/utils.py

[warning] 464-464: File path is request-/variable-derived; validate and normalize to prevent path traversal.
Context: open(path, encoding="utf-8")
Note: [CWE-22] Improper Limitation of a Pathname to a Restricted Directory ('Path Traversal').

(open-filename-from-request)

🔇 Additional comments (1)

src/nemo_safe_synthesizer/llm/utils.py (1)

439-454: LGTM!

Also applies to: 457-468, 500-534, 537-585

[binaryaaron]

Thanks! makes sense, will run a few things on our side and get back with proper review. we do run within kube quite a bit and i'll validate a few things against that env (or if you have availability to do so, feel free).

[mckornfield]

lol interesting

[mckornfield]

could do
from contextlib import suppress

with suppress(Exception):
  ...
return None