Skip to content

[rollout] feat: add owned train admission holds - #54

Draft
ashtonchew wants to merge 1 commit into
ashtoncw/fully-async-v4-11-terminal-cancellationfrom
ashtoncw/fully-async-v4-12-train-admission-holds
Draft

[rollout] feat: add owned train admission holds#54
ashtonchew wants to merge 1 commit into
ashtoncw/fully-async-v4-11-terminal-cancellationfrom
ashtoncw/fully-async-v4-12-train-admission-holds

Conversation

@ashtonchew

@ashtonchew ashtonchew commented Jul 31, 2026

Copy link
Copy Markdown

Summary

  • Defines rollout lifecycle and owner-scoped train-admission hold contracts, then implements them in FullyAsyncRolloutFn.
  • Blocks new source reservations and owned train-batch lease issuance while a hold is active, while each hold waits for the work admitted before its own terminal frontier.
  • Keeps an owned terminal task tracked until its result is published to the finished-group buffer or its receipt is rolled back, so its terminal frontier observes handoff and settlement failures. A legacy frontier ends when its execution task becomes terminal.
  • Sets the owned buffer's blocking capacity to completed-prefetch capacity P, preserving explicit eviction policy while guaranteeing room for every terminal result that already owns a completed slot.
  • Adds scheduler-local checkpoint preparation and lease-settlement checks. Ray manager and checkpoint-publication wiring remain in PR #56.

Context

PR #53 can prove when scheduled execution becomes terminal and can retry retained cleanup. Weight updates and checkpoint publication need a separate boundary: they must stop new training admission, allow work already admitted at that point to settle, and prevent a completed batch lease from crossing the lifecycle operation unnoticed.

A shared pause flag is insufficient when evaluation, checkpointing, and weight updates can overlap because one caller could reopen admission while another still owns its pause. This layer gives each caller a distinct hold, captures that hold's terminal frontier at acquisition, and reopens admission only after every active owner releases its claim.

The finished-group buffer composed from the current open PR #2030 has two modes. An explicit --async-buffer-max-groups uses receipt-aware eviction. With no explicit policy cap, the buffer blocks at its transport capacity. In owned mode that blocking capacity must be P, the value of --fully-async-max-completed-prefetch-groups. A terminal result acquires one of those P slots before publication, so at most P - 1 other slot-owning results can already be buffered. Publication therefore always has room.

This capacity invariant also closes a terminal-frontier deadlock. With a smaller blocking buffer, one terminal result could fill the buffer while another terminal task blocked during publication. The hold would correctly keep waiting because that task still owned an unpublished receipt, but no consumer remained to make room. Removing the task from _active_executions before publication would hide the receipt from frontier and cleanup checks. Matching the blocking capacity to P lets the task stay tracked through publication or rollback without blocking the hold.

Description

  • Defines TrainAdmissionHold and RolloutFnLifecycle contracts for terminal-frontier waits, checkpoint preparation, and repeatable lifecycle close.
  • Closes source reservation and owned lease admission before returning a hold, then snapshots the active owned or legacy executions that were admitted before that point.
  • Lets each hold wait for its fixed terminal frontier without cancelling execution, consuming completed groups, or issuing a train-batch lease.
  • Retains each terminal task in _active_executions until the worker publishes its result or rolls back its exact terminal receipt, then reports any canonical worker failure to the waiting hold.
  • Separates GroupBuffer's blocking transport capacity from its optional eviction cap. Legacy mode keeps the existing blocking bound, while owned mode uses completed-prefetch capacity P.
  • Keeps overlapping holds independent and reopens admission only after the final active hold releases; lifecycle close permanently dominates outstanding holds.
  • Blocks a completed owned batch before lease issuance and refreshes numeric staleness validation when its admission epoch changed while held. A required refresh failure leaves the batch owned and does not issue a lease.
  • Tracks open train-batch leases and requires an active hold, settled leases, settled rollback work, and a healthy worker before checkpoint preparation succeeds.
  • Attempts to requeue the batch within the same settlement attempt when lease commit fails; if that requeue or an explicit rollback fails, lifecycle cleanup retains the receipts for checkpoint preparation or close to retry.

Test Plan

  • At exact head 81a3bc5726b9ace62e3fe30a933e3f14b0eb1008, python -m pytest -q -p no:cacheprovider tests/fast/rollout/test_base_types.py tests/fast/rollout/test_fully_async_rollout.py passed 103 tests.
  • Validation head d6cd677a89759d6ea5316d2805a3d9cb6f392a1b passed the 8xB200 composition run for candidate 7d53849102b1804fd4a8b21515aa01ca23f89b48: one real actor optimization step, exact trained-weight reset and rebroadcast recovery, update-boundary rollback and replay, one nonempty admission hold, and clean inner and outer Ray completion. The validation head is two validation-only commits above the candidate and includes the router startup-deadline patch, so it is composition-level evidence; it does not establish exact branch-head GPU proof.

@ashtonchew
ashtonchew force-pushed the ashtoncw/fully-async-v4-11-terminal-cancellation branch from eefb379 to 7bc7e52 Compare July 31, 2026 19:00
@ashtonchew
ashtonchew force-pushed the ashtoncw/fully-async-v4-12-train-admission-holds branch 3 times, most recently from ba1a956 to 11ca220 Compare August 3, 2026 20:27
@ashtonchew
ashtonchew force-pushed the ashtoncw/fully-async-v4-11-terminal-cancellation branch 2 times, most recently from 311c9d4 to 85b694b Compare August 3, 2026 21:43
@ashtonchew
ashtonchew force-pushed the ashtoncw/fully-async-v4-12-train-admission-holds branch from 11ca220 to 5e266b5 Compare August 3, 2026 21:43
@ashtonchew
ashtonchew force-pushed the ashtoncw/fully-async-v4-11-terminal-cancellation branch from 85b694b to cffeb69 Compare August 3, 2026 22:50
@ashtonchew
ashtonchew force-pushed the ashtoncw/fully-async-v4-12-train-admission-holds branch from 5e266b5 to 563de22 Compare August 3, 2026 22:50
@ashtonchew
ashtonchew force-pushed the ashtoncw/fully-async-v4-11-terminal-cancellation branch from cffeb69 to d12c73b Compare August 4, 2026 01:00
@ashtonchew
ashtonchew force-pushed the ashtoncw/fully-async-v4-12-train-admission-holds branch from 563de22 to 536c35e Compare August 4, 2026 01:00
@ashtonchew
ashtonchew force-pushed the ashtoncw/fully-async-v4-12-train-admission-holds branch from 536c35e to 81a3bc5 Compare August 4, 2026 09:55
@ashtonchew
ashtonchew force-pushed the ashtoncw/fully-async-v4-11-terminal-cancellation branch from d12c73b to 4762f0b Compare August 4, 2026 09:55
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.

1 participant