chore: update ty to 0.0.44

[binaryaaron]

Summary

• Updates the pinned ty version to 0.0.44.
• Fixes diagnostics surfaced by the checker update.

Test plan

• mise run typecheck

Related issue: #614

Summary by CodeRabbit

• Bug Fixes
  • Dataset loading now verifies readers return a pandas.DataFrame and fails fast with a clear error when they don’t.
  • Date/datetime transformations and distribution formatting are more robust (precision/format behavior preserved; safer handling of optional formatting outputs).
  • Evaluation protection logic improvements for quasi-identifier generation, categorical detection, and entropy normalization.
  • Multi-modal figure generation now skips missing histogram figures instead of assuming they always exist.
• Documentation / Style
  • Updated style-guide guidance for TypeIs type narrowing.
  • Improved type-safety around logging and configuration handling.
• Telemetry
  • Telemetry JSON field serialization keys updated for consistent naming.
• Tests
  • Added/expanded unit tests covering the above behaviors.
• Chores
  • Bumped the ty tool version.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 6647cca9-7119-4b02-bfd2-171b838eef20

📥 Commits

Reviewing files that changed from the base of the PR and between d181b0c and dbb50c6.

📒 Files selected for processing (1)

• src/nemo_safe_synthesizer/cli/datasets.py

🚧 Files skipped from review as they are similar to previous changes (1)

• src/nemo_safe_synthesizer/cli/datasets.py

📜 Recent review details

⏰ Context from checks skipped due to timeout. (5)

• GitHub Check: Unit Tests (3.12)
• GitHub Check: Unit Tests (3.11)
• GitHub Check: Unit Tests (3.13)
• GitHub Check: Smoke Tests
• GitHub Check: Analyze (Python)

---

Walkthrough

The PR updates typing guidance and runtime validation across datasets, data processing, privacy, telemetry, and evaluation code, with matching tests and a tool version bump.

Changes

Type Safety Hardening

Layer / File(s)	Summary
Type checker upgrade and TypeIs convention .mise.toml, STYLE_GUIDE.md, src/nemo_safe_synthesizer/utils.py, tests/test_utils.py	ty is bumped to 0.0.44; the style guide adds TypeIs[T] guidance; log_stats narrows statistics input with _is_statistics_list; tests cover the helper.
Dataset reader DataFrame runtime validation src/nemo_safe_synthesizer/cli/datasets.py, tests/cli/test_datasets.py	DatasetInfo.fetch now wraps pandas readers through a widened callable and validates that the result is a pd.DataFrame; a test asserts the TypeError path.
Datetime distribution and date label typing src/nemo_safe_synthesizer/data_processing/actions/distributions.py, src/nemo_safe_synthesizer/data_processing/actions/dates.py, tests/data_processing/test_distributions.py, tests/data_processing/test_dates.py	The datetime helper now branches on format and builds typed output lists; date helpers widen keys to Hashable; tests cover precision/format output and non-string date labels.
DP linear grad sample helpers src/nemo_safe_synthesizer/privacy/dp_transformers/linear.py, tests/training/test_dp_linear.py	_contract_tensor and _linear_parameter validate contraction results and parameter types, and compute_linear_grad_sample uses them; tests cover the normal path and both error cases.
Telemetry serialization alias migration src/nemo_safe_synthesizer/telemetry.py	Several telemetry fields move from alias= to serialization_alias=, the model config becomes ConfigDict(populate_by_name=True), and _redact_endpoint returns the redacted string directly.
Evaluation component logic updates src/nemo_safe_synthesizer/evaluation/components/..., src/nemo_safe_synthesizer/evaluation/reports/multimodal/multimodal_report.py	Membership-inference types are tightened, attribute-inference quasi-identifier generation and entropy normalization change, histogram figures are appended only when present, and multimodal components are typed.
Evaluation tests and config dispatch tests/config/test_autoconfig.py, tests/evaluation/components/...	Tests cover the updated membership-inference and attribute-inference behavior, and AutoConfigTestCase.get_config now distinguishes instances from factory callables.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

• NVIDIA-NeMo/Safe-Synthesizer#505 — Both PRs touch src/nemo_safe_synthesizer/data_processing/actions/dates.py and change fit_and_transform_dates behavior/typing.

Suggested reviewers

• kendrickb-nvidia

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% 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
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly reflects the primary change: updating the pinned ty version to 0.0.44.
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.

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

✨ 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 binaryaaron/pr596-01-ty-bump

---

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[greptile-apps]

Greptile Summary

This PR bumps the ty type-checker from 0.0.32 to 0.0.44 and resolves all new diagnostics raised by the stricter version. Several of the fixes correct pre-existing runtime bugs that the type checker surfaced.

• distributions.py: _apply_universal_params was aliasing ret = samples and then appending to it inside for sample in samples:, producing an infinite loop whenever neither precision nor format was set. The rewrite builds fresh output lists instead.
• attribute_inference_protection.py: The wide-table quasi-identifier path had two bugs — an off-by-one in the range (len - count instead of len - count + 1, silently dropping the last window) and a bounds check that fired after the fetch, causing the last combo's attack to be skipped via continue. Both are fixed.
• All other changes are type-annotation corrections (Hashable column labels, dict key/value types, serialization_alias vs alias in telemetry) and new unit tests that directly cover the fixed code paths.

Confidence Score: 5/5

Safe to merge; the type-checker upgrade uncovered and fixed real bugs rather than introducing new ones.

Every source-code change is either a type-annotation correction that matches the runtime behavior already present, or an actual bug fix (infinite-loop risk in distributions, off-by-one and skipped attack iteration in AIA). The new unit tests directly cover the fixed paths. No logic was removed or re-architected in a way that could regress existing behavior.

No files require special attention.

Important Files Changed

Filename	Overview
.mise.toml	Bumps ty version pin from 0.0.32 to 0.0.44; no other changes.
mise.lock	Updates checksums and download URLs for the new ty 0.0.44 binaries across all supported platforms.
STYLE_GUIDE.md	Adds guidance for TypeIs[T] from typing_extensions, explaining when to prefer it over TypeGuard and bool.
src/nemo_safe_synthesizer/cli/datasets.py	Adds _require_dataframe runtime validator and _dynamic_callable type-eraser to satisfy ty's invalid-argument-type diagnostics on overloaded pandas readers; replaces ty:ignore comment with proper helpers.
src/nemo_safe_synthesizer/data_processing/actions/dates.py	Widens column key types from str to Hashable to correctly model pandas DataFrame column labels, fixing a latent type error for non-string column indices.
src/nemo_safe_synthesizer/data_processing/actions/distributions.py	Rewrites _apply_universal_params to build fresh output lists rather than appending to the input list; fixes a latent infinite-loop bug when both precision and format were None.
src/nemo_safe_synthesizer/evaluation/components/attribute_inference_protection.py	Multiple fixes: Hashable column label support in _is_really_categorical; wide-table path off-by-one corrected and set→tuple combo type; qi_index bounds check moved before the fetch; entropy weighting converts to np.array before min/max.
src/nemo_safe_synthesizer/evaluation/components/membership_inference_protection.py	Corrects stale type annotations on tp_cnts/fp_cnts (dict[str, list[int]] → dict[float, int]) to match actual runtime structure keyed by float threshold values.
src/nemo_safe_synthesizer/evaluation/components/multi_modal_figures.py	Adds a None guard before appending a figure to the list; adds explicit list[go.Figure] annotation.
src/nemo_safe_synthesizer/evaluation/reports/multimodal/multimodal_report.py	Imports Component and annotates the components list to satisfy ty's type inference for the heterogeneous append calls.
src/nemo_safe_synthesizer/privacy/dp_transformers/linear.py	Introduces _contract_tensor and _linear_parameter type-narrowing helpers so the grad-sample dict is correctly typed as dict[nn.Parameter, torch.Tensor]; logic is otherwise unchanged.
src/nemo_safe_synthesizer/telemetry.py	Converts all alias= fields to serialization_alias= (output-only) and removes the now-no-op populate_by_name=True model_config; removes unnecessary cast() call on urlunsplit.
src/nemo_safe_synthesizer/utils.py	Adds _is_statistics_list TypeIs predicate and uses it in log_stats to avoid the inline isinstance workaround; aligns with new style guide guidance.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[DatetimeDistribution.sample] --> B[sample_datetimes]
    B --> C[_apply_universal_params]
    C --> D{format set?}
    D -- Yes --> E[Build fresh str list with optional rounding and strftime]
    D -- No --> F{precision set?}
    F -- Yes --> G[Build fresh datetime list with rounding]
    F -- No --> H[Return samples in fresh list - no mutation]
    E --> I[Return list of str]
    G --> I2[Return list of datetime]
    H --> I2

    subgraph AIA[AIA Wide-Table Path]
        J[training_columns from df.columns] --> K{len lt 500?}
        K -- No --> L[range len minus qi_count plus 1 - fixed off-by-one]
        L --> M[Build tuples not sets - fixed from set]
        K -- Yes --> N[itertools.combinations]
        M --> O[shuffle qi_combos]
        N --> O
        O --> P{qi_index equals len qi_combos?}
        P -- Yes --> Q[Stop - bounds check before fetch]
        P -- No --> R[Fetch combo and run attack]
    end

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"}}}%%
flowchart TD
    A[DatetimeDistribution.sample] --> B[sample_datetimes]
    B --> C[_apply_universal_params]
    C --> D{format set?}
    D -- Yes --> E[Build fresh str list with optional rounding and strftime]
    D -- No --> F{precision set?}
    F -- Yes --> G[Build fresh datetime list with rounding]
    F -- No --> H[Return samples in fresh list - no mutation]
    E --> I[Return list of str]
    G --> I2[Return list of datetime]
    H --> I2

    subgraph AIA[AIA Wide-Table Path]
        J[training_columns from df.columns] --> K{len lt 500?}
        K -- No --> L[range len minus qi_count plus 1 - fixed off-by-one]
        L --> M[Build tuples not sets - fixed from set]
        K -- Yes --> N[itertools.combinations]
        M --> O[shuffle qi_combos]
        N --> O
        O --> P{qi_index equals len qi_combos?}
        P -- Yes --> Q[Stop - bounds check before fetch]
        P -- No --> R[Fetch combo and run attack]
    end

Loading

_{Reviews (4): Last reviewed commit: "docs: explain dynamic pandas reader typi..." | Re-trigger Greptile}

[coderabbitai]

Actionable comments posted: 6

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

src/nemo_safe_synthesizer/data_processing/actions/dates.py (1)

418-431: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

Preserve the original column label in date_min_dict.

Casting object_col to str changes the metadata key that later reversal code needs to map back onto the DataFrame. For non-string labels like 0 or tuples, this records "0" here while the transformed frame still carries 0, so the reverse transform will miss the column or fail. Keep the original label as the key and widen the annotation instead of stringifying it.

As per path instructions, review src/nemo_safe_synthesizer/data_processing/**/*.py for data-contract regressions.

Source: Path instructions

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 3070e10e-c84b-4a12-9c48-d92f2d5c395b

📥 Commits

Reviewing files that changed from the base of the PR and between a9848f7 and d899676.

⛔ Files ignored due to path filters (1)

• mise.lock is excluded by !**/*.lock, !mise.lock

📒 Files selected for processing (19)

• .mise.toml
• STYLE_GUIDE.md
• src/nemo_safe_synthesizer/cli/datasets.py
• src/nemo_safe_synthesizer/data_processing/actions/dates.py
• src/nemo_safe_synthesizer/data_processing/actions/distributions.py
• src/nemo_safe_synthesizer/evaluation/components/attribute_inference_protection.py
• src/nemo_safe_synthesizer/evaluation/components/membership_inference_protection.py
• src/nemo_safe_synthesizer/evaluation/components/multi_modal_figures.py
• src/nemo_safe_synthesizer/evaluation/reports/multimodal/multimodal_report.py
• src/nemo_safe_synthesizer/privacy/dp_transformers/linear.py
• src/nemo_safe_synthesizer/telemetry.py
• src/nemo_safe_synthesizer/utils.py
• tests/cli/test_datasets.py
• tests/config/test_autoconfig.py
• tests/data_processing/test_dates.py
• tests/data_processing/test_distributions.py
• tests/evaluation/components/test_attribute_inference_protection.py
• tests/evaluation/components/test_membership_inference_protection.py
• tests/training/test_dp_linear.py

📜 Review details

⏰ Context from checks skipped due to timeout. (3)

• GitHub Check: Unit Tests (3.13)
• GitHub Check: Unit Tests (3.11)
• GitHub Check: Unit Tests (3.12)

🧰 Additional context used

📓 Path-based instructions (17)

.mise.toml

⚙️ CodeRabbit configuration file

Treat .mise.toml as toolchain supply-chain configuration. Check pinned tool choices, install cadence, platform coverage, environment settings, and whether changes require regenerating mise.lock.

Files:

• .mise.toml

**/*.{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/privacy/dp_transformers/linear.py
• tests/training/test_dp_linear.py
• STYLE_GUIDE.md
• tests/data_processing/test_distributions.py
• tests/data_processing/test_dates.py
• tests/config/test_autoconfig.py
• tests/cli/test_datasets.py
• tests/evaluation/components/test_attribute_inference_protection.py
• tests/evaluation/components/test_membership_inference_protection.py
• src/nemo_safe_synthesizer/evaluation/reports/multimodal/multimodal_report.py
• src/nemo_safe_synthesizer/evaluation/components/multi_modal_figures.py
• src/nemo_safe_synthesizer/cli/datasets.py
• src/nemo_safe_synthesizer/data_processing/actions/distributions.py
• src/nemo_safe_synthesizer/utils.py
• src/nemo_safe_synthesizer/evaluation/components/membership_inference_protection.py
• src/nemo_safe_synthesizer/data_processing/actions/dates.py
• src/nemo_safe_synthesizer/evaluation/components/attribute_inference_protection.py
• src/nemo_safe_synthesizer/telemetry.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: Source code must remain Python 3.11 syntax-compatible; do not use Python 3.12-only syntax such as PEP 695 type statements or bracketed generic class/function parameters in shared package code
Use ruff for formatting and linting via mise run format and mise run check tasks; run formatting before committing
Use ty for type checking via mise run check task; ensure type hints are present and valid
New features must include tests; bug fixes must include regression tests

**/*.py: In Python config/ models, use NSSBaseModel for user-facing configuration/parameter models; use raw BaseModel or module-specific bases for DTOs and internal structures.
Use BaseSettings for environment/CLI settings, prefer AliasChoices for per-field aliases when a setting must respond to both its Python name and an environment variable name, and use env_prefix only for simple shared-prefix settings.
For Pydantic models, use Field(description=...) as the canonical field documentation and always include it for model fields.
Prefer assignment-style Pydantic fields (name: T = Field(...)) over Annotated[...] unless extra metadata beyond Field() is needed; use bare assignments for defaults with Annotated, except that default_factory still belongs in assignment-style Field(default_factory=...).
Use @dataclass(frozen=True) for immutable value objects and validators; mutable dataclasses are acceptable for builders, accumulators, and pipeline state.
Use field(default_factory=list) (or another factory) for mutable dataclass defaults; never use a mutable literal default such as = [].
Use StrEnum for string-valued enums used in configs or serialization, and plain Enum for internal-only named constants.
In Python cod...

Files:

• src/nemo_safe_synthesizer/privacy/dp_transformers/linear.py
• tests/training/test_dp_linear.py
• tests/data_processing/test_distributions.py
• tests/data_processing/test_dates.py
• tests/config/test_autoconfig.py
• tests/cli/test_datasets.py
• tests/evaluation/components/test_attribute_inference_protection.py
• tests/evaluation/components/test_membership_inference_protection.py
• src/nemo_safe_synthesizer/evaluation/reports/multimodal/multimodal_report.py
• src/nemo_safe_synthesizer/evaluation/components/multi_modal_figures.py
• src/nemo_safe_synthesizer/cli/datasets.py
• src/nemo_safe_synthesizer/data_processing/actions/distributions.py
• src/nemo_safe_synthesizer/utils.py
• src/nemo_safe_synthesizer/evaluation/components/membership_inference_protection.py
• src/nemo_safe_synthesizer/data_processing/actions/dates.py
• src/nemo_safe_synthesizer/evaluation/components/attribute_inference_protection.py
• src/nemo_safe_synthesizer/telemetry.py

**/*.{py,sh,yaml,yml,md}

📄 CodeRabbit inference engine (CONTRIBUTING.md)

All source files (.py, .sh, .yaml, .yml, .md) require SPDX copyright headers; mise run format adds them automatically

Files:

• src/nemo_safe_synthesizer/privacy/dp_transformers/linear.py
• tests/training/test_dp_linear.py
• STYLE_GUIDE.md
• tests/data_processing/test_distributions.py
• tests/data_processing/test_dates.py
• tests/config/test_autoconfig.py
• tests/cli/test_datasets.py
• tests/evaluation/components/test_attribute_inference_protection.py
• tests/evaluation/components/test_membership_inference_protection.py
• src/nemo_safe_synthesizer/evaluation/reports/multimodal/multimodal_report.py
• src/nemo_safe_synthesizer/evaluation/components/multi_modal_figures.py
• src/nemo_safe_synthesizer/cli/datasets.py
• src/nemo_safe_synthesizer/data_processing/actions/distributions.py
• src/nemo_safe_synthesizer/utils.py
• src/nemo_safe_synthesizer/evaluation/components/membership_inference_protection.py
• src/nemo_safe_synthesizer/data_processing/actions/dates.py
• src/nemo_safe_synthesizer/evaluation/components/attribute_inference_protection.py
• src/nemo_safe_synthesizer/telemetry.py

src/nemo_safe_synthesizer/**/*.py

📄 CodeRabbit inference engine (CONTRIBUTING.md)

API reference pages are auto-generated from Python docstrings using Google-style format; write docstrings in src/nemo_safe_synthesizer/ and they will appear in the reference/

Files:

• src/nemo_safe_synthesizer/privacy/dp_transformers/linear.py
• src/nemo_safe_synthesizer/evaluation/reports/multimodal/multimodal_report.py
• src/nemo_safe_synthesizer/evaluation/components/multi_modal_figures.py
• src/nemo_safe_synthesizer/cli/datasets.py
• src/nemo_safe_synthesizer/data_processing/actions/distributions.py
• src/nemo_safe_synthesizer/utils.py
• src/nemo_safe_synthesizer/evaluation/components/membership_inference_protection.py
• src/nemo_safe_synthesizer/data_processing/actions/dates.py
• src/nemo_safe_synthesizer/evaluation/components/attribute_inference_protection.py
• src/nemo_safe_synthesizer/telemetry.py

src/**/*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

Every directory under src/ that contains Python files must include an __init__.py file, even if empty.

Files:

• src/nemo_safe_synthesizer/privacy/dp_transformers/linear.py
• src/nemo_safe_synthesizer/evaluation/reports/multimodal/multimodal_report.py
• src/nemo_safe_synthesizer/evaluation/components/multi_modal_figures.py
• src/nemo_safe_synthesizer/cli/datasets.py
• src/nemo_safe_synthesizer/data_processing/actions/distributions.py
• src/nemo_safe_synthesizer/utils.py
• src/nemo_safe_synthesizer/evaluation/components/membership_inference_protection.py
• src/nemo_safe_synthesizer/data_processing/actions/dates.py
• src/nemo_safe_synthesizer/evaluation/components/attribute_inference_protection.py
• src/nemo_safe_synthesizer/telemetry.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/privacy/dp_transformers/linear.py
• src/nemo_safe_synthesizer/evaluation/reports/multimodal/multimodal_report.py
• src/nemo_safe_synthesizer/evaluation/components/multi_modal_figures.py
• src/nemo_safe_synthesizer/cli/datasets.py
• src/nemo_safe_synthesizer/data_processing/actions/distributions.py
• src/nemo_safe_synthesizer/utils.py
• src/nemo_safe_synthesizer/evaluation/components/membership_inference_protection.py
• src/nemo_safe_synthesizer/data_processing/actions/dates.py
• src/nemo_safe_synthesizer/evaluation/components/attribute_inference_protection.py
• src/nemo_safe_synthesizer/telemetry.py

**/*

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

**/*: Every source file must include an SPDX copyright header, with HTML comments for Markdown and hash comments for .py, .sh, .yaml, and .yml; Markdown files with YAML frontmatter must place the hash-comment SPDX header inside the frontmatter block.
Ensure files end with a newline and contain no trailing whitespace; keep single spacing between sentences.

Files:

• src/nemo_safe_synthesizer/privacy/dp_transformers/linear.py
• tests/training/test_dp_linear.py
• STYLE_GUIDE.md
• tests/data_processing/test_distributions.py
• tests/data_processing/test_dates.py
• tests/config/test_autoconfig.py
• tests/cli/test_datasets.py
• tests/evaluation/components/test_attribute_inference_protection.py
• tests/evaluation/components/test_membership_inference_protection.py
• src/nemo_safe_synthesizer/evaluation/reports/multimodal/multimodal_report.py
• src/nemo_safe_synthesizer/evaluation/components/multi_modal_figures.py
• src/nemo_safe_synthesizer/cli/datasets.py
• src/nemo_safe_synthesizer/data_processing/actions/distributions.py
• src/nemo_safe_synthesizer/utils.py
• src/nemo_safe_synthesizer/evaluation/components/membership_inference_protection.py
• src/nemo_safe_synthesizer/data_processing/actions/dates.py
• src/nemo_safe_synthesizer/evaluation/components/attribute_inference_protection.py
• src/nemo_safe_synthesizer/telemetry.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/privacy/dp_transformers/linear.py
• tests/training/test_dp_linear.py
• STYLE_GUIDE.md
• tests/data_processing/test_distributions.py
• tests/data_processing/test_dates.py
• tests/config/test_autoconfig.py
• tests/cli/test_datasets.py
• tests/evaluation/components/test_attribute_inference_protection.py
• tests/evaluation/components/test_membership_inference_protection.py
• src/nemo_safe_synthesizer/evaluation/reports/multimodal/multimodal_report.py
• src/nemo_safe_synthesizer/evaluation/components/multi_modal_figures.py
• src/nemo_safe_synthesizer/cli/datasets.py
• src/nemo_safe_synthesizer/data_processing/actions/distributions.py
• src/nemo_safe_synthesizer/utils.py
• src/nemo_safe_synthesizer/evaluation/components/membership_inference_protection.py
• src/nemo_safe_synthesizer/data_processing/actions/dates.py
• src/nemo_safe_synthesizer/evaluation/components/attribute_inference_protection.py
• src/nemo_safe_synthesizer/telemetry.py

**

⚙️ CodeRabbit configuration file

**:

AGENTS.md

Guide for AI agents (Cursor, Windsurf, Claude Code, etc.) working in the Safe-Synthesizer repo.

This project loads local developer preferences from @AGENTS.local.md. You MUST read this file if it exists and give its instructions top priority.

Skills

Repo-specific skills live in .agents/skills/; see .agents/README.md for the catalog. Read a skill when the task matches its scope instead of copying workflow details into this file.

Durable implementation guidance belongs with the code it describes: function and class docstrings for public contracts and source comments for local invariants. Test-suite guidance belongs in tests/TESTING.md.

Repo Conventions

See STYLE_GUIDE.md for detailed code style conventions (Python, markdown, Dockerfiles, shell scripts, testing, config files, docstrings).

Use uv for everything -- never pip or raw python. Python 3.11–3.13 with modern syntax (X | Y, list[str], Self). Python 3.14+ is not supported.

Common commands: mise run test (unit tests), mise run format (auto-fix formatting + lint + copyright), mise run check (read-only local quality checks), mise run validate (pre-PR quality, lock, and CI unit checks), mise run typecheck (ty only). Always use mise tasks or the wrapper scripts in tools/ instead of running ruff or ty directly. Use uv run for Python execution. When in doubt, inspect mise tasks and pytest --markers.

The canonical uv sync command for a full GPU/dev environment is:

uv sync --frozen --extra cu129 --extra engine --group dev

Bare uv sync --frozen (without extras) installs an incomplete environment -- ty, import checks, and GPU tests will fail.

Feature branches off main. Branch names often include an issue number prefix (e.g., <author>/123-short-name).

Do ...

Files:

• src/nemo_safe_synthesizer/privacy/dp_transformers/linear.py
• tests/training/test_dp_linear.py
• STYLE_GUIDE.md
• tests/data_processing/test_distributions.py
• tests/data_processing/test_dates.py
• tests/config/test_autoconfig.py
• tests/cli/test_datasets.py
• tests/evaluation/components/test_attribute_inference_protection.py
• tests/evaluation/components/test_membership_inference_protection.py
• src/nemo_safe_synthesizer/evaluation/reports/multimodal/multimodal_report.py
• src/nemo_safe_synthesizer/evaluation/components/multi_modal_figures.py
• src/nemo_safe_synthesizer/cli/datasets.py
• src/nemo_safe_synthesizer/data_processing/actions/distributions.py
• src/nemo_safe_synthesizer/utils.py
• src/nemo_safe_synthesizer/evaluation/components/membership_inference_protection.py
• src/nemo_safe_synthesizer/data_processing/actions/dates.py
• src/nemo_safe_synthesizer/evaluation/components/attribute_inference_protection.py
• src/nemo_safe_synthesizer/telemetry.py

src/nemo_safe_synthesizer/privacy/**/*.py

⚙️ CodeRabbit configuration file

Treat privacy changes as high-risk. Check DP accounting, parameter validation, data leakage, seed handling, model state persistence, and whether privacy guarantees are documented accurately.

Files:

• src/nemo_safe_synthesizer/privacy/dp_transformers/linear.py

**/test_*.py

📄 CodeRabbit inference engine (AGENTS.md)

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

Files:

• tests/training/test_dp_linear.py
• tests/data_processing/test_distributions.py
• tests/data_processing/test_dates.py
• tests/config/test_autoconfig.py
• tests/cli/test_datasets.py
• tests/evaluation/components/test_attribute_inference_protection.py
• tests/evaluation/components/test_membership_inference_protection.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/training/test_dp_linear.py
• tests/data_processing/test_distributions.py
• tests/data_processing/test_dates.py
• tests/config/test_autoconfig.py
• tests/cli/test_datasets.py
• tests/evaluation/components/test_attribute_inference_protection.py
• tests/evaluation/components/test_membership_inference_protection.py

---

⚙️ CodeRabbit configuration file

tests/**:

Testing Guide

Comprehensive testing reference for Safe-Synthesizer developers. Covers commands, markers, test data, fixtures, and gotchas.

Read First

1. tests/conftest.py -- auto-marking, load_test_dataset/load_test_dataframe, fixture_mock_processor pattern
2. pytest.ini -- markers, asyncio, timeout
3. tests/evaluation/conftest.py -- most complex: Faker-based make_df, nullable dtype conversion
4. tests/generation/conftest.py -- JSONL/schema fixtures, fixture_valid_iris_dataset_jsonl_and_schema

Running Tests

All mise test tasks, grouped by scope:

mise run test                              # Unit (excludes slow, e2e, and smoke)
mise run test:unit-slow                    # Unit tests including slow (excludes e2e and smoke)
mise run test:smoke                        # CPU smoke tests (~few min, no GPU required)
mise run test:smoke:gpu                    # All staged GPU smoke tests (requires CUDA)
mise run test:smoke:gpu:train-only
mise run test:smoke:gpu:generation
mise run test:smoke:gpu:resume
mise run test:smoke:gpu:structured-generation
mise run test:smoke:gpu:timeseries
mise run test:smoke:gpu:smollm2
mise run test:e2e                          # All e2e (requires CUDA) -- runs default + dp
mise run test:e2e:default                  # e2e default (no-DP) tests only
mise run test:e2e:dp                       # e2e DP tests only
mise run test:ci                           # CI unit tests with coverage (excludes slow, e2e, gpu, smoke)
mise run test:ci-slow                      # CI slow tests with coverage
mise run test:ci-container                 # CI tests in a Linux container (Docker/Podman)

Run a single test:

uv run --frozen pytest tests/path/test_file.py::test_name -vvs -n0

Test runner: uv run --frozen pytest -n auto --dist loadscope -vv...

Files:

• tests/training/test_dp_linear.py
• tests/data_processing/test_distributions.py
• tests/data_processing/test_dates.py
• tests/config/test_autoconfig.py
• tests/cli/test_datasets.py
• tests/evaluation/components/test_attribute_inference_protection.py
• tests/evaluation/components/test_membership_inference_protection.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/training/test_dp_linear.py
• tests/data_processing/test_distributions.py
• tests/data_processing/test_dates.py
• tests/config/test_autoconfig.py
• tests/cli/test_datasets.py
• tests/evaluation/components/test_attribute_inference_protection.py
• tests/evaluation/components/test_membership_inference_protection.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/training/test_dp_linear.py
• tests/data_processing/test_distributions.py
• tests/data_processing/test_dates.py
• tests/config/test_autoconfig.py
• tests/cli/test_datasets.py
• tests/evaluation/components/test_attribute_inference_protection.py
• tests/evaluation/components/test_membership_inference_protection.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:

• STYLE_GUIDE.md

**/*.md

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

**/*.md: Do not use decorative bold text in Markdown body text, list items, or docstrings; use headers, list markers, colons, and backticks for structure instead.
Use -- for asides in Markdown and single backticks for code identifiers, paths, and CLI commands.
For Mermaid diagrams in Markdown, use node IDs without spaces, quote labels that contain special characters, and do not use explicit colors or styles.

Files:

• STYLE_GUIDE.md

tests/cli/**/*.py

📄 CodeRabbit inference engine (tests/TESTING.md)

Use mock_workdir(tmp_path) helper in CLI tests to create temporary Workdir instances

Files:

• tests/cli/test_datasets.py

src/nemo_safe_synthesizer/evaluation/**/*.py

⚙️ CodeRabbit configuration file

Treat evaluation changes as correctness-sensitive. Check metric inputs, holdout usage, privacy metric semantics, report data shape, missing-data handling, and whether unavailable metrics fail or degrade intentionally.

Files:

• src/nemo_safe_synthesizer/evaluation/reports/multimodal/multimodal_report.py
• src/nemo_safe_synthesizer/evaluation/components/multi_modal_figures.py
• src/nemo_safe_synthesizer/evaluation/components/membership_inference_protection.py
• src/nemo_safe_synthesizer/evaluation/components/attribute_inference_protection.py

src/nemo_safe_synthesizer/data_processing/**/*.py

⚙️ CodeRabbit configuration file

Review for data-contract regressions. Check input/training/test/synthetic naming, group boundaries, token-budget math, record ordering, schema and column validation, nullable dtypes, and deterministic behavior.

Files:

• src/nemo_safe_synthesizer/data_processing/actions/distributions.py
• src/nemo_safe_synthesizer/data_processing/actions/dates.py

🧠 Learnings (2)

📓 Common learnings

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer

Timestamp: 2026-06-24T17:38:02.057Z
Learning: Use American English spelling in code and documentation (for example, "initialize" not "initialise", "recognize" not "recognise", "color" not "colour").

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer

Timestamp: 2026-06-24T17:38:02.057Z
Learning: Treat `__all__` as the public API surface, and treat identifiers with a leading `_` as private and changeable without notice.

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer

Timestamp: 2026-06-24T17:38:02.057Z
Learning: Follow local consistency in the surrounding code even when it differs from the style guide, and migrate legacy code toward the conventions here when practical.

📚 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/training/test_dp_linear.py
• tests/data_processing/test_distributions.py
• tests/data_processing/test_dates.py
• tests/config/test_autoconfig.py
• tests/cli/test_datasets.py
• tests/evaluation/components/test_attribute_inference_protection.py
• tests/evaluation/components/test_membership_inference_protection.py

🪛 LanguageTool

STYLE_GUIDE.md

[style] ~144-~144: Consider using “incompatible” to avoid wordiness.
Context: ...ot express because the narrowed type is not compatible with the input type. ```python from ty...

(NOT_ABLE_PREMIUM)

🪛 Ruff (0.15.18)

tests/training/test_dp_linear.py

[warning] 41-41: Pattern passed to match= contains metacharacters but is neither escaped nor raw

(RUF043)

---

[warning] 46-46: Pattern passed to match= contains metacharacters but is neither escaped nor raw

(RUF043)

src/nemo_safe_synthesizer/utils.py

[warning] 137-137: zip() without an explicit strict= parameter

Add explicit value for parameter strict=

(B905)

🔇 Additional comments (3)

src/nemo_safe_synthesizer/telemetry.py (1)

26-29: LGTM!

Also applies to: 74-81, 170-265

src/nemo_safe_synthesizer/data_processing/actions/distributions.py (1)

79-92: LGTM!

tests/data_processing/test_distributions.py (1)

9-27: LGTM!

[mckornfield]

did this just need to be defined earlier?

[binaryaaron]

Yes, the intent was to define training_columns earlier so the quasi-identifier combinations and the prediction counters use the same column identity instead of mixing original column objects with stringified names.