From bb43d90d67ed019397175fa412bf6fd34c5a4f43 Mon Sep 17 00:00:00 2001 From: Crane-Liu Date: Mon, 27 Jul 2026 22:55:58 +0800 Subject: [PATCH] Feat: add bounded whole-run FIFO admission MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reserve a generation-safe pipeline lease before graph construction, so a callback learns there is no capacity before it builds a DAG rather than after: `begin_run` blocks on the lease, and the depth it admits to is the minimum the child backends published rather than a constant. A backend without a depth-two contract therefore keeps serial behavior instead of being handed a slot-1 lease it cannot serve. A retiring run gives its lease back under the same mutex the waiter evaluates its predicate on. Releasing outside it can land between a waiter finding the pool full and its registering on the condition variable, and the notify then reaches nobody — at depth one that strands the only slot there is. Ready queues are partitioned by run and the scheduler pops only from the partition of the run holding the FIFO head, which is what keeps two admitted runs from interleaving device work. The head may execute while its own graph callback is still open — existing callbacks submit device work and wait for L2 communication before returning — while successors stay gated until every prior run is terminal. Popping a slot is not owning it. A run whose callback throws fails and consumes its own unstarted slots, and its fence can then release the run's lease, all while the scheduler sits between reading READY and launching. Every dispatch path now claims the slot with an atomic READY -> RUNNING exchange at the point of launch, and a lost claim leaves the slot to whoever moved it. The cancellation side retries its own exchange, because a producer completing concurrently moves a consumer PENDING -> READY and a single attempt would leave that slot dispatchable under a failing run. `dispatchable_run_id()` is what the scheduler asks for: the head must also be EXECUTING and still own its lease. Cancellation releases dependency references only for the slots it actually failed. Releasing one held by a slot that is still RUNNING would let its producer reach CONSUMED — and its HeapRing output be reclaimed — while the device is still reading it, and the consumer's real completion would then release the same reference twice. `malloc` / `free` / `copy_to` / `copy_from` reach a child directly rather than through a TaskSlot, so the ready-queue FIFO does not order them and the mailbox mutex only serialises one command at a time. They now block until the calling run holds the FIFO head, so a prepared successor cannot free or overwrite child memory the active run is still reading. That wait can be long, so their bindings release the GIL — holding it across the wait froze every other Python thread, including the one that would have let the active run finish. The lease reaches the runtime slot it names: `TaskSlot` carries `{slot_id, generation}` through the local chip mailbox into `run_from_blob`, so a production run executes under its own lease instead of unconditionally on slot 0. Every failure path claims a slot through the same exchange, and the winner is the only one that writes the failure message: a producer poisoning its consumers, a run cancelling its unstarted slots, and a submit that wires onto a producer that has already failed. A plain store could put a CONSUMED slot back to FAILED and then consume it, and release its dependency references, a second time. A submit charges its slot to the run before it can finish building it, so a throw part-way has to leave something the run's cancellation reclaims in full. The scope reference is registered before it is charged, and both happen before anything that can throw: charging first leaves a slot owing a release nothing can make, since cancellation contributes only the terminal release, the threshold is never reached, and the run's task count — and its fence — never resolve. A task is BUILDING until its submit publishes it. Wiring happens under each producer's fanout mutex, so a half-built task is reachable from its producers' fanout lists well before its own fanin and fanout counters are final. A producer that fails in that window claims the slot and stops there, leaving the propagation to the thread that knows the wiring is done; running it from both sides releases every producer reference twice, which reclaims a producer's output while the device is still reading it. A claim won from BUILDING records that the propagation is owed, so a submit that throws before paying it leaves the debt visible to run cancellation rather than a FAILED slot nobody finishes and a fence nothing reaches. The debt is settled at one point, and it is not the start of the propagation: every step up to the first reference release is ordered first and is idempotent — group marking skips already-terminal members, the run error is first-wins — so an exception there resumes, while the releases that follow are not repeatable. A producer that *completes* in that window is the counterpart case, and it is not solved by publishing halves in an order. It advances the fanin release against a count submit has not written, and zero is a count any release passes; acting on that after submit publishes launches a task whose other producers are still running. Readiness is therefore one decision: the count comparison and the state change happen under the same mutex, and publication writes the count together with the transition out of BUILDING under it. Whichever side arrives first is re-evaluated by the other, and exactly one enqueues the task. A task's terminal transition, its failure message and its fanout snapshot are one decision under that same mutex. Split, a consumer can wire itself in after the snapshot and never be released, or read COMPLETED, decline to count the fanin, and still be released — reaching READY one producer early. The group queue's head is observed before its target workers are checked, so a run cancelling in that window can consume the slot and erase its partition. That is a legal outcome, not a corrupt queue: throwing there would end the process, because the scheduler thread has no handler. A terminal run's ready-queue partitions stay erased. Both routes into `enqueue_ready` can race the erase, and re-inserting would rebuild a partition the scheduler will never read again. Direct device control is ordered against the run that issued it, named explicitly rather than read from `building_run_id_`: that field says a callback is open somewhere, not that this thread is inside it, so a public `Worker.copy_*` on another thread would otherwise be charged to whichever run happens to be building. The wait happens before any Python lock is taken, since a callback holding the provenance lock across it would block the paths that let the active run finish, and its binding releases the GIL. Control issued after that run's first `submit_*` is refused outright: the task travels the ready queue and the control travels the mailbox, so holding the FIFO head says nothing about which arrives first, and two such pairs on different chips can each hold the mailbox the other waits for. A run id names nothing on its own — run 1 exists on every Worker — so the callback context carries the Worker it belongs to, and a control call finds the frame for the Worker it is about to touch. Applying the innermost frame instead orders a call on Worker B against Worker A's run and skips B's own admission entirely, while an L4 callback driving its child's run on the same thread needs the parent's frame to still be reachable underneath. Control that belongs to no run is ordered only by being alone, and that has to cover the command rather than the moment before it — a sampled check leaves the caller free to write its mailbox after a submit was admitted behind its back. It therefore takes the serializer submission itself holds, for the whole call, re-entrant per Worker so a command built out of others does not deadlock on it while one that reaches a different Worker still owes that Worker a reservation. The policy is one function, and everything that reaches a child outside a TaskSlot goes through it: `malloc` / `free` / `copy_*`, `committed_device_memory` — a query, but one on the same mailbox — domain and region and queue creation, and every `remote_*` buffer call, which contends with remote task dispatch for the endpoint's command mutex. A `remote_free` behind a live slot reference is the exception that proves it: it only records the debt, sends nothing, and is not ordered. A run whose own cleanup touches the device degrades itself to depth one. The whole-run FIFO orders tasks; it cannot order cleanup, which happens after the native fence and reaches a child through mailbox control rather than a TaskSlot — N+1 allocating a domain while N is still releasing one can leave two collectives each holding a different chip's mailbox and waiting for the other. A run marks itself when it takes a CommDomain, an L3-L2 region or queue, or a remote slot reference, and the next submission waits for that run's cleanup before it is admitted. Successful control that leaves nothing to tear down does not mark it: runs that only dispatch tasks — and runs that merely malloc — keep the full depth, so a control-free hot path is unaffected. Graph callbacks are serialized, which is what makes that a fact rather than a prediction: a predecessor's callback has always returned before the next submission reaches the check, so whether it bears cleanup is already decided and at most one such run can be outstanding. The wait deliberately does not swallow its error the way pre-submission draining does — cleanup that failed leaves collective device state this process can neither describe nor reclaim, so the worker records it and refuses every later operation instead of admitting work on top of it. Every path that can leave a resource on a chip now either hands it to that cleanup or refuses further work itself. The domain sweeps and the deferred remote-buffer flush attempt every handle and then raise the first error, where they previously wrote to stderr and returned success. A domain allocation that reached only some of its chips registers the ones that committed under a handle the run owns, so the fence releases exactly those and not the chip that never allocated — and each chip publishes that commit itself, the instant its window exists, because the window is committed before the reply is written and an RPC that failed afterwards is not an allocation that never happened. A region rollback that cannot give the region back records the refusal directly, because the id was never tracked and no later cleanup can name it; an interrupt through the create, where the id never came back at all, records it for the same reason. Those refusals are re-read by every run-owned control call and every task submission, not only at admission: a callback that caught its own rollback failure would otherwise keep putting work behind device state this worker can no longer reclaim. And the paths that can be interrupted mid-flight close over it — the domain fan-out joins every thread it started even on an asynchronous unwind, because the caller unlinks the request and reply shms on its way out and a dispatch thread inside `control_*` is still reading them. A leak that reads as a clean run is the state the degradation exists to prevent. `Scheduler::Config::before_claim_cb` is a test seam, unset in production: the instant between the pop and the claim is unreachable from outside, so the losing side of the claim can only be exercised from there. --- docs/orchestrator.md | 133 ++- docs/scheduler.md | 17 +- docs/task-flow.md | 63 +- python/bindings/worker_bind.h | 32 +- python/simpler/orchestrator.py | 230 +++++- python/simpler/worker.py | 697 ++++++++++++---- src/common/hierarchical/orchestrator.cpp | 354 ++++++-- src/common/hierarchical/orchestrator.h | 38 +- src/common/hierarchical/scheduler.cpp | 150 ++-- src/common/hierarchical/scheduler.h | 26 + src/common/hierarchical/types.cpp | 152 +++- src/common/hierarchical/types.h | 127 ++- src/common/hierarchical/worker.cpp | 3 + src/common/hierarchical/worker.h | 5 + src/common/hierarchical/worker_manager.cpp | 1 + src/common/hierarchical/worker_manager.h | 11 +- src/common/worker/pipeline_slot_pool.h | 9 +- .../test_worker_async_fifo.py | 337 ++++++++ .../ut/cpp/hierarchical/test_orchestrator.cpp | 362 ++++++++- .../hierarchical/test_pipeline_contract.cpp | 10 + tests/ut/cpp/hierarchical/test_scheduler.cpp | 286 +++++++ tests/ut/py/test_callable_identity.py | 11 +- tests/ut/py/test_worker/test_host_worker.py | 757 +++++++++++++++++- 23 files changed, 3431 insertions(+), 380 deletions(-) create mode 100644 tests/st/a2a3/host_build_graph/worker_async_fifo/test_worker_async_fifo.py diff --git a/docs/orchestrator.md b/docs/orchestrator.md index 00d5dab8c..1bdb11abc 100644 --- a/docs/orchestrator.md +++ b/docs/orchestrator.md @@ -99,10 +99,13 @@ worker.submit(orchestration, args, config).wait() The current L2 backend is synchronous, so L2 `submit()` executes the existing blocking path and returns an already-completed handle. At L3 and above, graph -callbacks remain serialized, but a later submit waits only for every dispatch -in the prior run to cross its endpoint acceptance boundary. On A2A3 onboard, -that boundary is after both device kernels are enqueued and before stream -synchronization. Endpoints without an earlier signal fall back to completion. +callbacks remain serialized, and what admits a later submit is a free pipeline +slot rather than the prior run's acceptance: `begin_run` reserves a +generation-safe lease before the callback is invoked and blocks there when the +negotiated depth is already spent. Endpoint acceptance remains the launch fence +a run's own dispatches advance — on A2A3 onboard, after both device kernels are +enqueued and before stream synchronization, with endpoints lacking an earlier +signal falling back to completion — but it no longer gates the next callback. Each run still owns its completion error, keepalives, and cleanup independently. Remote L3 submit adds two hidden pieces of metadata: final eligible worker-id @@ -160,23 +163,28 @@ SubmitResult Orchestrator::submit_next_level(const CallableIdentity &callable, // NO_DEP: skip both } - // 4. Record fanin on self - s.fanin_count = static_cast(producers.size()); - s.fanin_released = 0; - - // 5. Register with scope (holds slot open until scope_end releases ref) + // 4. Register with scope (holds slot open until scope_end releases ref) scope_.register_task(sid); // increments s.fanout_total by 1 - // 6. Attach fanout edges under each producer's mutex. Producers already + // 5. Attach fanout edges under each producer's mutex. Producers already // completed do not count as live fanins; failed producers poison this - // slot. Route an immediately READY slot through enqueue_ready(). - attach_fanout_and_count_live_producers(sid, producers); - if (s.fanin_count == 0) { - s.state = TaskState::READY; - enqueue_ready(sid); - } else { - s.state = TaskState::PENDING; + // slot. + int32_t live = attach_fanout_and_count_live_producers(sid, producers); + + // 6. Publication: the one point where the slot leaves BUILDING. The count + // and the transition are published together under fanout_mu — see + // §8 "The BUILDING publication rule". + bool ready; + { + std::lock_guard lk(s.fanout_mu); + s.fanin_count.store(live); + ready = s.fanin_released >= live; // releases that landed while BUILDING + if (!s.state.compare_exchange_strong(BUILDING, ready ? READY : PENDING)) { + propagate_failure(sid); // a producer claimed us mid-wiring + return {sid}; + } } + if (ready) enqueue_ready(sid); // outside the lock: takes runs_mu_ // 7. Return handle return {sid}; @@ -556,13 +564,18 @@ insert, erase, and size operations across those threads. Each `TaskSlotState.state` progresses through: ```text -FREE ──► PENDING ──► READY ──► RUNNING ──► COMPLETED ──► CONSUMED ──► FREE - │ │ - └──────────────► FAILED ─────► CONSUMED ──► FREE +FREE ─► BUILDING ─► PENDING ──► READY ──► RUNNING ──► COMPLETED ──► CONSUMED ──► FREE + │ │ │ + └────────────────┴─────────────┴──► FAILED ─────► CONSUMED ──► FREE ``` - **FREE**: slot in the ring pool, not allocated -- **PENDING**: allocated; waiting on live fanin producers +- **BUILDING**: allocated; submit owns it. Step 3 registers its outputs in the + TensorMap and step 5 appends it to its producers' fanout lists, so it is + observable from other threads well before its fanin and fanout counters are + final. No other thread may advance it — see + [the publication rule](#the-building-publication-rule) below +- **PENDING**: published; waiting on live fanin producers - **READY**: all fanins satisfied; queued for Scheduler dispatch - **RUNNING**: dispatched to one or more endpoints - **COMPLETED**: worker(s) done; may still be referenced by fanout / scope @@ -573,12 +586,86 @@ FREE ──► PENDING ──► READY ──► RUNNING ──► COMPLETED ─ State transitions are driven by atomic operations: -- Orch: FREE → PENDING at submit time +- Orch: FREE → BUILDING at slot claim, BUILDING → PENDING/READY at publication - Scheduler: PENDING → READY → RUNNING during dispatch -- Scheduler: RUNNING → COMPLETED or RUNNING/PENDING/READY → FAILED during +- Scheduler: RUNNING → COMPLETED, or BUILDING/PENDING/READY → FAILED during completion / dependency poisoning - Scheduler/Orch cleanup: COMPLETED/FAILED → CONSUMED +### The BUILDING publication rule + +Wiring a task to a producer has to happen under that producer's `fanout_mu` — +otherwise a producer completing concurrently either misses the new consumer or +releases one that has not counted it. That makes a half-built task reachable +from another thread, and BUILDING is what tells that thread the counters are +not final: + +- A **producer that fails** claims the slot through `claim_task_failure` and + **stops there**. It does not mark group members, poison consumers, or release + references — the submitting thread does, once its wiring is done. Running the + propagation from both sides releases every producer reference twice, which + reclaims a producer's output while the device is still reading it. +- A **producer that completes** advances `fanin_released` against a + `fanin_count` submit has not published, and zero is a count any release + passes. Readiness is therefore one decision, not two readable facts: + `try_mark_ready` compares the pair and changes the state under `fanout_mu`, + and publication writes `fanin_count` together with the transition out of + BUILDING under the same lock. A producer that arrives first is re-evaluated + by the publication; one that arrives after reads a published count. Neither + can act on half of the pair, and exactly one enqueues the task. +- **Dispatch** never claims a BUILDING slot: there are no final args or fanin + count to dispatch on. + +`cancel_unstarted_run` is the one caller that *does* propagate for a BUILDING +slot, because submission is closed before it runs — a slot still BUILDING there +is one whose submit threw part-way through wiring, and no submitting thread is +left to publish it. The same applies one step later: a claim won from BUILDING +sets `failure_propagation_pending`, and cancellation takes over any slot still +carrying it. Without that, a submit that threw *after* being claimed would +leave a FAILED slot cancellation skips, holding references the run's fence +waits on forever. + +That debt is settled at one point, and it is not the start of the propagation. +Everything up to the first reference release is ordered first and is +idempotent — group marking skips already-terminal members, the run error is +first-wins, the producer list is a copy — so an exception there leaves the debt +for the next attempt to resume from. The releases that follow are not +repeatable, which is why the debt is cleared immediately before them rather +than after. + +### What a submit that throws leaves behind + +A submit charges its slot to the run before it can finish building it, so a +throw part-way has to leave something the run's cancellation can fully reclaim. +Two orderings carry that: + +- **The scope reference is registered before it is charged.** `fanout_total` + counts a release that `scope_end` will make; charging it before + `scope_->register_task` succeeded would leave a slot owing a release nothing + can make. Cancellation contributes only the terminal release, the threshold + is never reached, and the slot — and with it the run's task count and its + fence — never resolves. +- **Both happen before anything that can throw**, and `fanout_total` is charged + before Step 2 publishes the slot's outputs, because a consumer that wires + onto it from that point on increments the same field. + +### The failure claim + +Every path that fails a task — a completion poisoning its consumers, a run +cancelling its unstarted slots, and a submit that wires onto an already-failed +producer — moves it to FAILED through `claim_task_failure`, and the winner is +the only one that writes `failure_message`. A plain store could put a CONSUMED +slot back to FAILED and then consume it, releasing its dependency references a +second time. + +The claim, the message and the fanout snapshot are all taken under `fanout_mu`, +so a task wiring itself onto a producer either registers before the claim and is +poisoned by it, or observes FAILED — with its reason — and poisons itself. The +same lock covers the ordinary terminal transition in `on_task_complete`: split +from the snapshot, a consumer can wire itself in afterwards and never be +released, or read COMPLETED, decline to count the fanin, and still be released, +reaching READY one producer early. + ### Fanout-release threshold Both paths that can trigger COMPLETED/FAILED → CONSUMED (the scheduler's diff --git a/docs/scheduler.md b/docs/scheduler.md index 25622cfa5..ec996dfb4 100644 --- a/docs/scheduler.md +++ b/docs/scheduler.md @@ -37,10 +37,19 @@ Both immediately-ready submissions and dependency-released consumers use this function. A directed task therefore cannot re-enter a shared NEXT_LEVEL queue after waiting in PENDING. -Each `ReadyQueue` is a mutex-protected non-blocking FIFO. Root submission, -worker completion, and stop requests notify the Scheduler condition variable; -its wait predicate checks the completion FIFO, every ready queue, and the stop -flag. Ready queues have no blocking pop or shutdown state. +Each `ReadyQueue` is a mutex-protected non-blocking FIFO, partitioned by run: +a task is enqueued under its own `run_id`, and the Scheduler pops only from the +partition of the run that currently holds the FIFO head. That is what keeps two +admitted runs from interleaving their device work while both are live. Root +submission, worker completion, and stop requests notify the Scheduler condition +variable; its wait predicate checks the completion FIFO, the dispatchable run's +partitions, and the stop flag. Ready queues have no blocking pop or shutdown +state. + +Popping a slot is not the same as owning it. A run whose graph callback throws +fails and consumes its own unstarted slots, so the Scheduler claims each slot +with an atomic `READY -> RUNNING` compare-exchange at the moment it dispatches; +a lost claim means the cancelling path already owns that slot. ## 2. Scheduler loop diff --git a/docs/task-flow.md b/docs/task-flow.md index f9b109ea9..7a7231772 100644 --- a/docs/task-flow.md +++ b/docs/task-flow.md @@ -64,8 +64,9 @@ to C++: | `w3.submit_sub(handle, …)` dispatched to a SUB child | `LOCAL_PYTHON` | child resolves digest to a Python callable and calls `fn(args)` | All three paths share one mailbox wire format: `MAILBOX_OFF_CALLABLE` is -reserved, and the 32-byte digest prefixes the args blob. The receiving child -does the digest-to-slot resolve in its own address space. +reserved, the run's generation-safe pipeline lease follows `CallConfig`, and +the 32-byte digest prefixes the args blob. The receiving child does the +digest-to-slot resolve in its own address space. The proposed remote L3 path keeps the same callable identity contract, but sends it in a versioned TASK frame. The remote endpoint resolves the digest @@ -292,12 +293,52 @@ own AICore stream and retires it on every exit path, and no record of which image a stream last ran is load-bearing. The AICPU stream carries no such state and stays with its slot. -This is dormant capacity at this layer. The ordinary synchronous entry point -continues to use slot 0, and the chip child's mailbox loop passes no lease, so -every production run is unleased. This contract does not enable a second -mailbox frame, a second device execution, or cross-run publication overlap. -Carrying a lease across the mailbox, and deciding when slot 1 may be leased at -all, belong to whole-run admission. +#### Whole-run FIFO admission + +L3 graph callbacks remain synchronous and serialized, and how many live run +reservations native admission allows is the depth the child backends +negotiated — not a constant. At the negotiated depth two: + +```text +run N: EXECUTING +run N+1: PREPARED +run N+2: blocked in begin_run before its graph callback +``` + +Where a backend publishes depth one — A5, or any runtime without a depth-two +contract — the *second* submission is the one that blocks in `begin_run`, and +there is no prepared successor at all. Code that relies on a later callback to +unblock an earlier run deadlocks on such a backend. + +`begin_run` acquires a generation-safe lease before invoking the callback. +The FIFO head may enter `EXECUTING` while its callback is still building, which +preserves orchestration callbacks that submit device work and wait for L2 +communication before returning. A non-head run remains `BUILDING` or becomes +`PREPARED` when graph construction closes; it cannot execute until every prior +run is terminal. The scheduler observes only the ready-queue partition belonging +to that single active FIFO head. A run's device effects therefore cannot +interleave with another run even when both graphs contain ready tasks. TensorMap +keys remain `(run_id, tensor_key)`, so adjacent runs may reuse the same tensor +address without creating cross-run dependencies. + +The terminal transition releases the reservation and lease exactly once, +wakes whichever submission was blocked on capacity, and activates the next prepared run. Empty +runs take the same transition immediately. If graph construction fails, every +unstarted slot is poisoned and consumed, its ready-queue partition is erased, +and the lease is returned without dispatching device work. + +Each direct chip child publishes its runtime contract's `pipeline_depth` in +the startup mailbox before `INIT_READY`. The parent configures admission to the +minimum published depth. Backends without a depth-two contract therefore keep +depth-one serial behavior instead of receiving an invalid slot-1 lease. + +Whole-run admission decides when a slot may be leased and carries the lease +from `TaskSlot` through the chip mailbox into the runtime slot, so a production +run now executes under the lease its run holds rather than unconditionally on +slot 0. What this contract still does not enable is a second mailbox frame, a +second device execution, or cross-run publication overlap: the endpoint remains +a single synchronous round trip, and the scheduler dispatches only the run that +holds the FIFO head and still owns its lease. Simulation implements the same depth, so the contract means the same thing on both platforms: its runner owns one arena bank and one retained temporary @@ -382,6 +423,9 @@ Where the data goes after submit: Tags are consumed during the same submit call for dep inference and **never carried further**. 3. `CallConfig` — copied into `slot.config` (parent heap, POD) +4. `PipelineSlotLease` — copied from the owning run into + `slot.pipeline_lease`; local chip mailboxes forward `{slot_id, generation}` + to `ChipWorker::run_with_lease`. For the full submit mechanics (ring alloc, TensorMap lookup/insert, scope ref, fanout wiring), see [orchestrator.md](orchestrator.md). @@ -390,7 +434,8 @@ fanout wiring), see [orchestrator.md](orchestrator.md). For local endpoints, after the Scheduler resolves the submitted NEXT_LEVEL target (or chooses an idle SUB worker), `LocalMailboxEndpoint` encodes -`(callable digest, CallConfig, TaskArgs)` into the per-worker shm mailbox and +`(callable digest, CallConfig, PipelineSlotLease, TaskArgs)` into the +per-worker shm mailbox and the forked child decodes it. Remote NEXT_LEVEL dispatch through `RemoteL3Endpoint` serializes the same logical payload into a framed TASK request instead. diff --git a/python/bindings/worker_bind.h b/python/bindings/worker_bind.h index 8b5a84441..39b3f5990 100644 --- a/python/bindings/worker_bind.h +++ b/python/bindings/worker_bind.h @@ -245,6 +245,7 @@ inline void bind_worker(nb::module_ &m) { // --- TaskState --- nb::enum_(m, "TaskState") .value("FREE", TaskState::FREE) + .value("BUILDING", TaskState::BUILDING) .value("PENDING", TaskState::PENDING) .value("READY", TaskState::READY) .value("RUNNING", TaskState::RUNNING) @@ -307,19 +308,30 @@ inline void bind_worker(nb::module_ &m) { nb::arg("digest"), nb::arg("kind"), nb::arg("target_namespace"), nb::arg("args_list"), "Submit a group of SUB tasks: N args -> N workers, 1 DAG node." ) + .def( + "await_run_admission", + [](Orchestrator &self, uint64_t run_id) { + self.await_run_admission(static_cast(run_id)); + }, + nb::arg("run_id"), nb::call_guard(), + "Block until run_id holds the whole-run FIFO head, or is terminal. Releases " + "the GIL: the wait is unbounded, and holding it would stop the very threads " + "that let the active run finish." + ) .def( "malloc", [](Orchestrator &self, int worker_id, size_t size) { return self.malloc(worker_id, size); }, - nb::arg("worker_id"), nb::arg("size"), "Allocate memory on next-level worker." + nb::arg("worker_id"), nb::arg("size"), nb::call_guard(), + "Allocate memory on next-level worker." ) .def( "committed_device_memory", [](Orchestrator &self, int worker_id) { return self.committed_device_memory(worker_id); }, - nb::arg("worker_id"), + nb::arg("worker_id"), nb::call_guard(), "Total device HBM (bytes) committed by next-level worker's MemoryAllocator " "(tensors + pooled arenas + runtime buffers)." ) @@ -328,21 +340,24 @@ inline void bind_worker(nb::module_ &m) { [](Orchestrator &self, int worker_id, uint64_t ptr) { self.free(worker_id, ptr); }, - nb::arg("worker_id"), nb::arg("ptr"), "Free memory on next-level worker." + nb::arg("worker_id"), nb::arg("ptr"), nb::call_guard(), + "Free memory on next-level worker." ) .def( "copy_to", [](Orchestrator &self, int worker_id, uint64_t dst, uint64_t src, size_t size) { self.copy_to(worker_id, dst, src, size); }, - nb::arg("worker_id"), nb::arg("dst"), nb::arg("src"), nb::arg("size"), "Copy host src to worker dst." + nb::arg("worker_id"), nb::arg("dst"), nb::arg("src"), nb::arg("size"), + nb::call_guard(), "Copy host src to next-level worker." ) .def( "copy_from", [](Orchestrator &self, int worker_id, uint64_t dst, uint64_t src, size_t size) { self.copy_from(worker_id, dst, src, size); }, - nb::arg("worker_id"), nb::arg("dst"), nb::arg("src"), nb::arg("size"), "Copy worker src to host dst." + nb::arg("worker_id"), nb::arg("dst"), nb::arg("src"), nb::arg("size"), + nb::call_guard(), "Copy worker src to next-level worker." ) .def( "alloc", @@ -359,7 +374,7 @@ inline void bind_worker(nb::module_ &m) { .def("scope_end", &Orchestrator::scope_end, "Close the innermost scope. Non-blocking.") .def("_scope_begin", &Orchestrator::scope_begin) .def("_scope_end", &Orchestrator::scope_end) - .def("_begin_run", &Orchestrator::begin_run) + .def("_begin_run", &Orchestrator::begin_run, nb::call_guard()) .def("_close_run_submission", &Orchestrator::close_run_submission, nb::arg("run_id")) .def( "_fail_run_submission", @@ -419,6 +434,10 @@ inline void bind_worker(nb::module_ &m) { nb::arg("worker_id"), nb::arg("mailbox_ptr"), nb::arg("child_pid") = -1, "Add a NEXT_LEVEL sub-worker with an explicit worker id." ) + .def( + "configure_pipeline_depth", &Worker::configure_pipeline_depth, nb::arg("depth"), + "Set run admission depth from the minimum direct-chip runtime capability before init." + ) .def( "add_sub_worker", [](Worker &self, uint64_t mailbox_ptr, int child_pid) { @@ -775,6 +794,7 @@ inline void bind_worker(nb::module_ &m) { m.attr("MAILBOX_SIZE") = static_cast(MAILBOX_SIZE); m.attr("MAILBOX_OFF_ERROR_MSG") = static_cast(MAILBOX_OFF_ERROR_MSG); m.attr("MAILBOX_ERROR_MSG_SIZE") = static_cast(MAILBOX_ERROR_MSG_SIZE); + m.attr("PTO_PIPELINE_MAX_DEPTH") = static_cast(PTO_PIPELINE_MAX_DEPTH); m.attr("MAX_RING_DEPTH") = static_cast(MAX_RING_DEPTH); m.attr("MAX_SCOPE_DEPTH") = static_cast(MAX_SCOPE_DEPTH); diff --git a/python/simpler/orchestrator.py b/python/simpler/orchestrator.py index f0e71aceb..9e4c29be6 100644 --- a/python/simpler/orchestrator.py +++ b/python/simpler/orchestrator.py @@ -36,6 +36,7 @@ def my_orch(orch, args, cfg): import contextlib import operator +import threading from collections.abc import Iterator, Sequence from typing import Any @@ -140,6 +141,116 @@ def _remote_data_eligible_worker_ids( return final_worker_ids +# Which graph callbacks the *current thread* is inside. Direct device control +# has to be ordered against the run that issued it, and a process-wide "a run is +# building" flag cannot say that: a public ``Worker.copy_*`` on another thread +# would be charged to whichever run happens to be building and blocked behind +# it. +# +# A run id alone cannot say it either, because it names nothing on its own — +# run 1 exists on every Worker. An L4 callback drives its children's runs on its +# own thread, so the frames nest and a control call has to find the one for the +# Worker it is about to touch. Applying the innermost frame instead would order +# a call on Worker B against Worker A's run, and skip B's own reservation +# entirely. +_CALLBACK_RUN = threading.local() + + +class _CallbackFrame: + __slots__ = ("has_submitted_task", "run_id", "worker") + + def __init__(self, worker: Any, run_id: int) -> None: + self.worker = worker + self.run_id = int(run_id) + self.has_submitted_task = False + + +def _callback_frames() -> list[_CallbackFrame]: + frames = getattr(_CALLBACK_RUN, "frames", None) + if frames is None: + frames = [] + _CALLBACK_RUN.frames = frames + return frames + + +def _callback_frame_for(worker: Any) -> _CallbackFrame | None: + """The innermost open callback on this thread that belongs to *worker*.""" + for frame in reversed(_callback_frames()): + if frame.worker is worker: + return frame + return None + + +@contextlib.contextmanager +def _callback_run(run_id: int, worker: Any = None): + """Mark this thread as executing *worker*'s *run_id* graph callback.""" + frames = _callback_frames() + frames.append(_CallbackFrame(worker, run_id)) + try: + yield + finally: + frames.pop() + + +def _admit_task_submission(worker: Any = None) -> None: + """Gate one task submission and record that it put work in flight. + + Direct device control after this point cannot be ordered against it: the + task reaches its child through the ready queue and the control through the + mailbox, and waiting for the run to hold the FIFO head says nothing about + which of the two arrives first. Two of them on different chips can each hold + one mailbox and wait for the other. + + The sticky refusal is re-read here for the same reason control re-reads it: + an open callback that caught a failed rollback would otherwise keep + submitting on top of device state this worker can no longer reclaim. + """ + if worker is not None: + worker._require_no_ordered_cleanup_failure("submit") + frame = _callback_frame_for(worker) + if frame is not None: + frame.has_submitted_task = True + + +@contextlib.contextmanager +def direct_control(worker: Any, native_orch: Any, api: str): + """Order one command that reaches a child outside any TaskSlot. + + `malloc`, `copy_*`, domain and region creation, and every `remote_*` + buffer call travel the mailbox rather than the ready queue, so the + whole-run FIFO does not sequence them. Two cases, and the reservation is + held for the *whole* call in both — a check that only samples state leaves + the command itself outside the decision it just made. + + A call issued inside a graph callback belongs to that run and waits for it + to hold the FIFO head. A call that belongs to no run is ordered only by + being alone: it takes the same serializer submission uses, so no run can be + admitted between the check and the command. + """ + frame = _callback_frame_for(worker) + if frame is not None: + if frame.has_submitted_task: + raise RuntimeError( + f"{api}: direct device control cannot follow a task submission in the same run — the task " + "travels the ready queue and this travels the mailbox, so their order is not defined and two " + "such pairs can deadlock across chips. Issue all control before the run's first submit_*()" + ) + # A run that is already known to have left device state behind is not a + # valid owner for more of it, even though it is still open: its callback + # may have caught the rollback failure and carried on. + if worker is not None: + worker._require_no_ordered_cleanup_failure(api) + if native_orch is not None: + native_orch.await_run_admission(frame.run_id) + yield + return + if worker is None: + yield + return + with worker._control_reservation(api): + yield + + class Orchestrator: """DAG builder. Valid only inside the orch function passed to Worker.run(). @@ -177,6 +288,7 @@ def submit_next_level(self, callable_handle: Any, args: TaskArgs, config: CallCo ``worker`` is the exact stable NEXT_LEVEL worker id that runs the task. For L3 chip dispatch, these are the existing chip worker ids. """ + _admit_task_submission(self._worker) cfg = config if config is not None else CallConfig() cpp_worker_id = _require_next_level_worker_id(worker, argument="worker") expected_namespace = ( @@ -246,6 +358,7 @@ def submit_next_level_group( # noqa: PLR0912 -- linear per-member sidecar + eli ``workers`` contains the exact stable NEXT_LEVEL worker id for each member. For L3 chip dispatch, these are the existing chip worker ids. """ + _admit_task_submission(self._worker) cfg = config if config is not None else CallConfig() worker_ids = [_require_next_level_worker_id(value, argument="workers entries") for value in workers] if len(worker_ids) != len(args_list): @@ -341,6 +454,7 @@ def submit_sub(self, callable_handle: Any, args: TaskArgs | None = None): ``args`` may be omitted for a tag-less task (no dependencies, no outputs). """ + _admit_task_submission(self._worker) if args is None: args = TaskArgs() digest, kind, target_namespace, _eligible_worker_ids = _require_handle( @@ -354,6 +468,7 @@ def submit_sub(self, callable_handle: Any, args: TaskArgs | None = None): def submit_sub_group(self, callable_handle: Any, args_list: list): """Submit a group of SUB tasks (N TaskArgs → N workers, 1 DAG node).""" + _admit_task_submission(self._worker) digest, kind, target_namespace, _eligible_worker_ids = _require_handle( callable_handle, kind="orch.submit_sub_group", @@ -403,12 +518,16 @@ def allocate_domain( """ if self._worker is None: raise RuntimeError("allocate_domain requires an Orchestrator bound to a Worker") - return self._worker._allocate_domain( - name=str(name), - workers=tuple(int(w) for w in workers), - window_size=int(window_size), - buffers=list(buffers), - ) + # Collective domain setup drives CTRL_COMM_INIT / CTRL_ALLOC_DOMAIN on + # every member chip, so it is a device effect and must not overtake the + # active run's mailbox traffic. + with self._control_admission("allocate_domain"): + return self._worker._allocate_domain( + name=str(name), + workers=tuple(int(w) for w in workers), + window_size=int(window_size), + buffers=list(buffers), + ) def release_domain(self, handle: CommDomainHandle) -> None: """Collective release. Equivalent to ``handle.release()``.""" @@ -418,7 +537,8 @@ def create_l3_l2_region(self, *, worker_id: int, payload_bytes: int, counter_byt """Create an L3-L2 communication region on one NEXT_LEVEL chip worker.""" if self._worker is None: raise RuntimeError("create_l3_l2_region requires an Orchestrator bound to a Worker") - return self._worker._create_l3_l2_region(int(worker_id), int(payload_bytes), int(counter_bytes)) + with self._control_admission("create_l3_l2_region"): + return self._worker._create_l3_l2_region(int(worker_id), int(payload_bytes), int(counter_bytes)) def create_l3_l2_queue(self, *, worker_id: int, depth: int, input_arena_bytes: int, output_arena_bytes: int): """Create an L3-L2 message queue backed by one L3-L2 communication region.""" @@ -426,13 +546,18 @@ def create_l3_l2_queue(self, *, worker_id: int, depth: int, input_arena_bytes: i raise RuntimeError("create_l3_l2_queue requires an Orchestrator bound to a Worker") from .l3_l2_message_queue import create_l3_l2_queue # noqa: PLC0415 - return create_l3_l2_queue( - self, - worker_id=int(worker_id), - depth=int(depth), - input_arena_bytes=int(input_arena_bytes), - output_arena_bytes=int(output_arena_bytes), - ) + # Reserved across the whole build, not just the region creation it + # nests: the descriptor writes that follow are device effects too. The + # reservation is re-entrant, so the inner create_l3_l2_region joins this + # one rather than deadlocking on it. + with self._control_admission("create_l3_l2_queue"): + return create_l3_l2_queue( + self, + worker_id=int(worker_id), + depth=int(depth), + input_arena_bytes=int(input_arena_bytes), + output_arena_bytes=int(output_arena_bytes), + ) # ------------------------------------------------------------------ # Nested scope (Strict-1 per-scope rings) @@ -478,6 +603,15 @@ def scope(self) -> Iterator[Orchestrator]: finally: self._o.scope_end() + def _control_admission(self, api: str): + """Hold this call's ordering against runs for its whole duration. + + Entered before any Worker lock: the wait can be long, and a callback + holding ``_child_prov_lock`` across it would block the very paths that + let the active run finish. + """ + return direct_control(self._worker, self._o, f"Orchestrator.{api}") + def malloc(self, worker_id: int, size: int) -> int: """Allocate memory on next-level worker *worker_id*. Returns a pointer. @@ -486,24 +620,36 @@ def malloc(self, worker_id: int, size: int) -> int: returned pointer's ``(worker_id, ptr)`` provenance is recorded so a later free / copy / kind4 dispatch to the wrong worker is rejected. """ - wid, sz = int(worker_id), int(size) - if self._worker is None: - return int(self._o.malloc(wid, sz)) - with self._worker._child_prov_lock: - ptr = int(self._o.malloc(wid, sz)) - self._worker._child_prov_record_malloc(wid, ptr, sz) - return ptr + with self._control_admission("malloc"): + wid, sz = int(worker_id), int(size) + if self._worker is None: + return int(self._o.malloc(wid, sz)) + with self._worker._child_prov_lock: + ptr = int(self._o.malloc(wid, sz)) + self._worker._child_prov_record_malloc(wid, ptr, sz) + return ptr def committed_device_memory(self, worker_id: int) -> int: - """Total device HBM (bytes) committed by next-level worker *worker_id*'s ``MemoryAllocator``.""" - return int(self._o.committed_device_memory(int(worker_id))) + """Total device HBM (bytes) committed by next-level worker *worker_id*'s ``MemoryAllocator``. + + A query, but one that travels the same chip mailbox as every other + command, so it takes the same ordering: read behind a run that is still + allocating and the number is a snapshot of neither side. + """ + with self._control_admission("committed_device_memory"): + return int(self._o.committed_device_memory(int(worker_id))) def free(self, worker_id: int, ptr: int) -> None: """Free memory on next-level worker *worker_id*.""" - wid, p = int(worker_id), int(ptr) - if self._worker is None: - self._o.free(wid, p) - return + with self._control_admission("free"): + wid, p = int(worker_id), int(ptr) + if self._worker is None: + self._o.free(wid, p) + return + self._free_locked(wid, p) + + def _free_locked(self, wid: int, p: int) -> None: + assert self._worker is not None with self._worker._child_prov_lock: # Safety-first commit barrier: revoke provenance BEFORE the native # free. If the native free succeeds and an async unwind (e.g. a @@ -518,23 +664,25 @@ def free(self, worker_id: int, ptr: int) -> None: def copy_to(self, worker_id: int, dst: int, src: int, size: int) -> None: """Copy *size* bytes from host *src* to worker *dst*.""" - wid, d = int(worker_id), int(dst) - if self._worker is None: - self._o.copy_to(wid, d, int(src), int(size)) - return - with self._worker._child_prov_lock: - self._worker._child_prov_require_live_range(wid, d, int(size), api="copy_to") - self._o.copy_to(wid, d, int(src), int(size)) + with self._control_admission("copy_to"): + wid, d = int(worker_id), int(dst) + if self._worker is None: + self._o.copy_to(wid, d, int(src), int(size)) + return + with self._worker._child_prov_lock: + self._worker._child_prov_require_live_range(wid, d, int(size), api="copy_to") + self._o.copy_to(wid, d, int(src), int(size)) def copy_from(self, worker_id: int, dst: int, src: int, size: int) -> None: """Copy *size* bytes from worker *src* to host *dst*.""" - wid, s = int(worker_id), int(src) - if self._worker is None: - self._o.copy_from(wid, int(dst), s, int(size)) - return - with self._worker._child_prov_lock: - self._worker._child_prov_require_live_range(wid, s, int(size), api="copy_from") - self._o.copy_from(wid, int(dst), s, int(size)) + with self._control_admission("copy_from"): + wid, s = int(worker_id), int(src) + if self._worker is None: + self._o.copy_from(wid, int(dst), s, int(size)) + return + with self._worker._child_prov_lock: + self._worker._child_prov_require_live_range(wid, s, int(size), api="copy_from") + self._o.copy_from(wid, int(dst), s, int(size)) def alloc(self, shape: Sequence[int], dtype: DataType) -> Tensor: """Allocate a runtime-managed intermediate buffer. diff --git a/python/simpler/worker.py b/python/simpler/worker.py index d3fd1f1bd..ac015eae4 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -82,6 +82,7 @@ def my_l4_orch(orch, args, config): import cloudpickle from _task_interface import ( # pyright: ignore[reportMissingImports] MAX_REGISTERED_CALLABLE_IDS, + PTO_PIPELINE_MAX_DEPTH, RUNTIME_ENV_RING_COUNT, TENSOR_CHILD_MEMORY_OFFSET, WorkerType, @@ -125,7 +126,7 @@ def my_l4_orch(orch, args, config): peek_region_create_reply_region_id, validate_region_create_reply, ) -from .orchestrator import Orchestrator +from .orchestrator import Orchestrator, _callback_run, direct_control from .task_interface import ( MAILBOX_ERROR_MSG_SIZE, MAILBOX_OFF_ERROR_MSG, @@ -176,10 +177,13 @@ def my_l4_orch(orch, args, config): # travels separately via ChipWorker.init(log_level) — not on per-task wire. _RUNTIME_ENV_UINT64_FIELD_COUNT = 3 * RUNTIME_ENV_RING_COUNT _CFG_FMT = struct.Struct("=iiiiii" + ("Q" * _RUNTIME_ENV_UINT64_FIELD_COUNT) + "1024s") -# Args region starts after CONFIG, rounded up to 8 bytes so the first +# The generation-safe pipeline lease follows CONFIG. Args start after the +# lease, rounded up to 8 bytes so the first # Tensor.data (uint64_t at OFF_ARGS+8) is 8-byte aligned, avoiding # SIGBUS on strict-alignment platforms (aarch64 atomics, some ARM cores). -_OFF_ARGS = (_OFF_CONFIG + _CFG_FMT.size + 7) & ~7 +_PIPELINE_LEASE_FMT = struct.Struct("=IIQ") +_OFF_PIPELINE_LEASE = (_OFF_CONFIG + _CFG_FMT.size + 7) & ~7 +_OFF_ARGS = (_OFF_PIPELINE_LEASE + _PIPELINE_LEASE_FMT.size + 7) & ~7 assert _OFF_ARGS % 8 == 0, "_OFF_ARGS must be 8-aligned for Tensor.data" _OFF_TASK_CALLABLE_HASH = _OFF_ARGS _OFF_TASK_ARGS_BLOB = _OFF_TASK_CALLABLE_HASH + CALLABLE_HASH_DIGEST_BYTES @@ -341,11 +345,17 @@ def my_l4_orch(orch, args, config): # window_size (u64), buffer_count (u32), padding (4 bytes) assert _DOMAIN_REQ_HEADER.size == 32 -# Domain-allocation reply shm layout: 24-byte header + buffer_ptrs (u64). -_DOMAIN_REPLY_HEADER = struct.Struct(" set[int]: + held = getattr(_CONTROL_RESERVATION, "workers", None) + if held is None: + held = set() + _CONTROL_RESERVATION.workers = held + return held + + +def _domain_reply_committed(reply_shm: SharedMemory | None) -> bool: + """Whether the chip owning *reply_shm* published a committed window. + + The parent creates the shm zero-filled, so a chip that never ran, never + allocated, or died before its window existed reads as not committed. + """ + if reply_shm is None: + return False + buf = reply_shm.buf + if buf is None: + return False + return struct.unpack_from(" None: + """Apply ``step`` to every item, then raise the first error it produced. + + Teardown of a resource set is terminal: one handle that cannot be reclaimed + must not strand the rest, and it must not be reported as success either — + an unreclaimed device resource is what poisons the worker for its successor. + """ + first_error: BaseException | None = None + for item in items: + try: + step(item) + except BaseException as exc: # noqa: BLE001 + if first_error is None: + first_error = exc + if first_error is not None: + raise first_error + + def _rewrite_blob_host_addrs(buf: memoryview, blob_off: int, ranges: list[tuple[int, int, int]]) -> None: """Redirect registered host pointers in a task-args blob to child mappings. @@ -1215,32 +1272,37 @@ def _handle_ctrl_alloc_domain(cw: ChipWorker, buf: memoryview) -> None: req_buf.release() req_shm.close() - handle = _comm_base_handle(cw) # base communicator handle (cached on the ChipWorker) - device_ctx, local_window_base = cw._impl.comm_alloc_domain_windows( - int(handle), - int(allocation_id), - rank_ids, - int(domain_rank), - int(window_size), - ) - - # Carve buffer pointers sequentially inside the local window. - buffer_ptrs: list[int] = [] - offset = 0 - for nbytes in buffer_nbytes: - if offset + nbytes > window_size: - raise ValueError( - f"alloc_domain: buffer #{len(buffer_ptrs)} (nbytes={nbytes}) at offset={offset} " - f"overflows window_size {window_size}" - ) - buffer_ptrs.append(int(local_window_base) + offset) - offset += int(nbytes) - + # Opened before the collective so the commit can be published the instant + # the window exists. Everything after that point — the carving bounds check, + # the pack — can fail with the allocation already made, and a parent that + # inferred "not allocated" from the failure would leak it. reply_shm = SharedMemory(name=reply_shm_name) reply_buf = reply_shm.buf assert reply_buf is not None try: - _DOMAIN_REPLY_HEADER.pack_into(reply_buf, 0, int(device_ctx), int(local_window_base), int(buffer_count)) + handle = _comm_base_handle(cw) # base communicator handle (cached on the ChipWorker) + device_ctx, local_window_base = cw._impl.comm_alloc_domain_windows( + int(handle), + int(allocation_id), + rank_ids, + int(domain_rank), + int(window_size), + ) + struct.pack_into(" window_size: + raise ValueError( + f"alloc_domain: buffer #{len(buffer_ptrs)} (nbytes={nbytes}) at offset={offset} " + f"overflows window_size {window_size}" + ) + buffer_ptrs.append(int(local_window_base) + offset) + offset += int(nbytes) + + _DOMAIN_REPLY_HEADER.pack_into(reply_buf, 0, 1, int(device_ctx), int(local_window_base), int(buffer_count)) if buffer_ptrs: struct.pack_into(f"<{len(buffer_ptrs)}Q", reply_buf, _DOMAIN_REPLY_HEADER.size, *buffer_ptrs) finally: @@ -1559,9 +1621,14 @@ def handle_task() -> tuple[int, str]: # No-op when nothing is registered. if host_buf_ranges: _rewrite_blob_host_addrs(buf, _OFF_TASK_ARGS_BLOB, host_buf_ranges) + pipeline_slot, pipeline_reserved, pipeline_generation = _PIPELINE_LEASE_FMT.unpack_from( + buf, _OFF_PIPELINE_LEASE + ) + if pipeline_reserved != 0: + raise RuntimeError(f"chip_process dev={device_id}: invalid pipeline lease reserved field") # Hand the mailbox bytes straight to C++ (zero-copy zero-decode): # the blob layout is what `write_blob` already wrote, so re-parsing - # it in Python is N×40B of avoidable work and a permanent + # it in Python is N x 40B of avoidable work and a permanent # opportunity to drop a field. C++ reinterpret_cast # is the source of truth. cw._impl.run_from_blob( @@ -1571,6 +1638,8 @@ def handle_task() -> tuple[int, str]: cfg, mailbox_addr + _OFF_ACCEPTED, _TASK_ACCEPTED, + pipeline_slot, + pipeline_generation, ) except Exception as e: # noqa: BLE001 code = 1 @@ -1770,6 +1839,7 @@ def _chip_process_loop( # noqa: PLR0913 -- fork-child entry: all context (bins, # child to reach _INIT_READY before dispatching the first task, so the # per-rank host-side stream sync budget only covers actual op execution # rather than absorbing peer-rank init skew. + _PIPELINE_LEASE_FMT.pack_into(buf, _OFF_PIPELINE_LEASE, int(cw.pipeline_depth), 0, 0) _mailbox_store_i32(state_addr, _INIT_READY) sys.stderr.write(f"[chip_process pid={os.getpid()} dev={device_id}] ready\n") sys.stderr.flush() @@ -1982,6 +2052,17 @@ class _RunResources: # unreachable from the fence and from close(), which the release itself # made blind by dropping the handle from `_live_domains`. domain_lock: threading.Lock = field(default_factory=threading.Lock) + # Sticky: this run's cleanup itself touches the device, so a successor may + # not be admitted until that cleanup has finished. + # + # The whole-run FIFO orders *tasks*. It cannot order cleanup, which runs + # after the native fence and reaches a child through mailbox control rather + # than a TaskSlot: N+1 allocating a domain while N is still releasing one + # can leave two collectives each holding a different chip's mailbox and + # waiting for the other. Runs that only dispatch tasks keep the full + # admission depth; a run that acquires any of these degrades itself to + # depth one. Set, never cleared. + requires_ordered_cleanup: bool = False class RunHandle: @@ -2009,6 +2090,12 @@ def __init__( self._launch_accepted = False self._terminal = False self._error: BaseException | None = None + # Set only after this run's fence-owned cleanup has actually run. The + # native fence firing is not the same event: it says the device is + # drained, not that the CommDomain / L3-L2 / remote-slot teardown that + # follows it has happened. A successor keyed on the native answer would + # be admitted while that teardown is still outstanding. + self._cleanup_published = False @classmethod def _completed(cls, worker: Worker) -> RunHandle: @@ -2023,6 +2110,8 @@ def _completed(cls, worker: Worker) -> RunHandle: handle._launch_accepted = True handle._terminal = True handle._error = None + # Nothing was ever admitted, so there is no cleanup owing. + handle._cleanup_published = True return handle @staticmethod @@ -2036,11 +2125,17 @@ def _deadline(timeout: float | None) -> float | None: @property def done(self) -> bool: - """True once this run is terminal and its cleanup has been published. + """True once waiting on this run would not block. + + This answers the native fence — the device is drained and every task + has reached its terminal state — plus this handle's own terminal flag. + It does **not** say the run's fence-owned cleanup has happened: the + CommDomain, L3-L2 and remote-slot teardown runs after the fence, in + whichever thread waits. Anything that has to be ordered behind that + teardown keys on ``_cleanup_published`` instead. - Reads False while another waiter is still publishing cleanup, so a run - whose native fence has already fired can report False until that waiter - finishes. + Reads False while another waiter is crossing the fence, because the + native run identity can disappear in that interval. """ with self._cv: if self._terminal: @@ -2151,6 +2246,24 @@ def _wait_for_serialization(self) -> None: # The result remains cached for this handle's public wait/result. pass + def _wait_for_handoff(self) -> None: + """Drain this run *and* its device-touching cleanup before a successor. + + The run's own failure is not this caller's problem — a failed kernel + says nothing about whether cleanup succeeded — so a task error is + swallowed here exactly as pre-submission draining does. What must not be + swallowed is a *cleanup* failure, and that is published by the cleanup + itself rather than inferred from this exception: whoever ran the cleanup + recorded it on the worker before dropping the handle, and any waiter + could have been the one to run it. So this waits, then lets the + admission check see the poison. + """ + try: + self.wait() + except Exception: + # Cached on the handle for its public wait()/result(). + pass + def _wait_for_acceptance(self) -> None: """Wait until this run's dispatches cross their acceptance boundary.""" with self._cv: @@ -2336,13 +2449,22 @@ def __init__( self._last_rollback: dict[str, list[int]] | None = None self._hierarchical_start_mu = threading.Lock() self._hierarchical_start_cv = threading.Condition(self._hierarchical_start_mu) - # Asynchronous completion currently admits only one live device run. A - # new submit drains the prior accepted handle under this gate. + # Device *execution* is one run at a time — the whole-run FIFO admits a + # prepared successor but dispatches only the head — while how many runs + # may be admitted at once is the depth the child backends negotiated. + # This gate serialises graph construction itself, which stays + # synchronous on the submitting caller at any depth. self._submit_mu = threading.Lock() # Guarded by _hierarchical_start_cv. Handles are installed before their # orchestration callback can enqueue work and retired after fence-owned # cleanup, so close() can drain the exact accepted set. self._accepted_run_handles: set[RunHandle] = set() + # Set once a cleanup-bearing run's ordered cleanup failed. That cleanup + # is collective control on the children, so a failure leaves device + # state this process can neither describe nor reclaim; admitting more + # work would build on it. Sticky — there is no recovery short of close. + # Guarded by _hierarchical_start_cv. + self._ordered_cleanup_error: BaseException | None = None # submit graph construction is serialized by _submit_mu. Resource # creation helpers use this pointer to bind new objects to the handle # being built; the pointer is cleared before submit() returns. @@ -2838,7 +2960,7 @@ def remote_malloc(self, *, worker: int, nbytes: int) -> RemoteBufferHandle: size = int(nbytes) if size <= 0: raise ValueError("Worker.remote_malloc nbytes must be positive") - with self._operation_lease("remote_malloc"): + with self._operation_lease("remote_malloc"), self._control_admission("remote_malloc"): self._require_remote_worker_started(worker_id) assert self._worker is not None fields = self._worker.remote_malloc(worker_id, size) @@ -2875,11 +2997,15 @@ def remote_free(self, handle: RemoteBufferHandle) -> None: # must fence admission itself rather than lean on the transport gate.) with self._operation_lease("remote_free"): if handle._live_slot_refs > 0 or handle._live_import_refs > 0: + # Deferred: nothing reaches the owner here, so there is no + # command to order against a run. The fence flushes it, and the + # ordering that matters is the fence's own. handle._mark_released() if handle not in self._pending_remote_buffer_frees: self._pending_remote_buffer_frees.append(handle) return - self._send_remote_free(handle) + with self._control_admission("remote_free"): + self._send_remote_free(handle) handle._mark_released() def remote_copy_to(self, handle: RemoteBufferHandle, host_ptr: Any, nbytes: int, *, offset: int = 0) -> None: @@ -2888,7 +3014,7 @@ def remote_copy_to(self, handle: RemoteBufferHandle, host_ptr: Any, nbytes: int, Requires an owner handle, not an imported one. ``offset + nbytes`` must fall within ``handle.nbytes``. """ - with self._operation_lease("remote_copy_to"): + with self._operation_lease("remote_copy_to"), self._control_admission("remote_copy_to"): self._require_live_remote_buffer(handle) if handle.is_imported: raise ValueError("Worker.remote_copy_to expects an owner remote buffer handle") @@ -2915,7 +3041,7 @@ def remote_copy_from(self, handle: RemoteBufferHandle, host_ptr: Any, nbytes: in Requires an owner handle, not an imported one. ``offset + nbytes`` must fall within ``handle.nbytes``. """ - with self._operation_lease("remote_copy_from"): + with self._operation_lease("remote_copy_from"), self._control_admission("remote_copy_from"): self._require_live_remote_buffer(handle) if handle.is_imported: raise ValueError("Worker.remote_copy_from expects an owner remote buffer handle") @@ -2951,7 +3077,7 @@ def remote_export( requested ``access`` must be a subset of the handle's own access flags — an export can narrow permissions but never widen them. """ - with self._operation_lease("remote_export"): + with self._operation_lease("remote_export"), self._control_admission("remote_export"): return self._remote_export_locked( handle, offset=offset, nbytes=nbytes, access=access, transport_profile=transport_profile ) @@ -3023,7 +3149,7 @@ def remote_import( raise ValueError("Worker.remote_import rejects forged or different Worker RemoteBufferExport values") if exported._owner_handle is not None and exported._owner_handle.released: raise ValueError("Worker.remote_import rejects stale RemoteBufferExport values for released buffers") - with self._operation_lease("remote_import"): + with self._operation_lease("remote_import"), self._control_admission("remote_import"): return self._remote_import_locked(exported, worker=worker, access=access) def _remote_import_locked( @@ -3099,16 +3225,18 @@ def remote_release_import(self, handle: RemoteBufferHandle) -> None: # accepts CLOSED for teardown, so fence admission here). with self._operation_lease("remote_release_import"): if handle._live_slot_refs > 0: + # Deferred: see remote_free — nothing is sent from here. handle._mark_released() if handle not in self._pending_remote_import_releases: self._pending_remote_import_releases.append(handle) return - self._send_remote_release_import(handle) - if handle._owner_handle_ref is not None: - handle._owner_handle_ref._release_import_ref() - handle._owner_handle_ref = None - handle._mark_released() - self._flush_pending_remote_frees() + with self._control_admission("remote_release_import"): + self._send_remote_release_import(handle) + if handle._owner_handle_ref is not None: + handle._owner_handle_ref._release_import_ref() + handle._owner_handle_ref = None + handle._mark_released() + self._flush_pending_remote_frees() def _capture_remote_sidecar_refs(self, remote_sidecar: Any) -> list[RemoteBufferHandle]: captured: list[RemoteBufferHandle] = [] @@ -3136,6 +3264,10 @@ def _adopt_remote_slot_refs(self, handles: list[RemoteBufferHandle]) -> None: self._active_remote_slot_refs.extend(handles) else: resources.remote_slot_refs.extend(handles) + # Releasing a remote slot ref is an RPC to its owner, so it is + # cleanup that reaches off this process. + if handles: + resources.requires_ordered_cleanup = True def _release_remote_slot_refs(self, handles: list[RemoteBufferHandle]) -> None: for handle in handles: @@ -3185,11 +3317,14 @@ def _flush_pending_remote_frees(self) -> None: continue self._pending_remote_buffer_frees.extend(remaining) if errors: - sys.stderr.write( - "Worker._flush_pending_remote_frees(): deferred remote buffer cleanup after control error. " - f"First error: {errors[0]}\n" + # Every handle is attempted and the ones that failed stay pending, + # but the failure is the caller's answer: this runs inside the run + # fence's cleanup, and swallowing it would publish a run as cleanly + # finished while a remote allocation it owns is still held. + raise RuntimeError( + f"Worker._flush_pending_remote_frees: {len(errors)} deferred remote buffer cleanup(s) " + f"failed and remain owed. First error: {errors[0]}" ) - sys.stderr.flush() # ------------------------------------------------------------------ # Callable registration (before init) @@ -3325,6 +3460,11 @@ def _operation_lease(self, api: str): with self._hierarchical_start_cv: if self._lifecycle is not _Lifecycle.READY: raise RuntimeError(f"Worker.{api}: requires an initialized (READY) worker") from self._startup_error + if self._ordered_cleanup_error is not None: + raise RuntimeError( + f"Worker.{api}: a prior run's ordered cleanup failed, so this worker's device state is " + "unreclaimed and no further work is admitted; close() it" + ) from self._ordered_cleanup_error self._active_ops += 1 self._lease_depth[tid] = self._lease_depth.get(tid, 0) + 1 try: @@ -4492,6 +4632,7 @@ def _start_hierarchical(self) -> None: # noqa: PLR0912 -- three parallel fork l device_ids = self._config.get("device_ids", []) n_sub = self._config.get("num_sub_workers", 0) deadline = self._startup_deadline + direct_chip_pipeline_depth = PTO_PIPELINE_MAX_DEPTH # Freeze the startup registry snapshot. init() already holds the epoch in # the INITIALIZING state, so a concurrent register/unregister is blocked @@ -4597,6 +4738,14 @@ def _setup(): # documented in issue #897. A chip that fails or dies during init # raises here rather than spinning forever. self._await_children_ready(self._chip_shms, self._chip_pids, "chip", deadline) + chip_depths = [] + for shm in self._chip_shms: + buf = shm.buf + assert buf is not None + chip_depths.append(_PIPELINE_LEASE_FMT.unpack_from(buf, _OFF_PIPELINE_LEASE)[0]) + if any(depth <= 0 or depth > PTO_PIPELINE_MAX_DEPTH for depth in chip_depths): + raise RuntimeError(f"chip worker published invalid pipeline depths: {chip_depths}") + direct_chip_pipeline_depth = min(chip_depths) # Fork next-level Worker children (L4+ with Worker children). # Each child process eagerly inits the inner Worker, which forks its own @@ -4660,6 +4809,7 @@ def _setup(inner=inner_worker): # the unified mailbox. dw = self._worker assert dw is not None + dw.configure_pipeline_depth(direct_chip_pipeline_depth) # Register chip workers as NEXT_LEVEL (L3). The child pid lets the C++ # endpoint fail a dispatch whose child died instead of spinning on a @@ -5032,6 +5182,7 @@ def _create_l3_l2_region(self, worker_id: int, payload_bytes: int, counter_bytes reply_buf = cast(memoryview, reply_shm.buf) region_id = 0 l3_host_mapping = None + dispatched = False try: L3L2RegionCreateRequest( magic_version=_REGION_MAGIC_VERSION, @@ -5041,6 +5192,11 @@ def _create_l3_l2_region(self, worker_id: int, payload_bytes: int, counter_bytes ).encode_into(req_buf) worker = self._worker assert worker is not None + # From here the chip may have created a region. The create releases + # the GIL, so an interrupt can land after the child committed and + # before the id is read back — the rollback below cannot assume + # "no id" means "nothing exists". + dispatched = True worker.control_l3_l2_region_create(int(worker_id), req_shm.name, reply_shm.name) # Peek before decode: decode rejects malformed replies, but the # child has already created the region and the rollback below @@ -5077,20 +5233,52 @@ def _create_l3_l2_region(self, worker_id: int, payload_bytes: int, counter_bytes self._live_l3_l2_regions.append(region) if resources is not None: resources.l3_l2_regions.append(region) + # Region teardown is mailbox control on its owning chip. + resources.requires_ordered_cleanup = True return region - except Exception: + except BaseException as exc: if l3_host_mapping is not None: try: l3_host_mapping.close() except RuntimeError: pass + if not region_id: + # The failure may have landed after the child created its + # region. The reply shm is parent-owned and zero-filled and the + # child writes the id as part of committing, so a nonzero id + # here is the child saying the region exists. + with contextlib.suppress(BaseException): + region_id = peek_region_create_reply_region_id(reply_buf) + if not region_id and dispatched and not isinstance(exc, Exception): + # An asynchronous unwind through the create. The child may still + # be finishing a region, and the id it would write lands in a + # reply this frame is about to unlink — leaving something on the + # chip that nothing here can name. An ordinary failure is not + # this case: the child releases its own region before reporting + # one, so a zero id there really does mean nothing exists. + self._record_unreclaimable( + f"create_l3_l2_region: interrupted on worker {int(worker_id)} before the region id was " + "read back; a region may be live on the chip and no further work is admitted", + exc, + ) if region_id: try: assert self._worker is not None self._worker.control_l3_l2_region_release(int(worker_id), int(region_id)) - except RuntimeError: - pass - raise + except BaseException as release_exc: # noqa: BLE001 + # The chip created the region and this call could not give + # it back. Nothing else knows the id — it was never tracked + # — so no later cleanup can reclaim it, and a rollback that + # swallows this reports an ordinary failure over a region + # that stays allocated for the worker's life. That is the + # unreclaimed-device-state condition, recorded here directly + # because there is no handle for a fence to fail on. + raise self._record_unreclaimable( + f"create_l3_l2_region: rollback could not release region {region_id} on worker " + f"{int(worker_id)}; it is leaked and no further work is admitted", + release_exc, + ) + raise exc finally: del req_buf del reply_buf @@ -5292,36 +5480,56 @@ def _allocate_domain( # noqa: PLR0912 -- linear input-validation + per-chip shm reply_shms[chip_idx] = SharedMemory(create=True, size=reply_size) - self._dispatch_control_domain( - workers=workers, - request_shms=request_shms, - reply_shms=reply_shms, - op="alloc", - allocation_id=allocation_id, - ) + # A chip that committed its window holds it whether or not this call + # returns a handle, and whether or not its RPC reported success — + # the window exists before the reply is written. Each chip says so + # itself, and the chips that did are registered under a handle the + # run owns, so the fence sweep releases exactly those and poisons + # the worker if that release fails too. Inferring it from the RPC + # outcome leaks an allocation whose reply could not be delivered. + try: + self._dispatch_control_domain( + workers=workers, + request_shms=request_shms, + reply_shms=reply_shms, + op="alloc", + allocation_id=allocation_id, + ) - contexts: dict[int, ChipDomainContext] = {} - for chip_idx in workers: - reply_buf = reply_shms[chip_idx].buf - assert reply_buf is not None - (device_ctx, local_window_base, reply_buffer_count) = _DOMAIN_REPLY_HEADER.unpack_from(reply_buf, 0) - if reply_buffer_count != buffer_count: - raise RuntimeError( - f"allocate_domain: chip {chip_idx} reply buffer_count={reply_buffer_count} " - f"!= requested {buffer_count}" + contexts: dict[int, ChipDomainContext] = {} + for chip_idx in workers: + reply_buf = reply_shms[chip_idx].buf + assert reply_buf is not None + (committed, device_ctx, local_window_base, reply_buffer_count) = _DOMAIN_REPLY_HEADER.unpack_from( + reply_buf, 0 ) - ptrs: list[int] = [] - if buffer_count: - ptrs = list(struct.unpack_from(f"<{buffer_count}Q", reply_buf, _DOMAIN_REPLY_HEADER.size)) - contexts[chip_idx] = ChipDomainContext( - name=name, - domain_rank=worker_to_rank[chip_idx], - domain_size=len(workers), - device_ctx=int(device_ctx), - local_window_base=int(local_window_base), - actual_window_size=int(window_size), - buffer_ptrs={b.name: ptrs[i] for i, b in enumerate(buffers)}, + if not committed: + raise RuntimeError(f"allocate_domain: chip {chip_idx} reported success without committing") + if reply_buffer_count != buffer_count: + raise RuntimeError( + f"allocate_domain: chip {chip_idx} reply buffer_count={reply_buffer_count} " + f"!= requested {buffer_count}" + ) + ptrs: list[int] = [] + if buffer_count: + ptrs = list(struct.unpack_from(f"<{buffer_count}Q", reply_buf, _DOMAIN_REPLY_HEADER.size)) + contexts[chip_idx] = ChipDomainContext( + name=name, + domain_rank=worker_to_rank[chip_idx], + domain_size=len(workers), + device_ctx=int(device_ctx), + local_window_base=int(local_window_base), + actual_window_size=int(window_size), + buffer_ptrs={b.name: ptrs[i] for i, b in enumerate(buffers)}, + ) + except BaseException: + self._register_unclaimed_domain( + name, + tuple(w for w in workers if _domain_reply_committed(reply_shms.get(w))), + allocation_id, + resources, ) + raise finally: # Close + unlink local copies regardless of outcome. Children # have already finished reading by the time CONTROL_DONE fires. @@ -5348,6 +5556,9 @@ def _allocate_domain( # noqa: PLR0912 -- linear input-validation + per-chip shm self._live_domains[name] = handle if resources is not None: resources.live_domains[name] = handle + # Releasing this domain is a collective on every member chip, so a + # successor cannot be admitted until it has run. + resources.requires_ordered_cleanup = True # The backend windows are now live: record each chip's window base and # every carved buffer pointer so a later kind4 (child_memory) dispatch of # one of them is validated against its owning chip. Revoked by @@ -5363,6 +5574,30 @@ def _allocate_domain( # noqa: PLR0912 -- linear input-validation + per-chip shm self._child_prov_record_domain(chip_idx, int(buf_ptr), allocation_id, buf_nbytes[buf_name]) return handle + def _register_unclaimed_domain( + self, name: str, committed: tuple[int, ...], allocation_id: int, resources: _RunResources + ) -> None: + """Give a failed allocation an owner, so its committed windows are reclaimed. + + Only the chips that committed are listed: driving CTRL_RELEASE_DOMAIN + at a chip that never allocated this id would fail on a debt it does not + hold, and poison the worker for a partial failure it handled correctly. + The handle carries no contexts because there is nothing usable to hand + a caller — release needs only the member list and the allocation id. + """ + if not committed: + return + handle = CommDomainHandle( + name=name, + workers=committed, + contexts={}, + allocation_id=allocation_id, + _release_fn=lambda released, owner=resources: self._release_domain_handle(released, owner), + ) + self._live_domains[name] = handle + resources.live_domains[name] = handle + resources.requires_ordered_cleanup = True + def _release_domain_handle(self, handle: CommDomainHandle, resources: _RunResources) -> None: """Mark a handle for release. Actual backend free is deferred. @@ -5405,8 +5640,7 @@ def _retire_run_domains(self, resources: _RunResources) -> None: resources.retired = True stragglers = list(resources.pending_release_domains) resources.pending_release_domains.clear() - for handle in stragglers: - self._free_domain_after_fence(handle) + _raise_first(self._free_domain_after_fence, stragglers) def _free_domain_after_fence(self, handle: CommDomainHandle) -> None: """Back-end free for a handle whose owning run has retired. @@ -5414,45 +5648,36 @@ def _free_domain_after_fence(self, handle: CommDomainHandle) -> None: Tolerates a handle the end-of-run sweep already freed: the sweep works from a snapshot, so a release taken after that snapshot reaches here for a domain that is already gone. + + A failed release propagates. Its run carries + ``requires_ordered_cleanup`` — that is set when the domain is created — + so the error is what stops a successor from starting on top of device + state nobody can describe. """ if handle.freed: return - try: - self._release_domain_now(handle) - handle._freed = True # noqa: SLF001 -- runtime owns this transition - except Exception as e: # noqa: BLE001 - sys.stderr.write( - f"Worker._free_domain_after_fence: {handle.name!r} " - f"(allocation_id={handle.allocation_id}) released after its run's fence, " - f"immediate free failed: {type(e).__name__}: {e}\n" - ) - sys.stderr.flush() + self._release_domain_now(handle) + handle._freed = True # noqa: SLF001 -- runtime owns this transition def _execute_pending_domain_releases(self, resources: _RunResources) -> None: """Drive CTRL_RELEASE_DOMAIN for every queued handle. Must run after ``self._orch._wait_run()`` so chip-side tasks have completed their use of the domain memory. + + Per-handle best-effort: one failing release never strands the rest, and + the first error is raised once all are attempted. """ pending_releases = resources.pending_release_domains if not pending_releases: return pending = list(pending_releases) pending_releases.clear() - for handle in pending: - try: - self._release_domain_now(handle) - handle._freed = True # noqa: SLF001 -- runtime owns this transition - except Exception as e: # noqa: BLE001 - # A failed release should not block other handles' frees or - # the rest of Worker.run() shutdown. Drop from any residual - # tracking and log; the kernel-side memory may have already - # been reclaimed by the device_ctx-owning chip's finalize. - sys.stderr.write( - f"Worker._execute_pending_domain_releases: {handle.name!r} " - f"(allocation_id={handle.allocation_id}) failed: " - f"{type(e).__name__}: {e}\n" - ) - sys.stderr.flush() + + def _release(handle: CommDomainHandle) -> None: + self._release_domain_now(handle) + handle._freed = True # noqa: SLF001 -- runtime owns this transition + + _raise_first(_release, pending) def _release_domain_now(self, handle: CommDomainHandle) -> None: """Synchronous backend release for one handle, exactly once. @@ -5566,10 +5791,28 @@ def dispatch(chip_idx: int) -> None: errors[chip_idx] = e threads = [threading.Thread(target=dispatch, args=(w,), name=f"{op}_domain_chip_{w}") for w in workers] - for t in threads: - t.start() - for t in threads: - t.join() + started: list[threading.Thread] = [] + interrupt: BaseException | None = None + try: + for t in threads: + t.start() + started.append(t) + finally: + # Every started thread is joined even on an asynchronous unwind. The + # caller unlinks the request and reply shms on its way out, and a + # dispatch thread still inside control_* is still reading them. An + # async BaseException is delivered to this frame once, so one retry + # per join always completes it. + for t in started: + for _ in range(2): + try: + t.join() + break + except BaseException as exc: # noqa: BLE001, PERF203 + if interrupt is None: + interrupt = exc + if interrupt is not None: + raise interrupt if errors: first = next(iter(errors.items())) @@ -5585,29 +5828,27 @@ def _release_all_live_domains(self, resources: _RunResources | None = None) -> N releases) and from ``Worker.close``. Skips the deferred-release queue because by the time this runs, drain has already happened — synchronous release of leftover handles is safe. Falls back to - immediate backend free + drop from ``_live_domains`` on each handle; - logs and moves on if one fails. + immediate backend free + drop from ``_live_domains`` on each handle. + + Every handle is attempted and the first error is raised once they all + have been. A failed handle stays in ``_live_domains`` so the leak stays + detectable: close() reports it as a terminal residual instead of + returning success (terminal — it is not retried). """ live_domains = self._live_domains if resources is None else resources.live_domains - for handle in list(live_domains.values())[::-1]: - try: - # Mark released first (flips handle._released so further - # indexing raises), then synchronously free. The handle is - # not in the deferred-release queue, so we use the direct path. - if not handle.released: - handle._released = True # noqa: SLF001 -- runtime owns the transition - self._release_domain_now(handle) - handle._freed = True # noqa: SLF001 - if live_domains.get(handle.name) is handle: - live_domains.pop(handle.name) - except Exception as e: # noqa: BLE001 - sys.stderr.write( - f"Worker._release_all_live_domains: {handle.name!r} release failed: {type(e).__name__}: {e}\n" - ) - sys.stderr.flush() - # Keep the un-freed handle in _live_domains so the leak stays - # detectable: close() reports it as a terminal residual instead - # of returning success (terminal — it is not retried). + + def _release(handle: CommDomainHandle) -> None: + # Mark released first (flips handle._released so further indexing + # raises), then synchronously free. The handle is not in the + # deferred-release queue, so we use the direct path. + if not handle.released: + handle._released = True # noqa: SLF001 -- runtime owns the transition + self._release_domain_now(handle) + handle._freed = True # noqa: SLF001 + if live_domains.get(handle.name) is handle: + live_domains.pop(handle.name) + + _raise_first(_release, list(live_domains.values())[::-1]) # ------------------------------------------------------------------ # memory management — forward to C++ Orchestrator, which holds @@ -6170,9 +6411,14 @@ def submit(self, callable, args=None, config=None) -> RunHandle: ``args`` : TaskArgs (optional) ``config``: CallConfig (optional, default-constructed if None) - Graph construction remains serialized. A later submission waits until - prior dispatches are accepted; completion and cleanup stay attached to - each returned handle. + Graph construction remains serialized. How many runs may be admitted is + the depth this worker's backends negotiated, not a constant: at the + negotiated depth two, one active plus one prepared run are permitted and + a third submission blocks before invoking its graph callback; where a + backend publishes depth one — A5, or any runtime without a depth-two + contract — the *second* submission already blocks there. A caller whose + first run only completes because a later callback runs would deadlock on + such a backend. Completion and cleanup stay attached to each handle. """ with self._operation_lease("submit"): return self._submit_locked(callable, args, config) @@ -6199,14 +6445,118 @@ def _submit_locked(self, callable, args, config) -> RunHandle: return RunHandle._completed(self) with self._submit_mu: - # Graph callbacks are serialized, but accepted runs may remain live: - # their Python resources are isolated in each RunHandle. - with self._hierarchical_start_cv: - prior_handles = tuple(self._accepted_run_handles) - for handle in prior_handles: - handle._wait_for_acceptance() + # Graph callbacks stay serialized, so a predecessor's callback has + # always returned by the time we get here — which is what makes the + # decision below a fact about that run rather than a guess about + # this one. + # + # Native begin_run owns depth-two backpressure and slot generations + # for runs that only dispatch tasks. A predecessor whose cleanup + # itself touches the device gets stricter treatment: its teardown is + # mailbox control that the whole-run FIFO cannot order against this + # run's control, so we wait for that teardown and this worker + # degrades to depth one for exactly those runs. + predecessor = self._cleanup_bearing_predecessor() + if predecessor is not None: + predecessor._wait_for_handoff() + # Re-check under _submit_mu, after the wait. Two submissions can + # both pass the lease's admission check before either reaches here, + # and the poison may have been published by the cleanup this call + # just waited on — the check at lease time cannot have seen it. + self._require_no_ordered_cleanup_failure("submit") return self._submit_l3_locked(callable, args, cfg) + def _record_unreclaimable(self, message: str, cause: BaseException | None = None) -> RuntimeError: + """Refuse all further work on this worker, and return the reason. + + For device state that no cleanup can reach: nothing tracks it, so there + is no handle for a run's fence to fail on and no sweep that could + retry it. First-wins, like every other terminal error here. + """ + leaked = RuntimeError(message) + if cause is not None: + leaked.__cause__ = cause + with self._hierarchical_start_cv: + if self._ordered_cleanup_error is None: + self._ordered_cleanup_error = leaked + return leaked + + def _require_no_ordered_cleanup_failure(self, api: str) -> None: + """Refuse if a prior run's ordered cleanup failed.""" + with self._hierarchical_start_cv: + if self._ordered_cleanup_error is not None: + raise RuntimeError( + f"Worker.{api}: a prior run's ordered cleanup failed, so this worker's device state is " + "unreclaimed and no further work is admitted; close() it" + ) from self._ordered_cleanup_error + + def _control_admission(self, api: str): + """The direct-control ordering policy for a Worker-level command. + + Same policy the Orchestrator facade applies: a call inside a graph + callback belongs to that run and waits for the FIFO head; one that + belongs to no run reserves the worker for its duration. + """ + native = self._orch._o if self._orch is not None else None + return direct_control(self, native, f"Worker.{api}") + + @contextlib.contextmanager + def _control_reservation(self, api: str): + """Reserve this worker for one command that belongs to no run. + + A public `Worker.malloc/free/copy_*/remote_*` outside a graph callback + never becomes a task, so the whole-run FIFO cannot order it against one: + against a live run it can copy into a buffer that run is reading, or + free one it still holds. It is ordered only by being alone. + + "Alone" has to cover the command, not just the moment before it. A + sampled check leaves the caller free to send its mailbox command after a + submit admitted a run behind its back, so this takes the serializer + submission itself holds — no run can be admitted between the check and + the command. Re-entrant per thread: one control call may be built out of + others (a queue out of a region), and the second must join the + reservation rather than deadlock on it. + """ + held = _held_control_reservations() + if id(self) in held: + yield + return + with self._submit_mu: + with self._hierarchical_start_cv: + if self._ordered_cleanup_error is not None: + raise RuntimeError( + f"{api}: a prior run's ordered cleanup failed, so this worker's device state is " + "unreclaimed and no further work is admitted; close() it" + ) from self._ordered_cleanup_error + live = [h for h in self._accepted_run_handles if not h._cleanup_published] + if live: + raise RuntimeError( + f"{api}: {len(live)} run(s) still in flight. Device control that belongs to no " + "run is only ordered when nothing else is: wait on the outstanding handles first, " + "or issue it from inside a run's orchestration callback" + ) + held.add(id(self)) + try: + yield + finally: + held.discard(id(self)) + + def _cleanup_bearing_predecessor(self) -> RunHandle | None: + """A live handle whose ordered cleanup must finish before admission. + + At most one can be outstanding: graph callbacks are serialized, so a + run only acquires cleanup-bearing resources while it holds + ``_submit_mu``, and the next submission waits for it here. + """ + with self._hierarchical_start_cv: + for handle in self._accepted_run_handles: + # `done` answers the native fence, which fires before this + # run's device-touching cleanup runs — keying on it would admit + # a successor while that cleanup is still outstanding. + if handle._resources.requires_ordered_cleanup and not handle._cleanup_published: + return handle + return None + def _submit_l3_locked(self, callable, args, cfg: CallConfig) -> RunHandle: assert self._orch is not None assert self._worker is not None @@ -6222,7 +6572,8 @@ def _submit_l3_locked(self, callable, args, cfg: CallConfig) -> RunHandle: try: self._orch._scope_begin() scope_open = True - callable(self._orch, args, cfg) + with _callback_run(run_id, self): + callable(self._orch, args, cfg) scope_open = False self._orch._scope_end() self._orch._close_run_submission(run_id) @@ -6269,15 +6620,21 @@ def _finalize_run_handle( self, handle: RunHandle, run_id: int, native_error: BaseException | None ) -> BaseException | None: """Run fence-owned cleanup exactly once and return the cached result.""" - result = native_error + # Two different failures, deliberately not merged. A task that failed is + # this run's business and says nothing about the worker; a cleanup that + # failed leaves collective device state nobody can describe, and poisons + # the worker. Reporting a task failure as a cleanup failure would + # permanently poison a worker whose cleanup was fine. + cleanup_error: BaseException | None = None def _step(fn) -> None: - nonlocal result + """Run one cleanup step. Every step runs even if an earlier failed.""" + nonlocal cleanup_error try: fn() except BaseException as exc: # noqa: BLE001 - if result is None: - result = exc + if cleanup_error is None: + cleanup_error = exc resources = handle._resources if native_error is not None: @@ -6294,26 +6651,50 @@ def _step(fn) -> None: _step(lambda: self._retire_run_domains(resources)) orch = self._orch if orch is None: - if result is None: - result = RuntimeError("RunHandle cleanup lost its native Orchestrator") + if cleanup_error is None: + cleanup_error = RuntimeError("RunHandle cleanup lost its native Orchestrator") else: _step(lambda: orch._release_run(run_id)) + # Retirement is not optional: a handle left in the accepted set keeps # every later close() from completing its drain, and the lock is # mandatory (the drain and the submit serializer both snapshot the set # under it). An async BaseException in the interruptible acquire is # delivered to this frame once, so one retry always retires the handle. - try: + # Publish the poison and drop the handle under one acquisition, and in + # that order. A submitter that saw neither — handle already gone, poison + # not yet set — would find no cleanup-bearing predecessor and no reason + # to refuse, and would be admitted on top of unreclaimed device state. + # `_cleanup_published` is set here too: it is what a successor keys on, + # and it must not become true before the poison it implies. + def _publish() -> None: with self._hierarchical_start_cv: + if cleanup_error is not None and self._ordered_cleanup_error is None: + self._ordered_cleanup_error = cleanup_error + handle._cleanup_published = True self._accepted_run_handles.discard(handle) self._hierarchical_start_cv.notify_all() + + publish_error: BaseException | None = None + try: + _publish() except BaseException as exc: # noqa: BLE001 - if result is None: - result = exc - with self._hierarchical_start_cv: - self._accepted_run_handles.discard(handle) - self._hierarchical_start_cv.notify_all() - return result + # Retirement is not optional: a handle left in the accepted set + # keeps every later close() from completing its drain. An async + # BaseException in the interruptible acquire is delivered to this + # frame once, so one retry always publishes. It is reported as this + # run's result rather than raised — raising would escape into + # whichever waiter happened to run the cleanup. + publish_error = exc + _publish() + # Precedence: the run's own failure is what its waiters came for. A + # cleanup failure is reported when the run itself succeeded, since it is + # the reason the worker is now shut. + if native_error is not None: + return native_error + if cleanup_error is not None: + return cleanup_error + return publish_error @property def aicpu_dlopen_count(self) -> int: diff --git a/src/common/hierarchical/orchestrator.cpp b/src/common/hierarchical/orchestrator.cpp index 20a17493c..14c82be4e 100644 --- a/src/common/hierarchical/orchestrator.cpp +++ b/src/common/hierarchical/orchestrator.cpp @@ -11,6 +11,7 @@ #include "orchestrator.h" +#include #include #include #include @@ -61,32 +62,112 @@ std::shared_ptr Orchestrator::current_building_run() const { } RunId Orchestrator::begin_run() { - std::lock_guard lk(runs_mu_); - if (building_run_id_ != INVALID_RUN_ID) { - throw std::logic_error("Orchestrator::begin_run: another run is still building"); - } - if (next_run_id_ == INVALID_RUN_ID || next_run_id_ == std::numeric_limits::max()) { - throw std::overflow_error("Orchestrator::begin_run: run id space exhausted"); + RunId run_id = INVALID_RUN_ID; + { + std::unique_lock lk(runs_mu_); + if (building_run_id_ != INVALID_RUN_ID) { + throw std::logic_error("Orchestrator::begin_run: another run is still building"); + } + std::optional lease; + runs_cv_.wait(lk, [this, &lease] { + lease = pipeline_slots_.try_acquire(admission_depth_); + return lease.has_value(); + }); + if (next_run_id_ == INVALID_RUN_ID || next_run_id_ == std::numeric_limits::max()) { + pipeline_slots_.release(*lease); + throw std::overflow_error("Orchestrator::begin_run: run id space exhausted"); + } + run_id = next_run_id_++; + auto run = std::make_shared(run_id, *lease); + runs_.emplace(run_id, run); + run_fifo_.push_back(run_id); + building_run_id_ = run_id; + run->phase.store(RunPhase::BUILDING, std::memory_order_release); } - RunId run_id = next_run_id_++; - runs_.emplace(run_id, std::make_shared(run_id)); - building_run_id_ = run_id; + + // The FIFO head may execute while its graph callback is still building. + // Existing orchestration callbacks use this path to submit device work and + // wait for L2 communication before returning. A successor remains gated by + // active_run_id_ until the prior run reaches its terminal fence. + activate_fifo_head(); return run_id; } +void Orchestrator::configure_pipeline_depth(uint32_t depth) { + if (depth == 0 || depth > PTO_PIPELINE_MAX_DEPTH) { + throw std::invalid_argument("Orchestrator: pipeline depth is outside the supported range"); + } + std::lock_guard lk(runs_mu_); + if (!runs_.empty() || building_run_id_ != INVALID_RUN_ID || active_run_id_ != INVALID_RUN_ID) { + throw std::logic_error("Orchestrator: pipeline depth cannot change after admission starts"); + } + admission_depth_ = depth; +} + void Orchestrator::finish_run_if_ready(const std::shared_ptr &run) { bool notify = false; { std::lock_guard lk(run->completion_mu); - if (!run->submission_closed || run->active_tasks.load(std::memory_order_acquire) != 0 || - is_terminal(run->phase.load(std::memory_order_acquire))) { + RunPhase phase = run->phase.load(std::memory_order_acquire); + if (!run->submission_closed || run->active_tasks.load(std::memory_order_acquire) != 0 || is_terminal(phase) || + (phase != RunPhase::EXECUTING && !run->submission_failed)) { return; } bool failed = run->submission_failed || static_cast(run->first_error); run->phase.store(failed ? RunPhase::FAILED : RunPhase::COMPLETED, std::memory_order_release); notify = true; } - if (notify) run->completion_cv.notify_all(); + if (!notify) return; + run->completion_cv.notify_all(); + retire_terminal_run(run); +} + +void Orchestrator::clear_run_ready_queues(RunId run_id) { + if (ready_sub_queue_ != nullptr) ready_sub_queue_->erase_run(run_id); + if (ready_next_level_queues_ != nullptr) ready_next_level_queues_->erase_run(run_id); +} + +void Orchestrator::retire_terminal_run(const std::shared_ptr &run) { + { + std::lock_guard lk(runs_mu_); + if (active_run_id_ == run->id) active_run_id_ = INVALID_RUN_ID; + auto pos = std::find(run_fifo_.begin(), run_fifo_.end(), run->id); + if (pos != run_fifo_.end()) run_fifo_.erase(pos); + // The lease goes back under runs_mu_, not after it. begin_run evaluates + // "is a slot free" as its wait predicate while holding this mutex, so a + // release that lands outside it can fall between that evaluation and the + // waiter registering on the condition variable — the notify finds nobody + // and the free slot is stranded for good. At depth one that is the only + // slot there is. The pool takes its own mutex under this one, which is + // the order try_acquire already uses from inside the predicate. + if (!run->lease_released) { + run->lease_released = true; + pipeline_slots_.release(run->lease); + } + } + clear_run_ready_queues(run->id); + runs_cv_.notify_all(); + activate_fifo_head(); +} + +void Orchestrator::activate_fifo_head() { + std::shared_ptr run; + { + std::lock_guard lk(runs_mu_); + if (active_run_id_ != INVALID_RUN_ID || run_fifo_.empty()) return; + auto it = runs_.find(run_fifo_.front()); + if (it == runs_.end()) return; + run = it->second; + RunPhase phase = run->phase.load(std::memory_order_acquire); + if (phase != RunPhase::BUILDING && phase != RunPhase::PREPARED) return; + active_run_id_ = run->id; + run->phase.store(RunPhase::EXECUTING, std::memory_order_release); + } + // Direct device control waits on this, not only admission does: a prepared + // successor blocked in copy_to must wake when it becomes the active run. + runs_cv_.notify_all(); + if (ready_notify_cb_) ready_notify_cb_(); + finish_run_if_ready(run); } void Orchestrator::close_run_submission(RunId run_id) { @@ -104,16 +185,25 @@ void Orchestrator::close_run_submission(RunId run_id) { throw std::logic_error("Orchestrator::close_run_submission: submission already closed"); } run->submission_closed = true; - if (run->active_tasks.load(std::memory_order_acquire) != 0) { - run->phase.store(RunPhase::EXECUTING, std::memory_order_release); + if (run->phase.load(std::memory_order_acquire) == RunPhase::BUILDING) { + run->phase.store(RunPhase::PREPARED, std::memory_order_release); } } run->completion_cv.notify_all(); + activate_fifo_head(); finish_run_if_ready(run); } void Orchestrator::fail_run_submission(RunId run_id, std::exception_ptr error) { auto run = get_run(run_id); + std::string message = "graph construction failed"; + if (error) { + try { + std::rethrow_exception(error); + } catch (const std::exception &e) { + message = e.what(); + } catch (...) {} + } if (error) record_run_error(run_id, std::move(error)); { std::lock_guard runs_lk(runs_mu_); @@ -126,10 +216,11 @@ void Orchestrator::fail_run_submission(RunId run_id, std::exception_ptr error) { std::lock_guard lk(run->completion_mu); run->submission_failed = true; run->submission_closed = true; - if (run->active_tasks.load(std::memory_order_acquire) != 0) { - run->phase.store(RunPhase::EXECUTING, std::memory_order_release); - } } + // The FIFO head may already be executing while graph construction is + // open. Cancel only slots that have not started; concurrently running work + // remains part of the failed run's terminal fence. + cancel_unstarted_run(run, message); run->completion_cv.notify_all(); finish_run_if_ready(run); } @@ -190,6 +281,31 @@ bool Orchestrator::run_failed(RunId run_id) const { return run->submission_failed || static_cast(run->first_error); } +bool Orchestrator::dispatchable_locked(RunId run_id) const { + if (active_run_id_ != run_id) return false; + auto it = runs_.find(run_id); + return it != runs_.end() && it->second->phase.load(std::memory_order_acquire) == RunPhase::EXECUTING && + pipeline_slots_.owns(it->second->lease); +} + +bool Orchestrator::can_dispatch_run(RunId run_id) const { + std::lock_guard lk(runs_mu_); + return dispatchable_locked(run_id); +} + +RunId Orchestrator::dispatchable_run_id() const { + // The scheduler asks for a run to dispatch, not merely which run is at the + // FIFO head: a head whose lease has been released is not dispatchable, and + // answering with its id would send the scheduler looking for its tasks. + std::lock_guard lk(runs_mu_); + return dispatchable_locked(active_run_id_) ? active_run_id_ : INVALID_RUN_ID; +} + +RunId Orchestrator::active_run_id() const { + std::lock_guard lk(runs_mu_); + return active_run_id_; +} + bool Orchestrator::quiescent_locked() const { return runs_.empty() && building_run_id_ == INVALID_RUN_ID; } void Orchestrator::compact_if_quiescent() { @@ -218,8 +334,67 @@ void Orchestrator::release_run(RunId run_id) { compact_if_quiescent(); } -void Orchestrator::increment_run_tasks(RunId run_id) { - get_run(run_id)->active_tasks.fetch_add(1, std::memory_order_relaxed); +void Orchestrator::register_run_slot(const std::shared_ptr &run, TaskSlot slot) { + std::lock_guard lk(run->completion_mu); + run->task_slots.push_back(slot); + run->active_tasks.fetch_add(1, std::memory_order_relaxed); +} + +void Orchestrator::cancel_unstarted_run(const std::shared_ptr &run, const std::string &message) { + std::vector slots; + { + std::lock_guard lk(run->completion_mu); + slots.assign(run->task_slots.begin(), run->task_slots.end()); + } + + // Only the slots this call actually cancelled. A slot already RUNNING is + // owned by the device: it keeps reading its producers' outputs until it + // completes, and its normal completion path releases those references. + // + // A slot still BUILDING here is one whose submit threw part-way through + // wiring: submission is closed before this runs, so no submitting thread + // is left to publish it and this call owns its propagation. That is the + // opposite of the scheduler's poison path, which can hit a slot whose + // submit is still in flight and must leave it alone. + std::vector cancelled; + cancelled.reserve(slots.size()); + for (TaskSlot slot : slots) { + TaskSlotState &state = slot_state(slot); + std::optional claimed = claim_task_failure(state, message); + // A slot a producer already claimed mid-wiring is FAILED but not + // propagated: its submit was supposed to finish the job and threw + // instead. Taking that debt over is what keeps its references from + // stranding the run fence — skipping it for being FAILED would not. + // The debt is only observed here, not settled: everything up to the + // first reference release can throw, and a debt cleared before that + // point leaves a FAILED slot every later attempt skips. + bool owed = state.failure_propagation_pending.load(std::memory_order_acquire); + if (!claimed.has_value() && !owed) continue; + mark_group_members_skipped(state, message); + cancelled.push_back(slot); + } + // Every step that could still fail recoverably is done, and each of them is + // idempotent, so an exception above leaves the debt for a retry to resume + // from. From here the propagation is reference releases, which are not + // repeatable — so the debt is settled first. + for (TaskSlot slot : cancelled) + slot_state(slot).failure_propagation_pending.store(false, std::memory_order_release); + for (TaskSlot slot : cancelled) + try_consume(slot); + // Releasing a producer reference held by a slot that is still RUNNING would + // let the producer reach CONSUMED — and its HeapRing output be reclaimed — + // while the device is still reading it, and its real completion would then + // release the same reference a second time. + for (TaskSlot slot : cancelled) { + TaskSlotState &state = slot_state(slot); + std::vector producers; + { + std::lock_guard lk(state.fanout_mu); + producers = state.fanin_producers; + } + for (TaskSlot producer : producers) + try_consume(producer); + } } void Orchestrator::decrement_run_tasks(RunId run_id) { @@ -295,6 +470,27 @@ void Orchestrator::mark_task_accepted(TaskSlot slot) { decrement_run_accepts(task->run_id); } +void Orchestrator::await_run_admission(RunId run_id) { + // Direct device control — malloc / free / copy_* / domain and region + // creation — reaches a child without a TaskSlot, so the ready-queue FIFO + // does not order it and the mailbox mutex only serialises one command at a + // time. A prepared successor that freed or overwrote child memory the + // active run is still reading would break the whole-run ordering the FIFO + // exists to provide. + // + // The run is named by the caller rather than read from building_run_id_: + // that field says a graph callback is open *somewhere*, not that this + // thread is inside it, and a public Worker.copy_* on another thread must + // not be charged to whichever run happens to be building. + if (run_id == INVALID_RUN_ID) return; + std::unique_lock lk(runs_mu_); + runs_cv_.wait(lk, [this, run_id] { + if (active_run_id_ == run_id) return true; + auto it = runs_.find(run_id); + return it == runs_.end() || is_terminal(it->second->phase.load(std::memory_order_acquire)); + }); +} + uint64_t Orchestrator::malloc(int worker_id, size_t size) { auto *wt = manager_->get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); if (!wt) throw std::runtime_error("Orchestrator::malloc: invalid worker_id"); @@ -372,6 +568,8 @@ Tensor Orchestrator::alloc(const std::vector &shape, DataType dtype) { TaskSlotState &s = slot_state(ar.slot); s.reset(); s.run_id = run->id; + s.pipeline_lease = run->lease; + register_run_slot(run, ar.slot); uint64_t ptr = reinterpret_cast(ar.heap_ptr); if (ptr != 0) { @@ -381,7 +579,7 @@ Tensor Orchestrator::alloc(const std::vector &shape, DataType dtype) { } // No fanin — alloc has no work to wait on. - s.fanin_count = 0; + s.fanin_count.store(0, std::memory_order_relaxed); s.fanin_released.store(0, std::memory_order_relaxed); // Initial fanout_total = scope_ref. Consumers that wire on this slot @@ -400,8 +598,6 @@ Tensor Orchestrator::alloc(const std::vector &shape, DataType dtype) { s.state.store(TaskState::COMPLETED, std::memory_order_release); - increment_run_tasks(run->id); - // Build a contiguous external Tensor over the allocated buffer. ptr may be // 0 for a 0-byte request (a shape with a zero dim), in which case // init_external sets buffer.addr == 0 — the "no tensor" sentinel honored by @@ -485,6 +681,28 @@ SubmitResult Orchestrator::submit_impl( TaskSlotState &s = slot_state(slot); s.reset(); s.run_id = run->id; + s.pipeline_lease = run->lease; + // BUILDING until the publication at the end of this function. Step 2 + // registers this slot's outputs in the tensormap and Step 5 appends it to + // its producers' fanout lists, so it is observable well before its + // bookkeeping is final; BUILDING is what tells every other thread the + // counters below cannot be acted on yet. + s.state.store(TaskState::BUILDING, std::memory_order_release); + register_run_slot(run, slot); + // The scope reference is registered before it is charged, and both happen + // before anything below can throw. Charging first would leave a slot owing + // a release that no scope_end can make: cancellation contributes only the + // terminal release, the threshold is never reached, the slot never reaches + // CONSUMED, and the run's task count — and its fence — never resolve. + // fanout_total is charged here rather than later for the opposite reason: + // once Step 2 publishes this slot's outputs, a consumer can wire onto it + // and increment fanout_total, and a later assignment would drop that. + int32_t scope_ref = (scope_->depth() > 0) ? 1 : 0; + if (scope_ref > 0) scope_->register_task(slot); + { + std::lock_guard lk(s.fanout_mu); + s.fanout_total = scope_ref; + } s.worker_type = worker_type; s.callable = callable; @@ -541,26 +759,60 @@ SubmitResult Orchestrator::submit_impl( } } - s.fanin_count = live_fanins; - s.fanin_released.store(0, std::memory_order_relaxed); + increment_run_accepts(run->id, s.group_size()); - int32_t scope_ref = (scope_->depth() > 0) ? 1 : 0; - { + // --- Step 6: Publication — the one point where the slot leaves BUILDING. + // + // fanin_count and the transition out of BUILDING are published together, + // under the lock every completing producer takes to decide readiness. A + // producer that completed while this slot was BUILDING already advanced + // fanin_released and found nothing it could do with it; the comparison + // here is what picks that release up. A producer that arrives afterwards + // reads a count that is already published. Neither can act on one half. + bool ready_now = false; + if (!poisoned_by_failed_producer) { std::lock_guard lk(s.fanout_mu); - s.fanout_total = scope_ref; + s.fanin_count.store(live_fanins, std::memory_order_release); + TaskState expected = TaskState::BUILDING; + bool satisfied = s.fanin_released.load(std::memory_order_acquire) >= live_fanins; + TaskState published = satisfied ? TaskState::READY : TaskState::PENDING; + if (s.state.compare_exchange_strong( + expected, published, std::memory_order_acq_rel, std::memory_order_acquire + )) { + ready_now = satisfied; + } else { + // Only a failure claim moves a BUILDING slot, and it stops there + // precisely so this thread — the one that knows the wiring is + // final — runs the propagation. + poisoned_by_failed_producer = true; + } } - s.fanout_released.store(0, std::memory_order_relaxed); - - if (scope_ref > 0) scope_->register_task(slot); - increment_run_tasks(run->id); - increment_run_accepts(run->id, s.group_size()); if (poisoned_by_failed_producer) { - if (poison_message.empty()) poison_message = "producer task failed"; - s.failure_message = poison_message; - record_run_error(run->id, std::make_exception_ptr(std::runtime_error(poison_message))); - s.state.store(TaskState::FAILED, std::memory_order_release); + // Publishing our own failure goes through the shared claim, so exactly + // one party writes failure_message. Losing the claim to a producer that + // failed while we were wiring is fine: its message stands, and the + // propagation is ours either way. + (void)claim_task_failure(s, poison_message.empty() ? "producer task failed" : poison_message); + std::string message; + { + std::lock_guard lk(s.fanout_mu); + message = s.failure_message; + } + if (message.empty()) message = "producer task failed"; + // Ordered so that everything which can throw runs while the propagation + // debt is still owed, and every one of those steps is idempotent: group + // marking skips already-terminal members, the error is first-wins, and + // the producer list is a copy. An exception here therefore leaves a + // FAILED slot whose debt run cancellation still sees, and it resumes + // from the top. Settling the debt first would leave that slot FAILED + // with nothing recorded, and cancellation skips exactly those. + mark_group_members_skipped(s, message); + record_run_error(run->id, std::make_exception_ptr(std::runtime_error(message))); std::vector fanin_producers = s.fanin_producers; + // Past the resumable part: the releases below are the propagation, and + // repeating one of them would release a reference twice. + s.failure_propagation_pending.store(false, std::memory_order_release); try_consume(slot); for (TaskSlot prod : fanin_producers) { try_consume(prod); @@ -568,13 +820,12 @@ SubmitResult Orchestrator::submit_impl( return SubmitResult{slot}; } - // --- Step 6: If no live fanins → READY and route by final placement --- - if (live_fanins == 0) { - s.state.store(TaskState::READY, std::memory_order_release); + // Enqueued outside the publication lock: enqueue_ready takes runs_mu_, and + // a producer completing between the two finds the slot already READY and + // does nothing, so the queue receives it exactly once. + if (ready_now) { enqueue_ready(slot); if (ready_notify_cb_) ready_notify_cb_(); - } else { - s.state.store(TaskState::PENDING, std::memory_order_release); } return SubmitResult{slot}; @@ -582,17 +833,28 @@ SubmitResult Orchestrator::submit_impl( void Orchestrator::enqueue_ready(TaskSlot slot) { TaskSlotState &s = slot_state(slot); + if (s.worker_type == WorkerType::NEXT_LEVEL && ready_next_level_queues_ == nullptr) { + throw std::runtime_error("Orchestrator::enqueue_ready: NEXT_LEVEL queues are not initialized"); + } + // retire_terminal_run erases this run's partitions, and the scheduler only + // ever reads the partition of the run holding the FIFO head — so a push + // that lands after the erase builds one nothing will drain. Holding + // runs_mu_ across both the test and the push is what makes them one + // decision: releasing it in between lets retire erase the partition we + // just decided was live. Every requeue route comes through here for the + // same reason. + std::lock_guard lk(runs_mu_); + auto it = runs_.find(s.run_id); + if (it == runs_.end() || is_terminal(it->second->phase.load(std::memory_order_acquire))) return; if (s.worker_type == WorkerType::NEXT_LEVEL) { - if (ready_next_level_queues_ == nullptr) - throw std::runtime_error("Orchestrator::enqueue_ready: NEXT_LEVEL queues are not initialized"); if (s.is_group()) { - ready_next_level_queues_->push_group(slot); + ready_next_level_queues_->push_group(s.run_id, slot); } else { - ready_next_level_queues_->push_single(s.target_worker_id(0), slot); + ready_next_level_queues_->push_single(s.target_worker_id(0), s.run_id, slot); } return; } - ready_sub_queue_->push(slot); + ready_sub_queue_->push(s.run_id, slot); } void Orchestrator::validate_worker_eligibility( diff --git a/src/common/hierarchical/orchestrator.h b/src/common/hierarchical/orchestrator.h index 5461b6491..d575d44ed 100644 --- a/src/common/hierarchical/orchestrator.h +++ b/src/common/hierarchical/orchestrator.h @@ -33,6 +33,7 @@ #include #include #include +#include #include #include #include @@ -43,6 +44,7 @@ #include "../task_interface/data_type.h" #include "../task_interface/task_args.h" #include "../task_interface/tensor.h" +#include "../worker/pipeline_slot_pool.h" #include "ring.h" #include "scope.h" #include "tensormap.h" @@ -119,6 +121,7 @@ class Orchestrator { // Only the calling orchestration thread builds a run at a time. RunId begin_run(); + void configure_pipeline_depth(uint32_t depth); void close_run_submission(RunId run_id); void fail_run_submission(RunId run_id, std::exception_ptr error = nullptr); void wait_run_accepted(RunId run_id); @@ -127,6 +130,27 @@ class Orchestrator { bool wait_run_for(RunId run_id, double timeout_seconds); bool run_done(RunId run_id) const; bool run_failed(RunId run_id) const; + bool can_dispatch_run(RunId run_id) const; + + /** + * The run the scheduler may dispatch from, or INVALID_RUN_ID when none is. + * + * Stricter than `active_run_id()`: the FIFO head must also be EXECUTING and + * still own its pipeline lease, so a head whose lease was released cannot + * have further work handed down under it. + */ + RunId dispatchable_run_id() const; + + /** + * Block until `run_id` holds the whole-run FIFO head, or is terminal. + * + * Direct device control bypasses the ready-queue FIFO, so a caller acting + * on behalf of a run must wait here before touching a child. The run is + * explicit: only the thread actually inside that run's graph callback may + * be ordered against it. INVALID_RUN_ID returns immediately. + */ + void await_run_admission(RunId run_id); + RunId active_run_id() const; void release_run(RunId run_id); // Open a nested scope. Every task submitted between this call and the @@ -177,9 +201,14 @@ class Orchestrator { NextLevelReadyQueues *ready_next_level_queues_ = nullptr; mutable std::mutex runs_mu_; + std::condition_variable runs_cv_; std::unordered_map> runs_; + std::deque run_fifo_; + PipelineSlotPool pipeline_slots_{PTO_PIPELINE_MAX_DEPTH}; + uint32_t admission_depth_{PTO_PIPELINE_MAX_DEPTH}; RunId next_run_id_{1}; RunId building_run_id_{INVALID_RUN_ID}; + RunId active_run_id_{INVALID_RUN_ID}; // Scheduler's loop mutex (not owned). Held across optional quiescent // compaction so the scheduler cannot retain a slot pointer being removed. @@ -191,13 +220,18 @@ class Orchestrator { std::shared_ptr find_run(RunId run_id) const; std::shared_ptr get_run(RunId run_id) const; std::shared_ptr current_building_run() const; - static void finish_run_if_ready(const std::shared_ptr &run); + void finish_run_if_ready(const std::shared_ptr &run); static bool is_terminal(RunPhase phase); static bool acceptance_ready(const std::shared_ptr &run); // Callers hold runs_mu_. bool quiescent_locked() const; + bool dispatchable_locked(RunId run_id) const; + void activate_fifo_head(); + void retire_terminal_run(const std::shared_ptr &run); + void cancel_unstarted_run(const std::shared_ptr &run, const std::string &message); + void register_run_slot(const std::shared_ptr &run, TaskSlot slot); + void clear_run_ready_queues(RunId run_id); void compact_if_quiescent(); - void increment_run_tasks(RunId run_id); void decrement_run_tasks(RunId run_id); void increment_run_accepts(RunId run_id, int32_t count); void decrement_run_accepts(RunId run_id); diff --git a/src/common/hierarchical/scheduler.cpp b/src/common/hierarchical/scheduler.cpp index 8949b6aac..aef1838b4 100644 --- a/src/common/hierarchical/scheduler.cpp +++ b/src/common/hierarchical/scheduler.cpp @@ -11,6 +11,7 @@ #include "scheduler.h" +#include #include #include @@ -175,8 +176,15 @@ void Scheduler::run() { { std::unique_lock lk(completion_mu_); completion_cv_.wait(lk, [this] { - return !completion_queue_.empty() || !cfg_.ready_next_level_queues->empty() || - !cfg_.ready_sub_queue->empty() || stop_requested_.load(std::memory_order_acquire); + bool ready = false; + if (cfg_.active_run_cb) { + RunId active = cfg_.active_run_cb(); + ready = active != INVALID_RUN_ID && + (!cfg_.ready_next_level_queues->empty(active) || !cfg_.ready_sub_queue->empty(active)); + } else { + ready = !cfg_.ready_next_level_queues->empty() || !cfg_.ready_sub_queue->empty(); + } + return !completion_queue_.empty() || ready || stop_requested_.load(std::memory_order_acquire); }); } @@ -229,17 +237,20 @@ void Scheduler::on_task_complete(const WorkerCompletion &completion) { TaskSlot slot = completion.task_slot; TaskSlotState &s = *cfg_.ring->slot_state(slot); bool failed = is_failure(completion.outcome); - if (failed) { - s.failure_message = completion.error_message; - s.state.store(TaskState::FAILED, std::memory_order_release); - } else { - s.state.store(TaskState::COMPLETED, std::memory_order_release); - } - // Release fanin on downstream consumers + // The transition and the consumer snapshot are one decision, taken under + // fanout_mu. A task wiring itself onto this slot holds the same lock while + // it appends itself and reads the state, so it either lands in `consumers` + // — and is released or poisoned below — or sees the terminal state and + // handles it in submit. Splitting the two lets a consumer register as a + // live fanin after the snapshot and never be released, or observe + // COMPLETED, decline to count the fanin, and still be released here, which + // makes it READY one producer early. std::vector consumers; { std::lock_guard lk(s.fanout_mu); + if (failed) s.failure_message = completion.error_message; + s.state.store(failed ? TaskState::FAILED : TaskState::COMPLETED, std::memory_order_release); consumers = s.fanout_consumers; } for (TaskSlot consumer : consumers) { @@ -248,13 +259,12 @@ void Scheduler::on_task_complete(const WorkerCompletion &completion) { continue; } TaskSlotState &cs = *cfg_.ring->slot_state(consumer); - int32_t released = cs.fanin_released.fetch_add(1, std::memory_order_acq_rel) + 1; - if (released >= cs.fanin_count) { - TaskState expected = TaskState::PENDING; - if (cs.state.compare_exchange_strong(expected, TaskState::READY, std::memory_order_acq_rel)) { - cfg_.enqueue_ready_cb(consumer); - completion_cv_.notify_one(); - } + cs.fanin_released.fetch_add(1, std::memory_order_acq_rel); + // A consumer still BUILDING is not readiable, and this release is not + // lost: submit's publication compares the pair under the same lock. + if (try_mark_ready(cs)) { + cfg_.enqueue_ready_cb(consumer); + completion_cv_.notify_one(); } } @@ -273,33 +283,17 @@ void Scheduler::on_task_complete(const WorkerCompletion &completion) { void Scheduler::poison_task(TaskSlot slot, const std::string &root_message) { TaskSlotState &s = *cfg_.ring->slot_state(slot); - TaskState state = s.state.load(std::memory_order_acquire); - if (state == TaskState::FAILED || state == TaskState::CONSUMED || state == TaskState::FREE) return; - if (state == TaskState::RUNNING || state == TaskState::COMPLETED) return; - - s.failure_message = root_message; - s.state.store(TaskState::FAILED, std::memory_order_release); - - if (s.is_group()) { - std::lock_guard lk(s.group_mu); - const int32_t group_size = s.group_size(); - if (s.group_member_states.size() != static_cast(group_size)) { - s.group_member_states.assign(static_cast(group_size), GroupMemberState::NOT_DISPATCHED); - s.group_member_outcomes.assign(static_cast(group_size), EndpointOutcome::SKIPPED); - } - if (!s.group_failed) { - s.group_failed = true; - s.group_first_failure_index = -1; - s.group_first_failure_message = root_message; - } - for (int32_t i = 0; i < group_size; ++i) { - GroupMemberState &member_state = s.group_member_states[static_cast(i)]; - if (is_terminal_group_state(member_state)) continue; - member_state = GroupMemberState::SKIPPED; - s.group_member_outcomes[static_cast(i)] = EndpointOutcome::SKIPPED; - s.group_terminal_count.fetch_add(1, std::memory_order_acq_rel); - } - } + std::optional claimed = claim_task_failure(s, root_message); + // Not claimable: RUNNING or COMPLETED is owned by the device until its own + // completion arrives, and FAILED/CONSUMED/FREE is already someone else's — + // resurrecting a CONSUMED slot would release its dependency references a + // second time. + if (!claimed.has_value()) return; + // Claimed while submit was still wiring it: submit finishes the + // propagation, using the counters only it can know are final. + if (*claimed == TaskState::BUILDING) return; + + mark_group_members_skipped(s, root_message); std::vector consumers; { @@ -354,11 +348,24 @@ void Scheduler::dispatch_ready() { dispatch_sub_ready(); } +bool claim_for_dispatch(TaskSlotState &s) { + TaskState expected = TaskState::READY; + return s.state.compare_exchange_strong( + expected, TaskState::RUNNING, std::memory_order_acq_rel, std::memory_order_acquire + ); +} + void Scheduler::dispatch_sub_ready() { + RunId active_run = cfg_.active_run_cb ? cfg_.active_run_cb() : INVALID_RUN_ID; + if (cfg_.active_run_cb && active_run == INVALID_RUN_ID) return; TaskSlot slot; - while (cfg_.ready_sub_queue->try_pop(slot)) { + while (cfg_.active_run_cb ? cfg_.ready_sub_queue->try_pop(active_run, slot) : cfg_.ready_sub_queue->try_pop(slot)) { TaskSlotState &s = *cfg_.ring->slot_state(slot); if (s.state.load(std::memory_order_acquire) != TaskState::READY) continue; + if (cfg_.active_run_cb && s.run_id != active_run) { + cfg_.enqueue_ready_cb(slot); + return; + } if (s.worker_type != WorkerType::SUB) { throw std::runtime_error("Scheduler::dispatch_sub_ready: misrouted task slot"); } @@ -369,13 +376,17 @@ void Scheduler::dispatch_sub_ready() { for (int32_t i = 0; i < group_size; ++i) { WorkerThread *worker = cfg_.manager->pick_idle_sub_excluding(workers); if (worker == nullptr) { - cfg_.ready_sub_queue->push(slot); + // Through the same entry point as every other requeue: pushing + // straight onto the queue skips the check that this run is + // still live, and would rebuild a partition retire has erased. + cfg_.enqueue_ready_cb(slot); return; } workers.push_back(worker); } - s.state.store(TaskState::RUNNING, std::memory_order_release); + if (cfg_.before_claim_cb) cfg_.before_claim_cb(slot); + if (!claim_for_dispatch(s)) continue; if (s.is_group()) { reset_group_state(s, group_size, GroupMemberState::NOT_DISPATCHED); } @@ -392,14 +403,27 @@ void Scheduler::dispatch_sub_ready() { } std::unordered_set Scheduler::dispatch_next_level_group() { + RunId active_run = cfg_.active_run_cb ? cfg_.active_run_cb() : INVALID_RUN_ID; + if (cfg_.active_run_cb && active_run == INVALID_RUN_ID) return {}; TaskSlot slot; - while (cfg_.ready_next_level_queues->try_front_group(slot)) { + while (cfg_.active_run_cb ? cfg_.ready_next_level_queues->try_front_group(active_run, slot) : + cfg_.ready_next_level_queues->try_front_group(slot)) { TaskSlotState &s = *cfg_.ring->slot_state(slot); if (s.state.load(std::memory_order_acquire) != TaskState::READY) { TaskSlot stale; - cfg_.ready_next_level_queues->try_pop_group(stale); + if (cfg_.active_run_cb) { + cfg_.ready_next_level_queues->try_pop_group(active_run, stale); + } else { + cfg_.ready_next_level_queues->try_pop_group(stale); + } continue; } + if (cfg_.active_run_cb && s.run_id != active_run) { + TaskSlot misplaced; + cfg_.ready_next_level_queues->try_pop_group(active_run, misplaced); + cfg_.enqueue_ready_cb(slot); + return {}; + } if (s.worker_type != WorkerType::NEXT_LEVEL || !s.is_group()) { throw std::runtime_error("Scheduler::dispatch_next_level_group: misrouted task slot"); } @@ -424,13 +448,25 @@ std::unordered_set Scheduler::dispatch_next_level_group() { } if (!all_workers_idle) return target_worker_ids; + // The head was observed before the worker checks above, so a run + // cancelling in that window can consume the slot and erase its whole + // partition. That is legal, and throwing here would end the process: + // this runs on the scheduler thread, which has no handler. TaskSlot popped; - if (!cfg_.ready_next_level_queues->try_pop_group(popped) || popped != slot) { - throw std::runtime_error("Scheduler::dispatch_next_level_group: group queue changed unexpectedly"); + bool popped_ok = cfg_.active_run_cb ? cfg_.ready_next_level_queues->try_pop_group(active_run, popped) : + cfg_.ready_next_level_queues->try_pop_group(popped); + if (!popped_ok) return {}; + if (popped != slot) { + // A different group reached the head meanwhile. Put back what we + // took and re-decide from the top rather than dispatching a group + // whose workers we never checked. + cfg_.enqueue_ready_cb(popped); + return {}; } + if (cfg_.before_claim_cb) cfg_.before_claim_cb(slot); + if (!claim_for_dispatch(s)) continue; reset_group_state(s, group_size, GroupMemberState::RUNNING); - s.state.store(TaskState::RUNNING, std::memory_order_release); for (int32_t i = 0; i < group_size; ++i) { workers[static_cast(i)]->dispatch(WorkerDispatch{slot, i}); } @@ -439,6 +475,8 @@ std::unordered_set Scheduler::dispatch_next_level_group() { } void Scheduler::dispatch_next_level_singles(const std::unordered_set &reserved_worker_ids) { + RunId active_run = cfg_.active_run_cb ? cfg_.active_run_cb() : INVALID_RUN_ID; + if (cfg_.active_run_cb && active_run == INVALID_RUN_ID) return; for (int32_t worker_id : cfg_.ready_next_level_queues->worker_ids()) { if (reserved_worker_ids.find(worker_id) != reserved_worker_ids.end()) continue; @@ -451,13 +489,19 @@ void Scheduler::dispatch_next_level_singles(const std::unordered_set &r if (!worker->idle()) continue; TaskSlot slot; - while (cfg_.ready_next_level_queues->try_pop_single(worker_id, slot)) { + while (cfg_.active_run_cb ? cfg_.ready_next_level_queues->try_pop_single(worker_id, active_run, slot) : + cfg_.ready_next_level_queues->try_pop_single(worker_id, slot)) { TaskSlotState &s = *cfg_.ring->slot_state(slot); if (s.state.load(std::memory_order_acquire) != TaskState::READY) continue; + if (cfg_.active_run_cb && s.run_id != active_run) { + cfg_.enqueue_ready_cb(slot); + break; + } if (s.worker_type != WorkerType::NEXT_LEVEL || s.is_group() || s.target_worker_id(0) != worker_id) { throw std::runtime_error("Scheduler::dispatch_next_level_singles: misrouted task slot"); } - s.state.store(TaskState::RUNNING, std::memory_order_release); + if (cfg_.before_claim_cb) cfg_.before_claim_cb(slot); + if (!claim_for_dispatch(s)) continue; worker->dispatch(WorkerDispatch{slot, 0}); break; } diff --git a/src/common/hierarchical/scheduler.h b/src/common/hierarchical/scheduler.h index 9eec55be1..bafc56b33 100644 --- a/src/common/hierarchical/scheduler.h +++ b/src/common/hierarchical/scheduler.h @@ -51,6 +51,22 @@ class Ring; // forward decl // Scheduler — DAG engine (no worker pool ownership) // ============================================================================= +/** + * Take ownership of a READY slot for dispatch, atomically. + * + * Reading READY and later storing RUNNING are not the same thing. A run whose + * graph callback throws fails its own unstarted slots and consumes them, and + * its fence can then release the run's pipeline lease — all while the scheduler + * sits between those two points picking workers. A plain store would overwrite + * that cancelled state and dispatch a task whose run is already terminal and + * whose slot may have been reused. A failed claim means the slot is no longer + * ours: whoever moved it out of READY owns its consume. + * + * Call this at the point of dispatch, after every other admission check — the + * window this closes is exactly the code between the queue pop and the launch. + */ +bool claim_for_dispatch(TaskSlotState &s); + class Scheduler { public: struct Config { @@ -60,11 +76,21 @@ class Scheduler { WorkerManager *manager; // not owned — Scheduler calls manager for dispatch // Shared READY routing path owned by Orchestrator. std::function enqueue_ready_cb; + // Production workers expose exactly one whole-run FIFO head. Tests + // that omit this callback retain the legacy unpartitioned queue path. + std::function active_run_cb; // Called when a task reaches CONSUMED (TensorMap cleanup + ring release). std::function on_consumed_cb; // Called as soon as an endpoint reports failure so the error is // attached to the task's run even when a group has other members live. std::function on_task_failed_cb; + // Test seam. Invoked immediately before the dispatch claim, which is + // the one instant a cancelling run can still take a slot away. The + // window is unreachable from outside — every other observable point is + // either before the pop or after the launch — so a test that wants to + // exercise the losing side of the claim has to be let in here. + // Unset in production. + std::function before_claim_cb; }; void start(const Config &cfg); diff --git a/src/common/hierarchical/types.cpp b/src/common/hierarchical/types.cpp index 3fcebcd85..073ed59ad 100644 --- a/src/common/hierarchical/types.cpp +++ b/src/common/hierarchical/types.cpp @@ -11,6 +11,7 @@ #include "types.h" +#include #include // ============================================================================= @@ -20,8 +21,10 @@ void TaskSlotState::reset() { state.store(TaskState::FREE, std::memory_order_relaxed); run_id = INVALID_RUN_ID; - fanin_count = 0; + pipeline_lease = PipelineSlotLease{}; + fanin_count.store(0, std::memory_order_relaxed); fanin_released.store(0, std::memory_order_relaxed); + failure_propagation_pending.store(false, std::memory_order_relaxed); { std::lock_guard lk(fanout_mu); fanout_consumers.clear(); @@ -55,35 +58,141 @@ void TaskSlotState::reset() { group_terminal_count.store(0, std::memory_order_relaxed); } +bool try_mark_ready(TaskSlotState &s) { + std::lock_guard lk(s.fanout_mu); + if (s.state.load(std::memory_order_acquire) != TaskState::PENDING) return false; + if (s.fanin_released.load(std::memory_order_acquire) < s.fanin_count.load(std::memory_order_acquire)) return false; + s.state.store(TaskState::READY, std::memory_order_release); + return true; +} + +std::optional claim_task_failure(TaskSlotState &s, const std::string &message) { + std::lock_guard lk(s.fanout_mu); + TaskState current = s.state.load(std::memory_order_acquire); + // Retry rather than attempt once: a producer completing concurrently moves + // its consumer PENDING -> READY without holding fanout_mu, and a single + // failed exchange would leave that slot dispatchable under a run that is + // already failing. The loop ends as soon as the slot leaves the claimable + // set, which is either our own FAILED or a dispatch that beat us to it. + while (current == TaskState::PENDING || current == TaskState::READY || current == TaskState::BUILDING) { + if (s.state.compare_exchange_strong( + current, TaskState::FAILED, std::memory_order_acq_rel, std::memory_order_acquire + )) { + s.failure_message = message; + // The submitting thread owns the propagation for a slot claimed + // mid-wiring; this records that it is owed, so a submit that throws + // before running it leaves the debt visible to run cancellation + // rather than a FAILED slot nobody will finish. + if (current == TaskState::BUILDING) { + s.failure_propagation_pending.store(true, std::memory_order_release); + } + return current; + } + } + return std::nullopt; +} + +void mark_group_members_skipped(TaskSlotState &s, const std::string &message) { + if (!s.is_group()) return; + std::lock_guard lk(s.group_mu); + const int32_t group_size = s.group_size(); + if (s.group_member_states.size() != static_cast(group_size)) { + s.group_member_states.assign(static_cast(group_size), GroupMemberState::NOT_DISPATCHED); + s.group_member_outcomes.assign(static_cast(group_size), EndpointOutcome::SKIPPED); + } + if (!s.group_failed) { + s.group_failed = true; + s.group_first_failure_index = -1; + s.group_first_failure_message = message; + } + for (int32_t i = 0; i < group_size; ++i) { + GroupMemberState &member_state = s.group_member_states[static_cast(i)]; + if (member_state == GroupMemberState::SUCCESS || member_state == GroupMemberState::FAILED || + member_state == GroupMemberState::SKIPPED) { + continue; + } + member_state = GroupMemberState::SKIPPED; + s.group_member_outcomes[static_cast(i)] = EndpointOutcome::SKIPPED; + s.group_terminal_count.fetch_add(1, std::memory_order_acq_rel); + } +} + // ============================================================================= // ReadyQueue // ============================================================================= -void ReadyQueue::push(TaskSlot slot) { +void ReadyQueue::push(TaskSlot slot) { push(INVALID_RUN_ID, slot); } + +void ReadyQueue::push(RunId run_id, TaskSlot slot) { std::lock_guard lk(mu_); - q_.push(slot); + auto [it, inserted] = queues_.try_emplace(run_id); + if (inserted) run_order_.push_back(run_id); + it->second.push(slot); } bool ReadyQueue::try_pop(TaskSlot &out) { std::lock_guard lk(mu_); - if (q_.empty()) return false; - out = q_.front(); - q_.pop(); + if (run_order_.empty()) return false; + auto it = queues_.find(run_order_.front()); + if (it == queues_.end() || it->second.empty()) throw std::logic_error("ReadyQueue: corrupt run order"); + out = it->second.front(); + it->second.pop(); + if (it->second.empty()) { + queues_.erase(it); + run_order_.pop_front(); + } + return true; +} + +bool ReadyQueue::try_pop(RunId run_id, TaskSlot &out) { + std::lock_guard lk(mu_); + auto it = queues_.find(run_id); + if (it == queues_.end() || it->second.empty()) return false; + out = it->second.front(); + it->second.pop(); + if (it->second.empty()) { + queues_.erase(it); + auto order_it = std::find(run_order_.begin(), run_order_.end(), run_id); + if (order_it != run_order_.end()) run_order_.erase(order_it); + } return true; } bool ReadyQueue::empty() const { std::lock_guard lk(mu_); - return q_.empty(); + return run_order_.empty(); +} + +bool ReadyQueue::empty(RunId run_id) const { + std::lock_guard lk(mu_); + auto it = queues_.find(run_id); + return it == queues_.end() || it->second.empty(); } bool ReadyQueue::try_front(TaskSlot &out) { std::lock_guard lk(mu_); - if (q_.empty()) return false; - out = q_.front(); + if (run_order_.empty()) return false; + auto it = queues_.find(run_order_.front()); + if (it == queues_.end() || it->second.empty()) throw std::logic_error("ReadyQueue: corrupt run order"); + out = it->second.front(); return true; } +bool ReadyQueue::try_front(RunId run_id, TaskSlot &out) { + std::lock_guard lk(mu_); + auto it = queues_.find(run_id); + if (it == queues_.end() || it->second.empty()) return false; + out = it->second.front(); + return true; +} + +void ReadyQueue::erase_run(RunId run_id) { + std::lock_guard lk(mu_); + queues_.erase(run_id); + auto it = std::find(run_order_.begin(), run_order_.end(), run_id); + if (it != run_order_.end()) run_order_.erase(it); +} + // ============================================================================= // NextLevelReadyQueues // ============================================================================= @@ -114,15 +223,26 @@ size_t NextLevelReadyQueues::index_for(int32_t worker_id) const { void NextLevelReadyQueues::push_single(int32_t worker_id, TaskSlot slot) { queues_[index_for(worker_id)]->push(slot); } +void NextLevelReadyQueues::push_single(int32_t worker_id, RunId run_id, TaskSlot slot) { + queues_[index_for(worker_id)]->push(run_id, slot); +} + bool NextLevelReadyQueues::try_pop_single(int32_t worker_id, TaskSlot &out) { return queues_[index_for(worker_id)]->try_pop(out); } +bool NextLevelReadyQueues::try_pop_single(int32_t worker_id, RunId run_id, TaskSlot &out) { + return queues_[index_for(worker_id)]->try_pop(run_id, out); +} + void NextLevelReadyQueues::push_group(TaskSlot slot) { group_queue_.push(slot); } +void NextLevelReadyQueues::push_group(RunId run_id, TaskSlot slot) { group_queue_.push(run_id, slot); } bool NextLevelReadyQueues::try_front_group(TaskSlot &out) { return group_queue_.try_front(out); } +bool NextLevelReadyQueues::try_front_group(RunId run_id, TaskSlot &out) { return group_queue_.try_front(run_id, out); } bool NextLevelReadyQueues::try_pop_group(TaskSlot &out) { return group_queue_.try_pop(out); } +bool NextLevelReadyQueues::try_pop_group(RunId run_id, TaskSlot &out) { return group_queue_.try_pop(run_id, out); } bool NextLevelReadyQueues::empty() const { if (!group_queue_.empty()) return false; @@ -131,3 +251,17 @@ bool NextLevelReadyQueues::empty() const { } return true; } + +bool NextLevelReadyQueues::empty(RunId run_id) const { + if (!group_queue_.empty(run_id)) return false; + for (const auto &queue : queues_) { + if (!queue->empty(run_id)) return false; + } + return true; +} + +void NextLevelReadyQueues::erase_run(RunId run_id) { + group_queue_.erase_run(run_id); + for (const auto &queue : queues_) + queue->erase_run(run_id); +} diff --git a/src/common/hierarchical/types.h b/src/common/hierarchical/types.h index cdce0f120..c25a3b502 100644 --- a/src/common/hierarchical/types.h +++ b/src/common/hierarchical/types.h @@ -32,16 +32,20 @@ #include #include #include +#include #include #include #include #include +#include #include #include +#include #include #include "../task_interface/call_config.h" #include "../task_interface/task_args.h" +#include "../worker/pto_runtime_c_api.h" // ============================================================================= // TensorKey — compound key for TensorMap dependency tracking @@ -123,30 +127,35 @@ static constexpr int32_t INVALID_SLOT = -1; using RunId = uint64_t; static constexpr RunId INVALID_RUN_ID = 0; -// No transition assigns PREPARED while one run builds at a time: closing -// submission moves BUILDING straight to EXECUTING, or to a terminal phase when -// the run holds no tasks. Bounded prepared-run admission is what occupies it. +// Admission reserves a generation-safe pipeline slot before graph construction. +// Closing the graph publishes PREPARED; only the whole-run FIFO head may enter +// EXECUTING and reach a device endpoint. enum class RunPhase : int32_t { - BUILDING = 0, - PREPARED = 1, - EXECUTING = 2, - COMPLETED = 3, - FAILED = 4, + RESERVED = 0, + BUILDING = 1, + PREPARED = 2, + EXECUTING = 3, + COMPLETED = 4, + FAILED = 5, }; struct RunState { - explicit RunState(RunId run_id) : - id(run_id) {} + RunState(RunId run_id, PipelineSlotLease slot_lease) : + id(run_id), + lease(slot_lease) {} RunId id{INVALID_RUN_ID}; - std::atomic phase{RunPhase::BUILDING}; + PipelineSlotLease lease{}; + std::atomic phase{RunPhase::RESERVED}; std::atomic active_tasks{0}; std::atomic pending_accepts{0}; mutable std::mutex completion_mu; std::condition_variable completion_cv; std::exception_ptr first_error; + std::vector task_slots; bool submission_closed{false}; bool submission_failed{false}; + bool lease_released{false}; }; // ============================================================================= @@ -196,6 +205,14 @@ enum class TaskState : int32_t { COMPLETED = 4, // worker finished, outputs may still be referenced FAILED = 5, // worker failed or a producer poisoned this slot CONSUMED = 6, // all references released, slot may be reused + // Submit owns the slot. It is already reachable from its producers' fanout + // lists — wiring has to happen under each producer's fanout_mu, so the slot + // cannot be kept private until it is finished — but fanin_count, + // fanout_total and the ref counters are not final yet. No other thread may + // advance a BUILDING slot: a producer that completes leaves its fanin + // release for submit to observe, and a producer that fails parks the slot + // at FAILED for submit to propagate. See claim_task_failure(). + BUILDING = 7, }; enum class EndpointOutcome : int32_t { @@ -322,13 +339,28 @@ struct WorkerCompletion { struct TaskSlotState { std::atomic state{TaskState::FREE}; RunId run_id{INVALID_RUN_ID}; - - // --- Fanin (orch writes once; scheduler reads atomically) --- - int32_t fanin_count{0}; + PipelineSlotLease pipeline_lease{}; + + // --- Fanin --- + // `fanin_count` is written once by submit's publication and read by every + // completing producer. The two are concurrent — a producer can complete + // while the slot is still BUILDING — so both the write and every read live + // under `fanout_mu`, which also carries the state transition they decide + // on. See try_mark_ready() for why they cannot be separated. + std::atomic fanin_count{0}; std::atomic fanin_released{0}; // incremented by each completing producer - // --- Fanout (protected by fanout_mu) --- - // orch adds consumers; scheduler traverses on completion + // Set when a failure claim won this slot while it was BUILDING and left the + // propagation to its submitting thread. Cleared by whoever runs that + // propagation. It survives a submit that throws before it gets there, which + // is what lets run cancellation take the slot over instead of skipping it + // for being FAILED already. + std::atomic failure_propagation_pending{false}; + + // --- Fanout, plus the slot's pre-dispatch state transitions --- + // orch adds consumers; scheduler traverses on completion. The same mutex + // covers every BUILDING/PENDING -> {READY, FAILED} transition and the + // fields each of those decisions reads. std::mutex fanout_mu; std::vector fanout_consumers; int32_t fanout_total{0}; // 1 (scope ref) + fanout_consumers.size() @@ -351,6 +383,11 @@ struct TaskSlotState { // --- Failure state --- // Root worker failures and downstream poison both land here. The // originating RunState owns first-error-wins reporting. + // + // Guarded by fanout_mu, the same lock that guards the fanout snapshot a + // failure propagates over: whoever claims the failure writes the message + // and reads the consumer list under one acquisition, so a consumer that + // sees FAILED while wiring itself in also sees the reason. std::string failure_message; // --- Task data (stored on parent heap, lives until slot CONSUMED) --- @@ -416,6 +453,48 @@ struct TaskSlotState { void reset(); }; +// The one claim every failure path uses to move a task to FAILED: a device +// completion poisoning its consumers, a run cancellation, and a submit that +// wires onto a producer that has already failed. Returns the state the slot +// held when the claim was won, or std::nullopt when the slot was not claimable +// — already terminal, or RUNNING and therefore owned by the device until its +// own completion arrives. +// +// The claim and failure_message are published under fanout_mu, the lock a +// propagation snapshots the consumer list under. That is what makes the +// snapshot exact: a task wiring itself onto this producer either registers +// before the claim and is poisoned by it, or observes FAILED — with its +// message — and poisons itself. +// +// A claim won from BUILDING stops there. The submitting thread still owns that +// slot's fanin/fanout bookkeeping, so it completes the propagation itself; a +// caller that touched the reference counters would race with the wiring that is +// still in flight. Every other claim owns the full propagation. +std::optional claim_task_failure(TaskSlotState &s, const std::string &message); + +// Records every not-yet-terminal member of a group slot as SKIPPED. A no-op on +// a single-task slot. Called by whoever owns a failure's propagation, so a +// group that never dispatched still reports one terminal outcome per member. +void mark_group_members_skipped(TaskSlotState &s, const std::string &message); + +// The single PENDING -> READY transition, taken by submit's publication and by +// every completing producer. Returns true for the one caller that made it, who +// then owns enqueuing the slot. +// +// "Every live producer has released" is not two independently readable facts. +// `fanin_count` is published by submit and `fanin_released` is advanced by +// producers, so a producer that reads a count submit has not published yet — +// zero, which any release passes — and acts on the state afterwards would +// launch a task whose remaining producers are still running. Comparing the +// pair and changing the state under one acquisition is what ties the count a +// caller judges to the state it changes. +// +// Neither side can lose the other's half: submit takes the same lock to +// publish the count together with the transition out of BUILDING, so a +// producer that arrives first is re-evaluated by submit, and one that arrives +// after sees a published count. +bool try_mark_ready(TaskSlotState &s); + // ============================================================================= // ReadyQueue — Orch pushes, Scheduler pops // ============================================================================= @@ -423,17 +502,24 @@ struct TaskSlotState { class ReadyQueue { public: void push(TaskSlot slot); + void push(RunId run_id, TaskSlot slot); // Non-blocking: returns false immediately if empty. bool try_pop(TaskSlot &out); + bool try_pop(RunId run_id, TaskSlot &out); bool empty() const; + bool empty(RunId run_id) const; // Non-blocking: copies the front without removing it. bool try_front(TaskSlot &out); + bool try_front(RunId run_id, TaskSlot &out); + + void erase_run(RunId run_id); private: - std::queue q_; + std::unordered_map> queues_; + std::deque run_order_; mutable std::mutex mu_; }; @@ -443,11 +529,18 @@ class NextLevelReadyQueues { public: void reset(const std::vector &worker_ids); void push_single(int32_t worker_id, TaskSlot slot); + void push_single(int32_t worker_id, RunId run_id, TaskSlot slot); bool try_pop_single(int32_t worker_id, TaskSlot &out); + bool try_pop_single(int32_t worker_id, RunId run_id, TaskSlot &out); void push_group(TaskSlot slot); + void push_group(RunId run_id, TaskSlot slot); bool try_front_group(TaskSlot &out); + bool try_front_group(RunId run_id, TaskSlot &out); bool try_pop_group(TaskSlot &out); + bool try_pop_group(RunId run_id, TaskSlot &out); bool empty() const; + bool empty(RunId run_id) const; + void erase_run(RunId run_id); const std::vector &worker_ids() const { return worker_ids_; } private: diff --git a/src/common/hierarchical/worker.cpp b/src/common/hierarchical/worker.cpp index 90efac32a..288d8c177 100644 --- a/src/common/hierarchical/worker.cpp +++ b/src/common/hierarchical/worker.cpp @@ -122,6 +122,9 @@ void Worker::init() { cfg.enqueue_ready_cb = [this](TaskSlot slot) { orchestrator_.enqueue_ready(slot); }; + cfg.active_run_cb = [this] { + return orchestrator_.dispatchable_run_id(); + }; cfg.on_consumed_cb = [this](TaskSlot slot) { orchestrator_.on_consumed(slot); }; diff --git a/src/common/hierarchical/worker.h b/src/common/hierarchical/worker.h index f8e2e82b0..e6969f01c 100644 --- a/src/common/hierarchical/worker.h +++ b/src/common/hierarchical/worker.h @@ -92,6 +92,11 @@ class Worker { // otherwise be accidentally inherited across fork. void init(); + void configure_pipeline_depth(uint32_t depth) { + if (initialized_) throw std::logic_error("Worker: configure_pipeline_depth after init"); + orchestrator_.configure_pipeline_depth(depth); + } + // Shut down the Scheduler thread and release resources. void close(); diff --git a/src/common/hierarchical/worker_manager.cpp b/src/common/hierarchical/worker_manager.cpp index ff34e4b6e..68ccd90c5 100644 --- a/src/common/hierarchical/worker_manager.cpp +++ b/src/common/hierarchical/worker_manager.cpp @@ -386,6 +386,7 @@ WorkerCompletion LocalMailboxEndpoint::run_with_accept( // Write config as a single packed POD block (see call_config.h). std::memcpy(mbox() + MAILBOX_OFF_CONFIG, &s.config, sizeof(CallConfig)); + std::memcpy(mbox() + MAILBOX_OFF_PIPELINE_LEASE, &s.pipeline_lease, sizeof(PipelineSlotLease)); // Write length-prefixed TaskArgs blob: [T][S][tensors][scalars]. size_t blob_bytes = TASK_ARGS_BLOB_HEADER_SIZE + static_cast(view.tensor_count) * sizeof(Tensor) + diff --git a/src/common/hierarchical/worker_manager.h b/src/common/hierarchical/worker_manager.h index 622bcbf0a..2f678fbbb 100644 --- a/src/common/hierarchical/worker_manager.h +++ b/src/common/hierarchical/worker_manager.h @@ -92,7 +92,8 @@ static constexpr size_t MAILBOX_ERROR_MSG_SIZE = 256; // CallConfig is written/read as a single packed POD block (see call_config.h). // Both ends transfer it with one memcpy — no per-field offsets to keep in sync. // -// MAILBOX_OFF_ARGS is derived: round up CallConfig's end to 8 bytes so the +// The generation-safe run lease follows CallConfig. MAILBOX_OFF_ARGS is +// derived by rounding up the lease's end so the // args blob's first Tensor field (buffer.addr, a uint64_t at OFF_ARGS+8) is // 8-byte aligned, avoiding SIGBUS on strict-alignment platforms (aarch64 // atomics, some ARM cores). The control region (CTRL_OFF_ARG0..CTRL_OFF_RESULT) lives @@ -102,13 +103,19 @@ static constexpr ptrdiff_t MAILBOX_OFF_STATE = 0; static constexpr ptrdiff_t MAILBOX_OFF_ERROR = 4; static constexpr ptrdiff_t MAILBOX_OFF_CALLABLE = 8; // also: control sub-command (uint64) static constexpr ptrdiff_t MAILBOX_OFF_CONFIG = 16; -static constexpr ptrdiff_t MAILBOX_OFF_ARGS = +static constexpr ptrdiff_t MAILBOX_OFF_PIPELINE_LEASE = (MAILBOX_OFF_CONFIG + static_cast(sizeof(CallConfig)) + 7) & ~ptrdiff_t{7}; +static constexpr ptrdiff_t MAILBOX_OFF_ARGS = + (MAILBOX_OFF_PIPELINE_LEASE + static_cast(sizeof(PipelineSlotLease)) + 7) & ~ptrdiff_t{7}; static_assert(MAILBOX_OFF_ARGS % 8 == 0, "MAILBOX_OFF_ARGS must be 8-aligned for Tensor.buffer.addr"); static_assert( MAILBOX_OFF_CONFIG + static_cast(sizeof(CallConfig)) <= MAILBOX_OFF_ARGS, "CallConfig overflows reserved config region" ); +static_assert( + MAILBOX_OFF_PIPELINE_LEASE + static_cast(sizeof(PipelineSlotLease)) <= MAILBOX_OFF_ARGS, + "PipelineSlotLease overflows reserved lease region" +); static constexpr ptrdiff_t MAILBOX_OFF_ERROR_MSG = static_cast(MAILBOX_SIZE) - static_cast(MAILBOX_ERROR_MSG_SIZE); // Launch acceptance is sticky, not a MailboxState: a state word carrying it diff --git a/src/common/worker/pipeline_slot_pool.h b/src/common/worker/pipeline_slot_pool.h index 66f18a176..59c1cb08f 100644 --- a/src/common/worker/pipeline_slot_pool.h +++ b/src/common/worker/pipeline_slot_pool.h @@ -37,9 +37,14 @@ class PipelineSlotPool { } } - std::optional try_acquire() { + std::optional try_acquire() { return try_acquire(depth_); } + + std::optional try_acquire(uint32_t admission_depth) { + if (admission_depth == 0 || admission_depth > depth_) { + throw std::invalid_argument("pipeline admission depth is outside the pool range"); + } std::lock_guard lock(mu_); - for (uint32_t slot = 0; slot < depth_; ++slot) { + for (uint32_t slot = 0; slot < admission_depth; ++slot) { SlotState &state = slots_[slot]; if (state.in_use) continue; if (state.generation == std::numeric_limits::max()) { diff --git a/tests/st/a2a3/host_build_graph/worker_async_fifo/test_worker_async_fifo.py b/tests/st/a2a3/host_build_graph/worker_async_fifo/test_worker_async_fifo.py new file mode 100644 index 000000000..a58d23aa7 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/worker_async_fifo/test_worker_async_fifo.py @@ -0,0 +1,337 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Onboard validation for bounded whole-run FIFO admission. + +The first run completes real NPU work but remains active behind a SubTask +fence. The second run builds its graph into the other pipeline slot and must +not dispatch until the first run becomes terminal. A third submission must +block before its graph callback while both slots are admitted. +""" + +import multiprocessing +import threading +import time +from contextlib import suppress + +import pytest +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler_setup.scene_test import _build_l3_task_args + +_VECTOR_KERNELS = "../vector_example/kernels" +_SIZE = 128 * 128 +_SUB_ENTERED = multiprocessing.Event() +_SUB_RELEASE = multiprocessing.Event() + + +def _wait_for_release(_args): + _SUB_ENTERED.set() + if not _SUB_RELEASE.wait(30.0): + raise RuntimeError("whole-run FIFO test timed out waiting for the release fence") + + +@scene_test(level=3, runtime="host_build_graph") +class TestWorkerAsyncWholeRunFifo(SceneTestCase): + """A prepared run may build ahead but cannot dispatch ahead.""" + + CALLABLE = { + "callables": [ + { + "name": "vector", + "orchestration": { + "source": f"{_VECTOR_KERNELS}/orchestration/example_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": f"{_VECTOR_KERNELS}/aiv/kernel_add.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "source": f"{_VECTOR_KERNELS}/aiv/kernel_add_scalar.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT], + }, + { + "func_id": 2, + "source": f"{_VECTOR_KERNELS}/aiv/kernel_mul.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + ], + }, + {"name": "wait_for_release", "callable": _wait_for_release}, + ], + } + + CASES = [ + { + "name": "whole_run_fifo", + "platforms": ["a2a3"], + "config": {"device_count": 1, "num_sub_workers": 1, "aicpu_thread_num": 4}, + "params": {}, + }, + ] + + @staticmethod + def _tensor_from_host_buffer(worker, value): + buffer = worker.create_host_buffer(_SIZE * torch.float32.itemsize) + tensor = torch.frombuffer(buffer.buffer, dtype=torch.float32, count=_SIZE) + tensor.fill_(value) + return buffer, tensor + + def test_prepared_run_device_control_waits_for_the_active_run(self, st_platform, st_worker): + """A prepared run's malloc/free/copy must not overtake the active run. + + These reach a child directly rather than through a TaskSlot, so the + ready-queue FIFO does not order them and the mailbox mutex only + serialises one command at a time. Without an admission gate a prepared + successor could free or overwrite child memory the active run is still + reading, which is the ordering the whole-run FIFO exists to provide. + """ + if st_platform != "a2a3": + pytest.skip("whole-run FIFO leases require an a2a3 onboard worker") + + _SUB_ENTERED.clear() + _SUB_RELEASE.clear() + entered_callback = threading.Event() + control_returned = threading.Event() + result = {} + buffers = [] + submitter = None + first = None + try: + tensors = [] + for value in (2.0, 3.0, 0.0): + buffer, tensor = self._tensor_from_host_buffer(st_worker, value) + buffers.append(buffer) + tensors.append(tensor) + first_a, first_b, first_out = tensors + + vector_handle = type(self)._st_chip_handles["vector"] + vector_signature = type(self)._st_chip_handles["vector_sig"] + sub_handle = type(self)._st_sub_handles["wait_for_release"] + + def first_graph(orch, _args, _cfg): + builder = TaskArgsBuilder(Tensor("a", first_a), Tensor("b", first_b), Tensor("f", first_out)) + chip_args, _ = _build_l3_task_args(builder, vector_signature) + orch.submit_next_level(vector_handle, chip_args, self._build_config(self.CASES[0]["config"]), worker=0) + orch.submit_sub(sub_handle) + + first = st_worker.submit(first_graph) + assert _SUB_ENTERED.wait(30.0), "the first run never reached its SUB fence" + + # The successor is admitted into the second slot, so its callback + # runs — but its first device-control call must block until the + # first run releases the FIFO head. + def second_graph(orch, _args, _cfg): + entered_callback.set() + ptr = orch.malloc(0, 4096) + control_returned.set() + orch.free(0, ptr) + + submitter = threading.Thread( + target=lambda: result.setdefault("handle", st_worker.submit(second_graph)), daemon=True + ) + submitter.start() + + assert entered_callback.wait(10.0), "the prepared run's graph callback never entered" + assert not control_returned.wait(1.0), ( + "a prepared run's device control returned while the active run still held the FIFO head" + ) + + _SUB_RELEASE.set() + first.wait(30.0) + assert control_returned.wait(30.0), "device control stayed blocked after the active run became terminal" + submitter.join(30.0) + assert not submitter.is_alive() + result["handle"].wait(30.0) + finally: + _SUB_RELEASE.set() + if submitter is not None: + submitter.join(30.0) + handles = [first, result.get("handle")] + for handle in handles: + if handle is not None: + with suppress(Exception): + handle.wait(30.0) + if all(handle is None or handle.done for handle in handles): + for buffer in buffers: + st_worker.free_host_buffer(buffer) + + def test_a_run_whose_cleanup_touches_the_device_degrades_to_depth_one(self, st_platform, st_worker): + """A CommDomain release is mailbox control the whole-run FIFO cannot order. + + The FIFO orders tasks. A domain's teardown happens after the native + fence and reaches a child through the mailbox, so N+1 allocating a + domain while N is still releasing one can leave two collectives each + holding a different chip's mailbox. A run that takes such a resource + therefore degrades this worker to depth one: the successor's graph + callback does not start at all until that teardown has run. + + The sibling test above pins the other half — a run that only dispatches + tasks keeps the full depth, and its successor's callback does run while + it is still fenced. + """ + if st_platform != "a2a3": + pytest.skip("whole-run FIFO leases require an a2a3 onboard worker") + + _SUB_ENTERED.clear() + _SUB_RELEASE.clear() + entered_callback = threading.Event() + result = {} + submitter = None + first = None + try: + sub_handle = type(self)._st_sub_handles["wait_for_release"] + + def first_graph(orch, _args, _cfg): + # Control before any submit: a task travels the ready queue and + # this travels the mailbox, so the two are not ordered. + with orch.allocate_domain(name="degrade", workers=[0], window_size=4096, buffers=[]): + pass + orch.submit_sub(sub_handle) + + first = st_worker.submit(first_graph) + assert _SUB_ENTERED.wait(30.0), "the first run never reached its SUB fence" + + def second_graph(_orch, _args, _cfg): + entered_callback.set() + + submitter = threading.Thread( + target=lambda: result.setdefault("handle", st_worker.submit(second_graph)), daemon=True + ) + submitter.start() + + assert not entered_callback.wait(2.0), ( + "a successor's graph callback started while a cleanup-bearing run was still outstanding" + ) + + _SUB_RELEASE.set() + first.wait(30.0) + assert entered_callback.wait(30.0), "the successor stayed blocked after its predecessor's cleanup ran" + submitter.join(30.0) + assert not submitter.is_alive() + result["handle"].wait(30.0) + finally: + _SUB_RELEASE.set() + if submitter is not None: + submitter.join(30.0) + for handle in (first, result.get("handle")): + if handle is not None: + with suppress(Exception): + handle.wait(30.0) + + def test_run(self, st_platform, st_worker): + if st_platform != "a2a3": + pytest.skip("whole-run FIFO leases require an a2a3 onboard worker") + + _SUB_ENTERED.clear() + _SUB_RELEASE.clear() + third_callback = threading.Event() + third_result = {} + buffers = [] + tensors = [] + submitter = None + first = None + second = None + try: + for value in (2.0, 3.0, 0.0, 5.0, 7.0, 0.0): + buffer, tensor = self._tensor_from_host_buffer(st_worker, value) + buffers.append(buffer) + tensors.append(tensor) + first_a, first_b, first_out, second_a, second_b, second_out = tensors + + vector_handle = type(self)._st_chip_handles["vector"] + vector_signature = type(self)._st_chip_handles["vector_sig"] + sub_handle = type(self)._st_sub_handles["wait_for_release"] + + def submit_vector(orch, a, b, out, *, hold_open=False): + builder = TaskArgsBuilder(Tensor("a", a), Tensor("b", b), Tensor("f", out)) + chip_args, _ = _build_l3_task_args(builder, vector_signature) + orch.submit_next_level(vector_handle, chip_args, self._build_config(self.CASES[0]["config"]), worker=0) + if hold_open: + orch.submit_sub(sub_handle) + + first = st_worker.submit( + lambda orch, _args, _cfg: submit_vector(orch, first_a, first_b, first_out, hold_open=True) + ) + assert _SUB_ENTERED.wait(10.0), "the first run's SubTask did not start" + + first_expected = (first_a + first_b + 1) * (first_a + first_b + 2) + deadline = time.monotonic() + 10.0 + while not torch.allclose(first_out, first_expected) and time.monotonic() < deadline: + time.sleep(0.001) + assert torch.allclose(first_out, first_expected), "the first run's NPU task did not complete" + + second_graph_done = threading.Event() + + def second_graph(orch, _args, _cfg): + submit_vector(orch, second_a, second_b, second_out) + second_graph_done.set() + + second = st_worker.submit(second_graph) + assert second_graph_done.is_set(), "the second run did not build ahead" + assert torch.count_nonzero(second_out).item() == 0, ( + "the prepared run dispatched before the active run ended" + ) + + def third_graph(_orch, _args, _cfg): + third_callback.set() + + submitter = threading.Thread( + target=lambda: third_result.setdefault("handle", st_worker.submit(third_graph)), daemon=True + ) + submitter.start() + assert not third_callback.wait(0.1), "the third graph callback entered before admission capacity was free" + + # The check above only proves the prepared run had not dispatched at + # one instant. Re-check after the hold: the active run is still + # blocked in its SUB task for the whole window, so a prepared run + # that leaks past the FIFO gate at any point during it is caught + # here rather than passing on timing. + assert torch.count_nonzero(second_out).item() == 0, ( + "the prepared run dispatched while the active run was still holding the FIFO head" + ) + + _SUB_RELEASE.set() + first.wait(10.0) + assert third_callback.wait(10.0), "the third submission did not enter after the first run freed its slot" + second.wait(10.0) + second_expected = (second_a + second_b + 1) * (second_a + second_b + 2) + assert torch.allclose(second_out, second_expected), "the prepared run did not execute correctly on the NPU" + + submitter.join(10.0) + assert not submitter.is_alive() + third_result["handle"].wait(10.0) + finally: + _SUB_RELEASE.set() + if submitter is not None: + submitter.join(10.0) + handles = [first, second, third_result.get("handle")] + for handle in handles: + if handle is not None: + with suppress(Exception): + handle.wait(10.0) + tensors.clear() + first_a = first_b = first_out = second_a = second_b = second_out = None + if all(handle is None or handle.done for handle in handles): + for buffer in buffers: + st_worker.free_host_buffer(buffer) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/ut/cpp/hierarchical/test_orchestrator.cpp b/tests/ut/cpp/hierarchical/test_orchestrator.cpp index 8dd4791c1..871d149e4 100644 --- a/tests/ut/cpp/hierarchical/test_orchestrator.cpp +++ b/tests/ut/cpp/hierarchical/test_orchestrator.cpp @@ -12,6 +12,8 @@ #include #include +#include +#include #include "call_config.h" #include "ring.h" @@ -42,6 +44,7 @@ struct OrchestratorFixture : public ::testing::Test { struct WorkerQueueView { NextLevelReadyQueues *queues; bool try_pop(TaskSlot &out) { return queues->try_pop_single(0, out); } + bool try_pop(RunId run_id, TaskSlot &out) { return queues->try_pop_single(0, run_id, out); } } rq{&rq_next_level}; void SetUp() override { @@ -79,6 +82,24 @@ struct OrchestratorFixture : public ::testing::Test { // Tests // --------------------------------------------------------------------------- +TEST(ReadyQueueTest, UnscopedAccessPreservesRunInsertionOrder) { + ReadyQueue queue; + queue.push(/*run_id=*/7, /*slot=*/70); + queue.push(/*run_id=*/8, /*slot=*/80); + queue.push(/*run_id=*/7, /*slot=*/71); + + TaskSlot slot = INVALID_SLOT; + ASSERT_TRUE(queue.try_front(slot)); + EXPECT_EQ(slot, 70); + ASSERT_TRUE(queue.try_pop(slot)); + EXPECT_EQ(slot, 70); + ASSERT_TRUE(queue.try_pop(slot)); + EXPECT_EQ(slot, 71); + ASSERT_TRUE(queue.try_pop(slot)); + EXPECT_EQ(slot, 80); + EXPECT_TRUE(queue.empty()); +} + TEST_F(OrchestratorFixture, IndependentTaskIsImmediatelyReady) { auto a = single_tensor_args(0xCAFE, TensorArgType::OUTPUT); auto res = orch.submit_next_level(C(42), a, cfg, 0); @@ -494,11 +515,7 @@ TEST_F(OrchestratorFixture, RunAcceptanceWaitsForEveryDispatchedGroupMember) { TEST_F(OrchestratorFixture, TerminalFailureUnblocksRunAcceptance) { auto result = orch.submit_next_level(C(80), single_tensor_args(0x8030, TensorArgType::OUTPUT), cfg, 0); orch.fail_run_submission(run_id, std::make_exception_ptr(std::runtime_error("graph build failed"))); - EXPECT_FALSE(orch.run_accepted(run_id)); - - S(result.task_slot).state.store(TaskState::FAILED, std::memory_order_release); - EXPECT_TRUE(orch.on_consumed(result.task_slot)); - + EXPECT_EQ(S(result.task_slot).state.load(std::memory_order_acquire), TaskState::CONSUMED); EXPECT_TRUE(orch.run_accepted(run_id)); EXPECT_NO_THROW(orch.wait_run_accepted(run_id)); } @@ -675,3 +692,338 @@ TEST_F(OrchestratorFixture, FailedSubmissionCarriesItsMessageToTheFence) { } orch.release_run(run_id); } + +// Cancelling a run must not release a producer reference that a RUNNING +// consumer still holds. Releasing it early lets the producer reach CONSUMED — +// and its HeapRing output be reclaimed — while the device is still reading it, +// and the consumer's real completion then releases the same reference twice. +TEST_F(OrchestratorFixture, CancelKeepsProducerRefsHeldByARunningConsumer) { + orch.scope_begin(); + auto producer = orch.submit_next_level(C(70), single_tensor_args(0x7000, TensorArgType::OUTPUT), cfg, 0); + auto consumer = orch.submit_next_level(C(71), single_tensor_args(0x7000, TensorArgType::INPUT), cfg, 0); + ASSERT_EQ(S(consumer.task_slot).fanin_producers.size(), 1u); + ASSERT_EQ(S(consumer.task_slot).fanin_producers[0], producer.task_slot); + + // The consumer is on the device when the graph callback throws. + S(consumer.task_slot).state.store(TaskState::RUNNING, std::memory_order_release); + const int32_t released_before = S(producer.task_slot).fanout_released.load(std::memory_order_acquire); + + orch.fail_run_submission(run_id, std::make_exception_ptr(std::runtime_error("orchestration: boom"))); + + EXPECT_EQ(S(consumer.task_slot).state.load(std::memory_order_acquire), TaskState::RUNNING) + << "a RUNNING slot is owned by the device and must not be cancelled"; + // Exactly one release: the producer's own terminal self-release, because + // the cancel path failed it. The reference the RUNNING consumer holds is + // NOT among them — releasing that one too would make this 2, and would let + // the producer be reclaimed under a consumer still reading it. + EXPECT_EQ(S(producer.task_slot).fanout_released.load(std::memory_order_acquire), released_before + 1) + << "cancellation released a producer reference still held by a RUNNING consumer"; + + // The run stays live on purpose: only the consumer's real completion can + // release the last reference, and that path belongs to the scheduler. + S(consumer.task_slot).state.store(TaskState::COMPLETED, std::memory_order_release); + (void)orch.on_consumed(consumer.task_slot); + orch.scope_end(); +} + +// The converse: a consumer the cancel path *did* fail must release the producer +// it will never read, or an unstarted consumer pins its producer forever. +TEST_F(OrchestratorFixture, CancelReleasesProducerRefsHeldByAnUnstartedConsumer) { + orch.scope_begin(); + auto producer = orch.submit_next_level(C(72), single_tensor_args(0x7200, TensorArgType::OUTPUT), cfg, 0); + auto consumer = orch.submit_next_level(C(73), single_tensor_args(0x7200, TensorArgType::INPUT), cfg, 0); + ASSERT_EQ(S(consumer.task_slot).state.load(std::memory_order_acquire), TaskState::PENDING); + const int32_t released_before = S(producer.task_slot).fanout_released.load(std::memory_order_acquire); + + orch.fail_run_submission(run_id, std::make_exception_ptr(std::runtime_error("orchestration: boom"))); + + EXPECT_GT(S(producer.task_slot).fanout_released.load(std::memory_order_acquire), released_before) + << "a cancelled consumer must release the producer it will never read"; + + orch.scope_end(); +} + +// retire erases a run's ready-queue partitions. A late enqueue — a dependency +// releasing its consumer, or the scheduler requeueing a task whose SUB worker +// was busy — must not rebuild one: nothing drains a partition whose run will +// never hold the FIFO head again. +// A failure claim won while a slot was BUILDING leaves the propagation to the +// submitting thread. If that thread throws before it gets there, the slot is +// FAILED with its references still held — and skipping it for being FAILED is +// what would leave the run's fence unreachable forever. +TEST_F(OrchestratorFixture, ACancelledRunTakesOverAPropagationItsSubmitAbandoned) { + auto task = orch.submit_next_level(C(76), single_tensor_args(0x7600, TensorArgType::OUTPUT), cfg, 0); + TaskSlot slot = task.task_slot; + TaskSlot ready; + ASSERT_TRUE(rq.try_pop(run_id, ready)); + + // The state an abandoned submit leaves behind. + S(slot).state.store(TaskState::BUILDING, std::memory_order_release); + ASSERT_TRUE(claim_task_failure(S(slot), "producer failed").has_value()); + ASSERT_TRUE(S(slot).failure_propagation_pending.load(std::memory_order_acquire)); + + orch.fail_run_submission(run_id, std::make_exception_ptr(std::runtime_error("callback threw"))); + + EXPECT_EQ(S(slot).state.load(std::memory_order_acquire), TaskState::CONSUMED) + << "an abandoned propagation stranded the slot"; + EXPECT_FALSE(S(slot).failure_propagation_pending.load(std::memory_order_acquire)); + EXPECT_TRUE(orch.run_done(run_id)) << "the run fence never became reachable"; + orch.release_run(run_id); +} + +// Every route back into a ready queue goes through enqueue_ready — including +// the scheduler's requeue of a task whose target worker was busy, which is the +// one that can fire after its run is already terminal. +TEST_F(OrchestratorFixture, ATerminalRunsQueuePartitionIsNotRebuilt) { + auto task = orch.submit_next_level(C(74), single_tensor_args(0x7400, TensorArgType::OUTPUT), cfg, 0); + auto sub = orch.submit_sub(C(75), single_tensor_args(0x7500, TensorArgType::OUTPUT)); + TaskSlot ready; + ASSERT_TRUE(rq.try_pop(run_id, ready)); + ASSERT_EQ(ready, task.task_slot); + ASSERT_TRUE(rq_sub.try_pop(run_id, ready)); + ASSERT_EQ(ready, sub.task_slot); + + for (TaskSlot slot : {task.task_slot, sub.task_slot}) { + S(slot).state.store(TaskState::COMPLETED, std::memory_order_release); + ASSERT_TRUE(orch.on_consumed(slot)); + } + orch.close_run_submission(run_id); + ASSERT_TRUE(orch.run_done(run_id)); + + // The slots are long gone, but a stale enqueue can still arrive here. + orch.enqueue_ready(task.task_slot); + orch.enqueue_ready(sub.task_slot); + + TaskSlot leaked; + EXPECT_FALSE(rq.try_pop(run_id, leaked)) << "a terminal run's partition was rebuilt by a late enqueue"; + EXPECT_FALSE(rq_next_level.try_front_group(run_id, leaked)); + EXPECT_FALSE(rq_sub.try_pop(run_id, leaked)) << "a busy-worker requeue bypassed the terminal-run guard"; + EXPECT_TRUE(rq_sub.empty()); + orch.release_run(run_id); +} + +// A run whose lease has been released is not dispatchable, even while its id is +// still the FIFO head the scheduler last observed. +TEST_F(OrchestratorFixture, ReleasedLeaseMakesTheRunUndispatchable) { + EXPECT_EQ(orch.dispatchable_run_id(), run_id); + EXPECT_TRUE(orch.can_dispatch_run(run_id)); + + orch.close_run_submission(run_id); + EXPECT_TRUE(orch.run_done(run_id)); + + EXPECT_EQ(orch.dispatchable_run_id(), INVALID_RUN_ID) + << "a terminal run must not be offered to the scheduler as dispatchable"; + EXPECT_FALSE(orch.can_dispatch_run(run_id)); + orch.release_run(run_id); +} + +TEST_F(OrchestratorFixture, FifoHeadCanExecuteWhileGraphConstructionIsOpen) { + EXPECT_TRUE(orch.can_dispatch_run(run_id)); + auto task = orch.submit_next_level(C(89), single_tensor_args(0x8900, TensorArgType::OUTPUT), cfg, 0); + + TaskSlot ready; + ASSERT_TRUE(rq.try_pop(run_id, ready)); + EXPECT_EQ(ready, task.task_slot); + S(task.task_slot).state.store(TaskState::COMPLETED, std::memory_order_release); + ASSERT_TRUE(orch.on_consumed(task.task_slot)); + EXPECT_FALSE(orch.run_done(run_id)); + + orch.close_run_submission(run_id); + EXPECT_TRUE(orch.run_done(run_id)); + orch.release_run(run_id); +} + +TEST_F(OrchestratorFixture, BuildingSuccessorActivatesAfterPriorRunIsTerminal) { + auto first = orch.submit_next_level(C(90), single_tensor_args(0x9000, TensorArgType::OUTPUT), cfg, 0); + orch.close_run_submission(run_id); + + TaskSlot ready; + ASSERT_TRUE(rq.try_pop(run_id, ready)); + RunId second = orch.begin_run(); + auto interactive = orch.submit_next_level(C(91), single_tensor_args(0x9100, TensorArgType::OUTPUT), cfg, 0); + EXPECT_FALSE(orch.can_dispatch_run(second)); + + S(first.task_slot).state.store(TaskState::COMPLETED, std::memory_order_release); + ASSERT_TRUE(orch.on_consumed(first.task_slot)); + EXPECT_TRUE(orch.can_dispatch_run(second)); + ASSERT_TRUE(rq.try_pop(second, ready)); + EXPECT_EQ(ready, interactive.task_slot); + + S(interactive.task_slot).state.store(TaskState::COMPLETED, std::memory_order_release); + ASSERT_TRUE(orch.on_consumed(interactive.task_slot)); + EXPECT_FALSE(orch.run_done(second)); + orch.close_run_submission(second); + EXPECT_TRUE(orch.run_done(second)); + + orch.release_run(run_id); + orch.release_run(second); +} + +TEST_F(OrchestratorFixture, PreparedRunWaitsForActiveRunAndThirdAdmissionBlocks) { + auto first = orch.submit_next_level(C(90), single_tensor_args(0x9000, TensorArgType::OUTPUT), cfg, 0); + orch.close_run_submission(run_id); + + TaskSlot active_slot; + ASSERT_TRUE(rq.try_pop(run_id, active_slot)); + ASSERT_EQ(active_slot, first.task_slot); + + RunId second = orch.begin_run(); + auto prepared = orch.submit_next_level(C(91), single_tensor_args(0x9100, TensorArgType::OUTPUT), cfg, 0); + orch.close_run_submission(second); + + EXPECT_FALSE(orch.can_dispatch_run(second)); + + auto third = std::async(std::launch::async, [this] { + return orch.begin_run(); + }); + EXPECT_EQ(third.wait_for(std::chrono::milliseconds(50)), std::future_status::timeout) + << "third admission must block before graph construction"; + + S(first.task_slot).state.store(TaskState::COMPLETED, std::memory_order_release); + ASSERT_TRUE(orch.on_consumed(first.task_slot)); + + ASSERT_EQ(third.wait_for(std::chrono::seconds(1)), std::future_status::ready); + RunId third_id = third.get(); + ASSERT_TRUE(rq.try_pop(second, active_slot)); + EXPECT_EQ(active_slot, prepared.task_slot); + + S(prepared.task_slot).state.store(TaskState::COMPLETED, std::memory_order_release); + ASSERT_TRUE(orch.on_consumed(prepared.task_slot)); + orch.close_run_submission(third_id); + + EXPECT_NO_THROW(orch.wait_run(run_id)); + EXPECT_NO_THROW(orch.wait_run(second)); + EXPECT_NO_THROW(orch.wait_run(third_id)); + orch.release_run(run_id); + orch.release_run(second); + orch.release_run(third_id); +} + +// At depth one the pipeline slot a retiring run gives back is the only one +// there is, so the waiter in begin_run has exactly one wakeup to catch. It is +// caught because the release happens under the mutex whose predicate reads it: +// releasing outside it can land between a waiter finding "full" and its +// registering on the condition variable, and the notify then reaches nobody. +// +// The loop is what exercises the window — a single pass would rarely place the +// release there — so a stranded slot shows up as this test not finishing. +TEST(DepthOneAdmission, ARetiringRunAlwaysWakesTheOneWaiterItUnblocks) { + for (int round = 0; round < 200; ++round) { + TensorMap tm; + Ring allocator; + Scope scope; + NextLevelReadyQueues rq_next_level; + ReadyQueue rq_sub; + Orchestrator orch; + allocator.init(/*heap_bytes=*/1ULL << 20); + rq_next_level.reset({0}); + orch.init(&tm, &allocator, &scope, &rq_sub, &rq_next_level); + orch.configure_pipeline_depth(1); + + RunId first = orch.begin_run(); + CallConfig cfg; + TaskArgs args; + Tensor t{}; + t.buffer.addr = 0xD000 + static_cast(round); + t.ndims = 1; + t.shapes[0] = 1; + t.dtype = DataType::UINT8; + args.add_tensor(t, TensorArgType::OUTPUT); + CallableIdentity callable; + callable.digest.fill(77); + auto task = orch.submit_next_level(callable, args, cfg, 0); + orch.close_run_submission(first); + TaskSlot ready = INVALID_SLOT; + ASSERT_TRUE(rq_next_level.try_pop_single(0, first, ready)); + + // The successor races the retirement for the only slot in the pool. + auto successor = std::async(std::launch::async, [&orch] { + return orch.begin_run(); + }); + allocator.slot_state(task.task_slot)->state.store(TaskState::COMPLETED, std::memory_order_release); + ASSERT_TRUE(orch.on_consumed(task.task_slot)); + + ASSERT_EQ(successor.wait_for(std::chrono::seconds(5)), std::future_status::ready) + << "round " << round << ": the freed pipeline slot never woke its waiter"; + RunId second = successor.get(); + orch.close_run_submission(second); + orch.release_run(first); + orch.release_run(second); + allocator.shutdown(); + } +} + +// A submit that throws part-way has already charged its slot to the run. What +// the run's fence then waits on is that slot reaching CONSUMED, and that needs +// the scope reference it was charged to be one a scope_end can actually +// release — so the registration happens before the charge, and both before +// anything that can throw. Python closes the scope before cancelling, which is +// the order used here. +TEST(SubmitFailure, ASlotWhoseSubmitThrewIsFullyReclaimedByCancellation) { + TensorMap tm; + Ring allocator; + Scope scope; + ReadyQueue rq_sub; + Orchestrator orch; + allocator.init(/*heap_bytes=*/1ULL << 20); + // No NEXT_LEVEL queues: the publication at the end of submit throws, which + // is the one point in the sequence a test can fault from outside. + orch.init(&tm, &allocator, &scope, &rq_sub, nullptr); + + RunId run = orch.begin_run(); + orch.scope_begin(); + + CallConfig cfg; + TaskArgs args; + Tensor t{}; + t.buffer.addr = 0xE100; + t.ndims = 1; + t.shapes[0] = 1; + t.dtype = DataType::UINT8; + args.add_tensor(t, TensorArgType::OUTPUT); + CallableIdentity callable; + callable.digest.fill(88); + EXPECT_THROW((void)orch.submit_next_level(callable, args, cfg, 0), std::runtime_error); + + orch.scope_end(); + orch.fail_run_submission(run, std::make_exception_ptr(std::runtime_error("graph construction failed"))); + + EXPECT_EQ(allocator.slot_state(0)->state.load(std::memory_order_acquire), TaskState::CONSUMED) + << "a slot charged a scope reference it never registered stalls the run fence"; + EXPECT_TRUE(orch.run_done(run)); + EXPECT_THROW(orch.wait_run(run), std::runtime_error); + orch.release_run(run); + allocator.shutdown(); +} + +TEST_F(OrchestratorFixture, FailedPreparedConstructionReturnsAdmissionWithoutDispatch) { + auto active = orch.submit_next_level(C(92), single_tensor_args(0x9200, TensorArgType::OUTPUT), cfg, 0); + orch.close_run_submission(run_id); + TaskSlot active_slot; + ASSERT_TRUE(rq.try_pop(run_id, active_slot)); + + RunId failed = orch.begin_run(); + auto cancelled = orch.submit_next_level(C(93), single_tensor_args(0x9300, TensorArgType::OUTPUT), cfg, 0); + orch.fail_run_submission(failed, std::make_exception_ptr(std::runtime_error("graph build failed"))); + + EXPECT_TRUE(orch.run_done(failed)); + EXPECT_EQ(S(cancelled.task_slot).state.load(std::memory_order_acquire), TaskState::CONSUMED); + TaskSlot unexpected; + EXPECT_FALSE(rq.try_pop(failed, unexpected)); + + auto replacement = std::async(std::launch::async, [this] { + return orch.begin_run(); + }); + ASSERT_EQ(replacement.wait_for(std::chrono::seconds(1)), std::future_status::ready); + RunId replacement_id = replacement.get(); + + S(active.task_slot).state.store(TaskState::COMPLETED, std::memory_order_release); + ASSERT_TRUE(orch.on_consumed(active.task_slot)); + orch.close_run_submission(replacement_id); + + EXPECT_THROW(orch.wait_run(failed), std::runtime_error); + EXPECT_NO_THROW(orch.wait_run(replacement_id)); + orch.release_run(run_id); + orch.release_run(failed); + orch.release_run(replacement_id); +} diff --git a/tests/ut/cpp/hierarchical/test_pipeline_contract.cpp b/tests/ut/cpp/hierarchical/test_pipeline_contract.cpp index b4d8fedc3..b40aae485 100644 --- a/tests/ut/cpp/hierarchical/test_pipeline_contract.cpp +++ b/tests/ut/cpp/hierarchical/test_pipeline_contract.cpp @@ -217,6 +217,16 @@ TEST(PipelineSlotPool, DepthTwoProvidesExactlyTwoIndependentLeases) { EXPECT_TRUE(pool.owns(*second)); } +TEST(PipelineSlotPool, AdmissionDepthCanConservativelyLimitACapablePool) { + PipelineSlotPool pool(2); + auto first = pool.try_acquire(/*admission_depth=*/1); + ASSERT_TRUE(first.has_value()); + EXPECT_EQ(first->slot_id, 0u); + EXPECT_FALSE(pool.try_acquire(/*admission_depth=*/1).has_value()); + EXPECT_TRUE(pool.release(*first)); + EXPECT_THROW((void)pool.try_acquire(/*admission_depth=*/3), std::invalid_argument); +} + TEST(PipelineSlotPool, StaleGenerationCannotAccessOrReleaseAReusedSlot) { PipelineSlotPool pool(1); const PipelineSlotLease first = *pool.try_acquire(); diff --git a/tests/ut/cpp/hierarchical/test_scheduler.cpp b/tests/ut/cpp/hierarchical/test_scheduler.cpp index d56061ce4..443a3219e 100644 --- a/tests/ut/cpp/hierarchical/test_scheduler.cpp +++ b/tests/ut/cpp/hierarchical/test_scheduler.cpp @@ -272,6 +272,142 @@ static CallableIdentity C(uint8_t seed) { // Fixture // --------------------------------------------------------------------------- +// The claim is what makes "pop a READY slot" and "dispatch it" one decision. +// A cancelling run moves its unstarted slots out of READY and consumes them; +// anything the scheduler had already popped must lose the race rather than +// overwrite that state with RUNNING. +TEST(ClaimForDispatch, OnlyAReadySlotCanBeClaimed) { + TaskSlotState s; + + s.state.store(TaskState::READY, std::memory_order_release); + EXPECT_TRUE(claim_for_dispatch(s)); + EXPECT_EQ(s.state.load(std::memory_order_acquire), TaskState::RUNNING); + + // Already claimed by this scheduler: a second claim must not re-dispatch. + EXPECT_FALSE(claim_for_dispatch(s)); + EXPECT_EQ(s.state.load(std::memory_order_acquire), TaskState::RUNNING); + + // Cancelled between the pop and here — the cancelling path owns it now. + // BUILDING is in the list for a different reason: a slot whose submit has + // not published it has no final args or fanin count to dispatch on. + for (TaskState taken : + {TaskState::FAILED, TaskState::COMPLETED, TaskState::CONSUMED, TaskState::PENDING, TaskState::BUILDING}) { + s.state.store(taken, std::memory_order_release); + EXPECT_FALSE(claim_for_dispatch(s)) << "claimed a slot in state " << static_cast(taken); + EXPECT_EQ(s.state.load(std::memory_order_acquire), taken) << "claim overwrote a state it did not own"; + } +} + +// Every failing path — a device completion poisoning its consumers, a run +// cancellation, and a submit that wired onto a producer that had already +// failed — moves a task to FAILED through this one exchange, so exactly one of +// them writes the message and runs the propagation. +TEST(ClaimTaskFailure, ReportsThePriorStateAndIsWonOnce) { + TaskSlotState s; + + for (TaskState claimable : {TaskState::PENDING, TaskState::READY, TaskState::BUILDING}) { + s.state.store(claimable, std::memory_order_release); + s.failure_message.clear(); + + std::optional won = claim_task_failure(s, "first"); + ASSERT_TRUE(won.has_value()) << "refused a claimable slot in state " << static_cast(claimable); + EXPECT_EQ(*won, claimable) << "claim did not report the state it took the slot from"; + EXPECT_EQ(s.state.load(std::memory_order_acquire), TaskState::FAILED); + EXPECT_EQ(s.failure_message, "first"); + + // A second path reaching the same slot must not overwrite the reason + // the first one recorded. + EXPECT_FALSE(claim_task_failure(s, "second").has_value()); + EXPECT_EQ(s.failure_message, "first"); + } +} + +// Readiness is one decision over two fields that different threads own. +// Judging the count outside the transition lets a producer pass a comparison +// against a count submit has not published, then act on it after submit has — +// dispatching a task whose remaining producers are still running. +TEST(TryMarkReady, JudgesThePublishedCountNotTheOneItArrivedWith) { + TaskSlotState s; + s.state.store(TaskState::BUILDING, std::memory_order_release); + + // A producer completes while the slot is still building. The count it can + // see is zero, which any release count passes — but nothing is readiable + // yet, and the release is not lost either. + s.fanin_released.store(1, std::memory_order_release); + EXPECT_FALSE(try_mark_ready(s)); + + // Submit publishes two live producers alongside the transition. + { + std::lock_guard lk(s.fanout_mu); + s.fanin_count.store(2, std::memory_order_release); + s.state.store(TaskState::PENDING, std::memory_order_release); + } + EXPECT_FALSE(try_mark_ready(s)) << "one of two producers released and the task was marked ready"; + + s.fanin_released.store(2, std::memory_order_release); + EXPECT_TRUE(try_mark_ready(s)); + EXPECT_EQ(s.state.load(std::memory_order_acquire), TaskState::READY); + + // Exactly one caller owns the enqueue. + EXPECT_FALSE(try_mark_ready(s)); +} + +// The publication is not a moment another thread can slip through. A producer +// that has already completed is held outside it for its whole duration, so +// there is no instant at which the slot is PENDING with a count only half the +// deciders have seen — which is the state that dispatches a task whose +// remaining producers are still running. +TEST(TryMarkReady, NoProducerTransitionsTheSlotWhileThePublicationHoldsIt) { + TaskSlotState s; + s.state.store(TaskState::BUILDING, std::memory_order_release); + s.fanin_count.store(0, std::memory_order_release); + + std::atomic producer_marked_ready{false}; + std::unique_lock publishing(s.fanout_mu); + + std::thread producer([&] { + s.fanin_released.fetch_add(1, std::memory_order_acq_rel); + producer_marked_ready.store(try_mark_ready(s), std::memory_order_release); + }); + + // The producer's release has landed and its decision is now in flight. + while (s.fanin_released.load(std::memory_order_acquire) == 0) + std::this_thread::yield(); + + // Publish two live producers alongside the transition, then stay in the + // critical section: a decider that judged the count before entering it + // would take PENDING to READY right here. + s.fanin_count.store(2, std::memory_order_release); + s.state.store(TaskState::PENDING, std::memory_order_release); + for (int i = 0; i < 1000; ++i) + std::this_thread::yield(); + EXPECT_EQ(s.state.load(std::memory_order_acquire), TaskState::PENDING) + << "a producer transitioned the slot from inside the publication"; + + publishing.unlock(); + producer.join(); + EXPECT_FALSE(producer_marked_ready.load(std::memory_order_acquire)); + EXPECT_EQ(s.state.load(std::memory_order_acquire), TaskState::PENDING) + << "a task was made ready with a live producer still running"; +} + +TEST(ClaimTaskFailure, RefusesASlotItDoesNotOwn) { + TaskSlotState s; + + // RUNNING is the device's until its completion arrives; the rest are + // already terminal, and resurrecting one would release its dependency + // references a second time. + for (TaskState owned : + {TaskState::RUNNING, TaskState::COMPLETED, TaskState::FAILED, TaskState::CONSUMED, TaskState::FREE}) { + s.state.store(owned, std::memory_order_release); + s.failure_message.clear(); + EXPECT_FALSE(claim_task_failure(s, "cancelled").has_value()) + << "claimed a slot in state " << static_cast(owned); + EXPECT_EQ(s.state.load(std::memory_order_acquire), owned); + EXPECT_TRUE(s.failure_message.empty()); + } +} + struct SchedulerFixture : public ::testing::Test { TensorMap tm; Ring allocator; @@ -288,6 +424,10 @@ struct SchedulerFixture : public ::testing::Test { std::vector consumed_slots; std::mutex consumed_mu; + // Set by a test before the Scheduler reaches a claim; see + // Scheduler::Config::before_claim_cb. + std::function before_claim_hook; + TaskSlotState &S(TaskSlot id) { return *allocator.slot_state(id); } void SetUp() override { @@ -318,6 +458,13 @@ struct SchedulerFixture : public ::testing::Test { c.enqueue_ready_cb = [this](TaskSlot slot) { orch.enqueue_ready(slot); }; + // Same gate Worker::start installs: an active run that is also the + // EXECUTING FIFO head and still owns its pipeline lease. Testing + // against the weaker active_run_id() would let a slot dispatch here + // that production refuses. + c.active_run_cb = [this] { + return orch.dispatchable_run_id(); + }; c.on_consumed_cb = [this](TaskSlot s) { orch.on_consumed(s); std::lock_guard lk(consumed_mu); @@ -326,6 +473,9 @@ struct SchedulerFixture : public ::testing::Test { c.on_task_failed_cb = [this](TaskSlot s, const std::string &message) { orch.report_task_error(s, message); }; + c.before_claim_cb = [this](TaskSlot slot) { + if (before_claim_hook) before_claim_hook(slot); + }; sched.start(c); } @@ -387,6 +537,7 @@ TEST(WorkerManagerTest, LocalMailboxPublishesAcceptanceBeforeCompletion) { ASSERT_NE(slot, nullptr); slot->reset(); slot->callable.digest[0] = 0x42; + slot->pipeline_lease = PipelineSlotLease{1, 0, 7}; LocalMailboxEndpoint endpoint(/*worker_id=*/0, child.mailbox_ptr()); std::promise result; @@ -400,6 +551,12 @@ TEST(WorkerManagerTest, LocalMailboxPublishesAcceptanceBeforeCompletion) { child.wait_running(); EXPECT_TRUE(child.is_running.load(std::memory_order_acquire)); + PipelineSlotLease wire_lease{}; + std::memcpy( + &wire_lease, static_cast(child.mailbox_ptr()) + MAILBOX_OFF_PIPELINE_LEASE, sizeof(PipelineSlotLease) + ); + EXPECT_EQ(wire_lease.slot_id, 1u); + EXPECT_EQ(wire_lease.generation, 7u); child.write_task_accepted(); auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(3); while (!accepted.load(std::memory_order_acquire) && std::chrono::steady_clock::now() < deadline) {} @@ -533,6 +690,37 @@ TEST(WorkerManagerTest, ControlPrepareUsesStableNextLevelWorkerId) { EXPECT_EQ(worker3_prepares.load(std::memory_order_relaxed), 1); } +// The losing side of the dispatch claim, driven by the real cancellation path +// rather than a simulated state write. `before_claim_cb` is the only point that +// can observe the window: everything else is either before the queue pop or +// after the launch. +TEST_F(SchedulerFixture, ACancellationThatWinsTheClaimStopsTheDispatch) { + std::atomic hook_calls{0}; + before_claim_hook = [this, &hook_calls](TaskSlot) { + if (hook_calls.fetch_add(1) != 0) return; + orch.fail_run_submission(run_id, std::make_exception_ptr(std::runtime_error("cancelled mid-dispatch"))); + }; + + auto task = orch.submit_next_level(C(60), single_tensor_args(0x6000, TensorArgType::OUTPUT), cfg, 0); + + auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (hook_calls.load() == 0 && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + ASSERT_GT(hook_calls.load(), 0) << "the Scheduler never reached a dispatch claim"; + + // Give the Scheduler its whole loop iteration; a lost claim must leave the + // slot alone rather than continue into the launch. + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + + EXPECT_NE(S(task.task_slot).state.load(std::memory_order_acquire), TaskState::RUNNING) + << "the dispatch overwrote a slot the cancellation already owned"; + { + std::lock_guard lk(mock_worker.dispatched_mu); + EXPECT_TRUE(mock_worker.dispatched.empty()) << "a cancelled task was still handed to a worker"; + } +} + TEST_F(SchedulerFixture, IndependentTaskDispatchedAndConsumed) { auto args_a = single_tensor_args(0xCAFE, TensorArgType::OUTPUT); auto res = orch.submit_next_level(C(42), args_a, cfg, 0); @@ -682,6 +870,12 @@ struct GroupSchedulerFixture : public ::testing::Test { c.enqueue_ready_cb = [this](TaskSlot slot) { orch.enqueue_ready(slot); }; + // Same gate Worker::start installs. Without it the scheduler takes the + // unpartitioned branch, which is not the one #1565's group reservation + // and placement run through. + c.active_run_cb = [this] { + return orch.dispatchable_run_id(); + }; c.on_consumed_cb = [this](TaskSlot s) { orch.on_consumed(s); std::lock_guard lk(consumed_mu); @@ -1086,6 +1280,69 @@ TEST_F(GroupSchedulerFixture, DependencyReleaseUsesConsumerWorkerQueue) { wait_consumed(consumer.task_slot); } +// A producer that fails while its consumer is still being submitted. Wiring +// happens under each producer's fanout_mu, so the consumer is reachable from +// the failing producer's fanout list well before its own fanin/fanout counters +// are final — which is exactly what BUILDING marks. The poison must stop at the +// claim and leave the propagation to the submitting thread; running it from +// both sides releases every producer reference the consumer holds twice. +// +// The window is opened by holding the *second* producer's fanout_mu: submit +// wires the first producer, then parks on the second, and the failure is +// injected in between. +TEST_F(GroupSchedulerFixture, APoisonThatLandsMidSubmitLeavesThePropagationToSubmit) { + auto failing = orch.submit_next_level(C(70), single_tensor_args(0xF100, TensorArgType::OUTPUT), cfg, 0); + auto blocking = orch.submit_next_level(C(71), single_tensor_args(0xB200, TensorArgType::OUTPUT), cfg, 1); + worker_a.wait_running(); + worker_b.wait_running(); + + TaskArgs consumer_args; + for (uint64_t key : {0xF100ULL, 0xB200ULL}) { + Tensor t{}; + t.buffer.addr = key; + t.ndims = 1; + t.shapes[0] = 1; + t.dtype = DataType::UINT8; + consumer_args.add_tensor(t, TensorArgType::INPUT); + } + + std::unique_lock parked(S(blocking.task_slot).fanout_mu); + std::thread submitter([&] { + (void)orch.submit_next_level(C(72), consumer_args, cfg, 2); + }); + + // Wired into `failing` and now parked on `blocking`: the exact window. + TaskSlot consumer = INVALID_SLOT; + auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (consumer == INVALID_SLOT && std::chrono::steady_clock::now() < deadline) { + std::lock_guard lk(S(failing.task_slot).fanout_mu); + if (!S(failing.task_slot).fanout_consumers.empty()) consumer = S(failing.task_slot).fanout_consumers[0]; + } + ASSERT_NE(consumer, INVALID_SLOT) << "submit never reached the failing producer's fanout list"; + ASSERT_EQ(S(consumer).state.load(std::memory_order_acquire), TaskState::BUILDING); + + worker_a.complete_with_error("producer boom"); + while (S(consumer).state.load(std::memory_order_acquire) == TaskState::BUILDING && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + ASSERT_EQ(S(consumer).state.load(std::memory_order_acquire), TaskState::FAILED) + << "the poison did not claim the consumer while it was building"; + + parked.unlock(); + submitter.join(); + + // One release per producer reference the consumer holds, from the one + // thread that knows the wiring is final. `failing` reaches its threshold + // (fanout_total 1 + the terminal self release) and no further. + EXPECT_EQ(S(consumer).state.load(std::memory_order_acquire), TaskState::CONSUMED); + EXPECT_EQ(S(failing.task_slot).fanout_released.load(std::memory_order_acquire), 2); + EXPECT_EQ(S(blocking.task_slot).fanout_released.load(std::memory_order_acquire), 1); + + worker_b.complete(); + wait_consumed(blocking.task_slot); +} + TEST_F(GroupSchedulerFixture, TargetMustBeInEligibleEndpointSet) { TaskArgs args = single_tensor_args(0xE1, TensorArgType::OUTPUT); EXPECT_THROW((void)orch.submit_next_level(C(56), args, cfg, 0, {1}), std::invalid_argument); @@ -1273,6 +1530,13 @@ struct MixedTypeSchedulerFixture : public ::testing::Test { c.enqueue_ready_cb = [this](TaskSlot slot) { orch.enqueue_ready(slot); }; + // Same gate Worker::start installs: an active run that is also the + // EXECUTING FIFO head and still owns its pipeline lease. Testing + // against the weaker active_run_id() would let a slot dispatch here + // that production refuses. + c.active_run_cb = [this] { + return orch.dispatchable_run_id(); + }; c.on_consumed_cb = [this](TaskSlot s) { orch.on_consumed(s); std::lock_guard lk(consumed_mu); @@ -1336,6 +1600,28 @@ TEST_F(MixedTypeSchedulerFixture, SubTaskDispatchesWhileNextLevelPoolSaturated) wait_consumed(chip.task_slot); } +TEST_F(MixedTypeSchedulerFixture, BusySubWorkerRequeuesWithinTheActiveRun) { + auto first = orch.submit_sub(C(8), single_tensor_args(0xC01, TensorArgType::OUTPUT)); + sub_worker.wait_running(); + ASSERT_TRUE(sub_worker.is_running.load()); + + auto second = orch.submit_sub(C(9), single_tensor_args(0xC02, TensorArgType::OUTPUT)); + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + EXPECT_EQ(sub_worker.dispatched_count(), 1); + + sub_worker.complete(); + wait_consumed(first.task_slot); + auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(500); + while (sub_worker.dispatched_count() < 2 && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + ASSERT_EQ(sub_worker.dispatched_count(), 2); + EXPECT_TRUE(sub_worker.is_running.load()); + + sub_worker.complete(); + wait_consumed(second.task_slot); +} + TEST_F(GroupSchedulerFixture, GroupDependencyChain) { // Group A (2 workers) produces an OUTPUT at key 0xCAFE. // Task B reads INPUT at the same key -- depends on group A. diff --git a/tests/ut/py/test_callable_identity.py b/tests/ut/py/test_callable_identity.py index e8ce43449..7e2880fe4 100644 --- a/tests/ut/py/test_callable_identity.py +++ b/tests/ut/py/test_callable_identity.py @@ -540,6 +540,7 @@ class FakeCWorker: def __init__(self, *args): self.remote_worker_ids = [] self.closed = False + self.pipeline_depth = None def add_remote_l3_socket(self, worker_id, *args): self.remote_worker_ids.append(worker_id) @@ -555,6 +556,9 @@ def add_next_level_worker(self, *args): def add_next_level_worker_at(self, *args): pass + def configure_pipeline_depth(self, depth): + self.pipeline_depth = depth + def init(self): pass @@ -597,6 +601,7 @@ def fake_open_remote_session(self, *, spec, worker_id, session_id, deadline): assert worker._resolve_handle(handle).eligible_worker_ids == (remote_worker_id,) worker.init() + assert fake_c_worker.pipeline_depth == 2 assert opened_worker_ids == [remote_worker_id] assert fake_c_worker.remote_worker_ids == [remote_worker_id] @@ -1605,7 +1610,11 @@ def remote_free(self, *args): worker._lifecycle = worker_mod._Lifecycle.READY worker._pending_remote_buffer_frees = [owner] - worker._flush_pending_remote_frees() + # Retained for a retry, and reported: this runs inside a run's fence-owned + # cleanup, where returning normally publishes the run as cleanly finished + # over a remote allocation it still owns. + with pytest.raises(RuntimeError, match="remain owed"): + worker._flush_pending_remote_frees() assert worker._pending_remote_buffer_frees == [owner] diff --git a/tests/ut/py/test_worker/test_host_worker.py b/tests/ut/py/test_worker/test_host_worker.py index ac93af168..95acfd5c0 100644 --- a/tests/ut/py/test_worker/test_host_worker.py +++ b/tests/ut/py/test_worker/test_host_worker.py @@ -19,9 +19,12 @@ import time import weakref from multiprocessing.shared_memory import SharedMemory +from types import SimpleNamespace from typing import Any, cast +from unittest.mock import patch import pytest +import simpler.orchestrator as orch_mod import simpler.worker as worker_mod from _task_interface import MAX_REGISTERED_CALLABLE_IDS # pyright: ignore[reportMissingImports] from simpler.callable_identity import ( @@ -130,8 +133,11 @@ def _chip_payload_shm(callable_obj: ChipCallable) -> SharedMemory: def test_chip_process_loop_inits_runs_and_finalizes(monkeypatch): events: list[tuple] = [] + published_depths: list[int] = [] class FakeChipWorker: + pipeline_depth = 2 + def init(self, device_id, bins, *, log_level, prewarm_config=None, enable_sdma=False): events.append(("init", device_id, bins, log_level, prewarm_config, enable_sdma)) @@ -139,6 +145,7 @@ def finalize(self) -> None: events.append(("finalize",)) def fake_run_chip_main_loop(cw, *_args, chip_platform, chip_runtime, prepared=None): + published_depths.append(worker_mod._PIPELINE_LEASE_FMT.unpack_from(_args[0], worker_mod._OFF_PIPELINE_LEASE)[0]) events.append(("main_loop", cw, chip_platform, chip_runtime)) monkeypatch.setattr(worker_mod, "ChipWorker", FakeChipWorker) @@ -165,6 +172,7 @@ def fake_run_chip_main_loop(cw, *_args, chip_platform, chip_runtime, prepared=No assert events[1][0] == "main_loop" assert events[1][2:] == ("a2a3", "tensormap_and_ringbuffer") assert events[2] == ("finalize",) + assert published_depths == [2] def _chip_digest(callable_obj: ChipCallable, *, platform: str = "", runtime: str = "") -> bytes: @@ -1435,6 +1443,77 @@ def orch(o, _args, _cfg): state_shm.close() state_shm.unlink() + def test_depth_two_prepares_next_run_and_blocks_third_callback(self): + state_shm = SharedMemory(create=True, size=16) + state_buf = state_shm.buf + assert state_buf is not None + for offset in range(0, 16, 4): + _set_flag(state_buf, offset, 0) + + third_callback = threading.Event() + third_result: dict[str, RunHandle] = {} + hw = Worker(level=3, num_sub_workers=2) + try: + + def first_task(_args): + _set_flag(state_buf, 0, 1) + while _get_flag(state_buf, 4) == 0: + time.sleep(0.001) + + def second_task(_args): + _set_flag(state_buf, 8, 1) + while _get_flag(state_buf, 12) == 0: + time.sleep(0.001) + + first_target = hw.register(first_task) + second_target = hw.register(second_task) + hw.init() + + first = hw.submit(lambda o, _args, _cfg: o.submit_sub(first_target)) + deadline = time.monotonic() + 3.0 + while _get_flag(state_buf, 0) == 0 and time.monotonic() < deadline: + time.sleep(0.001) + assert _get_flag(state_buf, 0) == 1 + + second_callback_done = False + + def second_graph(o, _args, _cfg): + nonlocal second_callback_done + o.submit_sub(second_target) + second_callback_done = True + + second = hw.submit(second_graph) + assert second_callback_done + time.sleep(0.05) + assert _get_flag(state_buf, 8) == 0, "prepared run dispatched before the active run became terminal" + + def third_graph(_o, _args, _cfg): + third_callback.set() + + submitter = threading.Thread(target=lambda: third_result.setdefault("handle", hw.submit(third_graph))) + submitter.start() + assert not third_callback.wait(0.05), "third callback ran before depth-two admission freed a slot" + + _set_flag(state_buf, 4, 1) + assert third_callback.wait(3.0) + deadline = time.monotonic() + 3.0 + while _get_flag(state_buf, 8) == 0 and time.monotonic() < deadline: + time.sleep(0.001) + assert _get_flag(state_buf, 8) == 1 + + _set_flag(state_buf, 12, 1) + submitter.join(5.0) + assert not submitter.is_alive() + first.wait(5.0) + second.wait(5.0) + third_result["handle"].wait(5.0) + finally: + _set_flag(state_buf, 4, 1) + _set_flag(state_buf, 12, 1) + hw.close() + state_shm.close() + state_shm.unlink() + def test_graph_error_is_synchronous(self): hw = Worker(level=3, num_sub_workers=0) hw.init() @@ -2105,8 +2184,10 @@ def test_failed_domain_release_is_replayed_to_a_second_caller(self): assert not contested.freed # The sweep is the second caller: it must see the failure, keep the - # handle, and leave `freed` false. - worker._release_all_live_domains() + # handle, leave `freed` false, and report the failure to its own caller + # rather than returning as if the domain were reclaimed. + with pytest.raises(RuntimeError, match="backend release failed"): + worker._release_all_live_domains() assert not contested.freed, "a failed release must not be reported as freed" assert "contested" in worker._live_domains, "a failed release must stay a detectable residual" @@ -2119,6 +2200,678 @@ def test_allocate_domain_outside_graph_construction_is_rejected(self): worker._allocate_domain(name="d", workers=(0,), window_size=4096, buffers=[]) +# --------------------------------------------------------------------------- +# Test: conditional serial degradation +# +# The whole-run FIFO orders tasks. It cannot order a run's *cleanup*, which +# happens after the native fence and reaches a child through mailbox control +# rather than a TaskSlot. A run that acquires device-touching cleanup therefore +# degrades this worker to depth one for exactly that run; runs that only +# dispatch tasks keep the full pipeline depth. +# --------------------------------------------------------------------------- + + +class TestOrderedCleanupDegradation: + @staticmethod + def _worker(): + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + return worker + + def _accepted(self, worker, *, bears_cleanup: bool): + resources = worker_mod._RunResources() + resources.requires_ordered_cleanup = bears_cleanup + handle = RunHandle(worker, 1, (), resources) + worker._accepted_run_handles.add(handle) + return handle, resources + + def test_only_device_touching_resources_make_a_run_cleanup_bearing(self): + """The sticky flag is set where teardown reaches a child, and nowhere else. + + Marking every control-using run would drag a run that merely mallocs + down to depth one, which is the overlap the pipeline exists to buy. + """ + worker = self._worker() + + # A CommDomain: its release drives CTRL_RELEASE_DOMAIN on every member. + resources = worker_mod._RunResources() + assert not resources.requires_ordered_cleanup + handle = worker_mod.CommDomainHandle( + name="d", workers=(), contexts={}, allocation_id=1, _release_fn=lambda released: None + ) + resources.live_domains[handle.name] = handle + resources.requires_ordered_cleanup = True # set by _allocate_domain at this point + assert resources.requires_ordered_cleanup + + # A remote slot reference: releasing it is an RPC to the owning worker. + # `create_l3_l2_queue` is not a fourth set-point — it builds its region + # through `create_l3_l2_region`, and inherits the flag from there. + remote = worker_mod._RunResources() + worker._building_run_resources = remote + ref = cast(Any, object()) + worker._adopt_remote_slot_refs([ref]) + assert remote.requires_ordered_cleanup + assert remote.remote_slot_refs == [ref] + + # Adopting nothing is not an acquisition. + empty = worker_mod._RunResources() + worker._building_run_resources = empty + worker._adopt_remote_slot_refs([]) + assert not empty.requires_ordered_cleanup + + def test_a_task_only_run_does_not_block_the_next_submission(self): + worker = self._worker() + self._accepted(worker, bears_cleanup=False) + assert worker._cleanup_bearing_predecessor() is None + + def test_a_cleanup_bearing_run_blocks_until_its_cleanup_is_published(self): + worker = self._worker() + handle, _ = self._accepted(worker, bears_cleanup=True) + assert worker._cleanup_bearing_predecessor() is handle + + handle._cleanup_published = True + assert worker._cleanup_bearing_predecessor() is None + + def test_a_fired_native_fence_does_not_release_the_successor(self): + """`done` is the device draining; it is not the cleanup boundary. + + A successor keyed on the native answer would be admitted while the + CommDomain / L3-L2 / remote-slot teardown is still outstanding. + """ + worker = self._worker() + handle, _ = self._accepted(worker, bears_cleanup=True) + with handle._cv: + handle._terminal = True + assert handle.done + assert not handle._cleanup_published + + assert worker._cleanup_bearing_predecessor() is handle + + def test_a_failed_task_does_not_poison_the_worker(self): + """A kernel that failed says nothing about whether cleanup succeeded. + + Merging the two would shut a worker permanently for an ordinary task + error, and there is no reopening it. + """ + worker = self._worker() + handle, _ = self._accepted(worker, bears_cleanup=True) + worker._orch = cast(Any, _StubOrch()) + + native = RuntimeError("kernel failed") + assert worker._finalize_run_handle(handle, 1, native) is native + assert worker._ordered_cleanup_error is None + assert handle._cleanup_published + assert handle not in worker._accepted_run_handles + worker._require_no_ordered_cleanup_failure("submit") # admits + + def test_a_failed_cleanup_poisons_the_worker_and_refuses_admission(self): + worker = self._worker() + handle, resources = self._accepted(worker, bears_cleanup=True) + worker._orch = cast(Any, _StubOrch()) + boom = RuntimeError("domain release failed") + + def failing_release(res): + raise boom + + worker._release_all_live_domains = cast(Any, failing_release) + resources.live_domains["d"] = cast(Any, object()) + + assert worker._finalize_run_handle(handle, 1, None) is boom + assert worker._ordered_cleanup_error is boom + with pytest.raises(RuntimeError, match="no further work is admitted"): + worker._require_no_ordered_cleanup_failure("submit") + + def test_the_poison_is_published_before_the_handle_is_dropped(self): + """Ordering, not just outcome. + + A submitter that saw neither — handle already gone, poison not yet set + — would find no cleanup-bearing predecessor and no reason to refuse, + and would be admitted on top of unreclaimed device state. + """ + worker = self._worker() + handle, resources = self._accepted(worker, bears_cleanup=True) + worker._orch = cast(Any, _StubOrch()) + observed: list[tuple[bool, bool]] = [] + + def watching_discard(item): + observed.append((worker._ordered_cleanup_error is not None, item._cleanup_published)) + + worker._accepted_run_handles = cast(Any, _WatchedSet(worker._accepted_run_handles, watching_discard)) + resources.live_domains["d"] = cast(Any, object()) + worker._release_all_live_domains = cast(Any, _raiser(RuntimeError("cleanup failed"))) + + worker._finalize_run_handle(handle, 1, None) + assert observed == [(True, True)], "the handle was dropped before its poison was visible" + + def test_a_submission_waiting_on_cleanup_cannot_bypass_its_poison(self): + """The lease-time check cannot have seen a poison this call just waited on. + + Two submissions can both pass native admission before either reaches + the serializer, so the refusal is re-tested under `_submit_mu`, after + the handoff wait. + """ + worker = self._worker() + handle, _ = self._accepted(worker, bears_cleanup=True) + boom = RuntimeError("cleanup failed") + + def handoff(): + # Whoever ran the cleanup published the poison before dropping the + # handle; this waiter only has to let the check see it. + worker._ordered_cleanup_error = boom + handle._cleanup_published = True + worker._accepted_run_handles.discard(handle) + + handle._wait_for_handoff = cast(Any, handoff) + worker._submit_l3_locked = cast(Any, _raiser(AssertionError("admitted on top of a failed cleanup"))) + + with pytest.raises(RuntimeError, match="no further work is admitted") as excinfo: + worker._submit_locked(lambda *a: None, None, None) + assert excinfo.value.__cause__ is boom + + +class TestDirectControlOrdering: + """`malloc` / `free` / `copy_*` / `committed_device_memory` / domain and + region creation / every `remote_*` buffer call reach a child directly + rather than through a TaskSlot, so the ready-queue FIFO does not order + them. Two boundaries close that: control issued inside a run waits for that + run to hold the FIFO head, and control issued outside every run reserves + the worker for the whole call. + """ + + @staticmethod + def _orch(worker): + from unittest.mock import MagicMock # noqa: PLC0415 + + from simpler.orchestrator import Orchestrator # noqa: PLC0415 + + native = MagicMock() + native.malloc.return_value = 0x1000 + return Orchestrator(native, worker), native + + def test_control_inside_a_run_waits_for_that_run_to_hold_the_fifo_head(self): + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + orch, native = self._orch(worker) + + with orch_mod._callback_run(7, worker): + orch.malloc(0, 64) + native.await_run_admission.assert_called_once_with(7) + + def test_control_after_a_task_submission_in_the_same_run_is_refused(self): + """A task travels the ready queue and control travels the mailbox. + + Their order is undefined, and two such pairs on different chips can + each hold the mailbox the other is waiting for. + """ + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + orch, native = self._orch(worker) + + with orch_mod._callback_run(7, worker): + orch.malloc(0, 64) # before any submit: fine + orch_mod._admit_task_submission(worker) + with pytest.raises(RuntimeError, match="cannot follow a task submission"): + orch.malloc(0, 64) + assert native.malloc.call_count == 1 + + def test_each_run_starts_with_no_submissions_of_its_own(self): + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + orch, _native = self._orch(worker) + + with orch_mod._callback_run(7, worker): + orch_mod._admit_task_submission(worker) + with orch_mod._callback_run(8, worker): + orch.malloc(0, 64) + + def test_a_nested_run_restores_its_callers_marker(self): + """An L4 callback drives its children's runs on its own thread. + + The inner run's context is not the outer one's: it must neither inherit + the outer submission nor erase it on the way out. + """ + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + orch, _native = self._orch(worker) + + with orch_mod._callback_run(7, worker): + orch_mod._admit_task_submission(worker) + with orch_mod._callback_run(8, worker): + orch.malloc(0, 64) # the inner run has submitted nothing + with pytest.raises(RuntimeError, match="cannot follow a task submission"): + orch.malloc(0, 64) + assert orch_mod._callback_frames() == [] + + def test_the_reservation_spans_the_call_not_just_the_check(self): + """A sampled check leaves the command itself outside the decision. + + Between "no run is in flight" and the mailbox write, a submit can be + admitted and dispatch a task that races it — which is the whole thing + the check was for. The serializer submission holds is held here too. + """ + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + orch, _native = self._orch(worker) + + in_control = threading.Event() + release_control = threading.Event() + submit_entered = threading.Event() + + def _blocking_native_malloc(*_args): + in_control.set() + assert release_control.wait(5.0), "test never released the control call" + return 0x1000 + + _native.malloc = _blocking_native_malloc + worker._submit_l3_locked = cast(Any, lambda *a: submit_entered.set()) + + control = threading.Thread(target=lambda: orch.malloc(0, 64), daemon=True) + control.start() + assert in_control.wait(5.0), "the control call never reached the child" + + submitter = threading.Thread(target=lambda: worker._submit_locked(lambda *a: None, None, None), daemon=True) + submitter.start() + assert not submit_entered.wait(0.5), "a run was admitted while a control call was still in flight" + + release_control.set() + control.join(5.0) + assert submit_entered.wait(5.0), "the submission stayed blocked after the control call finished" + submitter.join(5.0) + + def test_the_reservation_is_reentrant_within_one_thread(self): + """One control call can be built out of others — a queue out of a region. + + The serializer it takes is not re-entrant, so the inner call has to + join the outer reservation rather than deadlock on it. + """ + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + + with worker._control_reservation("outer"), worker._control_reservation("inner"): + pass + # And the reservation is released, not leaked, on the way out. + with worker._control_reservation("again"): + pass + + def test_a_run_on_one_worker_does_not_admit_control_on_another(self): + """A run id names nothing on its own — run 1 exists on every Worker. + + Inside Worker A's run, a call on Worker B is B's business: it must take + B's own reservation, not wait on whatever B happens to call run 1, and + certainly not skip B's admission because A has a run open. + """ + worker_a = Worker(level=3, num_sub_workers=0) + worker_a._worker = cast(Any, object()) + worker_b = Worker(level=3, num_sub_workers=0) + worker_b._worker = cast(Any, object()) + orch_b, native_b = self._orch(worker_b) + + # B has its own run 1 in flight — colliding id, unrelated run. + b_handle = RunHandle(worker_b, 1, (), worker_mod._RunResources()) + worker_b._accepted_run_handles.add(b_handle) + + with orch_mod._callback_run(1, worker_a): + with pytest.raises(RuntimeError, match="still in flight"): + orch_b.malloc(0, 64) + native_b.await_run_admission.assert_not_called() + + def test_a_reservation_on_one_worker_does_not_cover_another(self): + worker_a = Worker(level=3, num_sub_workers=0) + worker_a._worker = cast(Any, object()) + worker_b = Worker(level=3, num_sub_workers=0) + worker_b._worker = cast(Any, object()) + orch_b, _native_b = self._orch(worker_b) + + b_handle = RunHandle(worker_b, 1, (), worker_mod._RunResources()) + worker_b._accepted_run_handles.add(b_handle) + + with worker_a._control_reservation("Worker.malloc"): + with pytest.raises(RuntimeError, match="still in flight"): + orch_b.malloc(0, 64) + + def test_a_nested_worker_callback_keeps_its_callers_ordering(self): + """An L4 callback drives its child's run on its own thread. + + The inner frame belongs to the child, so control on the parent still + finds the parent's frame rather than falling through to a reservation + the parent's own open callback would deadlock on. + """ + parent = Worker(level=4, num_sub_workers=0) + parent._worker = cast(Any, object()) + child = Worker(level=3, num_sub_workers=0) + child._worker = cast(Any, object()) + orch_parent, native_parent = self._orch(parent) + orch_child, native_child = self._orch(child) + + with orch_mod._callback_run(5, parent): + with orch_mod._callback_run(9, child): + orch_child.malloc(0, 64) + orch_parent.malloc(0, 64) + native_child.await_run_admission.assert_called_once_with(9) + native_parent.await_run_admission.assert_called_once_with(5) + + def test_run_owned_control_rechecks_the_sticky_poison(self): + """A callback can catch a rollback failure and carry on. + + Holding the FIFO head says nothing about whether this worker still has + reclaimable device state, so the refusal is re-read on every call. + """ + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + orch, native = self._orch(worker) + + with orch_mod._callback_run(7, worker): + orch.malloc(0, 64) + worker._ordered_cleanup_error = RuntimeError("region rollback leaked") + with pytest.raises(RuntimeError, match="no further work is admitted"): + orch.malloc(0, 64) + assert native.malloc.call_count == 1 + + def test_task_submission_rechecks_the_sticky_poison(self): + """Same reason control does: the callback may have caught the failure. + + A run whose own graph construction leaked device state must not keep + putting work behind it. + """ + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + worker._ordered_cleanup_error = RuntimeError("region rollback leaked") + + with orch_mod._callback_run(7, worker): + with pytest.raises(RuntimeError, match="no further work is admitted"): + orch_mod._admit_task_submission(worker) + + def test_owner_less_control_is_refused_after_a_cleanup_failure(self): + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + orch, _native = self._orch(worker) + + worker._ordered_cleanup_error = RuntimeError("domain release failed") + with pytest.raises(RuntimeError, match="no further work is admitted"): + orch.malloc(0, 64) + + def test_a_mailbox_query_takes_the_same_ordering_as_a_command(self): + """`committed_device_memory` is a read, but it travels the same mailbox. + + Answered from behind a run that is still allocating, the number + describes neither the state before nor the state after. + """ + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + orch, native = self._orch(worker) + native.committed_device_memory.return_value = 4096 + + with orch_mod._callback_run(11, worker): + assert orch.committed_device_memory(0) == 4096 + native.await_run_admission.assert_called_once_with(11) + + handle = RunHandle(worker, 1, (), worker_mod._RunResources()) + worker._accepted_run_handles.add(handle) + with pytest.raises(RuntimeError, match="still in flight"): + orch.committed_device_memory(0) + + def test_owner_less_control_is_refused_while_a_run_is_in_flight(self): + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + orch, native = self._orch(worker) + + orch.malloc(0, 64) # quiescent: admitted + handle = RunHandle(worker, 1, (), worker_mod._RunResources()) + worker._accepted_run_handles.add(handle) + + with pytest.raises(RuntimeError, match="still in flight"): + orch.malloc(0, 64) + native.await_run_admission.assert_not_called() + + # A run whose cleanup has been published is no longer in flight, even + # though nothing removed the handle here. + handle._cleanup_published = True + orch.malloc(0, 64) + assert native.malloc.call_count == 2 + + +class TestRemoteControlOrdering: + """`remote_*` buffer commands and remote task dispatch both end up + contending for the endpoint's command mutex, so which arrives first is the + scheduler's choice. They take the same admission local control does. + """ + + @staticmethod + def _worker(): + from unittest.mock import MagicMock # noqa: PLC0415 + + worker = Worker(level=4, num_sub_workers=0) + native = MagicMock() + native.remote_malloc.return_value = (0, 1, 1, 2, 4, 0, 0, 0) + worker._worker = native + worker._lifecycle = worker_mod._Lifecycle.READY + worker._require_remote_worker_started = cast(Any, lambda _wid: None) + native_orch = MagicMock() + worker._orch = cast(Any, SimpleNamespace(_o=native_orch)) + return worker, native, native_orch + + def test_a_remote_command_inside_a_run_waits_for_the_fifo_head(self): + worker, _native, native_orch = self._worker() + with orch_mod._callback_run(3, worker): + worker.remote_malloc(worker=0, nbytes=4) + native_orch.await_run_admission.assert_called_once_with(3) + + def test_a_remote_command_after_a_task_submission_is_refused(self): + worker, native, _native_orch = self._worker() + with orch_mod._callback_run(3, worker): + orch_mod._admit_task_submission(worker) + with pytest.raises(RuntimeError, match="cannot follow a task submission"): + worker.remote_malloc(worker=0, nbytes=4) + + def test_an_owner_less_remote_command_is_refused_while_a_run_is_in_flight(self): + worker, native, _native_orch = self._worker() + worker.remote_malloc(worker=0, nbytes=4) # quiescent: admitted + + handle = RunHandle(worker, 1, (), worker_mod._RunResources()) + worker._accepted_run_handles.add(handle) + with pytest.raises(RuntimeError, match="still in flight"): + worker.remote_malloc(worker=0, nbytes=4) + + def test_a_deferred_free_sends_nothing_and_is_not_ordered(self): + """A free behind a live slot ref only records the debt. + + Nothing reaches the owner, so there is no command for admission to + order — and refusing it would break the ordinary shape of freeing an + input right after the task that reads it. + """ + worker, native, _native_orch = self._worker() + buffer = worker.remote_malloc(worker=0, nbytes=4) + buffer._acquire_slot_ref() + + with orch_mod._callback_run(3, worker): + orch_mod._admit_task_submission(worker) + worker.remote_free(buffer) + + assert buffer.released + assert worker._pending_remote_buffer_frees == [buffer] + native.remote_free.assert_not_called() + + +class TestUnreclaimedDeviceStateIsNeverSilent: + """Every path that leaves a resource on a chip either hands it to the run's + cleanup or refuses further work. Reporting a plain error over it lets the + next run start on top of state nothing can name. + """ + + @staticmethod + def _worker(): + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + return worker + + def test_a_partial_domain_allocation_is_handed_to_the_runs_cleanup(self): + """Two chips of three committed a window and no handle exists. + + The release has to reach exactly the two that allocated: driving it at + the third would fail on a debt it does not hold, and poison the worker + for a partial failure that was handled correctly. + """ + worker = self._worker() + resources = worker_mod._RunResources() + + worker._register_unclaimed_domain("d", (0, 2), 7, resources) + + handle = resources.live_domains["d"] + assert handle.workers == (0, 2), "release would reach a chip that never allocated" + assert handle.allocation_id == 7 + assert worker._live_domains["d"] is handle + assert resources.requires_ordered_cleanup, "a leaked window did not degrade the successor" + + def test_a_chip_that_committed_before_its_reply_failed_is_still_reclaimed(self): + """The window exists before the reply is written. + + A chip whose RPC failed after `comm_alloc_domain_windows` returned is + holding an allocation, so "which RPCs failed" is the wrong question — + each chip publishes its own commit and the parent reads that. + """ + reply = SharedMemory(create=True, size=worker_mod._DOMAIN_REPLY_HEADER.size) + try: + assert not worker_mod._domain_reply_committed(reply), "a zero-filled reply must not read as committed" + struct.pack_into(" None: + pass + + +class _WatchedSet(set): + """A set that reports each discard to `on_discard` before performing it.""" + + def __init__(self, source, on_discard): + super().__init__(source) + self._on_discard = on_discard + + def discard(self, item) -> None: + self._on_discard(item) + super().discard(item) + + +def _raiser(exc): + def _raise(*args, **kwargs): + raise exc + + return _raise + + # --------------------------------------------------------------------------- # Test: multiple SUB workers execute in parallel # ---------------------------------------------------------------------------