Runs WriteTasks and the step device→host staging pool.
The manager owns request/slot identity and when to submit. This class owns the staging pool, the GPU-byte budget, the copy queues, and the single committer that writes into the CPU block pool.
Two device→host paths — the step path does not write chunk.host in the committer:
JOIN_NEXT_STEP save already launched a whole-step device→host into
a staging slot and set `chunk.host` as views.
Committer waits that `step_d2h_event`, then writes
the pool.
JOIN_ON_FINISH committer copies the device clone → owned host
tensors, then writes the pool.
Async: high-priority then low-priority queues, then pool write. Eager: submit() does wait+write inline (CPU tests / no CUDA).
Staging pool + committer. Step device→host is launched at save; this thread waits that event (JOIN_NEXT_STEP) or copies deferred rows (JOIN_ON_FINISH), then writes into the CPU pool.
append_chunk(
task: WriteTask,
chunk: _WriteChunk,
freeze_event: Event | None = None,
) -> TaskState | None
Append to a pending task. None when appended, else the closing state.
Hand registered, queued tasks to the copy path. Threaded: enqueue. Eager: the copy + pool write run here, inline — never call this under the manager's state lock.
Pop pool-written tasks from _completed and drop them from _tasks. Staging holders were already released at the pool write.
Move pending tasks to the front of the high-priority queue.
PENDING -> QUEUED at the head; QUEUED on the low-priority queue moves up; QUEUED already high-priority is a no-op. COPYING and later belong to the worker and are untouched: the worker claims under _wake too, so a task cannot be popped and re-queued behind its back.
Rows for slots of one not-yet-done JOIN_ON_FINISH task.
_slot_ref puts JOIN_NEXT_STEP tids in join_tids (wait then pool). This path reads committer-written chunk.host, or the device clone if that device→host has not landed.
Block until each task has finished the CPU-pool write (or failed).
Block until each task's device→host is complete (host_ready).
Staging: committer has waited step_d2h_event. Deferred: committer has written chunk.host. Does not wait for the CPU-pool write.
pin_budget(ticket: _BudgetTicket, tid: int) -> None
Record that tid holds a view of the clone ticket charges.
Make a task visible (registry + QUEUED) without running anything; safe under the manager's state lock. queued=False (deferred tasks) stays PENDING on the GPU clone until finish/abort or the GPU-byte budget forces a copy. Queued tasks must then go through dispatch.
Caller must reserve() the task bytes first (budget flush can block; the manager does that outside the state lock) and pin the task on its budget ticket(s) before register.
reserve(nbytes: int) -> None
Reserve GPU-clone bytes; blocking flush happens here, so callers must not hold the manager's state lock.
stage_step_host(
tensors: dict[str, Tensor],
n: int,
freeze_event: Event | None,
step_holder: StagingBufferHolder,
) -> StepD2HClaim
Claim a staging slot and, when tensors is non-empty, launch ONE whole-step device→host into it.
Leftover-only saves pass empty tensors and still take a slot (empty views) so every step id shares this bound. A full pool waits for materialize/discard; timeout then errors. A step larger than the page overflows the next slot — that is a config break. The caller binds tasks after submit; step_holder is released by materialize/discard via staging_release.
Bases: NamedTuple
One owner of a staging-buffer slot. The slot is free when none remain.
Not a buffer state — concurrent owners share the same slot: - for_step: claimed at save, released when materialize/discard consumes the ctx - for_task: bound before WriteTask submit, released when that task completes
Reusable pinned landing zone for ONE whole-step device→host at save.
Per-task chunk.host is a row-range view into a slot, so the committer skips a per-task device→host. Slots recycle; this is not the CPU block pool.
A slot stays busy while anyone still holds it: the step (until materialize/discard) and each immediate write that views the page (until its pool write). Prefix hits do not hold a slot — they wait for the pool write and read the durable pool.
Saves with only leftover mm still claim a slot (empty views) so every step id shares this bound. A full pool waits; timeout then errors.
claim(holder: StagingBufferHolder, timeout: float) -> int
Grab a free slot for holder. Waits until one is free, then times out. timeout<=0 fails immediately if none are free.
dataclass ¶
One whole-step landing in StagingBufferPool.
Return of stage_step_host. The manager stores this on _StepContext until materialize/discard releases the step holder.
Bases: Enum
WriteTask lifecycle: a strict chain, plus FAILED from any non-terminal.
PENDING -> QUEUED -> COPYING -> HOST_READY -> WRITTEN
PENDING registered, not on a copy queue (deferred: waiting for finish / budget pressure) QUEUED on a copy queue COPYING one thread owns the copy stage (staging: waits the step device→host; deferred: copies the device clone) HOST_READY host rows ready (host_ready set); device clone dropped WRITTEN in the CPU pool (done set) FAILED committer could not finish (host_ready + done set; manager raises from every later entry)
Moves only through WriteTask.transition; skipping a step is illegal.
dataclass ¶
One write of (slot, key) rows for a single request.
Identity: tid is the handle in the manager's (slot, key) tables. req_id + write_n mark whose write this is and the nth time that request opened a write. One write may cover several keys.
state is the single source of truth for where the write is (see TaskState). host_ready / done are wait primitives set by the transitions into HOST_READY / WRITTEN / FAILED, never directly.
How HOST_READY is reached: - JOIN_NEXT_STEP: chunk.host is a staging view set at save; the copy stage only waits step_d2h_event. - JOIN_ON_FINISH: the copy stage copies device→host into chunk.host.
JOIN_NEXT_STEP is queued at submit. JOIN_ON_FINISH stays PENDING until finish or GPU-byte-budget pressure escalates it; budget flush takes the unfinished task with the oldest enqueued_time.
Concurrent readers/writers: - Staging readers (materialize clone, committer pool write) all wait the same step_d2h_event before touching the view. - A later task taking the same (slot, key) records those rows in reassigned; the old pool write skips them. - append_chunk is refused once the copy stage is claimed; the caller opens a fresh task rather than mutating a closed one. - lock covers state / reassigned / append_chunk / host↔freeze / slot_to_row. scatter_rows snapshots reassigned.
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
reassigned: dict[TensorName, Tensor] = field(
default_factory=dict
)
append_chunk(
chunk: _WriteChunk, freeze_event: Event | None = None
) -> TaskState | None
Grow this write with one save's rows. Returns None when appended, else the state that closed the task (COPYING or later).
freeze_event is stored in the same snapshot: events on one compute stream are ordered, so the newest also covers every earlier clone.
budget_tickets() -> list[_BudgetTicket]
Distinct tickets across this task's chunks (by identity; one per save).
claim_copy() -> bool
QUEUED -> COPYING. Only one thread may run the copy stage; True if this caller won.
Drop the device freeze. Host is unchanged (staging wait / fail).
chunks[si] host if written, else device freeze. One snapshot.
-> FAILED from any non-terminal state; unblocks joiners. False if already terminal. defer_wake: the caller publishes the failure record first, then calls wake_failed_waiters.
COPYING -> HOST_READY for staging views: wait the step device→host, drop freeze.
reassigned_intersects(chunk: _WriteChunk) -> bool
True if this write already lost ownership of any (key, slot) the chunk would add; the caller opens a fresh WriteTask instead.
scatter_rows() -> list[tuple[TensorName, Tensor, Tensor]]
(key, slots, host) to write, one entry per key. Omits slots in reassigned. A slot written twice by this task (preempt + resume onto the same block) keeps the later chunk's row.
COPYING -> HOST_READY: write these host tensors, drop the device freeze.
transition(to: TaskState) -> None
Move to to; raises unless it is the next chain step or FAILED.
Move to to if legal from the current state. True if moved.
Set the events after the failure record is published (see mark_failed(defer_wake=True)): a woken joiner must find it.