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/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/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/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/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/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_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/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) 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 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.""" 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,