From 4edb96ded5ad3551b2ca5cdb14f89b36bcebdad2 Mon Sep 17 00:00:00 2001 From: "Eric W. Tramel" Date: Wed, 7 Jan 2026 23:21:30 -0500 Subject: [PATCH 1/8] Expose shutdown features to `create(...)` --- .../dataset_builders/column_wise_builder.py | 8 +++++ src/data_designer/interface/data_designer.py | 30 +++++++++++++++++-- 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/src/data_designer/engine/dataset_builders/column_wise_builder.py b/src/data_designer/engine/dataset_builders/column_wise_builder.py index 77b6459b6..7a0433edc 100644 --- a/src/data_designer/engine/dataset_builders/column_wise_builder.py +++ b/src/data_designer/engine/dataset_builders/column_wise_builder.py @@ -60,6 +60,9 @@ def __init__( processor_configs: list[ProcessorConfig], resource_provider: ResourceProvider, registry: DataDesignerRegistry | None = None, + enable_early_shutdown: bool = True, + shutdown_error_rate: float = 0.5, + shutdown_error_window: int = 10, ): self.batch_manager = DatasetBatchManager(resource_provider.artifact_storage) self._resource_provider = resource_provider @@ -67,6 +70,9 @@ def __init__( self._registry = registry or DataDesignerRegistry() self._column_configs = column_configs self._processors: dict[BuildStage, list[Processor]] = self._initialize_processors(processor_configs) + self._enable_early_shutdown = enable_early_shutdown + self._shutdown_error_rate = shutdown_error_rate + self._shutdown_error_window = shutdown_error_window self._validate_column_configs() @property @@ -222,6 +228,8 @@ def _fan_out_with_threads(self, generator: WithModelGeneration, max_workers: int column_name=generator.config.name, result_callback=self._worker_result_callback, error_callback=self._worker_error_callback, + shutdown_error_rate=self._shutdown_error_rate if self._enable_early_shutdown else 1.0, + shutdown_error_window=self._shutdown_error_window, ) as executor: for i, record in self.batch_manager.iter_current_batch(): executor.submit(lambda record: generator.generate(record), record, context={"index": i}) diff --git a/src/data_designer/interface/data_designer.py b/src/data_designer/interface/data_designer.py index 0458dc60f..da72f3485 100644 --- a/src/data_designer/interface/data_designer.py +++ b/src/data_designer/interface/data_designer.py @@ -154,6 +154,9 @@ def create( *, num_records: int = DEFAULT_NUM_RECORDS, dataset_name: str = "dataset", + enable_early_shutdown: bool = True, + shutdown_error_rate: float = 0.5, + shutdown_error_window: int = 10, ) -> DatasetCreationResults: """Create dataset and save results to the local artifact storage. @@ -171,6 +174,15 @@ def create( a datetime stamp. For example, if the dataset name is "awesome_dataset" and a directory with the same name already exists, the dataset will be saved to a new directory with the name "awesome_dataset_2025-01-01_12-00-00". + enable_early_shutdown: If True (default), dataset generation will terminate + early if the error rate exceeds `shutdown_error_rate` after + `shutdown_error_window` tasks complete. Set to False to disable + early shutdown entirely (ignores `shutdown_error_rate` and + `shutdown_error_window`). + shutdown_error_rate: Error rate threshold (0.0-1.0) that triggers early + shutdown. Only used when `enable_early_shutdown=True`. Default is 0.5 (50%). + shutdown_error_window: Minimum number of completed tasks before error rate + monitoring begins. Only used when `enable_early_shutdown=True`. Default is 10. Returns: DatasetCreationResults object with methods for loading the generated dataset, @@ -184,7 +196,13 @@ def create( resource_provider = self._create_resource_provider(dataset_name, config_builder) - builder = self._create_dataset_builder(config_builder, resource_provider) + builder = self._create_dataset_builder( + config_builder, + resource_provider, + enable_early_shutdown=enable_early_shutdown, + shutdown_error_rate=shutdown_error_rate, + shutdown_error_window=shutdown_error_window, + ) try: builder.build(num_records=num_records, buffer_size=self._buffer_size) @@ -334,12 +352,20 @@ def _resolve_model_providers(self, model_providers: list[ModelProvider] | None) return model_providers or [] def _create_dataset_builder( - self, config_builder: DataDesignerConfigBuilder, resource_provider: ResourceProvider + self, + config_builder: DataDesignerConfigBuilder, + resource_provider: ResourceProvider, + enable_early_shutdown: bool = True, + shutdown_error_rate: float = 0.5, + shutdown_error_window: int = 10, ) -> ColumnWiseDatasetBuilder: return ColumnWiseDatasetBuilder( column_configs=compile_dataset_builder_column_configs(config_builder.build(raise_exceptions=True)), processor_configs=config_builder.get_processor_configs(), resource_provider=resource_provider, + enable_early_shutdown=enable_early_shutdown, + shutdown_error_rate=shutdown_error_rate, + shutdown_error_window=shutdown_error_window, ) def _create_dataset_profiler( From 69109b5071d9533953bc5c8d8552db99905da8c3 Mon Sep 17 00:00:00 2001 From: "Eric W. Tramel" Date: Wed, 7 Jan 2026 23:36:14 -0500 Subject: [PATCH 2/8] Tests to verify passthrough of early shutdown properties --- .../test_column_wise_builder.py | 76 +++++++++++++++++++ 1 file changed, 76 insertions(+) diff --git a/tests/engine/dataset_builders/test_column_wise_builder.py b/tests/engine/dataset_builders/test_column_wise_builder.py index dd8c84db6..6f0aef70b 100644 --- a/tests/engine/dataset_builders/test_column_wise_builder.py +++ b/tests/engine/dataset_builders/test_column_wise_builder.py @@ -9,6 +9,7 @@ from data_designer.config.column_configs import LLMTextColumnConfig, SamplerColumnConfig from data_designer.config.dataset_builders import BuildStage from data_designer.config.processors import DropColumnsProcessorConfig +from data_designer.engine.column_generators.generators.base import GenerationStrategy from data_designer.engine.dataset_builders.column_wise_builder import ( MAX_CONCURRENCY_PER_NON_LLM_GENERATOR, ColumnWiseDatasetBuilder, @@ -306,3 +307,78 @@ def test_emit_batch_inference_events_handles_multiple_models( events = [call[0][0] for call in mock_handler_instance.enqueue.call_args_list] model_names = {e.model for e in events} assert model_names == {"model-a", "model-b"} + + +@pytest.mark.parametrize( + "enable_early_shutdown,shutdown_error_rate,shutdown_error_window", + [ + (True, 0.5, 10), # defaults + (False, 0.8, 25), # custom values + ], +) +def test_column_wise_dataset_builder_stores_early_shutdown_params( + stub_resource_provider: Mock, + stub_test_column_configs: list, + stub_test_processor_configs: list, + enable_early_shutdown: bool, + shutdown_error_rate: float, + shutdown_error_window: int, +) -> None: + """Test that ColumnWiseDatasetBuilder stores early shutdown parameters.""" + builder = ColumnWiseDatasetBuilder( + column_configs=stub_test_column_configs, + processor_configs=stub_test_processor_configs, + resource_provider=stub_resource_provider, + enable_early_shutdown=enable_early_shutdown, + shutdown_error_rate=shutdown_error_rate, + shutdown_error_window=shutdown_error_window, + ) + + assert builder._enable_early_shutdown is enable_early_shutdown + assert builder._shutdown_error_rate == shutdown_error_rate + assert builder._shutdown_error_window == shutdown_error_window + + +@pytest.mark.parametrize( + "enable_early_shutdown,configured_rate,expected_rate", + [ + (True, 0.7, 0.7), # enabled: use configured rate + (False, 0.7, 1.0), # disabled: use 1.0 to effectively disable + ], +) +@patch("data_designer.engine.dataset_builders.column_wise_builder.ConcurrentThreadExecutor") +def test_fan_out_with_threads_respects_enable_early_shutdown_flag( + mock_executor_class: Mock, + stub_resource_provider: Mock, + stub_test_column_configs: list, + stub_test_processor_configs: list, + enable_early_shutdown: bool, + configured_rate: float, + expected_rate: float, +) -> None: + """Test that _fan_out_with_threads passes correct shutdown_error_rate based on enable flag.""" + builder = ColumnWiseDatasetBuilder( + column_configs=stub_test_column_configs, + processor_configs=stub_test_processor_configs, + resource_provider=stub_resource_provider, + enable_early_shutdown=enable_early_shutdown, + shutdown_error_rate=configured_rate, + shutdown_error_window=20, + ) + + mock_executor_class.return_value.__enter__ = Mock(return_value=Mock()) + mock_executor_class.return_value.__exit__ = Mock(return_value=False) + + mock_generator = Mock() + mock_generator.generation_strategy = GenerationStrategy.CELL_BY_CELL + mock_generator.config.name = "test" + mock_generator.config.column_type = "llm_text" + + builder.batch_manager = Mock() + builder.batch_manager.iter_current_batch.return_value = [] + + builder._fan_out_with_threads(mock_generator, max_workers=4) + + call_kwargs = mock_executor_class.call_args[1] + assert call_kwargs["shutdown_error_rate"] == expected_rate + assert call_kwargs["shutdown_error_window"] == 20 From e2f4c446b9c80f22cbe8671d620a47b0339eefb7 Mon Sep 17 00:00:00 2001 From: "Eric W. Tramel" Date: Thu, 8 Jan 2026 13:38:43 -0500 Subject: [PATCH 3/8] Change to a config type and helper settings --- src/data_designer/config/run_settings.py | 35 +++++++++++ .../generators/validation.py | 3 + .../dataset_builders/column_wise_builder.py | 11 +--- .../engine/resources/resource_provider.py | 2 + src/data_designer/essentials/__init__.py | 2 + src/data_designer/interface/data_designer.py | 43 ++++++-------- tests/engine/conftest.py | 2 + .../test_column_wise_builder.py | 57 ++++++------------ tests/interface/test_data_designer.py | 58 +++++++++++++++++++ 9 files changed, 140 insertions(+), 73 deletions(-) create mode 100644 src/data_designer/config/run_settings.py diff --git a/src/data_designer/config/run_settings.py b/src/data_designer/config/run_settings.py new file mode 100644 index 000000000..e83e023cb --- /dev/null +++ b/src/data_designer/config/run_settings.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import Self + +from pydantic import Field, model_validator + +from data_designer.config.base import ConfigBase + + +class RunConfig(ConfigBase): + """Runtime configuration for dataset generation. + + Groups configuration options that control generation behavior but aren't + part of the dataset configuration itself. + + Attributes: + disable_early_shutdown: If True, disables early shutdown entirely. Generation + will continue regardless of error rate. Default is False. + shutdown_error_rate: Error rate threshold (0.0-1.0) that triggers early shutdown. + When early shutdown is disabled, this value is normalized to 1.0. Default is 0.5. + shutdown_error_window: Minimum number of completed tasks before error rate + monitoring begins. Must be >= 0. Default is 10. + """ + + disable_early_shutdown: bool = False + shutdown_error_rate: float = Field(default=0.5, ge=0.0, le=1.0) + shutdown_error_window: int = Field(default=10, ge=0) + + @model_validator(mode="after") + def normalize_shutdown_settings(self) -> Self: + """Set shutdown_error_rate to 1.0 when early shutdown is disabled.""" + if self.disable_early_shutdown: + self.shutdown_error_rate = 1.0 + return self diff --git a/src/data_designer/engine/column_generators/generators/validation.py b/src/data_designer/engine/column_generators/generators/validation.py index f46296b99..a2e2b3c98 100644 --- a/src/data_designer/engine/column_generators/generators/validation.py +++ b/src/data_designer/engine/column_generators/generators/validation.py @@ -123,11 +123,14 @@ def result_callback(result: ValidationResult, context: dict): def error_callback(error: Exception, context: dict): outputs[context["index"]] = ValidationResult.empty(size=len(batched_records[context["index"]])) + settings = self.resource_provider.run_config with ConcurrentThreadExecutor( max_workers=self.config.validator_params.max_parallel_requests, column_name=self.config.name, result_callback=result_callback, error_callback=error_callback, + shutdown_error_rate=settings.shutdown_error_rate, + shutdown_error_window=settings.shutdown_error_window, ) as executor: for i, batch in enumerate(batched_records): executor.submit(lambda batch: self._validate_batch(validator, batch), batch, context={"index": i}) diff --git a/src/data_designer/engine/dataset_builders/column_wise_builder.py b/src/data_designer/engine/dataset_builders/column_wise_builder.py index 7a0433edc..0f54ddf11 100644 --- a/src/data_designer/engine/dataset_builders/column_wise_builder.py +++ b/src/data_designer/engine/dataset_builders/column_wise_builder.py @@ -60,9 +60,6 @@ def __init__( processor_configs: list[ProcessorConfig], resource_provider: ResourceProvider, registry: DataDesignerRegistry | None = None, - enable_early_shutdown: bool = True, - shutdown_error_rate: float = 0.5, - shutdown_error_window: int = 10, ): self.batch_manager = DatasetBatchManager(resource_provider.artifact_storage) self._resource_provider = resource_provider @@ -70,9 +67,6 @@ def __init__( self._registry = registry or DataDesignerRegistry() self._column_configs = column_configs self._processors: dict[BuildStage, list[Processor]] = self._initialize_processors(processor_configs) - self._enable_early_shutdown = enable_early_shutdown - self._shutdown_error_rate = shutdown_error_rate - self._shutdown_error_window = shutdown_error_window self._validate_column_configs() @property @@ -223,13 +217,14 @@ def _fan_out_with_threads(self, generator: WithModelGeneration, max_workers: int f"🐙 Processing {generator.config.column_type} column '{generator.config.name}' " f"with {max_workers} concurrent workers" ) + settings = self._resource_provider.run_config with ConcurrentThreadExecutor( max_workers=max_workers, column_name=generator.config.name, result_callback=self._worker_result_callback, error_callback=self._worker_error_callback, - shutdown_error_rate=self._shutdown_error_rate if self._enable_early_shutdown else 1.0, - shutdown_error_window=self._shutdown_error_window, + shutdown_error_rate=settings.shutdown_error_rate, + shutdown_error_window=settings.shutdown_error_window, ) as executor: for i, record in self.batch_manager.iter_current_batch(): executor.submit(lambda record: generator.generate(record), record, context={"index": i}) diff --git a/src/data_designer/engine/resources/resource_provider.py b/src/data_designer/engine/resources/resource_provider.py index 98c6576fa..2a7dd7fd0 100644 --- a/src/data_designer/engine/resources/resource_provider.py +++ b/src/data_designer/engine/resources/resource_provider.py @@ -3,6 +3,7 @@ from data_designer.config.base import ConfigBase from data_designer.config.models import ModelConfig +from data_designer.config.run_settings import RunConfig from data_designer.config.utils.type_helpers import StrEnum from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage from data_designer.engine.model_provider import ModelProviderRegistry @@ -23,6 +24,7 @@ class ResourceProvider(ConfigBase): blob_storage: ManagedBlobStorage | None = None datastore: SeedDatasetDataStore | None = None model_registry: ModelRegistry | None = None + run_config: RunConfig = RunConfig() def create_resource_provider( diff --git a/src/data_designer/essentials/__init__.py b/src/data_designer/essentials/__init__.py index 6c2d021e6..6933ee0c8 100644 --- a/src/data_designer/essentials/__init__.py +++ b/src/data_designer/essentials/__init__.py @@ -3,6 +3,7 @@ from data_designer.config.default_model_settings import resolve_seed_default_model_settings from data_designer.config.exports import * # noqa: F403 +from data_designer.config.run_settings import RunConfig from data_designer.config.validator_params import LocalCallableValidatorParams from data_designer.interface.data_designer import DataDesigner from data_designer.logging import LoggingConfig, configure_logging @@ -21,6 +22,7 @@ def get_essentials_exports() -> list[str]: local = [ DataDesigner.__name__, LocalCallableValidatorParams.__name__, + RunConfig.__name__, ] return logging + local + get_config_exports() # noqa: F405 diff --git a/src/data_designer/interface/data_designer.py b/src/data_designer/interface/data_designer.py index da72f3485..9e92fee58 100644 --- a/src/data_designer/interface/data_designer.py +++ b/src/data_designer/interface/data_designer.py @@ -20,6 +20,7 @@ ModelProvider, ) from data_designer.config.preview_results import PreviewResults +from data_designer.config.run_settings import RunConfig from data_designer.config.seed import LocalSeedDatasetReference from data_designer.config.utils.constants import ( DEFAULT_NUM_RECORDS, @@ -97,6 +98,7 @@ def __init__( self._secret_resolver = secret_resolver or CompositeResolver([EnvironmentResolver(), PlaintextResolver()]) self._artifact_path = Path(artifact_path) if artifact_path is not None else Path.cwd() / "artifacts" self._buffer_size = DEFAULT_BUFFER_SIZE + self._run_config = RunConfig() self._managed_assets_path = Path(managed_assets_path or MANAGED_ASSETS_PATH) self._model_providers = self._resolve_model_providers(model_providers) self._model_provider_registry = resolve_model_provider_registry( @@ -154,9 +156,6 @@ def create( *, num_records: int = DEFAULT_NUM_RECORDS, dataset_name: str = "dataset", - enable_early_shutdown: bool = True, - shutdown_error_rate: float = 0.5, - shutdown_error_window: int = 10, ) -> DatasetCreationResults: """Create dataset and save results to the local artifact storage. @@ -174,15 +173,6 @@ def create( a datetime stamp. For example, if the dataset name is "awesome_dataset" and a directory with the same name already exists, the dataset will be saved to a new directory with the name "awesome_dataset_2025-01-01_12-00-00". - enable_early_shutdown: If True (default), dataset generation will terminate - early if the error rate exceeds `shutdown_error_rate` after - `shutdown_error_window` tasks complete. Set to False to disable - early shutdown entirely (ignores `shutdown_error_rate` and - `shutdown_error_window`). - shutdown_error_rate: Error rate threshold (0.0-1.0) that triggers early - shutdown. Only used when `enable_early_shutdown=True`. Default is 0.5 (50%). - shutdown_error_window: Minimum number of completed tasks before error rate - monitoring begins. Only used when `enable_early_shutdown=True`. Default is 10. Returns: DatasetCreationResults object with methods for loading the generated dataset, @@ -196,13 +186,7 @@ def create( resource_provider = self._create_resource_provider(dataset_name, config_builder) - builder = self._create_dataset_builder( - config_builder, - resource_provider, - enable_early_shutdown=enable_early_shutdown, - shutdown_error_rate=shutdown_error_rate, - shutdown_error_window=shutdown_error_window, - ) + builder = self._create_dataset_builder(config_builder, resource_provider) try: builder.build(num_records=num_records, buffer_size=self._buffer_size) @@ -336,6 +320,20 @@ def set_buffer_size(self, buffer_size: int) -> None: raise InvalidBufferValueError("Buffer size must be greater than 0.") self._buffer_size = buffer_size + def set_run_config(self, run_config: RunConfig) -> None: + """Set the runtime configuration for dataset generation. + + Args: + run_config: A RunConfig instance containing runtime settings such as + early shutdown behavior. Import RunConfig from data_designer.essentials. + + Example: + >>> from data_designer.essentials import DataDesigner, RunConfig + >>> dd = DataDesigner() + >>> dd.set_run_config(RunConfig(disable_early_shutdown=True)) + """ + self._run_config = run_config + def _resolve_model_providers(self, model_providers: list[ModelProvider] | None) -> list[ModelProvider]: if model_providers is None: model_providers = get_default_providers() @@ -355,17 +353,11 @@ def _create_dataset_builder( self, config_builder: DataDesignerConfigBuilder, resource_provider: ResourceProvider, - enable_early_shutdown: bool = True, - shutdown_error_rate: float = 0.5, - shutdown_error_window: int = 10, ) -> ColumnWiseDatasetBuilder: return ColumnWiseDatasetBuilder( column_configs=compile_dataset_builder_column_configs(config_builder.build(raise_exceptions=True)), processor_configs=config_builder.get_processor_configs(), resource_provider=resource_provider, - enable_early_shutdown=enable_early_shutdown, - shutdown_error_rate=shutdown_error_rate, - shutdown_error_window=shutdown_error_window, ) def _create_dataset_profiler( @@ -400,6 +392,7 @@ def _create_resource_provider( token=settings.token, ) ), + run_config=self._run_config, ) def _get_interface_info(self, model_providers: list[ModelProvider]) -> InterfaceInfo: diff --git a/tests/engine/conftest.py b/tests/engine/conftest.py index dc30ba256..5f2135a4a 100644 --- a/tests/engine/conftest.py +++ b/tests/engine/conftest.py @@ -6,6 +6,7 @@ import pandas as pd import pytest +from data_designer.config.run_settings import RunConfig from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage from data_designer.engine.models.facade import ModelFacade from data_designer.engine.models.registry import ModelRegistry @@ -36,6 +37,7 @@ def stub_resource_provider(tmp_path, stub_model_facade): mock_provider.artifact_storage = ArtifactStorage(artifact_path=tmp_path) mock_provider.blob_storage = Mock(spec=ManagedBlobStorage) mock_provider.datastore = Mock() + mock_provider.run_config = RunConfig() return mock_provider diff --git a/tests/engine/dataset_builders/test_column_wise_builder.py b/tests/engine/dataset_builders/test_column_wise_builder.py index 6f0aef70b..0156fba7a 100644 --- a/tests/engine/dataset_builders/test_column_wise_builder.py +++ b/tests/engine/dataset_builders/test_column_wise_builder.py @@ -310,60 +310,37 @@ def test_emit_batch_inference_events_handles_multiple_models( @pytest.mark.parametrize( - "enable_early_shutdown,shutdown_error_rate,shutdown_error_window", + "disable_early_shutdown,configured_rate,expected_rate,shutdown_error_window", [ - (True, 0.5, 10), # defaults - (False, 0.8, 25), # custom values - ], -) -def test_column_wise_dataset_builder_stores_early_shutdown_params( - stub_resource_provider: Mock, - stub_test_column_configs: list, - stub_test_processor_configs: list, - enable_early_shutdown: bool, - shutdown_error_rate: float, - shutdown_error_window: int, -) -> None: - """Test that ColumnWiseDatasetBuilder stores early shutdown parameters.""" - builder = ColumnWiseDatasetBuilder( - column_configs=stub_test_column_configs, - processor_configs=stub_test_processor_configs, - resource_provider=stub_resource_provider, - enable_early_shutdown=enable_early_shutdown, - shutdown_error_rate=shutdown_error_rate, - shutdown_error_window=shutdown_error_window, - ) - - assert builder._enable_early_shutdown is enable_early_shutdown - assert builder._shutdown_error_rate == shutdown_error_rate - assert builder._shutdown_error_window == shutdown_error_window - - -@pytest.mark.parametrize( - "enable_early_shutdown,configured_rate,expected_rate", - [ - (True, 0.7, 0.7), # enabled: use configured rate - (False, 0.7, 1.0), # disabled: use 1.0 to effectively disable + (False, 0.7, 0.7, 20), # enabled: use configured rate + (True, 0.7, 1.0, 20), # disabled: use 1.0 to effectively disable + (False, 0.5, 0.5, 10), # defaults ], ) @patch("data_designer.engine.dataset_builders.column_wise_builder.ConcurrentThreadExecutor") -def test_fan_out_with_threads_respects_enable_early_shutdown_flag( +def test_fan_out_with_threads_uses_early_shutdown_settings_from_resource_provider( mock_executor_class: Mock, stub_resource_provider: Mock, stub_test_column_configs: list, stub_test_processor_configs: list, - enable_early_shutdown: bool, + disable_early_shutdown: bool, configured_rate: float, expected_rate: float, + shutdown_error_window: int, ) -> None: - """Test that _fan_out_with_threads passes correct shutdown_error_rate based on enable flag.""" + """Test that _fan_out_with_threads uses run settings from resource_provider.""" + from data_designer.config.run_settings import RunConfig + + stub_resource_provider.run_config = RunConfig( + disable_early_shutdown=disable_early_shutdown, + shutdown_error_rate=configured_rate, + shutdown_error_window=shutdown_error_window, + ) + builder = ColumnWiseDatasetBuilder( column_configs=stub_test_column_configs, processor_configs=stub_test_processor_configs, resource_provider=stub_resource_provider, - enable_early_shutdown=enable_early_shutdown, - shutdown_error_rate=configured_rate, - shutdown_error_window=20, ) mock_executor_class.return_value.__enter__ = Mock(return_value=Mock()) @@ -381,4 +358,4 @@ def test_fan_out_with_threads_respects_enable_early_shutdown_flag( call_kwargs = mock_executor_class.call_args[1] assert call_kwargs["shutdown_error_rate"] == expected_rate - assert call_kwargs["shutdown_error_window"] == 20 + assert call_kwargs["shutdown_error_window"] == shutdown_error_window diff --git a/tests/interface/test_data_designer.py b/tests/interface/test_data_designer.py index c8e94000e..74a08d2db 100644 --- a/tests/interface/test_data_designer.py +++ b/tests/interface/test_data_designer.py @@ -13,6 +13,7 @@ from data_designer.config.dataset_builders import BuildStage from data_designer.config.errors import InvalidFileFormatError from data_designer.config.processors import DropColumnsProcessorConfig +from data_designer.config.run_settings import RunConfig from data_designer.config.seed import LocalSeedDatasetReference from data_designer.engine.model_provider import ModelProvider from data_designer.engine.secret_resolver import CompositeResolver, EnvironmentResolver, PlaintextResolver @@ -268,6 +269,63 @@ def test_set_buffer_size_raises_error_for_invalid_buffer_size(stub_artifact_path data_designer.set_buffer_size(0) +def test_run_config_setting_persists(stub_artifact_path, stub_model_providers): + """Test that run config setting persists across multiple calls.""" + data_designer = DataDesigner(artifact_path=stub_artifact_path, model_providers=stub_model_providers) + + # Test default values + assert data_designer._run_config.disable_early_shutdown is False + assert data_designer._run_config.shutdown_error_rate == 0.5 + assert data_designer._run_config.shutdown_error_window == 10 + + # Test setting custom values (note: shutdown_error_rate is normalized to 1.0 when disabled) + data_designer.set_run_config( + RunConfig( + disable_early_shutdown=True, + shutdown_error_rate=0.8, + shutdown_error_window=25, + ) + ) + assert data_designer._run_config.disable_early_shutdown is True + assert data_designer._run_config.shutdown_error_rate == 1.0 # normalized when disabled + assert data_designer._run_config.shutdown_error_window == 25 + + # Test updating values + data_designer.set_run_config( + RunConfig( + disable_early_shutdown=False, + shutdown_error_rate=0.3, + shutdown_error_window=5, + ) + ) + assert data_designer._run_config.disable_early_shutdown is False + assert data_designer._run_config.shutdown_error_rate == 0.3 + assert data_designer._run_config.shutdown_error_window == 5 + + +def test_run_config_normalizes_error_rate_when_disabled(stub_artifact_path, stub_model_providers): + """Test that shutdown_error_rate is normalized to 1.0 when disabled.""" + data_designer = DataDesigner(artifact_path=stub_artifact_path, model_providers=stub_model_providers) + + # When enabled (default), shutdown_error_rate should use the configured value + data_designer.set_run_config( + RunConfig( + disable_early_shutdown=False, + shutdown_error_rate=0.7, + ) + ) + assert data_designer._run_config.shutdown_error_rate == 0.7 + + # When disabled, shutdown_error_rate should be normalized to 1.0 + data_designer.set_run_config( + RunConfig( + disable_early_shutdown=True, + shutdown_error_rate=0.7, + ) + ) + assert data_designer._run_config.shutdown_error_rate == 1.0 + + def test_multiple_seed_references_can_be_created(): """Test that multiple seed references can be created from different sources.""" with tempfile.TemporaryDirectory() as temp_dir: From ca4c1877546be807dc9cab1f41a5534218972752 Mon Sep 17 00:00:00 2001 From: "Eric W. Tramel" Date: Thu, 8 Jan 2026 13:44:29 -0500 Subject: [PATCH 4/8] Update tests/interface/test_data_designer.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- tests/interface/test_data_designer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/interface/test_data_designer.py b/tests/interface/test_data_designer.py index 74a08d2db..0305795f9 100644 --- a/tests/interface/test_data_designer.py +++ b/tests/interface/test_data_designer.py @@ -278,7 +278,7 @@ def test_run_config_setting_persists(stub_artifact_path, stub_model_providers): assert data_designer._run_config.shutdown_error_rate == 0.5 assert data_designer._run_config.shutdown_error_window == 10 - # Test setting custom values (note: shutdown_error_rate is normalized to 1.0 when disabled) + # Test setting custom values data_designer.set_run_config( RunConfig( disable_early_shutdown=True, From 467c1a86662b179509ab067b97f6ff1393069082 Mon Sep 17 00:00:00 2001 From: "Eric W. Tramel" Date: Thu, 8 Jan 2026 13:55:21 -0500 Subject: [PATCH 5/8] Update module naming --- src/data_designer/config/{run_settings.py => run_config.py} | 0 src/data_designer/engine/resources/resource_provider.py | 2 +- src/data_designer/essentials/__init__.py | 2 +- src/data_designer/interface/data_designer.py | 2 +- tests/engine/conftest.py | 2 +- tests/engine/dataset_builders/test_column_wise_builder.py | 2 +- tests/interface/test_data_designer.py | 2 +- 7 files changed, 6 insertions(+), 6 deletions(-) rename src/data_designer/config/{run_settings.py => run_config.py} (100%) diff --git a/src/data_designer/config/run_settings.py b/src/data_designer/config/run_config.py similarity index 100% rename from src/data_designer/config/run_settings.py rename to src/data_designer/config/run_config.py diff --git a/src/data_designer/engine/resources/resource_provider.py b/src/data_designer/engine/resources/resource_provider.py index 3f2c4b8b9..9c7c85ccf 100644 --- a/src/data_designer/engine/resources/resource_provider.py +++ b/src/data_designer/engine/resources/resource_provider.py @@ -3,7 +3,7 @@ from data_designer.config.base import ConfigBase from data_designer.config.models import ModelConfig -from data_designer.config.run_settings import RunConfig +from data_designer.config.run_config import RunConfig from data_designer.config.seed_source import SeedSource from data_designer.config.utils.type_helpers import StrEnum from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage diff --git a/src/data_designer/essentials/__init__.py b/src/data_designer/essentials/__init__.py index 6933ee0c8..5ef89a13a 100644 --- a/src/data_designer/essentials/__init__.py +++ b/src/data_designer/essentials/__init__.py @@ -3,7 +3,7 @@ from data_designer.config.default_model_settings import resolve_seed_default_model_settings from data_designer.config.exports import * # noqa: F403 -from data_designer.config.run_settings import RunConfig +from data_designer.config.run_config import RunConfig from data_designer.config.validator_params import LocalCallableValidatorParams from data_designer.interface.data_designer import DataDesigner from data_designer.logging import LoggingConfig, configure_logging diff --git a/src/data_designer/interface/data_designer.py b/src/data_designer/interface/data_designer.py index 79370097a..5e9d14c48 100644 --- a/src/data_designer/interface/data_designer.py +++ b/src/data_designer/interface/data_designer.py @@ -20,7 +20,7 @@ ModelProvider, ) from data_designer.config.preview_results import PreviewResults -from data_designer.config.run_settings import RunConfig +from data_designer.config.run_config import RunConfig from data_designer.config.utils.constants import ( DEFAULT_NUM_RECORDS, MANAGED_ASSETS_PATH, diff --git a/tests/engine/conftest.py b/tests/engine/conftest.py index 707575651..c6fdc1727 100644 --- a/tests/engine/conftest.py +++ b/tests/engine/conftest.py @@ -6,7 +6,7 @@ import pandas as pd import pytest -from data_designer.config.run_settings import RunConfig +from data_designer.config.run_config import RunConfig from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage from data_designer.engine.models.facade import ModelFacade from data_designer.engine.models.registry import ModelRegistry diff --git a/tests/engine/dataset_builders/test_column_wise_builder.py b/tests/engine/dataset_builders/test_column_wise_builder.py index 0156fba7a..38bbd006d 100644 --- a/tests/engine/dataset_builders/test_column_wise_builder.py +++ b/tests/engine/dataset_builders/test_column_wise_builder.py @@ -329,7 +329,7 @@ def test_fan_out_with_threads_uses_early_shutdown_settings_from_resource_provide shutdown_error_window: int, ) -> None: """Test that _fan_out_with_threads uses run settings from resource_provider.""" - from data_designer.config.run_settings import RunConfig + from data_designer.config.run_config import RunConfig stub_resource_provider.run_config = RunConfig( disable_early_shutdown=disable_early_shutdown, diff --git a/tests/interface/test_data_designer.py b/tests/interface/test_data_designer.py index 1e0eac73c..27fa21002 100644 --- a/tests/interface/test_data_designer.py +++ b/tests/interface/test_data_designer.py @@ -13,7 +13,7 @@ from data_designer.config.errors import InvalidConfigError from data_designer.config.models import ModelProvider from data_designer.config.processors import DropColumnsProcessorConfig -from data_designer.config.run_settings import RunConfig +from data_designer.config.run_config import RunConfig from data_designer.config.sampler_params import CategorySamplerParams, SamplerType from data_designer.config.seed_source import HuggingFaceSeedSource from data_designer.engine.secret_resolver import CompositeResolver, EnvironmentResolver, PlaintextResolver From 8efcc716020a60c3ccde8c0138be301e3b1813d9 Mon Sep 17 00:00:00 2001 From: "Eric W. Tramel" Date: Thu, 8 Jan 2026 13:57:14 -0500 Subject: [PATCH 6/8] Formatting --- tests/interface/test_data_designer.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/interface/test_data_designer.py b/tests/interface/test_data_designer.py index 27fa21002..745d11f7a 100644 --- a/tests/interface/test_data_designer.py +++ b/tests/interface/test_data_designer.py @@ -161,7 +161,6 @@ def test_run_config_normalizes_error_rate_when_disabled(stub_artifact_path, stub assert data_designer._run_config.shutdown_error_rate == 1.0 - def test_create_dataset_e2e_using_only_sampler_columns( stub_sampler_only_config_builder, stub_artifact_path, stub_model_providers, stub_managed_assets_path ): From b68d89ff6724e17cd6cb36c9e04a8b4672162a78 Mon Sep 17 00:00:00 2001 From: "Eric W. Tramel" Date: Thu, 8 Jan 2026 13:59:31 -0500 Subject: [PATCH 7/8] Backwards Compat for python 3.10 --- src/data_designer/config/run_config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/data_designer/config/run_config.py b/src/data_designer/config/run_config.py index e83e023cb..90bbcc16b 100644 --- a/src/data_designer/config/run_config.py +++ b/src/data_designer/config/run_config.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from typing import Self +from typing_extensions import Self from pydantic import Field, model_validator From a05f116720a90e9fc7a05930d8ab8be3e50985a7 Mon Sep 17 00:00:00 2001 From: "Eric W. Tramel" Date: Thu, 8 Jan 2026 14:41:46 -0500 Subject: [PATCH 8/8] Import ordering --- src/data_designer/config/run_config.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/data_designer/config/run_config.py b/src/data_designer/config/run_config.py index 90bbcc16b..2af5a67ab 100644 --- a/src/data_designer/config/run_config.py +++ b/src/data_designer/config/run_config.py @@ -1,9 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from typing_extensions import Self - from pydantic import Field, model_validator +from typing_extensions import Self from data_designer.config.base import ConfigBase