Skip to content
Draft
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
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,7 @@
LocalFileSeedSource,
)
from data_designer.config.seed_source_dataframe import DataFrameSeedSource # noqa: F401
from data_designer.config.terminal_failure import TerminalTaskFailure # noqa: F401
from data_designer.config.utils.code_lang import CodeLang # noqa: F401
from data_designer.config.utils.info import InfoType # noqa: F401
from data_designer.config.utils.media_helpers import AudioFormat, ImageFormat, VideoFormat # noqa: F401
Expand Down Expand Up @@ -194,6 +195,7 @@
"ResumeMode": (f"{_MOD_BASE}.run_config", "ResumeMode"),
"RunConfig": (f"{_MOD_BASE}.run_config", "RunConfig"),
"ThrottleConfig": (f"{_MOD_BASE}.run_config", "ThrottleConfig"),
"TerminalTaskFailure": (f"{_MOD_BASE}.terminal_failure", "TerminalTaskFailure"),
# script_params
"DataDesignerScriptParams": (f"{_MOD_BASE}.script_params", "DataDesignerScriptParams"),
# scheduling metadata
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ def create(
config_builder: DataDesignerConfigBuilder,
*,
num_records: int = DEFAULT_NUM_RECORDS,
capture_terminal_failures: bool = False,
) -> ResultsT: ...

@abstractmethod
Expand All @@ -41,6 +42,7 @@ def preview(
config_builder: DataDesignerConfigBuilder,
*,
num_records: int = DEFAULT_NUM_RECORDS,
capture_terminal_failures: bool = False,
) -> PreviewResults: ...

@abstractmethod
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
from data_designer.config.config_builder import DataDesignerConfigBuilder
from data_designer.config.dataset_metadata import DatasetMetadata
from data_designer.config.seed_source_dataframe import DataFrameSeedSource
from data_designer.config.terminal_failure import TerminalTaskFailure
from data_designer.config.utils.visualization import WithRecordSamplerMixin

if TYPE_CHECKING:
Expand All @@ -25,6 +26,8 @@ def __init__(
analysis: DatasetProfilerResults | None = None,
processor_artifacts: dict[str, list[dict]] | None = None,
task_traces: list[Any] | None = None,
terminal_failures: list[TerminalTaskFailure] | None = None,
early_shutdown: bool = False,
):
"""Creates a new instance with results from a Data Designer preview run.

Expand All @@ -35,12 +38,18 @@ def __init__(
analysis: Analysis of the preview run.
processor_artifacts: Artifacts generated by the processors.
task_traces: Async scheduler task traces (when DATA_DESIGNER_ASYNC_TRACE=1).
terminal_failures: Terminal column failures captured for omitted seed rows.
Check ``early_shutdown`` before treating this list as complete.
early_shutdown: Whether generation stopped at the global error-rate threshold.
Cancelled rows are not included in ``terminal_failures``.
"""
self.dataset: pd.DataFrame | None = dataset
self.analysis: DatasetProfilerResults | None = analysis
self.processor_artifacts: dict[str, list[dict]] | None = processor_artifacts
self.dataset_metadata: DatasetMetadata | None = dataset_metadata
self.task_traces: list[Any] | None = task_traces
self.terminal_failures: list[TerminalTaskFailure] = list(terminal_failures or [])
self.early_shutdown = early_shutdown
self._config_builder = config_builder

def to_config_builder(self, columns: list[str] | None = None) -> DataDesignerConfigBuilder:
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

from dataclasses import dataclass


@dataclass(frozen=True, order=True, slots=True)
class TerminalTaskFailure:
"""Terminal column failure for an omitted seed row.

``seed_row_index`` is the zero-based position in the requested generation
sequence. It is not necessarily the raw source index for shuffled, selected,
or cycled seed datasets.
"""

seed_row_index: int
column: str
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

import data_designer.lazy_heavy_imports as lazy
from data_designer.config.column_configs import ExpressionColumnConfig, GenerationStrategy
from data_designer.config.terminal_failure import TerminalTaskFailure
from data_designer.engine.capacity import (
AsyncCapacityConfigured,
AsyncCapacityObservedMaxima,
Expand Down Expand Up @@ -214,6 +215,7 @@ def __init__(
adaptive_row_group_initial_target: int = 1,
request_pressure_provider: RequestPressureSnapshotProvider | None = None,
request_pressure_advisory: bool = False,
capture_terminal_failures: bool = False,
) -> None:
self._generators = generators
self._graph = graph
Expand Down Expand Up @@ -339,6 +341,7 @@ def __init__(
# context naturally because the from_scratch task raised; the async
# engine drops rows and continues, losing the cause unless we capture it.
self._first_non_retryable_error: Exception | None = None
self._terminal_failures: list[TerminalTaskFailure] | None = [] if capture_terminal_failures else None
self._fatal_worker_error: BaseException | None = None
self._cancel_requested = Event()
self._run_loop: asyncio.AbstractEventLoop | None = None
Expand Down Expand Up @@ -446,6 +449,11 @@ def first_non_retryable_error(self) -> Exception | None:
"""
return self._first_non_retryable_error

@property
def terminal_failures(self) -> list[TerminalTaskFailure]:
"""Terminal column failures captured for omitted seed rows."""
return sorted(self._terminal_failures or [])

@property
def retryable_outcome_metrics(self) -> dict[str, object]:
"""Return sanitized rolling and cumulative model-task outcome counts."""
Expand Down Expand Up @@ -1542,6 +1550,7 @@ async def _salvage_stalled_row_groups(
already_dropped = task.row_index is not None and self._tracker.is_dropped(task.row_group, task.row_index)
if not already_dropped and self._reporter:
self._reporter.record_failure(task.column)
self._record_terminal_failure(task)
if task.row_index is not None:
self._drop_row(task.row_group, task.row_index, exclude_columns={task.column})
else:
Expand Down Expand Up @@ -1767,6 +1776,22 @@ def _drop_row(self, row_group: int, row_index: int, *, exclude_columns: set[str]
if self._buffer_manager:
self._buffer_manager.drop_row(row_group, row_index)

def _record_terminal_failure(self, task: Task) -> None:
if self._terminal_failures is None:
return

start_offset = self._get_rg_start_offset(task.row_group)
if start_offset is None:
return
row_indices = (task.row_index,) if task.row_index is not None else range(self._get_rg_size(task.row_group))
for row_index in row_indices:
# Preserve the failure that actually caused the row to be omitted.
if self._tracker.is_dropped(task.row_group, row_index):
continue
self._terminal_failures.append(
TerminalTaskFailure(seed_row_index=start_offset + row_index, column=task.column)
)

def _drop_row_group(self, row_group: int, row_group_size: int, *, exclude_columns: set[str] | None = None) -> None:
for row_index in range(row_group_size):
self._drop_row(row_group, row_index, exclude_columns=exclude_columns)
Expand Down Expand Up @@ -2091,6 +2116,7 @@ async def _execute_task_inner_impl(self, task: Task, lease: TaskAdmissionLease,
logger.error("Unexpected %s", log_message, exc_info=True)
# Non-retryable data/user/provider failures drop the affected row(s);
# internal bug-shaped failures above abort the run instead.
self._record_terminal_failure(task)
if task.row_index is not None:
self._drop_row(task.row_group, task.row_index, exclude_columns={task.column})
else:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
ProcessorConfig,
ProcessorType,
)
from data_designer.config.terminal_failure import TerminalTaskFailure
from data_designer.config.utils.type_helpers import StrEnum
from data_designer.config.version import get_library_version
from data_designer.engine.column_generators.generators.base import (
Expand Down Expand Up @@ -193,6 +194,7 @@ def __init__(
# async run, if any. Used by the interface to surface the original cause
# when a run produces 0 records due to deterministic failures.
self._first_non_retryable_error: Exception | None = None
self._terminal_failures: list[TerminalTaskFailure] = []

self._data_designer_config = compile_data_designer_config(data_designer_config, resource_provider)
self._column_configs = compile_dataset_builder_column_configs(self._data_designer_config)
Expand Down Expand Up @@ -239,6 +241,11 @@ def first_non_retryable_error(self) -> Exception | None:
"""First non-retryable error captured by the scheduler in the most recent run."""
return self._first_non_retryable_error

@property
def terminal_failures(self) -> list[TerminalTaskFailure]:
"""Terminal column failures captured during the most recent run."""
return list(self._terminal_failures)

@functools.cached_property
def single_column_configs(self) -> list[ColumnConfigT]:
configs = []
Expand All @@ -256,6 +263,7 @@ def build(
on_batch_complete: Callable[[Path], None] | None = None,
save_multimedia_to_disk: bool = True,
resume: ResumeMode = ResumeMode.NEVER,
capture_terminal_failures: bool = False,
) -> Path:
"""Build the dataset.

Expand All @@ -279,6 +287,7 @@ def build(

In all resume modes, in-flight partial results from the interrupted run are
discarded before generation continues.
capture_terminal_failures: Capture the terminal column for omitted seed rows.

Returns:
Path to the generated dataset directory.
Expand Down Expand Up @@ -351,7 +360,14 @@ def build(
resume = ResumeMode.NEVER
self.artifact_storage.resume = ResumeMode.NEVER

self._build_async(generators, num_records, buffer_size, on_batch_complete, resume=resume)
self._build_async(
generators,
num_records,
buffer_size,
on_batch_complete,
resume=resume,
capture_terminal_failures=capture_terminal_failures,
)

# After-generation processors run unconditionally on the on-disk dataset
# (not gated on ``generated``). When resume sees every row group already
Expand Down Expand Up @@ -537,7 +553,7 @@ def _load_resume_state(self, num_records: int, buffer_size: int) -> _ResumeState
completed_row_groups=completed_row_groups,
)

def build_preview(self, *, num_records: int) -> pd.DataFrame:
def build_preview(self, *, num_records: int, capture_terminal_failures: bool = False) -> pd.DataFrame:
self._reset_run_state()
run_readiness_check(
self.single_column_configs,
Expand All @@ -551,7 +567,11 @@ def build_preview(self, *, num_records: int) -> pd.DataFrame:
generators, self._graph = self._initialize_generators_and_graph()
start_time = time.perf_counter()

dataset = self._build_async_preview(generators, num_records)
dataset = self._build_async_preview(
generators,
num_records,
capture_terminal_failures=capture_terminal_failures,
)

self._resource_provider.model_registry.log_model_usage(time.perf_counter() - start_time)

Expand All @@ -564,8 +584,15 @@ def _reset_run_state(self) -> None:
self._actual_num_records = -1
self._first_non_retryable_error = None
self._task_traces = []
self._terminal_failures = []

def _build_async_preview(self, generators: list[ColumnGenerator], num_records: int) -> pd.DataFrame:
def _build_async_preview(
self,
generators: list[ColumnGenerator],
num_records: int,
*,
capture_terminal_failures: bool = False,
) -> pd.DataFrame:
"""Async preview path - single row group, no disk writes, returns in-memory DataFrame."""
logger.info("⚑ Using async task-queue preview")

Expand All @@ -578,6 +605,7 @@ def _build_async_preview(self, generators: list[ColumnGenerator], num_records: i
buffer_size=num_records,
run_post_batch_in_scheduler=False,
trace=trace_enabled,
capture_terminal_failures=capture_terminal_failures,
)

loop = ensure_async_engine_loop()
Expand All @@ -590,6 +618,7 @@ def _build_async_preview(self, generators: list[ColumnGenerator], num_records: i
self._partial_row_groups = scheduler.partial_row_groups
self._actual_num_records = buffer_manager.actual_num_records
self._first_non_retryable_error = scheduler.first_non_retryable_error
self._terminal_failures = scheduler.terminal_failures

if not buffer_manager.has_row_group(0):
return lazy.pd.DataFrame()
Expand Down Expand Up @@ -746,6 +775,7 @@ def _build_async(
on_batch_complete: Callable[[Path], None] | None = None,
*,
resume: ResumeMode = ResumeMode.NEVER,
capture_terminal_failures: bool = False,
) -> bool:
"""Async task-queue builder path - dispatches tasks based on dependency readiness.

Expand Down Expand Up @@ -851,6 +881,7 @@ def on_complete(final_path: Path | str | None) -> None:
initial_actual_num_records=initial_actual_num_records,
initial_total_num_batches=initial_total_num_batches,
scheduler_event_sink=scheduler_event_sink,
capture_terminal_failures=capture_terminal_failures,
)

# Run on background event loop. Capture scheduler state in `finally`
Expand All @@ -867,6 +898,7 @@ def on_complete(final_path: Path | str | None) -> None:
self._partial_row_groups = scheduler.partial_row_groups
self._actual_num_records = buffer_manager.actual_num_records
self._first_non_retryable_error = scheduler.first_non_retryable_error
self._terminal_failures = scheduler.terminal_failures

# Emit telemetry
try:
Expand Down Expand Up @@ -917,6 +949,7 @@ def _prepare_async_run(
initial_actual_num_records: int = 0,
initial_total_num_batches: int = 0,
scheduler_event_sink: SchedulerAdmissionEventSink | None = None,
capture_terminal_failures: bool = False,
) -> tuple[AsyncTaskScheduler, RowGroupBufferManager]:
"""Build a fully-wired scheduler and buffer manager for async generation.

Expand Down Expand Up @@ -1008,6 +1041,7 @@ def on_before_checkpoint(rg_id: int, rg_size: int) -> None:
),
request_pressure_provider=self._resource_provider.model_registry.request_admission,
request_pressure_advisory=True,
capture_terminal_failures=capture_terminal_failures,
)
return scheduler, buffer_manager

Expand Down
Loading
Loading