chore: persist train/test/transformed_training csvs before preflight check by nina-xu · Pull Request #660 · NVIDIA-NeMo/Safe-Synthesizer

nina-xu · 2026-07-17T17:07:37Z

Summary

A quick adjustment to the order of operations around preflight:
Current behavior:
train/test split -> PII replacement -> preflight check -> persisting train/test/transformed_train
New behavior:
train/test split -> PII replacement -> persisting train/test/transformed_train -> preflight check

This is so that if preflight fails, we have the transformed training split to inspect. This happened to me because a faulty PII replacement prototype altered the shape of the input data significantly, thus failing the preflight. I suspect this would happen less frequently for main, but maybe it doesn't hurt to have? I'm open to closing this PR if we think this is not worth the while.

Pre-Review Checklist

Ensure that the following pass:

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

Pre-Merge Checklist

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

Other Notes

Closes #

Summary by CodeRabbit

Bug Fixes
- Processed training and test datasets are now saved before validation completes, making them available for inspection even when validation encounters an error.
- Validation-only checks no longer write processed datasets or transformed training data to disk.

Signed-off-by: nina-xu <19981858+nina-xu@users.noreply.github.com>

coderabbitai · 2026-07-17T17:09:15Z

Walkthrough

SafeSynthesizer.process_data() now persists processed dataset splits before preflight when not running in check-only mode. The previous end-of-function persistence block was removed, so validation avoids CSV and transformed-training writes.

Changes

Dataset persistence flow

Layer / File(s)	Summary
Guarded split persistence `src/nemo_safe_synthesizer/sdk/library_builder.py`	Processed `training.csv`, applicable `transformed_training`, and `test.csv` are written before preflight when `check_only` is false; the former unconditional end-of-function writes are removed.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested labels: bug

Suggested reviewers: binaryaaron

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and accurately summarizes the main change: moving CSV persistence before the preflight check.

✨ 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 nina-xu/save-transformed-training-before-preflight

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

codecov · 2026-07-17T17:13:57Z

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

greptile-apps · 2026-07-17T17:22:43Z

Greptile Summary

This PR reorders the CSV persistence step in process_data() to run before the preflight check (instead of after), so that the train/test splits and transformed training data are available on disk for inspection even when preflight fails. The check_only path is unchanged and still skips all CSV writes.

CSV writes moved earlier: training.csv, test.csv, and transformed_training.csv are now written immediately after the train/test split and PII replacement, and before run_preflight() is invoked on the training split.
check_only path unaffected: The if not check_only: guard preserves the existing behavior where validation-mode runs produce no CSV output.
_data_processed flag unaffected: The flag is still only set after a successful preflight, so a failed preflight on a non-check-only run leaves the builder in a re-runnable state.

Confidence Score: 5/5

Safe to merge — the reordering is tightly scoped and well-guarded, with no change to the check_only path or the data_processed flag semantics.

The change moves a block of CSV writes earlier in the method, wrapped in an existing if not check_only: guard that was already used for PII replacement. All three branches (check_only, preflight failure, preflight success) produce the correct on-disk and in-memory state. The _data_processed flag is still only set after a clean preflight, preserving the retry and validate→full-run patterns documented in the comments.

No files require special attention.

Important Files Changed

Filename	Overview
src/nemo_safe_synthesizer/sdk/library_builder.py	Moves CSV persistence block from after run_preflight() to before it (inside a if not check_only: guard); logic, guards, and flag semantics are all preserved correctly.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant C as Caller
    participant PD as process_data()
    participant FS as Filesystem
    participant PF as run_preflight()

    Note over C,PF: Before this PR (non-check_only path)
    C->>PD: process_data()
    PD->>PD: train_test_split()
    PD->>PD: PII replacement
    PD->>PF: run_preflight(training_df)
    alt preflight error
        PF-->>PD: raises ParameterError
        PD-->>C: ParameterError (no CSVs on disk)
    else preflight OK
        PF-->>PD: ok
        PD->>FS: write training.csv
        PD->>FS: write transformed_training.csv
        PD->>FS: write test.csv
        PD->>PD: "_data_processed = True"
    end

    Note over C,PF: After this PR (non-check_only path)
    C->>PD: process_data()
    PD->>PD: train_test_split()
    PD->>PD: PII replacement
    PD->>FS: write training.csv
    PD->>FS: write transformed_training.csv
    PD->>FS: write test.csv
    PD->>PF: run_preflight(training_df)
    alt preflight error
        PF-->>PD: raises ParameterError
        PD-->>C: ParameterError (CSVs available for inspection)
    else preflight OK
        PF-->>PD: ok
        PD->>PD: "_data_processed = True"
    end

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant C as Caller
    participant PD as process_data()
    participant FS as Filesystem
    participant PF as run_preflight()

    Note over C,PF: Before this PR (non-check_only path)
    C->>PD: process_data()
    PD->>PD: train_test_split()
    PD->>PD: PII replacement
    PD->>PF: run_preflight(training_df)
    alt preflight error
        PF-->>PD: raises ParameterError
        PD-->>C: ParameterError (no CSVs on disk)
    else preflight OK
        PF-->>PD: ok
        PD->>FS: write training.csv
        PD->>FS: write transformed_training.csv
        PD->>FS: write test.csv
        PD->>PD: "_data_processed = True"
    end

    Note over C,PF: After this PR (non-check_only path)
    C->>PD: process_data()
    PD->>PD: train_test_split()
    PD->>PD: PII replacement
    PD->>FS: write training.csv
    PD->>FS: write transformed_training.csv
    PD->>FS: write test.csv
    PD->>PF: run_preflight(training_df)
    alt preflight error
        PF-->>PD: raises ParameterError
        PD-->>C: ParameterError (CSVs available for inspection)
    else preflight OK
        PF-->>PD: ok
        PD->>PD: "_data_processed = True"
    end

_{Reviews (1): Last reviewed commit: "persist train/test/transformed_training ..." | Re-trigger Greptile}

coderabbitai

Actionable comments posted: 1

🧹 Nitpick comments (1)

src/nemo_safe_synthesizer/sdk/library_builder.py (1)
435-437: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win

Replace the runtime assertion with an explicit invariant check.

If the workdir invariant is ever broken, optimized Python removes the assert and the next line fails with an opaque AttributeError. Raise a precise RuntimeError instead.

As per coding guidelines: src/**/*.py must not use assert for validation in library code; raise an appropriate exception instead.
[details]

Proposed fix
 if not check_only:
-    assert self._workdir is not None
+    if self._workdir is None:
+        raise RuntimeError("Workdir must be initialized before persisting processed datasets")
     self._workdir.ensure_directories()
Source: Coding guidelines

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 60c84a90-dd08-4be6-96cb-1fbd03be7a8c

📥 Commits

Reviewing files that changed from the base of the PR and between 44e9d54 and 4b831a8.

📒 Files selected for processing (1)

src/nemo_safe_synthesizer/sdk/library_builder.py

📜 Review details

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

GitHub Check: Greptile Review

🧰 Additional context used

📓 Path-based instructions (6)

**/*.{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/sdk/library_builder.py

**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

**/*.py: Place durable implementation guidance in function and class docstrings for public contracts and source comments for local invariants
Target Python 3.11–3.13 with modern syntax (X | Y, list[str], Self). Python 3.14+ is not supported

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

Files:

src/nemo_safe_synthesizer/sdk/library_builder.py

src/**/*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

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

Files:

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/sdk/library_builder.py

**/*

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

**/*: Every source file requires the SPDX copyright and license header appropriate to its file format.
End files with a newline, remove trailing whitespace, use one space between sentences, and keep code, comments, and docstrings within 120 characters.

**/*: All merged commits must use Conventional Commits format: <type>(<scope>): <description> or <type>: <description>, with a lowercase valid type and a description no longer than 100 characters.
All contributions require a DCO Signed-off-by trailer and a cryptographic commit signature.
Before submitting a pull request, run formatting, checks, and tests using mise run format, mise run check, and mise run test.
Branches other than main must use lowercase author-prefixed names in one of the documented forms, optionally including an issue ID and category.
Release tags must use a v prefix and PEP 440 stable or release-candidate versions, such as v1.0.0 or v0.1.0rc0; alpha versions and dashed -rc suffixes are not used.
Do not move a published release tag; create and validate a new release candidate when code changes.
Use mise run <task> for project tasks; the Makefile only bootstraps mise and provides deprecated compatibility messages.

Files:

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/sdk/library_builder.py

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

📄 CodeRabbit inference engine (CONTRIBUTING.md)

All source files with .py, .sh, .yaml, .yml, or .md extensions must include SPDX copyright headers, except files listed in .copyrightignore.

Files:

src/nemo_safe_synthesizer/sdk/library_builder.py

src/nemo_safe_synthesizer/**/*.py

📄 CodeRabbit inference engine (CONTRIBUTING.md)

Write durable module-level guidance in Python docstrings and source comments so it appears in the generated API reference.

Files:

src/nemo_safe_synthesizer/sdk/library_builder.py

🔇 Additional comments (1)

src/nemo_safe_synthesizer/sdk/library_builder.py (1)

438-446: 🔒 Security & Privacy

Confirm raw-data retention for failed runs.

training.csv and test.csv contain the original, unredacted splits, but they are now written before full preflight. A non-check-only preflight failure therefore leaves raw data on disk. Confirm that failed workdirs have the intended access controls and cleanup/retention policy; otherwise use a protected temporary artifact lifecycle.

As per path instructions: privacy changes involving PII replacement and persisted datasets are high-risk and require explicit review.

Source: Path instructions

coderabbitai · 2026-07-17T17:22:52Z

+        # Persist the train/test splits now -- before the full preflight run
+        # below -- so that a preflight failure (e.g. group_exceeds_context)
+        # still leaves the processed datasets on disk for inspection. Skipped on
+        # the ``--validate`` path (``check_only``), which intentionally elides
+        # CSV writes.
+        if not check_only:
+            assert self._workdir is not None
+            self._workdir.ensure_directories()
+            # ``training.csv`` is the canonical persisted original training
+            # split -- reloaded by load_from_save_path and used for evaluation.
+            self._original_training_df.to_csv(self._workdir.dataset.training, index=False)
+            if not self._training_df.equals(self._original_training_df):
+                # The transformed (e.g. PII-replaced) training data is saved for
+                # inspection only -- we don't need it in generation or evaluation.
+                self._training_df.to_csv(self._workdir.dataset.transformed_training, index=False)
+            if self._test_df is not None:
+                self._test_df.to_csv(self._workdir.dataset.test, index=False)
+


🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Add regression coverage for persistence ordering and check_only.

Add focused tests using tmp_path that verify full-run persistence remains available when the later preflight raises, while process_data(check_only=True) creates no split CSVs. These are the two user-visible contracts introduced by this change.

As per path instructions: new behavior needs focused tests near the affected subsystem.

Source: Path instructions

binaryaaron · 2026-07-17T21:41:24Z

+                # inspection only -- we don't need it in generation or evaluation.
+                self._training_df.to_csv(self._workdir.dataset.transformed_training, index=False)
+            if self._test_df is not None:
+                self._test_df.to_csv(self._workdir.dataset.test, index=False)


agent-assisted (review-pr): Suggested addition addressing the raw-data-retention question raised by CodeRabbit.

The reordering itself looks correct, verified the failure-path behavior directly (wrote and ran a throwaway test against this exact commit: training.csv does persist when the full preflight check fails). The one open question was CodeRabbit's: raw training.csv/test.csv now land on disk even on a failed non-check_only run, where previously nothing was written.

Rather than blocking on that (it seems like the intended tradeoff, and workdirs already have no separate retention policy), a small log line would at least make the new behavior visible to the user in the moment it matters -- when preflight fails right after this write, they'd otherwise have no signal that their data was still saved.

Suggested change

self._test_df.to_csv(self._workdir.dataset.test, index=False)

self._test_df.to_csv(self._workdir.dataset.test, index=False)

logger.user.info(

f"Processed training data saved to {self._workdir.dataset.training} "

"(available for inspection even if the preflight check below fails)",

extra={"path": str(self._workdir.dataset.training)},

)

Verified this passes ruff check/ruff format and all 22 existing tests in tests/sdk/test_process_data.py.

yeah that's a good suggestion. Really I think what needs to be saved is the transformed dataset, if any, because the user should have a copy of the original dataset they can inspect. I can change to only save early if transformation, i.e. PII replacement is done. That might just make the code a bit wacky looking though. What do you think is a good behavior to have here?

persist train/test/transformed_training csvs before preflight check

4b831a8

Signed-off-by: nina-xu <19981858+nina-xu@users.noreply.github.com>

github-actions Bot added the area:sdk-cli label Jul 17, 2026

nina-xu changed the title ~~persist train/test/transformed_training csvs before preflight check~~ chore: persist train/test/transformed_training csvs before preflight check Jul 17, 2026

nina-xu marked this pull request as ready for review July 17, 2026 17:19

nina-xu requested a review from a team as a code owner July 17, 2026 17:19

coderabbitai Bot added the bug Defects in shipped behavior label Jul 17, 2026

coderabbitai Bot reviewed Jul 17, 2026

View reviewed changes

binaryaaron reviewed Jul 17, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

chore: persist train/test/transformed_training csvs before preflight check#660

chore: persist train/test/transformed_training csvs before preflight check#660
nina-xu wants to merge 1 commit into
mainfrom
nina-xu/save-transformed-training-before-preflight

nina-xu commented Jul 17, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented Jul 17, 2026 •

edited

Loading

Uh oh!

codecov Bot commented Jul 17, 2026

Uh oh!

greptile-apps Bot commented Jul 17, 2026

Sequence Diagram

Uh oh!

coderabbitai Bot left a comment

Uh oh!

coderabbitai Bot Jul 17, 2026

Uh oh!

binaryaaron Jul 17, 2026 •

edited

Loading

Uh oh!

nina-xu Jul 20, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

-                self._test_df.to_csv(self._workdir.dataset.test, index=False)
+                self._test_df.to_csv(self._workdir.dataset.test, index=False)
+            logger.user.info(
+                f"Processed training data saved to {self._workdir.dataset.training} "
+                "(available for inspection even if the preflight check below fails)",
+                extra={"path": str(self._workdir.dataset.training)},
+            )

Uh oh!

Conversation

nina-xu commented Jul 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Pre-Review Checklist

Pre-Merge Checklist

Other Notes

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Jul 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Uh oh!

codecov Bot commented Jul 17, 2026

Codecov Report

Uh oh!

greptile-apps Bot commented Jul 17, 2026

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Sequence Diagram

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

binaryaaron Jul 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nina-xu Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

nina-xu commented Jul 17, 2026 •

edited

Loading

coderabbitai Bot commented Jul 17, 2026 •

edited

Loading

binaryaaron Jul 17, 2026 •

edited

Loading