Skip to content

fix: add a mechanism for allowing more max tokens#662

Open
mckornfield wants to merge 1 commit into
mainfrom
max-token-multiplier/mck
Open

fix: add a mechanism for allowing more max tokens#662
mckornfield wants to merge 1 commit into
mainfrom
max-token-multiplier/mck

Conversation

@mckornfield

@mckornfield mckornfield commented Jul 17, 2026

Copy link
Copy Markdown
Collaborator

Summary

Pre-Review Checklist

Ensure that the following pass:

  • mise run format && mise run check or via prek validation.
  • mise run test passes locally
  • mise run test:e2e passes locally
  • mise run test:ci-container passes locally (recommended)
  • GPU CI status check passes -- comment /sync on this PR to trigger a run (auto-triggers on ready-for-review)

Pre-Merge Checklist

  • New or updated tests for any fix or new behavior
  • Updated documentation for new features and behaviors, including docstrings for API docs.

Other Notes

  • Closes #

Summary by CodeRabbit

  • New Features

    • Added a configurable token-budget multiplier for generation, defaulting to a safety margin of 1.2.
    • Generation budgets can now be widened for longer free-text outputs while remaining within the model’s context window.
    • Invalid multiplier values, including zero and negative numbers, are rejected.
  • Bug Fixes

    • Improved max-token calculation consistency across supported generation backends.
    • Preserved prompt-length safeguards to prevent exceeding the available context window.

Signed-off-by: Matt Kornfield <mkornfield@nvidia.com>
@coderabbitai

coderabbitai Bot commented Jul 17, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Walkthrough

Changes

Generation token budget

Layer / File(s) Summary
Budget configuration and calculation
src/nemo_safe_synthesizer/config/generate.py, src/nemo_safe_synthesizer/llm/metadata.py, tests/config/test_generate.py, tests/llm/test_metadata.py
Adds a validated max_tokens_multiplier, applies it to stat-derived token budgets, preserves the default multiplier behavior, and clamps results to the model context window.
Backend multiplier wiring
src/nemo_safe_synthesizer/generation/timeseries_backend.py, src/nemo_safe_synthesizer/generation/vllm_backend.py, tests/generation/test_vllm_backend.py
Passes the configured multiplier through both generation backends and updates cached-prompt token-count assertions.

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

Suggested labels: feature, test

Suggested reviewers: binaryaaron

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly matches the main change: adding a configurable way to increase max token limits.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch max-token-multiplier/mck

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

@coderabbitai coderabbitai Bot added feature New feature or request test Test-only addition or change labels Jul 17, 2026
@greptile-apps

greptile-apps Bot commented Jul 17, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR adds a user-configurable max_tokens_multiplier field to GenerateParameters that controls how aggressively the per-sample generation budget is sized relative to the longest tokenized training example. The change is backward-compatible — the new parameter defaults to 1.2 (matching the existing GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER constant), and legacy callers such as the training eval callback are unaffected since they continue to hit the None-default path in the updated method.

  • generation_max_tokens_for gains an optional multiplier: float | None parameter; None resolves to the existing safety-margin constant, preserving all non-generation call sites.
  • vllm_backend.py and timeseries_backend.py now forward config.generation.max_tokens_multiplier to generation_max_tokens_for, giving users a runtime knob to widen the token budget for long, unbounded free-text columns without modifying model or training artefacts.
  • Tests cover validation rejection of non-positive values, widen-budget behavior, context-window clamping, and a synchronization check that the config literal matches the metadata constant.

Confidence Score: 5/5

Safe to merge — the change is additive, backward-compatible, and all existing call sites retain identical behavior through the None-default path.

The new max_tokens_multiplier knob is cleanly threaded from config through both generation backends with no changes to the training eval callback or any other caller. The context-window clamp in generation_max_tokens_for still bounds the value regardless of how large the multiplier is set, preventing any vLLM overrun. The literal-vs-constant duplication in generate.py is intentional (import cycle) and is actively guarded by a dedicated synchronization test. Test coverage is thorough across validation, widening, clamping, and plumbing scenarios.

No files require special attention.

Important Files Changed

Filename Overview
src/nemo_safe_synthesizer/llm/metadata.py Adds optional multiplier parameter to generation_max_tokens_for; defaults to None → GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER so all pre-existing callers retain identical behavior.
src/nemo_safe_synthesizer/config/generate.py Adds max_tokens_multiplier field with ValueValidator(v > 0), defaulting to literal 1.2 to avoid a config→llm import cycle; sync guarded by a dedicated unit test.
src/nemo_safe_synthesizer/generation/vllm_backend.py Threads config.generation.max_tokens_multiplier into the generation_max_tokens_for call; no other logic change.
src/nemo_safe_synthesizer/generation/timeseries_backend.py Same one-line plumbing change as vllm_backend.py; consistent with parent class change.
tests/llm/test_metadata.py Adds three well-targeted tests: None multiplier reproduces legacy sizing, custom multiplier widens budget, and aggressive multiplier is still clamped to the context window.
tests/config/test_generate.py Adds validation tests for the new field: default sync with metadata constant, acceptance of widened values, and rejection of zero/negative inputs.
tests/generation/test_vllm_backend.py Updates two existing plumbing assertions to verify generation_max_tokens_for is called with both the prompt token count and the configured multiplier keyword argument.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant User
    participant GenerateParameters
    participant VllmBackend
    participant TimeseriesBackend
    participant ModelMetadata

    User->>GenerateParameters: "max_tokens_multiplier=1.8 (or default 1.2)"
    GenerateParameters-->>VllmBackend: config.generation.max_tokens_multiplier
    GenerateParameters-->>TimeseriesBackend: config.generation.max_tokens_multiplier

    VllmBackend->>VllmBackend: _get_prompt_token_count()
    VllmBackend->>ModelMetadata: "generation_max_tokens_for(prompt_len, multiplier=1.8)"
    ModelMetadata->>ModelMetadata: "sized = int(max_tokens_per_example * multiplier)"
    ModelMetadata->>ModelMetadata: clamp to max_seq_length - prompt_len
    ModelMetadata-->>VllmBackend: max_tokens (safe value)

    TimeseriesBackend->>TimeseriesBackend: _get_prompt_token_count()
    TimeseriesBackend->>ModelMetadata: "generation_max_tokens_for(prompt_len, multiplier=1.8)"
    ModelMetadata-->>TimeseriesBackend: max_tokens (safe value)

    Note over ModelMetadata: Training callback still calls generation_max_tokens_for(prompt_len) with no multiplier → uses legacy 1.2 constant
Loading
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant User
    participant GenerateParameters
    participant VllmBackend
    participant TimeseriesBackend
    participant ModelMetadata

    User->>GenerateParameters: "max_tokens_multiplier=1.8 (or default 1.2)"
    GenerateParameters-->>VllmBackend: config.generation.max_tokens_multiplier
    GenerateParameters-->>TimeseriesBackend: config.generation.max_tokens_multiplier

    VllmBackend->>VllmBackend: _get_prompt_token_count()
    VllmBackend->>ModelMetadata: "generation_max_tokens_for(prompt_len, multiplier=1.8)"
    ModelMetadata->>ModelMetadata: "sized = int(max_tokens_per_example * multiplier)"
    ModelMetadata->>ModelMetadata: clamp to max_seq_length - prompt_len
    ModelMetadata-->>VllmBackend: max_tokens (safe value)

    TimeseriesBackend->>TimeseriesBackend: _get_prompt_token_count()
    TimeseriesBackend->>ModelMetadata: "generation_max_tokens_for(prompt_len, multiplier=1.8)"
    ModelMetadata-->>TimeseriesBackend: max_tokens (safe value)

    Note over ModelMetadata: Training callback still calls generation_max_tokens_for(prompt_len) with no multiplier → uses legacy 1.2 constant
Loading

Reviews (1): Last reviewed commit: "fix: add a mechanism for allowing more m..." | Re-trigger Greptile

@codecov

codecov Bot commented Jul 17, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

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

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

Files:

  • src/nemo_safe_synthesizer/generation/timeseries_backend.py
  • tests/generation/test_vllm_backend.py
  • src/nemo_safe_synthesizer/generation/vllm_backend.py
  • tests/config/test_generate.py
  • src/nemo_safe_synthesizer/llm/metadata.py
  • tests/llm/test_metadata.py
  • src/nemo_safe_synthesizer/config/generate.py
src/**/*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

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

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

Files:

  • src/nemo_safe_synthesizer/generation/timeseries_backend.py
  • src/nemo_safe_synthesizer/generation/vllm_backend.py
  • src/nemo_safe_synthesizer/llm/metadata.py
  • src/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.py
  • src/nemo_safe_synthesizer/generation/vllm_backend.py
  • src/nemo_safe_synthesizer/llm/metadata.py
  • src/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 than main must follow <author>/<description>, optionally including an issue ID or type; branch names must use lowercase alphanumeric characters and hyphens.
Commits merged to main must follow Conventional Commits, using a valid lowercase type and a description of at most 100 characters.

Files:

  • src/nemo_safe_synthesizer/generation/timeseries_backend.py
  • tests/generation/test_vllm_backend.py
  • src/nemo_safe_synthesizer/generation/vllm_backend.py
  • tests/config/test_generate.py
  • src/nemo_safe_synthesizer/llm/metadata.py
  • tests/llm/test_metadata.py
  • src/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.py
  • tests/generation/test_vllm_backend.py
  • src/nemo_safe_synthesizer/generation/vllm_backend.py
  • tests/config/test_generate.py
  • src/nemo_safe_synthesizer/llm/metadata.py
  • tests/llm/test_metadata.py
  • src/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.py
  • tests/generation/test_vllm_backend.py
  • src/nemo_safe_synthesizer/generation/vllm_backend.py
  • tests/config/test_generate.py
  • src/nemo_safe_synthesizer/llm/metadata.py
  • tests/llm/test_metadata.py
  • src/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 type statements or bracketed generic class/function parameters.

Files:

  • src/nemo_safe_synthesizer/generation/timeseries_backend.py
  • tests/generation/test_vllm_backend.py
  • src/nemo_safe_synthesizer/generation/vllm_backend.py
  • tests/config/test_generate.py
  • src/nemo_safe_synthesizer/llm/metadata.py
  • tests/llm/test_metadata.py
  • src/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 mise formatting and checking tasks.

Files:

  • src/nemo_safe_synthesizer/generation/timeseries_backend.py
  • tests/generation/test_vllm_backend.py
  • src/nemo_safe_synthesizer/generation/vllm_backend.py
  • tests/config/test_generate.py
  • src/nemo_safe_synthesizer/llm/metadata.py
  • tests/llm/test_metadata.py
  • src/nemo_safe_synthesizer/config/generate.py
**/*.{py,sh}

📄 CodeRabbit inference engine (CONTRIBUTING.md)

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

Files:

  • src/nemo_safe_synthesizer/generation/timeseries_backend.py
  • tests/generation/test_vllm_backend.py
  • src/nemo_safe_synthesizer/generation/vllm_backend.py
  • tests/config/test_generate.py
  • src/nemo_safe_synthesizer/llm/metadata.py
  • tests/llm/test_metadata.py
  • src/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.py
  • src/nemo_safe_synthesizer/generation/vllm_backend.py
**/test_*.py

📄 CodeRabbit inference engine (AGENTS.md)

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

Files:

  • tests/generation/test_vllm_backend.py
  • tests/config/test_generate.py
  • tests/llm/test_metadata.py
tests/**

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

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

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

Files:

  • tests/generation/test_vllm_backend.py
  • tests/config/test_generate.py
  • tests/llm/test_metadata.py
tests/**/*.py

📄 CodeRabbit inference engine (tests/TESTING.md)

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

Files:

  • tests/generation/test_vllm_backend.py
  • tests/config/test_generate.py
  • tests/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.py
  • tests/config/test_generate.py
  • tests/llm/test_metadata.py
src/**/config/**/*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

Use NSSBaseModel for 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.py
  • tests/config/test_generate.py
  • tests/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!

Comment on lines +244 to +264
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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 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 -S

Repository: 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))
PY

Repository: 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 -S

Repository: 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.py

Repository: 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

Comment on lines +924 to +927
max_tokens=self.model_metadata.generation_max_tokens_for(
self._get_prompt_token_count(),
multiplier=self.config.generation.max_tokens_multiplier,
),

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 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-L1049
  • tests/generation/test_vllm_backend.py#L1079-L1081

Comment on lines +234 to +238
@pytest.mark.parametrize("value", [0, -0.5])
def test_rejects_non_positive(self, value: float) -> None:
"""Non-positive multipliers are rejected by the validator."""
with pytest.raises(ValidationError):
GenerateParameters(max_tokens_multiplier=value)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== tests/config/test_generate.py ==\n'
wc -l tests/config/test_generate.py
sed -n '1,320p' tests/config/test_generate.py

printf '\n== GenerateParameters definition search ==\n'
rg -n "class GenerateParameters|max_tokens_multiplier|ValidationError" -S . --glob '!**/.git/**'

printf '\n== candidate source files ==\n'
# show only likely source files containing the validator
for f in $(rg -l "max_tokens_multiplier" -S . --glob '!tests/**' --glob '!**/.git/**'); do
  echo "--- $f"
  wc -l "$f"
  sed -n '1,260p' "$f"
done

Repository: NVIDIA-NeMo/Safe-Synthesizer

Length of output: 50384


🏁 Script executed:

#!/bin/bash
set -euo pipefail
rg -n "max_tokens_multiplier|class GenerateParameters|ValidationError" -S tests/config/test_generate.py . --glob '!**/.git/**'

Repository: NVIDIA-NeMo/Safe-Synthesizer

Length of output: 13132


🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== generate.py around max_tokens_multiplier ==\n'
sed -n '230,280p' src/nemo_safe_synthesizer/config/generate.py

printf '\n== range_validator definition ==\n'
rg -n "def range_validator|class ValueValidator|max_tokens_multiplier" src/nemo_safe_synthesizer/configurator/validators.py src/nemo_safe_synthesizer/config/generate.py
sed -n '1,260p' src/nemo_safe_synthesizer/configurator/validators.py

Repository: NVIDIA-NeMo/Safe-Synthesizer

Length of output: 9581


Assert the field-specific validation error. Use match="max_tokens_multiplier" so the test can’t pass on an unrelated ValidationError.

Source: Coding guidelines

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:config area:generation area:llm area:tests feature New feature or request test Test-only addition or change

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant