NVIDIA-NeMo · eric-tramel · Jul 9, 2026 · Jul 9, 2026 · Jul 9, 2026 · Jul 9, 2026
@@ -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,15 +76,15 @@ 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.
 
 Resume relies on stable row-group boundaries within a run. It treats datasets that have completed `process_after_generation()` as terminal: after-generation processors operate on the whole dataset and can re-chunk rows or change schema, invalidating row-group identity for later resume/extension. The terminal-state check raises a clear `DatasetGenerationError` (not a `TypeError`) when the persisted metadata is missing required fields such as `target_num_records`.
 
 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.
 

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

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

@@ -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,15 +136,36 @@ 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,
         row_group_size: int,
         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):

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

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

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

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

@@ -256,28 +256,35 @@ 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,
             resource_limits={"llm_wait": 8},
             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:

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