Skip to content

perf(scheduler): improve utilization and scale#810

Open
eric-tramel wants to merge 6 commits into
mainfrom
codex/telemetry-scheduler-optimization
Open

perf(scheduler): improve utilization and scale#810
eric-tramel wants to merge 6 commits into
mainfrom
codex/telemetry-scheduler-optimization

Conversation

@eric-tramel

@eric-tramel eric-tramel commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

Summary

Improves async generation throughput, model-slot utilization, and large-run scheduler scaling with four targeted changes. The implementation is net +16 nonblank production LoC and adds no public API, dependency, feedback controller, or user tuning knob.

Changes

  • Make bounded-borrow admission work-conserving: a sole runnable scheduling group can use all resource slots, while existing borrow debt gives a newly waiting peer priority as slots are released.
  • Bound fair-queue priority state by active scheduling groups instead of retaining historical task entries.
  • Compact checkpointed completion state after a row group is durably written and all of its workers finish.
  • Write resume metadata after the first durable row group and refresh it at completion; parquet files remain the source of truth for resume progress.

Performance vs. main

End-to-end throughput and utilization

Full-job reservation measurements count each physical endpoint once, including shared-model endpoints. The reported unit is endpoint concurrent-service-capacity slot-hours, a GPU-hours proxy until the inference owner supplies the endpoint-to-GPU allocation.

Workload main This PR Change
Heavy-root/late-peer simulation, wall 36.70 s 34.00 s -7.36%
Heavy-root/late-peer simulation, model-slot utilization 91.55% 98.82% +7.27 pp
Heavy-root/late-peer simulation, idle slot-seconds 24.80 s 3.20 s -87.1%
In-process slow→fast DAG, 32 records 71.18 ms / 449.6 records/s 65.85 ms / 486.0 records/s wall -7.49%, throughput +8.09%, utilization +9.81 pp
In-process slow→fast DAG, 512 records 798.02 ms / 641.6 records/s 783.90 ms / 653.1 records/s wall -1.77%, throughput +1.80%
OpenAI-compatible HTTP DAG, 32 records 0.930 s / 34.4 records/s 0.889 s / 36.0 records/s wall -4.38%, throughput +4.60%, utilization +9.36 pp
OpenAI-compatible HTTP DAG, 256 records 4.044 s / 63.3 records/s 4.026 s / 63.6 records/s wall -0.45%, throughput +0.45%, utilization +3.95 pp
Endpoint-queued HTTP DAG, 64 records 1.950 s / 32.8 records/s 1.967 s / 32.5 records/s wall +0.86%, throughput -0.86%, utilization +0.52 pp; statistically neutral
Calibrated slow→fast corpus, 160,032 requests 105.919 s / 0.7061 slot-h 105.907 s / 0.7060 slot-h wall/reserved -0.01%, full-job utilization 32.45%→32.54%; practical equivalence
Calibrated fast→slow corpus, 160,032 requests 105.839 s / 0.7056 slot-h 106.230 s / 0.7082 slot-h wall/reserved +0.37%, full-job utilization 32.09%→32.06%; practical equivalence
Heterogeneous fork/join corpus, wall 113.485 s 85.979 s -24.238%, slow-endpoint utilization +21.931 pp
Heterogeneous fork/join, full-job endpoint reservation 1.1349 slot-h / 10.227% utilized 0.8598 slot-h / 13.498% utilized reserved -24.238%, idle -26.999%
Five-wide → slow-join → five-wide corpus 212.510 s / 1.06255 slot-h 100.393 s / 0.50196 slot-h wall/reserved -52.760%, idle -56.462%; 99.17% CI [-53.429%, -52.083%]
Five-wide → slow-join → five-wide endpoint flow 6.557% utilized / 23.506% all-provider blackout 13.880% utilized / 3.773% blackout utilization +7.323 pp, blackout -19.733 pp, fast gap p95 -33.8%
Shared-endpoint fairness corpus, wall 122.846 s 123.052 s +0.168%, utilization -0.011 pp; practical equivalence
Dynamic-capacity mixed roots, endpoint cap 2→8 169.564 s / 0.53187 slot-h 171.183 s / 0.53727 slot-h wall +0.956%, reserved +1.016%, idle +3.497%, utilization -0.473 pp; 99.17% CI [-1.209%, +3.168%], inconclusive

Work-conserving admission removes speculative idle capacity while remaining practically equivalent for both slow-before-fast and fast-before-slow calibrated flows. It cannot preempt work already running: in the measured worst case, a late peer started up to 350 ms later when it arrived behind a 400 ms task; borrow debt still gives it the next released slot.

Large-run scheduler and memory scaling

Workload main This PR Change
Completion state, 1M records in 1,000 groups 217.180 MiB 0.059 MiB -99.97%
Completion state, 100k one-row groups 151.899 MiB 8.045 MiB -94.70%
Fair-queue scheduler, 1k records 0.240 s 0.042 s 5.69× faster
Fair-queue scheduler, 10k records 22.794 s 0.462 s 49.34× faster
Fair-queue scheduler, 100k records unfinished at 369.6 s 4.935 s >74.9× faster to cutoff
Fair-queue scheduler, 1M records impractical on the historical-state path 49.085 s 20,373 records/s, no retained priority state

The pathological one-row-group-per-record case retains one compact terminal marker per group: 70.511 MiB at one million groups. With normal buffer sizes, retained completion state is effectively constant relative to record count.

Metadata checkpointing

Workload Change vs. main
Isolated 1M-record metadata loop 8.705 s → 0.336 s (-96.1%)
Heterogeneous fork/join DAG wall -1.630%, throughput +1.657%, utilization +1.289 pp
Shared-endpoint fairness DAG wall -0.054%, throughput +0.054%, utilization +0.032 pp
Five-wide → slow-join → five-wide DAG wall -0.222%, throughput +0.222%, utilization +0.260 pp
Incremental metadata blocking across complex DAGs -99.50% to -99.85%

The shared-endpoint and hourglass cells are practically equivalent on wall time within a ±2% margin, while the heterogeneous fork/join cell improves both throughput and utilization.

Signed-off-by: Eric W. Tramel <eric.tramel@gmail.com>
Signed-off-by: Eric W. Tramel <eric.tramel@gmail.com>
Signed-off-by: Eric W. Tramel <eric.tramel@gmail.com>
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.
@andreatnvidia
andreatnvidia marked this pull request as ready for review July 13, 2026 18:00
@andreatnvidia
andreatnvidia requested a review from a team as a code owner July 13, 2026 18:00
@greptile-apps

greptile-apps Bot commented Jul 13, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR improves async dataset scheduling throughput and large-run state handling. The main changes are:

  • Work-conserving bounded-borrow admission for solo runnable groups.
  • Active-group bounded fair-queue priority state.
  • Row-group completion compaction after durable checkpointing.
  • Reduced resume metadata writes with final metadata refresh.

Confidence Score: 5/5

This looks safe to merge.

  • No blocking issues found in the changed code.

Important Files Changed

Filename Overview
packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/completion.py Adds compacted dropped-row state for checkpointed row groups while preserving terminal drop and completion queries.
packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py Waits for row-group workers to finish before checkpointing and compacts tracker state only after checkpoint success.
packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py Writes resume metadata after the first durable row group, refreshes it at completion, and updates already-complete resumes from disk.
packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/queue.py Bounds fair-queue ordering state to active task groups instead of retaining historical heap entries.
packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/task_policies.py Lets solo groups borrow the full resource capacity while preserving peer priority through borrow debt.

Reviews (2): Last reviewed commit: "docs(scheduler): align checkpoint and bo..." | Re-trigger Greptile

Comment on lines 138 to 143
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())

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Compacted Drops Shadow Later Drops

When a row group has been compacted, is_dropped() always reads _compacted_dropped before _dropped. If a late salvage or failure path calls drop_row() for the same compacted group, the new drop is added to _dropped but this changed branch keeps returning the old bitset, so the row can remain visible as not dropped while the tracker still reports the group complete.

Context Used: Do not suggest defensive coding patterns such as a... (source)

Prompt To Fix With AI
This is a comment left during a code review.
Path: packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/completion.py
Line: 138-143

Comment:
**Compacted Drops Shadow Later Drops**

When a row group has been compacted, `is_dropped()` always reads `_compacted_dropped` before `_dropped`. If a late salvage or failure path calls `drop_row()` for the same compacted group, the new drop is added to `_dropped` but this changed branch keeps returning the old bitset, so the row can remain visible as not dropped while the tracker still reports the group complete.

**Context Used:** Do not suggest defensive coding patterns such as a... ([source](greptile.json))

How can I resolve this? If you propose a fix, please make it concise.

@github-actions

Copy link
Copy Markdown
Contributor

Code Review: PR #810perf(scheduler): improve utilization and scale

Summary

This PR makes four targeted performance changes to the async scheduler, netting ~+16 production LoC with no public API, dependency, or tuning-knob additions:

  1. Work-conserving bounded borrow (task_policies.py): default DEFAULT_DYNAMIC_BORROW_RESERVE_FRACTION drops from 0.1250.0, and _dynamic_reserved_slots loses its max(1, …) floor. A solo scheduling group can now use the full resource; borrow debt still hands the next released slot to a newly-queued peer.
  2. Bounded fair-queue state (queue.py): the heapq + _active_heap_keys/_active_heap_entries triple is replaced by a single _active_entries: dict[key, (finish, sequence)], sorted on demand in select_next. State is now bounded by active groups rather than historical task entries.
  3. Compacted completion state (completion.py + async_scheduler.py): after a row group is durably checkpointed and all its workers finish (in_flight_count == 0), compact_row_group releases per-cell _completed/_batch_complete/_dropped/frontier state, retaining only a compact dropped-row bitset and an implicit "complete" terminal marker.
  4. Reduced metadata checkpointing (dataset_builder.py): incremental metadata.json is written only after the first durable row group (plus the final refresh), instead of after every row group. Parquet files remain the source of truth for resume progress.

The change is well-scoped, each commit is atomic, and test coverage is added for every behavior change. Verdict details below.

Findings

Correctness

  • [Verified — sound] Compaction ordering guard. _checkpoint_completed_row_groups now requires state.in_flight_count == 0 before considering a group complete (async_scheduler.py:1596), and compact_row_group is only called after the row group is removed from self._rg_states (del at line 1613) and successfully checkpointed (line 1632–1633). This correctly prevents releasing per-cell state while a sibling worker is still running — the exact hazard test_scheduler_compacts_only_after_row_group_workers_finish covers. Good defensive design.

  • [Verified — sound] Post-compaction query surface. compact_row_group pops _completed, _batch_complete, and prunes _frontier for the group, but the narrower query methods is_complete, is_all_complete, and is_column_complete_for_rg were not updated to consult _compacted_dropped. This is safe only because a compacted group has already been removed from _rg_states and has no in-flight workers, so the scheduler loops that call those methods (async_scheduler.py:1692, 1711, 1767) never touch a compacted group. This invariant is load-bearing but implicit — see the maintainability note below.

  • [Verified — correct] Bitset sizing. bytearray((max(dropped, default=-1) // 8) + 1) yields length 0 for an empty drop set (-1 // 8 == -1), length 1 for max index 0–7, length 2 for index 8, etc. is_dropped guards with 0 <= byte_index < len(dropped_mask), so out-of-range indices (including the empty-mask case) correctly return False. Idempotent re-compaction is handled by the early if row_group in self._compacted_dropped: return.

  • [Verified — no regression] Metadata-write reduction and resume. _recover_progress_from_disk (dataset_builder.py:620) derives num_completed_batches / actual_num_records from parquet files on disk, not from metadata.json counters — the architecture doc already declared the filesystem the source of truth for progress. Writing incremental metadata only once therefore does not affect crash-resume progress recovery; a mid-run crash leaves stale counters in metadata that resume ignores. _check_resume_config_compatibility only reads the config fingerprint from metadata, which is written on the first durable checkpoint. The doc and comment updates (architecture/dataset-builders.md, the "resume-configuration checkpoint" rewording) are consistent with this.

    • [Minor — worth confirming] First-write timing when early groups are fully dropped. finalize_row_group (which performs the gated first metadata write) is only invoked for groups that finalize; a fully-dropped row group takes the free_row_group path and does not call _on_finalize_row_group. So the fingerprint-bearing metadata.json is not written until the first non-fully-dropped group completes. This matches pre-PR behavior (the old code also only wrote after row groups, never at run start), so it is not a regression — but if a run crashes after only fully-dropped groups, there is still no metadata.json. Resume already tolerates a missing metadata file, so this is acceptable; flagging only for confirmation that no consumer expects metadata to exist after the first group regardless of drop status.

Efficiency

  • [Verified — positive] Queue state. Replacing the lazy-deletion heap (which retained stale (finish, sequence, key) tuples across select_next scans and required heapify(list(self._heap)) copies) with a dict keyed by group, sorted on demand, bounds memory by active-group count and eliminates the O(historical-entries) heap growth. Since active groups are typically small, the per-call sorted(...) is cheap and the reported 5–75× speedups on large runs are plausible. _remove_queued_item now eagerly drops the _active_entries slot when a group's queue empties (queue.py:195–196), keeping the dict tight.

  • [Verified — positive] Completion memory. Compaction collapses per-cell sets to a single bitset per checkpointed group, matching the reported ~99.97% reduction on the 1M-record/1000-group workload. The pathological one-row-group-per-record case is correctly acknowledged in the PR body as retaining one marker per group.

Maintainability / Style

  • [Minor] Load-bearing implicit invariant. The safety of the un-updated is_complete/is_all_complete/is_column_complete_for_rg methods depends entirely on "compacted groups are never queried by the live scheduler." A one-line comment on compact_row_group (or an assertion) stating that these narrow queries are not compaction-aware and must not be called for compacted groups would prevent a future caller from silently getting wrong answers (a compacted-but-completed group would report is_complete → False). Consider documenting the invariant.

  • [Nit] Non-descriptive comment marker. task_policies.py:32 — the comment # ponytail: a future peer can take the next released slot; do not reserve idle capacity. reads like a stray codename/tag. The explanation itself is useful; recommend dropping the ponytail: prefix for a clean rationale comment.

  • [Observation — not blocking] _dynamic_reserved_slots floor removal. Dropping max(1, …) means any configured reserve_fraction that rounds down to 0 (not just the new 0.0 default) now reserves 0 slots instead of being floored to 1. This is arguably more correct (honor the configured fraction), and the explicit-ceiling paths are unaffected, but it is a behavior change for any downstream config that relied on the implicit floor. Covered by the updated test_bounded_borrow_dynamic_ceiling_uses_full_solo_capacity_then_yields.

Testing

  • Coverage is added for each behavior change: compaction correctness and idempotency (test_compact_row_group_releases_details_and_preserves_terminal_queries), the checkpoint-failure-skips-compaction path (test_scheduler_does_not_compact_row_group_when_checkpoint_fails), the sibling-worker ordering guard (test_scheduler_compacts_only_after_row_group_workers_finish), bounded queue state (test_fair_task_queue_ordering_state_stays_bounded_by_active_groups), work-conserving admission (test_bounded_borrow_dynamic_ceiling_uses_full_solo_capacity_then_yields), and the two-write metadata cadence (test_build_writes_metadata_after_first_checkpoint_and_at_end). Tests assert on private attributes (_completed, _active_entries, _dropped) — acceptable for white-box scheduler internals and consistent with the existing test style in these files.
  • Note: I was unable to execute the test suite in this CI environment — pytest is not importable and no .venv was present on PATH at review time, so the make install-dev step's environment was not available to this reviewer. The review above is based on static analysis of the diff and surrounding code. The added tests should be confirmed green by CI before merge.

Structural Impact

Structural Impact (graphify, 2.6s)

Risk: HIGH (1 core abstraction(s) modified; 25 import direction violation(s))

  • 11 Python files, 306 AST entities, 2/79 clusters

Import Direction Violations (25)

Legal direction: interface -> engine -> config

  • _string_keyed_counts() (engine) --calls--> .items() (interface)
  • .__init__() (engine) --calls--> .items() (interface)
  • ._record_observed_task_state() (engine) --calls--> .items() (interface)
  • ._request_pressure_diagnostics() (engine) --calls--> .items() (interface)
  • ._record_preserved_retryables() (engine) --calls--> .items() (interface)
  • +20 more

Core Abstractions Modified

High-Connectivity Changes

  • AsyncTaskScheduler (159 deps) in packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py
  • CompletionTracker (99 deps) in packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/completion.py
  • DatasetBuilder (75 deps) in packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py
  • FrontierDelta (72 deps) in packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/completion.py
  • QueueView (65 deps) in packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/queue.py
  • BoundedBorrowTaskAdmissionPolicyConfig (60 deps) in packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/task_policies.py
  • FairTaskQueue (60 deps) in packages/data-designer-engine/src/data_designer/engine/dataset_builders/scheduling/queue.py
  • _DispatchOutcome (49 deps) in packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py
  • +195 more

Cross-Package Dependencies

  • mcp_command() (interface) --calls--> .run() (engine)
  • models_command() (interface) --calls--> .run() (engine)
  • providers_command() (interface) --calls--> .run() (engine)
  • tools_command() (interface) --calls--> .run() (engine)
  • .download_persona_dataset() (interface) --calls--> .run() (engine)
  • ._evaluate_target() (interface) --calls--> .evaluate() (engine)
  • +309 more

Reviewer note on structural impact: The HIGH rating is driven by the blast radius of AsyncTaskScheduler/CompletionTracker/FairTaskQueue (three of the most-connected engine abstractions), which warrants the extra scrutiny applied above on the compaction ordering invariant and resume-progress recovery. The reported import-direction violations (engine --calls--> .items() / .run() on interface) appear to be graph false positives from generic method names (.items(), .run()) resolving to interface-package symbols rather than genuine reverse imports — none of the changed lines introduce a new cross-package import, and the diff stays entirely within data_designer.engine plus a docs file. No new backward-compatibility surface is introduced (no public API, no config schema change beyond an internal default constant).

Verdict

Approve with minor comments (do not block). The four changes are correct, well-isolated, and each is backed by targeted tests. The compaction ordering guard and the parquet-as-source-of-truth resume contract are the two highest-risk areas given the HIGH structural rating, and both hold up under static analysis. Recommended before merge:

  1. Confirm the added tests pass in CI (this reviewer could not run them locally).
  2. Add a one-line comment documenting that is_complete/is_all_complete/is_column_complete_for_rg are not compaction-aware and must not be called for compacted groups (maintainability guard against future misuse).
  3. Drop the ponytail: prefix on the task_policies.py:32 comment.

None of these are blocking.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants