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..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 @@ -154,6 +157,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}) @@ -259,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 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/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 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):