Skip to content

fix: upgrade xgrammar and use new structured gen method - #541

Merged
mckornfield merged 7 commits into
mainfrom
xgrammar-upgrade-structured-tags/mck
Jun 3, 2026
Merged

fix: upgrade xgrammar and use new structured gen method#541
mckornfield merged 7 commits into
mainfrom
xgrammar-upgrade-structured-tags/mck

Conversation

@mckornfield

@mckornfield mckornfield commented May 29, 2026

Copy link
Copy Markdown
Collaborator
  • Also address loading of params tweak
  • benchmark for methods using pretrained paths

Summary

Pre-Review Checklist

Ensure that the following pass:

  • make format && make check or via prek validation.
  • make test passes locally
  • make test-e2e passes locally
  • make 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

Summary by CodeRabbit

Release Notes

  • New Features

    • Added "auto" schema method for structured generation, now the default. Intelligently selects "structural_tag" on compatible backends or falls back to "regex".
    • Added "structural_tag" method option for structured generation with XGrammar support.
  • Documentation

    • Updated configuration reference and guides to reflect new "auto" default and additional schema method options.
  • Chores

    • Added XGrammar dependency (Linux-only) to optional dependencies.

Signed-off-by: mkornfield <mkornfield@nvidia.com>
@mckornfield
mckornfield requested review from a team as code owners May 29, 2026 04:51
@coderabbitai

coderabbitai Bot commented May 29, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review

Walkthrough

This PR adds XGrammar Structural Tag support for structured generation, enabling a new "structural_tag" schema method that resolves automatically on xgrammar-capable backends while falling back to "regex" elsewhere. The implementation includes configuration schema updates, structural-tag format builders, vLLM backend integration, comprehensive validation and testing, benchmarking infrastructure, and user documentation.

Changes

Structural Tag Feature Implementation

Layer / File(s) Summary
Configuration Schema, Type Aliases, & Resolution Logic
src/nemo_safe_synthesizer/config/generate.py
New exported types StructuredGenerationSchemaMethod, StructuredGenerationBackend, ResolvedStructuredGenerationSchemaMethod; constant STRUCTURAL_TAG_COMPATIBLE_BACKENDS; exported functions resolve_structured_generation_schema_method() and structural_tag_backend_error_message(); expanded __all__ to export resolvers.
GenerateParameters Field & Backend Validation
src/nemo_safe_synthesizer/config/generate.py
structured_generation_schema_method field updated to new type with default "auto"; Pydantic model_validator(mode="after") enforces backend compatibility when schema method is explicitly "structural_tag", raising ParameterError for incompatible backends.
XGrammar Dependency Addition
pyproject.toml
Added xgrammar>=0.2.0; sys_platform=='linux' to both cpu and cu129 optional extras.
Structural Tag Format Builder
src/nemo_safe_synthesizer/generation/regex_manager.py
Helper constructors for structural-tag dict fragments (constant, sequence, plus); build_json_structural_tag() function composing format for JSON-schema-constrained newline-delimited records with optional BOS/EOS wrapping and single-sequence/plus-wrapped selection based on config flags.
vLLM Backend Structural Tag Integration
src/nemo_safe_synthesizer/generation/vllm_backend.py
Updated _build_structured_output_params() to resolve schema method and add "structural_tag" handling: validates backend compatibility, logs selection, builds params["structural_tag"] using build_json_structural_tag() with schema and BOS/EOS tokens; updated imports.
E2E Test Backend Update
tests/e2e/test_safe_synthesizer.py
DP e2e test structured_generation_backend changed from "outlines" to "xgrammar".

Configuration, Generation, & Integration Tests

Layer / File(s) Summary
Schema Method Resolution Unit Tests
tests/config/test_generate.py
Tests for resolve_structured_generation_schema_method(): "auto" resolves to "structural_tag" for xgrammar-capable backends and "regex" for others; explicit methods pass through unchanged.
GenerateParameters Validation Unit Tests
tests/config/test_generate.py
Validation tests for GenerateParameters structural-tag constraints: compatible backends succeed; "structural_tag" with incompatible backends raises ValidationError; "auto" with any backend succeeds; default is "auto"; validation skipped when structured generation disabled; SafeSynthesizerParameters.from_params() rejects incompatible backends.
Structural Tag Test Helper & Unit Tests
tests/generation/structural_tag_helpers.py, tests/generation/test_regex_manager.py
Test helper module structural_tag_accepts_text() using xgrammar to validate text acceptance; unit tests for build_json_structural_tag() verifying JSON structure for four config scenarios (JSONL, single-sequence, grouped, grouped+single-sequence).
Structural Tag XGrammar Acceptance Tests
tests/generation/test_regex_manager.py
Parametrized round-trip tests using structural_tag_accepts_text() validating four training-text shapes; negative test rejecting invalid JSONL content.
vLLM Backend Structural Tag Tests
tests/generation/test_vllm_backend.py
Fixtures for "auto" and "structural_tag" schema methods; "structural_tag" method test verifying build_json_structural_tag invocation with BOS/EOS tokens; "auto" resolution tests verifying routing to structural-tag vs regex; negative test asserting "structural_tag" with non-xgrammar backends raises ParameterError.

Documentation & Benchmarking

Layer / File(s) Summary
Configuration & Running Guide Documentation
docs/user-guide/configuration.md, docs/user-guide/running.md
Updated default structured_generation_schema_method to "auto" and expanded allowed values to include "auto" and "structural_tag"; added option descriptions for backend-dependent resolution and XGrammar Structural Tag behavior.
Structured Generation Benchmark Test Module
tests/benchmarks/test_generation_structured_methods.py
Benchmark module with GenerationMethod and GenerationBenchmarkResult dataclasses; environment-driven method selection; CUDA gating; configurable input path resolution relative to pytestconfig.rootpath; CLI command building wiring structured-generation flags; subprocess execution with timeout, merged logging, and wall-clock duration; CSV record counting; per-method JSON summary writing; parametrized pytest test asserting at least one output record.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes


Suggested reviewers

  • binaryaaron
  • kendrickb-nvidia
🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 51.56% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title concisely summarizes the main change: upgrading xgrammar and implementing the new structured generation method (Structural Tags).
Linked Issues check ✅ Passed The PR implements all coding requirements from issue #540: xgrammar2 upgrade with Structural Tag support, backend resolution logic, new schema method type, validation, and benchmarking.
Out of Scope Changes check ✅ Passed All changes are directly aligned with #540 objectives: xgrammar2 integration, Structural Tags, schema method resolution, validation, documentation updates, and benchmarks.
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch xgrammar-upgrade-structured-tags/mck

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

@coderabbitai coderabbitai Bot added docs Documentation-only change feature New feature or request test Test-only addition or change chore Maintenance not tied to a user-visible change labels May 29, 2026
@greptile-apps

greptile-apps Bot commented May 29, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR upgrades xgrammar from 0.1.32 to 0.2.1 and introduces a new "structural_tag" structured-generation method alongside an "auto" default that selects structural_tag on xgrammar-capable backends and falls back to regex elsewhere, preserving existing behavior for outlines/guidance configs.

  • config/generate.py: Adds resolve_structured_generation_schema_method, structural_tag_backend_error_message, and a Pydantic model_validator that rejects structural_tag + incompatible backend combinations at config-load time.
  • regex_manager.py: Adds build_json_structural_tag covering all four output shapes (plain JSONL, single-sequence, grouped, grouped single-sequence), mirroring the existing build_json_based_regex contract.
  • vllm_backend.py: Resolves auto at call time and dispatches to the new structural_tag branch; includes a runtime defense check for backends that bypass Pydantic validation.

Confidence Score: 5/5

Safe to merge — the auto default preserves regex for existing outlines/guidance/lm-format-enforcer configs and the new structural_tag path is thoroughly tested at both the unit and round-trip xgrammar levels.

The auto resolver correctly degrades to regex for any non-xgrammar backend, so existing user configs are not silently broken. The new structural_tag path has shape tests, XGrammar GrammarMatcher acceptance/rejection tests, and vLLM backend unit tests. The Pydantic model validator correctly wraps ParameterError (a ValueError subclass) into ValidationError, confirmed by the test suite. No logic errors found in the four-shape structural tag builder.

No files require special attention.

Important Files Changed

Filename Overview
src/nemo_safe_synthesizer/config/generate.py Adds auto/structural_tag schema method types, resolver function, backend compatibility check, and Pydantic model validator; default changes from regex to auto
src/nemo_safe_synthesizer/generation/regex_manager.py Adds build_json_structural_tag function that composes XGrammar Structural Tag JSON for all four output shapes (plain JSONL, single-sequence, grouped, grouped single-sequence)
src/nemo_safe_synthesizer/generation/vllm_backend.py Resolves auto schema method at call time and adds structural_tag branch to _build_structured_output_params, including a defense-in-depth backend compatibility check
tests/config/test_generate.py New unit tests covering auto resolution, structural-tag backend validation, and SafeSynthesizerParameters.from_params rejection of incompatible configs
tests/generation/test_regex_manager.py Adds structural-tag shape tests (JSON match + XGrammar round-trip acceptance/rejection) for all four output configurations
tests/generation/test_vllm_backend.py Adds fixtures and tests for structural_tag and auto schema methods, covering both xgrammar and non-xgrammar backend paths
tests/benchmarks/test_generation_structured_methods.py New benchmark comparing unstructured, regex, json_schema, structural_tag, and three alternative-backend methods; requires pre-trained run path and GPU, skips when paths are absent
tests/generation/structural_tag_helpers.py Helper module that compiles a structural-tag string with XGrammar and checks whether a text is fully accepted; mirrors regex re.fullmatch round-trip tests
tests/e2e/test_safe_synthesizer.py Switches the DP e2e test backend from outlines to xgrammar, aligning it with the new default structured-generation method
pyproject.toml Adds xgrammar>=0.2.0 (Linux-only) to both cpu and cu129 optional extras; lock file pins to 0.2.1

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["VllmBackend._build_structured_output_params()"] --> B{"use_structured_generation?"}
    B -- No --> C["return None"]
    B -- Yes --> D["resolve_structured_generation_schema_method(schema_method, backend)"]
    D --> E{"schema_method == 'auto'?"}
    E -- No --> F["return schema_method as-is"]
    E -- Yes --> G{"backend in STRUCTURAL_TAG_COMPATIBLE_BACKENDS?"}
    G -- Yes --> H["return 'structural_tag'"]
    G -- No --> I["return 'regex'"]
    F --> J["resolved_method"]
    H --> J
    I --> J
    J --> K{"resolved_method?"}
    K -- regex --> L["build_json_based_regex()"]
    K -- json_schema --> M["params['json'] = schema"]
    K -- structural_tag --> N{"structural_tag_backend_error_message(backend)?"}
    N -- error --> O["raise ParameterError"]
    N -- None --> P["build_json_structural_tag()"]
    L --> Q["return StructuredOutputsParams(**params)"]
    M --> Q
    P --> Q
Loading

Reviews (6): Last reviewed commit: "chore: address latest CodeRabbit review ..." | Re-trigger Greptile

Comment thread tests/e2e/test_safe_synthesizer.py Outdated
Comment thread tests/benchmarks/test_generation_structured_methods.py Outdated
Comment thread tests/benchmarks/test_generation_structured_methods.py Outdated
@codecov

codecov Bot commented May 29, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 81.33333% with 56 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
...s/benchmarks/test_generation_structured_methods.py 50.44% 56 Missing ⚠️

📢 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: 2


ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: ec9a86c4-abdc-4ae2-b966-48a29404293a

📥 Commits

Reviewing files that changed from the base of the PR and between 3d9f82a and 15abb24.

⛔ Files ignored due to path filters (1)
  • uv.lock is excluded by !**/*.lock, !uv.lock
📒 Files selected for processing (14)
  • docs/user-guide/configuration.md
  • docs/user-guide/running.md
  • pyproject.toml
  • src/nemo_safe_synthesizer/cli/run.py
  • src/nemo_safe_synthesizer/config/generate.py
  • src/nemo_safe_synthesizer/generation/regex_manager.py
  • src/nemo_safe_synthesizer/generation/vllm_backend.py
  • src/nemo_safe_synthesizer/sdk/library_builder.py
  • tests/benchmarks/test_generation_structured_methods.py
  • tests/cli/test_run.py
  • tests/e2e/test_safe_synthesizer.py
  • tests/generation/test_regex_manager.py
  • tests/generation/test_vllm_backend.py
  • tests/sdk/test_process_data.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.12)
  • GitHub Check: Unit Tests (3.13)
  • GitHub Check: Unit Tests (3.11)
  • GitHub Check: Smoke Tests
🧰 Additional context used
📓 Path-based instructions (16)
**/*.{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/cli/run.py
  • src/nemo_safe_synthesizer/generation/vllm_backend.py
  • tests/e2e/test_safe_synthesizer.py
  • tests/sdk/test_process_data.py
  • src/nemo_safe_synthesizer/config/generate.py
  • tests/generation/test_vllm_backend.py
  • src/nemo_safe_synthesizer/generation/regex_manager.py
  • docs/user-guide/running.md
  • tests/generation/test_regex_manager.py
  • tests/cli/test_run.py
  • docs/user-guide/configuration.md
  • tests/benchmarks/test_generation_structured_methods.py
  • src/nemo_safe_synthesizer/sdk/library_builder.py
**/*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

**/*.py: Use American English spelling: 'initialize' not 'initialise', 'recognize' not 'recognise', 'color' not 'colour'.
Use observability.get_logger(__name__) for logging, never logging.getLogger() or structlog.get_logger() directly.
Use category loggers: .runtime for internals, .user for progress/results, .system for system events.
Never use print() for operational output. Use click.echo() for CLI output or sys.stdout.write() for raw output in tools.
Use extra={} 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.
Use NSSBaseModel for config/parameter models in config/ which define user-facing configuration. Use raw BaseModel or module-specific bases for data transfer objects and internal structures.
Use BaseSettings for env/CLI settings. Prefer AliasChoices on individual fields when a field needs to respond to both its Python name and an env var name.
Include Field(description=...) for Pydantic model fields as the canonical field docstring for API documentation and CLI help text.
Use assignment-style type = Field(default=..., description="...") as the default for Pydantic model fields because type checkers understand default, default_factory, and alias in assignment style.
Use Annotated only when the field carries additional metadata beyond Field() -- ValueValidator, AutoParam, DependsOnValidator, reusable constrained type aliases, nested-type constraints, or discriminated unions.
Put defaults as bare assignment (= value), not inside Field(default=...), when using Annotated. Exception: use assignment-style Field(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/cli/run.py
  • src/nemo_safe_synthesizer/generation/vllm_backend.py
  • tests/e2e/test_safe_synthesizer.py
  • tests/sdk/test_process_data.py
  • src/nemo_safe_synthesizer/config/generate.py
  • tests/generation/test_vllm_backend.py
  • src/nemo_safe_synthesizer/generation/regex_manager.py
  • tests/generation/test_regex_manager.py
  • tests/cli/test_run.py
  • tests/benchmarks/test_generation_structured_methods.py
  • src/nemo_safe_synthesizer/sdk/library_builder.py
src/**/*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

src/**/*.py: Use relative imports in src/ (e.g., from ..observability import get_logger).
Do not use print() statements in library code. Use get_logger(__name__) from observability.py or click.echo() for CLI.
Do not use assert for validation in library code. Use if/raise for input validation. assert statements can be stripped by -O and must never guard correctness.

Files:

  • src/nemo_safe_synthesizer/cli/run.py
  • src/nemo_safe_synthesizer/generation/vllm_backend.py
  • src/nemo_safe_synthesizer/config/generate.py
  • src/nemo_safe_synthesizer/generation/regex_manager.py
  • src/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/cli/run.py
  • src/nemo_safe_synthesizer/generation/vllm_backend.py
  • src/nemo_safe_synthesizer/config/generate.py
  • src/nemo_safe_synthesizer/generation/regex_manager.py
  • src/nemo_safe_synthesizer/sdk/library_builder.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. The make format command handles this automatically.

Files:

  • src/nemo_safe_synthesizer/cli/run.py
  • src/nemo_safe_synthesizer/generation/vllm_backend.py
  • tests/e2e/test_safe_synthesizer.py
  • tests/sdk/test_process_data.py
  • src/nemo_safe_synthesizer/config/generate.py
  • tests/generation/test_vllm_backend.py
  • src/nemo_safe_synthesizer/generation/regex_manager.py
  • tests/generation/test_regex_manager.py
  • tests/cli/test_run.py
  • tests/benchmarks/test_generation_structured_methods.py
  • src/nemo_safe_synthesizer/sdk/library_builder.py
**/*

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

**/*: Include a newline at the end of all files, never trailing whitespace. This is enforced by pre-commit.
Use line length of 120 characters for code, comments, and docstrings (configured in ruff.toml).

Files:

  • src/nemo_safe_synthesizer/cli/run.py
  • src/nemo_safe_synthesizer/generation/vllm_backend.py
  • tests/e2e/test_safe_synthesizer.py
  • tests/sdk/test_process_data.py
  • src/nemo_safe_synthesizer/config/generate.py
  • tests/generation/test_vllm_backend.py
  • src/nemo_safe_synthesizer/generation/regex_manager.py
  • docs/user-guide/running.md
  • tests/generation/test_regex_manager.py
  • pyproject.toml
  • tests/cli/test_run.py
  • docs/user-guide/configuration.md
  • tests/benchmarks/test_generation_structured_methods.py
  • src/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/cli/run.py
  • src/nemo_safe_synthesizer/generation/vllm_backend.py
  • tests/e2e/test_safe_synthesizer.py
  • tests/sdk/test_process_data.py
  • src/nemo_safe_synthesizer/config/generate.py
  • tests/generation/test_vllm_backend.py
  • src/nemo_safe_synthesizer/generation/regex_manager.py
  • docs/user-guide/running.md
  • tests/generation/test_regex_manager.py
  • pyproject.toml
  • tests/cli/test_run.py
  • docs/user-guide/configuration.md
  • tests/benchmarks/test_generation_structured_methods.py
  • src/nemo_safe_synthesizer/sdk/library_builder.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
  • src/nemo_safe_synthesizer/generation/regex_manager.py
tests/**

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

Tests should mirror the src/ directory structure in tests/

Files:

  • tests/e2e/test_safe_synthesizer.py
  • tests/sdk/test_process_data.py
  • tests/generation/test_vllm_backend.py
  • tests/generation/test_regex_manager.py
  • tests/cli/test_run.py
  • tests/benchmarks/test_generation_structured_methods.py
tests/**/*.py

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

Tests in tests/e2e/ should be auto-marked with e2e marker, tests in tests/smoke/ with smoke marker, others with unit marker

tests/**/*.py: Use absolute imports in tests/ (e.g., from nemo_safe_synthesizer.observability import get_logger).
Use fixture_ prefix convention for fixtures for grep-ability and to separate fixtures from test functions. 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.
Use bare assert as the primary assertion style; pytest.raises() with match= for exceptions; pytest.approx() for floating-point comparisons.
Mark CUDA-dependent tests with @pytest.mark.e2e, @pytest.mark.smoke, or @pytest.mark.requires_gpu.
Mock only external boundaries, not internal implementation details.
Ensure test isolation: no shared mutable state or execution-order dependencies between tests. If something must be run first before executing a test, include it in the test or a fixture.
Use @pytest.mark.parametrize for testing multiple input combinations rather than copy-pasting similar tests.

Organize tests using pytest following the structure in tests/TESTING.md with support for unit tests, smoke tests, and end-to-end tests

Files:

  • tests/e2e/test_safe_synthesizer.py
  • tests/sdk/test_process_data.py
  • tests/generation/test_vllm_backend.py
  • tests/generation/test_regex_manager.py
  • tests/cli/test_run.py
  • tests/benchmarks/test_generation_structured_methods.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/e2e/test_safe_synthesizer.py
  • tests/sdk/test_process_data.py
  • tests/generation/test_vllm_backend.py
  • tests/generation/test_regex_manager.py
  • tests/cli/test_run.py
  • tests/benchmarks/test_generation_structured_methods.py
**/*test*.py

📄 CodeRabbit inference engine (AGENTS.md)

The unit_test marker is deprecated; use unit instead

Files:

  • tests/e2e/test_safe_synthesizer.py
  • tests/sdk/test_process_data.py
  • tests/generation/test_vllm_backend.py
  • tests/generation/test_regex_manager.py
  • tests/cli/test_run.py
  • tests/benchmarks/test_generation_structured_methods.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
**/*.{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.md
  • docs/user-guide/configuration.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 .md files 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.md
  • docs/user-guide/configuration.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 the nav: section of mkdocs.yml for 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.md
  • docs/user-guide/configuration.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/running.md
  • docs/user-guide/configuration.md
pyproject.toml

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

Package metadata, dependencies, and extras (cpu/cu128/engine) should be configured in pyproject.toml

Order sections in pyproject.toml as: [project], [dependency-groups], [project.optional-dependencies], [tool.uv], [build-system], [tool.*].

Python version support must be 3.11–3.13 as specified in .python-version and pyproject.toml; Python 3.14+ is not supported

Files:

  • pyproject.toml

⚙️ CodeRabbit configuration file

Treat pyproject.toml as high-risk. Check package metadata, uv indexes, dependency groups, optional extras, Python version bounds, hatch config, ty config, script entry points, dependency consistency, and whether changes require regenerating uv.lock.

Files:

  • pyproject.toml
**/*.toml

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

**/*.toml: Use spaces around = for key-value pairs in TOML files.
Use # comment format for comments in TOML files; use inline comments for dependency pins.

Files:

  • pyproject.toml
🧠 Learnings (1)
📚 Learning: 2026-05-27T22:20:37.354Z
Learnt from: kendrickb-nvidia
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 520
File: tests/generation/test_vllm_backend.py:556-587
Timestamp: 2026-05-27T22:20:37.354Z
Learning: In NVIDIA-NeMo/Safe-Synthesizer, `tests/conftest.py`’s `pytest_collection_modifyitems` hook applies pytest category markers automatically based on each test file’s path: tests under `/e2e/` get `pytest.mark.e2e`, tests under `/smoke/` get `pytest.mark.smoke`, and all other tests get `pytest.mark.unit`. Therefore, when reviewing pytest tests outside `tests/e2e/` and `tests/smoke/`, do not flag missing explicit `pytest.mark.unit` decorators on test classes/functions as an issue (the hook will add them during collection). If a new test directory/category is introduced, ensure the hook is updated so it’s categorized correctly.

Applied to files:

  • tests/e2e/test_safe_synthesizer.py
  • tests/sdk/test_process_data.py
  • tests/generation/test_vllm_backend.py
  • tests/generation/test_regex_manager.py
  • tests/cli/test_run.py
  • tests/benchmarks/test_generation_structured_methods.py
🪛 Ruff (0.15.14)
tests/benchmarks/test_generation_structured_methods.py

[error] 208-208: subprocess call: check for execution of untrusted input

(S603)

🔇 Additional comments (12)
tests/benchmarks/test_generation_structured_methods.py (5)

206-221: Static analysis S603 is a false positive here.

The subprocess call uses sys.executable (the current Python interpreter) and constructs the command from controlled test fixtures and environment variables. This is standard practice for integration/benchmark tests that invoke the project's CLI. The warning about "untrusted input" doesn't apply to this controlled test scenario.


1-39: LGTM!


42-130: LGTM!


132-166: LGTM!


169-248: LGTM!

tests/generation/test_vllm_backend.py (1)

117-124: LGTM!

Also applies to: 220-248

src/nemo_safe_synthesizer/cli/run.py (1)

571-571: LGTM!

tests/cli/test_run.py (1)

623-650: LGTM!

tests/sdk/test_process_data.py (1)

442-478: LGTM!

docs/user-guide/configuration.md (1)

163-163: LGTM!

docs/user-guide/running.md (1)

886-891: LGTM!

tests/e2e/test_safe_synthesizer.py (1)

75-75: ⚡ Quick win

Audit pretrained_model fixture/parameter usage in tests/e2e/test_safe_synthesizer.py: Locate where pretrained_model is defined (fixture or imported symbol) and confirm the e2e test actually uses it; if unused, remove the fixture/parameter to avoid unnecessary setup/side effects.

Comment thread src/nemo_safe_synthesizer/generation/vllm_backend.py Outdated
Comment thread tests/generation/test_regex_manager.py
@mckornfield mckornfield self-assigned this May 29, 2026
@mckornfield
mckornfield requested a review from binaryaaron May 29, 2026 18:35
Comment thread src/nemo_safe_synthesizer/generation/vllm_backend.py Outdated
Signed-off-by: mkornfield <mkornfield@nvidia.com>
Comment thread tests/benchmarks/test_generation_structured_methods.py

@kendrickb-nvidia kendrickb-nvidia left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

A few comments, but looks promising.

Comment thread src/nemo_safe_synthesizer/generation/vllm_backend.py Outdated
Comment thread src/nemo_safe_synthesizer/sdk/library_builder.py Outdated
Comment thread tests/benchmarks/test_generation_structured_methods.py
Comment thread tests/benchmarks/test_generation_structured_methods.py Outdated
Comment thread tests/e2e/test_safe_synthesizer.py Outdated
Comment thread tests/generation/test_regex_manager.py
Signed-off-by: mkornfield <mkornfield@nvidia.com>
Comment thread src/nemo_safe_synthesizer/sdk/library_builder.py Outdated
Comment thread tests/generation/test_regex_manager.py
Signed-off-by: mkornfield <mkornfield@nvidia.com>
nina-xu
nina-xu previously approved these changes Jun 1, 2026

@nina-xu nina-xu 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.

lgtm! agree with the other comments

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

@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: 2


ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: e20f69c8-4514-401c-99c2-b93c4fd625bc

📥 Commits

Reviewing files that changed from the base of the PR and between 0bccefe and fcd382e.

📒 Files selected for processing (10)
  • src/nemo_safe_synthesizer/config/generate.py
  • src/nemo_safe_synthesizer/generation/vllm_backend.py
  • src/nemo_safe_synthesizer/preflight/__init__.py
  • src/nemo_safe_synthesizer/preflight/checks/__init__.py
  • src/nemo_safe_synthesizer/preflight/checks/generation.py
  • tests/benchmarks/test_generation_structured_methods.py
  • tests/config/test_generate.py
  • tests/generation/structural_tag_helpers.py
  • tests/generation/test_regex_manager.py
  • tests/preflight/test_preflight.py
🚧 Files skipped from review as they are similar to previous changes (2)
  • src/nemo_safe_synthesizer/generation/vllm_backend.py
  • tests/benchmarks/test_generation_structured_methods.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.12)
  • GitHub Check: Unit Tests (3.11)
  • GitHub Check: Unit Tests (3.13)
  • GitHub Check: Smoke Tests
🧰 Additional context used
📓 Path-based instructions (10)
**/*.{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/preflight/checks/generation.py
  • src/nemo_safe_synthesizer/preflight/checks/__init__.py
  • src/nemo_safe_synthesizer/preflight/__init__.py
  • tests/preflight/test_preflight.py
  • src/nemo_safe_synthesizer/config/generate.py
  • tests/generation/structural_tag_helpers.py
  • tests/config/test_generate.py
  • tests/generation/test_regex_manager.py
**/*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

**/*.py: Use American English spelling: 'initialize' not 'initialise', 'recognize' not 'recognise', 'color' not 'colour'.
Use observability.get_logger(__name__) for logging, never logging.getLogger() or structlog.get_logger() directly.
Use category loggers: .runtime for internals, .user for progress/results, .system for system events.
Never use print() for operational output. Use click.echo() for CLI output or sys.stdout.write() for raw output in tools.
Use extra={} 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.
Use NSSBaseModel for config/parameter models in config/ which define user-facing configuration. Use raw BaseModel or module-specific bases for data transfer objects and internal structures.
Use BaseSettings for env/CLI settings. Prefer AliasChoices on individual fields when a field needs to respond to both its Python name and an env var name.
Include Field(description=...) for Pydantic model fields as the canonical field docstring for API documentation and CLI help text.
Use assignment-style type = Field(default=..., description="...") as the default for Pydantic model fields because type checkers understand default, default_factory, and alias in assignment style.
Use Annotated only when the field carries additional metadata beyond Field() -- ValueValidator, AutoParam, DependsOnValidator, reusable constrained type aliases, nested-type constraints, or discriminated unions.
Put defaults as bare assignment (= value), not inside Field(default=...), when using Annotated. Exception: use assignment-style Field(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/preflight/checks/generation.py
  • src/nemo_safe_synthesizer/preflight/checks/__init__.py
  • src/nemo_safe_synthesizer/preflight/__init__.py
  • tests/preflight/test_preflight.py
  • src/nemo_safe_synthesizer/config/generate.py
  • tests/generation/structural_tag_helpers.py
  • tests/config/test_generate.py
  • tests/generation/test_regex_manager.py
src/**/*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

src/**/*.py: Use relative imports in src/ (e.g., from ..observability import get_logger).
Do not use print() statements in library code. Use get_logger(__name__) from observability.py or click.echo() for CLI.
Do not use assert for validation in library code. Use if/raise for input validation. assert statements can be stripped by -O and must never guard correctness.

Files:

  • src/nemo_safe_synthesizer/preflight/checks/generation.py
  • src/nemo_safe_synthesizer/preflight/checks/__init__.py
  • src/nemo_safe_synthesizer/preflight/__init__.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/preflight/checks/generation.py
  • src/nemo_safe_synthesizer/preflight/checks/__init__.py
  • src/nemo_safe_synthesizer/preflight/__init__.py
  • src/nemo_safe_synthesizer/config/generate.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. The make format command handles this automatically.

Files:

  • src/nemo_safe_synthesizer/preflight/checks/generation.py
  • src/nemo_safe_synthesizer/preflight/checks/__init__.py
  • src/nemo_safe_synthesizer/preflight/__init__.py
  • tests/preflight/test_preflight.py
  • src/nemo_safe_synthesizer/config/generate.py
  • tests/generation/structural_tag_helpers.py
  • tests/config/test_generate.py
  • tests/generation/test_regex_manager.py
**/*

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

**/*: Include a newline at the end of all files, never trailing whitespace. This is enforced by pre-commit.
Use line length of 120 characters for code, comments, and docstrings (configured in ruff.toml).

Files:

  • src/nemo_safe_synthesizer/preflight/checks/generation.py
  • src/nemo_safe_synthesizer/preflight/checks/__init__.py
  • src/nemo_safe_synthesizer/preflight/__init__.py
  • tests/preflight/test_preflight.py
  • src/nemo_safe_synthesizer/config/generate.py
  • tests/generation/structural_tag_helpers.py
  • tests/config/test_generate.py
  • tests/generation/test_regex_manager.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/preflight/checks/generation.py
  • src/nemo_safe_synthesizer/preflight/checks/__init__.py
  • src/nemo_safe_synthesizer/preflight/__init__.py
  • tests/preflight/test_preflight.py
  • src/nemo_safe_synthesizer/config/generate.py
  • tests/generation/structural_tag_helpers.py
  • tests/config/test_generate.py
  • tests/generation/test_regex_manager.py
src/**/__init__.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

Include __init__.py files in every directory under src/ that contains Python files, even if empty, to ensure the directory is recognized as a Python package.

Files:

  • src/nemo_safe_synthesizer/preflight/checks/__init__.py
  • src/nemo_safe_synthesizer/preflight/__init__.py
tests/**

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

Tests should mirror the src/ directory structure in tests/

Files:

  • tests/preflight/test_preflight.py
  • tests/generation/structural_tag_helpers.py
  • tests/config/test_generate.py
  • tests/generation/test_regex_manager.py
tests/**/*.py

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

Tests in tests/e2e/ should be auto-marked with e2e marker, tests in tests/smoke/ with smoke marker, others with unit marker

tests/**/*.py: Use absolute imports in tests/ (e.g., from nemo_safe_synthesizer.observability import get_logger).
Use fixture_ prefix convention for fixtures for grep-ability and to separate fixtures from test functions. 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.
Use bare assert as the primary assertion style; pytest.raises() with match= for exceptions; pytest.approx() for floating-point comparisons.
Mark CUDA-dependent tests with @pytest.mark.e2e, @pytest.mark.smoke, or @pytest.mark.requires_gpu.
Mock only external boundaries, not internal implementation details.
Ensure test isolation: no shared mutable state or execution-order dependencies between tests. If something must be run first before executing a test, include it in the test or a fixture.
Use @pytest.mark.parametrize for testing multiple input combinations rather than copy-pasting similar tests.

Organize tests using pytest following the structure in tests/TESTING.md with support for unit tests, smoke tests, and end-to-end tests

Files:

  • tests/preflight/test_preflight.py
  • tests/generation/structural_tag_helpers.py
  • tests/config/test_generate.py
  • tests/generation/test_regex_manager.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/preflight/test_preflight.py
  • tests/generation/structural_tag_helpers.py
  • tests/config/test_generate.py
  • tests/generation/test_regex_manager.py
**/*test*.py

📄 CodeRabbit inference engine (AGENTS.md)

The unit_test marker is deprecated; use unit instead

Files:

  • tests/preflight/test_preflight.py
  • tests/config/test_generate.py
  • tests/generation/test_regex_manager.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 (1)
📚 Learning: 2026-05-27T22:20:37.354Z
Learnt from: kendrickb-nvidia
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 520
File: tests/generation/test_vllm_backend.py:556-587
Timestamp: 2026-05-27T22:20:37.354Z
Learning: In NVIDIA-NeMo/Safe-Synthesizer, `tests/conftest.py`’s `pytest_collection_modifyitems` hook applies pytest category markers automatically based on each test file’s path: tests under `/e2e/` get `pytest.mark.e2e`, tests under `/smoke/` get `pytest.mark.smoke`, and all other tests get `pytest.mark.unit`. Therefore, when reviewing pytest tests outside `tests/e2e/` and `tests/smoke/`, do not flag missing explicit `pytest.mark.unit` decorators on test classes/functions as an issue (the hook will add them during collection). If a new test directory/category is introduced, ensure the hook is updated so it’s categorized correctly.

Applied to files:

  • tests/preflight/test_preflight.py
  • tests/generation/structural_tag_helpers.py
  • tests/config/test_generate.py
  • tests/generation/test_regex_manager.py
🪛 Ruff (0.15.15)
src/nemo_safe_synthesizer/config/generate.py

[warning] 25-30: __all__ is not sorted

Apply an isort-style sorting to __all__

(RUF022)

🔇 Additional comments (12)
tests/preflight/test_preflight.py (1)

42-42: LGTM!

Also applies to: 265-309

tests/config/test_generate.py (1)

27-40: ⚡ Quick win

Pydantic wrapping of ParameterError is appropriate

ParameterError is defined as class ParameterError(UserError, ValueError), so Pydantic v2 will wrap it as pydantic.ValidationError; these tests’ pytest.raises(ValidationError, ...) expectations align with the error hierarchy.

src/nemo_safe_synthesizer/config/generate.py (4)

221-229: LGTM!


33-46: LGTM!


23-23: LGTM!


174-185: LGTM!

src/nemo_safe_synthesizer/preflight/checks/generation.py (1)

1-33: LGTM!

src/nemo_safe_synthesizer/preflight/__init__.py (1)

27-27: LGTM!

Also applies to: 81-81

src/nemo_safe_synthesizer/preflight/checks/__init__.py (1)

33-33: LGTM!

Also applies to: 47-47, 64-64

tests/generation/test_regex_manager.py (2)

562-611: LGTM!


614-626: LGTM!

tests/generation/structural_tag_helpers.py (1)

31-35: xgrammar 0.2.0 Structural Tag round-trip API matches the helper’s usage

GrammarCompiler.compile_structural_tag() in xgrammar 0.2.0 accepts the structural_tag argument as a JSON string (as well as a dict/StructuralTag), and GrammarMatcher provides accept_string() plus is_completed() for the exact matcher round-trip pattern used here.

Comment thread src/nemo_safe_synthesizer/config/generate.py
Comment thread tests/generation/test_regex_manager.py Outdated
@greptile-apps

greptile-apps Bot commented Jun 1, 2026

Copy link
Copy Markdown
Contributor

Want your agent to iterate on Greptile's feedback? Try greploops.

Comment thread src/nemo_safe_synthesizer/config/generate.py Outdated
Signed-off-by: mkornfield <mkornfield@nvidia.com>

@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: 2

🧹 Nitpick comments (1)
tests/config/test_generate.py (1)

70-72: ⚡ Quick win

Strengthen the assertion to match the precise validation message.

match="outlines" only checks that the backend value is echoed somewhere in the error, so this test can pass even if the rejection comes from an unrelated validation path. The sibling test at Line 52 already asserts the precise condition message; mirroring it here makes the regression coverage meaningful.

♻️ Proposed change
     def test_from_params_rejects_incompatible_backend(self) -> None:
-        with pytest.raises(ValidationError, match="outlines"):
+        with pytest.raises(ValidationError, match="requires `structured_generation_backend`"):
             SafeSynthesizerParameters.from_params(**self._generation_kwargs(backend="outlines"))

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: f20564a9-1a16-45ac-ad01-7196ec409e92

📥 Commits

Reviewing files that changed from the base of the PR and between fcd382e and 1996bf6.

📒 Files selected for processing (7)
  • docs/user-guide/configuration.md
  • docs/user-guide/running.md
  • src/nemo_safe_synthesizer/config/generate.py
  • src/nemo_safe_synthesizer/generation/vllm_backend.py
  • tests/config/test_generate.py
  • tests/generation/test_regex_manager.py
  • tests/generation/test_vllm_backend.py
💤 Files with no reviewable changes (1)
  • tests/generation/test_regex_manager.py
✅ Files skipped from review due to trivial changes (2)
  • docs/user-guide/configuration.md
  • docs/user-guide/running.md
🚧 Files skipped from review as they are similar to previous changes (1)
  • src/nemo_safe_synthesizer/generation/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). (4)
  • GitHub Check: Unit Tests (3.12)
  • GitHub Check: Unit Tests (3.13)
  • GitHub Check: Unit Tests (3.11)
  • GitHub Check: Smoke Tests
🧰 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:

  • src/nemo_safe_synthesizer/config/generate.py
  • tests/generation/test_vllm_backend.py
  • tests/config/test_generate.py
**/*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

**/*.py: Use American English spelling: 'initialize' not 'initialise', 'recognize' not 'recognise', 'color' not 'colour'.
Use observability.get_logger(__name__) for logging, never logging.getLogger() or structlog.get_logger() directly.
Use category loggers: .runtime for internals, .user for progress/results, .system for system events.
Never use print() for operational output. Use click.echo() for CLI output or sys.stdout.write() for raw output in tools.
Use extra={} 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.
Use NSSBaseModel for config/parameter models in config/ which define user-facing configuration. Use raw BaseModel or module-specific bases for data transfer objects and internal structures.
Use BaseSettings for env/CLI settings. Prefer AliasChoices on individual fields when a field needs to respond to both its Python name and an env var name.
Include Field(description=...) for Pydantic model fields as the canonical field docstring for API documentation and CLI help text.
Use assignment-style type = Field(default=..., description="...") as the default for Pydantic model fields because type checkers understand default, default_factory, and alias in assignment style.
Use Annotated only when the field carries additional metadata beyond Field() -- ValueValidator, AutoParam, DependsOnValidator, reusable constrained type aliases, nested-type constraints, or discriminated unions.
Put defaults as bare assignment (= value), not inside Field(default=...), when using Annotated. Exception: use assignment-style Field(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/config/generate.py
  • tests/generation/test_vllm_backend.py
  • tests/config/test_generate.py
src/**/*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

src/**/*.py: Use relative imports in src/ (e.g., from ..observability import get_logger).
Do not use print() statements in library code. Use get_logger(__name__) from observability.py or click.echo() for CLI.
Do not use assert for validation in library code. Use if/raise for input validation. assert statements can be stripped by -O and must never guard correctness.

Files:

  • 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/config/generate.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. The make format command handles this automatically.

Files:

  • src/nemo_safe_synthesizer/config/generate.py
  • tests/generation/test_vllm_backend.py
  • tests/config/test_generate.py
**/*

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

**/*: Include a newline at the end of all files, never trailing whitespace. This is enforced by pre-commit.
Use line length of 120 characters for code, comments, and docstrings (configured in ruff.toml).

Files:

  • src/nemo_safe_synthesizer/config/generate.py
  • tests/generation/test_vllm_backend.py
  • tests/config/test_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/config/generate.py
  • tests/generation/test_vllm_backend.py
  • tests/config/test_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
tests/**

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

Tests should mirror the src/ directory structure in tests/

Files:

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

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

Tests in tests/e2e/ should be auto-marked with e2e marker, tests in tests/smoke/ with smoke marker, others with unit marker

tests/**/*.py: Use absolute imports in tests/ (e.g., from nemo_safe_synthesizer.observability import get_logger).
Use fixture_ prefix convention for fixtures for grep-ability and to separate fixtures from test functions. 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.
Use bare assert as the primary assertion style; pytest.raises() with match= for exceptions; pytest.approx() for floating-point comparisons.
Mark CUDA-dependent tests with @pytest.mark.e2e, @pytest.mark.smoke, or @pytest.mark.requires_gpu.
Mock only external boundaries, not internal implementation details.
Ensure test isolation: no shared mutable state or execution-order dependencies between tests. If something must be run first before executing a test, include it in the test or a fixture.
Use @pytest.mark.parametrize for testing multiple input combinations rather than copy-pasting similar tests.

Organize tests using pytest following the structure in tests/TESTING.md with support for unit tests, smoke tests, and end-to-end tests

Files:

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

📄 CodeRabbit inference engine (AGENTS.md)

The unit_test marker is deprecated; use unit instead

Files:

  • tests/generation/test_vllm_backend.py
  • tests/config/test_generate.py
🧠 Learnings (1)
📚 Learning: 2026-05-27T22:20:37.354Z
Learnt from: kendrickb-nvidia
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 520
File: tests/generation/test_vllm_backend.py:556-587
Timestamp: 2026-05-27T22:20:37.354Z
Learning: In NVIDIA-NeMo/Safe-Synthesizer, `tests/conftest.py`’s `pytest_collection_modifyitems` hook applies pytest category markers automatically based on each test file’s path: tests under `/e2e/` get `pytest.mark.e2e`, tests under `/smoke/` get `pytest.mark.smoke`, and all other tests get `pytest.mark.unit`. Therefore, when reviewing pytest tests outside `tests/e2e/` and `tests/smoke/`, do not flag missing explicit `pytest.mark.unit` decorators on test classes/functions as an issue (the hook will add them during collection). If a new test directory/category is introduced, ensure the hook is updated so it’s categorized correctly.

Applied to files:

  • tests/generation/test_vllm_backend.py
  • tests/config/test_generate.py
🪛 Ruff (0.15.15)
src/nemo_safe_synthesizer/config/generate.py

[warning] 29-38: __all__ is not sorted

Apply an isort-style sorting to __all__

(RUF022)

🔇 Additional comments (8)
tests/config/test_generate.py (2)

4-34: LGTM!


38-68: LGTM!

src/nemo_safe_synthesizer/config/generate.py (5)

29-38: __all__ sorting issue already flagged.

Ruff RUF022 still reports that __all__ is not sorted using isort-style conventions. This was raised in a previous review comment. Running make format should auto-fix this.


23-26: LGTM!


41-55: LGTM!


199-211: LGTM!


247-255: LGTM!

tests/generation/test_vllm_backend.py (1)

100-106: LGTM!

Comment thread tests/generation/test_vllm_backend.py
Comment thread tests/generation/test_vllm_backend.py
Verify full argument contracts in auto schema-method tests and fix
isort-style __all__ ordering in generate config.

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

@kendrickb-nvidia kendrickb-nvidia left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

So the defaults are backend="auto" and method="auto" which will use xgrammar and structural tags, right?

Comment thread src/nemo_safe_synthesizer/config/generate.py
@mckornfield
mckornfield added this pull request to the merge queue Jun 3, 2026
@mckornfield
mckornfield removed this pull request from the merge queue due to a manual request Jun 3, 2026
@mckornfield
mckornfield added this pull request to the merge queue Jun 3, 2026
Merged via the queue into main with commit 37f33f7 Jun 3, 2026
19 checks passed
@mckornfield
mckornfield deleted the xgrammar-upgrade-structured-tags/mck branch June 3, 2026 19:57
seayang-nv pushed a commit that referenced this pull request Jun 3, 2026
* Also address loading of params tweak
* benchmark for methods using pretrained paths

# Summary
<!-- Brief description of changes -->

## Pre-Review Checklist

<!-- These checks should be completed before a PR is reviewed, -->
<!-- but you can submit a draft early to indicate that the issue is
being worked on. -->

Ensure that the following pass:

- [x] `make format && make check` or via prek validation.
- [x] `make test` passes locally
- [x] `make test-e2e` passes locally
- [ ] `make 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

<!-- These checks need to be completed before a PR is merged, -->
<!-- but as PRs often change significantly during review, -->
<!-- it's OK for them to be incomplete when review is first requested.
-->

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

## Other Notes

<!-- Please add the issue number that should be closed when this PR is
merged. -->
- Closes #540

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

## Release Notes

* **New Features**
* Added `"auto"` schema method for structured generation, now the
default. Intelligently selects `"structural_tag"` on compatible backends
or falls back to `"regex"`.
* Added `"structural_tag"` method option for structured generation with
XGrammar support.

* **Documentation**
* Updated configuration reference and guides to reflect new `"auto"`
default and additional schema method options.

* **Chores**
  * Added XGrammar dependency (Linux-only) to optional dependencies.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: mkornfield <mkornfield@nvidia.com>
Signed-off-by: Sean Yang <seayang@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

chore Maintenance not tied to a user-visible change docs Documentation-only change feature New feature or request test Test-only addition or change

Projects

None yet

Development

Successfully merging this pull request may close these issues.

upgrade to xgrammar2 and benchmark

4 participants