From 284152547d7c5a6d31cdc21dfcf3a94be1f4e9bb Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Mon, 10 Aug 2026 11:40:29 -0300 Subject: [PATCH 1/7] feat: define public generation contract --- .../src/data_designer/config/__init__.py | 3 + .../tests/config/test_public_contract.py | 11 + .../data_designer/interface/data_designer.py | 8 +- .../src/data_designer/interface/results.py | 28 ++ .../tests/interface/test_data_designer.py | 49 ++- .../tests/interface/test_results.py | 32 ++ plans/850/data-designer-contract.md | 301 ++++++++++++++++++ 7 files changed, 430 insertions(+), 2 deletions(-) create mode 100644 packages/data-designer-config/tests/config/test_public_contract.py create mode 100644 plans/850/data-designer-contract.md diff --git a/packages/data-designer-config/src/data_designer/config/__init__.py b/packages/data-designer-config/src/data_designer/config/__init__.py index c3fc1195f..2eeed4fde 100644 --- a/packages/data-designer-config/src/data_designer/config/__init__.py +++ b/packages/data-designer-config/src/data_designer/config/__init__.py @@ -32,6 +32,7 @@ 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 InvalidConfigError # noqa: F401 from data_designer.config.mcp import ( # noqa: F401 LocalStdioMCPProvider, MCPProvider, @@ -160,6 +161,8 @@ "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"), # mcp "LocalStdioMCPProvider": (_MOD_MCP, "LocalStdioMCPProvider"), "MCPProvider": (_MOD_MCP, "MCPProvider"), diff --git a/packages/data-designer-config/tests/config/test_public_contract.py b/packages/data-designer-config/tests/config/test_public_contract.py new file mode 100644 index 000000000..f0d30b5e1 --- /dev/null +++ b/packages/data-designer-config/tests/config/test_public_contract.py @@ -0,0 +1,11 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from data_designer.config import InvalidConfigError +from data_designer.config.errors import InvalidConfigError as InvalidConfigErrorDefinition + + +def test_invalid_config_error_is_publicly_exported() -> None: + assert InvalidConfigError is InvalidConfigErrorDefinition diff --git a/packages/data-designer/src/data_designer/interface/data_designer.py b/packages/data-designer/src/data_designer/interface/data_designer.py index e1b6961f7..b6a295661 100644 --- a/packages/data-designer/src/data_designer/interface/data_designer.py +++ b/packages/data-designer/src/data_designer/interface/data_designer.py @@ -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") @@ -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( diff --git a/packages/data-designer/src/data_designer/interface/results.py b/packages/data-designer/src/data_designer/interface/results.py index 6c5b076a3..512150916 100644 --- a/packages/data-designer/src/data_designer/interface/results.py +++ b/packages/data-designer/src/data_designer/interface/results.py @@ -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 @@ -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 = False, + requested_resume_mode: ResumeMode = ResumeMode.NEVER, + effective_resume_mode: ResumeMode | None = None, ): """Creates a new instance with results based on a dataset creation run. @@ -62,12 +67,35 @@ 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. + requested_resume_mode: Resume mode requested for this invocation. + effective_resume_mode: Resume mode selected after compatibility checks. """ 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 or requested_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: + """Return whether the result contains fewer records than requested.""" + return self.requested_num_records is not None and self.actual_num_records < self.requested_num_records def load_analysis(self) -> DatasetProfilerResults: """Load the profiling analysis results for the generated dataset. diff --git a/packages/data-designer/tests/interface/test_data_designer.py b/packages/data-designer/tests/interface/test_data_designer.py index 0a0ce5a9c..0ea707771 100644 --- a/packages/data-designer/tests/interface/test_data_designer.py +++ b/packages/data-designer/tests/interface/test_data_designer.py @@ -26,7 +26,7 @@ from data_designer.config.errors import InvalidConfigError from data_designer.config.models import ChatCompletionInferenceParams, ModelConfig, ModelProvider from data_designer.config.processors import DropColumnsProcessorConfig -from data_designer.config.run_config import JinjaRenderingEngine, RequestAdmissionTuningConfig, RunConfig +from data_designer.config.run_config import JinjaRenderingEngine, RequestAdmissionTuningConfig, ResumeMode, RunConfig from data_designer.config.sampler_params import CategorySamplerParams, DatetimeSamplerParams, SamplerType from data_designer.config.seed import IndexRange, PartitionBlock, SamplingStrategy from data_designer.config.seed_source import ( @@ -671,11 +671,58 @@ def test_create_dataset_e2e_using_only_sampler_columns( analysis = results.load_analysis() assert analysis.target_num_records == num_records + assert results.dataset_path == stub_artifact_path / "dataset" + assert results.requested_num_records == num_records + assert results.actual_num_records == num_records + assert results.is_partial is False + assert results.early_shutdown is False + assert results.requested_resume_mode == ResumeMode.NEVER + assert results.effective_resume_mode == ResumeMode.NEVER # display report with no errors analysis.to_report() +def test_create_reports_effective_resume_mode( + stub_sampler_only_config_builder, + stub_artifact_path, + stub_model_providers, + stub_managed_assets_path, +) -> None: + data_designer = DataDesigner( + artifact_path=stub_artifact_path, + model_providers=stub_model_providers, + secret_resolver=PlaintextResolver(), + managed_assets_path=stub_managed_assets_path, + ) + data_designer.set_run_config(RunConfig(buffer_size=1, otel_metrics_port=None)) + + first = data_designer.create(stub_sampler_only_config_builder, num_records=2, dataset_name="resumable") + resumed = data_designer.create( + stub_sampler_only_config_builder, + num_records=3, + dataset_name="resumable", + resume=ResumeMode.IF_POSSIBLE, + ) + + assert resumed.requested_resume_mode == ResumeMode.IF_POSSIBLE + assert resumed.effective_resume_mode == ResumeMode.ALWAYS + assert resumed.dataset_path == first.dataset_path + assert resumed.actual_num_records == 3 + + _add_irrelevant_sampler_column(stub_sampler_only_config_builder) + restarted = data_designer.create( + stub_sampler_only_config_builder, + num_records=3, + dataset_name="resumable", + resume=ResumeMode.IF_POSSIBLE, + ) + + assert restarted.requested_resume_mode == ResumeMode.IF_POSSIBLE + assert restarted.effective_resume_mode == ResumeMode.NEVER + assert restarted.dataset_path != first.dataset_path + + def test_create_with_drop_true_can_skip_dropped_column_artifacts( stub_artifact_path, stub_model_providers, diff --git a/packages/data-designer/tests/interface/test_results.py b/packages/data-designer/tests/interface/test_results.py index 200d43810..2bbe154af 100644 --- a/packages/data-designer/tests/interface/test_results.py +++ b/packages/data-designer/tests/interface/test_results.py @@ -15,6 +15,7 @@ from data_designer.config.dataset_metadata import DatasetMetadata from data_designer.config.errors import InvalidFileFormatError from data_designer.config.preview_results import PreviewResults +from data_designer.config.run_config import ResumeMode from data_designer.config.utils.errors import DatasetSampleDisplayError from data_designer.config.utils.visualization import display_sample_record as display_fn from data_designer.engine.dataset_builders.errors import ArtifactStorageError @@ -389,6 +390,37 @@ def test_count_records(stub_dataset_creation_results, stub_dataframe, stub_batch assert stub_dataset_creation_results.count_records() == len(stub_dataframe) +def test_public_creation_outcome_fields( + stub_artifact_storage, + stub_dataset_profiler_results, + stub_complete_builder, + stub_dataset_metadata, + stub_dataframe, + stub_batch_dir, + tmp_path, +) -> None: + dataset_path = tmp_path / "dataset" + stub_artifact_storage.base_dataset_path = dataset_path + stub_artifact_storage.final_dataset_path = stub_batch_dir + results = DatasetCreationResults( + artifact_storage=stub_artifact_storage, + analysis=stub_dataset_profiler_results, + config_builder=stub_complete_builder, + dataset_metadata=stub_dataset_metadata, + requested_num_records=len(stub_dataframe) + 1, + early_shutdown=True, + requested_resume_mode=ResumeMode.IF_POSSIBLE, + effective_resume_mode=ResumeMode.NEVER, + ) + + assert results.dataset_path == dataset_path + assert results.actual_num_records == len(stub_dataframe) + assert results.is_partial is True + assert results.early_shutdown is True + assert results.requested_resume_mode == ResumeMode.IF_POSSIBLE + assert results.effective_resume_mode == ResumeMode.NEVER + + def test_export_uppercase_extension_is_recognised(stub_dataset_creation_results, stub_batch_dir, tmp_path) -> None: """export() treats file extensions case-insensitively (e.g. .JSONL → jsonl).""" stub_dataset_creation_results.artifact_storage.final_dataset_path = stub_batch_dir diff --git a/plans/850/data-designer-contract.md b/plans/850/data-designer-contract.md new file mode 100644 index 000000000..e193c2d67 --- /dev/null +++ b/plans/850/data-designer-contract.md @@ -0,0 +1,301 @@ +--- +date: 2026-08-10 +authors: + - andreatnvidia +issue: https://github.com/NVIDIA-NeMo/DataDesigner/issues/851 +epic: https://github.com/NVIDIA-NeMo/DataDesigner/issues/850 +status: proposal +--- + +# Public Data Designer invocation contract + +## Summary + +An optional execution package can configure, validate, and invoke Data Designer without importing +`data_designer.engine`. The supported boundary consists of serialized builder input, public config models, +`DataDesigner`, `DatasetCreationResults`, and public errors. + +Plugin-independent envelope validation may run before optional plugin packages are installed. Full builder validation +must run in a fresh client process after those packages are installed and available on `sys.path`. + +## Public imports + +The execution package may import these public symbols: + +```python +from data_designer.config import ( + DataDesignerConfigBuilder, + InvalidConfigError, + LocalStdioMCPProvider, + MCPProvider, + ModelConfig, + ModelProvider, + PartitionBlock, + ResumeMode, + RunConfig, + SamplingStrategy, +) +from data_designer.interface import ( + DataDesigner, + DataDesignerEarlyShutdownError, + DataDesignerGenerationError, + DataDesignerProfilingError, + DatasetCreationResults, +) +from data_designer.plugins import Plugin +``` + +The package must not depend on: + +- Any `data_designer.engine` import. +- `data_designer.config.config_builder.BuilderConfig`. +- `data_designer.config.mcp.MCPProviderT`. +- `DatasetCreationResults.artifact_storage` or `task_traces`. +- The plugin registry implementation. +- Typed errors documented by `DataDesigner.check_models()` outside the public interface error module. + +`DataDesigner` may use engine modules internally. The restriction applies to imports and types crossing the package +boundary. + +## Serialized builder boundary + +Authored input contains exactly one of: + +- A path to a complete JSON or YAML builder config. +- The same builder config as an inline mapping. + +The payload is the format written by `DataDesignerConfigBuilder.write_config()`. The public loader is: + +```python +builder = DataDesignerConfigBuilder.from_config(builder_payload_or_path) +``` + +`BuilderConfig` is an implementation type used by the serializer. It is not part of the cross-package API. +Invocation controls such as record count, dataset name, resume, runtime settings, model endpoints, and output format +remain outside the builder payload. + +## Validation ownership + +### Submission environment + +The submission environment validates only information that does not require plugin-specific config classes: + +- Outer execution schema and version. +- Builder source location, format, digest, and top-level mapping shape. +- Unique aliases declared in the raw `data_designer.model_configs` list. +- Deployment coverage for those declared aliases. +- `RunConfig`, output, image, dependency, and resource fields owned by the execution package. + +It must not instantiate the complete builder when referenced plugins are unavailable. Plugin columns may declare +additional model references, so referenced-alias validation is deferred with full builder validation. + +### Client environment + +The package overlay must be installed before the client process imports Data Designer config unions. Plugin discovery +is process-global and happens on first use, so installing a package after config import is not a supported refresh +path. + +Client preflight runs in a fresh process and performs these steps: + +1. Enumerate the expected `data_designer.plugins` entry points with `importlib.metadata`. +2. Load each entry point and verify that it returns a public `Plugin` object. +3. Resolve `Plugin.config_cls` and `Plugin.impl_cls` so import failures stop preflight. +4. Load the builder with `DataDesignerConfigBuilder.from_config()`. +5. Materialize model bindings, per-model concurrency, seed input bindings, and MCP providers. +6. Verify referenced model aliases through each column config's public `get_model_aliases()` method. +7. Construct `DataDesigner`, apply `RunConfig`, and call `DataDesigner.validate(builder)`. + +The explicit entry-point load is required because normal plugin discovery logs and skips a failing entry point. +`DataDesigner.validate()` validates config structure and seed-dependent compilation without contacting model endpoints. +Endpoint readiness belongs after services are available; generation performs its normal readiness check. + +## Model binding + +`ModelConfig.alias` is the workload identity. The original `model` string is not an alias and must not be used to +match deployments. + +For every declared alias, the client creates a stable `ModelProvider` for the resolved logical endpoint and replaces +the matching `ModelConfig` while preserving unrelated inference parameters: + +```python +originals = list(builder.model_configs) +for model_config in originals: + builder.delete_model_config(model_config.alias) + +for model_config in originals: + binding = bindings_by_alias[model_config.alias] + inference_parameters = model_config.inference_parameters.model_copy( + update={"max_parallel_requests": binding.max_parallel_requests} + ) + builder.add_model_config( + model_config.model_copy( + update={ + "model": binding.served_model, + "provider": binding.provider_name, + "inference_parameters": inference_parameters, + } + ) + ) +``` + +Deleting all originals before adding replacements preserves alias uniqueness and original ordering. Full preflight +rejects missing or duplicate declared aliases, missing referenced aliases, missing providers, and deployment aliases +that do not match `ModelConfig.alias`. + +Per-model concurrency is `ModelConfig.inference_parameters.max_parallel_requests`. It is not a `RunConfig` field. + +## Runtime configuration + +The execution package merges its compatibility defaults with the raw authored `run_config` mapping, then validates +the effective mapping once: + +```python +effective_run_config = RunConfig.model_validate(compatibility_defaults | authored_run_config) +data_designer.set_run_config(effective_run_config) +``` + +Merging must occur before constructing the authored `RunConfig`. Model validators normalize related fields, so a +validated model's field-set metadata is not the authoritative record of which keys the user wrote. + +The effective default for every public field is: + +| Field | Data Designer default | Execution-package default | +| --- | ---: | ---: | +| `disable_early_shutdown` | `False` | `True` when no early-shutdown control is authored | +| `shutdown_error_rate` | `0.5` | `1.0` after disabled-shutdown normalization | +| `shutdown_error_window` | `10` | `10` | +| `buffer_size` | `1000` | `16384` | +| `max_concurrent_row_groups` | `3` | `3` | +| `max_in_flight_tasks` | `1024` | `1024` | +| `non_inference_max_parallel_workers` | `4` | `4` | +| `max_conversation_restarts` | `5` | `0` | +| `max_conversation_correction_steps` | `0` | `0` | +| `async_trace` | `False` | `False` | +| `write_scheduler_events` | `False` | `False` | +| `display_tui` | `False` | `False` | +| `progress_interval` | `5.0` | `5.0` | +| `otel_metrics_port` | `9464` | `None` | +| `preserve_dropped_columns` | `True` | `True` | +| `jinja_rendering_engine` | `secure` | `secure` | +| `request_admission` | `None` | `None` | + +Every authored key overrides the corresponding execution-package default. Early-shutdown controls form one related +group: if the user authors `shutdown_error_rate` or `shutdown_error_window` without `disable_early_shutdown`, the +package retains Data Designer's enabled-shutdown default instead of injecting `disable_early_shutdown=True` and +discarding the authored threshold. + +`non_inference_max_parallel_workers` is a public field but is not currently consumed by generation. The execution +package must not claim CPU-derived worker control until Data Designer implements that behavior. + +## Invocation + +The client materializes provider and MCP secrets immediately before constructing public config objects. Secret +values do not enter authored or persisted config. + +```python +data_designer = DataDesigner( + artifact_path=dataset_workspace, + model_providers=model_providers, + managed_assets_path=managed_assets_path, + mcp_providers=mcp_providers, + auto_configure_logging=False, +) +data_designer.set_run_config(effective_run_config) +data_designer.validate(builder) + +results = data_designer.create( + builder, + num_records=requested_num_records, + dataset_name=dataset_name, + resume=resume_mode, +) +``` + +`DataDesigner.acreate()` is the equivalent non-blocking entry point. It delegates generation to a worker thread and +returns the same `DatasetCreationResults` type. + +Remote MCP connections use `MCPProvider`; local subprocess connections use `LocalStdioMCPProvider`. The execution +package owns any secret-reference schema and passes only resolved strings to Data Designer. + +## Results and errors + +A normal return means generation and profiling completed without a public exception. It does not guarantee that the +dataset contains exactly the requested number of records. + +`DatasetCreationResults` exposes: + +| Field or method | Meaning | +| --- | --- | +| `dataset_path` | Resolved dataset directory, including collision or resume resolution. | +| `requested_num_records` | Target passed to the current `create()` invocation. | +| `actual_num_records` | Current total rows in the final dataset, including rows from an earlier resumed invocation. | +| `is_partial` | `True` when actual records are fewer than requested records. | +| `early_shutdown` | Whether the current invocation stopped through the early-shutdown gate. | +| `requested_resume_mode` | Resume mode passed to `create()`. | +| `effective_resume_mode` | `always` when the invocation resumed, otherwise `never` after compatibility resolution. | +| `count_records()` | Metadata-only row count, equivalent to `actual_num_records`. | +| `export(path, format=...)` | Stream the result to one JSONL, CSV, or Parquet file. | + +The caller derives its own exact-count policy from `actual_num_records == requested_num_records`. A partial result may +be caused by dropped rows or early shutdown; `early_shutdown` distinguishes those cases. + +Public failure behavior is: + +| Condition | Public behavior | +| --- | --- | +| Invalid serialized builder shape | `pydantic.ValidationError` or `ValueError` from `DataDesignerConfigBuilder.from_config()`. | +| Invalid compiled config | `InvalidConfigError` from `data_designer.config`. | +| Generation failure | `DataDesignerGenerationError`. | +| Profiling failure after generation | `DataDesignerProfilingError`. No successful result is returned. | +| Early shutdown with zero records | `DataDesignerEarlyShutdownError`, a `DataDesignerGenerationError` subclass. | +| Early shutdown with some records | A partial `DatasetCreationResults` with `early_shutdown=True`. | + +The caller must not inspect engine storage or task-trace types to classify an outcome. + +## Output, resume, seed, and telemetry + +- `artifact_path`, `dataset_name`, and `ResumeMode` define the resumable workspace. +- `dataset_path` is the resolved public location. This is important when `ResumeMode.IF_POSSIBLE` starts fresh or a + non-resumable name collides. +- Export format is explicit. Data Designer does not publish a default export-format constant. +- Seed datasets are applied with `with_seed_dataset()` and can be partitioned with public `PartitionBlock`. +- `PartitionBlock` does not define a deterministic generation random seed. No public generation-seed API exists. +- OpenTelemetry metrics are configured with `RunConfig.otel_metrics_port`; `None` disables metrics for the invocation. +- Embedded callers should use `auto_configure_logging=False` when they own process logging. + +## Callable and plugin limits + +Entry-point plugins with serializable config models are supported after installation in the client environment. + +`CustomColumnConfig.generator_function` and `LocalCallableValidatorParams.validation_function` are not portable +serialized references. Their serializers write only a function name while validation requires an in-memory callable. +Installing a package later does not resolve that string. A workload requiring these callables needs a separately +approved installed builder factory or a future qualified-callable-reference feature. + +## Import baseline + +The base import budget remains the existing `make perf-import CLEAN=1` measurement: + +```text +import data_designer.config as dd +from data_designer.interface import DataDesigner +``` + +The average of one cold and four warm runs must remain below three seconds. Optional command discovery may read +distribution metadata, but must not import the optional package until its command is selected. CLI extension discovery +is tracked separately by #853. + +## Contract tests + +The owning packages must cover: + +- Public root imports and absence of external `data_designer.engine` imports. +- Builder file and inline mapping loading after plugin installation. +- Plugin entry-point load failure before model services start. +- Model binding by alias, order preservation, and per-model concurrency. +- Compatibility-default merging from raw authored keys. +- Complete, partial, partial early-shutdown, zero-record early-shutdown, generation-error, and profiling-error outcomes. +- Resolved dataset paths for fresh, colliding, resumed, and incompatible `if_possible` invocations. +- Explicit JSONL, CSV, and Parquet exports. +- Remote and stdio MCP providers without persisted secret values. +- The existing import-performance threshold. From 63040a97a358b34d18fdfb5cc83ea49dcbd59c4e Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Mon, 10 Aug 2026 12:08:38 -0300 Subject: [PATCH 2/7] fix: complete public generation contract --- .../src/data_designer/config/__init__.py | 8 ++- .../data_designer/config/config_builder.py | 19 +++++- .../data_designer/config/utils/io_helpers.py | 2 +- .../tests/config/test_public_contract.py | 33 ++++++++++- .../src/data_designer/interface/__init__.py | 2 + .../interface/composite_workflow.py | 5 ++ .../src/data_designer/interface/results.py | 19 +++--- .../interface/test_composite_workflow.py | 33 +++++++++++ .../tests/interface/test_data_designer.py | 59 +++++++++++++++++++ .../tests/interface/test_results.py | 2 +- plans/850/data-designer-contract.md | 19 ++++-- 11 files changed, 180 insertions(+), 21 deletions(-) diff --git a/packages/data-designer-config/src/data_designer/config/__init__.py b/packages/data-designer-config/src/data_designer/config/__init__.py index 2eeed4fde..09c908474 100644 --- a/packages/data-designer-config/src/data_designer/config/__init__.py +++ b/packages/data-designer-config/src/data_designer/config/__init__.py @@ -32,7 +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 InvalidConfigError # noqa: F401 + from data_designer.config.errors import ( # noqa: F401 + InvalidConfigError, + InvalidFileFormatError, + InvalidFilePathError, + ) from data_designer.config.mcp import ( # noqa: F401 LocalStdioMCPProvider, MCPProvider, @@ -163,6 +167,8 @@ "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"), diff --git a/packages/data-designer-config/src/data_designer/config/config_builder.py b/packages/data-designer-config/src/data_designer/config/config_builder.py index f69e63706..428a111c9 100644 --- a/packages/data-designer-config/src/data_designer/config/config_builder.py +++ b/packages/data-designer-config/src/data_designer/config/config_builder.py @@ -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 @@ -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 @@ -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 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} diff --git a/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py b/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py index e71ec3a10..a9c91a5e9 100644 --- a/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py +++ b/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py @@ -190,7 +190,7 @@ 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): + if not from_url and yaml_in.endswith((".yaml", ".yml", ".json")) and not os.path.isfile(yaml_in): raise FileNotFoundError(f"File not found: {yaml_in}") else: yaml_out = yaml.safe_load(yaml_in) diff --git a/packages/data-designer-config/tests/config/test_public_contract.py b/packages/data-designer-config/tests/config/test_public_contract.py index f0d30b5e1..6ff326def 100644 --- a/packages/data-designer-config/tests/config/test_public_contract.py +++ b/packages/data-designer-config/tests/config/test_public_contract.py @@ -3,9 +3,36 @@ from __future__ import annotations -from data_designer.config import InvalidConfigError -from data_designer.config.errors import InvalidConfigError as InvalidConfigErrorDefinition +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 InvalidConfigErrorDefinition + 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"]) +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) + + +def test_from_config_normalizes_malformed_config_error() -> None: + with pytest.raises(InvalidFileFormatError): + DataDesignerConfigBuilder.from_config("data_designer: [") diff --git a/packages/data-designer/src/data_designer/interface/__init__.py b/packages/data-designer/src/data_designer/interface/__init__.py index d4a112f74..9717473da 100644 --- a/packages/data-designer/src/data_designer/interface/__init__.py +++ b/packages/data-designer/src/data_designer/interface/__init__.py @@ -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, @@ -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"), diff --git a/packages/data-designer/src/data_designer/interface/composite_workflow.py b/packages/data-designer/src/data_designer/interface/composite_workflow.py index 408083be5..ea7d30b5e 100644 --- a/packages/data-designer/src/data_designer/interface/composite_workflow.py +++ b/packages/data-designer/src/data_designer/interface/composite_workflow.py @@ -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 @@ -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 @@ -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, ) diff --git a/packages/data-designer/src/data_designer/interface/results.py b/packages/data-designer/src/data_designer/interface/results.py index 512150916..113b0dafe 100644 --- a/packages/data-designer/src/data_designer/interface/results.py +++ b/packages/data-designer/src/data_designer/interface/results.py @@ -52,8 +52,8 @@ def __init__( dataset_metadata: DatasetMetadata, task_traces: list[TaskTrace] | None = None, requested_num_records: int | None = None, - early_shutdown: bool = False, - requested_resume_mode: ResumeMode = ResumeMode.NEVER, + 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. @@ -69,8 +69,11 @@ def __init__( retained. requested_num_records: Number of records requested for this invocation. early_shutdown: Whether generation stopped at the early-shutdown gate. - requested_resume_mode: Resume mode requested for this invocation. - effective_resume_mode: Resume mode selected after compatibility checks. + ``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 @@ -80,7 +83,7 @@ def __init__( 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 or requested_resume_mode + self.effective_resume_mode = effective_resume_mode @property def dataset_path(self) -> Path: @@ -93,9 +96,11 @@ def actual_num_records(self) -> int: return self.count_records() @property - def is_partial(self) -> bool: + def is_partial(self) -> bool | None: """Return whether the result contains fewer records than requested.""" - return self.requested_num_records is not None and self.actual_num_records < self.requested_num_records + 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. diff --git a/packages/data-designer/tests/interface/test_composite_workflow.py b/packages/data-designer/tests/interface/test_composite_workflow.py index eb9a7abad..1968f1c2e 100644 --- a/packages/data-designer/tests/interface/test_composite_workflow.py +++ b/packages/data-designer/tests/interface/test_composite_workflow.py @@ -590,6 +590,39 @@ def test_composite_workflow_resume_if_possible_skips_completed_stages( assert create_mock.call_count == 0 assert results.count_records() == 3 assert results.load_dataset()["category"].tolist() == ["alpha", "alpha", "alpha"] + base_result = results["base"] + assert isinstance(base_result, DatasetCreationResults) + assert base_result.requested_num_records == 3 + assert base_result.actual_num_records == 3 + assert base_result.is_partial is False + assert base_result.early_shutdown is None + assert base_result.requested_resume_mode is None + assert base_result.effective_resume_mode is None + + +def test_composite_workflow_reused_stage_reports_partial_result( + stub_artifact_path: Path, + stub_model_providers: list[ModelProvider], + stub_model_configs: list[ModelConfig], + stub_dataset_profiler_results, +) -> None: + data_designer = _data_designer(stub_artifact_path, stub_model_providers) + create_mock = _patch_create(data_designer, stub_dataset_profiler_results) + workflow = data_designer.compose_workflow(name="resume-partial") + workflow.add_stage("base", _category_builder(stub_model_configs), num_records=3) + first = workflow.run() + batch_path = first["base"].artifact_storage.final_dataset_path / "batch_00000.parquet" + lazy.pd.read_parquet(batch_path).head(2).to_parquet(batch_path, index=False) + create_mock.reset_mock() + + resumed = data_designer.compose_workflow(name="resume-partial") + resumed.add_stage("base", _category_builder(stub_model_configs), num_records=3) + results = resumed.run(resume=ResumeMode.IF_POSSIBLE) + + assert create_mock.call_count == 0 + assert results["base"].requested_num_records == 3 + assert results["base"].actual_num_records == 2 + assert results["base"].is_partial is True def test_composite_workflow_resume_if_possible_skips_stage_with_output_processors( diff --git a/packages/data-designer/tests/interface/test_data_designer.py b/packages/data-designer/tests/interface/test_data_designer.py index 0ea707771..70f6bb355 100644 --- a/packages/data-designer/tests/interface/test_data_designer.py +++ b/packages/data-designer/tests/interface/test_data_designer.py @@ -36,6 +36,7 @@ FileContentsSeedSource, HuggingFaceSeedSource, ) +from data_designer.engine.dataset_builders.dataset_builder import DatasetBuilder from data_designer.engine.models.clients.adapters.http_model_client import ClientConcurrencyMode from data_designer.engine.models.errors import ( RETRYABLE_MODEL_ERRORS, @@ -683,6 +684,64 @@ def test_create_dataset_e2e_using_only_sampler_columns( analysis.to_report() +def test_create_reports_partial_early_shutdown( + stub_sampler_only_config_builder, + stub_artifact_path, + stub_model_providers, + stub_managed_assets_path, +) -> None: + data_designer = DataDesigner( + artifact_path=stub_artifact_path, + model_providers=stub_model_providers, + secret_resolver=PlaintextResolver(), + managed_assets_path=stub_managed_assets_path, + ) + data_designer.set_run_config(RunConfig(buffer_size=1, otel_metrics_port=None)) + original_build = DatasetBuilder.build + + def build_partial(builder: DatasetBuilder, **kwargs: Any) -> Path: + kwargs["num_records"] = 3 + return original_build(builder, **kwargs) + + with ( + patch.object(DatasetBuilder, "build", new=build_partial), + patch.object(DatasetBuilder, "early_shutdown", new_callable=PropertyMock, return_value=True), + ): + results = data_designer.create( + stub_sampler_only_config_builder, + num_records=10, + dataset_name="partial", + ) + + assert results.requested_num_records == 10 + assert results.actual_num_records == 3 + assert results.is_partial is True + assert results.early_shutdown is True + + +def test_create_reports_resolved_colliding_dataset_path( + stub_sampler_only_config_builder, + stub_artifact_path, + stub_model_providers, + stub_managed_assets_path, +) -> None: + data_designer = DataDesigner( + artifact_path=stub_artifact_path, + model_providers=stub_model_providers, + secret_resolver=PlaintextResolver(), + managed_assets_path=stub_managed_assets_path, + ) + data_designer.set_run_config(RunConfig(buffer_size=1, otel_metrics_port=None)) + + first = data_designer.create(stub_sampler_only_config_builder, num_records=1, dataset_name="collision") + second = data_designer.create(stub_sampler_only_config_builder, num_records=1, dataset_name="collision") + + assert first.dataset_path == stub_artifact_path / "collision" + assert second.dataset_path.parent == stub_artifact_path + assert second.dataset_path.name.startswith("collision_") + assert second.dataset_path != first.dataset_path + + def test_create_reports_effective_resume_mode( stub_sampler_only_config_builder, stub_artifact_path, diff --git a/packages/data-designer/tests/interface/test_results.py b/packages/data-designer/tests/interface/test_results.py index 2bbe154af..41a57c7b9 100644 --- a/packages/data-designer/tests/interface/test_results.py +++ b/packages/data-designer/tests/interface/test_results.py @@ -18,8 +18,8 @@ from data_designer.config.run_config import ResumeMode from data_designer.config.utils.errors import DatasetSampleDisplayError from data_designer.config.utils.visualization import display_sample_record as display_fn -from data_designer.engine.dataset_builders.errors import ArtifactStorageError from data_designer.engine.storage.artifact_storage import ArtifactStorage +from data_designer.interface import ArtifactStorageError from data_designer.interface.results import DatasetCreationResults diff --git a/plans/850/data-designer-contract.md b/plans/850/data-designer-contract.md index e193c2d67..8ebe005d5 100644 --- a/plans/850/data-designer-contract.md +++ b/plans/850/data-designer-contract.md @@ -26,6 +26,8 @@ The execution package may import these public symbols: from data_designer.config import ( DataDesignerConfigBuilder, InvalidConfigError, + InvalidFileFormatError, + InvalidFilePathError, LocalStdioMCPProvider, MCPProvider, ModelConfig, @@ -36,6 +38,7 @@ from data_designer.config import ( SamplingStrategy, ) from data_designer.interface import ( + ArtifactStorageError, DataDesigner, DataDesignerEarlyShutdownError, DataDesignerGenerationError, @@ -227,12 +230,12 @@ dataset contains exactly the requested number of records. | Field or method | Meaning | | --- | --- | | `dataset_path` | Resolved dataset directory, including collision or resume resolution. | -| `requested_num_records` | Target passed to the current `create()` invocation. | +| `requested_num_records` | Target passed to `create()`, or the persisted target for a reconstructed workflow result. | | `actual_num_records` | Current total rows in the final dataset, including rows from an earlier resumed invocation. | -| `is_partial` | `True` when actual records are fewer than requested records. | -| `early_shutdown` | Whether the current invocation stopped through the early-shutdown gate. | -| `requested_resume_mode` | Resume mode passed to `create()`. | -| `effective_resume_mode` | `always` when the invocation resumed, otherwise `never` after compatibility resolution. | +| `is_partial` | `True` when actual records are fewer than requested records, or `None` when the target is unavailable. | +| `early_shutdown` | Whether the current invocation stopped through the early-shutdown gate, or `None` when no generation invocation produced the result object. | +| `requested_resume_mode` | Resume mode passed to `create()`, or `None` when no generation invocation produced the result object. | +| `effective_resume_mode` | `always` when the invocation resumed, `never` when it started fresh, or `None` when no generation invocation produced the result object. | | `count_records()` | Metadata-only row count, equivalent to `actual_num_records`. | | `export(path, format=...)` | Stream the result to one JSONL, CSV, or Parquet file. | @@ -243,12 +246,16 @@ Public failure behavior is: | Condition | Public behavior | | --- | --- | -| Invalid serialized builder shape | `pydantic.ValidationError` or `ValueError` from `DataDesignerConfigBuilder.from_config()`. | +| Missing or unreadable local builder file | `InvalidFilePathError` from `data_designer.config`. | +| Malformed local or inline builder data | `InvalidFileFormatError` from `data_designer.config`. | +| Invalid serialized builder shape or remote source | `pydantic.ValidationError` or `ValueError` from `DataDesignerConfigBuilder.from_config()`. | | Invalid compiled config | `InvalidConfigError` from `data_designer.config`. | | Generation failure | `DataDesignerGenerationError`. | | Profiling failure after generation | `DataDesignerProfilingError`. No successful result is returned. | | Early shutdown with zero records | `DataDesignerEarlyShutdownError`, a `DataDesignerGenerationError` subclass. | | Early shutdown with some records | A partial `DatasetCreationResults` with `early_shutdown=True`. | +| Invalid export format or incompatible Parquet schemas | `InvalidFileFormatError` from `data_designer.config`. | +| Missing dataset or processor artifacts | `ArtifactStorageError` from `data_designer.interface`. | The caller must not inspect engine storage or task-trace types to classify an outcome. From 4c9c794a4c7498301e7e0a6e11aa243da8ea166e Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Mon, 10 Aug 2026 12:40:12 -0300 Subject: [PATCH 3/7] docs: clarify public invocation limits --- plans/850/data-designer-contract.md | 28 +++++++++++++++++++++++++--- 1 file changed, 25 insertions(+), 3 deletions(-) diff --git a/plans/850/data-designer-contract.md b/plans/850/data-designer-contract.md index 8ebe005d5..21d2d93e9 100644 --- a/plans/850/data-designer-contract.md +++ b/plans/850/data-designer-contract.md @@ -105,12 +105,14 @@ Client preflight runs in a fresh process and performs these steps: 3. Resolve `Plugin.config_cls` and `Plugin.impl_cls` so import failures stop preflight. 4. Load the builder with `DataDesignerConfigBuilder.from_config()`. 5. Materialize model bindings, per-model concurrency, seed input bindings, and MCP providers. -6. Verify referenced model aliases through each column config's public `get_model_aliases()` method. +6. Verify referenced model aliases through each column config's public `get_model_aliases()` method and each profiler + config's public `model_alias` field when present. 7. Construct `DataDesigner`, apply `RunConfig`, and call `DataDesigner.validate(builder)`. The explicit entry-point load is required because normal plugin discovery logs and skips a failing entry point. `DataDesigner.validate()` validates config structure and seed-dependent compilation without contacting model endpoints. -Endpoint readiness belongs after services are available; generation performs its normal readiness check. +It does not verify that every referenced column or profiler alias has a matching model config, so client preflight owns +that check. Endpoint readiness belongs after services are available; generation performs its normal readiness check. ## Model binding @@ -259,6 +261,10 @@ Public failure behavior is: The caller must not inspect engine storage or task-trace types to classify an outcome. +Failure exceptions do not carry a `DatasetCreationResults` object. The resolved dataset path, actual record count, +early-shutdown state, and effective resume mode may therefore be unavailable after a failed invocation. A semantic +failure record owned by an embedding package must make those facts optional rather than inspect engine storage. + ## Output, resume, seed, and telemetry - `artifact_path`, `dataset_name`, and `ResumeMode` define the resumable workspace. @@ -270,6 +276,20 @@ The caller must not inspect engine storage or task-trace types to classify an ou - OpenTelemetry metrics are configured with `RunConfig.otel_metrics_port`; `None` disables metrics for the invocation. - Embedded callers should use `auto_configure_logging=False` when they own process logging. +`DatasetCreationResults` does not expose per-invocation usage or phase timing. The public OpenTelemetry endpoint +provides create duration, generated/dropped record counters, and request-duration histograms, but not token totals or +separate generation and profiling durations. Callers must treat unavailable values as optional unless a future public +result summary provides them; logs and engine usage types are not a supported substitute. + +## Sharded execution limits + +Public builder methods expose column, processor, and profiler configs through `get_column_configs()`, +`get_processor_configs()`, and `get_profilers()`. Data Designer does not publish capability metadata that identifies +whether a processor, profiler, media output, or plugin is safe to merge across independently generated partitions. +An embedding package must use a conservative policy and reject multi-partition execution when those semantics are +present or unknown. Ordered seed input can use `PartitionBlock`; `SamplingStrategy.SHUFFLE` has no deterministic +partition-then-shuffle contract. + ## Callable and plugin limits Entry-point plugins with serializable config models are supported after installation in the client environment. @@ -299,10 +319,12 @@ The owning packages must cover: - Public root imports and absence of external `data_designer.engine` imports. - Builder file and inline mapping loading after plugin installation. - Plugin entry-point load failure before model services start. -- Model binding by alias, order preservation, and per-model concurrency. +- Model binding by alias, order preservation, per-model concurrency, profiler aliases, and plugin secondary aliases. - Compatibility-default merging from raw authored keys. - Complete, partial, partial early-shutdown, zero-record early-shutdown, generation-error, and profiling-error outcomes. +- Failure records that allow result-only path, count, early-shutdown, and resume facts to be unavailable. - Resolved dataset paths for fresh, colliding, resumed, and incompatible `if_possible` invocations. - Explicit JSONL, CSV, and Parquet exports. - Remote and stdio MCP providers without persisted secret values. +- Conservative multi-partition rejection for unsafe or unknown processor, profiler, media, and plugin semantics. - The existing import-performance threshold. From a5cac7ee28b6f2bf31023ee2347fe7877872b6e3 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Mon, 10 Aug 2026 14:08:44 -0300 Subject: [PATCH 4/7] fix: handle mixed-case config extensions Signed-off-by: Andre Manoel --- .../src/data_designer/config/utils/io_helpers.py | 2 +- .../data-designer-config/tests/config/test_public_contract.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py b/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py index a9c91a5e9..b990947ae 100644 --- a/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py +++ b/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py @@ -190,7 +190,7 @@ 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", ".json")) and not os.path.isfile(yaml_in): + if not from_url and yaml_in.lower().endswith((".yaml", ".yml", ".json")) and not os.path.isfile(yaml_in): raise FileNotFoundError(f"File not found: {yaml_in}") else: yaml_out = yaml.safe_load(yaml_in) diff --git a/packages/data-designer-config/tests/config/test_public_contract.py b/packages/data-designer-config/tests/config/test_public_contract.py index 6ff326def..0ae3d8f43 100644 --- a/packages/data-designer-config/tests/config/test_public_contract.py +++ b/packages/data-designer-config/tests/config/test_public_contract.py @@ -25,7 +25,7 @@ def test_builder_file_errors_are_publicly_exported() -> None: assert InvalidFilePathError is config_errors.InvalidFilePathError -@pytest.mark.parametrize("filename", ["missing.yaml", "missing.json"]) +@pytest.mark.parametrize("filename", ["missing.yaml", "missing.json", "missing.YAML", "missing.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)) From cf0f0bdf4a40d323f5ca821725d36e6dbf78e81b Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Thu, 13 Aug 2026 11:02:37 -0300 Subject: [PATCH 5/7] fix: complete generation contract edge cases Signed-off-by: Andre Manoel --- .../concepts/architecture-and-performance.mdx | 16 +++++++--- .../data_designer/config/config_builder.py | 2 +- .../data_designer/config/utils/io_helpers.py | 5 ++- .../tests/config/test_public_contract.py | 26 ++++++++++++++++ .../src/data_designer/interface/results.py | 26 +++++++++++++--- .../tests/interface/test_results.py | 31 ++++++++++++++++--- 6 files changed, 89 insertions(+), 17 deletions(-) diff --git a/fern/versions/latest/pages/concepts/architecture-and-performance.mdx b/fern/versions/latest/pages/concepts/architecture-and-performance.mdx index 60fb54b08..fb42f3d64 100644 --- a/fern/versions/latest/pages/concepts/architecture-and-performance.mdx +++ b/fern/versions/latest/pages/concepts/architecture-and-performance.mdx @@ -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: @@ -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. Only resume datasets from trusted artifact directories. Resume reads local `metadata.json`, `builder_config.json`, and parquet files to determine checkpoint state. @@ -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 diff --git a/packages/data-designer-config/src/data_designer/config/config_builder.py b/packages/data-designer-config/src/data_designer/config/config_builder.py index 428a111c9..08a1e1b17 100644 --- a/packages/data-designer-config/src/data_designer/config/config_builder.py +++ b/packages/data-designer-config/src/data_designer/config/config_builder.py @@ -128,7 +128,7 @@ def from_config(cls, config: dict | str | Path | BuilderConfig) -> Self: loaded_config = smart_load_yaml(config) except OSError as e: raise InvalidFilePathError(f"Failed to load builder config: {e}") from e - except yaml.YAMLError as 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 diff --git a/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py b/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py index b990947ae..76feacee9 100644 --- a/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py +++ b/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py @@ -190,10 +190,9 @@ 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.lower().endswith((".yaml", ".yml", ".json")) and not os.path.isfile(yaml_in): + yaml_out = yaml.safe_load(yaml_in) + if not from_url and not isinstance(yaml_out, dict) and yaml_in.lower().endswith((".yaml", ".yml", ".json")): raise FileNotFoundError(f"File not found: {yaml_in}") - else: - yaml_out = yaml.safe_load(yaml_in) else: raise ValueError( f"'{yaml_in}' is an invalid yaml config format. Valid options are: dict, yaml string, or yaml file path." diff --git a/packages/data-designer-config/tests/config/test_public_contract.py b/packages/data-designer-config/tests/config/test_public_contract.py index 0ae3d8f43..20ad2d8db 100644 --- a/packages/data-designer-config/tests/config/test_public_contract.py +++ b/packages/data-designer-config/tests/config/test_public_contract.py @@ -33,6 +33,32 @@ def test_from_config_normalizes_missing_file_error(tmp_path: Path, filename: str 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""" +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: [") diff --git a/packages/data-designer/src/data_designer/interface/results.py b/packages/data-designer/src/data_designer/interface/results.py index 113b0dafe..7e41bf429 100644 --- a/packages/data-designer/src/data_designer/interface/results.py +++ b/packages/data-designer/src/data_designer/interface/results.py @@ -116,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. @@ -141,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. @@ -211,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": diff --git a/packages/data-designer/tests/interface/test_results.py b/packages/data-designer/tests/interface/test_results.py index 41a57c7b9..deb11e50a 100644 --- a/packages/data-designer/tests/interface/test_results.py +++ b/packages/data-designer/tests/interface/test_results.py @@ -202,10 +202,7 @@ def test_display_sample_record_with_empty_dataset(): dataset_metadata=DatasetMetadata(), ) - # Empty DataFrame is still a valid DataFrame, so accessing _record_sampler_dataset succeeds - # but display_sample_record fails when trying to access index 0 - # Note: Currently raises UnboundLocalError due to bug in error handling, but tests that it fails - with pytest.raises((DatasetSampleDisplayError, UnboundLocalError)): + with pytest.raises(ArtifactStorageError, match="No batch parquet files found"): results.display_sample_record() @@ -390,6 +387,32 @@ def test_count_records(stub_dataset_creation_results, stub_dataframe, stub_batch assert stub_dataset_creation_results.count_records() == len(stub_dataframe) +def test_missing_dataset_artifacts_raise_public_error(stub_dataset_creation_results, tmp_path) -> None: + empty_dir = tmp_path / "parquet-files" + empty_dir.mkdir() + stub_dataset_creation_results.artifact_storage.final_dataset_path = empty_dir + stub_dataset_creation_results.artifact_storage.load_dataset.return_value = lazy.pd.DataFrame() + stub_dataset_creation_results.requested_num_records = 1 + + with pytest.raises(ArtifactStorageError, match="No batch parquet files found"): + stub_dataset_creation_results.load_dataset() + with pytest.raises(ArtifactStorageError, match="No batch parquet files found"): + stub_dataset_creation_results.count_records() + with pytest.raises(ArtifactStorageError, match="No batch parquet files found"): + _ = stub_dataset_creation_results.actual_num_records + with pytest.raises(ArtifactStorageError, match="No batch parquet files found"): + _ = stub_dataset_creation_results.is_partial + + +def test_load_dataset_normalizes_missing_artifact_error(stub_dataset_creation_results) -> None: + stub_dataset_creation_results.artifact_storage.load_dataset.side_effect = FileNotFoundError("missing") + + with pytest.raises(ArtifactStorageError, match="Failed to load dataset artifacts") as exc_info: + stub_dataset_creation_results.load_dataset() + + assert isinstance(exc_info.value.__cause__, FileNotFoundError) + + def test_public_creation_outcome_fields( stub_artifact_storage, stub_dataset_profiler_results, From 40daaa966ca593bd4845d144383ef8f870c4458b Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Thu, 13 Aug 2026 11:25:55 -0300 Subject: [PATCH 6/7] test: make inline config regression self-contained Signed-off-by: Andre Manoel --- .../data-designer-config/tests/config/test_public_contract.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/packages/data-designer-config/tests/config/test_public_contract.py b/packages/data-designer-config/tests/config/test_public_contract.py index 20ad2d8db..ec1f877fe 100644 --- a/packages/data-designer-config/tests/config/test_public_contract.py +++ b/packages/data-designer-config/tests/config/test_public_contract.py @@ -37,6 +37,10 @@ def test_from_config_normalizes_missing_file_error(tmp_path: Path, filename: str 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 From 1945750e8ef77ea61a7ddd04cd436cd643db475b Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Thu, 13 Aug 2026 13:46:22 -0300 Subject: [PATCH 7/7] fix: classify malformed config paths Signed-off-by: Andre Manoel --- .../src/data_designer/config/utils/io_helpers.py | 10 ++++++++-- .../tests/config/test_public_contract.py | 2 +- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py b/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py index 76feacee9..0e8878248 100644 --- a/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py +++ b/packages/data-designer-config/src/data_designer/config/utils/io_helpers.py @@ -190,8 +190,14 @@ 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): - yaml_out = yaml.safe_load(yaml_in) - if not from_url and not isinstance(yaml_out, dict) and yaml_in.lower().endswith((".yaml", ".yml", ".json")): + 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( diff --git a/packages/data-designer-config/tests/config/test_public_contract.py b/packages/data-designer-config/tests/config/test_public_contract.py index ec1f877fe..0f3347ff8 100644 --- a/packages/data-designer-config/tests/config/test_public_contract.py +++ b/packages/data-designer-config/tests/config/test_public_contract.py @@ -25,7 +25,7 @@ def test_builder_file_errors_are_publicly_exported() -> None: assert InvalidFilePathError is config_errors.InvalidFilePathError -@pytest.mark.parametrize("filename", ["missing.yaml", "missing.json", "missing.YAML", "missing.JsOn"]) +@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))