Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
216 changes: 93 additions & 123 deletions packages/data-designer-config/src/data_designer/config/run_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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."


Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -220,20 +180,32 @@ 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=(
"Template renderer used for engine-side Jinja evaluation. "
"`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
Expand Down Expand Up @@ -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
Original file line number Diff line number Diff line change
@@ -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,
)
Loading
Loading