From bc467190cd79e27ffa6663d79997e5997bf027bd Mon Sep 17 00:00:00 2001 From: ChethanUK Date: Fri, 11 Sep 2026 04:34:45 +0200 Subject: [PATCH 1/3] refactor(config): RunConfig hygiene pass With disable_early_shutdown=True, an after-validator overwrote shutdown_error_rate with 1.0 and marked it as set. The value did not round-trip, and a run-config YAML that only disabled early shutdown replaced the baseline rate during the CLI overlay. Replace the validator with a read-only effective_shutdown_error_rate property and point the two engine readers at it, so enforcement gets the same value as before while shutdown_error_rate keeps the user's input. Make Field(description=...) the single place field docs live: the 12 undescribed fields take their text from the Attributes block, which is removed. Group fields with section comments without reordering them. Move the deprecated ThrottleConfig and its message to run_config_deprecated.py and import them back, so every existing import path and old pickles still resolve. No public field names, types, defaults or bounds change. Signed-off-by: ChethanUK --- .../src/data_designer/config/run_config.py | 216 ++++++++---------- .../config/run_config_deprecated.py | 72 ++++++ .../tests/config/test_run_config.py | 15 ++ .../generators/validation.py | 2 +- .../dataset_builders/dataset_builder.py | 2 +- .../generators/test_validation.py | 12 +- .../data_designer/interface/data_designer.py | 6 +- .../controllers/test_generation_controller.py | 5 +- .../tests/cli/utils/test_config_loader.py | 16 +- .../tests/interface/test_data_designer.py | 12 +- 10 files changed, 217 insertions(+), 141 deletions(-) create mode 100644 packages/data-designer-config/src/data_designer/config/run_config_deprecated.py diff --git a/packages/data-designer-config/src/data_designer/config/run_config.py b/packages/data-designer-config/src/data_designer/config/run_config.py index 30aa05d5b..1febecd92 100644 --- a/packages/data-designer-config/src/data_designer/config/run_config.py +++ b/packages/data-designer-config/src/data_designer/config/run_config.py @@ -11,6 +11,7 @@ from typing_extensions import Self from data_designer.config.base import ConfigBase +from data_designer.config.run_config_deprecated import _THROTTLE_DEPRECATION_MESSAGE, ThrottleConfig from data_designer.config.utils.type_helpers import StrEnum from data_designer.config.utils.warning_helpers import warn_at_caller @@ -28,10 +29,6 @@ class ResumeMode(StrEnum): IF_POSSIBLE = "if_possible" -_THROTTLE_DEPRECATION_MESSAGE = ( - "RunConfig.throttle and ThrottleConfig are deprecated. Use RunConfig.request_admission with " - "RequestAdmissionTuningConfig for supported advanced request-admission tuning." -) _PROGRESS_BAR_DEPRECATION_MESSAGE = "RunConfig.progress_bar is deprecated. Use RunConfig.display_tui instead." @@ -74,122 +71,45 @@ class RequestAdmissionTuningConfig(ConfigBase): ) -class ThrottleConfig(ConfigBase): - """Deprecated compatibility DTO for request-admission tuning. - - Use ``RequestAdmissionTuningConfig`` via ``RunConfig.request_admission`` - instead. ``ceiling_overshoot`` is accepted for compatibility but is not - forwarded because request admission no longer exposes an overshoot knob. - """ - - reduce_factor: float = Field( - default=0.75, - gt=0.0, - lt=1.0, - description="Deprecated alias for RequestAdmissionTuningConfig.multiplicative_decrease_factor.", - ) - additive_increase: int = Field( - default=1, - ge=1, - description="Deprecated alias for RequestAdmissionTuningConfig.additive_increase_step.", - ) - success_window: int = Field( - default=25, - ge=1, - description="Deprecated alias for RequestAdmissionTuningConfig.successes_until_increase.", - ) - cooldown_seconds: float = Field( - default=2.0, - gt=0.0, - description="Deprecated alias for RequestAdmissionTuningConfig.cooldown_seconds.", - ) - ceiling_overshoot: float = Field( - default=0.10, - ge=0.0, - description="Deprecated compatibility field; not forwarded to request admission.", - ) - rampup_seconds: float = Field( - default=0.0, - ge=0.0, - description=("Deprecated alias for RequestAdmissionTuningConfig.startup_ramp_seconds."), - ) - - def to_request_admission_tuning(self) -> RequestAdmissionTuningConfig: - """Translate legacy throttle tuning into the request-admission DTO.""" - return RequestAdmissionTuningConfig( - multiplicative_decrease_factor=self.reduce_factor, - additive_increase_step=self.additive_increase, - successes_until_increase=self.success_window, - cooldown_seconds=self.cooldown_seconds, - startup_ramp_seconds=self.rampup_seconds, - ) - - 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 the executor's early-shutdown behavior entirely. - Generation will continue regardless of error rate, and the early-shutdown exception - will never be raised. Error counts and summaries are still collected. Default is False. - shutdown_error_rate: Error rate threshold (0.0-1.0) that triggers early shutdown when - early shutdown is enabled. Default is 0.5. - shutdown_error_window: Minimum number of completed tasks before error rate - monitoring begins. Must be >= 1. Default is 10. - buffer_size: Number of records in each row group during dataset generation. - Must be > 0. Default is 1000. - max_concurrent_row_groups: Maximum number of row groups the async scheduler may - keep active at once. Must be >= 1. Default is 3. - max_in_flight_tasks: Maximum number of async scheduler tasks that may hold task - leases at once. Tasks may be executing, awaiting I/O, or waiting on model - request admission. Model API request concurrency is controlled separately by - ``max_parallel_requests``. Must be >= 1. Default is 1024. - non_inference_max_parallel_workers: Maximum number of worker threads used for non-inference - cell-by-cell generators. Must be >= 1. Default is 4. - max_conversation_restarts: Maximum number of full conversation restarts permitted when - generation tasks call `ModelFacade.generate(...)`. Must be >= 0. Default is 5. - max_conversation_correction_steps: Maximum number of correction rounds permitted within a - single conversation when generation tasks call `ModelFacade.generate(...)`. Must be >= 0. - Default is 0. - async_trace: If True, collect per-task tracing data. Default is False. - write_scheduler_events: If True, create runs write structured scheduler diagnostics to - ``scheduler_events.jsonl`` in the dataset directory. The file is JSONL, not direct - Perfetto input, and may contain sensitive column, provider, model, task, and resource - labels. Each event is flushed to disk, so enabling this option adds file I/O overhead. - Preview runs do not write it. Default is False. - display_tui: If True, display the terminal throughput TUI instead of periodic - log lines during generation. Requires a TTY; falls back to log lines in - non-TTY environments. Default is False. - progress_interval: How often (in seconds) the async progress reporter emits a - consolidated log block. Must be > 0. Default is 5.0. - otel_metrics_port: Loopback port for the pull-based OpenTelemetry metrics endpoint. - Set to None to disable instrumentation for this create invocation. An endpoint - already opened by the process may remain available for metrics from prior runs. - The endpoint exposes metrics, not raw log records. Default is 9464. - preserve_dropped_columns: If True, write columns removed by drop processors to - separate dropped-column parquet files. Set to False to omit those artifacts - while still removing dropped columns from the final dataset. Default is True. - jinja_rendering_engine: Template renderer used for engine-side Jinja evaluation. - ``native`` uses Jinja2's built-in sandbox with the standard filter set and - fewer Data Designer-specific restrictions. ``secure`` uses Data Designer's - hardened sandbox with additional AST, filter, and output guards. - Default is ``secure``. - request_admission: Advanced AIMD request-admission tuning for provider/model calls. - Most users should leave this unset and tune ``max_parallel_requests`` instead. - Notes: - Request-admission controller internals remain engine-owned. This field + Request-admission controller internals remain engine-owned. ``request_admission`` exposes only the supported tuning DTO and does not expose controller mutation APIs, leases, queues, or pressure snapshots. """ - 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=1) - buffer_size: int = Field(default=1000, gt=0) + # Early shutdown + disable_early_shutdown: bool = Field( + default=False, + description=( + "If True, disables the executor's early-shutdown behavior entirely. Generation will continue " + "regardless of error rate, and the early-shutdown exception will never be raised. Error counts " + "and summaries are still collected." + ), + ) + shutdown_error_rate: float = Field( + default=0.5, + ge=0.0, + le=1.0, + description="Error rate threshold (0.0-1.0) that triggers early shutdown when early shutdown is enabled.", + ) + shutdown_error_window: int = Field( + default=10, + ge=1, + description="Minimum number of completed tasks before error rate monitoring begins.", + ) + + # Scheduling and concurrency + buffer_size: int = Field( + default=1000, + gt=0, + description="Number of records in each row group during dataset generation.", + ) max_concurrent_row_groups: int = Field( default=3, ge=1, @@ -203,13 +123,53 @@ class RunConfig(ConfigBase): "Model API request concurrency is controlled separately by max_parallel_requests." ), ) - non_inference_max_parallel_workers: int = Field(default=4, ge=1) - max_conversation_restarts: int = Field(default=5, ge=0) - max_conversation_correction_steps: int = Field(default=0, ge=0) - async_trace: bool = False - write_scheduler_events: bool = False - display_tui: bool = False - progress_interval: float = Field(default=5.0, gt=0.0) + non_inference_max_parallel_workers: int = Field( + default=4, + ge=1, + description="Maximum number of worker threads used for non-inference cell-by-cell generators.", + ) + + # Conversation recovery + max_conversation_restarts: int = Field( + default=5, + ge=0, + description=( + "Maximum number of full conversation restarts permitted when generation tasks call " + "`ModelFacade.generate(...)`." + ), + ) + max_conversation_correction_steps: int = Field( + default=0, + ge=0, + description=( + "Maximum number of correction rounds permitted within a single conversation when generation " + "tasks call `ModelFacade.generate(...)`." + ), + ) + + # Observability + async_trace: bool = Field(default=False, description="If True, collect per-task tracing data.") + write_scheduler_events: bool = Field( + default=False, + description=( + "If True, create runs write structured scheduler diagnostics to ``scheduler_events.jsonl`` in the " + "dataset directory. The file is JSONL, not direct Perfetto input, and may contain sensitive column, " + "provider, model, task, and resource labels. Each event is flushed to disk, so enabling this option " + "adds file I/O overhead. Preview runs do not write it." + ), + ) + display_tui: bool = Field( + default=False, + description=( + "If True, display the terminal throughput TUI instead of periodic log lines during generation. " + "Requires a TTY; falls back to log lines in non-TTY environments." + ), + ) + progress_interval: float = Field( + default=5.0, + gt=0.0, + description="How often (in seconds) the async progress reporter emits a consolidated log block.", + ) otel_metrics_port: int | None = Field( default=9464, ge=1, @@ -220,12 +180,16 @@ class RunConfig(ConfigBase): "available for prior metrics. Raw log records are not exposed." ), ) + + # Output preserve_dropped_columns: bool = Field( default=True, description=( "Whether columns removed by drop processors are preserved in separate dropped-column parquet files." ), ) + + # Templating jinja_rendering_engine: JinjaRenderingEngine = Field( default=JinjaRenderingEngine.SECURE, description=( @@ -233,7 +197,15 @@ class RunConfig(ConfigBase): "`native` uses Jinja2's built-in sandbox; `secure` uses Data Designer's hardened sandbox." ), ) - request_admission: RequestAdmissionTuningConfig | None = None + + # Request admission + request_admission: RequestAdmissionTuningConfig | None = Field( + default=None, + description=( + "Advanced AIMD request-admission tuning for provider/model calls. " + "Most users should leave this unset and tune ``max_parallel_requests`` instead." + ), + ) @model_validator(mode="before") @classmethod @@ -301,9 +273,7 @@ def model_copy(self, *, update: Mapping[str, Any] | None = None, deep: bool = Fa update = normalized_update return super().model_copy(update=update, deep=deep) - @model_validator(mode="after") - def normalize_shutdown_settings(self) -> Self: - """Normalize shutdown settings for compatibility.""" - if self.disable_early_shutdown: - self.shutdown_error_rate = 1.0 - return self + @property + def effective_shutdown_error_rate(self) -> float: + """Error rate handed to early-shutdown checks: 1.0 when early shutdown is disabled.""" + return 1.0 if self.disable_early_shutdown else self.shutdown_error_rate diff --git a/packages/data-designer-config/src/data_designer/config/run_config_deprecated.py b/packages/data-designer-config/src/data_designer/config/run_config_deprecated.py new file mode 100644 index 000000000..ac0ac5ded --- /dev/null +++ b/packages/data-designer-config/src/data_designer/config/run_config_deprecated.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from pydantic import Field + +from data_designer.config.base import ConfigBase + +if TYPE_CHECKING: + from data_designer.config.run_config import RequestAdmissionTuningConfig + +_THROTTLE_DEPRECATION_MESSAGE = ( + "RunConfig.throttle and ThrottleConfig are deprecated. Use RunConfig.request_admission with " + "RequestAdmissionTuningConfig for supported advanced request-admission tuning." +) + + +class ThrottleConfig(ConfigBase): + """Deprecated compatibility DTO for request-admission tuning. + + Use ``RequestAdmissionTuningConfig`` via ``RunConfig.request_admission`` + instead. ``ceiling_overshoot`` is accepted for compatibility but is not + forwarded because request admission no longer exposes an overshoot knob. + """ + + reduce_factor: float = Field( + default=0.75, + gt=0.0, + lt=1.0, + description="Deprecated alias for RequestAdmissionTuningConfig.multiplicative_decrease_factor.", + ) + additive_increase: int = Field( + default=1, + ge=1, + description="Deprecated alias for RequestAdmissionTuningConfig.additive_increase_step.", + ) + success_window: int = Field( + default=25, + ge=1, + description="Deprecated alias for RequestAdmissionTuningConfig.successes_until_increase.", + ) + cooldown_seconds: float = Field( + default=2.0, + gt=0.0, + description="Deprecated alias for RequestAdmissionTuningConfig.cooldown_seconds.", + ) + ceiling_overshoot: float = Field( + default=0.10, + ge=0.0, + description="Deprecated compatibility field; not forwarded to request admission.", + ) + rampup_seconds: float = Field( + default=0.0, + ge=0.0, + description=("Deprecated alias for RequestAdmissionTuningConfig.startup_ramp_seconds."), + ) + + def to_request_admission_tuning(self) -> RequestAdmissionTuningConfig: + """Translate legacy throttle tuning into the request-admission DTO.""" + # Runtime import: module-level would cause circular import + from data_designer.config.run_config import RequestAdmissionTuningConfig + + return RequestAdmissionTuningConfig( + multiplicative_decrease_factor=self.reduce_factor, + additive_increase_step=self.additive_increase, + successes_until_increase=self.success_window, + cooldown_seconds=self.cooldown_seconds, + startup_ramp_seconds=self.rampup_seconds, + ) diff --git a/packages/data-designer-config/tests/config/test_run_config.py b/packages/data-designer-config/tests/config/test_run_config.py index f92e2b077..b4abd5103 100644 --- a/packages/data-designer-config/tests/config/test_run_config.py +++ b/packages/data-designer-config/tests/config/test_run_config.py @@ -154,6 +154,21 @@ def test_run_config_rejects_invalid_max_in_flight_tasks() -> None: RunConfig(max_in_flight_tasks=0) +@pytest.mark.parametrize( + ("disable", "rate", "expected_effective"), + [(False, 0.2, 0.2), (True, 0.2, 1.0), (True, 0.0, 1.0), (False, 1.0, 1.0)], + ids=["enabled-keeps-rate", "disabled-overrides-rate", "disabled-overrides-zero-rate", "enabled-max-rate"], +) +def test_run_config_keeps_rate_and_derives_effective_rate( + disable: bool, rate: float, expected_effective: float +) -> None: + run_config = RunConfig(disable_early_shutdown=disable, shutdown_error_rate=rate) + + assert run_config.shutdown_error_rate == rate + assert run_config.effective_shutdown_error_rate == expected_effective + assert RunConfig.model_validate(run_config.model_dump()) == run_config + + def test_run_config_throttle_shim_rejects_unknown_legacy_fields() -> None: with pytest.raises(ValidationError, match="max_concurrent_requests"): RunConfig(throttle={"max_concurrent_requests": 1}) diff --git a/packages/data-designer-engine/src/data_designer/engine/column_generators/generators/validation.py b/packages/data-designer-engine/src/data_designer/engine/column_generators/generators/validation.py index cec75d247..19a542074 100644 --- a/packages/data-designer-engine/src/data_designer/engine/column_generators/generators/validation.py +++ b/packages/data-designer-engine/src/data_designer/engine/column_generators/generators/validation.py @@ -119,7 +119,7 @@ def error_callback(error: Exception, context: dict): column_name=self.config.name, result_callback=result_callback, error_callback=error_callback, - shutdown_error_rate=settings.shutdown_error_rate, + shutdown_error_rate=settings.effective_shutdown_error_rate, shutdown_error_window=settings.shutdown_error_window, disable_early_shutdown=settings.disable_early_shutdown, ) as executor: diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py index 5615571f9..43c277e32 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py @@ -843,7 +843,7 @@ def on_complete(final_path: Path | str | None) -> None: num_records, buffer_size, on_finalize_row_group=finalize_row_group, - shutdown_error_rate=settings.shutdown_error_rate, + shutdown_error_rate=settings.effective_shutdown_error_rate, shutdown_error_window=settings.shutdown_error_window, disable_early_shutdown=settings.disable_early_shutdown, trace=trace_enabled, diff --git a/packages/data-designer-engine/tests/engine/column_generators/generators/test_validation.py b/packages/data-designer-engine/tests/engine/column_generators/generators/test_validation.py index cd4befad1..d546b4f6a 100644 --- a/packages/data-designer-engine/tests/engine/column_generators/generators/test_validation.py +++ b/packages/data-designer-engine/tests/engine/column_generators/generators/test_validation.py @@ -9,6 +9,7 @@ import data_designer.lazy_heavy_imports as lazy from data_designer.config.column_configs import ValidationColumnConfig +from data_designer.config.run_config import RunConfig from data_designer.config.utils.code_lang import CodeLang from data_designer.config.validator_params import ( CodeValidatorParams, @@ -203,8 +204,16 @@ def test_validation_column_generator_generate_with_different_strategies( assert len(result["validation_column"]) == len(df) +@pytest.mark.parametrize( + ("disable", "rate", "expected_rate"), + [(True, 0.2, 1.0), (False, 0.2, 0.2)], + ids=["disabled-passes-effective-rate", "enabled-passes-configured-rate"], +) @patch("data_designer.engine.column_generators.generators.validation.get_validator_from_params", autospec=True) -def test_validation_column_generator_validate_in_parallel_failure(mock_get_validator, stub_resource_provider): +def test_validation_column_generator_validate_in_parallel_failure( + mock_get_validator: Mock, stub_resource_provider: Mock, disable: bool, rate: float, expected_rate: float +) -> None: + stub_resource_provider.run_config = RunConfig(disable_early_shutdown=disable, shutdown_error_rate=rate) mock_validator = Mock() mock_validator.run_validation.return_value = ValidationResult(data=[ValidationOutput(is_valid=True)]) mock_get_validator.return_value = mock_validator @@ -242,3 +251,4 @@ def mock_submit(func, batch, context): call_kwargs = mock_executor_class.call_args[1] assert call_kwargs["disable_early_shutdown"] == stub_resource_provider.run_config.disable_early_shutdown + assert call_kwargs["shutdown_error_rate"] == expected_rate 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..1ab77e60d 100644 --- a/packages/data-designer/src/data_designer/interface/data_designer.py +++ b/packages/data-designer/src/data_designer/interface/data_designer.py @@ -623,11 +623,7 @@ def run_config(self) -> RunConfig: """Get the runtime configuration applied to dataset generation. Returns: - The active RunConfig instance. Note that ``RunConfig`` normalizes - some fields on construction (e.g., ``shutdown_error_rate`` becomes - ``1.0`` when ``disable_early_shutdown=True``), so the returned - object may not exactly equal the one originally passed to - ``set_run_config``. + The active RunConfig instance. """ return self._run_config diff --git a/packages/data-designer/tests/cli/controllers/test_generation_controller.py b/packages/data-designer/tests/cli/controllers/test_generation_controller.py index 9cfd2a51d..a35c9e2cc 100644 --- a/packages/data-designer/tests/cli/controllers/test_generation_controller.py +++ b/packages/data-designer/tests/cli/controllers/test_generation_controller.py @@ -866,7 +866,7 @@ def test_run_create_applies_run_config_precedence( } ) mock_dd = MagicMock() - mock_dd.run_config = RunConfig(buffer_size=64, progress_interval=11.0, display_tui=True) + mock_dd.run_config = RunConfig(buffer_size=64, progress_interval=11.0, display_tui=True, shutdown_error_rate=0.3) mock_dd.create.return_value = _make_mock_create_results() mock_dd_cls.return_value = mock_dd @@ -886,7 +886,8 @@ def test_run_create_applies_run_config_precedence( assert effective.buffer_size == 1000 assert effective.progress_interval == 11.0 assert effective.preserve_dropped_columns is False - assert effective.shutdown_error_rate == 1.0 + assert effective.shutdown_error_rate == 0.3 + assert effective.effective_shutdown_error_rate == 1.0 assert effective.display_tui is False assert [item[0] for item in mock_dd.method_calls] == ["set_run_config", "create"] assert "Run config: run.yaml" in capsys.readouterr().out diff --git a/packages/data-designer/tests/cli/utils/test_config_loader.py b/packages/data-designer/tests/cli/utils/test_config_loader.py index 1e7944f6c..0c68b5d78 100644 --- a/packages/data-designer/tests/cli/utils/test_config_loader.py +++ b/packages/data-designer/tests/cli/utils/test_config_loader.py @@ -68,13 +68,23 @@ def test_load_run_config_accepts_empty_mapping(tmp_path: Path) -> None: assert loaded.model_dump(exclude_unset=True) == {} -def test_load_run_config_preserves_explicit_partial_fields(tmp_path: Path) -> None: +@pytest.mark.parametrize( + ("yaml_text", "expected"), + [ + ("buffer_size: 250\ndisplay_tui: true\n", {"buffer_size": 250, "display_tui": True}), + ("disable_early_shutdown: true\n", {"disable_early_shutdown": True}), + ], + ids=["buffer-size-and-tui", "disable-early-shutdown-only"], +) +def test_load_run_config_preserves_explicit_partial_fields( + tmp_path: Path, yaml_text: str, expected: dict[str, object] +) -> None: run_config_file = tmp_path / "run-config.yaml" - run_config_file.write_text("buffer_size: 250\ndisplay_tui: true\n") + run_config_file.write_text(yaml_text) loaded = load_run_config(str(run_config_file)) - assert loaded.model_dump(exclude_unset=True) == {"buffer_size": 250, "display_tui": True} + assert loaded.model_dump(exclude_unset=True) == expected def test_load_run_config_preserves_partial_nested_fields(tmp_path: Path) -> None: diff --git a/packages/data-designer/tests/interface/test_data_designer.py b/packages/data-designer/tests/interface/test_data_designer.py index 0a0ce5a9c..59dfafa88 100644 --- a/packages/data-designer/tests/interface/test_data_designer.py +++ b/packages/data-designer/tests/interface/test_data_designer.py @@ -491,7 +491,8 @@ def test_run_config_setting_persists(stub_artifact_path, stub_model_providers): ) ) 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_rate == 0.8 + assert data_designer.run_config.effective_shutdown_error_rate == 1.0 assert data_designer.run_config.shutdown_error_window == 25 assert data_designer.run_config.buffer_size == 500 assert data_designer.run_config.max_in_flight_tasks == 1536 @@ -561,8 +562,8 @@ def test_resource_provider_uses_otel_sink_only_when_enabled( assert create_provider.call_args.kwargs["scheduler_event_sink"] is None -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.""" +def test_run_config_keeps_error_rate_when_disabled(stub_artifact_path, stub_model_providers): + """Test that shutdown_error_rate round-trips and only the effective rate is 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 @@ -574,14 +575,15 @@ def test_run_config_normalizes_error_rate_when_disabled(stub_artifact_path, stub ) assert data_designer.run_config.shutdown_error_rate == 0.7 - # When disabled, shutdown_error_rate should be normalized to 1.0 + # When disabled, shutdown_error_rate keeps the configured value; the effective rate is 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 + assert data_designer.run_config.shutdown_error_rate == 0.7 + assert data_designer.run_config.effective_shutdown_error_rate == 1.0 def test_get_models_uses_sync_clients(stub_artifact_path, stub_model_providers): From 766fafd941b9401e0a4ea02914293d378b0bb3c4 Mon Sep 17 00:00:00 2001 From: ChethanUK Date: Sat, 12 Sep 2026 01:07:24 +0200 Subject: [PATCH 2/3] test(cli): cover create --run-config overlay The run-config overlay was only exercised with load_run_config mocked, so the YAML-file hop that produces the reported symptom was never covered through the command a user actually runs. Invoke the real create command with a run-config YAML on disk and assert the baseline shutdown_error_rate survives the overlay, with the derived effective rate still 1.0 when early shutdown is disabled. Signed-off-by: ChethanUK --- .../tests/cli/commands/test_create_command.py | 54 +++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/packages/data-designer/tests/cli/commands/test_create_command.py b/packages/data-designer/tests/cli/commands/test_create_command.py index 6b07c7a1e..1f4c0f76c 100644 --- a/packages/data-designer/tests/cli/commands/test_create_command.py +++ b/packages/data-designer/tests/cli/commands/test_create_command.py @@ -3,11 +3,20 @@ from __future__ import annotations +from pathlib import Path from unittest.mock import MagicMock, patch +import pytest +from typer.testing import CliRunner + from data_designer.cli.commands.create import create_command +from data_designer.cli.main import app +from data_designer.config.run_config import RunConfig from data_designer.engine.storage.artifact_storage import ResumeMode +runner = CliRunner() +_CTRL = "data_designer.cli.controllers.generation_controller" + # --------------------------------------------------------------------------- # create_command delegation tests # --------------------------------------------------------------------------- @@ -229,3 +238,48 @@ def test_create_command_passes_tui_override(mock_ctrl_cls: MagicMock) -> None: tui=False, script_args=None, ) + + +# --------------------------------------------------------------------------- +# create --run-config overlay tests +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("yaml_text", "expected_rate", "expected_effective_rate"), + [ + ("disable_early_shutdown: true\n", 0.3, 1.0), + ("disable_early_shutdown: true\nshutdown_error_rate: 0.45\n", 0.45, 1.0), + ("shutdown_error_rate: 0.45\n", 0.45, 0.45), + ], + ids=["flag-only-keeps-baseline-rate", "flag-with-explicit-rate", "rate-only"], +) +@patch(f"{_CTRL}.DataDesigner") +@patch(f"{_CTRL}.load_config_builder") +def test_create_run_config_overlay_preserves_shutdown_error_rate( + mock_load_config: MagicMock, + mock_dd_cls: MagicMock, + tmp_path: Path, + yaml_text: str, + expected_rate: float, + expected_effective_rate: float, +) -> None: + """`create --run-config` overlays the YAML file without discarding the baseline error rate.""" + run_config_file = tmp_path / "run-config.yaml" + run_config_file.write_text(yaml_text) + + mock_load_config.return_value = MagicMock() + mock_results = MagicMock() + mock_results.count_records.return_value = 10 + mock_results.load_analysis.return_value = None + mock_dd = MagicMock() + mock_dd.run_config = RunConfig(shutdown_error_rate=0.3) + mock_dd.create.return_value = mock_results + mock_dd_cls.return_value = mock_dd + + result = runner.invoke(app, ["create", "dataset.yaml", "--run-config", str(run_config_file)]) + + assert result.exit_code == 0, result.output + effective = mock_dd.set_run_config.call_args.args[0] + assert effective.shutdown_error_rate == expected_rate + assert effective.effective_shutdown_error_rate == expected_effective_rate From 458ae4e36a00680f5314aa5c3b5aa2a0deb0c796 Mon Sep 17 00:00:00 2001 From: ChethanUK Date: Sat, 12 Sep 2026 03:40:48 +0200 Subject: [PATCH 3/3] test(config): cover legacy ThrottleConfig pickle The move to run_config_deprecated changed the module ThrottleConfig reports. The compatibility claim that pickles written before the move still load rested on the module-level re-import in run_config, with nothing covering it. Rebuild a pre-move payload by pickling with __module__ set to the old path and assert it loads. Dropping the re-import makes this fail with AttributeError, so the test guards that binding instead of restating it. Signed-off-by: ChethanUK --- .../tests/config/test_run_config.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/packages/data-designer-config/tests/config/test_run_config.py b/packages/data-designer-config/tests/config/test_run_config.py index b4abd5103..2124e16ae 100644 --- a/packages/data-designer-config/tests/config/test_run_config.py +++ b/packages/data-designer-config/tests/config/test_run_config.py @@ -3,6 +3,9 @@ from __future__ import annotations +import pickle +from unittest.mock import patch + import pytest from pydantic import ValidationError @@ -274,6 +277,17 @@ def test_deprecated_throttle_config_is_exported_from_config_package() -> None: assert namespace["ThrottleConfig"] is ThrottleConfig +def test_deprecated_throttle_config_unpickles_from_pre_move_module_path() -> None: + # Pickles written before ThrottleConfig moved to run_config_deprecated record + # data_designer.config.run_config; patching __module__ reproduces those bytes. Unpickling + # resolves them through the module-level re-import in run_config. + with patch.object(ThrottleConfig, "__module__", "data_designer.config.run_config"): + legacy_payload = pickle.dumps(ThrottleConfig(reduce_factor=0.5)) + + assert b"data_designer.config.run_config_deprecated" not in legacy_payload + assert pickle.loads(legacy_payload) == ThrottleConfig(reduce_factor=0.5) + + def test_throttle_config_accepts_rampup_seconds() -> None: config = ThrottleConfig(rampup_seconds=30.0) assert config.rampup_seconds == 30.0