Skip to content

Commit 650134f

Browse files
committed
fix: claim dataset folders atomically
Prevent fresh generation runs from sharing an artifact directory when their folder timestamps collide. Keep resumed runs in their existing directories and refresh media paths after claiming a fresh directory. Fixes #1050 Signed-off-by: Andre Manoel <amanoel@nvidia.com>
1 parent c8aa954 commit 650134f

5 files changed

Lines changed: 160 additions & 4 deletions

File tree

‎packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -328,6 +328,7 @@ def build(
328328
if self._post_generation_processed_resume_result(resume, num_records) is not None:
329329
return self.artifact_storage.final_dataset_path
330330

331+
self.artifact_storage.create_dataset_directory(resume=resume)
331332
self._write_builder_config()
332333

333334
# Set media storage mode based on parameters

‎packages/data-designer-engine/src/data_designer/engine/storage/artifact_storage.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
from functools import cached_property
1212
from pathlib import Path
1313
from typing import TYPE_CHECKING
14+
from uuid import uuid4
1415

1516
from pydantic import BaseModel, ConfigDict, PrivateAttr, field_validator, model_validator
1617

@@ -165,6 +166,30 @@ def refresh_media_storage_path(self) -> None:
165166
self._media_storage.base_path = self.base_dataset_path
166167
self._media_storage.images_dir = self.base_dataset_path / images_subdir
167168

169+
def create_dataset_directory(self, *, resume: ResumeMode) -> None:
170+
"""Exclusively claim a fresh dataset directory, or reuse it when resuming."""
171+
if resume != ResumeMode.NEVER:
172+
self.mkdir_if_needed(self.base_dataset_path)
173+
return
174+
175+
dataset_path = self.base_dataset_path
176+
while True:
177+
try:
178+
dataset_path.mkdir(exist_ok=False)
179+
except FileExistsError:
180+
dataset_path = self.artifact_path / (
181+
f"{self.dataset_name}_{datetime.now().strftime('%m-%d-%Y_%H%M%S')}_{uuid4().hex[:8]}"
182+
)
183+
except OSError as exc:
184+
raise ArtifactStorageError(f"Cannot create dataset directory {str(dataset_path)!r}: {exc}") from exc
185+
else:
186+
break
187+
188+
if dataset_path != self.base_dataset_path:
189+
logger.info(f"📂 Dataset directory already claimed. This run will be saved to {str(dataset_path)!r}.")
190+
self.__dict__["resolved_dataset_name"] = dataset_path.name
191+
self.refresh_media_storage_path()
192+
168193
@staticmethod
169194
def mkdir_if_needed(path: Path | str) -> Path:
170195
"""Create the directory if it does not exist."""

‎packages/data-designer-engine/tests/engine/storage/test_artifact_storage.py‎

Lines changed: 94 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,12 @@
44
from __future__ import annotations
55

66
import json
7+
import multiprocessing as mp
78
from datetime import datetime
8-
from unittest.mock import patch
9+
from multiprocessing.queues import Queue
10+
from multiprocessing.synchronize import Barrier
11+
from pathlib import Path
12+
from unittest.mock import MagicMock, patch
913

1014
import pytest
1115
from pyarrow import ArrowNotImplementedError
@@ -502,3 +506,92 @@ def test_clear_partial_results_is_noop_when_no_partial_folder(tmp_path):
502506
storage = ArtifactStorage(artifact_path=tmp_path)
503507
assert not storage.partial_results_path.exists()
504508
storage.clear_partial_results() # must not raise
509+
510+
511+
@pytest.mark.parametrize("existing", [False, True], ids=["fresh", "occupied"])
512+
@patch("data_designer.engine.storage.artifact_storage.datetime")
513+
def test_create_dataset_directory_is_exclusive(mock_datetime: MagicMock, tmp_path: Path, existing: bool) -> None:
514+
mock_datetime.now.return_value = datetime(2025, 1, 1, 12, 3, 4)
515+
if existing:
516+
(tmp_path / "dataset").mkdir()
517+
(tmp_path / "dataset" / "sentinel").write_text("original")
518+
storages = [ArtifactStorage(artifact_path=tmp_path) for _ in range(3)]
519+
520+
for storage in storages:
521+
storage.create_dataset_directory(resume=ResumeMode.NEVER)
522+
assert storage.base_dataset_path.is_dir()
523+
assert storage.media_storage.base_path == storage.base_dataset_path
524+
assert storage.media_storage.images_dir == storage.base_dataset_path / "images"
525+
526+
assert len({storage.base_dataset_path for storage in storages}) == 3
527+
assert len(list(tmp_path.iterdir())) == 3 + existing
528+
if existing:
529+
assert (tmp_path / "dataset" / "sentinel").read_text() == "original"
530+
531+
532+
def test_create_dataset_directory_does_not_reuse_empty_directory(tmp_path: Path) -> None:
533+
storage = ArtifactStorage(artifact_path=tmp_path)
534+
(tmp_path / "dataset").mkdir()
535+
536+
storage.create_dataset_directory(resume=ResumeMode.NEVER)
537+
538+
assert storage.base_dataset_path != tmp_path / "dataset"
539+
assert storage.base_dataset_path.is_dir()
540+
assert list((tmp_path / "dataset").iterdir()) == []
541+
542+
543+
@pytest.mark.parametrize("resume", [ResumeMode.ALWAYS, ResumeMode.IF_POSSIBLE])
544+
def test_create_dataset_directory_reuses_resumed_directory(tmp_path: Path, resume: ResumeMode) -> None:
545+
dataset_path = tmp_path / "dataset"
546+
dataset_path.mkdir()
547+
(dataset_path / "sentinel").write_text("original")
548+
storage = ArtifactStorage(artifact_path=tmp_path, resume=resume)
549+
550+
storage.create_dataset_directory(resume=resume)
551+
552+
assert storage.base_dataset_path == dataset_path
553+
assert list(tmp_path.iterdir()) == [dataset_path]
554+
assert (dataset_path / "sentinel").read_text() == "original"
555+
556+
557+
def test_create_dataset_directory_reports_filesystem_error(tmp_path: Path) -> None:
558+
storage = ArtifactStorage(artifact_path=tmp_path)
559+
with patch.object(Path, "mkdir", side_effect=PermissionError("permission denied")):
560+
with pytest.raises(ArtifactStorageError, match="Cannot create dataset directory"):
561+
storage.create_dataset_directory(resume=ResumeMode.NEVER)
562+
563+
564+
def _claim_dataset_directory(artifact_path: Path, barrier: Barrier, output: Queue) -> None:
565+
with patch("data_designer.engine.storage.artifact_storage.datetime") as clock:
566+
clock.now.return_value = datetime(2025, 1, 1, 12, 3, 4)
567+
storage = ArtifactStorage(artifact_path=artifact_path)
568+
barrier.wait(timeout=30)
569+
storage.create_dataset_directory(resume=ResumeMode.NEVER)
570+
output.put(str(storage.base_dataset_path))
571+
572+
573+
@pytest.mark.parametrize("existing", [False, True], ids=["fresh", "occupied"])
574+
def test_create_dataset_directory_is_exclusive_across_processes(tmp_path: Path, existing: bool) -> None:
575+
if existing:
576+
(tmp_path / "dataset").mkdir()
577+
(tmp_path / "dataset" / "sentinel").write_text("original")
578+
context = mp.get_context("spawn")
579+
barrier = context.Barrier(2)
580+
output = context.Queue()
581+
processes = [context.Process(target=_claim_dataset_directory, args=(tmp_path, barrier, output)) for _ in range(2)]
582+
try:
583+
for process in processes:
584+
process.start()
585+
for process in processes:
586+
process.join(timeout=45)
587+
assert process.exitcode == 0
588+
paths = [Path(output.get(timeout=5)) for _ in processes]
589+
assert len(set(paths)) == 2
590+
assert all(path.is_dir() for path in paths)
591+
finally:
592+
for process in processes:
593+
if process.is_alive():
594+
process.terminate()
595+
process.join(timeout=5)
596+
output.close()
597+
output.join_thread()

‎packages/data-designer/src/data_designer/interface/data_designer.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -235,11 +235,12 @@ def create(
235235
configuration (columns, constraints, seed data, etc.).
236236
num_records: Number of records to generate.
237237
dataset_name: Name of the dataset. This name will be used as the dataset
238-
folder name in the artifact path directory. If a non-empty directory with the
238+
folder name in the artifact path directory. If a directory with the
239239
same name already exists, dataset will be saved to a new directory with
240240
a datetime stamp. For example, if the dataset name is "awesome_dataset" and a directory
241241
with the same name already exists, the dataset will be saved to a new directory
242-
with the name "awesome_dataset_2025-01-01_12-00-00".
242+
with the name "awesome_dataset_01-01-2025_120000". A random suffix may be added
243+
to avoid collisions.
243244
resume: Controls how interrupted runs are handled.
244245
245246
- ``ResumeMode.NEVER`` (default): always start a fresh generation run.

‎packages/data-designer/tests/interface/test_data_designer.py‎

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
from data_designer.config.errors import InvalidConfigError
2828
from data_designer.config.models import ChatCompletionInferenceParams, ModelConfig, ModelProvider
2929
from data_designer.config.processors import DropColumnsProcessorConfig
30-
from data_designer.config.run_config import JinjaRenderingEngine, RequestAdmissionTuningConfig, RunConfig
30+
from data_designer.config.run_config import JinjaRenderingEngine, RequestAdmissionTuningConfig, ResumeMode, RunConfig
3131
from data_designer.config.sampler_constraints import ScalarInequalityConstraint
3232
from data_designer.config.sampler_params import (
3333
CategorySamplerParams,
@@ -2468,3 +2468,39 @@ def test_create_raises_error_when_all_trace_files_are_skipped(
24682468

24692469
with pytest.raises(DataDesignerGenerationError, match="did not produce any rows"):
24702470
data_designer.create(builder, num_records=1, dataset_name="invalid-trace-seed")
2471+
2472+
2473+
@pytest.mark.parametrize("resume", [ResumeMode.NEVER, ResumeMode.IF_POSSIBLE])
2474+
@patch("data_designer.engine.storage.artifact_storage.datetime")
2475+
def test_create_same_timestamp_keeps_runs_isolated(
2476+
mock_datetime: MagicMock, tmp_path: Path, stub_model_providers: list[ModelProvider], resume: ResumeMode
2477+
) -> None:
2478+
mock_datetime.now.return_value = datetime(2025, 1, 1, 12, 3, 4)
2479+
designer = DataDesigner(artifact_path=tmp_path, model_providers=stub_model_providers, auto_configure_logging=False)
2480+
designer.set_run_config(RunConfig(buffer_size=10, display_tui=False))
2481+
results = []
2482+
for tag, num_records in [("A", 20), ("B", 10), ("C", 10)]:
2483+
builder = DataDesignerConfigBuilder(model_configs=[])
2484+
builder.add_column(
2485+
SamplerColumnConfig(name="tag", sampler_type="category", params=CategorySamplerParams(values=[tag]))
2486+
)
2487+
result = designer.create(builder, num_records=num_records, resume=resume)
2488+
results.append((tag, num_records, result))
2489+
2490+
assert len({result.artifact_storage.base_dataset_path for _, _, result in results}) == 3
2491+
assert len(list(tmp_path.iterdir())) == 3
2492+
for tag, num_records, result in results:
2493+
assert result.load_dataset()["tag"].tolist() == [tag] * num_records
2494+
assert result.artifact_storage.media_storage.base_path == result.artifact_storage.base_dataset_path
2495+
2496+
2497+
def test_preview_does_not_claim_dataset_directory(tmp_path: Path, stub_model_providers: list[ModelProvider]) -> None:
2498+
builder = DataDesignerConfigBuilder(model_configs=[])
2499+
builder.add_column(
2500+
SamplerColumnConfig(name="tag", sampler_type="category", params=CategorySamplerParams(values=["A"]))
2501+
)
2502+
designer = DataDesigner(artifact_path=tmp_path, model_providers=stub_model_providers, auto_configure_logging=False)
2503+
designer.set_run_config(RunConfig(display_tui=False))
2504+
2505+
assert designer.preview(builder, num_records=1).dataset["tag"].tolist() == ["A"]
2506+
assert list(tmp_path.iterdir()) == []

0 commit comments

Comments
 (0)