Change to a config type and helper settings

Author: eric-tramel

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

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})

src/data_designer/engine/dataset_builders/column_wise_builder.py

@@ -60,19 +60,13 @@ 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
         self._records_to_drop: set[int] = set()
         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})

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(

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

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:

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
 
 

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

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: