fix(async): salvage partial row groups on early shutdown

Before: when the early-shutdown gate fired, any row group still in
flight stayed in `_rg_states` un-checkpointed. The buffer manager
later raised `FileNotFoundError` when the builder tried to read the
finalized parquet. User-visible result: `0 records produced`.

After: a new `_finalize_after_shutdown` step runs in `run()`'s finally
block, after `_cancel_workers` has drained in-flight tasks (Codex
caveat: in-flight `from_scratch`/`batch` tasks must not be allowed to
write into a buffer that's being finalized). For each remaining row
group it drops rows that aren't fully complete, then delegates to the
existing `_checkpoint_completed_row_groups` so the buffer manager's
zero-survivor handling (skip empty parquet, free buffer) kicks in
unchanged.

Also surfaces partial completion as a structured signal: scheduler
exposes `early_shutdown: bool` and `partial_row_groups: tuple[int, ...]`
properties so callers can detect partial completion programmatically
rather than parsing log lines. Builder uses this to emit a more
specific WARN distinguishing early shutdown from non-shutdown drops.

Refs #575.

Author: andreatgretel

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

@@ -195,6 +195,11 @@ def __init__(
         self._recent_retryable: deque[bool] = deque(maxlen=degraded_warn_window)
         self._last_degraded_warn_at: float = 0.0
 
+        # Row groups that were partially salvaged after early shutdown
+        # (i.e., some rows complete, some incomplete-then-dropped). Surfaced
+        # via the partial_row_groups property as a structured signal.
+        self._partial_row_groups: list[int] = []
+
         # Pre-compute row-group sizes for O(1) lookup
         self._rg_size_map: dict[int, int] = dict(row_groups)
 
@@ -239,6 +244,20 @@ def _setup_async_progress_reporter(
     def active_worker_count(self) -> int:
         return sum(1 for t in self._worker_tasks if not t.done())
 
+    @property
+    def early_shutdown(self) -> bool:
+        """True if the run terminated via the early-shutdown gate."""
+        return self._early_shutdown
+
+    @property
+    def partial_row_groups(self) -> tuple[int, ...]:
+        """Row group ids that were partially salvaged after early shutdown.
+
+        Empty unless ``early_shutdown`` is True. Each id had some rows
+        complete and the rest dropped before checkpointing.
+        """
+        return tuple(self._partial_row_groups)
+
     def _spawn_worker(self, coro: Coroutine[Any, Any, None]) -> asyncio.Task:
         """Create a tracked worker task that auto-removes itself on completion."""
         task = asyncio.create_task(coro)
@@ -304,6 +323,11 @@ async def run(self) -> None:
                     with contextlib.suppress(asyncio.CancelledError):
                         await admission_task
                 await asyncio.shield(self._cancel_workers())
+                # Salvage partially-complete row groups left over from early
+                # shutdown. Must run AFTER _cancel_workers - in-flight tasks
+                # could otherwise write into a buffer that's being finalized.
+                if self._early_shutdown and self._rg_states:
+                    self._finalize_after_shutdown(all_columns)
 
             if self._reporter:
                 self._reporter.log_final()
@@ -552,6 +576,37 @@ def _checkpoint_completed_row_groups(self, all_columns: list[str]) -> None:
             checkpointed = {rg_id for rg_id, _ in completed}
             self._deferred = [t for t in self._deferred if t.row_group not in checkpointed]
 
+    def _finalize_after_shutdown(self, all_columns: list[str]) -> None:
+        """Salvage row groups left in flight when early shutdown fired.
+
+        For each remaining row group, drop rows that aren't fully complete
+        (and weren't already dropped); after that, ``is_row_group_complete``
+        is true by construction over the surviving rows, so delegating to
+        ``_checkpoint_completed_row_groups`` writes survivors and frees
+        zero-survivor groups via the buffer manager's existing logic.
+        """
+        for rg_id in list(self._rg_states.keys()):
+            rg_size = self._rg_states[rg_id].size
+            had_incomplete = False
+            for ri in range(rg_size):
+                if self._tracker.is_dropped(rg_id, ri):
+                    continue
+                if all(
+                    self._tracker.is_complete(SliceRef(column=col, row_group=rg_id, row_index=ri))
+                    for col in all_columns
+                ):
+                    continue
+                had_incomplete = True
+                self._drop_row(rg_id, ri)
+            if had_incomplete:
+                survivors = sum(1 for ri in range(rg_size) if not self._tracker.is_dropped(rg_id, ri))
+                if survivors > 0:
+                    self._partial_row_groups.append(rg_id)
+                    logger.warning(f"Row group {rg_id}: salvaging {survivors} of {rg_size} rows after early shutdown.")
+                else:
+                    logger.warning(f"Row group {rg_id}: 0 of {rg_size} rows survived early shutdown - skipping write.")
+        self._checkpoint_completed_row_groups(all_columns)
+
     def _run_seeds_complete_check(self, seed_cols: frozenset[str]) -> None:
         """Run pre-batch callbacks for row groups whose seeds just completed."""
         for rg_id, state in list(self._rg_states.items()):

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

@@ -341,10 +341,18 @@ def on_complete(final_path: Path | str | None) -> None:
         actual = buffer_manager.actual_num_records
         if actual < num_records:
             pct = actual / num_records * 100 if num_records > 0 else 0
-            logger.warning(
-                f"⚠️ Generated {actual} of {num_records} requested records ({pct:.0f}%). "
-                "The dataset may be incomplete due to errors or early shutdown."
-            )
+            base = f"⚠️ Generated {actual} of {num_records} requested records ({pct:.0f}%). "
+            if scheduler.early_shutdown:
+                partial = scheduler.partial_row_groups
+                detail = (
+                    f"Early shutdown was triggered (non-retryable error rate exceeded threshold); "
+                    f"{len(partial)} row group(s) salvaged with partial rows."
+                    if partial
+                    else "Early shutdown was triggered (non-retryable error rate exceeded threshold)."
+                )
+                logger.warning(base + detail)
+            else:
+                logger.warning(base + "The dataset may be incomplete due to dropped rows.")
 
     def _prepare_async_run(
         self,

packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py

@@ -173,6 +173,52 @@ def generate(self, data: dict) -> dict:
         return data
 
 
+class MockSelectiveFailGenerator(ColumnGenerator[ExpressionColumnConfig]):
+    """Cell generator with deterministic per-seed behavior.
+
+    - Seeds in ``fail_on_seeds``: raise a non-retryable ``ValueError`` immediately.
+    - Seeds in ``slow_seeds``: block on ``slow_event`` (or ``asyncio.sleep``) so
+      they remain in-flight when the early-shutdown gate fires.
+    - All others: succeed.
+    """
+
+    def __init__(
+        self,
+        *args: Any,
+        fail_on_seeds: set[int] = frozenset(),
+        slow_seeds: set[int] = frozenset(),
+        slow_timeout_s: float = 5.0,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self._fail = set(fail_on_seeds)
+        self._slow = set(slow_seeds)
+        self._slow_timeout_s = slow_timeout_s
+
+    @staticmethod
+    def get_generation_strategy() -> GenerationStrategy:
+        return GenerationStrategy.CELL_BY_CELL
+
+    async def agenerate(self, data: dict) -> dict:
+        seed = data.get("seed")
+        if seed in self._fail:
+            raise ValueError(f"non-retryable on seed={seed}")
+        if seed in self._slow:
+            try:
+                await asyncio.sleep(self._slow_timeout_s)
+            except asyncio.CancelledError:
+                raise
+        data[self.config.name] = f"ok_{seed}"
+        return data
+
+    def generate(self, data: dict) -> dict:
+        seed = data.get("seed")
+        if seed in self._fail:
+            raise ValueError(f"non-retryable on seed={seed}")
+        data[self.config.name] = f"ok_{seed}"
+        return data
+
+
 class MockRetryableErrorGenerator(ColumnGenerator[ExpressionColumnConfig]):
     """Generator that raises a parametrizable retryable error then succeeds."""
 
@@ -722,6 +768,141 @@ async def test_scheduler_error_rate_shutdown() -> None:
 
     # Early shutdown: not all rows should be checkpointed (some row groups incomplete)
     assert buffer_mgr.actual_num_records < 10
+    # No leftover unfinished row groups (finalize-after-shutdown drains them).
+    assert not scheduler._rg_states
+
+
+@pytest.mark.asyncio(loop_scope="session")
+async def test_partial_row_group_salvaged_after_early_shutdown() -> None:
+    """Mid-run shutdown drops incomplete rows and checkpoints survivors."""
+    provider = _mock_provider()
+    configs = [
+        SamplerColumnConfig(name="seed", sampler_type=SamplerType.CATEGORY, params={"values": ["A"]}),
+        LLMTextColumnConfig(name="cell_out", prompt="{{ seed }}", model_alias=MODEL_ALIAS),
+    ]
+    strategies = {
+        "seed": GenerationStrategy.FULL_COLUMN,
+        "cell_out": GenerationStrategy.CELL_BY_CELL,
+    }
+    # 3 succeed (0,1,2), 3 fail non-retryable (5,6,7), 4 stay in-flight (3,4,8,9)
+    # until cancellation. Window=4, rate=0.5 → gate trips after ~3-5 outcomes.
+    generators = {
+        "seed": MockSeedGenerator(config=_expr_config("seed"), resource_provider=provider),
+        "cell_out": MockSelectiveFailGenerator(
+            config=_expr_config("cell_out"),
+            resource_provider=provider,
+            fail_on_seeds={5, 6, 7},
+            slow_seeds={3, 4, 8, 9},
+        ),
+    }
+
+    graph = ExecutionGraph.create(configs, strategies)
+    row_groups = [(0, 10)]
+    tracker = CompletionTracker.with_graph(graph, row_groups)
+
+    storage = MagicMock()
+    storage.dataset_name = "test"
+    storage.get_file_paths.return_value = {}
+    storage.write_batch_to_parquet_file.return_value = "/fake.parquet"
+    storage.move_partial_result_to_final_file_path.return_value = "/fake_final.parquet"
+    buffer_mgr = RowGroupBufferManager(storage)
+
+    finalized: list[int] = []
+
+    def on_finalize(rg_id: int) -> None:
+        buffer_mgr.checkpoint_row_group(rg_id)
+        finalized.append(rg_id)
+
+    scheduler = AsyncTaskScheduler(
+        generators=generators,
+        graph=graph,
+        tracker=tracker,
+        row_groups=row_groups,
+        buffer_manager=buffer_mgr,
+        on_finalize_row_group=on_finalize,
+        shutdown_error_rate=0.5,
+        shutdown_error_window=4,
+    )
+    await scheduler.run()
+
+    assert scheduler.early_shutdown
+    # The row group survived with the 3 fast successes; the in-flight rows were
+    # cancelled and dropped by _finalize_after_shutdown.
+    assert 0 in finalized
+    assert scheduler.partial_row_groups == (0,)
+    # Exactly 3 rows survived (seeds 0, 1, 2).
+    assert buffer_mgr.actual_num_records == 3
+
+
+@pytest.mark.asyncio(loop_scope="session")
+async def test_zero_survivor_shutdown_does_not_raise() -> None:
+    """If every row is dropped at shutdown, the row group is freed without writing parquet."""
+    provider = _mock_provider()
+    configs = [
+        SamplerColumnConfig(name="seed", sampler_type=SamplerType.CATEGORY, params={"values": ["A"]}),
+        LLMTextColumnConfig(name="cell_out", prompt="{{ seed }}", model_alias=MODEL_ALIAS),
+    ]
+    strategies = {
+        "seed": GenerationStrategy.FULL_COLUMN,
+        "cell_out": GenerationStrategy.CELL_BY_CELL,
+    }
+    # All 5 seeds fail non-retryable → all rows dropped before any can complete.
+    generators = {
+        "seed": MockSeedGenerator(config=_expr_config("seed"), resource_provider=provider),
+        "cell_out": MockSelectiveFailGenerator(
+            config=_expr_config("cell_out"),
+            resource_provider=provider,
+            fail_on_seeds=set(range(5)),
+        ),
+    }
+
+    graph = ExecutionGraph.create(configs, strategies)
+    row_groups = [(0, 5)]
+    tracker = CompletionTracker.with_graph(graph, row_groups)
+
+    storage = MagicMock()
+    storage.dataset_name = "test"
+    storage.get_file_paths.return_value = {}
+    storage.write_batch_to_parquet_file.return_value = "/fake.parquet"
+    storage.move_partial_result_to_final_file_path.return_value = "/fake_final.parquet"
+    buffer_mgr = RowGroupBufferManager(storage)
+
+    finalized: list[int] = []
+
+    def on_finalize(rg_id: int) -> None:
+        buffer_mgr.checkpoint_row_group(rg_id)
+        finalized.append(rg_id)
+
+    scheduler = AsyncTaskScheduler(
+        generators=generators,
+        graph=graph,
+        tracker=tracker,
+        row_groups=row_groups,
+        buffer_manager=buffer_mgr,
+        on_finalize_row_group=on_finalize,
+        shutdown_error_rate=0.5,
+        shutdown_error_window=2,
+    )
+    # Must not raise (no FileNotFoundError, no DataDesignerGenerationError).
+    await scheduler.run()
+
+    assert scheduler.early_shutdown
+    assert buffer_mgr.actual_num_records == 0
+    # All rows dropped → checkpoint path frees buffer without writing; on_finalize
+    # is *not* called because every row was dropped before survivors could exist.
+    assert finalized == []
+    # No partial-row-groups recorded — there were no incomplete-but-not-dropped rows.
+    assert scheduler.partial_row_groups == ()
+    storage.write_batch_to_parquet_file.assert_not_called()
+
+
+@pytest.mark.asyncio(loop_scope="session")
+async def test_healthy_run_has_no_partial_signal() -> None:
+    """Successful run leaves early_shutdown=False and partial_row_groups empty."""
+    scheduler, _tracker = _build_simple_pipeline(num_records=3)
+    await scheduler.run()
+    assert not scheduler.early_shutdown
+    assert scheduler.partial_row_groups == ()
 
 
 @pytest.mark.asyncio(loop_scope="session")