fix: prevent dataset folder collisions

[andreatnvidia]

Summary

Fresh create() calls can select the same artifact directory and overwrite each other's batches, causing results to include another run's rows. Claim each fresh directory atomically after the resume decision so repeated or concurrent calls cannot share it.

Related Issue

Fixes #1050

Follow-up refactor: #1056 (low priority).

Changes

• Retry directory collisions with a timestamp and random suffix, then update the cached directory name and media paths.
• Keep resumed runs in their existing directories and avoid allocating directories for previews.
• Correct the folder-naming example in the create() docstring.

Testing

• Storage compatibility checks on Pydantic 2.9.2: fresh, empty, and occupied roots; result isolation; media paths; resume reuse.
• .venv/bin/pytest -q packages/data-designer-engine/tests: 2,489 passed.
• .venv/bin/pytest -q packages/data-designer/tests: 1,186 passed, 3 skipped.
• Regression coverage for same-timestamp calls, concurrent processes on fresh and occupied roots, empty directories, resume fallback, filesystem errors, and previews.
• Ruff lint and formatting checks passed for packages/, scripts/, and tests_e2e/; commit hooks passed.
• E2E tests: N/A; the regression tests exercise real filesystem allocation and generation without model endpoints.

Checklist

• Follows commit message conventions.
• Commits are signed off (DCO).
• External-contributor triage: N/A; repository collaborator.
• Architecture docs: N/A; the existing fresh-run and resume contracts remain accurate.

[github-actions]

Fern preview: https://nvidia-preview-pr-1051.docs.buildwithfern.com/nemo/datadesigner

Fern previews include the docs-website version archive with PR changes synced into latest. Notebook tutorials are rendered without execution outputs in previews.

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no new actionable issues remain.

Summary

Fresh runs now claim dataset directories atomically instead of sharing a folder when names collide.

• Each fresh dataset run claims its own folder before saving artifacts.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Dataset build] --> B{Resume decision}
    B -->|Fresh run| C[Claim directory exclusively]
    C -->|Already exists| D[Choose timestamp and random suffix]
    D --> C
    C -->|Claimed| E[Update cached name and media paths]
    B -->|Resume| F[Reuse existing directory]
    E --> G[Write config and generate rows]
    F --> G

Loading

_{Reviews (4) · Last reviewed commit: "Merge branch 'main' into andreatnvidia/f..." · Reviewed by Greptile}

[github-actions]

Thanks for putting this together, @andreatnvidia. It's a tight, well-scoped fix for a nasty data-mixing bug.

Summary

Before this PR, a fresh create() run could pick the same artifact directory as an earlier or concurrent run: the timestamped name only had one-second resolution, and an empty directory got reused without checks. The two runs then overwrote each other's batch files. The PR adds ArtifactStorage.create_dataset_directory(), which claims the directory with mkdir(exist_ok=False). On a collision it retries with a timestamp plus a random suffix, then updates the cached resolved_dataset_name and the media storage paths. It's called once in DatasetBuilder.build(), after the resume decision and the post-generation short-circuit. Resumed runs keep using their existing directory, and previews create no directory. The code matches what the PR description says.

Findings

Suggestions — Take it or leave it

packages/data-designer-engine/src/data_designer/engine/storage/artifact_storage.py:169 — resume is passed in even though the storage already holds it

• What: create_dataset_directory(resume=...) takes a resume argument, but ArtifactStorage already has a resume field. DatasetBuilder.build() keeps the two in sync on the IF_POSSIBLE paths.
• Why: Two sources of truth can drift. Say the storage was built with resume=NEVER (so resolved_dataset_name may already be a timestamped name) and the builder calls this with resume=ALWAYS. The method would then silently mkdir_if_needed a path that was resolved under fresh-run rules. DataDesigner.create() always passes matching values, so this can't happen today. It's just a sharp edge for direct callers of DatasetBuilder.
• Suggestion: We could read self.resume instead of taking the argument. Or keep the argument as the explicit decision and assert or sync self.resume = resume at the top, so the cached name and the claim always use the same rules.

packages/data-designer-engine/src/data_designer/engine/storage/artifact_storage.py:180 — Timestamp naming is duplicated

• What: The f"{self.dataset_name}_{datetime.now().strftime('%m-%d-%Y_%H%M%S')}" pattern now appears both in resolved_dataset_name and in the retry branch (the retry adds the _uuid suffix).
• Why: If someone changes the folder format later (for example to match the docstring example that this PR just corrected), it's easy to update one copy and miss the other.
• Suggestion: Pull it into a small private helper, e.g. _timestamped_dataset_name(self, suffix: str | None = None) -> str, and use it in both places.

packages/data-designer-engine/src/data_designer/engine/storage/artifact_storage.py:175 — The retry loop has no upper bound

• What: while True keeps retrying for as long as mkdir raises FileExistsError.
• Why: With 8 hex characters of UUID this is practically unreachable. Still, if a filesystem reports FileExistsError for some other reason (some network or FUSE mounts behave oddly), the run would spin forever instead of failing loudly.
• Suggestion: Cap the retries (say, 10 attempts) and then raise ArtifactStorageError. It's a cheap safety net.

packages/data-designer-engine/src/data_designer/engine/storage/artifact_storage.py:170 — The public method's docstring is minimal

• What: create_dataset_directory is a public method on a core abstraction, but its docstring is one line with no Args/Raises sections.
• Why: Per STYLEGUIDE.md, public methods get Google-style docstrings. Callers would benefit from knowing that it raises ArtifactStorageError on filesystem failures and that it changes base_dataset_path as a side effect.
• Suggestion: Add Args: resume, a Raises: ArtifactStorageError line, and a note that resolved_dataset_name and the media storage paths may be updated.

packages/data-designer-engine/src/data_designer/engine/storage/artifact_storage.py:190 — Writes the cached_property slot directly

• What: self.__dict__["resolved_dataset_name"] = ... relies on how functools.cached_property stores its value internally.
• Why: The builder already pops this key elsewhere, so this follows an existing local idiom. It does spread that coupling across a second module, though.
• Suggestion: Optional: a tiny private setter or a "claimed name" PrivateAttr that resolved_dataset_name checks first would keep the cache handling in one place. Fine to defer.

What Looks Good

• Correct primitive. mkdir(exist_ok=False) is atomic on POSIX filesystems, which makes it the right way to claim a directory across processes, not just across threads. The spawn-based multiprocess test with a Barrier exercises that race for real instead of simulating it with mocks.
• Careful placement. The claim happens after the IF_POSSIBLE downgrade/upgrade and after the post-generation short-circuit. Resumed runs, already-complete datasets, and previews therefore create no directory, and a test covers each case. Wrapping OSError in ArtifactStorageError follows the "normalize errors at boundaries" principle.
• Thorough tests. The suite covers same-timestamp sequential runs end-to-end through DataDesigner.create(), including a row-level check that each run's tag stays isolated. It also covers occupied and empty directories, the resume reuse path, filesystem errors, and the media path refresh. The docstring fix (awesome_dataset_01-01-2025_120000) now matches the real format, and architecture/dataset-builders.md and the Fern resume docs still describe the behavior accurately.

Residual Risk

• Behavior change: under ResumeMode.NEVER (and when IF_POSSIBLE falls back to a fresh run), an empty existing directory is no longer reused, so it gets a suffixed sibling instead. This is intended and documented in the updated create() docstring. Users who pre-create empty output folders, or who rely on a fixed folder name, will now see a timestamp-and-suffix folder. Worth a line in the release notes.
• Composite workflows: they rebuild stage storage from the fixed name stage_dir_name (_stage_result_from_metadata). If two workflow processes ever ran concurrently under the same name, a renamed stage directory wouldn't be found on a later skip or resume. That's outside this PR's scope, and the sequential case is unaffected because the stage directory is rmtree'd before each fresh run.
• Local validation: Ruff check and format pass on all 5 changed files. The 11 new tests pass, as do the existing suites for test_dataset_builder.py, engine/storage/, and data-designer/tests/interface/, including the composite workflow tests (387 passed).

Structural Impact (graphify 0.9.90, 8.1s)

Risk: HIGH (3 core abstraction(s) modified)

• 5 Python files, 158 AST entities, 15/220 clusters

Core Abstractions Modified

• DataDesigner - CustomColumnGenerator for quick customization #6 most connected (66 deps)
• DatasetBuilder - Publish to dataset HF Hub  #7 most connected (63 deps)
• ArtifactStorage - Make column generators/profilers/processors more uniform in config building and validation #10 most connected (60 deps)

High-Connectivity Changes

• dataset_builder.py (84 deps) in packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py
• DataDesigner (66 deps) in packages/data-designer/src/data_designer/interface/data_designer.py
• DatasetBuilder (63 deps) in packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py
• ArtifactStorage (60 deps) in packages/data-designer-engine/src/data_designer/engine/storage/artifact_storage.py
• data_designer.py (53 deps) in packages/data-designer/src/data_designer/interface/data_designer.py
• artifact_storage.py (29 deps) in packages/data-designer-engine/src/data_designer/engine/storage/artifact_storage.py
• ._prepare_async_run() (18 deps) in packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py
• .build() (18 deps) in packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py
• +52 more

Cross-Package Dependencies

• _stage_result_from_metadata() (interface) --uses--> ArtifactStorage (engine)
• _load_stage_analysis() (interface) --uses--> ArtifactStorage (engine)
• DataDesigner (interface) --uses--> DatasetProfilerResults (config)
• DataDesigner (interface) --uses--> DataDesignerConfigBuilder (config)
• DataDesigner (interface) --uses--> DataDesignerConfig (config)
• DataDesigner (interface) --uses--> DataDesignerInterface (config)
• +92 more

The HIGH rating comes from touching central classes. The actual change is additive: one new method on ArtifactStorage and one call site in DatasetBuilder.build(). The DataDesigner edit only touches a docstring. No existing signatures change and there are no import-direction violations.

Verdict

Ship it (with nits): there are no blocking issues. The suggestions above (one source of truth for resume, a shared timestamp helper, a retry cap, and a fuller docstring) are optional polish.

---

This review was generated by an AI assistant.

[greptile-apps]

Fresh dataset creation fails

resolved_dataset_name is a cached_property, not a Pydantic field. Pydantic 2.9 and 2.10 reject this direct assignment with ValueError, and both versions are allowed by pydantic>=2.9.2,<3. Every fresh create() call therefore fails after creating its directory in those environments.

Keep the direct __dict__ update, or require a Pydantic version that supports assigning to cached properties.

Suggested change

	self.resolved_dataset_name = dataset_path.name
	self.__dict__["resolved_dataset_name"] = dataset_path.name

Prompt To Fix With AI

This is a comment left during a code review.
Path: packages/data-designer-engine/src/data_designer/engine/storage/artifact_storage.py
Line: 190

Comment:
**Fresh dataset creation fails**

`resolved_dataset_name` is a `cached_property`, not a Pydantic field. Pydantic 2.9 and 2.10 reject this direct assignment with `ValueError`, and both versions are allowed by `pydantic>=2.9.2,<3`. Every fresh `create()` call therefore fails after creating its directory in those environments.

Keep the direct `__dict__` update, or require a Pydantic version that supports assigning to cached properties.

```suggestion
        self.__dict__["resolved_dataset_name"] = dataset_path.name
```

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

[nabinchha]

Looks good overall. Please restore the __dict__ assignment for compatibility with Pydantic 2.9.2—my earlier suggestion was verified only against a newer version.

Once that’s addressed, I’m happy to approve. Could we also track a low-priority refactor to finalize the dataset directory after the resume decision and before initializing media storage? That would remove the need to rewrite the cached name and refresh media paths.

[andreatnvidia]

Restored the __dict__ assignment in 30674c0 and verified it with Pydantic 2.9.2. The directory lifecycle refactor is tracked in #1056.