fix: add a mechanism for allowing more max tokens#662
Conversation
Signed-off-by: Matt Kornfield <mkornfield@nvidia.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📜 Recent review details⏰ Context from checks skipped due to timeout. (1)
WalkthroughThe change adds a validated generation token multiplier, applies it to metadata-derived token budgets with context-window clamping, and forwards it through the timeseries and vLLM backends. ChangesGeneration token budget
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested labels: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
coderabbitai
Bot
added
feature
Greptile SummaryThis PR adds a user-configurable
Confidence Score: 5/5Safe to merge — the change is additive, backward-compatible, and well-tested. All existing callers that omit the multiplier argument continue to use the legacy constant via the No files require special attention. Important Files Changed
|
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d93b73ce-05d8-41dd-b23b-b2c8b3cab87b
📒 Files selected for processing (7)
src/nemo_safe_synthesizer/config/generate.pysrc/nemo_safe_synthesizer/generation/timeseries_backend.pysrc/nemo_safe_synthesizer/generation/vllm_backend.pysrc/nemo_safe_synthesizer/llm/metadata.pytests/config/test_generate.pytests/generation/test_vllm_backend.pytests/llm/test_metadata.py
📜 Review details
⏰ Context from checks skipped due to timeout. (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: End-user Wheel Install
- GitHub Check: Greptile Review
🧰 Additional context used
📓 Path-based instructions (14)
**/*.{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/generation/timeseries_backend.pytests/generation/test_vllm_backend.pysrc/nemo_safe_synthesizer/generation/vllm_backend.pytests/config/test_generate.pysrc/nemo_safe_synthesizer/llm/metadata.pytests/llm/test_metadata.pysrc/nemo_safe_synthesizer/config/generate.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.
UseField(description=...)for every Pydantic model field.
Use assignment-styleField()by default; useAnnotatedonly 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.
Usefield(default_factory=list)instead of mutable list defaults.
UseStrEnumfor string-valued configuration or serialization enums and plainEnumfor internal constants.
Obtain loggers withobservability.get_logger(__name__); do not calllogging.getLogger()orstructlog.get_logger()directly.
Use.runtime,.user, and.systemcategory loggers appropriately.
Do not useprint()for operational library output; use the approved logger,click.echo(), orsys.stdout.write()where appropriate.
Useextra={}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.
PreferX | Y, built-in collection generics, andSelfoverOptional,Union, and legacy typing collections.
Use collection ABCs for function arguments and concrete collection types for return values.
UseProtocolfor structural subtyping and avoidAnywhenobject, generics, or protocols are suitable.
UseTYPE_CHECKINGguards for heavy imports such as pandas, torch, and transformers.
...
Files:
src/nemo_safe_synthesizer/generation/timeseries_backend.pytests/generation/test_vllm_backend.pysrc/nemo_safe_synthesizer/generation/vllm_backend.pytests/config/test_generate.pysrc/nemo_safe_synthesizer/llm/metadata.pytests/llm/test_metadata.pysrc/nemo_safe_synthesizer/config/generate.py
src/**/*.py
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
src/**/*.py: Use relative imports for package code undersrc/.
Do not useassertfor 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/generation/timeseries_backend.pysrc/nemo_safe_synthesizer/generation/vllm_backend.pysrc/nemo_safe_synthesizer/llm/metadata.pysrc/nemo_safe_synthesizer/config/generate.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/generation/timeseries_backend.pysrc/nemo_safe_synthesizer/generation/vllm_backend.pysrc/nemo_safe_synthesizer/llm/metadata.pysrc/nemo_safe_synthesizer/config/generate.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 thanmainmust follow<author>/<description>, optionally including an issue ID or type; branch names must use lowercase alphanumeric characters and hyphens.
Commits merged tomainmust follow Conventional Commits, using a valid lowercase type and a description of at most 100 characters.
Files:
src/nemo_safe_synthesizer/generation/timeseries_backend.pytests/generation/test_vllm_backend.pysrc/nemo_safe_synthesizer/generation/vllm_backend.pytests/config/test_generate.pysrc/nemo_safe_synthesizer/llm/metadata.pytests/llm/test_metadata.pysrc/nemo_safe_synthesizer/config/generate.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/generation/timeseries_backend.pytests/generation/test_vllm_backend.pysrc/nemo_safe_synthesizer/generation/vllm_backend.pytests/config/test_generate.pysrc/nemo_safe_synthesizer/llm/metadata.pytests/llm/test_metadata.pysrc/nemo_safe_synthesizer/config/generate.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/generation/timeseries_backend.pytests/generation/test_vllm_backend.pysrc/nemo_safe_synthesizer/generation/vllm_backend.pytests/config/test_generate.pysrc/nemo_safe_synthesizer/llm/metadata.pytests/llm/test_metadata.pysrc/nemo_safe_synthesizer/config/generate.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
typestatements or bracketed generic class/function parameters.
Files:
src/nemo_safe_synthesizer/generation/timeseries_backend.pytests/generation/test_vllm_backend.pysrc/nemo_safe_synthesizer/generation/vllm_backend.pytests/config/test_generate.pysrc/nemo_safe_synthesizer/llm/metadata.pytests/llm/test_metadata.pysrc/nemo_safe_synthesizer/config/generate.py
**/*.{py,md}
📄 CodeRabbit inference engine (CONTRIBUTING.md)
Use the repository's documented Python and Markdown style conventions and validate changes with the pinned
miseformatting and checking tasks.
Files:
src/nemo_safe_synthesizer/generation/timeseries_backend.pytests/generation/test_vllm_backend.pysrc/nemo_safe_synthesizer/generation/vllm_backend.pytests/config/test_generate.pysrc/nemo_safe_synthesizer/llm/metadata.pytests/llm/test_metadata.pysrc/nemo_safe_synthesizer/config/generate.py
**/*.{py,sh}
📄 CodeRabbit inference engine (CONTRIBUTING.md)
Use the repository's pinned
misetasks for formatting, linting, type checking, and testing rather than invokingruffortydirectly for project-wide checks.
Files:
src/nemo_safe_synthesizer/generation/timeseries_backend.pytests/generation/test_vllm_backend.pysrc/nemo_safe_synthesizer/generation/vllm_backend.pytests/config/test_generate.pysrc/nemo_safe_synthesizer/llm/metadata.pytests/llm/test_metadata.pysrc/nemo_safe_synthesizer/config/generate.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/timeseries_backend.pysrc/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.pytests/config/test_generate.pytests/llm/test_metadata.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.pytests/config/test_generate.pytests/llm/test_metadata.py
tests/**/*.py
📄 CodeRabbit inference engine (tests/TESTING.md)
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 fromtests/stub_datasets/as HuggingFaceDatasetobjects
Useload_test_dataframe(filename)helper to load test data files fromtests/stub_datasets/as pandas DataFrames
Convert pandas columns to nullable dtypes (pd.Int64Dtype(),pd.BooleanDtype()) before assigningnp.nanvalues
Usefake.seed_instance(seed)andrandom.seed(seed)together for Faker-based test data reproducibility
When sharing methods across multiple test files, define them inconftest.pyand import them using relative imports (e.g.,from .conftest import train_with_sdk); note that importing from other test files liketests/cli/helpers.pydoes not work
Usefixture_mock_processororfixture_mock_processor_without_valid_recordsfor mocking ParsedResponse objects withvalid_records,invalid_records,errors, andprompt_numberfields
Usepytest.importorskipto 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 (ruffT201is suppressed fortests/directory) and should...
Files:
tests/generation/test_vllm_backend.pytests/config/test_generate.pytests/llm/test_metadata.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.pytests/config/test_generate.pytests/llm/test_metadata.py
src/**/config/**/*.py
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
Use
NSSBaseModelfor user-facing configuration and parameter models.
Files:
src/nemo_safe_synthesizer/config/generate.py
src/nemo_safe_synthesizer/config/**/*.py
⚙️ CodeRabbit configuration file
Treat config changes as user-facing API changes. Check Pydantic field descriptions, defaults, validators, aliases, override behavior, CLI help text impact, YAML compatibility, and documented parameter semantics.
Files:
src/nemo_safe_synthesizer/config/generate.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.pytests/config/test_generate.pytests/llm/test_metadata.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/llm/metadata.py (1)
453-489: LGTM!tests/config/test_generate.py (1)
222-232: LGTM!tests/llm/test_metadata.py (1)
692-717: LGTM!src/nemo_safe_synthesizer/generation/vllm_backend.py (1)
741-744: LGTM!
| max_tokens_multiplier: Annotated[ | ||
| float, | ||
| ValueValidator(value_func=lambda v: v > 0), | ||
| Field( | ||
| title="max_tokens_multiplier", | ||
| description=( | ||
| "Margin applied to the longest tokenized training example " | ||
| "(``max_tokens_per_example``) when sizing the per-sample generation " | ||
| "budget (``SamplingParams.max_tokens``). The effective budget is " | ||
| "``int(max_tokens_per_example * max_tokens_multiplier)``, still clamped " | ||
| "to the remaining context window. The default adds only a small jitter " | ||
| "margin, which suffices for most tables but is too tight for long, " | ||
| "unbounded free-text columns: a model that over-generates slightly past " | ||
| "the longest training example truncates mid-JSON and produces no " | ||
| "parseable record (``length``-dominated finish reasons, zero valid " | ||
| "records). Raise this (e.g. to 1.5-2.0) to give the model room to emit " | ||
| "the closing tokens on such datasets; there is no benefit to values that " | ||
| "push the budget past the context window. Must be > 0." | ||
| ), | ||
| ), | ||
| ] = 1.2 # mirrors llm.metadata.GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER (kept a literal to avoid a config->llm import cycle) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- generate.py around max_tokens_multiplier ---'
sed -n '220,285p' src/nemo_safe_synthesizer/config/generate.py
echo
echo '--- metadata.py around budget sizing ---'
sed -n '470,505p' src/nemo_safe_synthesizer/llm/metadata.py
echo
echo '--- tests around vLLM backend generation token propagation ---'
sed -n '1020,1105p' tests/generation/test_vllm_backend.py
echo
echo '--- search for related multiplier / finite validation ---'
rg -n "max_tokens_multiplier|ValueValidator|finite|isfinite|OverflowError|GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER" src/nemo_safe_synthesizer tests -SRepository: NVIDIA-NeMo/Safe-Synthesizer
Length of output: 22515
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import math
cases = [float('inf'), float('-inf'), float('nan'), 1e308, 1e309]
for v in cases:
print("case:", v)
try:
scaled = int(v * 1.2)
print(" int(v * 1.2) =>", scaled)
except Exception as e:
print(" int(v * 1.2) raised", type(e).__name__, e)
try:
print(" v > 0 =>", v > 0)
except Exception as e:
print(" v > 0 raised", type(e).__name__, e)
print(" isfinite =>", math.isfinite(v))
PYRepository: NVIDIA-NeMo/Safe-Synthesizer
Length of output: 1029
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n "class ValueValidator|def .*ValueValidator|allow_inf_nan|finite float|isfinite|math\.isfinite" src -S
echo
echo '--- model base class / config settings ---'
rg -n "class NSSBaseModel|ConfigDict|allow_inf_nan|validate_default" src/nemo_safe_synthesizer -SRepository: NVIDIA-NeMo/Safe-Synthesizer
Length of output: 4633
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- NSSBaseModel ---'
rg -n "class NSSBaseModel|allow_inf_nan|ConfigDict" src/nemo_safe_synthesizer -S
echo
echo '--- max_tokens_multiplier field line ---'
sed -n '236,270p' src/nemo_safe_synthesizer/config/generate.py
echo
echo '--- exact int conversion site ---'
sed -n '486,498p' src/nemo_safe_synthesizer/llm/metadata.pyRepository: NVIDIA-NeMo/Safe-Synthesizer
Length of output: 6423
Reject non-finite multipliers before sizing the budget.
ValueValidator(lambda v: v > 0) still allows inf, and int(max_tokens_per_example * multiplier) raises OverflowError before the context clamp. Tighten max_tokens_multiplier to finite-positive values, or clamp the product before int(). Add a regression for float("inf").
📍 Affects 2 files
src/nemo_safe_synthesizer/config/generate.py#L244-L264(this comment)src/nemo_safe_synthesizer/llm/metadata.py#L490-L496
| max_tokens=self.model_metadata.generation_max_tokens_for( | ||
| self._get_prompt_token_count(), | ||
| multiplier=self.config.generation.max_tokens_multiplier, | ||
| ), |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Exercise non-default multiplier propagation in both backends.
Current backend tests use the default multiplier, so they do not prove that a user-selected value reaches the budget helper.
src/nemo_safe_synthesizer/generation/timeseries_backend.py#L924-L927: add coverage with a non-default multiplier.tests/generation/test_vllm_backend.py#L1046-L1049: override the fixture multiplier before the fallback-path assertion.tests/generation/test_vllm_backend.py#L1079-L1081: override the fixture multiplier before the initialized-engine assertion.
📍 Affects 2 files
src/nemo_safe_synthesizer/generation/timeseries_backend.py#L924-L927(this comment)tests/generation/test_vllm_backend.py#L1046-L1049tests/generation/test_vllm_backend.py#L1079-L1081
Signed-off-by: Matt Kornfield <mkornfield@nvidia.com>
Signed-off-by: Matt Kornfield <mkornfield@nvidia.com>
1 participant
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
max_tokens_multiplierto control the per-request generation token budget, defaulting to 1.2.