feat: expose terminal failure locations

Add opt-in terminal failure capture for create and preview results, including zero-row generation errors. Surface early shutdown so callers can distinguish cancelled rows from attributable terminal failures.

Closes #860

Signed-off-by: Andre Manoel <amanoel@nvidia.com>

Author: andreatnvidia

packages/data-designer-config/src/data_designer/config/__init__.py

@@ -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
@@ -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

packages/data-designer-config/src/data_designer/config/interface.py

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

packages/data-designer-config/src/data_designer/config/preview_results.py

@@ -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:
@@ -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.
 
@@ -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:

packages/data-designer-config/src/data_designer/config/terminal_failure.py

@@ -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

packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py

@@ -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,
@@ -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
@@ -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
@@ -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."""
@@ -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:
@@ -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)
@@ -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:

packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py

@@ -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 (
@@ -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)
@@ -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 = []
@@ -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.
 
@@ -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.
@@ -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
@@ -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,
@@ -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)
 
@@ -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")
 
@@ -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()
@@ -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()
@@ -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.
 
@@ -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`
@@ -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:
@@ -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.
 
@@ -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