feat: account for batch size in preflight - #480
Conversation
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
binaryaaron
changed the title
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughPer-device GPU allocation utilities and a Hugging Face max_memory map were added; VRAM demand is estimated by components (base weights, optional activations, overhead) and compared to device headroom in the preflight check with a hard-fail threshold. Backend loading, Click Literal handling, tests, and docs were updated accordingly. ChangesGPU Memory Management & VRAM Validation Improvements
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
coderabbitai
Bot
added
docs
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/nemo_safe_synthesizer/llm/utils.py (1)
436-466:⚠️ Potential issue | 🟠 Major | ⚡ Quick winClamp
max_vram_fractionand log the applied utilization.
prepare_config(max_vram_fraction=...)now allows direct overrides, so values outside[0, 1]can yield negative or over-capacitymemory_byteshere. Also, when the 2 GiB safety buffer becomes the limiting factor, the log still reports the configured fraction instead of the fraction you actually apply.Suggested fix
def _get_vram_allocations(max_vram_fraction: float | None = None) -> dict[int, _VRAMAllocation]: @@ import torch if max_vram_fraction is None: max_vram_fraction = 0.8 + max_vram_fraction = min(max(max_vram_fraction, 0.0), 1.0) allocations = {} @@ safe_free = max(free - (2 * 1024**3), 0) gpu_memory_utilization = min(max_vram_fraction, safe_free / total) if total > 0 else 0.0 memory_bytes = int(gpu_memory_utilization * total) memory_gib = memory_bytes / (1024**3) allocations[i] = _VRAMAllocation(utilization=gpu_memory_utilization, memory_bytes=memory_bytes) logger.info( - f"GPU {i}: Will allocate {memory_gib:.2f}GiB ({max_vram_fraction * 100}% of {total / (1024**3):.2f}GiB)" + f"GPU {i}: Will allocate {memory_gib:.2f}GiB " + f"({gpu_memory_utilization * 100:.1f}% of {total / (1024**3):.2f}GiB)" )As per coding guidelines, "Prefer clamping/saturation over raising when out-of-range inputs shouldn't crash the system. Return a bounded value with a log warning instead of raising."
tests/training/test_huggingface_backend.py (1)
707-750:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winMark these tests with a category.
This class is unmarked, so the new
prepare_configcoverage doesn't meet the repo's exactly-one test-category rule and can be missed by marker-based selection. Add@pytest.mark.uniton the class or on each test method.As per coding guidelines, "Assign exactly one category marker (
unit,smoke, ore2e) to every test function."
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f37eff63-cf70-4b24-94aa-435af178414a
📒 Files selected for processing (10)
docs/user-guide/running.mddocs/user-guide/troubleshooting.mdsrc/nemo_safe_synthesizer/configurator/pydantic_click_options.pysrc/nemo_safe_synthesizer/llm/utils.pysrc/nemo_safe_synthesizer/preflight/checks/environment.pysrc/nemo_safe_synthesizer/training/huggingface_backend.pytests/configurator/test_pydantic_click_options.pytests/llm/test_utils.pytests/preflight/test_preflight.pytests/training/test_huggingface_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.11)
- GitHub Check: Unit Tests (3.12)
- GitHub Check: Unit Tests (3.13)
- GitHub Check: Smoke Tests
🧰 Additional context used
📓 Path-based instructions (15)
**/*.md
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
**/*.md: Do not use decorative**bold**in body text, list items, or docstrings. Use headers, list markers, colons, and backticks for structure. Bold is acceptable in table header-like cells and MkDocs Material card grid titles.
Use--(em-dash) for asides, not-(hyphen).
Use single backticks for code identifiers, paths, and CLI commands in Markdown.
UseMermaiddiagrams with no spaces in node IDs, quoted labels with special characters, and no explicit colors or styles.
Include SPDX copyright header at the top using HTML-comment syntax (<!-- ... -->).
If a Markdown file starts with YAML frontmatter (---), place the SPDX copyright header inside the frontmatter block using hash-comment syntax, not HTML comments.
Limit line length to 120 characters in Markdown files.
Files:
docs/user-guide/troubleshooting.mddocs/user-guide/running.md
docs/**/*.md
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
Classify documentation pages as tutorial, how-to, explanation, or reference per the Diataxis framework. Use MkDocs Material syntax: admonitions (
!!! note), tabs (===), code blocks with titles and highlights.
docs/**/*.md: Documentation must be written in Markdown underdocs/directory and organized following the Diataxis framework with subdirectories: getting-started/, user-guide/, architecture/, reference/, blog/
Use MkDocs Material Markdown extensions in documentation: admonitions (!!! note,!!! warning,??? tip), content tabs (=== "label"), code blocks with syntax highlighting, Mermaid diagrams, task lists, footnotes, definition lists
docs/**/*.md: Classify documentation content using the Diataxis framework (TUTORIAL, HOW-TO, EXPLANATION, or REFERENCE) before writing, ensuring each page fits exactly ONE type
Use cross-links between different Diataxis content types (TUTORIAL, HOW-TO, EXPLANATION, REFERENCE) to connect related documentation
Use MkDocs Material admonitions syntax for notes, warnings, and collapsible tips:!!! note,!!! warning,??? tip
Use MkDocs Material tabs syntax (=== "Label") for presenting multiple examples or implementations side-by-side
Use code blocks with metadata (title, hl_lines) to highlight relevant code snippets in documentation examples
Use Mermaid diagrams (flowchart, sequence diagrams, etc.) for visualizing architecture, workflows, and concepts in documentation
Write documentation following high signal-to-noise principles: every sentence must earn its place by providing essential information
Use progressive disclosure in documentation: start with simple concepts, then layer complexity for advanced readers
Include working code examples in documentation; ensure all code snippets are tested and actually work
List all prerequisites at the top of documentation pages before diving into main content
End documentation pages with 'Next steps' section containing links to related content and logical progression poi...
Files:
docs/user-guide/troubleshooting.mddocs/user-guide/running.md
**/*
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
**/*: Remove all trailing whitespace from files.
Use single space between sentences, never two spaces.
Files:
docs/user-guide/troubleshooting.mdsrc/nemo_safe_synthesizer/training/huggingface_backend.pytests/configurator/test_pydantic_click_options.pytests/llm/test_utils.pydocs/user-guide/running.mdsrc/nemo_safe_synthesizer/configurator/pydantic_click_options.pysrc/nemo_safe_synthesizer/llm/utils.pysrc/nemo_safe_synthesizer/preflight/checks/environment.pytests/preflight/test_preflight.pytests/training/test_huggingface_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:
docs/user-guide/troubleshooting.mdsrc/nemo_safe_synthesizer/training/huggingface_backend.pytests/configurator/test_pydantic_click_options.pytests/llm/test_utils.pydocs/user-guide/running.mdsrc/nemo_safe_synthesizer/configurator/pydantic_click_options.pysrc/nemo_safe_synthesizer/llm/utils.pysrc/nemo_safe_synthesizer/preflight/checks/environment.pytests/preflight/test_preflight.pytests/training/test_huggingface_backend.py
**/*.{py,sh,yaml,yml,md}
📄 CodeRabbit inference engine (CONTRIBUTING.md)
All source files (
.py,.sh,.yaml,.yml,.md) require SPDX copyright headers, automatically added bymake format
Files:
docs/user-guide/troubleshooting.mdsrc/nemo_safe_synthesizer/training/huggingface_backend.pytests/configurator/test_pydantic_click_options.pytests/llm/test_utils.pydocs/user-guide/running.mdsrc/nemo_safe_synthesizer/configurator/pydantic_click_options.pysrc/nemo_safe_synthesizer/llm/utils.pysrc/nemo_safe_synthesizer/preflight/checks/environment.pytests/preflight/test_preflight.pytests/training/test_huggingface_backend.py
**/*.{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:
docs/user-guide/troubleshooting.mdsrc/nemo_safe_synthesizer/training/huggingface_backend.pytests/configurator/test_pydantic_click_options.pytests/llm/test_utils.pydocs/user-guide/running.mdsrc/nemo_safe_synthesizer/configurator/pydantic_click_options.pysrc/nemo_safe_synthesizer/llm/utils.pysrc/nemo_safe_synthesizer/preflight/checks/environment.pytests/preflight/test_preflight.pytests/training/test_huggingface_backend.py
**/*.{md,markdown}
📄 CodeRabbit inference engine (.cursor/rules/agent-markdown-style.mdc)
**/*.{md,markdown}: Bold is acceptable only in markdown tables where it's the conventional way to mark header-like cells in the body
Use##headers to segment markdown sections instead of bold text
Use--(em-dash) instead of-(hyphen) for asides in markdown
Files:
docs/user-guide/troubleshooting.mddocs/user-guide/running.md
docs/**
⚙️ CodeRabbit configuration file
Review documentation as MkDocs Material content. Check Diataxis fit, accurate commands, internal links, code fences, and markdown style from STYLE_GUIDE.md.
Files:
docs/user-guide/troubleshooting.mddocs/user-guide/running.md
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
**/*.py: Use Python 3.11–3.13 with modern syntax (X | Y, list[str], Self). Python 3.14+ is not supported
Durable implementation guidance belongs with the code it describes: function and class docstrings for public contracts and source comments for local invariants
**/*.py: Use American English spelling: 'initialize' not 'initialise', 'recognize' not 'recognise', 'color' not 'colour'.
Useobservability.get_logger(__name__)for logging -- neverlogging.getLogger()orstructlog.get_logger()directly.
Category loggers: use.runtimefor internals,.userfor progress/results,.systemfor system events.
Never useprint()for operational output in library code. Useclick.echo()for CLI output orsys.stdout.write()for raw output in tools.
Useextra={}for data that downstream tools should query or aggregate (metrics, counts, durations). f-strings are fine for human-readable context.
UseX | Yinstead ofOptional[X]orUnion[X, Y]. Uselist[str]instead ofList[str],dict[str, int]instead ofDict[str, int].
UseSelffor fluent method returns that return the same instance type.
Use collection ABCs for function arguments (Sequence,Mapping,Iterable) so callers can pass any compatible container; use concrete types for return values so callers know exactly what they get.
UseProtocolfor structural subtyping when you need duck-typing boundaries. AvoidAny-- preferobject, generics, orProtocol.
Prefermatch/casefor dispatch on types or tagged values. Useif/eliffor simple boolean predicates.
Prefer comprehensions over imperative loops where intent is clearer. Avoid multipleforclauses -- optimize for readability, not conciseness.
Prefer clamping/saturation over raising when out-of-range inputs shouldn't crash the system. Return a bounded value with a log warning instead of raising.
If a function has more than two levels of indentation beyonddef, it needs refactoring. Use guard clauses, extract inner loops,...
Files:
src/nemo_safe_synthesizer/training/huggingface_backend.pytests/configurator/test_pydantic_click_options.pytests/llm/test_utils.pysrc/nemo_safe_synthesizer/configurator/pydantic_click_options.pysrc/nemo_safe_synthesizer/llm/utils.pysrc/nemo_safe_synthesizer/preflight/checks/environment.pytests/preflight/test_preflight.pytests/training/test_huggingface_backend.py
src/**/*.py
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
src/**/*.py: PreferNSSBaseModelfor config/parameter models inconfig/which define user-facing configuration. Use rawBaseModelor module-specific bases for data transfer objects and internal structures.
UseBaseSettingsfor env/CLI settings. PreferAliasChoiceson individual fields when a field should respond to both its Python name and an env var name (e.g.,validation_alias=AliasChoices('config_path', 'NSS_CONFIG')).
IncludeField(description=...)on all Pydantic model fields. This is the canonical field docstring extracted by griffe-pydantic for the API reference and used by the configurator for CLI help text.
Use assignment-styleField()as the default for model fields because type checkers understanddefault,default_factory, andaliasin this style.
UseAnnotatedonly when the field carries additional metadata beyondField()-- such asValueValidator,AutoParam,DependsOnValidator, reusable constrained type aliases, nested-type constraints, or discriminated unions.
For Pydantic fields withAnnotatedtypes that have defaults, put the default as a bare assignment (= value), not insideField(default=...). Exception:default_factoryhas no bare-assignment equivalent, so use assignment-styleField(default_factory=...).
Include valid ranges in thedescriptionwhen constraints won't appear in the rendered API docs (e.g., 'Must be in (0, 1).').
Use@dataclass(frozen=True)for immutable value objects and validators. Mutable@dataclassis acceptable for builders, accumulators, and pipeline state.
Usefield(default_factory=list)for mutable defaults in dataclasses, never= [].
UseStrEnumfor string-valued enums used in configs/serialization. Use plainEnumfor internal-only named constants.
UseTYPE_CHECKINGguards for heavy imports (pandas,torch,transformers); not needed for stdlib or lightweight imports.
Addfrom __future__ import annotationsto every module to make all annotations strings consis...
Files:
src/nemo_safe_synthesizer/training/huggingface_backend.pysrc/nemo_safe_synthesizer/configurator/pydantic_click_options.pysrc/nemo_safe_synthesizer/llm/utils.pysrc/nemo_safe_synthesizer/preflight/checks/environment.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/training/huggingface_backend.pysrc/nemo_safe_synthesizer/configurator/pydantic_click_options.pysrc/nemo_safe_synthesizer/llm/utils.pysrc/nemo_safe_synthesizer/preflight/checks/environment.py
src/nemo_safe_synthesizer/training/**/*.py
⚙️ CodeRabbit configuration file
Review training changes for dataset preprocessing, model path handling, artifact writes, LoRA/DP behavior, GPU memory usage, reproducibility, and cleanup on failure.
Files:
src/nemo_safe_synthesizer/training/huggingface_backend.py
**/test_*.py
📄 CodeRabbit inference engine (AGENTS.md)
**/test_*.py: In pytest.ini configuration, asyncio_mode is set to auto so async tests work without@pytest.mark.asyncio
Use the unit test marker instead of the deprecated unit_test marker
Files:
tests/configurator/test_pydantic_click_options.pytests/llm/test_utils.pytests/preflight/test_preflight.pytests/training/test_huggingface_backend.py
tests/**/*.py
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
tests/**/*.py: Use absolute imports intests/(e.g.,from nemo_safe_synthesizer.observability import get_logger).
Name test files astest_*.py, test classes asTest*, test functions astest_<module>_<expected_behavior>.
Usefixture_prefix convention for pytest fixtures 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 -- not based on assumptions about cost.
Use bareassertas the primary assertion style in tests. Usepytest.raises()withmatch=for exceptions;pytest.approx()for floating-point comparisons.
Docstrings are optional for simple tests, recommended for complex/e2e tests explaining purpose.
Markers are auto-assigned by path viapytest_collection_modifyitems(/e2e/->e2e,/smoke/->smoke, default ->unit). Explicit markers:@pytest.mark.slow,@pytest.mark.requires_gpu,@pytest.mark.timeout().
Usetmp_pathfixture for file operations in tests. 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.
Maintain test isolation: no shared mutable state or execution-order dependencies between tests. If something must run first, include it in the test or a fixture.
Use@pytest.mark.parametrizefor testing multiple input combinations rather than copy-pasting similar tests.Tests in
tests/e2e/should be auto-marked withe2emarker, tests intests/smoke/withsmokemarker, others withunitmarker
Files:
tests/configurator/test_pydantic_click_options.pytests/llm/test_utils.pytests/preflight/test_preflight.pytests/training/test_huggingface_backend.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/configurator/test_pydantic_click_options.pytests/llm/test_utils.pytests/preflight/test_preflight.pytests/training/test_huggingface_backend.py
tests/**/test_*.py
📄 CodeRabbit inference engine (tests/TESTING.md)
tests/**/test_*.py: Assign exactly one category marker (unit,smoke, ore2e) to every test function
Gate optional dependencies usingpytest.importorskipfor tests requiring packages that need specific extras (e.g.,sentence_transformers,vllmforcu128extra)
Usepytest.mark.vllmon test files that call vLLM.generate()method and run them with process isolation (-n 0) to prevent GPU memory exhaustion across test processes
Use relative imports from conftest.py when sharing methods across multiple test files (e.g.,from .conftest import train_with_sdk); direct imports from other test files undertests/do not work
Create mockParsedResponseobjects withvalid_records=[...],invalid_records=[...],errors=[...], andprompt_number=intfields; usefixture_mock_processororfixture_mock_processor_without_valid_recordsfixtures
Useprint()statements freely in test functions for debug output (ruffT201is suppressed fortests/directory)
Files:
tests/configurator/test_pydantic_click_options.pytests/llm/test_utils.pytests/preflight/test_preflight.pytests/training/test_huggingface_backend.py
tests/**
📄 CodeRabbit inference engine (.cursor/rules/repo-navigation.mdc)
Tests should mirror the
src/directory structure intests/
Files:
tests/configurator/test_pydantic_click_options.pytests/llm/test_utils.pytests/preflight/test_preflight.pytests/training/test_huggingface_backend.py
src/nemo_safe_synthesizer/configurator/**/*.py
⚙️ CodeRabbit configuration file
Review Pydantic-to-Click mapping carefully. Check option names, type conversion, nullable sub-config behavior, validation errors, help text, and compatibility with parse_overrides().
Files:
src/nemo_safe_synthesizer/configurator/pydantic_click_options.py
🔇 Additional comments (1)
src/nemo_safe_synthesizer/configurator/pydantic_click_options.py (1)
200-208: Good fix for numericLiteralClick type inference.This normalization cleanly fixes
Literal[int]handling while keeping string-sentinel routing intact.Also applies to: 242-242
| if training_cfg.quantize_model: | ||
| return training_cfg.quantization_bits / 8 + 0.1 |
There was a problem hiding this comment.
Mirror the runtime 8-bit fallback in the preflight estimate.
When training.quantize_model is True and quantization_bits is unset, the training path falls back to 8-bit, but this branch now does None / 8. That turns a valid runtime config into a gpu.vram check crash or bad estimate.
Suggested fix
if training_cfg.quantize_model:
- return training_cfg.quantization_bits / 8 + 0.1
+ quantization_bits = training_cfg.quantization_bits or 8
+ return quantization_bits / 8 + 0.1
return 2.0| def test_click_type_int_literal_returns_int(): | ||
| """Literal[4, 8] must parse CLI values as ints before Pydantic validation.""" | ||
| from typing import Literal | ||
|
|
||
| assert _click_type(Literal[4, 8]) == click.INT | ||
|
|
||
|
|
||
| def test_click_type_string_literal_plus_int_literal_returns_string(): | ||
| """String literals still need STRING so Pydantic can validate sentinel values.""" | ||
| from typing import Literal | ||
|
|
||
| assert _click_type(Literal["disabled", 4]) == click.STRING | ||
|
|
There was a problem hiding this comment.
Add required category markers to new tests.
These new tests should carry exactly one category marker (likely unit) to match repo test selection conventions.
Suggested patch
+@pytest.mark.unit
def test_click_type_int_literal_returns_int():
"""Literal[4, 8] must parse CLI values as ints before Pydantic validation."""
from typing import Literal
assert _click_type(Literal[4, 8]) == click.INT
+@pytest.mark.unit
def test_click_type_string_literal_plus_int_literal_returns_string():
"""String literals still need STRING so Pydantic can validate sentinel values."""
from typing import Literal
assert _click_type(Literal["disabled", 4]) == click.STRING+@pytest.mark.unit
def test_literal_int_override_end_to_end_via_click_runner():
"""A numeric Literal option must reach Pydantic as an int, not a string."""
captured: dict = {}As per coding guidelines, tests/**/test_*.py must “Assign exactly one category marker (unit, smoke, or e2e) to every test function.”
Also applies to: 455-467
Greptile SummaryThis PR adds batch-aware VRAM preflight estimation, splitting the estimate into three components (base weights, bf16 activations, and overhead), and introduces a hard-fail
Confidence Score: 5/5Safe to merge — all changed production paths are well-tested and the refactoring is structurally clean. The core VRAM estimation logic is well-exercised by boundary tests, the get_max_vram/get_max_memory_map split is straightforward, and the Literal CLI fix has both unit and end-to-end test coverage. The only gap is an incomplete assertion in one new test that does not affect production behavior. tests/preflight/test_preflight.py — the test_gradient_accumulation_steps_does_not_fake_batch test assertion is incomplete (see inline comment). Important Files Changed
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| def test_ample_vram_is_silent(self, default_config): | ||
| """Same config on an 80 GiB GPU must not warn (QLoRA ~5 GiB base + overhead).""" | ||
| metadata = MagicMock(spec=ModelMetadata, autoconfig=self._autoconfig()) | ||
| metadata = self._metadata(autoconfig=self._autoconfig()) | ||
| fake_props = MagicMock(total_memory=80 * 1024**3) | ||
| with ( | ||
| patch("torch.cuda.is_available", return_value=True), | ||
| patch("nemo_safe_synthesizer.llm.utils.get_max_vram", return_value={0: 1.0}), | ||
| patch("torch.cuda.get_device_properties", return_value=fake_props), | ||
| ): | ||
| issues = VRAMHeadroomCheck().run(make_ctx(config=default_config, metadata=metadata)) | ||
| assert not any(i.code == "low_vram" for i in issues) |
There was a problem hiding this comment.
test_ample_vram_is_silent
The docstring still says "QLoRA ~5 GiB base + overhead", but with the PR's new logic the default config has quantize_model=False, so bytes_per_base_weight returns 2.0 (bf16) and the base-weight estimate is ~15 GiB, not ~5 GiB. The test still passes because 80 GiB is generous, but the comment misrepresents what's being exercised.
Additionally, the assertion not any(i.code == "low_vram" ...) does not guard against vram_exceeds_capacity. If the estimate ever crossed the 1.5× hard-fail threshold it would be silently ignored, contradicting the "is_silent" intent of the test. A paired assertion like assert not any(i.code == "vram_exceeds_capacity" for i in issues) (or just assert issues == []) would close the gap.
8518a48 to
89b1f05
Compare
coderabbitai
Bot
added
feature
| ) | ||
| qualifier = "" if method == "exact" else " (approximate base param count; shape-heuristic fallback)" | ||
| ratio = estimated_gib / max_free_gib if max_free_gib > 0 else float("inf") | ||
| hard_fail = ratio >= _VRAM_HARD_FAIL_RATIO |
There was a problem hiding this comment.
question: How confident are we that hard_fail=True is a guarantee the run will crash? Is the estimated vram usage a lower bound?
If someone still wants to run and see what happens they can skip the vram check, though takes a bit to find how to do that. We could give a pointer to skip on the hard fail I guess, though that's probably more confusing at this point.
Good to go with what's here and we'll see if we notice any issues where the estimate is too large and the check fails when the run can actually complete.
89b1f05 to
f26e9b5
Compare
coderabbitai
Bot
added
docs
There was a problem hiding this comment.
Actionable comments posted: 4
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d1d4cf7b-9f1a-4ebd-a456-40d6bd269e99
📒 Files selected for processing (10)
docs/user-guide/running.mddocs/user-guide/troubleshooting.mdsrc/nemo_safe_synthesizer/configurator/pydantic_click_options.pysrc/nemo_safe_synthesizer/llm/utils.pysrc/nemo_safe_synthesizer/preflight/checks/environment.pysrc/nemo_safe_synthesizer/training/huggingface_backend.pytests/configurator/test_pydantic_click_options.pytests/llm/test_utils.pytests/preflight/test_preflight.pytests/training/test_huggingface_backend.py
🚧 Files skipped from review as they are similar to previous changes (7)
- src/nemo_safe_synthesizer/training/huggingface_backend.py
- tests/configurator/test_pydantic_click_options.py
- tests/training/test_huggingface_backend.py
- tests/llm/test_utils.py
- tests/preflight/test_preflight.py
- src/nemo_safe_synthesizer/configurator/pydantic_click_options.py
- src/nemo_safe_synthesizer/preflight/checks/environment.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.11)
- GitHub Check: Unit Tests (3.12)
- GitHub Check: Unit Tests (3.13)
- 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:
docs/user-guide/running.mddocs/user-guide/troubleshooting.mdsrc/nemo_safe_synthesizer/llm/utils.py
**/*.{md,markdown}
📄 CodeRabbit inference engine (.cursor/rules/agent-markdown-style.mdc)
**/*.{md,markdown}: Bold is acceptable only in markdown tables where it's the conventional way to mark header-like cells in the body
Use##headers to segment markdown sections instead of bold text
Use--(em-dash) instead of-(hyphen) for asides in markdown
Files:
docs/user-guide/running.mddocs/user-guide/troubleshooting.md
**/*.md
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
**/*.md: No decorative**bold**in body text, list items, or docstrings. Use headers, list markers, colons, and backticks for structure.
Use--(em-dash) for asides, not-(hyphen).
Use single backticks for code identifiers, paths, and CLI commands in Markdown.
Use Mermaid diagrams with no spaces in node IDs, quote labels with special characters, no explicit colors or styles.
Include SPDX copyright header in Markdown files using HTML comments:<!-- SPDX-FileCopyrightText: ... -->and<!-- SPDX-License-Identifier: Apache-2.0 -->. Exception: for.mdfiles with YAML frontmatter, include hash-comment headers inside the frontmatter block.All Markdown files require SPDX copyright headers, automatically added by
make format
Files:
docs/user-guide/running.mddocs/user-guide/troubleshooting.md
docs/**/*.md
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
Classify documentation pages as tutorial, how-to, explanation, or reference per the Diataxis framework. Use MkDocs Material syntax for admonitions (
!!! note), tabs (===), and code blocks with titles and highlights.
docs/**/*.md: Classify documentation content using the Diataxis framework (TUTORIAL, HOW-TO, EXPLANATION, or REFERENCE) and ensure each page fits ONE type only
Use MkDocs Material admonition syntax (!!! note, !!! warning, ??? tip) for callouts and collapsible content
Use MkDocs Material tab syntax (=== "Tab Name") to present multiple variations or language-specific examples
Include code block metadata in MkDocs Material format: use title attribute for filenames and hl_lines for syntax highlighting of specific lines
Use Mermaid diagram syntax for flowcharts and visual representations in documentation
List prerequisites at the top of each documentation page before main content
End documentation pages with 'Next steps' section containing links to related content
docs/**/*.md: Documentation pages must follow Diataxis framework organization: getting-started/ for tutorials, user-guide/ for how-tos and reference, architecture/ for explanations, reference/ for API docs (auto-generated), dev-notes/ for release notes
Add new documentation pages to thenav:section ofmkdocs.ymlfor sidebar appearance
Use MkDocs Material Markdown extensions including admonitions (!!! note, !!! warning), content tabs (===), code blocks with syntax highlighting, mermaid diagrams, task lists, footnotes, and definition lists
Files:
docs/user-guide/running.mddocs/user-guide/troubleshooting.md
**/*
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
**/*: Include a newline at the end of all files, never trailing whitespace. This is enforced bypre-commit.
Use line length of 120 characters for code, comments, and docstrings (configured inruff.toml).
Files:
docs/user-guide/running.mddocs/user-guide/troubleshooting.mdsrc/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:
docs/user-guide/running.mddocs/user-guide/troubleshooting.mdsrc/nemo_safe_synthesizer/llm/utils.py
docs/**
⚙️ CodeRabbit configuration file
Review documentation as MkDocs Material content. Check Diataxis fit, accurate commands, internal links, code fences, and markdown style from STYLE_GUIDE.md.
Files:
docs/user-guide/running.mddocs/user-guide/troubleshooting.md
**/*.py
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
**/*.py: Use American English spelling: 'initialize' not 'initialise', 'recognize' not 'recognise', 'color' not 'colour'.
Useobservability.get_logger(__name__)for logging, neverlogging.getLogger()orstructlog.get_logger()directly.
Use category loggers:.runtimefor internals,.userfor progress/results,.systemfor system events.
Never useprint()for operational output. Useclick.echo()for CLI output orsys.stdout.write()for raw output in tools.
Useextra={}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.
UseNSSBaseModelfor config/parameter models inconfig/which define user-facing configuration. Use rawBaseModelor module-specific bases for data transfer objects and internal structures.
UseBaseSettingsfor env/CLI settings. PreferAliasChoiceson individual fields when a field needs to respond to both its Python name and an env var name.
IncludeField(description=...)for Pydantic model fields as the canonical field docstring for API documentation and CLI help text.
Use assignment-styletype = Field(default=..., description="...")as the default for Pydantic model fields because type checkers understanddefault,default_factory, andaliasin assignment style.
UseAnnotatedonly when the field carries additional metadata beyondField()--ValueValidator,AutoParam,DependsOnValidator, reusable constrained type aliases, nested-type constraints, or discriminated unions.
Put defaults as bare assignment (= value), not insideField(default=...), when usingAnnotated. Exception: use assignment-styleField(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/llm/utils.py
src/**/*.py
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
src/**/*.py: Use relative imports insrc/(e.g.,from ..observability import get_logger).
Do not useprint()statements in library code. Useget_logger(__name__)fromobservability.pyorclick.echo()for CLI.
Do not useassertfor validation in library code. Useif/raisefor input validation.assertstatements can be stripped by-Oand must never guard correctness.
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
**/*.{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. Themake formatcommand handles this automatically.
Files:
src/nemo_safe_synthesizer/llm/utils.py
🪛 LanguageTool
docs/user-guide/troubleshooting.md
[style] ~75-~75: Consider using “incompatible” to avoid wordiness.
Context: ...des several early 5.x releases that are not compatible with its runtime. Keep vLLM's exclusion...
(NOT_ABLE_PREMIUM)
🪛 markdownlint-cli2 (0.22.1)
docs/user-guide/running.md
[warning] 128-128: Link fragments should be valid
(MD051, link-fragments)
🪛 Ruff (0.15.15)
src/nemo_safe_synthesizer/llm/utils.py
[warning] 609-609: Do not catch blind exception: Exception
(BLE001)
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
Actionable comments posted: 4
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d1d4cf7b-9f1a-4ebd-a456-40d6bd269e99
📒 Files selected for processing (10)
docs/user-guide/running.mddocs/user-guide/troubleshooting.mdsrc/nemo_safe_synthesizer/configurator/pydantic_click_options.pysrc/nemo_safe_synthesizer/llm/utils.pysrc/nemo_safe_synthesizer/preflight/checks/environment.pysrc/nemo_safe_synthesizer/training/huggingface_backend.pytests/configurator/test_pydantic_click_options.pytests/llm/test_utils.pytests/preflight/test_preflight.pytests/training/test_huggingface_backend.py
🚧 Files skipped from review as they are similar to previous changes (7)
- src/nemo_safe_synthesizer/training/huggingface_backend.py
- tests/configurator/test_pydantic_click_options.py
- tests/training/test_huggingface_backend.py
- tests/llm/test_utils.py
- tests/preflight/test_preflight.py
- src/nemo_safe_synthesizer/configurator/pydantic_click_options.py
- src/nemo_safe_synthesizer/preflight/checks/environment.py
🛑 Comments failed to post (4)
docs/user-guide/running.md (1)
128-129:
⚠️ Potential issue | 🟡 Minor | ⚡ Quick winFix broken internal anchor to
run --validate.The fragment
#run-validatedoes not match the heading slug for### `run --validate`and fails markdown link-fragment validation.🧰 Tools
🪛 markdownlint-cli2 (0.22.1)
[warning] 128-128: Link fragments should be valid
(MD051, link-fragments)
docs/user-guide/troubleshooting.md (1)
104-105:
⚠️ Potential issue | 🟡 Minor | ⚡ Quick winUse single backticks for inline identifiers.
tokenizer.modelandtokenizer.jsonare wrapped in double backticks; switch these to single backticks to match docs style rules.As per coding guidelines, "Use single backticks for code identifiers, paths, and CLI commands in Markdown."
src/nemo_safe_synthesizer/llm/utils.py (2)
467-469:
⚠️ Potential issue | 🟡 Minor | ⚡ Quick winLog the actual allocated fraction, not the requested fraction.
The message reports
max_vram_fraction * 100%, but allocation is clamped by the 2 GiB safety buffer and may be much lower. This can mislead operators during OOM triage.Proposed fix
- logger.info( - f"GPU {i}: Will allocate {memory_gib:.2f}GiB ({max_vram_fraction * 100}% of {total / (1024**3):.2f}GiB)" - ) + logger.info( + f"GPU {i}: Will allocate {memory_gib:.2f}GiB " + f"({gpu_memory_utilization * 100:.1f}% of {total / (1024**3):.2f}GiB)" + )
603-610:
⚠️ Potential issue | 🟡 Minor | ⚡ Quick win🧩 Analysis chain
🏁 Script executed:
#!/bin/bash set -euo pipefail # Locate the function and show surrounding context (including logger setup if any) rg -n "def get_device_name" -S src/nemo_safe_synthesizer/llm/utils.py python - <<'PY' import itertools, pathlib path = pathlib.Path("src/nemo_safe_synthesizer/llm/utils.py") start = 560 end = 680 lines = path.read_text().splitlines() for i in range(start-1, min(end, len(lines))): print(f"{i+1:4d}: {lines[i]}") PY # Find call sites/usages to understand expected behavior (should it ever raise?) rg -n "get_device_name\(" -S src/nemo_safe_synthesizer | head -n 50Repository: NVIDIA-NeMo/Safe-Synthesizer
Length of output: 5772
🏁 Script executed:
#!/bin/bash set -euo pipefail python - <<'PY' import pathlib path = pathlib.Path("src/nemo_safe_synthesizer/llm/utils.py") text = path.read_text().splitlines() # show top 120 lines for logger definition for i,line in enumerate(text[:140], start=1): print(f"{i:4d}: {line}") PYRepository: NVIDIA-NeMo/Safe-Synthesizer
Length of output: 6376
🏁 Script executed:
#!/bin/bash set -euo pipefail python - <<'PY' import pathlib path = pathlib.Path("src/nemo_safe_synthesizer/sdk/library_builder.py") lines = path.read_text().splitlines() for i in range(80, 160): if 1 <= i <= len(lines): print(f"{i:4d}: {lines[i-1]}") PYRepository: NVIDIA-NeMo/Safe-Synthesizer
Length of output: 3717
🏁 Script executed:
#!/bin/bash set -euo pipefail python - <<'PY' import pathlib path = pathlib.Path("src/nemo_safe_synthesizer/sdk/library_builder.py") lines = path.read_text().splitlines() for i in range(80, 160): if 1 <= i <= len(lines): print(f"{i:4d}: {lines[i-1]}") PYRepository: NVIDIA-NeMo/Safe-Synthesizer
Length of output: 3717
Avoid swallowing all exceptions in
get_device_name()
get_device_name()catchesExceptionand returns"undefined"without logging; since this value feeds telemetry (_build_telemetry_event), unrelated failures become hard to diagnose (and telemetry may be dropped entirely because_emit_nss_telemetrysuppresses exceptions).
Handle expected failure modes explicitly (includeImportErrorfor missingtorch, and CUDA/driver issues likeRuntimeError/AssertionError), and log a debug message withexc_info=Truewhen falling back to"undefined".🧰 Tools
🪛 Ruff (0.15.15)
[warning] 609-609: Do not catch blind exception:
Exception(BLE001)
coderabbitai
Bot
added
feature
coderabbitai
Bot
added
bug
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
- llm/utils.py: log the actual allocated VRAM fraction instead of the requested max_vram_fraction (clamped by the 2 GiB safety buffer), and narrow get_device_name() to expected failures with a debug log on fallback. - troubleshooting.md: use single backticks for inline tokenizer identifiers. Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Pop max_vram_fraction unconditionally in prepare_config. Previously it was only popped inside the add_max_memory branch, so with add_max_memory=False the NSS-internal kwarg leaked into framework_load_params and reached AutoModelForCausalLM.from_pretrained(), raising TypeError. Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Default config uses quantize_model=False, so the base-weight estimate is bf16 (~15 GiB), not QLoRA ~5 GiB. Update the docstring accordingly and add a paired assertion against vram_exceeds_capacity so the "is_silent" test also guards the hard-fail threshold. Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
aa7fc5c to
98053d7
Compare
3 participants
Summary
training.quantize_modeland keep bf16 compute activation wording explicit.max_memorybyte maps, and fix numericLiteralCLI parsing for--training__quantization_bits.Test plan
uv run --frozen pytest tests/preflight/test_preflight.py::TestVRAMHeadroomCheck tests/training/test_huggingface_backend.py::TestPrepareConfigIntegration tests/llm/test_utils.py::test_vram_helpers_return_fraction_and_hf_memory_map -vvs -n0uv run --frozen pytest tests/configurator/test_pydantic_click_options.py::test_click_type_int_literal_returns_int tests/configurator/test_pydantic_click_options.py::test_click_type_string_literal_plus_int_literal_returns_string tests/configurator/test_pydantic_click_options.py::test_literal_int_override_end_to_end_via_click_runner -vvs -n0uv run --frozen safe-synthesizer run --data-source "$HOME/dev/financial_transactions.csv" --validate --training__batch_size 100000 --training__quantize_model true --training__quantization_bits 4Made with Cursor
Summary by CodeRabbit
New Features
Bug Fixes
Documentation
Tests