data-designer-engine 0.8.0__tar.gz → 0.9.0__tar.gz

{data_designer_engine-0.8.0 → data_designer_engine-0.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: data-designer-engine
-Version: 0.8.0
+Version: 0.9.0
 Summary: Generation engine for DataDesigner synthetic data generation
 License-Expression: Apache-2.0
 Classifier: Development Status :: 4 - Beta
@@ -18,7 +18,7 @@ Requires-Dist: anyascii<1,>=0.3.3
 Requires-Dist: asciichartpy<2,>=1.5.25
 Requires-Dist: chardet<6,>=3.0.2
 Requires-Dist: cryptography<=49,>=48.0.1
-Requires-Dist: data-designer-config==0.8.0
+Requires-Dist: data-designer-config==0.9.0
 Requires-Dist: duckdb<2,>=1.5.0
 Requires-Dist: faker<21,>=20.1.0
 Requires-Dist: fsspec<2026,>=2025.3.0
@@ -31,7 +31,7 @@ Requires-Dist: jsonpath-rust-bindings<2,>=1.0
 Requires-Dist: jsonschema<5,>=4.0.0
 Requires-Dist: lxml<7,>=6.1.0
 Requires-Dist: marko<3,>=2.1.2
-Requires-Dist: mcp<2,>=1.26.0
+Requires-Dist: mcp<2,>=1.29.0
 Requires-Dist: networkx<4,>=3.0
 Requires-Dist: numpy<3,>=1.23.5
 Requires-Dist: pandas<3,>=2.3.3

{data_designer_engine-0.8.0 → data_designer_engine-0.9.0}/pyproject.toml

@@ -50,7 +50,7 @@ dependencies = [
     "jsonschema>=4.0.0,<5",
     "lxml>=6.1.0,<7",
     "marko>=2.1.2,<3",
-    "mcp>=1.26.0,<2",
+    "mcp>=1.29.0,<2", # 1.29.0 fixes MCP security advisories
     "networkx>=3.0,<4",
     "numpy>=1.23.5,<3",
     "pandas>=2.3.3,<3",

{data_designer_engine-0.8.0 → data_designer_engine-0.9.0}/src/data_designer/engine/dataset_builders/async_scheduler.py

@@ -68,8 +68,11 @@ from data_designer.engine.models.clients.errors import ProviderError
 from data_designer.engine.models.errors import (
     RETRYABLE_MODEL_ERRORS,
     GenerationValidationFailureError,
+    ModelAPIConnectionError,
+    ModelInternalServerError,
     ModelRateLimitError,
     ModelRequestAdmissionTimeoutError,
+    ModelTimeoutError,
 )
 from data_designer.engine.models.request_admission.config import RequestAdmissionConfig
 from data_designer.engine.models.request_admission.resources import RequestResourceKey
@@ -163,6 +166,13 @@ class _DeferredAdmissionAnalysis:
     independent_candidate_columns: tuple[str, ...]
 
 
+@dataclass
+class _GeneratorCallState:
+    """Whether task execution reached user or provider generation code."""
+
+    invoked: bool = False
+
+
 class AsyncTaskScheduler:
     """Dependency-aware async task scheduler for the dataset builder.
 
@@ -308,6 +318,10 @@ class AsyncTaskScheduler:
         self._degraded_warn_window = degraded_warn_window
         self._degraded_warn_interval_s = degraded_warn_interval_s
         self._recent_retryable: deque[bool] = deque(maxlen=degraded_warn_window)
+        self._recent_retryable_kinds: deque[str] = deque(maxlen=degraded_warn_window)
+        self._retryable_outcome_counts: Counter[str] = Counter()
+        self._retryable_detail_counts: Counter[str] = Counter()
+        self._logged_retryable_kinds: set[str] = set()
         # Initialize to -inf so the first WARN is always emitted regardless of
         # the monotonic clock's absolute value (which can be near-zero on freshly
         # booted CI runners).
@@ -432,6 +446,18 @@ class AsyncTaskScheduler:
         """
         return self._first_non_retryable_error
 
+    @property
+    def retryable_outcome_metrics(self) -> dict[str, object]:
+        """Return sanitized rolling and cumulative model-task outcome counts."""
+        rolling = Counter(self._recent_retryable_kinds)
+        return {
+            "window_size": self._degraded_warn_window,
+            "rolling_total": len(self._recent_retryable_kinds),
+            "rolling_counts": dict(sorted(rolling.items())),
+            "cumulative_counts": dict(sorted(self._retryable_outcome_counts.items())),
+            "retryable_details": dict(sorted(self._retryable_detail_counts.items())),
+        }
+
     def request_cancel(self) -> None:
         """Signal cancellation to scheduler tasks and bridged sync generator work."""
         self._cancel_requested.set()
@@ -580,6 +606,7 @@ class AsyncTaskScheduler:
             "request_pressure_advisory_skips": self._request_pressure_advisory_skips,
             "row_group_admission_blocked_reasons": dict(self._row_group_admission_blocked_reasons),
             "request_pressure": self._request_pressure_diagnostics(),
+            "retryable_outcomes": self.retryable_outcome_metrics,
         }
 
     def _scheduler_job_diagnostics(self) -> dict[str, object]:
@@ -1593,7 +1620,7 @@ class AsyncTaskScheduler:
         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 +1657,7 @@ class AsyncTaskScheduler:
                 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={
@@ -1778,7 +1806,44 @@ class AsyncTaskScheduler:
         if errors / self._shutdown_error_window >= self._shutdown_error_rate:
             self._early_shutdown = True
 
-    def _record_retryable_outcome(self, *, retryable: bool) -> None:
+    @staticmethod
+    def _retryable_error_kind(exc: Exception | None) -> str:
+        if isinstance(exc, ModelRequestAdmissionTimeoutError):
+            return "request_admission_timeout"
+        if isinstance(exc, ModelRateLimitError):
+            return "rate_limit"
+        if isinstance(exc, ModelTimeoutError):
+            return "timeout"
+        if isinstance(exc, ModelInternalServerError):
+            return "internal_server"
+        if isinstance(exc, ModelAPIConnectionError):
+            return "connection"
+        return "other_retryable"
+
+    @staticmethod
+    def _provider_error_details(exc: Exception | None) -> tuple[int | None, float | None]:
+        """Extract only safe transport metadata from a wrapped provider error.
+
+        Model-error normalization raises ``from None`` to hide provider details,
+        but Python retains the originating error in ``__context__``. Walking both
+        chains preserves access to bounded transport metadata without exposing
+        provider messages.
+        """
+        seen: set[int] = set()
+        current: BaseException | None = exc
+        while current is not None and id(current) not in seen:
+            seen.add(id(current))
+            if isinstance(current, ProviderError):
+                return current.status_code, current.retry_after
+            current = current.__cause__ or current.__context__
+        return None, None
+
+    def _record_retryable_outcome(
+        self,
+        *,
+        retryable: bool,
+        exc: Exception | None = None,
+    ) -> None:
         """Track retryable-error rate and emit a rate-limited WARN under provider degradation.
 
         Distinct from ``_check_error_rate``: every LLM-bound task outcome (success
@@ -1786,11 +1851,38 @@ class AsyncTaskScheduler:
         health, not just the error mix. The call site filters on ``is_llm`` so
         non-LLM tasks (samplers, expressions, non-LLM customs) don't dilute the
         rate. Only retryable errors (rate-limit, timeout, 5xx, connection) count
-        toward the rate; non-retryable failures register as 0.
+        toward the rate; non-retryable failures register as 0 and are categorized
+        separately from successful outcomes.
         """
+        if retryable:
+            kind = self._retryable_error_kind(exc)
+        elif exc is not None:
+            kind = "non_retryable_failure"
+        else:
+            kind = "success"
+        self._retryable_outcome_counts[kind] += 1
+        if retryable:
+            status_code, retry_after = self._provider_error_details(exc)
+            detail = kind
+            if status_code is not None:
+                detail = f"{detail}:http_{status_code}"
+            self._retryable_detail_counts[detail] += 1
+            if kind not in self._logged_retryable_kinds:
+                # Ongoing pressure is reported by the rate-limited degradation
+                # warning; this first-observation warning is once per scheduler.
+                self._logged_retryable_kinds.add(kind)
+                retry_after_text = f", retry_after={retry_after:g}s" if retry_after is not None else ""
+                status_text = f", http_status={status_code}" if status_code is not None else ""
+                logger.warning(
+                    "Observed retryable model-task error: kind=%s%s%s; the row task will be deferred.",
+                    kind,
+                    status_text,
+                    retry_after_text,
+                )
         if self._degraded_warn_window <= 0:
             return
         self._recent_retryable.append(retryable)
+        self._recent_retryable_kinds.append(kind)
         if len(self._recent_retryable) < self._degraded_warn_window:
             return
         rate = sum(self._recent_retryable) / self._degraded_warn_window
@@ -1801,10 +1893,20 @@ class AsyncTaskScheduler:
             return
         self._last_degraded_warn_at = now
         pct = int(round(rate * 100))
+        rolling = Counter(self._recent_retryable_kinds)
+        rolling_summary = ", ".join(f"{key}={rolling[key]}" for key in sorted(rolling))
+        cumulative_summary = ", ".join(
+            f"{key}={self._retryable_outcome_counts[key]}" for key in sorted(self._retryable_outcome_counts)
+        )
         logger.warning(
-            f"Provider showing degraded performance: {pct}% of last {self._degraded_warn_window} "
-            "task outcomes were retryable errors (rate-limit, timeout, 5xx, connection). "
-            "Run may take longer than expected; salvage will retry these."
+            "Provider showing degraded performance: %s%% of last %s task outcomes were "
+            "retryable errors (%s); cumulative outcomes: %s; deferred_tasks=%s. "
+            "Run may take longer than expected; salvage will retry these.",
+            pct,
+            self._degraded_warn_window,
+            rolling_summary,
+            cumulative_summary,
+            len(self._deferred),
         )
 
     async def _dispatch_seeds(self, rg_id: int, rg_size: int) -> None:
@@ -1874,6 +1976,7 @@ class AsyncTaskScheduler:
         # from the frontier (it was never completed, so it stays in the frontier).
         skipped = False
         uses_model_stage_resource = "llm_wait" in lease.resources
+        generator_call_state = _GeneratorCallState()
         stateful_lock_acquired = False
 
         try:
@@ -1893,11 +1996,11 @@ class AsyncTaskScheduler:
 
             cell_skipped = False
             if task.task_type == "from_scratch":
-                await self._run_from_scratch(task, generator)
+                await self._run_from_scratch(task, generator, call_state=generator_call_state)
             elif task.task_type == "cell":
-                _result, cell_skipped = await self._run_cell(task, generator)
+                _result, cell_skipped = await self._run_cell(task, generator, call_state=generator_call_state)
             elif task.task_type == "batch":
-                await self._run_batch(task, generator)
+                await self._run_batch(task, generator, call_state=generator_call_state)
             else:
                 raise ValueError(f"Unknown task type: {task.task_type}")
 
@@ -1915,7 +2018,7 @@ class AsyncTaskScheduler:
             # window from LLM-bound tasks so a healthy non-model task mix
             # (samplers, expressions, non-LLM customs) doesn't dilute the
             # rate and silence the WARN under genuine provider stress.
-            if uses_model_stage_resource:
+            if uses_model_stage_resource and generator_call_state.invoked:
                 self._record_retryable_outcome(retryable=False)
             if self._reporter:
                 if cell_skipped:
@@ -1934,14 +2037,17 @@ class AsyncTaskScheduler:
 
         except Exception as exc:
             retryable = self._is_retryable(exc)
+            if retryable:
+                self._deferred.append(task)
+                self._deferred_errors[task] = exc
             # Only non-retryable errors (auth, schema, code bugs) count toward
             # the early-shutdown gate. Retryable errors (rate-limit, timeout,
             # transient 5xx, connection blips) cluster under provider degradation
             # and would otherwise trip the gate even when salvage could recover.
             if not retryable:
                 self._check_error_rate(success=False)
-            if uses_model_stage_resource:
-                self._record_retryable_outcome(retryable=retryable)
+            if uses_model_stage_resource and generator_call_state.invoked:
+                self._record_retryable_outcome(retryable=retryable, exc=exc)
             if not retryable and self._reporter:
                 self._reporter.record_failure(task.column)
             if self._trace and trace:
@@ -1962,8 +2068,6 @@ class AsyncTaskScheduler:
                 raise
 
             if retryable:
-                self._deferred.append(task)
-                self._deferred_errors[task] = exc
                 self._emit_scheduler_event(
                     "retry_deferred", task=task, lease=lease, task_execution_id=task_execution_id
                 )
@@ -2043,8 +2147,16 @@ class AsyncTaskScheduler:
             self._record_observed_task_state()
             self._wake_event.set()
 
-    async def _run_generator_call(self, task: Task, operation: str, call: Coroutine[Any, Any, Any]) -> Any:
+    async def _run_generator_call(
+        self,
+        task: Task,
+        operation: str,
+        call: Coroutine[Any, Any, Any],
+        *,
+        call_state: _GeneratorCallState,
+    ) -> Any:
         """Run user/plugin generator code while preserving scheduler-owned failures."""
+        call_state.invoked = True
         try:
             return await call
         except Exception as exc:
@@ -2092,7 +2204,13 @@ class AsyncTaskScheduler:
     ) -> bool:
         return self._task_supports_row_drops(task, generator) and isinstance(exc, UserTemplateError)
 
-    async def _run_from_scratch(self, task: Task, generator: ColumnGenerator) -> Any:
+    async def _run_from_scratch(
+        self,
+        task: Task,
+        generator: ColumnGenerator,
+        *,
+        call_state: _GeneratorCallState,
+    ) -> Any:
         """Execute a from_scratch task."""
         rg_size = self._get_rg_size(task.row_group)
         # Runtime import: needed for isinstance check; module-level would cause circular import
@@ -2103,6 +2221,7 @@ class AsyncTaskScheduler:
                 task,
                 "from-scratch generation",
                 generator.agenerate_from_scratch(rg_size),
+                call_state=call_state,
             )
             result_operation = "From-scratch generator"
         else:
@@ -2114,7 +2233,12 @@ class AsyncTaskScheduler:
             # raises ``UserTemplateError`` instead of silently dropping the
             # whole row group.
             if self._task_supports_row_drops(task, generator):
-                return await self._run_full_column_from_scratch_with_row_drops(task, generator, rg_size)
+                return await self._run_full_column_from_scratch_with_row_drops(
+                    task,
+                    generator,
+                    rg_size,
+                    call_state=call_state,
+                )
 
             # Non-FromScratch generators dispatched as seeds (no upstream columns)
             # operate on existing buffer rows. Pass an ``rg_size``-row snapshot
@@ -2130,6 +2254,7 @@ class AsyncTaskScheduler:
                 task,
                 "full-column generation",
                 generator.agenerate(input_df),
+                call_state=call_state,
             )
             result_operation = "Full-column generator"
         result_df = self._require_dataframe_result(
@@ -2154,6 +2279,8 @@ class AsyncTaskScheduler:
         task: Task,
         generator: ColumnGenerator,
         rg_size: int,
+        *,
+        call_state: _GeneratorCallState,
     ) -> Any:
         """Execute a row-drop-capable seed task (root expression columns).
 
@@ -2179,6 +2306,7 @@ class AsyncTaskScheduler:
             task,
             "full-column generation",
             generator.agenerate(input_df),
+            call_state=call_state,
         )
 
         result_df = self._require_expression_row_drop_result(
@@ -2208,7 +2336,13 @@ class AsyncTaskScheduler:
 
         return result_df
 
-    async def _run_cell(self, task: Task, generator: ColumnGenerator) -> tuple[Any, bool]:
+    async def _run_cell(
+        self,
+        task: Task,
+        generator: ColumnGenerator,
+        *,
+        call_state: _GeneratorCallState,
+    ) -> tuple[Any, bool]:
         """Execute a cell-by-cell task. Returns ``(result, skipped)``."""
         if task.row_index is None:
             raise ValueError(f"Cell task requires a row_index, got None for column '{task.column}'")
@@ -2234,6 +2368,7 @@ class AsyncTaskScheduler:
             task,
             "cell generation",
             generator.agenerate(dict(record)),
+            call_state=call_state,
         )
 
         # Write back to buffer (include side-effect columns)
@@ -2266,7 +2401,13 @@ class AsyncTaskScheduler:
             side_effect_columns=self._graph.get_side_effect_columns(task.column),
         )
 
-    async def _run_batch(self, task: Task, generator: ColumnGenerator) -> Any:
+    async def _run_batch(
+        self,
+        task: Task,
+        generator: ColumnGenerator,
+        *,
+        call_state: _GeneratorCallState,
+    ) -> Any:
         """Execute a full-column/batch task."""
         rg_size = self._get_rg_size(task.row_group)
 
@@ -2312,6 +2453,7 @@ class AsyncTaskScheduler:
             task,
             "batch generation",
             generator.agenerate(batch_df),
+            call_state=call_state,
         )
         if self._task_supports_row_drops(task, generator):
             result_df = self._require_expression_row_drop_result(

{data_designer_engine-0.8.0 → data_designer_engine-0.9.0}/src/data_designer/engine/dataset_builders/dataset_builder.py

@@ -456,7 +456,7 @@ class DatasetBuilder:
         the source of truth for run *progress* (``num_completed_batches``,
         ``actual_num_records``). Splitting the two sources is what lets resume
         survive a crash between writing a batch and updating metadata: the
-        filesystem reflects the durable state even when metadata lags by a step.
+        filesystem reflects the durable state even when metadata lags behind.
 
         ``num_records`` must be >= the number of records already on disk (you may
         extend a dataset, but cannot shrink it below what has been written).
@@ -767,7 +767,7 @@ class DatasetBuilder:
             state = self._load_resume_state(num_records, buffer_size)
             # _load_resume_state already scans the filesystem for completed row groups
             # and exposes them via state.completed_row_groups. The filesystem is the
-            # source of truth for progress (metadata may lag by one row group between
+            # source of truth for progress (metadata may lag between
             # move_partial_result_to_final_file_path and write_metadata).
             completed_row_groups = state.completed_row_groups
             completed_ids = set(completed_row_groups)
@@ -788,6 +788,14 @@ class DatasetBuilder:
             remaining_row_group_count = len(resume_plan.remaining_row_groups)
             completed_row_group_count = resume_plan.total_row_groups - remaining_row_group_count
             if remaining_row_group_count == 0:
+                self.artifact_storage.update_metadata(
+                    {
+                        "actual_num_records": state.actual_num_records,
+                        "total_num_batches": state.num_completed_batches,
+                        "file_paths": self.artifact_storage.get_file_paths(),
+                        "num_completed_batches": state.num_completed_batches,
+                    }
+                )
                 logger.warning(
                     "⚠️ Dataset is already complete — all row groups were found in the existing artifact "
                     "directory. Nothing to resume. Use resume=ResumeMode.NEVER if you want to generate a new dataset."
@@ -801,18 +809,24 @@ class DatasetBuilder:
 
             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 +875,7 @@ class DatasetBuilder:
         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,

{data_designer_engine-0.8.0 → data_designer_engine-0.9.0}/src/data_designer/engine/dataset_builders/scheduling/completion.py

@@ -52,6 +52,8 @@ class CompletionTracker:
         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 @@ class CompletionTracker:
         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 @@ class CompletionTracker:
         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):

{data_designer_engine-0.8.0 → data_designer_engine-0.9.0}/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 @@ class FairTaskQueue:
         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 @@ class FairTaskQueue:
     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 @@ class FairTaskQueue:
 
         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 @@ class FairTaskQueue:
     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 @@ class FairTaskQueue:
         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:

{data_designer_engine-0.8.0 → data_designer_engine-0.9.0}/src/data_designer/engine/dataset_builders/scheduling/task_policies.py

@@ -29,7 +29,8 @@ TaskAdmissionDenyReason = Literal[
     "shutdown",
     "policy_denial",
 ]
-DEFAULT_DYNAMIC_BORROW_RESERVE_FRACTION = 0.125
+# 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(

{data_designer_engine-0.8.0 → data_designer_engine-0.9.0}/src/data_designer/engine/models/clients/adapters/anthropic_translation.py

@@ -15,6 +15,7 @@ from data_designer.config.utils.media_helpers import (
 from data_designer.engine.models.clients.parsing import extract_usage, fill_reasoning_token_count_from_content
 from data_designer.engine.models.clients.types import (
     AssistantMessage,
+    ChatCompletionChoice,
     ChatCompletionRequest,
     ChatCompletionResponse,
     ToolCall,
@@ -124,7 +125,16 @@ def parse_anthropic_response(response_json: dict[str, Any]) -> ChatCompletionRes
         usage = extract_usage(raw_usage)
         usage = fill_reasoning_token_count_from_content(usage, message.reasoning_content)
 
-    return ChatCompletionResponse(message=message, usage=usage, raw=response_json)
+    choice = ChatCompletionChoice(
+        message=message,
+        finish_reason=response_json.get("stop_reason"),
+    )
+    return ChatCompletionResponse(
+        message=message,
+        usage=usage,
+        raw=response_json,
+        choices=[choice],
+    )
 
 
 def translate_request_messages(