From 0107785dc0494d2effe014bed4e28e948421bfbf Mon Sep 17 00:00:00 2001 From: "Eric W. Tramel" Date: Thu, 9 Jul 2026 11:18:09 -0400 Subject: [PATCH 1/4] perf(scheduler): keep borrowing work-conserving Signed-off-by: Eric W. Tramel --- .../scheduling/task_policies.py | 10 +++---- .../scheduling/test_task_admission.py | 27 ++++++++++++------- .../scheduling/test_task_policies.py | 2 +- 3 files changed, 23 insertions(+), 16 deletions(-) diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/task_policies.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/task_policies.py index 227cf8658..9ae7d7888 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/task_policies.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/task_policies.py @@ -29,7 +29,8 @@ "shutdown", "policy_denial", ] -DEFAULT_DYNAMIC_BORROW_RESERVE_FRACTION = 0.125 +# ponytail: a future peer can take the next released slot; do not reserve idle capacity. +DEFAULT_DYNAMIC_BORROW_RESERVE_FRACTION = 0.0 DEFAULT_DYNAMIC_BORROW_MAX_RESERVED_SLOTS = 8 @@ -40,9 +41,8 @@ class BoundedBorrowTaskAdmissionPolicyConfig: Borrow debt is tracked by task group and scheduler resource. Any completed lease in the same group repays debt for the released resources; repayment is not tied to the specific lease that originally borrowed. When no explicit - borrow ceiling is configured, the policy reserves one slot per eight - resource slots, capped at eight reserved slots, and lets solo groups borrow - up to the remaining capacity. + borrow ceiling is configured, a solo group may use the full resource; + borrow debt gives a newly queued peer priority as leases complete. """ borrow_ceiling_by_group_resource: Mapping[tuple[TaskGroupKey, SchedulerResourceKey], int] = field( @@ -320,7 +320,7 @@ def _strict_share( def _dynamic_reserved_slots(resource_limit: int, *, reserve_fraction: float, max_reserved_slots: int) -> int: - return min(max_reserved_slots, max(1, math.ceil(resource_limit * reserve_fraction))) + return min(max_reserved_slots, math.ceil(resource_limit * reserve_fraction)) def _competing_group_specs( diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/scheduling/test_task_admission.py b/packages/data-designer-engine/tests/engine/dataset_builders/scheduling/test_task_admission.py index 2d12be946..c7a52df57 100644 --- a/packages/data-designer-engine/tests/engine/dataset_builders/scheduling/test_task_admission.py +++ b/packages/data-designer-engine/tests/engine/dataset_builders/scheduling/test_task_admission.py @@ -256,8 +256,9 @@ def test_bounded_borrow_prevents_solo_heavy_group_from_consuming_all_typed_capac assert controller.view().policy_debt_by_group_resource[(group.key, "llm_wait")] == 1 -def test_bounded_borrow_dynamic_ceiling_reserves_capacity() -> None: +def test_bounded_borrow_dynamic_ceiling_uses_full_solo_capacity_then_yields() -> None: group = TaskGroupSpec(TaskGroupKey(kind="model", identity=("provider", "hot")), weight=4.0, admitted_limit=8) + peer_group = TaskGroupSpec(TaskGroupKey(kind="model", identity=("provider", "peer")), admitted_limit=1) controller = TaskAdmissionController( TaskAdmissionConfig( submission_capacity=8, @@ -265,19 +266,25 @@ def test_bounded_borrow_dynamic_ceiling_reserves_capacity() -> None: bounded_borrow=BoundedBorrowTaskAdmissionPolicyConfig(), ) ) - items = [_item("hot", row, group=group, resources={"submission": 1, "llm_wait": 1}) for row in range(8)] - for index in range(7): + items = [_item("hot", row, group=group, resources={"submission": 1, "llm_wait": 1}) for row in range(9)] + leases = [] + for index in range(8): decision = controller.try_acquire(items[index], _queue_view(*items[index:])) assert isinstance(decision, TaskAdmissionLease) + leases.append(decision) - denied = controller.try_acquire(items[7], _queue_view(items[7])) + assert controller.view().resources_available["llm_wait"] == 0 + assert controller.view().policy_debt_by_group_resource[(group.key, "llm_wait")] == 4 - assert isinstance(denied, TaskAdmissionDenied) - assert denied.reason == "borrow_debt" - assert denied.diagnostics["ceiling"] == 3 - assert denied.diagnostics["strict_share"] == 4 - assert controller.view().resources_available["llm_wait"] == 1 - assert controller.view().policy_debt_by_group_resource[(group.key, "llm_wait")] == 3 + controller.release(leases[0]) + peer = _item("peer", 0, group=peer_group, resources={"submission": 1, "llm_wait": 1}) + queue = FairTaskQueue() + queue.enqueue((items[8], peer)) + + selection = queue.select_next(controller.is_eligible) + + assert selection is not None + assert selection.item.task_id == peer.task_id def test_bounded_borrow_explicit_ceiling_counts_marginal_borrowed_slots() -> None: diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/scheduling/test_task_policies.py b/packages/data-designer-engine/tests/engine/dataset_builders/scheduling/test_task_policies.py index 15c8deb28..18b7cbf16 100644 --- a/packages/data-designer-engine/tests/engine/dataset_builders/scheduling/test_task_policies.py +++ b/packages/data-designer-engine/tests/engine/dataset_builders/scheduling/test_task_policies.py @@ -107,7 +107,7 @@ def test_bounded_borrow_policy_defaults_to_ceil_strict_share_rounding() -> None: assert config.strict_share_rounding == "ceil" assert config.default_borrow_ceiling is None - assert config.dynamic_borrow_reserve_fraction == 0.125 + assert config.dynamic_borrow_reserve_fraction == 0.0 assert config.dynamic_borrow_max_reserved_slots == 8 From 40951840216abc8f7751b41827704965959990ff Mon Sep 17 00:00:00 2001 From: "Eric W. Tramel" Date: Thu, 9 Jul 2026 11:18:18 -0400 Subject: [PATCH 2/4] perf(scheduler): compact checkpointed state Signed-off-by: Eric W. Tramel --- .../dataset_builders/async_scheduler.py | 3 +- .../dataset_builders/scheduling/completion.py | 23 ++++++++ .../scheduling/test_completion.py | 17 ++++++ .../dataset_builders/test_async_scheduler.py | 55 ++++++++++++++++++- 4 files changed, 96 insertions(+), 2 deletions(-) diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py index 8b6be7889..523932858 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py @@ -1593,7 +1593,7 @@ def _checkpoint_completed_row_groups(self, all_columns: list[str]) -> None: completed = [ (rg_id, state.size) for rg_id, state in self._rg_states.items() - if self._tracker.is_row_group_complete(rg_id, state.size, all_columns) + if state.in_flight_count == 0 and self._tracker.is_row_group_complete(rg_id, state.size, all_columns) ] for rg_id, rg_size in completed: dropped_rows = sum(1 for ri in range(rg_size) if self._tracker.is_dropped(rg_id, ri)) @@ -1630,6 +1630,7 @@ def _checkpoint_completed_row_groups(self, all_columns: list[str]) -> None: self._rg_semaphore.release() self._row_group_admission_event.set() if checkpointed: + self._tracker.compact_row_group(rg_id) self._emit_scheduler_event( "row_group_checkpointed", diagnostics={ diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/completion.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/completion.py index 855c91642..b74b1f9af 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/completion.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/completion.py @@ -52,6 +52,8 @@ def __init__(self) -> None: self._row_group_plan: RowGroupPlanLike | None = None self._batch_complete: dict[int, set[str]] = defaultdict(set) self._frontier: set[Task] = set() + # Checkpointed row group → compact bitset of dropped row indices. + self._compacted_dropped: dict[int, bytes] = {} @classmethod def with_graph(cls, graph: ExecutionGraph, row_groups: RowGroupInput) -> CompletionTracker: @@ -134,8 +136,27 @@ def drop_row(self, row_group: int, row_index: int) -> FrontierDelta: return self._record_delta(added=added, removed=removed) def is_dropped(self, row_group: int, row_index: int) -> bool: + dropped_mask = self._compacted_dropped.get(row_group) + if dropped_mask is not None: + byte_index, bit_index = divmod(row_index, 8) + return 0 <= byte_index < len(dropped_mask) and bool(dropped_mask[byte_index] & (1 << bit_index)) return row_index in self._dropped.get(row_group, set()) + def compact_row_group(self, row_group: int) -> None: + """Release detailed state while retaining terminal group and drop queries.""" + if row_group in self._compacted_dropped: + return + self._validate_row_group(row_group) + dropped = self._dropped.pop(row_group, set()) + dropped_mask = bytearray((max(dropped, default=-1) // 8) + 1) + for row_index in dropped: + byte_index, bit_index = divmod(row_index, 8) + dropped_mask[byte_index] |= 1 << bit_index + self._compacted_dropped[row_group] = bytes(dropped_mask) + self._completed.pop(row_group, None) + self._batch_complete.pop(row_group, None) + self._frontier = {task for task in self._frontier if task.row_group != row_group} + def is_row_group_complete( self, row_group: int, @@ -143,6 +164,8 @@ def is_row_group_complete( all_columns: list[str], ) -> bool: """All non-dropped rows have all columns done.""" + if row_group in self._compacted_dropped: + return True dropped = self._dropped.get(row_group, set()) completed = self._completed.get(row_group, {}) for ri in range(row_group_size): diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/scheduling/test_completion.py b/packages/data-designer-engine/tests/engine/dataset_builders/scheduling/test_completion.py index e647d4ac6..98b070442 100644 --- a/packages/data-designer-engine/tests/engine/dataset_builders/scheduling/test_completion.py +++ b/packages/data-designer-engine/tests/engine/dataset_builders/scheduling/test_completion.py @@ -174,6 +174,23 @@ def test_row_group_not_complete_missing_non_dropped() -> None: assert not tracker.is_row_group_complete(0, 3, ["col_a", "col_b"]) +def test_compact_row_group_releases_details_and_preserves_terminal_queries() -> None: + tracker = CompletionTracker() + tracker.mark_row_range_complete("col_a", 0, 4) + tracker.mark_row_range_complete("col_b", 0, 4) + tracker.drop_row(0, 1) + tracker.drop_row(0, 3) + + tracker.compact_row_group(0) + tracker.compact_row_group(0) + + assert tracker.is_row_group_complete(0, 4, ["col_a", "col_b"]) + assert [tracker.is_dropped(0, row_index) for row_index in range(4)] == [False, True, False, True] + assert 0 not in tracker._completed + assert 0 not in tracker._dropped + assert 0 not in tracker._batch_complete + + # -- get_ready_tasks -------------------------------------------------------- diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py b/packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py index 8de0e4cee..c4a00982e 100644 --- a/packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py +++ b/packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py @@ -1789,7 +1789,7 @@ async def test_scheduler_on_before_checkpoint_callback() -> None: @pytest.mark.asyncio(loop_scope="session") -async def test_scheduler_on_finalize_row_group_callback_fires() -> None: +async def test_scheduler_on_finalize_row_group_callback_fires(monkeypatch: pytest.MonkeyPatch) -> None: """on_finalize_row_group is called for each completed row group.""" provider = _mock_provider() configs = [ @@ -1810,6 +1810,8 @@ async def test_scheduler_on_finalize_row_group_callback_fires() -> None: buffer_mgr = RowGroupBufferManager(storage) finalized: list[int] = [] + compact_row_group = MagicMock(wraps=tracker.compact_row_group) + monkeypatch.setattr(tracker, "compact_row_group", compact_row_group) def finalize_row_group(rg_id: int) -> None: buffer_mgr.checkpoint_row_group(rg_id) @@ -1827,6 +1829,22 @@ def finalize_row_group(rg_id: int) -> None: assert finalized == [0] assert storage.write_batch_to_parquet_file.call_count == 1 + compact_row_group.assert_called_once_with(0) + + +@pytest.mark.asyncio(loop_scope="session") +async def test_scheduler_does_not_compact_row_group_when_checkpoint_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + scheduler, tracker = _build_simple_pipeline(num_records=1) + compact_row_group = MagicMock(wraps=tracker.compact_row_group) + monkeypatch.setattr(tracker, "compact_row_group", compact_row_group) + scheduler._on_finalize_row_group = MagicMock(side_effect=RuntimeError("checkpoint failed")) + + await scheduler.run() + + compact_row_group.assert_not_called() + assert tracker.is_row_group_complete(0, 1, ["seed", "cell_out"]) @pytest.mark.asyncio(loop_scope="session") @@ -3173,6 +3191,41 @@ async def agenerate(self, data: dict) -> dict: return self.generate(data) +@pytest.mark.asyncio(loop_scope="session") +async def test_scheduler_compacts_only_after_row_group_workers_finish() -> None: + """A dropped row must not checkpoint while a sibling worker is still running.""" + provider = _mock_provider() + configs = [ + SamplerColumnConfig(name="seed", sampler_type=SamplerType.CATEGORY, params={"values": ["A"]}), + LLMTextColumnConfig(name="failing", prompt="{{ seed }}", model_alias=MODEL_ALIAS), + LLMTextColumnConfig(name="slow", prompt="{{ seed }}", model_alias=MODEL_ALIAS), + ] + strategies = { + "seed": GenerationStrategy.FULL_COLUMN, + "failing": GenerationStrategy.CELL_BY_CELL, + "slow": GenerationStrategy.CELL_BY_CELL, + } + generators = { + "seed": MockSeedGenerator(config=_expr_config("seed"), resource_provider=provider), + "failing": MockFailingGenerator(config=_expr_config("failing"), resource_provider=provider), + "slow": SlowCellGenerator(config=_expr_config("slow"), resource_provider=provider), + } + scheduler, tracker = _build_simple_pipeline( + num_records=1, + configs=configs, + strategies=strategies, + generators=generators, + ) + + await scheduler.run() + + assert tracker.is_row_group_complete(0, 1, ["seed", "failing", "slow"]) + assert tracker.is_dropped(0, 0) + assert 0 not in tracker._completed + assert 0 not in tracker._dropped + assert 0 not in tracker._batch_complete + + class SlowLLMBoundCellGenerator(SlowCellGenerator): """Slow cell generator that participates in model-stage scheduling.""" From 8085e5897dc7c9781015760b8e189709f0754258 Mon Sep 17 00:00:00 2001 From: "Eric W. Tramel" Date: Thu, 9 Jul 2026 11:38:01 -0400 Subject: [PATCH 3/4] perf(scheduler): bound fair queue state Signed-off-by: Eric W. Tramel --- .../dataset_builders/scheduling/queue.py | 28 ++++++------------- .../dataset_builders/scheduling/test_queue.py | 16 +++++++++++ 2 files changed, 24 insertions(+), 20 deletions(-) diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/queue.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/queue.py index 4166df4b8..2ae43cae7 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/queue.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/queue.py @@ -3,7 +3,6 @@ from __future__ import annotations -import heapq from collections import Counter, defaultdict, deque from collections.abc import Callable, Iterable, Mapping from dataclasses import dataclass @@ -50,9 +49,7 @@ def __init__(self) -> None: self._queued_resource_demand_by_group: dict[TaskGroupKey, Counter[SchedulerResourceKey]] = defaultdict(Counter) self._queued_peer_demand_by_resource: Counter[SchedulerResourceKey] = Counter() self._group_finish: dict[TaskGroupKey, float] = {} - self._heap: list[tuple[float, int, TaskGroupKey]] = [] - self._active_heap_keys: set[TaskGroupKey] = set() - self._active_heap_entries: dict[TaskGroupKey, tuple[float, int]] = {} + self._active_entries: dict[TaskGroupKey, tuple[float, int]] = {} self._sequence = 0 self._sequence_version = 0 self._virtual_time = 0.0 @@ -93,16 +90,8 @@ def discard_where(self, predicate: Callable[[SchedulableTask], bool]) -> None: def select_next(self, is_eligible: Callable[[SchedulableTask, QueueView], bool]) -> QueueSelection | None: """Return the next eligible task without mutating queue state.""" view = self.view() - heap_copy = list(self._heap) - heapq.heapify(heap_copy) - active_seen: set[TaskGroupKey] = set() - while heap_copy: - finish, sequence, key = heapq.heappop(heap_copy) - if key in active_seen: - continue - if self._active_heap_entries.get(key) != (finish, sequence): - continue - active_seen.add(key) + active_entries = sorted((finish, sequence, key) for key, (finish, sequence) in self._active_entries.items()) + for _finish, _sequence, key in active_entries: item = self._first_valid_item(key) if item is None: continue @@ -128,8 +117,7 @@ def commit(self, selection: QueueSelection) -> SchedulableTask | None: queue.popleft() self._remove_queued_item(item.task_id) - self._active_heap_keys.discard(key) - self._active_heap_entries.pop(key, None) + self._active_entries.pop(key, None) group = self._group_specs[key] finish = self._group_finish.get(key, self._virtual_time) self._virtual_time = max(self._virtual_time, finish) @@ -168,13 +156,11 @@ def view(self) -> QueueView: def _activate_group(self, key: TaskGroupKey) -> None: self._purge_queue_head(key) queue = self._queues.get(key) - if not queue or key in self._active_heap_keys: + if not queue or key in self._active_entries: return self._sequence += 1 finish = self._group_finish.get(key, self._virtual_time) - heapq.heappush(self._heap, (finish, self._sequence, key)) - self._active_heap_keys.add(key) - self._active_heap_entries[key] = (finish, self._sequence) + self._active_entries[key] = (finish, self._sequence) def _first_valid_item(self, key: TaskGroupKey) -> SchedulableTask | None: self._purge_queue_head(key) @@ -206,6 +192,8 @@ def _remove_queued_item(self, task_id: str) -> SchedulableTask | None: if item is None or key is None: return item self._decrement_queue_accounting(item, key) + if key not in self._queued_by_group: + self._active_entries.pop(key, None) return item def _decrement_queue_accounting(self, item: SchedulableTask, key: TaskGroupKey) -> None: diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/scheduling/test_queue.py b/packages/data-designer-engine/tests/engine/dataset_builders/scheduling/test_queue.py index 5e10fe5bd..0c157cf73 100644 --- a/packages/data-designer-engine/tests/engine/dataset_builders/scheduling/test_queue.py +++ b/packages/data-designer-engine/tests/engine/dataset_builders/scheduling/test_queue.py @@ -88,6 +88,22 @@ def test_fair_task_queue_weighted_groups() -> None: assert counts == {"a": 4, "b": 2} +def test_fair_task_queue_ordering_state_stays_bounded_by_active_groups() -> None: + queue = FairTaskQueue() + group = _group("a") + queue.enqueue(_item("a", row_index, group) for row_index in range(1_000)) + + for remaining in reversed(range(1_000)): + assert _select_and_commit(queue) is not None + assert len(queue._active_entries) == min(remaining, 1) + + discarded = _item("b", 0) + queue.enqueue([discarded]) + queue.discard(discarded.task_id) + + assert not queue._active_entries + + def test_select_next_is_non_mutating_until_commit() -> None: queue = FairTaskQueue() first = _item("a", 0) From e5cd032d0b898efc8c0b9e9ea3443bd7bb01c086 Mon Sep 17 00:00:00 2001 From: "Eric W. Tramel" Date: Thu, 9 Jul 2026 19:39:06 -0400 Subject: [PATCH 4/4] perf(scheduler): reduce metadata checkpointing Persist resume metadata after the first durable row group and refresh it at completion. Resume already recovers progress from parquet files, avoiding repeated metadata scans and writes without adding a tuning knob. Add focused cadence coverage and document the resume contract. --- architecture/dataset-builders.md | 6 ++-- .../dataset_builders/dataset_builder.py | 20 +++++++----- .../dataset_builders/test_dataset_builder.py | 31 +++++++++++++++++++ 3 files changed, 47 insertions(+), 10 deletions(-) diff --git a/architecture/dataset-builders.md b/architecture/dataset-builders.md index 0b69e5675..15f0aac0e 100644 --- a/architecture/dataset-builders.md +++ b/architecture/dataset-builders.md @@ -65,7 +65,7 @@ DAG edges are added for `skip.when` column references in both `topologically_sor Manages per-row-group DataFrames and persistence: - `checkpoint_row_group` → writes parquet via `ArtifactStorage` -- Updates dataset metadata between row groups +- Writes dataset metadata after the first durable row group and at completion - Tracks dropped rows and actual record counts for resume ### Resume Checkpointing @@ -76,7 +76,7 @@ Manages per-row-group DataFrames and persistence: - `ResumeMode.ALWAYS` resumes the existing dataset directory and raises on incompatible state. - `ResumeMode.IF_POSSIBLE` resumes when the persisted config fingerprint matches; otherwise it starts a fresh timestamped run. -Checkpoint state lives in `metadata.json`. Each metadata write includes the config fingerprint (`config_hash`, `config_hash_algo`, and `config_hash_version`) so compatibility checks do not need to deserialize `builder_config.json` for the common path. `builder_config.json` remains the human-readable record of the run configuration and the fallback for older datasets. +Resume configuration lives in `metadata.json`. Each metadata write includes the config fingerprint (`config_hash`, `config_hash_algo`, and `config_hash_version`) so compatibility checks do not need to deserialize `builder_config.json` for the common path. `builder_config.json` remains the human-readable record of the run configuration and the fallback for older datasets. Resume scans `parquet-files/batch_*.parquet` and reads parquet metadata to recover the completed row-group IDs and their actual persisted row counts. `metadata.json` remains the source of truth for the run *configuration* (`buffer_size`, `target_num_records`, `original_target_num_records`, config fingerprint), but the filesystem is the source of truth for *progress* (`num_completed_batches`, `actual_num_records`). Splitting the two sources is what lets resume survive a crash between writing a row-group parquet and updating metadata - the filesystem reflects the durable state even when metadata lags by a step. Reading actual row counts also matters for early-shutdown salvage, where a completed parquet file can contain fewer rows than the requested row-group size. Resume tolerates non-contiguous IDs because row groups can complete out of order. @@ -84,7 +84,7 @@ Resume relies on stable row-group boundaries within a run. It treats datasets th After-generation processors run unconditionally on the on-disk dataset whenever they are configured — including the case where resume sees every row group already on disk. This closes the crash window between the final row-group parquet write and the `post_generation_state="started"` marker write: in that window, the dataset is complete but post-generation never ran, and the on-disk parquet files are still clean (no processor has touched them). The `post_generation_state="started"` short-circuit still rejects the other direction (`process_after_generation()` crashed mid-rewrite, leaving the parquet files in an ambiguous state), so resume only re-runs after-generation when it is safe to do so. -Metadata writes are atomic (`tmp` file + `fsync` + `os.replace`) because `metadata.json` is the crash-recovery checkpoint. Corrupt or partially written metadata raises a clear `DatasetGenerationError` rather than falling through as a generic config mismatch. +Metadata writes are atomic (`tmp` file + `fsync` + `os.replace`) because `metadata.json` is the resume-configuration checkpoint. Corrupt or partially written metadata raises a clear `DatasetGenerationError` rather than falling through as a generic config mismatch. `DatasetCreationResults` from a resume invocation reflects the full on-disk dataset for anything that reads the artifact directory (`load_dataset`, `count_records`, `load_analysis`, `export`, `push_to_hub`), but per-run observability (`task_traces`, model-usage logs, telemetry events) is scoped to the current invocation — the original run's in-memory state is not persisted across process boundaries. 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 13d23e383..d0b1a780f 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 @@ -801,18 +801,24 @@ def _build_async( precomputed_row_groups = resume_plan.remaining_row_groups + incremental_metadata_written = False + def finalize_row_group(rg_id: int) -> None: + nonlocal incremental_metadata_written + def on_complete(final_path: Path | str | None) -> None: if final_path is not None and on_batch_complete: on_batch_complete(final_path) buffer_manager.checkpoint_row_group(rg_id, on_complete=on_complete) - # Write incremental metadata after each row group so interrupted runs can be resumed. - buffer_manager.write_metadata( - target_num_records=num_records, - original_target_num_records=original_target, - buffer_size=buffer_size, - ) + # Persist the first durable checkpoint; resume scans parquet files for later progress. + if not incremental_metadata_written: + buffer_manager.write_metadata( + target_num_records=num_records, + original_target_num_records=original_target, + buffer_size=buffer_size, + ) + incremental_metadata_written = True # Telemetry snapshot group_id = uuid.uuid4().hex @@ -861,7 +867,7 @@ def on_complete(final_path: Path | str | None) -> None: except Exception: logger.debug("Failed to emit batch telemetry for async run", exc_info=True) - # Write final metadata (overwrites the last incremental write with identical content). + # Refresh metadata with the final checkpoint state. buffer_manager.write_metadata( target_num_records=num_records, original_target_num_records=original_target, diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/test_dataset_builder.py b/packages/data-designer-engine/tests/engine/dataset_builders/test_dataset_builder.py index baca1ef2c..042a83283 100644 --- a/packages/data-designer-engine/tests/engine/dataset_builders/test_dataset_builder.py +++ b/packages/data-designer-engine/tests/engine/dataset_builders/test_dataset_builder.py @@ -1135,6 +1135,37 @@ def test_preview_does_not_write_scheduler_events(stub_resource_provider: Mock, t assert list(tmp_path.rglob("scheduler_events.jsonl")) == [] +def test_build_writes_metadata_after_first_checkpoint_and_at_end( + stub_resource_provider: Mock, + tmp_path: Path, +) -> None: + builder, storage = _make_sampler_only_builder( + stub_resource_provider, + tmp_path, + resume=ResumeMode.NEVER, + ) + payloads: list[dict[str, object]] = [] + original_write_metadata = ArtifactStorage.write_metadata + + def record_write_metadata(instance: ArtifactStorage, metadata: dict[str, object]) -> Path: + if instance is storage: + payloads.append(metadata) + return original_write_metadata(instance, metadata) + + with patch.object(ArtifactStorage, "write_metadata", autospec=True, side_effect=record_write_metadata): + builder.build(num_records=6, resume=ResumeMode.NEVER) + + assert [(payload["actual_num_records"], payload["num_completed_batches"]) for payload in payloads] == [ + (2, 1), + (6, 3), + ] + file_paths = [payload["file_paths"] for payload in payloads] + assert all(isinstance(paths, dict) for paths in file_paths) + assert [len(paths["parquet-files"]) for paths in file_paths if isinstance(paths, dict)] == [1, 3] + persisted = storage.read_metadata() + assert {key: persisted[key] for key in payloads[-1]} == payloads[-1] + + def test_resumed_build_appends_scheduler_event_segment(stub_resource_provider: Mock, tmp_path: Path) -> None: builder, _storage = _make_sampler_only_builder( stub_resource_provider,