Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,9 @@ results = designer.create(
dataset_name="training_data",
resume=ResumeMode.IF_POSSIBLE,
)

print(results.dataset_path)
print(results.requested_resume_mode, results.effective_resume_mode)
```

Resume modes:
Expand All @@ -241,7 +244,7 @@ Resume has a few important invariants:
- Once `process_after_generation()` has run, the dataset is considered terminal for resume. Re-running with the same target returns the existing dataset; extending requires a fresh run.
- If a run crashed after every row group was written but before `process_after_generation()` could start, resume runs after-generation on the existing on-disk dataset (the parquet files are still clean) and marks it terminal afterwards. A crash _during_ `process_after_generation()` still raises β€” the parquet files may have been partially rewritten and starting fresh is the only safe option.

The `DatasetCreationResults` returned by a resume invocation reflects the full dataset on disk for anything that reads the artifact directory (`load_dataset`, `count_records`, `load_analysis`, `export`, `push_to_hub`). Per-run observability β€” `task_traces`, model-usage logs, and telemetry events emitted during the call β€” is scoped to the resume invocation only; the original run's in-memory traces are not persisted across process boundaries.
The `DatasetCreationResults` returned by a resume invocation reflects the full dataset on disk for anything that reads the artifact directory (`load_dataset`, `count_records`, `load_analysis`, `export`, `push_to_hub`). `dataset_path` is the resolved output directory, which may differ from the requested name when Data Designer starts fresh. Compare `requested_resume_mode` with `effective_resume_mode` to determine whether the invocation resumed or started fresh. Per-run observability β€” `task_traces`, model-usage logs, and telemetry events emitted during the call β€” is scoped to the resume invocation only; the original run's in-memory traces are not persisted across process boundaries.

<Warning>
Only resume datasets from trusted artifact directories. Resume reads local `metadata.json`, `builder_config.json`, and parquet files to determine checkpoint state.
Expand Down Expand Up @@ -385,18 +388,23 @@ config_builder.add_model_config(

### Run outcomes

A run can finish with fewer records than requested when non-retryable errors drop rows. Inspect `len(result.load_dataset())` to detect.
A run can finish with fewer records than requested when non-retryable errors drop rows. Compare `actual_num_records` with `requested_num_records`, or inspect `is_partial`, without loading the dataset.

If the rate of non-retryable errors crosses `RunConfig.shutdown_error_rate`, generation stops early and raises `DataDesignerEarlyShutdownError` (a subclass of `DataDesignerGenerationError`). Catch it separately when a typed retry path is appropriate:
If the rate of non-retryable errors crosses `RunConfig.shutdown_error_rate`, generation stops early. A run with no surviving records raises `DataDesignerEarlyShutdownError` (a subclass of `DataDesignerGenerationError`). If some records survive, `create()` returns a partial result with `early_shutdown=True`.

```python
from data_designer.interface.errors import DataDesignerEarlyShutdownError
from data_designer.interface import DataDesignerEarlyShutdownError

try:
result = dd_instance.create(config_builder, num_records=1000)
except DataDesignerEarlyShutdownError:
# e.g. retry against a different model alias
...
else:
if result.early_shutdown:
print(f"Early shutdown preserved {result.actual_num_records} records")
if result.is_partial:
print(f"Generated {result.actual_num_records} of {result.requested_num_records} requested records")
```

## Local OpenTelemetry Metrics
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,11 @@
from data_designer.config.config_builder import DataDesignerConfigBuilder # noqa: F401
from data_designer.config.custom_column import custom_column_generator # noqa: F401
from data_designer.config.data_designer_config import DataDesignerConfig # noqa: F401
from data_designer.config.errors import ( # noqa: F401
InvalidConfigError,
InvalidFileFormatError,
InvalidFilePathError,
)
from data_designer.config.mcp import ( # noqa: F401
LocalStdioMCPProvider,
MCPProvider,
Expand Down Expand Up @@ -160,6 +165,10 @@
"custom_column_generator": (f"{_MOD_BASE}.custom_column", "custom_column_generator"),
# data_designer_config
"DataDesignerConfig": (f"{_MOD_BASE}.data_designer_config", "DataDesignerConfig"),
# errors
"InvalidConfigError": (f"{_MOD_BASE}.errors", "InvalidConfigError"),
"InvalidFileFormatError": (f"{_MOD_BASE}.errors", "InvalidFileFormatError"),
"InvalidFilePathError": (f"{_MOD_BASE}.errors", "InvalidFilePathError"),
# mcp
"LocalStdioMCPProvider": (_MOD_MCP, "LocalStdioMCPProvider"),
"MCPProvider": (_MOD_MCP, "MCPProvider"),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import logging
from pathlib import Path

import yaml
from pydantic import model_validator
from pygments import highlight
from pygments.formatters import HtmlFormatter
Expand All @@ -24,7 +25,13 @@
)
from data_designer.config.data_designer_config import DataDesignerConfig
from data_designer.config.default_model_settings import get_default_model_configs
from data_designer.config.errors import BuilderConfigurationError, BuilderSerializationError, InvalidColumnTypeError
from data_designer.config.errors import (
BuilderConfigurationError,
BuilderSerializationError,
InvalidColumnTypeError,
InvalidFileFormatError,
InvalidFilePathError,
)
from data_designer.config.exportable_config import ExportableConfigBase
from data_designer.config.mcp import ToolConfig
from data_designer.config.models import ModelConfig, load_model_configs
Expand Down Expand Up @@ -111,11 +118,19 @@ def from_config(cls, config: dict | str | Path | BuilderConfig) -> Self:
Raises:
ValueError: If the config format is invalid.
ValidationError: If the builder config loaded from the config is invalid.
InvalidFilePathError: If a local config path cannot be read.
InvalidFileFormatError: If a local or inline config contains malformed YAML or JSON.
"""
if isinstance(config, BuilderConfig):
builder_config = config
else:
json_config = json.loads(serialize_data(smart_load_yaml(config)))
try:
loaded_config = smart_load_yaml(config)
except OSError as e:
raise InvalidFilePathError(f"Failed to load builder config: {e}") from e
except (UnicodeError, yaml.YAMLError) as e:
raise InvalidFileFormatError(f"Failed to parse builder config: {e}") from e
json_config = json.loads(serialize_data(loaded_config))
# Normalize shorthand DataDesignerConfig into full BuilderConfig
if "columns" in json_config and "data_designer" not in json_config:
json_config = {"data_designer": json_config}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -190,10 +190,15 @@ def _smart_load_yaml_internal(yaml_in: str | Path | dict, *, from_url: bool) ->
with open(yaml_in) as file:
yaml_out = yaml.safe_load(file)
elif isinstance(yaml_in, str):
if not from_url and yaml_in.endswith((".yaml", ".yml")) and not os.path.isfile(yaml_in):
raise FileNotFoundError(f"File not found: {yaml_in}")
else:
is_local_config_path = not from_url and yaml_in.lower().endswith((".yaml", ".yml", ".json"))
try:
yaml_out = yaml.safe_load(yaml_in)
except yaml.YAMLError as e:
if is_local_config_path:
raise FileNotFoundError(f"File not found: {yaml_in}") from e
raise
if is_local_config_path and not isinstance(yaml_out, dict):
raise FileNotFoundError(f"File not found: {yaml_in}")
else:
raise ValueError(
f"'{yaml_in}' is an invalid yaml config format. Valid options are: dict, yaml string, or yaml file path."
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,68 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

from pathlib import Path

import pytest

import data_designer.config.errors as config_errors
from data_designer.config import (
DataDesignerConfigBuilder,
InvalidConfigError,
InvalidFileFormatError,
InvalidFilePathError,
)


def test_invalid_config_error_is_publicly_exported() -> None:
assert InvalidConfigError is config_errors.InvalidConfigError


def test_builder_file_errors_are_publicly_exported() -> None:
assert InvalidFileFormatError is config_errors.InvalidFileFormatError
assert InvalidFilePathError is config_errors.InvalidFilePathError


@pytest.mark.parametrize("filename", ["missing.yaml", "missing.json", "missing.YAML", "missing.JsOn", "[broken.JSON"])
def test_from_config_normalizes_missing_file_error(tmp_path: Path, filename: str) -> None:
with pytest.raises(InvalidFilePathError) as exc_info:
DataDesignerConfigBuilder.from_config(str(tmp_path / filename))

assert isinstance(exc_info.value.__cause__, FileNotFoundError)


@pytest.mark.parametrize("extension", ["json", "yaml", "yml"])
def test_from_config_accepts_inline_config_ending_in_supported_extension(extension: str) -> None:
builder = DataDesignerConfigBuilder.from_config(
f"""
model_configs:
- alias: stub-model
model: stub-model
provider: provider-1
columns:
- name: category
column_type: sampler
sampler_type: category
params:
values:
- value.{extension}"""
)

assert builder.get_column_config("category").params.values == [f"value.{extension}"]


def test_from_config_normalizes_undecodable_file_error(tmp_path: Path) -> None:
config_path = tmp_path / "invalid.yaml"
config_path.write_bytes(b"\xff")

with pytest.raises(InvalidFileFormatError) as exc_info:
DataDesignerConfigBuilder.from_config(config_path)

assert isinstance(exc_info.value.__cause__, UnicodeDecodeError)


def test_from_config_normalizes_malformed_config_error() -> None:
with pytest.raises(InvalidFileFormatError):
DataDesignerConfigBuilder.from_config("data_designer: [")
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@

if TYPE_CHECKING:
from data_designer.config.run_config import ResumeMode # noqa: F401
from data_designer.engine.dataset_builders.errors import ArtifactStorageError # noqa: F401
from data_designer.interface.composite_workflow import ( # noqa: F401
CompositeWorkflow,
CompositeWorkflowResults,
Expand All @@ -24,6 +25,7 @@
from data_designer.interface.results import DatasetCreationResults # noqa: F401

_LAZY_IMPORTS: dict[str, tuple[str, str]] = {
"ArtifactStorageError": ("data_designer.engine.dataset_builders.errors", "ArtifactStorageError"),
"CompositeWorkflow": ("data_designer.interface.composite_workflow", "CompositeWorkflow"),
"CompositeWorkflowResults": ("data_designer.interface.composite_workflow", "CompositeWorkflowResults"),
"DataDesigner": ("data_designer.interface.data_designer", "DataDesigner"),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -369,6 +369,9 @@ def run(
stage=stage,
stage_dir_name=stage_dir_name,
stage_builder=stage_builder,
requested_num_records=stage_metadata.get(
"num_records_actual" if stage.output_processors else "num_records_requested"
),
)
stage_results[stage.name] = output_result
stage_output_paths[stage.name] = output_seed_path
Expand Down Expand Up @@ -671,6 +674,7 @@ def _stage_result_from_metadata(
stage: _WorkflowStage,
stage_dir_name: str,
stage_builder: DataDesignerConfigBuilder,
requested_num_records: int | None,
) -> DatasetCreationResults:
main_storage = ArtifactStorage(artifact_path=workflow_path, dataset_name=stage_dir_name, resume=ResumeMode.ALWAYS)
result_storage = main_storage
Expand All @@ -691,6 +695,7 @@ def _stage_result_from_metadata(
analysis=_load_stage_analysis(result_storage),
config_builder=result_builder,
dataset_metadata=DatasetMetadata(),
requested_num_records=requested_num_records,
)


Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -266,10 +266,12 @@ def create(

Returns:
DatasetCreationResults object with methods for loading the generated dataset,
analysis results, and displaying sample records for inspection.
analysis results, displaying sample records, and inspecting completion,
artifact, and resume metadata.

Raises:
DataDesignerGenerationError: If an error occurs during dataset generation.
DataDesignerEarlyShutdownError: If early shutdown produces no records.
DataDesignerProfilingError: If an error occurs during dataset profiling.
"""
logger.info("🎨 Creating Data Designer dataset")
Expand Down Expand Up @@ -378,6 +380,10 @@ def create(
config_builder=config_builder,
dataset_metadata=dataset_metadata,
task_traces=task_traces,
requested_num_records=num_records,
early_shutdown=builder.early_shutdown,
requested_resume_mode=resume,
effective_resume_mode=builder.artifact_storage.resume,
)

async def acreate(
Expand Down
59 changes: 54 additions & 5 deletions packages/data-designer/src/data_designer/interface/results.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from data_designer.config.config_builder import DataDesignerConfigBuilder
from data_designer.config.dataset_metadata import DatasetMetadata
from data_designer.config.errors import InvalidFileFormatError
from data_designer.config.run_config import ResumeMode
from data_designer.config.seed_source_dataframe import DataFrameSeedSource
from data_designer.config.utils.visualization import WithRecordSamplerMixin
from data_designer.engine.dataset_builders.errors import ArtifactStorageError
Expand Down Expand Up @@ -50,6 +51,10 @@ def __init__(
config_builder: DataDesignerConfigBuilder,
dataset_metadata: DatasetMetadata,
task_traces: list[TaskTrace] | None = None,
requested_num_records: int | None = None,
early_shutdown: bool | None = None,
requested_resume_mode: ResumeMode | None = None,
effective_resume_mode: ResumeMode | None = None,
):
"""Creates a new instance with results based on a dataset creation run.

Expand All @@ -62,12 +67,40 @@ def __init__(
Resume note: only contains traces for the current invocation; traces
from earlier ``create()`` calls that this run resumed are not
retained.
requested_num_records: Number of records requested for this invocation.
early_shutdown: Whether generation stopped at the early-shutdown gate.
``None`` when no generation invocation produced this result object.
requested_resume_mode: Resume mode requested for this invocation, or
``None`` when no generation invocation produced this result object.
effective_resume_mode: Resume mode selected after compatibility checks,
or ``None`` when no generation invocation produced this result object.
"""
self.artifact_storage = artifact_storage
self._analysis = analysis
self._config_builder = config_builder
self.dataset_metadata = dataset_metadata
self.task_traces: list[TaskTrace] = task_traces or []
self.requested_num_records = requested_num_records
self.early_shutdown = early_shutdown
self.requested_resume_mode = requested_resume_mode
self.effective_resume_mode = effective_resume_mode

@property
def dataset_path(self) -> Path:
"""Return the resolved dataset directory for this result."""
return self.artifact_storage.base_dataset_path

@property
def actual_num_records(self) -> int:
"""Return the total number of records in the result dataset."""
return self.count_records()

@property
def is_partial(self) -> bool | None:
"""Return whether the result contains fewer records than requested."""
if self.requested_num_records is None:
return None
return self.actual_num_records < self.requested_num_records

def load_analysis(self) -> DatasetProfilerResults:
"""Load the profiling analysis results for the generated dataset.
Expand All @@ -83,8 +116,17 @@ def load_dataset(self) -> pd.DataFrame:

Returns:
A pandas DataFrame containing the full generated dataset.

Raises:
ArtifactStorageError: If the dataset artifacts are missing.
"""
return self.artifact_storage.load_dataset()
try:
dataset = self.artifact_storage.load_dataset()
except OSError as e:
raise ArtifactStorageError(f"Failed to load dataset artifacts: {e}") from e
if dataset is not None and dataset.empty:
self._get_batch_files()
return dataset

def to_config_builder(self, columns: list[str] | None = None) -> DataDesignerConfigBuilder:
"""Create a new config builder seeded from this result dataset.
Expand All @@ -108,10 +150,19 @@ def count_records(self) -> int:

Returns:
Total row count across all batch parquet files.

Raises:
ArtifactStorageError: If the dataset artifacts are missing.
"""
batch_files = sorted(self.artifact_storage.final_dataset_path.glob("batch_*.parquet"))
batch_files = self._get_batch_files()
return sum(lazy.pq.read_metadata(f).num_rows for f in batch_files)

def _get_batch_files(self) -> list[Path]:
batch_files = sorted(self.artifact_storage.final_dataset_path.glob("batch_*.parquet"))
if not batch_files:
raise ArtifactStorageError("No batch parquet files found.")
return batch_files

def load_processor_dataset(self, processor_name: str) -> pd.DataFrame:
"""Load the dataset generated by a processor.

Expand Down Expand Up @@ -178,9 +229,7 @@ def export(self, path: Path | str, *, format: ExportFormat | None = None) -> Pat
raise InvalidFileFormatError(
f"Unsupported export format: {resolved_format!r}. Choose one of: {', '.join(SUPPORTED_EXPORT_FORMATS)}."
)
batch_files = sorted(self.artifact_storage.final_dataset_path.glob("batch_*.parquet"))
if not batch_files:
raise ArtifactStorageError("No batch parquet files found to export.")
batch_files = self._get_batch_files()
if resolved_format == "jsonl":
_export_jsonl(batch_files, path)
elif resolved_format == "csv":
Expand Down
Loading
Loading