Tests to verify passthrough of early shutdown properties

Author: eric-tramel

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