fix: tighter memory checks between training and generate - #598
fix: tighter memory checks between training and generate#598mckornfield wants to merge 1 commit into
Conversation
WalkthroughVllmBackend initialization now enforces VRAM utilization from config instead of runtime queries, calling cleanup_memory with enhanced CUDA operations beforehand. SafeSynthesizer tracks training elapsed time across lifecycle methods and explicitly tears down the trainer instance with dedicated state management. ChangesResource Management Refactoring: VRAM Configuration and Training Lifecycle
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
coderabbitai
Bot
added
bug
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
| self._training_time_sec = result.elapsed_time | ||
| trainer.teardown() | ||
| del self.trainer | ||
| trainer = None |
Greptile SummaryThis PR addresses two memory-management problems that surface in train-then-generate pipelines:
Confidence Score: 4/5The change is safe to merge; all three fixes address real, reproducible failure modes in the train→generate→evaluate pipeline. The core logic — using a config-backed memory fraction instead of a dynamically computed one, capturing training time before trainer teardown, and hardening cleanup_memory — is straightforward and well-covered by the new test. The only open question is whether torch.cuda.synchronize() should target all devices rather than just the default, which matters only on multi-GPU nodes and leaves the memory-release intent intact even when only partially fulfilled. src/nemo_safe_synthesizer/llm/utils.py — the synchronize() call targets only the current CUDA device. Important Files Changed
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| with torch.no_grad(): | ||
| torch.cuda.synchronize() | ||
| torch.cuda.empty_cache() | ||
| torch.cuda.ipc_collect() |
There was a problem hiding this comment.
torch.cuda.synchronize() called without a device argument synchronizes only the currently selected CUDA device (device 0 by default). After multi-GPU training, streams on other devices remain unsynchronized, so their work may not be fully flushed before empty_cache() runs, potentially leaving memory unreleased on non-default devices.
| with torch.no_grad(): | |
| torch.cuda.synchronize() | |
| torch.cuda.empty_cache() | |
| torch.cuda.ipc_collect() | |
| with torch.no_grad(): | |
| for i in range(torch.cuda.device_count()): | |
| torch.cuda.synchronize(i) | |
| torch.cuda.empty_cache() | |
| torch.cuda.ipc_collect() |
| trainer.teardown() | ||
| del self.trainer | ||
| trainer = None |
There was a problem hiding this comment.
del self.trainer, assigning trainer = None reassigns the local variable but has no effect — the local is never read again and the attribute has already been removed from the instance. The assignment is a no-op that may imply to readers that trainer is used after this point.
| trainer.teardown() | |
| del self.trainer | |
| trainer = None | |
| trainer.teardown() | |
| del self.trainer |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ea827943-c74d-49b3-b00b-40ecfbde6e53
📒 Files selected for processing (4)
src/nemo_safe_synthesizer/generation/vllm_backend.pysrc/nemo_safe_synthesizer/llm/utils.pysrc/nemo_safe_synthesizer/sdk/library_builder.pytests/generation/test_vllm_backend.py
📜 Review details
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (6)
- GitHub Check: Unit Tests (3.12)
- GitHub Check: Unit Tests (3.11)
- GitHub Check: Unit Tests (3.13)
- GitHub Check: Smoke Tests
- GitHub Check: Analyze (Python)
- GitHub Check: Analyze (python)
🧰 Additional context used
📓 Path-based instructions (12)
**/*.{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.pysrc/nemo_safe_synthesizer/generation/vllm_backend.pysrc/nemo_safe_synthesizer/sdk/library_builder.pytests/generation/test_vllm_backend.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: UseField(description=...)as the canonical field docstring for Pydantic models.
Use assignment style (type = Field(default=..., description="...")) as the default for Pydantic model fields rather thanAnnotated-only patterns.
UseAnnotatedonly when the field carries additional metadata beyondField()--ValueValidator,AutoParam,DependsOnValidator, reusable constrained type aliases, nested-type constraints, or discriminated unions.
Use@dataclass(frozen=True)for immutable value objects and validators. Mutable@dataclassacceptable for builders, accumulators, and pipeline state.
Usefield(default_factory=list)for mutable defaults in dataclasses, never= [].
UseStrEnumfor string-valued enums used in configs/serialization. PlainEnumfor internal-only named constants.
UseX | YnotOptional[X]orUnion[X, Y]in type hints.
Uselist[str]notList[str],dict[str, int]notDict[str, int]in type hints.
UseSelffor fluent method returns in type hints.
Use collection ABCs for function arguments (Sequence,Mapping,Iterable) so callers can pass any compatible container; concrete types for return values.
UseProtocolfor structural subtyping when you need duck-typing boundaries.
AvoidAnyin type hints -- preferobject, generics, orProtocol.
UseTYPE_CHECKINGguards for heavy imports (pandas,torch,transformers); not needed for stdlib or lightweight imports.
Prefermatch/casefor dispatch on types or tagged values.if/elifis fine for simple boolean predicates.
Use comprehensions over imperative loops where intent is clearer. No multipleforclauses -- optimize for readability, not conciseness.
Use builder pattern withwith_*methods...
Files:
src/nemo_safe_synthesizer/llm/utils.pysrc/nemo_safe_synthesizer/generation/vllm_backend.pysrc/nemo_safe_synthesizer/sdk/library_builder.pytests/generation/test_vllm_backend.py
src/**/*.py
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
src/**/*.py: Useobservability.get_logger(__name__)-- 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={}for data that downstream tools should query or aggregate in log calls (metrics, counts, durations). f-strings are fine for human-readable context.
Raise from the custom error hierarchy with dual inheritance:SafeSynthesizerError(base),UserError,DataError,ParameterError,GenerationError,InternalError.
Prefer clamping/saturation over raising when out-of-range inputs shouldn't crash the system -- return a bounded value with a log warning.
Error messages must precisely match the actual error condition. Use!rfor repr of interpolated pieces to clearly identify them.
Use relative imports insrc/(from ..observability import get_logger), absolute imports intests/(from nemo_safe_synthesizer.observability import get_logger).
Usepathlib.Pathinstead ofos.path. Tolerateos.pathonly in vendored/tooling scripts.
Do not useprint()statements in library code. Useget_logger(__name__)fromobservability.pyorclick.echo()for CLI.print()is fine in tests, standalone scripts, and tooling.
Do not useassertfor validation in library code. Useif/raisefor input validation.assertis fine in tests wherepytestrelies on it.
Files:
src/nemo_safe_synthesizer/llm/utils.pysrc/nemo_safe_synthesizer/generation/vllm_backend.pysrc/nemo_safe_synthesizer/sdk/library_builder.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.pysrc/nemo_safe_synthesizer/generation/vllm_backend.pysrc/nemo_safe_synthesizer/sdk/library_builder.py
**/*.{py,sh,yaml,yml}
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
Include SPDX copyright headers in all source files using hash-comments for
.py,.sh,.yaml,.ymlfiles.
Files:
src/nemo_safe_synthesizer/llm/utils.pysrc/nemo_safe_synthesizer/generation/vllm_backend.pysrc/nemo_safe_synthesizer/sdk/library_builder.pytests/generation/test_vllm_backend.py
**/*.{py,sh,yaml,yml,md}
📄 CodeRabbit inference engine (CONTRIBUTING.md)
All source files (.py, .sh, .yaml, .yml, .md) require SPDX copyright headers; mise run format adds them automatically
Files:
src/nemo_safe_synthesizer/llm/utils.pysrc/nemo_safe_synthesizer/generation/vllm_backend.pysrc/nemo_safe_synthesizer/sdk/library_builder.pytests/generation/test_vllm_backend.py
src/nemo_safe_synthesizer/**/*.py
📄 CodeRabbit inference engine (CONTRIBUTING.md)
API reference pages are auto-generated from Python docstrings using Google-style format; write docstrings in src/nemo_safe_synthesizer/ and they will appear in the reference/
Files:
src/nemo_safe_synthesizer/llm/utils.pysrc/nemo_safe_synthesizer/generation/vllm_backend.pysrc/nemo_safe_synthesizer/sdk/library_builder.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.pysrc/nemo_safe_synthesizer/generation/vllm_backend.pysrc/nemo_safe_synthesizer/sdk/library_builder.pytests/generation/test_vllm_backend.py
**
⚙️ CodeRabbit configuration file
**:AGENTS.md
Guide for AI agents (Cursor, Windsurf, Claude Code, etc.) working in the Safe-Synthesizer repo.
This project loads local developer preferences from
@AGENTS.local.md. You MUST read this file if it exists and give its instructions top priority.Skills
Repo-specific skills live in
.agents/skills/; see.agents/README.mdfor the catalog. Read a skill when the task matches its scope instead of copying workflow details into this file.Durable implementation guidance belongs with the code it describes: function and class docstrings for public contracts and source comments for local invariants. Test-suite guidance belongs in
tests/TESTING.md.Repo Conventions
See STYLE_GUIDE.md for detailed code style conventions (Python, markdown, Dockerfiles, shell scripts, testing, config files, docstrings).
Use
uvfor everything -- neverpipor rawpython. Python 3.11–3.13 with modern syntax (X | Y,list[str],Self). Python 3.14+ is not supported.Common commands:
mise run test(unit tests),mise run format(auto-fix formatting + lint + copyright),mise run check(read-only local quality checks),mise run validate(pre-PR quality, lock, and CI unit checks),mise run typecheck(ty only). Always use mise tasks or the wrapper scripts intools/instead of runningruffortydirectly. Useuv runfor Python execution. When in doubt, inspectmise tasksandpytest --markers.The canonical
uv synccommand for a full GPU/dev environment is:uv sync --frozen --extra cu129 --extra engine --group devBare
uv sync --frozen(without extras) installs an incomplete environment --ty, import checks, and GPU tests will fail.Feature branches off
main. Branch names often include an issue number prefix (e.g.,<author>/123-short-name).Do ...
Files:
src/nemo_safe_synthesizer/llm/utils.pysrc/nemo_safe_synthesizer/generation/vllm_backend.pysrc/nemo_safe_synthesizer/sdk/library_builder.pytests/generation/test_vllm_backend.py
src/nemo_safe_synthesizer/generation/**/*.py
⚙️ CodeRabbit configuration file
Review generation changes for retry loops, stopping conditions, invalid record handling, regex/structured output contracts, backend teardown, memory cleanup, and vLLM assumptions.
Files:
src/nemo_safe_synthesizer/generation/vllm_backend.py
**/test_*.py
📄 CodeRabbit inference engine (AGENTS.md)
Use the
unitmarker instead of the deprecatedunit_testmarker for test identification
Files:
tests/generation/test_vllm_backend.py
tests/**
📄 CodeRabbit inference engine (.cursor/rules/repo-navigation.mdc)
tests/**: Mirrorsrc/directory structure intests/directory for test organization
Auto-mark tests by directory:tests/e2e/→e2e,tests/smoke/→smoke, otherwise default tounitMirror source code directory structure in tests directory (e.g.,
tests/training/,tests/generation/parallel to source structure)
Files:
tests/generation/test_vllm_backend.py
⚙️ CodeRabbit configuration file
tests/**:Testing Guide
Comprehensive testing reference for Safe-Synthesizer developers. Covers commands, markers, test data, fixtures, and gotchas.
Read First
tests/conftest.py-- auto-marking,load_test_dataset/load_test_dataframe,fixture_mock_processorpatternpytest.ini-- markers, asyncio, timeouttests/evaluation/conftest.py-- most complex: Faker-basedmake_df, nullable dtype conversiontests/generation/conftest.py-- JSONL/schema fixtures,fixture_valid_iris_dataset_jsonl_and_schemaRunning Tests
All mise test tasks, grouped by scope:
mise run test # Unit (excludes slow, e2e, and smoke) mise run test:unit-slow # Unit tests including slow (excludes e2e and smoke) mise run test:smoke # CPU smoke tests (~few min, no GPU required) mise run test:smoke:gpu # All staged GPU smoke tests (requires CUDA) mise run test:smoke:gpu:train-only mise run test:smoke:gpu:generation mise run test:smoke:gpu:resume mise run test:smoke:gpu:structured-generation mise run test:smoke:gpu:timeseries mise run test:smoke:gpu:smollm2 mise run test:e2e # All e2e (requires CUDA) -- runs default + dp mise run test:e2e:default # e2e default (no-DP) tests only mise run test:e2e:dp # e2e DP tests only mise run test:ci # CI unit tests with coverage (excludes slow, e2e, gpu, smoke) mise run test:ci-slow # CI slow tests with coverage mise run test:ci-container # CI tests in a Linux container (Docker/Podman)Run a single test:
uv run --frozen pytest tests/path/test_file.py::test_name -vvs -n0Test runner:
uv run --frozen pytest -n auto --dist loadscope -vv...
Files:
tests/generation/test_vllm_backend.py
tests/**/*.py
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
tests/**/*.py: Use file namingtest_*.py, class namingTest*, function namingtest_<module>_<expected_behavior>for test files.
Use function-scoped fixtures by default. Session scope only when empirically justified by test runtime.
Use bareassertas primary assertion style;pytest.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). Use 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 in tests, not internal implementation details.
Ensure test isolation: no shared mutable state or execution-order dependencies between tests. Include required setup in the test or a fixture.
Use@pytest.mark.parametrizefor testing multiple input combinations rather than copy-pasting similar tests.
tests/**/*.py: Auto-mark tests based on file path: tests under/e2e/gete2emarker, tests under/smoke/getsmokemarker, all others getunitmarker (only if no category marker already present)
Every test should have exactly one category marker:unit,smoke, ore2e
Usepytest.mark.requires_gpumodifier on tests that need CUDA hardware
Usepytest.mark.vllmon tests using vLLM generation backend and ensure each vLLM test file runs in its own process for GPU memory isolation
Usepytest.mark.slowon long-running tests
Usepytest.mark.smollm2for SmolLM2 Hub download tests to enable process isolation
Usepytest.mark.noautouseto skip autouse fixtures for specific tests
Useload_test_dataset(filename)helper to load test datasets from `...
Files:
tests/generation/test_vllm_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/generation/test_vllm_backend.py
🧠 Learnings (2)
📚 Learning: 2026-05-27T22:20:37.354Z
Learnt from: kendrickb-nvidia
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 520
File: tests/generation/test_vllm_backend.py:556-587
Timestamp: 2026-05-27T22:20:37.354Z
Learning: In NVIDIA-NeMo/Safe-Synthesizer, `tests/conftest.py`’s `pytest_collection_modifyitems` hook applies pytest category markers automatically based on each test file’s path: tests under `/e2e/` get `pytest.mark.e2e`, tests under `/smoke/` get `pytest.mark.smoke`, and all other tests get `pytest.mark.unit`. Therefore, when reviewing pytest tests outside `tests/e2e/` and `tests/smoke/`, do not flag missing explicit `pytest.mark.unit` decorators on test classes/functions as an issue (the hook will add them during collection). If a new test directory/category is introduced, ensure the hook is updated so it’s categorized correctly.
Applied to files:
tests/generation/test_vllm_backend.py
📚 Learning: 2026-06-04T16:14:09.868Z
Learnt from: binaryaaron
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 526
File: tests/generation/test_vllm_backend.py:399-509
Timestamp: 2026-06-04T16:14:09.868Z
Learning: In this repo, only apply `pytest.mark.vllm` to smoke tests under `tests/smoke/` that actually run real vLLM GPU generation and therefore require per-file process isolation (e.g., `test-smoke-gpu-*` Makefile targets). Do not apply `pytest.mark.vllm` to unit-style tests under `tests/generation/` that merely import `vllm_backend` but never instantiate a real vLLM engine and never call `.generate()` (GPU not required). Note that `tests/conftest.py` auto-marks these as `unit` via `pytest_collection_modifyitems`, and `vllm` is not among the auto-mark categories—so if a test in `tests/generation/` has `vllm`, it should be treated as a review issue unless it meets the real GPU generation criteria above.
Applied to files:
tests/generation/test_vllm_backend.py
🔇 Additional comments (4)
src/nemo_safe_synthesizer/generation/vllm_backend.py (1)
287-293: LGTM!src/nemo_safe_synthesizer/llm/utils.py (1)
409-415: LGTM!tests/generation/test_vllm_backend.py (2)
397-397: LGTM!
399-412: LGTM!
| if trainer is not None: | ||
| if result := getattr(trainer, "results", None): | ||
| self._training_time_sec = result.elapsed_time | ||
| trainer.teardown() | ||
| del self.trainer | ||
| trainer = None |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Clear self.trainer in a finally block.
If trainer.teardown() raises here, generate() exits before the generation backend is initialized and the trainer object stays attached to self, which defeats the memory-release step this PR is adding. Drop the attribute in finally so the large trainer graph is released even on teardown failures.
Suggested fix
trainer = getattr(self, "trainer", None)
if trainer is not None:
if result := getattr(trainer, "results", None):
self._training_time_sec = result.elapsed_time
- trainer.teardown()
- del self.trainer
- trainer = None
+ try:
+ trainer.teardown()
+ finally:
+ del self.trainer
+ trainer = NoneAs per coding guidelines, use try/finally for resource cleanup and keep teardown failures from blocking subsequent work.
Source: Coding guidelines
2 participants
Summary
Pre-Review Checklist
Ensure that the following pass:
mise run format && mise run checkor via prek validation.mise run testpasses locallymise run test:e2epasses locallymise run test:ci-containerpasses locally (recommended)/syncon this PR to trigger a run (auto-triggers on ready-for-review)Pre-Merge Checklist
Other Notes
Summary by CodeRabbit
Release Notes
Bug Fixes
Tests