diff --git a/docs/l3-l2-orch-comm.md b/docs/l3-l2-orch-comm.md index 7ec6b76ed..e47640463 100644 --- a/docs/l3-l2-orch-comm.md +++ b/docs/l3-l2-orch-comm.md @@ -148,7 +148,25 @@ addresses. On onboard platforms, region create allocates one child-owned VMM GM range, exports it through a shareable handle, and returns the import metadata to the parent. The parent imports that region and closes the mapped VMM import before -the child frees the physical allocation. +the child frees the physical allocation. The native import remains +provisionally owned until the Python region is published into its run. An +interruption during import, wrapper construction, or run publication closes +the parent mapping and rolls back the child region instead of leaving either +resource untracked. + +Every native payload, counter, and wait operation holds a lease on that parent +mapping for the complete GIL-released access. Closing first rejects new leases, +then waits for all active leases and the physical unmap to finish before it +returns; concurrent duplicate closes join that same completion. The Worker can +therefore release the child allocation only after no parent operation can still +touch it. + +If an unadopted native owner reports a cleanup failure while an import unwinds, +the diagnostic stays on the importing thread until the Worker consumes it and +poisons itself. It cannot be consumed by a Worker importing concurrently on a +different thread. If the diagnostic survives to a later create on the same +thread, its original Worker is no longer identifiable, so that Worker is +conservatively stopped with an explicit attribution message. ## 4. Signal Counters diff --git a/docs/orchestrator.md b/docs/orchestrator.md index 00d5dab8c..6bbb110ca 100644 --- a/docs/orchestrator.md +++ b/docs/orchestrator.md @@ -86,8 +86,10 @@ device or endpoint errors are attached to the handle and raised by `wait()` or bounded waits without cancelling or corrupting the run, and repeated waits replay the same terminal result. The handle keeps its `Worker`, callback arguments, configuration, and run-owned cleanup state alive until completion. -`Worker.close()` rejects new submissions and drains every accepted handle -before tearing down the worker tree. +`Worker.close()` rejects new submissions and drains operation leases plus +accepted handles within one cleanup budget before tearing down the worker +tree. If that budget expires, teardown remains unattempted and the closed +worker retains the handles/tree for a later `close()` retry. `Worker.run` remains source-compatible and blocking: @@ -99,10 +101,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 +165,32 @@ 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) - scope_.register_task(sid); // increments s.fanout_total by 1 + // 4. Register with scope (holds slot open until scope_end releases ref) + scope_.register_task(sid); + { + std::lock_guard lk(s.fanout_mu); + s.fanout_total += 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}; @@ -213,22 +227,30 @@ register this task as the new producer of the tensor's dependency key. For local tensors this key contains `tensor.data`; for remote sidecars it contains remote buffer identity and logical offset. -**Step 4 — fanin count**: The number of live producers. Decremented by -`fanin_released++` each time a producer completes; when `fanin_released == -fanin_count`, the slot is ready. A ready NEXT_LEVEL single task is routed to -the FIFO for its required stable worker id. The same routing function is used -for immediately-ready submissions and tasks released by Scheduler dependency -processing. A ready NEXT_LEVEL group is routed to the dedicated group FIFO; -SUB tasks remain on their shared queue. +Before each output mapping becomes visible, its key is appended to the slot's +cleanup journal. A failed journal append therefore publishes no mapping, while +a later failure leaves a key `on_consumed` can erase. Erasure remains +owner-checked: if publication failed before replacing an older producer, cleanup +of the failed slot does not remove that older producer's mapping. -**Step 5 — scope ref**: Each slot starts with one "scope reference" in its -fanout_total. Without this, a task with no downstream consumer would never be -reclaimable. See [§6 Scope](#6-scope). +**Step 4 — scope ref**: A slot submitted inside an open scope registers one +scope reference in `fanout_total`. Registration precedes the charge, so a +failed registration cannot leave a reference that no `scope_end` can release. +See [§6 Scope](#6-scope). -**Step 6 — fanout attachment and READY routing**: Submission synchronously +**Step 5 — fanout attachment and live-fanin count**: Submission synchronously locks each producer's `fanout_mu`, attaches the consumer, and counts only live -producers. An immediately READY task is routed to its exact NEXT_LEVEL worker -FIFO, the NEXT_LEVEL group FIFO, or the shared SUB FIFO. See +producers. The producer's forward edge and the consumer's reverse edge are one +transaction: failure to publish the reverse edge rolls the forward edge and +its reference charge back. A completing live producer advances +`fanin_released`; completed producers retain the reference edge but do not +contribute to `fanin_count`. + +**Step 6 — publication and READY routing**: `fanin_count` and the transition +from BUILDING to PENDING or READY are published together under `fanout_mu`. An +immediately READY task is routed to its exact NEXT_LEVEL worker FIFO, the +NEXT_LEVEL group FIFO, or the shared SUB FIFO. The same routing function is +used when Scheduler releases the final dependency. See [scheduler.md](scheduler.md) §1. --- @@ -259,7 +281,11 @@ SubmitResult Orchestrator::submit_next_level_group( `submit_impl` validates that `worker_ids` contains one unique, eligible target per group member before it performs shared dependency inference and READY -routing. +routing. It also prepares the member-state and member-outcome vectors as one +transaction while the slot is still BUILDING. Dispatch therefore never has to +allocate after claiming READY → RUNNING. A defensive size repair is likewise +prepared in local vectors before the claim and published only if the Scheduler +wins; a losing claim cannot overwrite cancellation's terminal bookkeeping. At dispatch time the Scheduler checks the group FIFO head and resolves every entry in `workers` to that exact stable worker ID. It dispatches only if the @@ -267,7 +293,8 @@ entire target set is idle. A blocked group reserves all of its targets against new singles but does not cause a scan past the FIFO head. Each WorkerThread runs `worker->run` with its own `task_args_list[i]`. Completion remains aggregated at the group slot, so downstream consumers are released once after -every member is terminal. +every member is terminal. Completion validates both bookkeeping-vector sizes; +repair preserves already-terminal and still-running members before indexing. --- @@ -556,13 +583,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 +605,104 @@ 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. + +Cancellation also records debt immediately after it wins a PENDING or READY +claim. Scheduler-owned PENDING/READY claims expose no takeover debt: their +claimant propagates exclusively, so cancellation cannot race it over stale slot +IDs. + +Debt settlement is not the start of cancellation propagation. Cancellation +first marks groups and snapshots the producer lists for **every** slot it owns; +all of that work is idempotent and may allocate. An exception therefore leaves +each claimed slot's debt set and any unvisited slot claimable by the next +attempt. Only after every snapshot succeeds does cancellation clear the debts, +then perform the non-repeatable self and producer reference releases without +any further allocation. + +### What a submit that throws leaves behind + +A submit publishes BUILDING before its slot is charged to the run, so a throw +part-way has to leave either an unowned Ring slot it releases directly or a +registered slot the run's cancellation can fully reclaim. These ownership +rules carry that: + +- **Run-slot registration is the ownership boundary.** If registration throws, + the slot is absent from the run and cancellation cannot discover it, so + submit marks it CONSUMED and directly releases the combined slot/HeapRing + reservation before rethrowing. Once registration succeeds, all remaining + fallible construction stays BUILDING and cancellation owns rollback. + +- **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. +- **The charge immediately follows successful registration** and precedes Step + 2 publishing the slot's outputs. A failure between registration and charge + leaves an extra scope release, which is safe; the opposite order would leave + an unreleasable charge. Once an output is published, downstream consumers + increment the same `fanout_total` field. +- **The output cleanup journal precedes TensorMap publication.** A failure can + never leave a mapping cancellation does not know to erase, and the map's + owner check preserves a previous producer when the replacement insert never + completed. + +### 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 @@ -614,29 +738,40 @@ Tensor Orchestrator::alloc(const std::vector &shape, DataType dtype) { AllocResult ar = allocator_.alloc(aligned); TaskSlotState &s = slots_[ar.slot]; s.reset(); - // 2. Register as this slot's output so downstream tensors with the same - // data pointer look up this slot as producer. + // 2. Publish cancellation ownership before registering active_tasks. If + // run-slot registration itself throws, release the unowned Ring slot. uint64_t key = reinterpret_cast(ar.heap_ptr); s.run_id = current_run_id; - tensormap_.insert(current_run_id, key, ar.slot); - s.output_keys.push_back(key); + s.state = TaskState::BUILDING; + current_run.register_slot(ar.slot); // 3. No fanin — alloc has no work to wait on. s.fanin_count = 0; - // 4. Initial fanout = scope_ref. Consumers that wire on this slot in - // infer_deps bump fanout_total; this slot's CONSUMED transition waits - // for all of them + scope_end. - s.fanout_total = (scope_.depth() > 0) ? 1 : 0; - if (s.fanout_total > 0) scope_.register_task(ar.slot); - // 5. Sim self-consume so the fanout-release threshold math aligns with + // 4. Register the scope reference before charging it. Consumers that wire + // on this slot in infer_deps later increment fanout_total. + int32_t scope_ref = (scope_.depth() > 0) ? 1 : 0; + if (scope_ref > 0) scope_.register_task(ar.slot); + s.fanout_total = scope_ref; + // 5. Journal the key before TensorMap publication so cancellation can + // erase every entry that became visible before a later failure. + s.output_keys.push_back(key); + tensormap_.insert(current_run_id, key, ar.slot); + // 6. Sim self-consume so the fanout-release threshold math aligns with // normal slots (see §8 Fanout-release threshold). s.fanout_released = 1; - // 6. Straight to COMPLETED — no dispatch needed. + // 7. Straight to COMPLETED — no dispatch needed. s.state = TaskState::COMPLETED; - current_run.active_tasks++; return Tensor{key, shape, dtype}; } ``` +All fallible setup after run registration happens while the synthetic slot is +BUILDING. If `alloc()` throws, graph-failure cancellation can therefore claim +and consume it just like an interrupted task submit. The synthetic self-release +is published only after those fallible steps; a failed BUILDING slot gets that +terminal release from cancellation instead. Run-slot registration is the one +earlier boundary cancellation cannot cover, so failure there directly releases +the Ring claim before rethrowing. + `on_consumed` runs the usual `tensormap.erase_task_outputs` and then calls `allocator_.release(sid)`. FIFO reclamation inside the allocator returns the slab to the heap's free region as `last_alive` advances; callers see no 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..f6d2c0a92 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. @@ -399,10 +444,11 @@ Every dispatched group member contributes one run-acceptance obligation. For an A2A3 onboard chip endpoint, the child-side native runner writes `TASK_ACCEPTED` after its AICore and AICPU kernels are both enqueued; the parent observes it without releasing the mailbox. Other endpoint paths satisfy the -same obligation conservatively when their completion returns. Once submission -is closed and all obligations are satisfied, the next serialized orchestration -callback may build its DAG even though the prior run has not reached its -completion fence. +same obligation conservatively when their completion returns. Acceptance is the +launch fence for that run's own dispatches; it does not admit the next graph +callback. Once the prior callback returns, `begin_run` may invoke the next +serialized callback whenever a negotiated pipeline lease is free, even while +the prior run remains below its acceptance or completion fence. Local mailbox path: diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index b89f523d2..1e2507d71 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -32,8 +32,10 @@ #include #include #include +#include #include #include +#include #include #include #include @@ -303,18 +305,57 @@ class AclRuntimeApi { AclRuntimeApi &acl_api() { static std::once_flag once; - static std::unique_ptr api; + static AclRuntimeApi *api{nullptr}; std::call_once(once, []() { auto candidate = std::make_unique(); candidate->load(); candidate->init(); - api = std::move(candidate); + api = candidate.release(); }); return *api; } +class L3HostMappedRegionCleanupErrors { +public: + void record(const std::string &message) noexcept { + try { + append_cleanup_error(error_, message); + } catch (...) {} + } + + std::string take() { return std::exchange(error_, {}); } + +private: + std::string error_; +}; + +L3HostMappedRegionCleanupErrors &l3_host_mapped_region_cleanup_errors() { + // A return-boundary owner destroyed by Python while unwinding an import is + // finalized on the importing thread. Keep its diagnostic thread-local so + // concurrent Workers cannot consume and misattribute one another's error. + // Leak the tiny sink to remain usable during late Python finalization. + static thread_local auto *errors = new L3HostMappedRegionCleanupErrors(); + return *errors; +} + class L3HostMappedRegion { public: + L3HostMappedRegion() = default; + L3HostMappedRegion(const L3HostMappedRegion &) = delete; + L3HostMappedRegion &operator=(const L3HostMappedRegion &) = delete; + + ~L3HostMappedRegion() noexcept { + try { + std::string cleanup_error; + close_collecting(cleanup_error); + if (!cleanup_error.empty()) { + l3_host_mapped_region_cleanup_errors().record(cleanup_error); + } + } catch (...) { + l3_host_mapped_region_cleanup_errors().record("L3-L2 mapped-region cleanup failed with an unknown error"); + } + } + L3L2RegionAccessProfile profile{L3L2RegionAccessProfile::SIM_POSIX_SHM}; int fd{-1}; uint64_t device_addr{0}; @@ -323,6 +364,14 @@ class L3HostMappedRegion { void *vmm_handle{nullptr}; uint64_t mapping_bytes{0}; + void close() { + std::string cleanup_error; + close_collecting(cleanup_error); + if (!cleanup_error.empty()) { + throw std::runtime_error(cleanup_error); + } + } + void bind_acl_device() const { if (device_id < 0) { throw std::runtime_error("L3-L2 onboard mapped-region handle has no device id"); @@ -420,6 +469,49 @@ class L3HostMappedRegion { } private: + void close_collecting(std::string &cleanup_error) { + uint64_t mapped_addr = std::exchange(device_addr, 0); + uint64_t mapped_bytes = std::exchange(mapping_bytes, 0); + void *physical_handle = std::exchange(vmm_handle, nullptr); + int mapped_device_id = std::exchange(device_id, -1); + int mapped_fd = std::exchange(fd, -1); + + if (profile == L3L2RegionAccessProfile::ONBOARD_VMM) { + if (mapped_addr == 0 && physical_handle == nullptr) { + return; + } + try { + if (mapped_device_id < 0) { + throw std::runtime_error("L3-L2 onboard mapped-region handle has no device id"); + } + AclRuntimeApi &api = acl_api(); + api.bind_device_with_check(mapped_device_id); + api.vmm_release_collecting( + reinterpret_cast(static_cast(mapped_addr)), physical_handle, cleanup_error + ); + } catch (const std::exception &exc) { + append_cleanup_error(cleanup_error, exc.what()); + } catch (...) { + append_cleanup_error(cleanup_error, "L3-L2 onboard mapped-region cleanup failed"); + } + return; + } + + if (mapped_addr != 0 && + munmap(reinterpret_cast(static_cast(mapped_addr)), mapped_bytes) != 0) { + int err = errno; + append_cleanup_error( + cleanup_error, std::string("L3-L2 sim L3 Host mapped-region munmap failed: ") + std::strerror(err) + ); + } + if (mapped_fd >= 0 && ::close(mapped_fd) != 0) { + int err = errno; + append_cleanup_error( + cleanup_error, std::string("L3-L2 sim L3 Host mapped-region close failed: ") + std::strerror(err) + ); + } + } + static constexpr int kWaitStatusOk = 0; static constexpr int kWaitStatusTimeout = -1; static constexpr int kWaitErrorNone = 0; @@ -450,42 +542,233 @@ struct L2ChildOnboardRegionExport { uint64_t registry_handle{0}; }; +class L3HostMappedRegionEntry { +public: + explicit L3HostMappedRegionEntry(std::unique_ptr mapping) : + mapping_(std::move(mapping)) {} + + void acquire() { + std::lock_guard lk(mu_); + if (state_ != State::OPEN) { + throw std::runtime_error("L3-L2 L3 Host mapped-region handle is closed or unknown"); + } + active_leases_ += 1; + } + + void release() noexcept { + std::lock_guard lk(mu_); + if (active_leases_ == 0) { + return; + } + active_leases_ -= 1; + if (active_leases_ == 0) { + idle_.notify_all(); + } + } + + L3HostMappedRegion &mapping() { return *mapping_; } + + size_t active_leases() const { + std::lock_guard lk(mu_); + return active_leases_; + } + + void close() { + std::unique_ptr mapping; + std::exception_ptr close_error; + { + std::unique_lock lk(mu_); + if (state_ != State::OPEN) { + idle_.wait(lk, [this]() { + return state_ == State::CLOSED; + }); + close_error = close_error_; + lk.unlock(); + if (close_error != nullptr) { + std::rethrow_exception(close_error); + } + return; + } + state_ = State::CLOSING; + idle_.wait(lk, [this]() { + return active_leases_ == 0; + }); + mapping = std::move(mapping_); + } + + try { + if (mapping != nullptr) { + mapping->close(); + } + } catch (...) { + close_error = std::current_exception(); + } + { + std::lock_guard lk(mu_); + close_error_ = close_error; + state_ = State::CLOSED; + } + idle_.notify_all(); + if (close_error != nullptr) { + std::rethrow_exception(close_error); + } + } + +private: + enum class State { OPEN, CLOSING, CLOSED }; + + std::unique_ptr mapping_; + mutable std::mutex mu_; + std::condition_variable idle_; + size_t active_leases_{0}; + State state_{State::OPEN}; + std::exception_ptr close_error_; +}; + +class L3HostMappedRegionLease { +public: + explicit L3HostMappedRegionLease(std::shared_ptr entry) : + entry_(std::move(entry)) { + entry_->acquire(); + } + L3HostMappedRegionLease(const L3HostMappedRegionLease &) = delete; + L3HostMappedRegionLease &operator=(const L3HostMappedRegionLease &) = delete; + L3HostMappedRegionLease(L3HostMappedRegionLease &&) noexcept = default; + L3HostMappedRegionLease &operator=(L3HostMappedRegionLease &&) = delete; + ~L3HostMappedRegionLease() { + if (entry_ != nullptr) { + entry_->release(); + } + } + + L3HostMappedRegion *operator->() { return &entry_->mapping(); } + +private: + std::shared_ptr entry_; +}; + class L3HostMappedRegionRegistry { public: - uint64_t emplace(L3HostMappedRegion mapping) { + uint64_t emplace(std::unique_ptr mapping) { + auto entry = std::make_shared(std::move(mapping)); std::lock_guard lk(mu_); - uint64_t handle = next_handle_++; - regions_.emplace(handle, std::move(mapping)); + if (std::exchange(fail_next_insert_for_test_, false)) { + throw std::runtime_error("injected mapped-region registry insertion failure"); + } + uint64_t handle = next_handle_; + auto result = regions_.emplace(handle, std::move(entry)); + if (!result.second) { + throw std::overflow_error("L3-L2 L3 Host mapped-region handle space is exhausted"); + } + next_handle_ += 1; + if (next_handle_ == 0) { + next_handle_ = 1; + } return handle; } - L3HostMappedRegion find(uint64_t handle) const { + L3HostMappedRegionLease lease(uint64_t handle) const { std::lock_guard lk(mu_); auto it = regions_.find(handle); if (it == regions_.end()) { throw std::runtime_error("L3-L2 L3 Host mapped-region handle is closed or unknown"); } - return it->second; + return L3HostMappedRegionLease(it->second); } - std::optional remove(uint64_t handle) { + size_t active_leases(uint64_t handle) const { std::lock_guard lk(mu_); auto it = regions_.find(handle); if (it == regions_.end()) { - return std::nullopt; + throw std::runtime_error("L3-L2 L3 Host mapped-region handle is closed or unknown"); + } + return it->second->active_leases(); + } + + void close(uint64_t handle) { + // Retain the entry while it closes: its state rejects new leases, and + // duplicate close callers join the same physical-cleanup completion. + std::shared_ptr entry; + { + std::lock_guard lk(mu_); + auto it = regions_.find(handle); + if (it == regions_.end()) { + return; + } + entry = it->second; + } + + std::exception_ptr close_error; + try { + entry->close(); + } catch (...) { + close_error = std::current_exception(); + } + + { + std::lock_guard lk(mu_); + auto it = regions_.find(handle); + if (it != regions_.end() && it->second == entry) { + regions_.erase(it); + } + } + if (close_error != nullptr) { + std::rethrow_exception(close_error); } - L3HostMappedRegion mapping = std::move(it->second); - regions_.erase(it); - return mapping; + } + + void fail_next_insert_for_test() { + std::lock_guard lk(mu_); + fail_next_insert_for_test_ = true; } private: mutable std::mutex mu_; - std::unordered_map regions_; + std::unordered_map> regions_; uint64_t next_handle_{1}; + bool fail_next_insert_for_test_{false}; }; -L3HostMappedRegionRegistry g_l3_host_mapped_regions; +L3HostMappedRegionRegistry &l3_host_mapped_region_registry() { + // Python owners may be finalized after ordinary C++ static destruction has + // begun. The registry and ACL dispatch table therefore have process + // lifetime; the OS reclaims any entries still open at process exit. + static auto *registry = new L3HostMappedRegionRegistry(); + return *registry; +} + +void close_l3_host_mapped_region(uint64_t handle) { l3_host_mapped_region_registry().close(handle); } + +class L3HostMappedRegionHandle { +public: + explicit L3HostMappedRegionHandle(uint64_t handle) : + handle_(handle) {} + L3HostMappedRegionHandle(const L3HostMappedRegionHandle &) = delete; + L3HostMappedRegionHandle &operator=(const L3HostMappedRegionHandle &) = delete; + L3HostMappedRegionHandle(L3HostMappedRegionHandle &&other) noexcept : + handle_(std::exchange(other.handle_, 0)) {} + L3HostMappedRegionHandle &operator=(L3HostMappedRegionHandle &&) = delete; + + ~L3HostMappedRegionHandle() noexcept { + if (handle_ == 0) { + return; + } + try { + close_l3_host_mapped_region(handle_); + } catch (const std::exception &exc) { + l3_host_mapped_region_cleanup_errors().record(exc.what()); + } catch (...) { + l3_host_mapped_region_cleanup_errors().record( + "L3-L2 mapped-region owner cleanup failed with an unknown error" + ); + } + } + + uint64_t value() const { return handle_; } + +private: + uint64_t handle_{0}; +}; class L2ChildOnboardRegionRegistry { public: @@ -1570,15 +1853,22 @@ NB_MODULE(_task_interface, m) { .def( "comm_alloc_domain_windows", [](ChipWorker &self, uint64_t comm_handle, uint64_t allocation_id, const std::vector &rank_ids, - uint32_t domain_rank, size_t window_size) { + uint32_t domain_rank, size_t window_size, uint64_t commit_flag_address) { + if (commit_flag_address == 0 || commit_flag_address % alignof(uint64_t) != 0) { + throw std::invalid_argument("comm_alloc_domain_windows: commit flag address is invalid"); + } auto [device_ctx, local_window_base] = self.comm_alloc_domain_windows(comm_handle, allocation_id, rank_ids, domain_rank, window_size); + __atomic_store_n( + reinterpret_cast(static_cast(commit_flag_address)), uint64_t{1}, + __ATOMIC_RELEASE + ); return nb::make_tuple(device_ctx, local_window_base); }, nb::arg("comm_handle"), nb::arg("allocation_id"), nb::arg("rank_ids"), nb::arg("domain_rank"), - nb::arg("window_size"), + nb::arg("window_size"), nb::arg("commit_flag_address"), "Collectively allocate a fresh per-rank pool for a subset; returns " - "(device_ctx, local_window_base) for this rank." + "(device_ctx, local_window_base) for this rank and publishes the commit flag before result conversion." ) .def( "comm_release_domain_windows", &ChipWorker::comm_release_domain_windows, nb::arg("comm_handle"), @@ -1617,64 +1907,59 @@ NB_MODULE(_task_interface, m) { .def_ro("shareable_handle", &L2ChildOnboardRegionExport::shareable_handle) .def_ro("registry_handle", &L2ChildOnboardRegionExport::registry_handle); + nb::class_(m, "_L3HostMappedRegionHandle") + .def("__int__", &L3HostMappedRegionHandle::value); + m.def( "_l3_host_mapped_region_import_sim", - [](const std::string &token, uint64_t mapping_bytes) -> uint64_t { + [](const std::string &token, uint64_t mapping_bytes) -> L3HostMappedRegionHandle { if (mapping_bytes == 0 || mapping_bytes > static_cast(std::numeric_limits::max())) { throw std::invalid_argument("L3-L2 sim L3 Host mapped-region import requires a positive mapping size"); } std::string name = shm_name_for_open(token); - int fd = shm_open(name.c_str(), O_RDWR, 0); - if (fd < 0) { + auto mapping = std::make_unique(); + mapping->fd = shm_open(name.c_str(), O_RDWR, 0); + if (mapping->fd < 0) { throw std::runtime_error("L3-L2 sim L3 Host mapped-region import shm_open failed"); } - void *base = mmap(nullptr, static_cast(mapping_bytes), PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0); + void *base = + mmap(nullptr, static_cast(mapping_bytes), PROT_READ | PROT_WRITE, MAP_SHARED, mapping->fd, 0); if (base == MAP_FAILED) { int err = errno; - ::close(fd); throw std::runtime_error( std::string("L3-L2 sim L3 Host mapped-region import mmap failed: ") + std::strerror(err) ); } - L3HostMappedRegion mapping{}; - mapping.profile = L3L2RegionAccessProfile::SIM_POSIX_SHM; - mapping.fd = fd; - mapping.device_addr = reinterpret_cast(base); - mapping.mapping_bytes = mapping_bytes; - return g_l3_host_mapped_regions.emplace(mapping); + mapping->profile = L3L2RegionAccessProfile::SIM_POSIX_SHM; + mapping->device_addr = reinterpret_cast(base); + mapping->mapping_bytes = mapping_bytes; + return L3HostMappedRegionHandle(l3_host_mapped_region_registry().emplace(std::move(mapping))); }, nb::arg("token"), nb::arg("mapping_bytes"), nb::call_guard(), "Import a sim L3-L2 POSIX shm region for L3 Host mapped-region access." ); m.def( "_l3_host_mapped_region_import_onboard", - [](int device_id, uint64_t shareable_handle, uint64_t mapping_bytes) -> uint64_t { + [](int device_id, uint64_t shareable_handle, uint64_t mapping_bytes) -> L3HostMappedRegionHandle { if (device_id < 0) { throw std::invalid_argument("L3-L2 onboard mapped-region import requires a non-negative device id"); } if (mapping_bytes == 0 || mapping_bytes > static_cast(std::numeric_limits::max())) { throw std::invalid_argument("L3-L2 onboard mapped-region import requires a positive mapping size"); } - L3HostMappedRegion mapping{}; - mapping.profile = L3L2RegionAccessProfile::ONBOARD_VMM; - mapping.device_id = device_id; - mapping.mapping_bytes = mapping_bytes; - mapping.bind_acl_device(); + auto mapping = std::make_unique(); + mapping->profile = L3L2RegionAccessProfile::ONBOARD_VMM; + mapping->device_id = device_id; + mapping->mapping_bytes = mapping_bytes; + mapping->bind_acl_device(); AclRuntimeApi &api = acl_api(); - mapping.shareable_handle = shareable_handle; - mapping.vmm_handle = api.vmm_import_shareable_with_check(shareable_handle, device_id); - void *mapped_addr = nullptr; - try { - mapped_addr = api.vmm_reserve_with_check(mapping_bytes); - api.vmm_map_with_check(mapped_addr, mapping_bytes, mapping.vmm_handle); - api.vmm_set_access_with_check(mapped_addr, mapping_bytes, device_id); - } catch (...) { - std::string cleanup_error; - api.vmm_release_collecting(mapped_addr, mapping.vmm_handle, cleanup_error); - throw; - } - mapping.device_addr = reinterpret_cast(mapped_addr); - return g_l3_host_mapped_regions.emplace(mapping); + mapping->shareable_handle = shareable_handle; + mapping->vmm_handle = api.vmm_import_shareable_with_check(shareable_handle, device_id); + void *mapped_addr = api.vmm_reserve_with_check(mapping_bytes); + mapping->device_addr = reinterpret_cast(mapped_addr); + api.vmm_map_with_check(mapped_addr, mapping_bytes, mapping->vmm_handle); + api.vmm_set_access_with_check(mapped_addr, mapping_bytes, device_id); + return L3HostMappedRegionHandle(l3_host_mapped_region_registry().emplace(std::move(mapping))); }, nb::arg("device_id"), nb::arg("shareable_handle"), nb::arg("mapping_bytes"), nb::call_guard(), "Import an onboard VMM L3-L2 region for L3 Host mapped-region access." @@ -1682,53 +1967,39 @@ NB_MODULE(_task_interface, m) { m.def( "_l3_host_mapped_region_close", [](uint64_t handle) { - std::optional removed = g_l3_host_mapped_regions.remove(handle); - if (!removed.has_value()) { - return; - } - L3HostMappedRegion mapping = *removed; - if (mapping.profile == L3L2RegionAccessProfile::ONBOARD_VMM) { - mapping.bind_acl_device(); - std::string cleanup_error; - acl_api().vmm_release_collecting( - reinterpret_cast(static_cast(mapping.device_addr)), mapping.vmm_handle, - cleanup_error - ); - if (!cleanup_error.empty()) { - throw std::runtime_error(cleanup_error); - } - return; - } - - std::string cleanup_error; - if (mapping.device_addr != 0 && - munmap(reinterpret_cast(static_cast(mapping.device_addr)), mapping.mapping_bytes) != - 0) { - int err = errno; - append_cleanup_error( - cleanup_error, std::string("L3-L2 sim L3 Host mapped-region munmap failed: ") + std::strerror(err) - ); - } - if (mapping.fd >= 0 && ::close(mapping.fd) != 0) { - int err = errno; - append_cleanup_error( - cleanup_error, std::string("L3-L2 sim L3 Host mapped-region close failed: ") + std::strerror(err) - ); - } - if (!cleanup_error.empty()) { - throw std::runtime_error(cleanup_error); - } + close_l3_host_mapped_region(handle); }, nb::arg("handle"), nb::call_guard(), "Close an L3 Host mapped-region handle." ); + m.def( + "_l3_host_mapped_region_active_leases", + [](uint64_t handle) { + return l3_host_mapped_region_registry().active_leases(handle); + }, + nb::arg("handle"), "Return the number of in-flight native operations holding this mapped region." + ); + m.def( + "_l3_host_mapped_region_take_cleanup_error", + []() { + return l3_host_mapped_region_cleanup_errors().take(); + }, + "Take a cleanup error recorded by an unadopted native mapped-region owner on this thread." + ); + m.def( + "_l3_host_mapped_region_fail_next_registry_insert_for_test", + []() { + l3_host_mapped_region_registry().fail_next_insert_for_test(); + }, + "Inject one mapped-region registry insertion failure after native acquisition." + ); m.def( "_l3_host_mapped_payload_write", [](uint64_t handle, uint64_t payload_offset, uint64_t host_ptr, uint64_t nbytes) { if (host_ptr == 0) { throw std::invalid_argument("L3-L2 payload_write host_ptr must be nonzero"); } - L3HostMappedRegion mapping = g_l3_host_mapped_regions.find(handle); - mapping.copy_to(payload_offset, reinterpret_cast(static_cast(host_ptr)), nbytes); + L3HostMappedRegionLease mapping = l3_host_mapped_region_registry().lease(handle); + mapping->copy_to(payload_offset, reinterpret_cast(static_cast(host_ptr)), nbytes); }, nb::arg("handle"), nb::arg("payload_offset"), nb::arg("host_ptr"), nb::arg("nbytes"), nb::call_guard(), "Copy L3 Host bytes into an imported L3-L2 payload range." @@ -1739,8 +2010,8 @@ NB_MODULE(_task_interface, m) { if (host_ptr == 0) { throw std::invalid_argument("L3-L2 payload_read host_ptr must be nonzero"); } - L3HostMappedRegion mapping = g_l3_host_mapped_regions.find(handle); - mapping.copy_from(reinterpret_cast(static_cast(host_ptr)), payload_offset, nbytes); + L3HostMappedRegionLease mapping = l3_host_mapped_region_registry().lease(handle); + mapping->copy_from(reinterpret_cast(static_cast(host_ptr)), payload_offset, nbytes); }, nb::arg("handle"), nb::arg("payload_offset"), nb::arg("host_ptr"), nb::arg("nbytes"), nb::call_guard(), "Copy imported L3-L2 payload bytes into L3 Host memory." @@ -1748,8 +2019,8 @@ NB_MODULE(_task_interface, m) { m.def( "_l3_host_mapped_counter_notify", [](uint64_t handle, uint64_t counter_offset, int32_t value, int op) { - L3HostMappedRegion mapping = g_l3_host_mapped_regions.find(handle); - mapping.notify_counter(counter_offset, value, checked_notify_op(op)); + L3HostMappedRegionLease mapping = l3_host_mapped_region_registry().lease(handle); + mapping->notify_counter(counter_offset, value, checked_notify_op(op)); }, nb::arg("handle"), nb::arg("counter_offset"), nb::arg("value"), nb::arg("op"), nb::call_guard(), "Store or add one L3 Host-side L3-L2 signal counter." @@ -1757,8 +2028,8 @@ NB_MODULE(_task_interface, m) { m.def( "_l3_host_mapped_counter_test", [](uint64_t handle, uint64_t counter_offset, int32_t operand, int cmp) -> std::tuple { - L3HostMappedRegion mapping = g_l3_host_mapped_regions.find(handle); - return mapping.test_counter(counter_offset, operand, checked_wait_cmp(cmp)); + L3HostMappedRegionLease mapping = l3_host_mapped_region_registry().lease(handle); + return mapping->test_counter(counter_offset, operand, checked_wait_cmp(cmp)); }, nb::arg("handle"), nb::arg("counter_offset"), nb::arg("operand"), nb::arg("cmp"), nb::call_guard(), "Load and compare one L3 Host-side L3-L2 signal counter." @@ -1767,8 +2038,8 @@ NB_MODULE(_task_interface, m) { "_l3_host_mapped_counter_wait", [](uint64_t handle, uint64_t counter_offset, int32_t operand, int cmp, uint64_t timeout_ns) -> std::tuple { - L3HostMappedRegion mapping = g_l3_host_mapped_regions.find(handle); - return mapping.wait_counter(counter_offset, operand, checked_wait_cmp(cmp), timeout_ns); + L3HostMappedRegionLease mapping = l3_host_mapped_region_registry().lease(handle); + return mapping->wait_counter(counter_offset, operand, checked_wait_cmp(cmp), timeout_ns); }, nb::arg("handle"), nb::arg("counter_offset"), nb::arg("operand"), nb::arg("cmp"), nb::arg("timeout_ns"), nb::call_guard(), "Poll one L3 Host-side L3-L2 signal counter until match or timeout." 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/l3_l2_orch_comm.py b/python/simpler/l3_l2_orch_comm.py index 4cb556606..e4c89fa0e 100644 --- a/python/simpler/l3_l2_orch_comm.py +++ b/python/simpler/l3_l2_orch_comm.py @@ -150,7 +150,7 @@ class L3HostRegionMapping: payload_bytes: int counter_offset: int counter_bytes: int - handle: int + handle: Any closed: bool = False def close(self) -> None: @@ -337,7 +337,7 @@ def payload_write(self, offset: int, host_buffer: Any, nbytes: int | None = None size = pinned.nbytes if nbytes is None else int(nbytes) self._validate_payload_range(offset, size, pinned.nbytes) try: - _l3_host_mapped_payload_write(self._l3_host_mapping.handle, int(offset), pinned.addr, size) + _l3_host_mapped_payload_write(int(self._l3_host_mapping.handle), int(offset), pinned.addr, size) except Exception: self._poison() raise @@ -348,7 +348,7 @@ def payload_read(self, offset: int, host_buffer: Any, nbytes: int | None = None) size = pinned.nbytes if nbytes is None else int(nbytes) self._validate_payload_range(offset, size, pinned.nbytes) try: - _l3_host_mapped_payload_read(self._l3_host_mapping.handle, int(offset), pinned.addr, size) + _l3_host_mapped_payload_read(int(self._l3_host_mapping.handle), int(offset), pinned.addr, size) except Exception: self._poison() raise @@ -402,7 +402,7 @@ def _direct_counter_notify(self, offset: int, value: int, op: NotifyOp) -> None: l3_host_mapping = self._l3_host_mapping mapping_offset = int(l3_host_mapping.counter_offset) + int(offset) try: - _l3_host_mapped_counter_notify(l3_host_mapping.handle, mapping_offset, int(value), int(op)) + _l3_host_mapped_counter_notify(int(l3_host_mapping.handle), mapping_offset, int(value), int(op)) except Exception: self._poison() raise @@ -412,7 +412,7 @@ def _direct_counter_test(self, offset: int, cmp_value: int, cmp: WaitCmp) -> Sig mapping_offset = int(l3_host_mapping.counter_offset) + int(offset) try: matched, observed = _l3_host_mapped_counter_test( - l3_host_mapping.handle, mapping_offset, int(cmp_value), int(cmp) + int(l3_host_mapping.handle), mapping_offset, int(cmp_value), int(cmp) ) except Exception: self._poison() @@ -424,7 +424,7 @@ def _direct_counter_wait(self, offset: int, cmp_value: int, cmp: WaitCmp, timeou mapping_offset = int(l3_host_mapping.counter_offset) + int(offset) try: status, error_kind, observed, _matched, message = _l3_host_mapped_counter_wait( - l3_host_mapping.handle, mapping_offset, int(cmp_value), int(cmp), int(timeout_ns) + int(l3_host_mapping.handle), mapping_offset, int(cmp_value), int(cmp), int(timeout_ns) ) except Exception: self._poison() diff --git a/python/simpler/orchestrator.py b/python/simpler/orchestrator.py index f0e71aceb..e341a871d 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 native task-submission attempt and record possible 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(). @@ -210,28 +321,22 @@ def submit_next_level(self, callable_handle: Any, args: TaskArgs, config: CallCo self._worker._stage_host_buffers_for_chip_submit(c_args) final_worker_ids = _remote_data_eligible_worker_ids(remote_sidecar, eligible_worker_ids) worker = self._worker - # Do the (fallible) kind4 provenance analysis BEFORE capturing remote slot - # refs, so an exception here can never leave captured refs neither - # released nor adopted (which would defer a remote free forever). Capture - # is the last step before the rollback try. + # Provenance validation precedes run ownership publication. Once a remote + # ref is published, only the run fence may release it because the native + # submit can commit before an exception reaches Python. child_ptrs = worker._child_ptrs_in_args(c_args) if worker is not None else [] prov_guard: Any = contextlib.nullcontext() if child_ptrs and worker is not None: prov_guard = worker._child_prov_lock - captured_refs = worker._capture_remote_sidecar_refs(remote_sidecar) if worker is not None else [] - try: - with prov_guard: - if child_ptrs and worker is not None: - worker._child_prov_check_dispatch(child_ptrs, cpp_worker_id, api="submit_next_level") - self._o.submit_next_level( - digest, kind, target_namespace, c_args, cfg, cpp_worker_id, final_worker_ids, remote_sidecar - ) - except BaseException: - if self._worker is not None: - self._worker._release_remote_slot_refs(captured_refs) - raise - if self._worker is not None: - self._worker._adopt_remote_slot_refs(captured_refs) + with prov_guard: + if child_ptrs and worker is not None: + worker._child_prov_check_dispatch(child_ptrs, cpp_worker_id, api="submit_next_level") + if worker is not None: + worker._adopt_remote_sidecar_refs((remote_sidecar,)) + _admit_task_submission(self._worker) + self._o.submit_next_level( + digest, kind, target_namespace, c_args, cfg, cpp_worker_id, final_worker_ids, remote_sidecar + ) def submit_next_level_group( # noqa: PLR0912 -- linear per-member sidecar + eligibility + kind4-provenance passes, one branch each self, @@ -304,8 +409,7 @@ def submit_next_level_group( # noqa: PLR0912 -- linear per-member sidecar + eli ) # Per-member kind4 dispatch guard: each member's child_memory pointers # must be live on that member's exact submitted target. - # Run this (fallible) analysis BEFORE capturing remote slot refs, so an - # exception here can never strand captured refs outside the rollback try. + # Run this fallible analysis before publishing remote-ref ownership. worker = self._worker member_checks: list[tuple[list[tuple[int, int]], int]] = [] if worker is not None: @@ -317,24 +421,16 @@ def submit_next_level_group( # noqa: PLR0912 -- linear per-member sidecar + eli prov_guard: Any = ( worker._child_prov_lock if (worker is not None and member_checks) else contextlib.nullcontext() ) - captured_refs: list[Any] = [] - if self._worker is not None and remote_sidecars is not None: - for sidecar in remote_sidecars: - captured_refs.extend(self._worker._capture_remote_sidecar_refs(sidecar)) - try: - with prov_guard: - for child_ptrs, target_worker_id in member_checks: - assert worker is not None # member_checks is only populated when worker is present - worker._child_prov_check_dispatch(child_ptrs, target_worker_id, api="submit_next_level_group") - self._o.submit_next_level_group( - digest, kind, target_namespace, c_args_list, cfg, worker_ids, worker_id_sets, remote_sidecars - ) - except BaseException: - if self._worker is not None: - self._worker._release_remote_slot_refs(captured_refs) - raise - if self._worker is not None: - self._worker._adopt_remote_slot_refs(captured_refs) + with prov_guard: + for child_ptrs, target_worker_id in member_checks: + assert worker is not None # member_checks is only populated when worker is present + worker._child_prov_check_dispatch(child_ptrs, target_worker_id, api="submit_next_level_group") + if worker is not None and remote_sidecars is not None: + worker._adopt_remote_sidecar_refs(remote_sidecars) + _admit_task_submission(self._worker) + self._o.submit_next_level_group( + digest, kind, target_namespace, c_args_list, cfg, worker_ids, worker_id_sets, remote_sidecars + ) def submit_sub(self, callable_handle: Any, args: TaskArgs | None = None): """Submit a SUB task by registered callable handle. @@ -350,6 +446,7 @@ def submit_sub(self, callable_handle: Any, args: TaskArgs | None = None): expected_namespace="LOCAL_PYTHON", ) _reject_remote_sidecar_args(args, kind="orch.submit_sub") + _admit_task_submission(self._worker) self._o.submit_sub(digest, kind, target_namespace, args) def submit_sub_group(self, callable_handle: Any, args_list: list): @@ -362,6 +459,7 @@ def submit_sub_group(self, callable_handle: Any, args_list: list): ) for args in args_list: _reject_remote_sidecar_args(args, kind="orch.submit_sub_group") + _admit_task_submission(self._worker) self._o.submit_sub_group(digest, kind, target_namespace, args_list) # ------------------------------------------------------------------ @@ -403,12 +501,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 +520,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 +529,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 +586,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 +603,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 +647,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/task_interface.py b/python/simpler/task_interface.py index 332b436eb..68e02f276 100644 --- a/python/simpler/task_interface.py +++ b/python/simpler/task_interface.py @@ -220,9 +220,10 @@ class RemoteBufferHandle: "_ub_ldst_va", "_access_flags", "_released", - "_live_slot_refs", - "_live_import_refs", + "_slot_ref_tokens", + "_import_ref_tokens", "_owner_handle_ref", + "_owner_import_ref_token", ) def __init__( # noqa: PLR0913 @@ -242,6 +243,7 @@ def __init__( # noqa: PLR0913 access_flags: int = 3, released: bool = False, owner_handle_ref: RemoteBufferHandle | None = None, + owner_import_ref_token: object | None = None, _internal_token: object | None = None, ) -> None: address_space = RemoteAddressSpace(int(address_space)) @@ -261,9 +263,10 @@ def __init__( # noqa: PLR0913 self._ub_ldst_va = int(ub_ldst_va) self._access_flags = int(access_flags) self._released = bool(released) - self._live_slot_refs = 0 - self._live_import_refs = 0 + self._slot_ref_tokens: set[object] = set() + self._import_ref_tokens: set[object] = set() self._owner_handle_ref = owner_handle_ref + self._owner_import_ref_token = owner_import_ref_token if self._worker_id < 0: raise ValueError("RemoteBufferHandle.worker_id must be non-negative") @@ -336,6 +339,7 @@ def _from_imported_mapping( # noqa: PLR0913 access_flags: int = 0, released: bool = False, owner_handle_ref: RemoteBufferHandle | None = None, + owner_import_ref_token: object | None = None, ) -> RemoteBufferHandle: return cls( worker_id=worker_id, @@ -352,6 +356,7 @@ def _from_imported_mapping( # noqa: PLR0913 access_flags=access_flags, released=released, owner_handle_ref=owner_handle_ref, + owner_import_ref_token=owner_import_ref_token, _internal_token=_REMOTE_BUFFER_HANDLE_TOKEN, ) @@ -395,28 +400,48 @@ def is_imported(self) -> bool: """Whether this came from ``remote_import`` rather than ``remote_malloc``.""" return self._import_id != 0 + @property + def _live_slot_refs(self) -> int: + return len(self._slot_ref_tokens) + + @property + def _live_import_refs(self) -> int: + return len(self._import_ref_tokens) + def _mark_released(self) -> None: self._released = True - def _acquire_slot_ref(self) -> None: + def _acquire_slot_ref(self, token: object | None = None) -> object: if self._released: raise RuntimeError("RemoteBufferHandle has already been released") - self._live_slot_refs += 1 - - def _release_slot_ref(self) -> None: - if self._live_slot_refs <= 0: + if token is None: + token = object() + self._slot_ref_tokens.add(token) + return token + + def _release_slot_ref(self, token: object | None = None) -> None: + if token is not None: + self._slot_ref_tokens.discard(token) + return + if not self._slot_ref_tokens: raise RuntimeError("RemoteBufferHandle live slot refs underflow") - self._live_slot_refs -= 1 + self._slot_ref_tokens.pop() - def _acquire_import_ref(self) -> None: + def _acquire_import_ref(self, token: object | None = None) -> object: if self._released: raise RuntimeError("RemoteBufferHandle has already been released") - self._live_import_refs += 1 - - def _release_import_ref(self) -> None: - if self._live_import_refs <= 0: + if token is None: + token = object() + self._import_ref_tokens.add(token) + return token + + def _release_import_ref(self, token: object | None = None) -> None: + if token is not None: + self._import_ref_tokens.discard(token) + return + if not self._import_ref_tokens: raise RuntimeError("RemoteBufferHandle live import refs underflow") - self._live_import_refs -= 1 + self._import_ref_tokens.pop() def __repr__(self) -> str: return ( @@ -932,7 +957,17 @@ class CommDomainHandle: ``released → freed`` transition is the runtime's job at end-of-run. """ - __slots__ = ("name", "workers", "contexts", "allocation_id", "_release_fn", "_released", "_freed") + __slots__ = ( + "name", + "workers", + "contexts", + "allocation_id", + "_domain_size", + "_domain_ranks", + "_release_fn", + "_released", + "_freed", + ) def __init__( self, @@ -942,12 +977,18 @@ def __init__( contexts: dict[int, ChipDomainContext], allocation_id: int, _release_fn, + _domain_size: int | None = None, + _domain_ranks: dict[int, int] | None = None, ) -> None: self.name = name self.workers = tuple(workers) # Frozen dict-ish — we don't expose mutation self.contexts: dict[int, ChipDomainContext] = dict(contexts) self.allocation_id = int(allocation_id) + self._domain_size = len(self.workers) if _domain_size is None else int(_domain_size) + self._domain_ranks = ( + {worker: rank for rank, worker in enumerate(self.workers)} if _domain_ranks is None else dict(_domain_ranks) + ) self._release_fn = _release_fn self._released = False self._freed = False diff --git a/python/simpler/worker.py b/python/simpler/worker.py index d3fd1f1bd..88a9e7e5e 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -76,19 +76,23 @@ def my_l4_orch(orch, args, config): import time import uuid from dataclasses import dataclass, field +from multiprocessing import resource_tracker from multiprocessing.shared_memory import SharedMemory from typing import Any, cast 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, _l3_child_onboard_region_close, _l3_child_onboard_region_create, + _l3_host_mapped_region_close, _l3_host_mapped_region_import_onboard, _l3_host_mapped_region_import_sim, + _l3_host_mapped_region_take_cleanup_error, _mailbox_load_i32, _mailbox_store_i32, read_args_from_blob, @@ -125,7 +129,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 +180,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 @@ -241,6 +248,10 @@ def my_l4_orch(orch, args, config): # terminal state and notifying), a parked waiter re-observes that state within # this interval instead of blocking forever. _RUN_HANDLE_WAIT_RECHECK_S = 1.0 +# Native cancellation may finish fallible fan-in preparation only on a retry. +# Never turn that contract into an unbounded Python loop: one retry is enough +# to either settle the fence or declare the worker unusable. +_RUN_CANCELLATION_ATTEMPTS = 2 # Control sub-commands (written at _OFF_CALLABLE as uint64) _CTRL_MALLOC = 0 @@ -341,11 +352,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 + + +@dataclass +class _ThreadFanoutState: + item: Any + claim_lock: Any = field(default_factory=threading.Lock) + claimed: bool = False + target_completed: threading.Event = field(default_factory=threading.Event) + launch_confirmed: bool = False + attempted_once: bool = False + terminal_start_failure: bool = False + + +@dataclass +class _ThreadFanoutCandidate: + completed: threading.Event = field(default_factory=threading.Event) + decided: threading.Event = field(default_factory=threading.Event) + admitted: bool | None = None + thread: threading.Thread | None = None + + +@dataclass +class _ThreadFanoutDrainCursor: + phases: tuple[Any, ...] + after_phase: Any | None + next_phase: int = 0 + + @property + def exhausted(self) -> bool: + return self.next_phase == len(self.phases) + + def advance(self) -> None: + phase = self.next_phase + self.phases[phase]() + self.next_phase = phase + 1 + if self.after_phase is not None: + self.after_phase(phase) + + +class _ThreadFanout: + """Own every possible launch until no target can retain caller state.""" + + def __init__(self, items, target, name_prefix: str, after_start, after_phase) -> None: + self._states = [_ThreadFanoutState(item) for item in items] + self._target = target + self._name_prefix = name_prefix + self._after_start = after_start + self._after_phase = after_phase + self._candidates: list[_ThreadFanoutCandidate] = [] + self._first_error: BaseException | None = None + + def _record_error(self, exc: BaseException) -> None: + if self._first_error is None: + self._first_error = exc + + def _set_event(self, event: threading.Event) -> None: + while not event.is_set(): + try: + event.set() + except BaseException as exc: # noqa: BLE001, PERF203 + self._record_error(exc) + + def _wait_event(self, event: threading.Event) -> None: + while not event.is_set(): + try: + event.wait() + except BaseException as exc: # noqa: BLE001, PERF203 + self._record_error(exc) + + def _publish_decision(self, candidate: _ThreadFanoutCandidate, value: bool) -> None: + # Admission is one-way. An interrupt can land after the candidate was + # admitted but before its rank is recorded as confirmed; cleanup must + # never retract that publication under a target leaving the gate. + if candidate.admitted is None: + candidate.admitted = value + self._set_event(candidate.decided) + + def _invoke(self, state: _ThreadFanoutState, candidate: _ThreadFanoutCandidate) -> None: + try: + self._wait_event(candidate.decided) + if candidate.admitted is not True: + return + with state.claim_lock: + if state.claimed: + return + state.claimed = True + try: + self._target(state.item) + finally: + self._set_event(state.target_completed) + finally: + self._set_event(candidate.completed) + + def _attempt_start(self, state: _ThreadFanoutState) -> bool | None: + candidate = _ThreadFanoutCandidate() + try: + thread = threading.Thread( + target=self._invoke, + args=(state, candidate), + name=f"{self._name_prefix}{state.item}", + ) + except Exception as exc: + self._record_error(exc) + state.terminal_start_failure = True + return None + except BaseException as exc: # noqa: BLE001 + self._record_error(exc) + return False + try: + candidate.thread = thread + self._candidates.append(candidate) + thread.start() + self._publish_decision(candidate, True) + state.launch_confirmed = True + return True + except Exception as exc: # noqa: BLE001 + self._record_error(exc) + self._publish_decision(candidate, False) + state.terminal_start_failure = True + return None + except BaseException as exc: # noqa: BLE001 + self._record_error(exc) + self._publish_decision(candidate, False) + return False + + def _finish_launches(self) -> None: + # Give every participant one opportunity before retrying an ambiguous + # BaseException boundary on any one participant. + for state in self._states: + if state.attempted_once: + continue + if state.launch_confirmed or state.terminal_start_failure: + state.attempted_once = True + continue + result = self._attempt_start(state) + if result is True and self._after_start is not None: + self._after_start(state.item) + state.attempted_once = True + + for state in self._states: + while not state.launch_confirmed and not state.terminal_start_failure: + result = self._attempt_start(state) + if result is True and self._after_start is not None: + self._after_start(state.item) + + def _cancel_undecided_candidates(self) -> None: + for candidate in self._candidates: + if not candidate.decided.is_set(): + self._publish_decision(candidate, False) + + def _wait_for_targets(self) -> None: + for state in self._states: + if state.launch_confirmed: + self._wait_event(state.target_completed) + + def _join_candidates(self) -> None: + for candidate in self._candidates: + thread = candidate.thread + assert thread is not None + if not thread._started.is_set(): # noqa: SLF001 -- a cancelled ambiguous start cannot access target state + continue + self._wait_event(candidate.completed) + while True: + try: + thread.join() + break + except BaseException as exc: # noqa: BLE001, PERF203 + self._record_error(exc) + + def _drain(self, cursor: _ThreadFanoutDrainCursor) -> None: + try: + while not cursor.exhausted: + cursor.advance() + except BaseException as exc: # noqa: BLE001 + self._record_error(exc) + self._drain(cursor) + + def run(self) -> None: + cursor = _ThreadFanoutDrainCursor( + phases=( + self._finish_launches, + self._cancel_undecided_candidates, + self._wait_for_targets, + self._join_candidates, + ), + after_phase=self._after_phase, + ) + self._drain(cursor) + if self._first_error is not None: + raise self._first_error + + +def _start_and_join_threads( + items, + target, + *, + name_prefix: str, + _after_start: Any | None = None, + _after_phase: Any | None = None, +) -> None: + """Start all targets and defer interruptions until every launch is drained. + + ``Thread.start`` may be interrupted after the OS thread exists but before + it returns. Candidates are owned before start and gated until the launch is + accepted, so an ambiguous candidate is harmless and can be retried. + """ + _ThreadFanout(items, target, name_prefix, _after_start, _after_phase).run() + + +@dataclass +class _IsolatedCallResult: + value: Any | None = None + error: BaseException | None = None + completed: bool = False + + +def _run_isolated_call( + result: _IsolatedCallResult, + operation, + *, + name_prefix: str, + after_success: Any | None = None, +) -> None: + """Run a control call outside the caller's async-exception boundary.""" + + def invoke(_item) -> None: + try: + result.value = operation() + if after_success is not None: + after_success() + result.completed = True + except BaseException as exc: # noqa: BLE001 + result.error = exc + + _start_and_join_threads((0,), invoke, name_prefix=name_prefix) + + +@dataclass +class _OwnedSharedMemorySlot: + shm: SharedMemory | None = None + owns_name: bool = False + close_buffer: Any | None = None + close_mmap: Any | None = None + + +@dataclass +class _SharedMemoryCreateResult: + shm: SharedMemory | None = None + error: BaseException | None = None + + +class _SharedMemoryOwner: + """Own fixed named shm slots before any create operation can publish one.""" + + def __init__(self, capacity: int) -> None: + self._slots = [_OwnedSharedMemorySlot() for _ in range(capacity)] + self._create_cursor = 0 + self._attempted = 0 + self._cleanup_step = 0 + self._cleanup_error: BaseException | None = None + + @staticmethod + def _create_in_helper(slot: _OwnedSharedMemorySlot, size: int, result: _SharedMemoryCreateResult) -> None: + while True: + # Several mailbox control payloads reserve 32 bytes including the + # trailing NUL for this token. Keep the explicit collision-safe + # name below that ABI limit while retaining 96 bits of entropy. + name = f"smp_{uuid.uuid4().hex[:24]}" + shm = SharedMemory.__new__(SharedMemory) + shm._name = f"/{name}" if SharedMemory._prepend_leading_slash else name # noqa: SLF001 + shm._fd = -1 # noqa: SLF001 + shm._mmap = None # noqa: SLF001 + shm._buf = None # noqa: SLF001 + slot.shm = shm + slot.owns_name = False + try: + SharedMemory.__init__(shm, name=name, create=True, size=size) + except FileExistsError as exc: + if _shm_slot_has_created_resource(slot): + slot.owns_name = True + result.error = exc + return + slot.shm = None + continue + except BaseException as exc: # noqa: BLE001 + slot.owns_name = _shm_slot_has_created_resource(slot) + result.error = exc + return + slot.owns_name = True + result.shm = shm + return + + def create(self, size: int) -> SharedMemory: + if self._create_cursor >= len(self._slots): + raise RuntimeError("shared-memory owner capacity exhausted") + slot = self._slots[self._create_cursor] + result = _SharedMemoryCreateResult() + self._attempted = self._create_cursor + 1 + + def create_in_helper(_item) -> None: + try: + self._create_in_helper(slot, size, result) + except BaseException as exc: # noqa: BLE001 + result.error = exc + + _start_and_join_threads((0,), create_in_helper, name_prefix="shm-create-") + if result.error is not None: + raise result.error + shm = result.shm + if shm is None: + raise RuntimeError("shared-memory helper completed without a result") + self._create_cursor += 1 + return shm + + +def _remember_cleanup_error(first_error: BaseException | None, exc: BaseException) -> BaseException: + return exc if first_error is None else first_error + + +def _shm_slot_has_created_resource(slot: _OwnedSharedMemorySlot) -> bool: + shm = slot.shm + if shm is None: + return False + if os.name == "posix" and getattr(shm, "_fd", -1) >= 0: + return True + return getattr(shm, "_mmap", None) is not None or getattr(shm, "_buf", None) is not None + + +def _release_owned_shm_buffer(slot: _OwnedSharedMemorySlot) -> None: + shm = slot.shm + if shm is None: + return + if slot.close_buffer is None: + slot.close_buffer = shm._buf # noqa: SLF001 + buffer = slot.close_buffer + if shm._buf is buffer: # noqa: SLF001 + shm._buf = None # noqa: SLF001 + if buffer is not None: + buffer.release() + slot.close_buffer = None + + +def _close_owned_shm_mmap(slot: _OwnedSharedMemorySlot) -> None: + shm = slot.shm + if shm is None: + return + if slot.close_mmap is None: + slot.close_mmap = shm._mmap # noqa: SLF001 + mapped = slot.close_mmap + if shm._mmap is mapped: # noqa: SLF001 + shm._mmap = None # noqa: SLF001 + if mapped is not None and not getattr(mapped, "closed", False): + mapped.close() + slot.close_mmap = None + + +def _close_owned_shm_fd(slot: _OwnedSharedMemorySlot) -> None: + if os.name != "posix" or slot.shm is None: + return + shm = slot.shm + fd = shm._fd # noqa: SLF001 + if fd < 0: + return + shm._fd = -1 # noqa: SLF001 + os.close(fd) + + +_SHM_CLEANUP_PHASES = 7 + + +class _SharedMemoryCleanupCursor: + def __init__(self, owner: _SharedMemoryOwner, after_error: Any | None, after_step: Any | None) -> None: + self._owner = owner + self._after_error = after_error + self._after_step = after_step + self._pending_error: BaseException | None = None + self._pending_advance = False + + @property + def exhausted(self) -> bool: + return self._owner._cleanup_step >= self._owner._attempted * _SHM_CLEANUP_PHASES + + def _finish_step(self) -> None: + self._owner._cleanup_step += 1 + if self._after_step is not None: + self._after_step() + + def _queue_error(self, exc: BaseException, *, advance: bool) -> None: + if self._pending_error is None: + self._pending_error = exc + self._pending_advance = advance + elif self._owner._cleanup_error is None: + self._owner._cleanup_error = self._pending_error + + def _capture_interruption(self, exc: BaseException) -> None: + seen: set[int] = set() + first = exc + while first.__context__ is not None and id(first) not in seen: + seen.add(id(first)) + first = first.__context__ + if self._pending_error is None: + self._pending_error = first + self._pending_advance = False + elif self._owner._cleanup_error is None: + self._owner._cleanup_error = self._pending_error + + def _record_pending_error(self) -> None: + pending = self._pending_error + assert pending is not None + self._owner._cleanup_error = _remember_cleanup_error(self._owner._cleanup_error, pending) + if self._after_error is not None: + self._after_error() + advance = self._pending_advance + self._pending_error = None + self._pending_advance = False + if advance: + self._finish_step() + + def _run_step(self) -> None: + slot_index, phase = divmod(self._owner._cleanup_step, _SHM_CLEANUP_PHASES) + slot = self._owner._slots[slot_index] + shm = slot.shm + try: + if phase == 0: + slot.owns_name = slot.owns_name or _shm_slot_has_created_resource(slot) + elif phase == 1: + if slot.owns_name and os.name == "posix" and shm is not None and shm._name is not None: # noqa: SLF001 + resource_tracker.register(shm._name, "shared_memory") # noqa: SLF001 + elif phase == 2: + if slot.owns_name and shm is not None: + shm.unlink() + elif phase == 3: + if slot.owns_name and os.name == "posix" and shm is not None and shm._name is not None: # noqa: SLF001 + resource_tracker.register(shm._name, "shared_memory") # noqa: SLF001 + resource_tracker.unregister(shm._name, "shared_memory") # noqa: SLF001 + elif phase == 4: + _release_owned_shm_buffer(slot) + elif phase == 5: + _close_owned_shm_mmap(slot) + else: + _close_owned_shm_fd(slot) + except FileNotFoundError: + self._finish_step() + except Exception as exc: # noqa: BLE001 + self._queue_error(exc, advance=True) + else: + self._finish_step() + + def drain(self, pending_interruption: BaseException | None = None) -> None: + if pending_interruption is not None: + self._capture_interruption(pending_interruption) + while not self.exhausted or self._pending_error is not None: + try: + if self._pending_error is not None: + self._record_pending_error() + else: + self._run_step() + except BaseException as exc: # noqa: BLE001 + self._capture_interruption(exc) + + +def _close_unlink_shms( + owner: _SharedMemoryOwner, + *, + _after_error: Any | None = None, + _after_step: Any | None = None, +) -> BaseException | None: + """Drain every owned shm and return the first non-benign cleanup error.""" + pending_interruption: BaseException | None = None + while True: + try: + cursor = _SharedMemoryCleanupCursor(owner, _after_error, _after_step) + if pending_interruption is None: + _start_and_join_threads((0,), lambda _item: cursor.drain(), name_prefix="shm-cleanup-") + else: + # A setup/fanout interruption is itself a cleanup error, but + # cannot be allowed to bypass the owner whose names are still + # live. The cursor records it and drains from the durable step. + cursor.drain(pending_interruption) + if not cursor.exhausted: + cursor.drain() + return owner._cleanup_error + except BaseException as exc: # noqa: BLE001, PERF203 + if pending_interruption is None: + pending_interruption = exc + + +def _run_with_owned_shared_memory( + capacity: int, + operation, + *, + name_prefix: str, + after_success: Any | None = None, +) -> Any: + """Run one complete shm lifecycle outside the caller's signal boundary. + + Python delivers ``KeyboardInterrupt`` to the main thread. Constructing the + owner inside the isolated helper means an interrupt can either prevent the + lifecycle from starting, or be deferred until every acquired name has been + unlinked and closed; it can never abort the caller immediately before its + cleanup call. + """ + result = _IsolatedCallResult() + + def run_lifecycle() -> Any: + owner = _SharedMemoryOwner(capacity) + value: Any | None = None + operation_error: BaseException | None = None + cleanup_error: BaseException | None = None + cleanup_done = False + try: + value = operation(owner) + except BaseException as exc: # noqa: BLE001 + operation_error = exc + while not cleanup_done: + try: + cleanup_error = _close_unlink_shms(owner) + cleanup_done = True + except BaseException as exc: # noqa: BLE001 + # This handles a synchronous failure at the cleanup call entry. + # Caller-directed asynchronous exceptions cannot reach this + # lifecycle thread. + operation_error = _remember_cleanup_error(operation_error, exc) + if operation_error is not None: + raise operation_error + if cleanup_error is not None: + raise cleanup_error + return value + + _run_isolated_call(result, run_lifecycle, name_prefix=name_prefix, after_success=after_success) + if result.error is not None: + raise result.error + if not result.completed: + raise RuntimeError("shared-memory lifecycle did not publish an outcome") + return result.value + + +def _validate_domain_allocation( + worker: Worker, + name: str, + workers: tuple[int, ...], + window_size: int, + buffers: list[CommBufferSpec], +) -> _RunResources: + """Validate a domain request before any communicator or shm side effect.""" + if worker.level < 3: + raise RuntimeError("allocate_domain requires level >= 3") + if worker._worker is None: + raise RuntimeError("allocate_domain requires a hierarchical Worker (_start_hierarchical ran)") + resources = worker._building_run_resources + if resources is None: + raise RuntimeError("allocate_domain is only valid while a run's graph is being built") + if not workers: + raise ValueError("allocate_domain: workers must be non-empty") + if len(set(workers)) != len(workers): + raise ValueError(f"allocate_domain: workers contains duplicates: {workers}") + device_ids = worker._config.get("device_ids", []) + for worker_id in workers: + if worker_id < 0 or worker_id >= len(device_ids): + raise ValueError(f"allocate_domain: worker_id {worker_id} outside [0, {len(device_ids)})") + if window_size <= 0: + raise ValueError("allocate_domain: window_size must be positive") + buffer_names = [buffer.name for buffer in buffers] + if len(set(buffer_names)) != len(buffer_names): + raise ValueError(f"allocate_domain: duplicate buffer names: {buffer_names}") + total_buffer_nbytes = sum(int(buffer.nbytes) for buffer in buffers) + if total_buffer_nbytes > window_size: + raise ValueError( + f"allocate_domain: buffers sum to {total_buffer_nbytes} bytes, exceeds window_size={window_size}" + ) + if name in worker._live_domains: + raise ValueError(f"allocate_domain: domain {name!r} already live") + return resources + + 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 +1852,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), + _buffer_field_addr(reply_buf, _OFF_DOMAIN_REPLY_COMMITTED), + ) + + # 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) + + _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 +2201,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 +2218,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 +2419,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() @@ -1930,6 +2580,14 @@ class _Lifecycle(enum.Enum): CLOSED = enum.auto() +@dataclass(frozen=True) +class _CloseOutcome: + """One immutable result published by a close attempt.""" + + error: BaseException | None + incomplete: bool + + class _CloseAttempt: """Private completion record for one close() teardown attempt. @@ -1942,16 +2600,57 @@ class _CloseAttempt: latches True), an un-reclaimed resource LEAKS — a later close() never re-drives a half-torn tree. Teardown stays UN-attempted, with the tree intact and a later close() free to drive drain+teardown, only where the drain itself did - not complete: in-flight leases outlived the cleanup budget, or an async - interruption left an accepted run fence undrained. + not complete: in-flight leases or accepted run fences outlived the shared + cleanup budget, or an async interruption left a run fence undrained. """ - __slots__ = ("done", "error", "incomplete") + __slots__ = ("_outcome",) def __init__(self) -> None: - self.done: bool = False - self.error: BaseException | None = None - self.incomplete: bool = False + self._outcome: _CloseOutcome | None = None + + @property + def done(self) -> bool: + return self._outcome is not None + + @property + def error(self) -> BaseException | None: + outcome = self._outcome + return None if outcome is None else outcome.error + + @property + def incomplete(self) -> bool: + outcome = self._outcome + return False if outcome is None else outcome.incomplete + + def publish(self, error: BaseException | None, incomplete: bool) -> _CloseOutcome: + """Install the attempt's only outcome despite pre-commit interruptions.""" + effective_error = error + effective_incomplete = incomplete + try: + while True: + try: + committed = self._outcome + if committed is not None: + return committed + candidate = _CloseOutcome(effective_error, effective_incomplete) + self._outcome = candidate + return candidate + except BaseException as exc: # noqa: BLE001, PERF203 + # An interruption after STORE_ATTR observes the immutable + # committed result and remains local to this publisher. An + # interruption before STORE_ATTR becomes the shared result. + if self._outcome is not None: + raise + if effective_error is None: + effective_error = exc + effective_incomplete = True + except BaseException as exc: # noqa: BLE001 + if self._outcome is not None: + raise + if effective_error is None: + effective_error = exc + return self.publish(effective_error, True) class _StartupCancelled(BaseException): @@ -1963,11 +2662,19 @@ class _StartupCancelled(BaseException): cancellable (``close()`` fails fast while INITIALIZING).""" +@dataclass(eq=False) +class _RemoteSlotRefClaim: + """One replayable remote-buffer reference owned by a run cleanup journal.""" + + handle: RemoteBufferHandle + token: object = field(default_factory=object) + + @dataclass class _RunResources: """Python resources whose lifetime ends at one native run fence.""" - remote_slot_refs: list[RemoteBufferHandle] = field(default_factory=list) + remote_slot_refs: list[_RemoteSlotRefClaim | RemoteBufferHandle] = field(default_factory=list) live_domains: dict[str, CommDomainHandle] = field(default_factory=dict) pending_release_domains: list[CommDomainHandle] = field(default_factory=list) l3_l2_regions: list[Any] = field(default_factory=list) @@ -1982,6 +2689,94 @@ 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 + + +@dataclass +class _PendingRemoteImportReleaseState: + """Durable local phases after a deferred import-release RPC starts.""" + + owner_ref: RemoteBufferHandle | None + owner_ref_token: object | None = None + rpc_complete: bool = False + owner_release_complete: bool = False + error: BaseException | None = None + + +@dataclass +class _RunFinalizationCursor: + """Own the boundary between each one-shot run cleanup operation.""" + + steps: tuple[tuple[str, Any], ...] + next_step: int = 0 + cleanup_error: BaseException | None = None + boundary_error: BaseException | None = None + incomplete: bool = False + + @property + def exhausted(self) -> bool: + return self.next_step == len(self.steps) + + def _remember_cleanup_error(self, exc: BaseException) -> None: + if self.cleanup_error is None: + self.cleanup_error = exc + + def remember_boundary_error(self, exc: BaseException) -> None: + if self.boundary_error is None: + self.boundary_error = exc + + def _advance(self, after_step) -> None: + index = self.next_step + name, operation = self.steps[index] + try: + operation() + except Exception as exc: # noqa: BLE001 + self._remember_cleanup_error(exc) + self.next_step = index + 1 + if after_step is not None: + after_step(name) + + def drain(self, after_step) -> None: + while not self.exhausted: + starting_step = self.next_step + try: + self._advance(after_step) + except BaseException as exc: # noqa: BLE001 + self.remember_boundary_error(exc) + if self.next_step == starting_step: + # An uncommitted operation has ambiguous ownership and is + # not replayable: its native side effect may have landed. + self.incomplete = True + return + + +class _AbandonedRunKeepaliveCursor: + """Drain retained references only after the native tree is gone.""" + + def __init__(self, handles: list[RunHandle]) -> None: + self._handles = handles + self.first_error: BaseException | None = None + + def drain(self, pending_error: BaseException | None = None) -> None: + try: + if pending_error is not None and self.first_error is None: + self.first_error = pending_error + while self._handles: + handle = self._handles[-1] + handle._keepalive = None + self._handles.pop() + except BaseException as exc: # noqa: BLE001 + self.drain(exc) class RunHandle: @@ -2009,6 +2804,14 @@ def __init__( self._launch_accepted = False self._terminal = False self._error: BaseException | None = None + self._finalization_error: BaseException | None = None + self._finalization_abandoned = False + # 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 +2826,10 @@ def _completed(cls, worker: Worker) -> RunHandle: handle._launch_accepted = True handle._terminal = True handle._error = None + handle._finalization_error = None + handle._finalization_abandoned = False + # Nothing was ever admitted, so there is no cleanup owing. + handle._cleanup_published = True return handle @staticmethod @@ -2036,11 +2843,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: @@ -2054,90 +2867,204 @@ def done(self) -> bool: assert run_id is not None return self._worker._run_handle_done(run_id) - def wait(self, timeout: float | None = None) -> None: + def _clear_wait_owner(self) -> None: + """Make an interrupted pre-finalization fence wait re-electable.""" + try: + while self._wait_in_progress: + try: + self._wait_in_progress = False + except BaseException: # noqa: BLE001, PERF203 + pass + try: + with self._cv: + self._cv.notify_all() + except BaseException: # noqa: BLE001 + # Parked waiters use bounded re-checks and observe the plain flag. + pass + except BaseException: # noqa: BLE001 + self._clear_wait_owner() + + def _clear_acceptance_owner(self) -> None: + """Make an interrupted acceptance wait re-electable.""" + try: + while self._accept_wait_in_progress: + try: + self._accept_wait_in_progress = False + except BaseException: # noqa: BLE001, PERF203 + pass + try: + with self._cv: + self._cv.notify_all() + except BaseException: # noqa: BLE001 + # Parked waiters use bounded re-checks and observe the plain flag. + pass + except BaseException: # noqa: BLE001 + self._clear_acceptance_owner() + + def _publish_acceptance(self) -> None: + """Publish launch acceptance before releasing its elected owner.""" + try: + while not self._launch_accepted or self._accept_wait_in_progress: + try: + self._launch_accepted = True + self._accept_wait_in_progress = False + except BaseException: # noqa: BLE001, PERF203 + pass + try: + with self._cv: + self._cv.notify_all() + except BaseException: # noqa: BLE001 + # Launch acceptance is visible before notification. + pass + except BaseException: # noqa: BLE001 + self._publish_acceptance() + + def _publish_terminal(self, error: BaseException | None) -> BaseException | None: + """Publish one terminal result despite interruptions between assigns.""" + effective_error = error + try: + while not self._terminal or self._wait_in_progress or self._accept_wait_in_progress: + try: + if not self._terminal: + self._error = effective_error + self._run_id = None + if not self._finalization_abandoned: + self._keepalive = None + self._launch_accepted = True + self._terminal = True + self._wait_in_progress = False + self._accept_wait_in_progress = False + except BaseException as exc: # noqa: BLE001, PERF203 + if not self._terminal and effective_error is None: + effective_error = exc + try: + with self._cv: + self._cv.notify_all() + except BaseException: # noqa: BLE001 + # Terminal is visible before notification; bounded re-checks + # cover a skipped wakeup. + pass + return self._error + except BaseException as exc: # noqa: BLE001 + if not self._terminal and effective_error is None: + effective_error = exc + return self._publish_terminal(effective_error) + + def _cache_finalization_error(self, error: BaseException | None) -> None: + """Store the finalizer outcome before control returns to its caller.""" + try: + while True: + try: + self._finalization_error = error + return + except BaseException: # noqa: BLE001, PERF203 + pass + except BaseException: # noqa: BLE001 + self._cache_finalization_error(error) + + def _recover_and_publish_terminal(self, error: BaseException) -> BaseException | None: + """Finish interrupted finalization without reopening its cleanup.""" + try: + while True: + try: + if self._terminal: + return self._error + recover = getattr(self._worker, "_recover_interrupted_run_finalization", None) + if recover is not None: + recover(self, error) + return self._publish_terminal(error) + except BaseException: # noqa: BLE001, PERF203 + # Recovery and terminal publication form one monotonic + # transition. Re-entering recovery only repeats ownership + # publication; it never replays a cleanup step. + pass + except BaseException: # noqa: BLE001 + return self._recover_and_publish_terminal(error) + + def wait(self, timeout: float | None = None) -> None: # noqa: PLR0912 -- one owner spans fence through publication """Wait for completion, raising ``TimeoutError`` or the run's error.""" - # Exactly one waiter is elected to cross the native fence; the rest park - # on the CV until it publishes. That election is unrecoverable if it is - # abandoned silently — no other path clears `_wait_in_progress`, and - # Worker.close() drains the handle with an untimed wait() — so every - # publication below (both abandonment paths and the terminal one) is a - # resilient publish in the shape close() uses for its _CloseAttempt: - # plain attribute assigns land BEFORE the CV acquire, whose only - # remaining work is notify_all(). An async BaseException in that - # interruptible acquire therefore cannot strand the handle, and a parked - # waiter's bounded re-check recovers the skipped notify. The only - # irreducible window is an async exception landing between the `_error` - # and `_terminal` assigns. + # One owner spans the native fence through finalization. An interruption + # before finalization clears that election; one after finalization starts + # never re-enters cleanup and conservatively abandons any unpublished + # ownership. Terminal fields are monotonic, and parked waiters use + # bounded re-checks when notification itself is interrupted. deadline = self._deadline(timeout) - with self._cv: - while not self._terminal and self._wait_in_progress: - remaining = None if deadline is None else deadline - time.monotonic() - if remaining is not None and remaining <= 0: - raise TimeoutError("RunHandle.wait() timed out") - recheck = ( - _RUN_HANDLE_WAIT_RECHECK_S if remaining is None else min(remaining, _RUN_HANDLE_WAIT_RECHECK_S) - ) - self._cv.wait(timeout=recheck) - if self._terminal: - error = self._error - if error is not None: - raise error - return - self._wait_in_progress = True - run_id = self._run_id - - assert run_id is not None + elected = False + finalization_started = False + final_error: BaseException | None = None native_error: BaseException | None = None try: - remaining = None if deadline is None else max(0.0, deadline - time.monotonic()) - completed = self._worker._wait_run_handle(run_id, remaining) - except Exception as exc: # native run failures are terminal - completed = True - native_error = exc - except BaseException: - self._wait_in_progress = False with self._cv: - self._cv.notify_all() - raise + while not self._terminal and self._wait_in_progress: + remaining = None if deadline is None else deadline - time.monotonic() + if remaining is not None and remaining <= 0: + raise TimeoutError("RunHandle.wait() timed out") + recheck = ( + _RUN_HANDLE_WAIT_RECHECK_S if remaining is None else min(remaining, _RUN_HANDLE_WAIT_RECHECK_S) + ) + self._cv.wait(timeout=recheck) + if self._terminal: + error = self._error + if error is not None: + raise error + return + run_id = self._run_id + elected = True + self._wait_in_progress = True - if not completed: - self._wait_in_progress = False - with self._cv: - self._cv.notify_all() - raise TimeoutError("RunHandle.wait() timed out") + boundary_hook = getattr(self, "_wait_boundary_hook", None) + if boundary_hook is not None: + boundary_hook("after_election") - # An acceptance waiter that already captured this run id must leave the - # native wait before finalize releases that id. It is terminal now, so - # this is only a short ownership hand-off and does not serialize device - # execution with acceptance waiting. - try: + assert run_id is not None + remaining = None if deadline is None else max(0.0, deadline - time.monotonic()) + try: + completed = self._worker._wait_run_handle(run_id, remaining) + except Exception as exc: # native run failures are terminal + completed = True + native_error = exc + + if not completed: + raise TimeoutError("RunHandle.wait() timed out") + + # An acceptance waiter that already captured this run id must leave + # the native wait before finalize releases that id. It is terminal + # now, so this ownership hand-off does not serialize device + # execution with acceptance waiting. with self._cv: while self._accept_wait_in_progress: - self._cv.wait(timeout=_RUN_HANDLE_WAIT_RECHECK_S) - except BaseException: - self._wait_in_progress = False - with self._cv: - self._cv.notify_all() - raise - - # Cleanup runs exactly once, on this waiter, and its outcome IS the - # handle's result: an interruption mid-finalize is cached like any other - # error rather than lost, so the publication below is unconditional. - try: - error = self._worker._finalize_run_handle(self, run_id, native_error) + remaining = None if deadline is None else deadline - time.monotonic() + if remaining is not None and remaining <= 0: + raise TimeoutError("RunHandle.wait() timed out") + recheck = ( + _RUN_HANDLE_WAIT_RECHECK_S if remaining is None else min(remaining, _RUN_HANDLE_WAIT_RECHECK_S) + ) + self._cv.wait(timeout=recheck) + + # Finalization owns the accepted-set retirement from this point. + # An escape after this flag takes conservative abandonment rather + # than re-entering cleanup or native run release. + finalization_started = True + final_error = self._worker._finalize_run_handle(self, run_id, native_error) + if boundary_hook is not None: + boundary_hook("after_finalize") + published_error = self._publish_terminal(final_error) except BaseException as exc: # noqa: BLE001 - error = exc - self._error = error - self._run_id = None - self._keepalive = None - self._launch_accepted = True - self._terminal = True - self._wait_in_progress = False - self._accept_wait_in_progress = False - with self._cv: - self._cv.notify_all() - if error is not None: - raise error + if not elected: + raise + if not finalization_started: + self._clear_wait_owner() + raise + + preferred_error = native_error + if preferred_error is None: + preferred_error = self._finalization_error + if preferred_error is None: + preferred_error = final_error if final_error is not None else exc + published_error = self._recover_and_publish_terminal(preferred_error) + + if published_error is not None: + raise published_error def result(self, timeout: float | None = None) -> None: """Alias for :meth:`wait`; successful runs have no return value.""" @@ -2151,30 +3078,57 @@ 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: - while not self._terminal and self._accept_wait_in_progress: - self._cv.wait(timeout=_RUN_HANDLE_WAIT_RECHECK_S) - if self._terminal or self._launch_accepted: - return - self._accept_wait_in_progress = True - run_id = self._run_id - - assert run_id is not None + elected = False + acceptance_completed = False try: + with self._cv: + while not self._terminal and self._accept_wait_in_progress: + self._cv.wait(timeout=_RUN_HANDLE_WAIT_RECHECK_S) + if self._terminal or self._launch_accepted: + return + run_id = self._run_id + elected = True + self._accept_wait_in_progress = True + + boundary_hook = getattr(self, "_wait_boundary_hook", None) + if boundary_hook is not None: + boundary_hook("after_acceptance_election") + + assert run_id is not None self._worker._wait_run_handle_accepted(run_id) + acceptance_completed = True + if boundary_hook is not None: + boundary_hook("after_acceptance_wait") + self._publish_acceptance() except BaseException: - self._accept_wait_in_progress = False - with self._cv: - self._cv.notify_all() + if not elected: + raise + if acceptance_completed: + self._publish_acceptance() + else: + self._clear_acceptance_owner() raise - self._launch_accepted = True - self._accept_wait_in_progress = False - with self._cv: - self._cv.notify_all() - def _forked_child_main(buf: memoryview, label: str, setup, serve, make_group_leader: bool = False) -> None: """Run a forked child to completion, always terminating via ``os._exit``. @@ -2336,13 +3290,27 @@ 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() + # A nonterminal cancellation or ambiguous finalization boundary has no + # fence-owned cleanup Python can safely replay. Keep that run's callback + # and resource references alive outside close()'s drain set; whole-tree + # teardown is the only remaining safe reclamation boundary. + self._abandoned_run_handles: list[RunHandle] = [] + # 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. @@ -2384,9 +3352,14 @@ def __init__( self._next_level_worker_id_count: int = 0 # Fallback ownership for private helpers used outside Worker.submit. # Normal orchestration-owned refs live in RunHandle._resources. - self._active_remote_slot_refs: list[RemoteBufferHandle] = [] + self._active_remote_slot_refs: list[_RemoteSlotRefClaim | RemoteBufferHandle] = [] self._pending_remote_buffer_frees: list[RemoteBufferHandle] = [] + # One publication/drain boundary for every remote release debt. In + # particular, an owner free must not append after a run fence has taken + # the queue's only snapshot, and two callers must not both send it. + self._remote_import_release_mu = threading.RLock() self._pending_remote_import_releases: list[RemoteBufferHandle] = [] + self._pending_remote_import_release_states: dict[RemoteBufferHandle, _PendingRemoteImportReleaseState] = {} # Dynamic CommDomain allocations. Keyed by user-facing name (unique # among live handles). ``orch.allocate_domain`` adds entries here; @@ -2838,7 +3811,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) @@ -2869,18 +3842,48 @@ def remote_free(self, handle: RemoteBufferHandle) -> None: if handle.is_imported: raise ValueError("remote_free is invalid for imported handles; use remote_release_import") if handle.released: - return + with self._remote_import_release_mu: + if handle not in self._pending_remote_buffer_frees: + return # Public admission: READY-only + drained. (The private _send_* transport # helper accepts CLOSED so teardown can flush pending frees, so remote_free # 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: - 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) - handle._mark_released() + with self._remote_import_release_mu: + if handle.released: + if handle not in self._pending_remote_buffer_frees: + return + if handle._live_slot_refs > 0 or handle._live_import_refs > 0: + return + 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. + self._publish_pending_remote_buffer_free(handle) + return + with self._control_admission("remote_free"): + # Admission can block behind a graph callback that acquires a + # slot/import ref, and another free can have completed while we + # waited. Re-read both facts at the linearization boundary. + with self._remote_import_release_mu: + if handle.released: + if handle not in self._pending_remote_buffer_frees: + return + if handle._live_slot_refs > 0 or handle._live_import_refs > 0: + return + self._flush_pending_remote_frees() + return + if handle._live_slot_refs > 0 or handle._live_import_refs > 0: + self._publish_pending_remote_buffer_free(handle) + return + # Publish logical release before the RPC. A Python async + # exception can otherwise land after the remote allocation + # is gone but before the handle is marked. FREE_REMOTE_BUFFER + # is idempotent by (buffer_id, generation), so a retained + # debt may safely be retried if queue retirement is + # interrupted after the reply. + self._publish_pending_remote_buffer_free(handle) + self._flush_pending_remote_frees() def remote_copy_to(self, handle: RemoteBufferHandle, host_ptr: Any, nbytes: int, *, offset: int = 0) -> None: """Copy ``nbytes`` from host memory into an owner remote buffer. @@ -2888,7 +3891,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 +3918,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 +3954,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,10 +4026,10 @@ 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( + def _remote_import_locked( # noqa: PLR0912 -- import and rollback phases share one ownership journal self, exported: RemoteBufferExport, *, worker: int, access: str | int | None = None ) -> RemoteBufferHandle: importer_worker_id = int(worker) @@ -3035,28 +4038,78 @@ def _remote_import_locked( if flags & ~exported._access_flags: raise ValueError("Worker.remote_import requested access is not a subset of export access") assert self._worker is not None + native_worker = self._worker owner_handle = exported._owner_handle - if owner_handle is not None: - owner_handle._acquire_import_ref() - fields: Any | None = None + owner_ref_token = object() if owner_handle is not None else None + owner_acquire = _IsolatedCallResult() + call_result = _IsolatedCallResult() + + def retire_owner_reference(message: str) -> None: + if owner_handle is None or owner_ref_token is None: + return + owner_release = _IsolatedCallResult() + + def release_owner_reference() -> None: + with self._remote_import_release_mu: + owner_handle._release_import_ref(owner_ref_token) + + try: + _run_isolated_call( + owner_release, + release_owner_reference, + name_prefix="simpler-remote-import-owner-release-", + ) + except BaseException as exc: # noqa: BLE001 + if owner_release.error is None: + owner_release.error = exc + if not owner_release.completed: + cleanup_error = owner_release.error or RuntimeError( + "Worker.remote_import: owner reference rollback did not settle" + ) + self._record_unreclaimable(message, cleanup_error) + try: - fields = self._worker.remote_import( - importer_worker_id, - exported._owner_worker_id, - exported._buffer_id, - exported._generation, - int(exported._address_space), - exported._offset, - exported._nbytes, - exported._export_id, - exported._remote_addr, - exported._rkey_or_token, - exported._ub_ldst_va, - exported._access_flags, - exported._transport_profile, - exported._transport_descriptor, - flags, + if owner_handle is not None: + + def acquire_owner_reference() -> object: + with self._remote_import_release_mu: + return owner_handle._acquire_import_ref(owner_ref_token) + + _run_isolated_call( + owner_acquire, + acquire_owner_reference, + name_prefix="simpler-remote-import-owner-acquire-", + ) + if owner_acquire.error is not None: + raise owner_acquire.error + if not owner_acquire.completed: + raise RuntimeError("Worker.remote_import: owner reference acquisition did not settle") + _run_isolated_call( + call_result, + lambda: native_worker.remote_import( + importer_worker_id, + exported._owner_worker_id, + exported._buffer_id, + exported._generation, + int(exported._address_space), + exported._offset, + exported._nbytes, + exported._export_id, + exported._remote_addr, + exported._rkey_or_token, + exported._ub_ldst_va, + exported._access_flags, + exported._transport_profile, + exported._transport_descriptor, + flags, + ), + name_prefix="simpler-remote-import-", ) + if call_result.error is not None: + raise call_result.error + fields = call_result.value + if fields is None: + raise RuntimeError("Worker.remote_import: remote transport returned no import descriptor") return RemoteBufferHandle._from_imported_mapping( worker_id=int(fields[0]), owner_worker_id=int(fields[1]), @@ -3071,16 +4124,89 @@ def _remote_import_locked( ub_ldst_va=int(fields[10]), access_flags=int(fields[11]), owner_handle_ref=owner_handle, + owner_import_ref_token=owner_ref_token, ) - except BaseException: + except BaseException as original_error: + fields = call_result.value if fields is not None: + release_result = _IsolatedCallResult() try: - self._send_remote_release_import_fields(fields) - except Exception: # noqa: BLE001 - pass - if owner_handle is not None: - owner_handle._release_import_ref() - raise + _run_isolated_call( + release_result, + lambda: self._send_remote_release_import_fields(fields), + name_prefix="simpler-remote-import-rollback-", + ) + except BaseException as exc: # noqa: BLE001 + if release_result.error is None: + release_result.error = exc + if not release_result.completed: + cleanup_error = release_result.error or RuntimeError( + "Worker.remote_import: rollback did not settle" + ) + self._record_unreclaimable( + "Worker.remote_import: a committed remote import could not be rolled back; " + "the mapping is retained until whole-tree teardown", + cleanup_error, + ) + else: + retire_owner_reference( + "Worker.remote_import: the remote mapping was rolled back but its owner reference " + "could not be retired" + ) + elif call_result.completed: + self._record_unreclaimable( + "Worker.remote_import: the transport completed without publishing an import descriptor; " + "remote ownership is retained until whole-tree teardown", + original_error, + ) + elif call_result.error is None: + retire_owner_reference( + "Worker.remote_import: the transport failed before publishing an import descriptor, but " + "its owner reference could not be retired" + ) + else: + self._record_unreclaimable( + "Worker.remote_import: import ownership is ambiguous after a failed transport call; " + "the owner reference is retained until whole-tree teardown", + call_result.error, + ) + raise original_error + + def _publish_pending_remote_import_release( + self, + handle: RemoteBufferHandle, + state: _PendingRemoteImportReleaseState, + pending_error: BaseException | None = None, + ) -> None: + """Publish release ownership before any import-release RPC can start.""" + with self._remote_import_release_mu: + try: + existing = self._pending_remote_import_release_states.get(handle) + if existing is None: + self._pending_remote_import_release_states[handle] = state + if handle not in self._pending_remote_import_releases: + self._pending_remote_import_releases.append(handle) + handle._mark_released() + except BaseException as exc: # noqa: BLE001 + self._publish_pending_remote_import_release(handle, state, pending_error or exc) + return + if pending_error is not None: + raise pending_error + + def _publish_pending_remote_buffer_free( + self, handle: RemoteBufferHandle, pending_error: BaseException | None = None + ) -> None: + """Atomically publish an owner-free debt and logical release.""" + with self._remote_import_release_mu: + try: + if handle not in self._pending_remote_buffer_frees: + self._pending_remote_buffer_frees.append(handle) + handle._mark_released() + except BaseException as exc: # noqa: BLE001 + self._publish_pending_remote_buffer_free(handle, pending_error or exc) + return + if pending_error is not None: + raise pending_error def remote_release_import(self, handle: RemoteBufferHandle) -> None: """Release an imported remote handle. @@ -3098,23 +4224,47 @@ def remote_release_import(self, handle: RemoteBufferHandle) -> None: # Public admission: READY-only + drained (the private _send_* transport # accepts CLOSED for teardown, so fence admission here). with self._operation_lease("remote_release_import"): - if handle._live_slot_refs > 0: - 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() - - def _capture_remote_sidecar_refs(self, remote_sidecar: Any) -> list[RemoteBufferHandle]: - captured: list[RemoteBufferHandle] = [] - if remote_sidecar is None: - return captured - try: + with self._remote_import_release_mu: + if handle.released: + return + if handle._live_slot_refs > 0: + # Deferred: see remote_free — nothing is sent from here. + self._publish_pending_remote_import_release( + handle, + _PendingRemoteImportReleaseState( + owner_ref=handle._owner_handle_ref, + owner_ref_token=handle._owner_import_ref_token, + ), + ) + return + with self._control_admission("remote_release_import"): + with self._remote_import_release_mu: + if handle.released: + return + if handle._live_slot_refs > 0: + self._publish_pending_remote_import_release( + handle, + _PendingRemoteImportReleaseState( + owner_ref=handle._owner_handle_ref, + owner_ref_token=handle._owner_import_ref_token, + ), + ) + return + self._publish_pending_remote_import_release( + handle, + _PendingRemoteImportReleaseState( + owner_ref=handle._owner_handle_ref, + owner_ref_token=handle._owner_import_ref_token, + ), + ) + self._flush_pending_remote_frees() + + @staticmethod + def _remote_sidecar_handles(remote_sidecars: Any) -> list[RemoteBufferHandle]: + handles: list[RemoteBufferHandle] = [] + for remote_sidecar in remote_sidecars: + if remote_sidecar is None: + continue for tensor_sidecar in getattr(remote_sidecar, "tensors", ()): if tensor_sidecar is None or not getattr(tensor_sidecar, "present", False): continue @@ -3123,73 +4273,209 @@ def _capture_remote_sidecar_refs(self, remote_sidecar: Any) -> list[RemoteBuffer continue if not isinstance(handle, RemoteBufferHandle): raise TypeError("remote sidecar handle must be a RemoteBufferHandle") - handle._acquire_slot_ref() - captured.append(handle) + handles.append(handle) + return handles + + def _capture_remote_sidecar_refs(self, remote_sidecar: Any) -> list[_RemoteSlotRefClaim]: + captured: list[_RemoteSlotRefClaim] = [] + try: + for handle in self._remote_sidecar_handles((remote_sidecar,)): + claim = _RemoteSlotRefClaim(handle) + captured.append(claim) + with self._remote_import_release_mu: + handle._acquire_slot_ref(claim.token) except BaseException: self._release_remote_slot_refs(captured) raise return captured - def _adopt_remote_slot_refs(self, handles: list[RemoteBufferHandle]) -> None: + def _adopt_remote_slot_refs(self, handles: list[Any]) -> None: resources = self._building_run_resources if resources is None: self._active_remote_slot_refs.extend(handles) else: resources.remote_slot_refs.extend(handles) - - def _release_remote_slot_refs(self, handles: list[RemoteBufferHandle]) -> None: + # 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 _adopt_remote_sidecar_refs(self, remote_sidecars: Any) -> None: + handles = self._remote_sidecar_handles(remote_sidecars) + if not handles: + return + resources = self._building_run_resources + refs = self._active_remote_slot_refs if resources is None else resources.remote_slot_refs + if resources is not None: + resources.requires_ordered_cleanup = True for handle in handles: - handle._release_slot_ref() + claim = _RemoteSlotRefClaim(handle) + refs.append(claim) + with self._remote_import_release_mu: + handle._acquire_slot_ref(claim.token) + + def _release_remote_slot_refs(self, refs: list[Any]) -> None: + while refs: + # Preserve acquisition order. If a release is interrupted, later + # claims must remain owned rather than being retired ahead of the + # ambiguous claim. Token releases are idempotent across the + # release/delete boundary. + ref = refs[0] + with self._remote_import_release_mu: + if isinstance(ref, _RemoteSlotRefClaim): + ref.handle._release_slot_ref(ref.token) + else: + ref._release_slot_ref() + del refs[0] def _release_active_remote_slot_refs(self, resources: _RunResources | None = None) -> None: if resources is None: refs = self._active_remote_slot_refs - self._active_remote_slot_refs = [] else: refs = resources.remote_slot_refs - resources.remote_slot_refs = [] self._release_remote_slot_refs(refs) - def _flush_pending_remote_frees(self) -> None: - errors: list[str] = [] - pending_imports = self._pending_remote_import_releases - self._pending_remote_import_releases = [] - remaining_imports: list[RemoteBufferHandle] = [] - for handle in pending_imports: - if handle._live_slot_refs > 0: - remaining_imports.append(handle) - continue - try: - self._send_remote_release_import(handle) - except Exception as exc: # noqa: BLE001 - remaining_imports.append(handle) - errors.append(f"release_import worker_id={handle.worker_id} import_id={handle.import_id}: {exc}") - continue - if handle._owner_handle_ref is not None: - handle._owner_handle_ref._release_import_ref() - handle._owner_handle_ref = None - self._pending_remote_import_releases.extend(remaining_imports) - - pending = self._pending_remote_buffer_frees - self._pending_remote_buffer_frees = [] - remaining: list[RemoteBufferHandle] = [] - for handle in pending: - if handle._live_slot_refs > 0 or handle._live_import_refs > 0: - remaining.append(handle) - continue - try: - self._send_remote_free(handle) - except Exception as exc: # noqa: BLE001 - remaining.append(handle) - errors.append(f"free worker_id={handle.worker_id} buffer_id={handle._buffer_id}: {exc}") - continue - self._pending_remote_buffer_frees.extend(remaining) + def _flush_pending_remote_frees(self) -> None: # noqa: PLR0912 + errors: list[str] = [] + first_async_error: BaseException | None = None + + def remember_error(exc: BaseException, context: str) -> None: + nonlocal first_async_error + if isinstance(exc, Exception): + errors.append(f"{context}: {exc}") + elif first_async_error is None: + first_async_error = exc + + with self._remote_import_release_mu: + release_states = self._pending_remote_import_release_states + pending_imports = self._pending_remote_import_releases + self._pending_remote_import_releases = [] + for handle in release_states: + if handle not in pending_imports: + pending_imports.append(handle) + remaining_imports: list[RemoteBufferHandle] = [] + for handle in pending_imports: + context = f"release_import worker_id={handle.worker_id} import_id={handle.import_id}" + state = release_states.get(handle) + if state is None: + state = _PendingRemoteImportReleaseState( + owner_ref=handle._owner_handle_ref, + owner_ref_token=handle._owner_import_ref_token, + ) + release_states[handle] = state + if handle._live_slot_refs > 0: + remaining_imports.append(handle) + continue + if state.error is not None: + remaining_imports.append(handle) + remember_error(state.error, context) + continue + + rpc_boundary_error: BaseException | None = None + if not state.rpc_complete: + rpc_result = _IsolatedCallResult() + try: + _run_isolated_call( + rpc_result, + lambda: self._send_remote_release_import(handle), + name_prefix="simpler-remote-import-release-", + after_success=lambda: setattr(state, "rpc_complete", True), + ) + except BaseException as exc: # noqa: BLE001 + rpc_boundary_error = exc + if rpc_result.error is not None: + rpc_boundary_error = rpc_result.error + if rpc_boundary_error is not None: + remember_error(rpc_boundary_error, context) + if not state.rpc_complete: + state.error = rpc_boundary_error + self._record_unreclaimable( + "Worker.remote_release_import: the release RPC did not publish completion; " + "the import is retained until whole-tree teardown", + rpc_boundary_error, + ) + remaining_imports.append(handle) + continue + + if state.owner_ref is None: + state.owner_release_complete = True + elif not state.owner_release_complete: + owner_boundary_error: BaseException | None = None + owner_result = _IsolatedCallResult() + owner_ref = state.owner_ref + owner_ref_token = state.owner_ref_token + owner_release = ( + owner_ref._release_import_ref + if owner_ref_token is None + else lambda: owner_ref._release_import_ref(owner_ref_token) + ) + try: + _run_isolated_call( + owner_result, + owner_release, + name_prefix="simpler-remote-import-owner-release-", + after_success=lambda: setattr(state, "owner_release_complete", True), + ) + except BaseException as exc: # noqa: BLE001 + owner_boundary_error = exc + if owner_result.error is not None: + owner_boundary_error = owner_result.error + if owner_boundary_error is not None: + remember_error(owner_boundary_error, context) + if not state.owner_release_complete: + state.error = owner_boundary_error + self._record_unreclaimable( + "Worker.remote_release_import: the release RPC completed but its owner " + "reference could not be retired", + owner_boundary_error, + ) + remaining_imports.append(handle) + continue + + while handle._owner_import_ref_token is not None: + try: + handle._owner_import_ref_token = None + except BaseException as exc: # noqa: BLE001, PERF203 + remember_error(exc, context) + while handle._owner_handle_ref is not None: + try: + handle._owner_handle_ref = None + except BaseException as exc: # noqa: BLE001, PERF203 + remember_error(exc, context) + release_states.pop(handle, None) + self._pending_remote_import_releases.extend(remaining_imports) + + with self._remote_import_release_mu: + # Keep the durable queue intact while RPCs are in flight. If an + # async exception escapes between a successful idempotent free and + # retirement below, the old debt is replayed rather than lost. + pending = list(self._pending_remote_buffer_frees) + remaining: list[RemoteBufferHandle] = [] + for handle in pending: + if handle._live_slot_refs > 0 or handle._live_import_refs > 0: + remaining.append(handle) + continue + try: + self._send_remote_free(handle) + except BaseException as exc: # noqa: BLE001 + remaining.append(handle) + if isinstance(exc, Exception): + errors.append(f"free worker_id={handle.worker_id} buffer_id={handle._buffer_id}: {exc}") + elif first_async_error is None: + first_async_error = exc + continue + self._pending_remote_buffer_frees[:] = remaining + if first_async_error is not None: + raise first_async_error 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 +4611,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 +5783,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 +5889,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 +5960,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 @@ -5025,13 +6326,25 @@ def _create_l3_l2_region(self, worker_id: int, payload_bytes: int, counter_bytes if counter_bytes <= 0 or counter_bytes % 4 != 0: raise ValueError("create_l3_l2_region: counter_bytes must be positive and a multiple of 4") self._validate_l3_l2_worker_id(int(worker_id)) + prior_native_cleanup_error = _l3_host_mapped_region_take_cleanup_error() + if prior_native_cleanup_error: + raise self._record_unreclaimable( + "create_l3_l2_region: a native L3 Host mapping finalized without explicit close on this " + "thread and could not be reclaimed; its owning Worker is no longer identifiable, so this " + "Worker is conservatively stopped and no further work is admitted", + RuntimeError(prior_native_cleanup_error), + ) resources = self._building_run_resources req_shm = SharedMemory(create=True, size=_REGION_CREATE_REQUEST_BYTES) reply_shm = SharedMemory(create=True, size=_REGION_CREATE_REPLY_BYTES) req_buf = cast(memoryview, req_shm.buf) reply_buf = cast(memoryview, reply_shm.buf) region_id = 0 + native_mapping_handle = None l3_host_mapping = None + region = None + dispatched = False + required_ordered_cleanup_before = resources.requires_ordered_cleanup if resources is not None else False try: L3L2RegionCreateRequest( magic_version=_REGION_MAGIC_VERSION, @@ -5041,6 +6354,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 @@ -5055,9 +6373,9 @@ def _create_l3_l2_region(self, worker_id: int, payload_bytes: int, counter_bytes ) counter_offset, total_bytes = validate_region_create_reply(reply, expected_access_profile) if platform.endswith("sim"): - handle = _l3_host_mapped_region_import_sim(reply.backing_shm, int(reply.mapping_bytes)) + native_mapping_handle = _l3_host_mapped_region_import_sim(reply.backing_shm, int(reply.mapping_bytes)) else: - handle = _l3_host_mapped_region_import_onboard( + native_mapping_handle = _l3_host_mapped_region_import_onboard( int(reply.device_id), int(reply.shareable_handle), int(reply.mapping_bytes), @@ -5071,26 +6389,94 @@ def _create_l3_l2_region(self, worker_id: int, payload_bytes: int, counter_bytes payload_bytes=int(reply.desc.payload_bytes), counter_offset=counter_offset, counter_bytes=int(reply.desc.counter_bytes), - handle=int(handle), + handle=native_mapping_handle, ) + native_mapping_handle = None region = L3L2OrchRegion(self, int(worker_id), reply.desc, l3_host_mapping) 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: + mapping_cleanup_error: BaseException | None = None + if region is not None: + try: + self._live_l3_l2_regions.remove(region) + except ValueError: + pass + if resources is not None: + try: + resources.l3_l2_regions.remove(region) + except ValueError: + pass + resources.requires_ordered_cleanup = required_ordered_cleanup_before + region._expire() if l3_host_mapping is not None: try: l3_host_mapping.close() - except RuntimeError: - pass + except BaseException as mapping_exc: # noqa: BLE001 + mapping_cleanup_error = mapping_exc + elif native_mapping_handle is not None: + try: + _l3_host_mapped_region_close(int(native_mapping_handle)) + except BaseException as mapping_exc: # noqa: BLE001 + mapping_cleanup_error = mapping_exc + deferred_native_cleanup_error = _l3_host_mapped_region_take_cleanup_error() + if deferred_native_cleanup_error: + deferred_exc = RuntimeError(deferred_native_cleanup_error) + if mapping_cleanup_error is not None: + combined_exc = RuntimeError( + f"{mapping_cleanup_error}; deferred native cleanup also failed: {deferred_native_cleanup_error}" + ) + combined_exc.__cause__ = mapping_cleanup_error + mapping_cleanup_error = combined_exc + else: + mapping_cleanup_error = deferred_exc + 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, + ) + if mapping_cleanup_error is not None: + raise self._record_unreclaimable( + f"create_l3_l2_region: rollback could not close the L3 Host mapping for region " + f"{region_id} on worker {int(worker_id)}; it is leaked and no further work is admitted", + mapping_cleanup_error, + ) + raise exc finally: del req_buf del reply_buf @@ -5102,28 +6488,47 @@ def _create_l3_l2_region(self, worker_id: int, payload_bytes: int, counter_bytes pass def _cleanup_l3_l2_regions(self, resources: _RunResources | None = None) -> None: - # Per-region best-effort: every region is attempted (and _expire()d) even - # if one raises, so a failing region never strands the rest; the first - # error is raised after all are attempted so close() reports the leak. + # Per-region best-effort: mapping close and child release are independent + # ownership debts, so both are attempted before the region is expired. tracked = self._live_l3_l2_regions if resources is None else resources.l3_l2_regions if not tracked: return regions = list(tracked) - tracked.clear() errors: list[BaseException] = [] for region in regions: try: - try: - region._close_l3_host_mapping() - if self._worker is not None: - self._worker.control_l3_l2_region_release(region._worker_id, region.region_id) - finally: - region._expire() + region._close_l3_host_mapping() except BaseException as exc: # noqa: BLE001 errors.append(exc) - finally: - if resources is not None: - self._live_l3_l2_regions[:] = [live for live in self._live_l3_l2_regions if live is not region] + try: + if self._worker is not None: + self._worker.control_l3_l2_region_release(region._worker_id, region.region_id) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + try: + region._expire() + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + + tracking_lists = (tracked,) if tracked is self._live_l3_l2_regions else (tracked, self._live_l3_l2_regions) + next_tracking_list = 0 + + def retire_tracking() -> None: + nonlocal next_tracking_list + try: + while next_tracking_list < len(tracking_lists): + owned_regions = tracking_lists[next_tracking_list] + owned_regions[:] = [owned for owned in owned_regions if owned is not region] + next_tracking_list += 1 + except Exception as exc: # noqa: BLE001 + errors.append(exc) + next_tracking_list += 1 + retire_tracking() + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + retire_tracking() + + retire_tracking() if errors: raise errors[0] @@ -5157,20 +6562,22 @@ def _ensure_comm_base(self) -> None: device_ids = self._config.get("device_ids", []) rootinfo_path = self._comm_plan_rootinfo_path() - request_shms: dict[int, SharedMemory] = {} # Layout: header (rank, nranks) + NUL-terminated rootinfo_path bytes. path_bytes = rootinfo_path.encode("utf-8") + b"\x00" req_size = _COMM_INIT_HEADER.size + len(path_bytes) - try: + + def initialize(request_owner: _SharedMemoryOwner) -> None: + request_shms: dict[int, SharedMemory] = {} for chip_idx, _device_id in enumerate(device_ids): - req = SharedMemory(create=True, size=req_size) + req = request_owner.create(req_size) + request_shms[chip_idx] = req req_buf = req.buf assert req_buf is not None _COMM_INIT_HEADER.pack_into(req_buf, 0, int(chip_idx), int(len(device_ids))) req_buf[_COMM_INIT_HEADER.size : _COMM_INIT_HEADER.size + len(path_bytes)] = path_bytes - request_shms[chip_idx] = req dw = self._worker + assert dw is not None errors: dict[int, BaseException] = {} def dispatch(chip_idx: int) -> None: @@ -5179,27 +6586,20 @@ def dispatch(chip_idx: int) -> None: except BaseException as e: # noqa: BLE001 errors[chip_idx] = e - threads = [ - threading.Thread(target=dispatch, args=(i,), name=f"comm_init_chip_{i}") for i in range(len(device_ids)) - ] - for t in threads: - t.start() - for t in threads: - t.join() + _start_and_join_threads(range(len(device_ids)), dispatch, name_prefix="comm_init_chip_") if errors: first = next(iter(errors.items())) raise RuntimeError( f"_ensure_comm_base failed on {len(errors)}/{len(device_ids)} chips; " f"first error chip={first[0]}: {first[1]}" ) - finally: - for shm in request_shms.values(): - try: - shm.close() - shm.unlink() - except Exception: # noqa: BLE001 - pass - self._comm_base_ready = True + + _run_with_owned_shared_memory( + len(device_ids), + initialize, + name_prefix="shm-comm-init-lifecycle-", + after_success=lambda: setattr(self, "_comm_base_ready", True), + ) def _allocate_domain( # noqa: PLR0912 -- linear input-validation + per-chip shm staging + dispatch + reply unpack; splitting obscures the fail-fast ordering self, @@ -5210,42 +6610,10 @@ def _allocate_domain( # noqa: PLR0912 -- linear input-validation + per-chip shm buffers: list[CommBufferSpec], ) -> CommDomainHandle: # Admission is the run() lease that the driving orchestrator holds; - # this checks resource presence (not the public lifecycle) so a domain - # allocation admitted before a concurrent close() published CLOSED still - # completes during the drain. The _worker check below is that gate. - if self.level < 3: - raise RuntimeError("allocate_domain requires level >= 3") - if self._worker is None: - raise RuntimeError("allocate_domain requires a hierarchical Worker (_start_hierarchical ran)") - # A domain's lifetime ends at the fence of the run that allocated it, - # so there is no owner to charge it to outside graph construction. - resources = self._building_run_resources - if resources is None: - raise RuntimeError("allocate_domain is only valid while a run's graph is being built") - if not workers: - raise ValueError("allocate_domain: workers must be non-empty") - if len(set(workers)) != len(workers): - raise ValueError(f"allocate_domain: workers contains duplicates: {workers}") - device_ids = self._config.get("device_ids", []) - for w in workers: - if w < 0 or w >= len(device_ids): - raise ValueError(f"allocate_domain: worker_id {w} outside [0, {len(device_ids)})") - if window_size <= 0: - raise ValueError("allocate_domain: window_size must be positive") - buffer_names = [b.name for b in buffers] - if len(set(buffer_names)) != len(buffer_names): - raise ValueError(f"allocate_domain: duplicate buffer names: {buffer_names}") - # Check buffer carving fits in window BEFORE dispatching: a chip-side - # overflow would still register the backend allocation (aclrtMalloc - # already succeeded) but never produce a Handle on the parent, so it - # would silently leak. Fail fast here instead. - total_buffer_nbytes = sum(int(b.nbytes) for b in buffers) - if total_buffer_nbytes > window_size: - raise ValueError( - f"allocate_domain: buffers sum to {total_buffer_nbytes} bytes, exceeds window_size={window_size}" - ) - if name in self._live_domains: - raise ValueError(f"allocate_domain: domain {name!r} already live") + # validation checks resource presence rather than the public lifecycle, + # so an allocation admitted before concurrent close still drains. + # Buffer carving is checked before communicator/device side effects. + resources = _validate_domain_allocation(self, name, workers, window_size, buffers) # Lazy base communicator: first orch.allocate_domain on this Worker # triggers HCCL RootInfo handshake + EnablePeerAccess on every chip. @@ -5267,100 +6635,139 @@ def _allocate_domain( # noqa: PLR0912 -- linear input-validation + per-chip shm # this, `workers.index(chip_idx)` makes the hot path quadratic. worker_to_rank = {w: r for r, w in enumerate(workers)} - request_shms: dict[int, SharedMemory] = {} - reply_shms: dict[int, SharedMemory] = {} - try: - for chip_idx in workers: - req = SharedMemory(create=True, size=req_size) - req_buf = req.buf - assert req_buf is not None - _DOMAIN_REQ_HEADER.pack_into( - req_buf, - 0, - int(allocation_id), - int(len(workers)), - int(worker_to_rank[chip_idx]), # domain_rank - int(window_size), - int(buffer_count), - ) - nbytes_off = _DOMAIN_REQ_HEADER.size - if buffer_count: - struct.pack_into(f"<{buffer_count}Q", req_buf, nbytes_off, *[int(b.nbytes) for b in buffers]) - rank_ids_off = nbytes_off + buffer_count * 8 - struct.pack_into(f"<{len(workers)}I", req_buf, rank_ids_off, *[int(w) for w in workers]) - request_shms[chip_idx] = req - - 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, - ) + handle: CommDomainHandle | None = None + contexts: dict[int, ChipDomainContext] = {} + previous_ordered_cleanup = resources.requires_ordered_cleanup - 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}" + def allocate(staged_shms: _SharedMemoryOwner) -> CommDomainHandle: + nonlocal handle + request_shms: dict[int, SharedMemory] = {} + reply_shms: dict[int, SharedMemory] = {} + dispatch_started = False + try: + for chip_idx in workers: + req = staged_shms.create(req_size) + request_shms[chip_idx] = req + req_buf = req.buf + assert req_buf is not None + _DOMAIN_REQ_HEADER.pack_into( + req_buf, + 0, + int(allocation_id), + int(len(workers)), + int(worker_to_rank[chip_idx]), # domain_rank + int(window_size), + int(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( + nbytes_off = _DOMAIN_REQ_HEADER.size + if buffer_count: + struct.pack_into(f"<{buffer_count}Q", req_buf, nbytes_off, *[int(b.nbytes) for b in buffers]) + rank_ids_off = nbytes_off + buffer_count * 8 + struct.pack_into(f"<{len(workers)}I", req_buf, rank_ids_off, *[int(w) for w in workers]) + reply = staged_shms.create(reply_size) + reply_shms[chip_idx] = reply + + # Ownership is complete before a chip can commit an allocation. + # Until the replies are sampled, every requested rank is retained + # conservatively; an interrupted sampler may cost a failed release, + # but cannot make an allocated window unreachable. + handle = CommDomainHandle( 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)}, + workers=workers, + contexts={}, + allocation_id=allocation_id, + _release_fn=lambda released, owner=resources: self._release_domain_handle(released, owner), + _domain_size=len(workers), + _domain_ranks=worker_to_rank, ) - finally: - # Close + unlink local copies regardless of outcome. Children - # have already finished reading by the time CONTROL_DONE fires. - for shm in request_shms.values(): - try: - shm.close() - shm.unlink() - except Exception: # noqa: BLE001 - pass - for shm in reply_shms.values(): + self._live_domains[name] = handle + resources.live_domains[name] = handle + resources.requires_ordered_cleanup = True + + # 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. + dispatch_started = True try: - shm.close() - shm.unlink() - except Exception: # noqa: BLE001 - pass + self._dispatch_control_domain( + workers=workers, + request_shms=request_shms, + reply_shms=reply_shms, + op="alloc", + allocation_id=allocation_id, + ) + + 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) + ) + 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)}, + ) + handle.contexts = contexts + finally: + committed = tuple(w for w in workers if _domain_reply_committed(reply_shms.get(w))) + handle.workers = committed + if not committed: + if self._live_domains.get(name) is handle: + self._live_domains.pop(name) + if resources.live_domains.get(name) is handle: + resources.live_domains.pop(name) + resources.requires_ordered_cleanup = previous_ordered_cleanup + except BaseException: # noqa: BLE001 + if handle is not None and not dispatch_started: + if self._live_domains.get(name) is handle: + self._live_domains.pop(name) + if resources.live_domains.get(name) is handle: + resources.live_domains.pop(name) + resources.requires_ordered_cleanup = previous_ordered_cleanup + raise + assert handle is not None + return handle - handle = CommDomainHandle( - name=name, - workers=workers, - contexts=contexts, - allocation_id=allocation_id, - _release_fn=lambda released, owner=resources: self._release_domain_handle(released, owner), + def publish_provenance() -> None: + assert handle is not None + # The backend windows are now live: record each chip's window base + # and every carved buffer pointer before the lifecycle publishes + # success back to the interruptible caller. + buf_nbytes = {b.name: int(b.nbytes) for b in buffers} + with self._child_prov_lock: + for chip_idx, ctx in contexts.items(): + self._child_prov_record_domain( + chip_idx, int(ctx.local_window_base), allocation_id, int(ctx.actual_window_size) + ) + for buf_name, buf_ptr in ctx.buffer_ptrs.items(): + self._child_prov_record_domain(chip_idx, int(buf_ptr), allocation_id, buf_nbytes[buf_name]) + + published_handle = _run_with_owned_shared_memory( + len(workers) * 2, + allocate, + name_prefix="shm-domain-alloc-lifecycle-", + after_success=publish_provenance, ) - self._live_domains[name] = handle - if resources is not None: - resources.live_domains[name] = handle - # 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 - # _release_domain_now just before the backend free (a commit barrier), - # not by the deferred marker — so the deferred window stays dispatchable. - buf_nbytes = {b.name: int(b.nbytes) for b in buffers} - with self._child_prov_lock: - for chip_idx, ctx in contexts.items(): - self._child_prov_record_domain( - chip_idx, int(ctx.local_window_base), allocation_id, int(ctx.actual_window_size) - ) - for buf_name, buf_ptr in ctx.buffer_ptrs.items(): - self._child_prov_record_domain(chip_idx, int(buf_ptr), allocation_id, buf_nbytes[buf_name]) + assert handle is not None and published_handle is handle return handle def _release_domain_handle(self, handle: CommDomainHandle, resources: _RunResources) -> None: @@ -5380,19 +6787,29 @@ def _release_domain_handle(self, handle: CommDomainHandle, resources: _RunResour if self._worker is None: return with resources.domain_lock: - resources.live_domains.pop(handle.name, None) - # Pop from live_domains so a subsequent allocate_domain(name=...) - # call within the same run can reuse the name. The actual memory - # is still live until _execute_pending_domain_releases runs. - if self._live_domains.get(handle.name) is handle: - self._live_domains.pop(handle.name) if not resources.retired: + # Publish the fence-owned claim before retiring either live + # registry. If append is interrupted, those registries still + # make the allocation reachable by the end-of-run/close sweep. resources.pending_release_domains.append(handle) + resources.live_domains.pop(handle.name, None) + # Pop from _live_domains so a subsequent allocation in this + # run can reuse the name. The pending claim keeps the old + # allocation alive until the fence drains it. + if self._live_domains.get(handle.name) is handle: + self._live_domains.pop(handle.name) return # Deferral exists so tasks that captured this domain still see live # memory; the owning run's fence has passed, so there is nothing left - # to defer behind. Freed outside the lock — it is a backend call. + # to defer behind. Keep both live registries as the durable owner while + # the backend call runs; exact-once free removes the global entry only + # after it commits, and the run entry is retired below. self._free_domain_after_fence(handle) + with resources.domain_lock: + if resources.live_domains.get(handle.name) is handle: + resources.live_domains.pop(handle.name) + if self._live_domains.get(handle.name) is handle: + self._live_domains.pop(handle.name) def _retire_run_domains(self, resources: _RunResources) -> None: """Close this run's deferred-release path and free anything left on it. @@ -5404,9 +6821,23 @@ def _retire_run_domains(self, resources: _RunResources) -> None: with resources.domain_lock: resources.retired = True stragglers = list(resources.pending_release_domains) - resources.pending_release_domains.clear() - for handle in stragglers: + self._drain_pending_domain_snapshot(resources, stragglers) + + def _drain_pending_domain_snapshot(self, resources: _RunResources, pending: list[CommDomainHandle]) -> None: + """Free a snapshot while each source-queue claim stays durable.""" + + def _release(handle: CommDomainHandle) -> None: self._free_domain_after_fence(handle) + # Retire only after exact-once backend success. An interruption + # before this deletion leaves a replayable claim; an interruption + # after it is safe because the allocation is already gone. + with resources.domain_lock: + for index, candidate in enumerate(resources.pending_release_domains): + if candidate is handle: + del resources.pending_release_domains[index] + break + + _raise_first(_release, pending) def _free_domain_after_fence(self, handle: CommDomainHandle) -> None: """Back-end free for a handle whose owning run has retired. @@ -5414,45 +6845,32 @@ 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() + # Snapshot under the same lock release() uses to append. Keep every + # claim in the source queue while the backend call is in flight: a + # release appended after this snapshot then remains for retirement's + # final drain instead of being erased by an unlocked clear(). + with resources.domain_lock: + pending = list(resources.pending_release_domains) + self._drain_pending_domain_snapshot(resources, pending) def _release_domain_now(self, handle: CommDomainHandle) -> None: """Synchronous backend release for one handle, exactly once. @@ -5476,48 +6894,71 @@ def _release_domain_now(self, handle: CommDomainHandle) -> None: if failure is not None: raise failure return - try: - self._release_domain_claimed(handle) - except BaseException as exc: - self._domain_free_results[handle.allocation_id] = exc - raise - self._domain_free_results[handle.allocation_id] = None + # Absence means no backend attempt has been claimed. Once the + # isolated target is admitted it publishes a conservative outcome + # before the backend call, so even failure to store the real + # post-RPC result cannot turn a committed free back into an + # apparently unclaimed allocation that a waiter replays. Publishing + # inside the target matters: an interruption before admission must + # leave the allocation retryable because no backend work began. + unpublished_outcome = RuntimeError( + "CommDomain backend release outcome publication was interrupted; refusing to replay the free" + ) + result = _IsolatedCallResult() + + def release_claim() -> None: + self._domain_free_results[handle.allocation_id] = unpublished_outcome + try: + self._release_domain_claimed(handle) + except BaseException as exc: + self._domain_free_results[handle.allocation_id] = exc + raise + + # The operation and its success publication run outside the + # caller's async-exception boundary while contenders remain behind + # _domain_free_mu. A KeyboardInterrupt after helper completion can + # propagate, but a retry already observes the committed outcome. + _run_isolated_call( + result, + release_claim, + name_prefix="simpler-domain-free-", + after_success=lambda: self._domain_free_results.__setitem__(handle.allocation_id, None), + ) + if result.error is not None: + raise result.error + if not result.completed: + raise RuntimeError("CommDomain backend release did not publish an outcome") def _release_domain_claimed(self, handle: CommDomainHandle) -> None: """Drive CTRL_RELEASE_DOMAIN. Caller holds this allocation's claim.""" - # Revoke provenance BEFORE the physical free: once release begins the - # domain's pointers are no longer dispatchable. Revoking after the - # backend free would leave a use-after-free window (a concurrent - # copy/dispatch could still validate the being-freed pointer as live), - # and a partial/failed release would strand a freed pointer as "live" - # forever. Dropping first is the safe direction — a leak (if the backend - # free later fails) is recoverable; a use-after-free is not. - with self._child_prov_lock: - self._child_prov_drop_domain(handle.allocation_id) workers = handle.workers # Release payload is just the fixed header — no rank_ids tail; the # backend looked them up from its own per-allocation record at # alloc time and doesn't need them again. req_size = _DOMAIN_REQ_HEADER.size - worker_to_rank = {w: r for r, w in enumerate(workers)} - request_shms: dict[int, SharedMemory] = {} - try: + def release(request_owner: _SharedMemoryOwner) -> None: + # Revoke provenance BEFORE the physical free: once release begins + # the domain's pointers are no longer dispatchable. Dropping first + # makes an interrupted/failed release a recoverable leak instead of + # leaving a use-after-free validation window. + with self._child_prov_lock: + self._child_prov_drop_domain(handle.allocation_id) + request_shms: dict[int, SharedMemory] = {} for chip_idx in workers: - req = SharedMemory(create=True, size=req_size) + req = request_owner.create(req_size) + request_shms[chip_idx] = req req_buf = req.buf assert req_buf is not None _DOMAIN_REQ_HEADER.pack_into( req_buf, 0, int(handle.allocation_id), - int(len(workers)), - int(worker_to_rank[chip_idx]), + int(handle._domain_size), # noqa: SLF001 -- backend release identity belongs to the handle + int(handle._domain_ranks[chip_idx]), # noqa: SLF001 -- preserve the allocation-time rank 0, # window_size — ignored on release 0, # buffer_count — ignored on release ) - request_shms[chip_idx] = req - self._dispatch_control_domain( workers=workers, request_shms=request_shms, @@ -5525,15 +6966,17 @@ def _release_domain_claimed(self, handle: CommDomainHandle) -> None: op="release", allocation_id=handle.allocation_id, ) - finally: - for shm in request_shms.values(): - try: - shm.close() - shm.unlink() - except Exception: # noqa: BLE001 - pass - if self._live_domains.get(handle.name) is handle: - self._live_domains.pop(handle.name) + + def retire_live_handle() -> None: + if self._live_domains.get(handle.name) is handle: + self._live_domains.pop(handle.name) + + _run_with_owned_shared_memory( + len(workers), + release, + name_prefix="shm-domain-release-lifecycle-", + after_success=retire_live_handle, + ) def _dispatch_control_domain( self, @@ -5565,11 +7008,7 @@ def dispatch(chip_idx: int) -> None: except BaseException as e: # noqa: BLE001 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() + _start_and_join_threads(workers, dispatch, name_prefix=f"{op}_domain_chip_") if errors: first = next(iter(errors.items())) @@ -5585,29 +7024,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 +7607,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 +7641,187 @@ 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 _recover_interrupted_run_finalization(self, handle: RunHandle, error: BaseException) -> BaseException: + """Retire an accepted handle after its finalizer escaped unexpectedly.""" + poison = RuntimeError( + "RunHandle finalization escaped before retirement; the run is retained until whole-tree teardown " + "and no further work is admitted" + ) + poison.__cause__ = error + self._publish_abandoned_run(handle, poison) + return error + + def _publish_abandoned_run(self, handle: RunHandle, poison: BaseException) -> None: + """Publish poison and abandonment before accepted-set removal.""" + try: + while True: + try: + with self._hierarchical_start_cv: + if handle in self._accepted_run_handles: + if self._ordered_cleanup_error is None: + self._ordered_cleanup_error = poison + handle._finalization_abandoned = True + if handle not in self._abandoned_run_handles: + self._abandoned_run_handles.append(handle) + handle._cleanup_published = True + self._accepted_run_handles.discard(handle) + self._hierarchical_start_cv.notify_all() + return + except BaseException: # noqa: BLE001, PERF203 + # Publication fields are monotonic; repetition never + # replays cleanup or a native release. + pass + except BaseException: # noqa: BLE001 + self._publish_abandoned_run(handle, poison) + + def _drain_abandoned_run_keepalives(self) -> BaseException | None: + """Drop retained run references after native teardown has succeeded.""" + cursor = _AbandonedRunKeepaliveCursor(self._abandoned_run_handles) + cursor.drain() + return cursor.first_error + + def _abandon_unsettled_run( + self, handle: RunHandle, message: str, cause: BaseException | None = None + ) -> RuntimeError: + """Retire a run whose native cancellation fence could not be settled. + + Waiting would be an unbounded operation because native never promised a + terminal fence. Retain every run-owned Python reference until whole-tree + teardown, publish the sticky poison before removing the handle from the + accepted drain set, and make the otherwise-private handle terminal. + """ + poison = handle._finalization_error + try: + while not handle._terminal: + try: + if poison is None: + poison = RuntimeError(message) + if cause is not None: + poison.__cause__ = cause + handle._cache_finalization_error(poison) + self._publish_abandoned_run(handle, poison) + handle._publish_terminal(poison) + except BaseException: # noqa: BLE001, PERF203 + # Every field above is monotonic. Retrying cannot restore + # acceptance or release retained ownership early. + pass + assert isinstance(handle._error, RuntimeError) + return handle._error + except BaseException: # noqa: BLE001 + return self._abandon_unsettled_run(handle, message, cause) + + 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 +7837,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) @@ -6232,19 +7848,31 @@ def _submit_l3_locked(self, callable, args, cfg: CallConfig) -> RunHandle: scope_open = False self._orch._scope_end() finally: - try: - self._orch._fail_run_submission(run_id, _format_exc("orchestration", e)) - except Exception: - # The orchestration cause is the useful one, so a failure - # while closing the run must not replace it; the handle's - # own wait still reports the run's terminal state. - pass - finally: - # Graph-construction failures remain synchronous, but any - # tasks already submitted still own their resources until - # the run fence fires. - self._building_run_resources = None + cancellation_error: BaseException | None = None + failure_text = _format_exc("orchestration", e) + for _ in range(_RUN_CANCELLATION_ATTEMPTS): + try: + self._orch._fail_run_submission(run_id, failure_text) + except BaseException as exc: # noqa: BLE001 + cancellation_error = exc + else: + cancellation_error = None + break + + # Graph-construction failures remain synchronous, but any + # tasks already submitted still own their resources until + # either the run fence fires or whole-tree teardown reclaims an + # unsettled cancellation. + self._building_run_resources = None + if cancellation_error is None: handle._wait_for_serialization() + else: + self._abandon_unsettled_run( + handle, + "Worker.submit(): native cancellation did not settle after its bounded retry; " + "the run fence is nonterminal and no further work is admitted", + cancellation_error, + ) raise finally: self._building_run_resources = None @@ -6266,54 +7894,103 @@ def _wait_run_handle_accepted(self, run_id: int) -> None: self._orch._wait_run_accepted(run_id) def _finalize_run_handle( - self, handle: RunHandle, run_id: int, native_error: BaseException | None + self, + handle: RunHandle, + run_id: int, + native_error: BaseException | None, + *, + _after_step: Any | None = None, ) -> BaseException | None: """Run fence-owned cleanup exactly once and return the cached result.""" - result = native_error - - def _step(fn) -> None: - nonlocal result - try: - fn() - except BaseException as exc: # noqa: BLE001 - if result is None: - result = exc - - resources = handle._resources - if native_error is not None: - _step(lambda: self._poison_l3_l2_region_from_endpoint_error(native_error, resources)) - _step(lambda: self._release_active_remote_slot_refs(resources)) - _step(self._flush_pending_remote_frees) - try: - _step(lambda: self._cleanup_l3_l2_regions(resources)) - finally: - resources.l3_l2_orch_comm_host_buffers.clear() - _step(lambda: self._execute_pending_domain_releases(resources)) - if resources.live_domains: - _step(lambda: self._release_all_live_domains(resources)) - _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") - 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. + # 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. + cursor: _RunFinalizationCursor | None = None + outer_error: BaseException | None = None try: - with self._hierarchical_start_cv: - self._accepted_run_handles.discard(handle) - self._hierarchical_start_cv.notify_all() + resources = handle._resources + orch = self._orch + + def _poison_endpoint() -> None: + if native_error is not None: + self._poison_l3_l2_region_from_endpoint_error(native_error, resources) + + def _release_native_run() -> None: + if orch is None: + raise RuntimeError("RunHandle cleanup lost its native Orchestrator") + orch._release_run(run_id) + + cursor = _RunFinalizationCursor( + steps=( + ("endpoint_poison", _poison_endpoint), + ("remote_slot_refs", lambda: self._release_active_remote_slot_refs(resources)), + ("remote_frees", self._flush_pending_remote_frees), + ("l3_l2_regions", lambda: self._cleanup_l3_l2_regions(resources)), + ("l3_l2_host_buffers", resources.l3_l2_orch_comm_host_buffers.clear), + ("pending_domains", lambda: self._execute_pending_domain_releases(resources)), + ("live_domains", lambda: self._release_all_live_domains(resources)), + ("retire_domains", lambda: self._retire_run_domains(resources)), + ("native_run", _release_native_run), + ) + ) + cursor.drain(_after_step) 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 + # Any escape outside a committed cursor edge has ambiguous cleanup + # ownership. Native teardown is its only safe reclamation boundary. + outer_error = exc + if cursor is not None: + cursor.remember_boundary_error(exc) + cursor.incomplete = True + + cleanup_error = None if cursor is None else cursor.cleanup_error + boundary_error = outer_error if cursor is None else cursor.boundary_error + incomplete = cursor is None or cursor.incomplete or not cursor.exhausted + abandonment_error: RuntimeError | None = None + if incomplete: + abandonment_error = RuntimeError( + "RunHandle finalization stopped at an ambiguous cleanup boundary; the run is retained until " + "whole-tree teardown and no further work is admitted" + ) + abandonment_error.__cause__ = boundary_error if boundary_error is not None else cleanup_error + + # Poison/cleanup publication precedes accepted-set removal under the + # same lifecycle lock. A successor therefore sees either the accepted + # predecessor or the sticky poison, including the conservative path + # whose keepalives stay owned until native teardown succeeds. + while True: + try: + if native_error is not None: + result = native_error + elif cleanup_error is not None: + result = cleanup_error + elif boundary_error is not None: + result = boundary_error + else: + result = abandonment_error + handle._cache_finalization_error(result) + with self._hierarchical_start_cv: + if incomplete: + assert abandonment_error is not None + if self._ordered_cleanup_error is None: + self._ordered_cleanup_error = abandonment_error + handle._finalization_abandoned = True + if handle not in self._abandoned_run_handles: + self._abandoned_run_handles.append(handle) + elif 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() + break + except BaseException as exc: # noqa: BLE001, PERF203 + if boundary_error is None: + boundary_error = exc + # 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. + return handle._finalization_error @property def aicpu_dlopen_count(self) -> int: @@ -6414,9 +8091,9 @@ def close(self) -> None: # noqa: PLR0912, PLR0915 -- lifecycle linearization: r - Teardown is single-shot. Once it runs, a later ``close()`` re-raises the same result rather than re-driving a half-torn tree. A retry is permitted only where the drain did not complete and so left teardown - un-attempted and the tree intact: either in-flight leases outlived - the cleanup budget, or an asynchronous interruption left an accepted - run fence undrained. + un-attempted and the tree intact: either in-flight leases or accepted + run fences outlived the shared cleanup budget, or an asynchronous + interruption left a run fence undrained. - Native teardown runs on the ``init()``-owner thread, being device-bound. """ # close() is a permanent commitment against a resource, not a reversible @@ -6432,25 +8109,28 @@ def close(self) -> None: # noqa: PLR0912, PLR0915 -- lifecycle linearization: r # resource leaks and a later close() re-raises the same result — it # never re-drives a half-torn tree. A later close() may retry only # where the drain left teardown un-attempted and the tree intact — - # leases outliving the budget, or an async interruption leaving an - # accepted run fence undrained; a tree with a live op is never torn - # down; + # leases or accepted run fences outliving the shared budget, or an + # async interruption leaving a run fence undrained; a tree with a + # live op is never torn down; # - native teardown runs only on the init-owner thread (device-bound). # `attempt` is None until the claim installs it. The pre-claim checks # raise/return before that, so the finally skips completion for them. From # the claim on, the attempt is completed in an innermost resilient finally - # whose only work is three plain attribute assigns (`error`, `incomplete`, - # then `done`) followed by a locked `notify_all()`. Every fallible step — + # by one immutable outcome-reference publication followed by a locked + # `notify_all()`. Every fallible step — # drain, teardown, residual synthesis, registry detach — runs before it # and folds its error into `result`. `done` is set BEFORE the CV acquire, # so an async BaseException in the (interruptible) acquire or the notify # cannot strand a joiner — the joiner's bounded re-check recovers a - # skipped notify. The only irreducible window is an async exception landing - # between the `error`/`incomplete` and `done` plain assigns. + # skipped notify. A pre-publication interruption becomes that immutable + # outcome's error; an interruption after publication cannot retract it. attempt: _CloseAttempt | None = None result: BaseException | None = None teardown_tree = False + teardown_completed = False drain_complete = False + drain_deadline: float | None = None + drain_deadline_expired = False handles_to_drain: tuple[RunHandle, ...] = () try: with self._hierarchical_start_cv: @@ -6503,6 +8183,10 @@ def close(self) -> None: # noqa: PLR0912, PLR0915 -- lifecycle linearization: r attempt = _CloseAttempt() self._close_completion = attempt self._hierarchical_start_cv.notify_all() + # One absolute budget covers both kinds of admitted work. A + # lease that consumes most of it must not be followed by a + # fresh full-budget wait on an accepted run fence. + drain_deadline = time.monotonic() + _ROLLBACK_GRACEFUL_TIMEOUT_S # Drain in-flight leases before touching the tree. CLOSED already # rejects new leases; a tree with a live op is never torn down. # If an op outlives the budget, teardown stays UN-attempted and @@ -6510,7 +8194,6 @@ def close(self) -> None: # noqa: PLR0912, PLR0915 -- lifecycle linearization: r # one of the two paths that keep close() retryable (the other is # an async interruption leaving an accepted run fence undrained). if self._active_ops > 0: - drain_deadline = time.monotonic() + _ROLLBACK_GRACEFUL_TIMEOUT_S while self._active_ops > 0: remaining = drain_deadline - time.monotonic() if remaining <= 0: @@ -6528,17 +8211,31 @@ def close(self) -> None: # noqa: PLR0912, PLR0915 -- lifecycle linearization: r # CLOSED prevents new admissions and active operations have # drained, so this is the complete accepted set. Wait outside # the lifecycle CV: handle retirement acquires the same lock. + assert drain_deadline is not None for handle in handles_to_drain: + remaining = drain_deadline - time.monotonic() + if remaining <= 0: + drain_deadline_expired = True + break try: - handle.wait() + handle.wait(remaining) except BaseException as exc: # noqa: BLE001 if result is None: result = exc + if isinstance(exc, TimeoutError) and time.monotonic() >= drain_deadline: + drain_deadline_expired = True + break with self._hierarchical_start_cv: if self._accepted_run_handles: - # An asynchronous interruption left at least one fence - # undrained. Keep teardown retryable and the tree intact. + # A timeout or asynchronous interruption left at least + # one fence undrained. Keep teardown retryable and the + # tree intact. drain_complete = False + if drain_deadline_expired and (result is None or isinstance(result, TimeoutError)): + result = TimeoutError( + "Worker.close(): run fence(s) still in flight after the cleanup budget " + f"({_ROLLBACK_GRACEFUL_TIMEOUT_S}s); teardown deferred (worker stays CLOSED)" + ) else: teardown_tree = self._has_live_resources() # Fence drain makes this close terminal even when the @@ -6547,6 +8244,7 @@ def close(self) -> None: # noqa: PLR0912, PLR0915 -- lifecycle linearization: r self._teardown_attempted = teardown_tree or result is not None if teardown_tree: self._teardown_ready_tree() + teardown_completed = True except BaseException as exc: # noqa: BLE001 if result is None: result = exc @@ -6555,6 +8253,14 @@ def close(self) -> None: # noqa: PLR0912, PLR0915 -- lifecycle linearization: r had_live = True # conservative default if a read below is interrupted detached_registry: tuple[dict, dict, dict] | None = None try: + # A successful tree teardown is the reclamation boundary + # for abandoned run ownership. Drain it before the residual + # inventory: that diagnostic probe is interruptible and an + # interruption must not retain already-safe references. + if teardown_completed: + retained_error = self._drain_abandoned_run_keepalives() + if result is None and retained_error is not None: + result = retained_error had_live = self._has_live_resources() # Terminal teardown is single-shot and best-effort: a resource # it could not reclaim LEAKS. Never return success with a @@ -6586,17 +8292,22 @@ def close(self) -> None: # noqa: PLR0912, PLR0915 -- lifecycle linearization: r if result is None: result = exc finally: - # Innermost, resilient publish: all plain attribute assigns - # (error/incomplete first, then `done`), so a joiner that - # observes `done` always observes the result. `done` is set - # BEFORE acquiring the CV — the only remaining work under the - # lock is notify_all(). A BaseException during the - # (interruptible, possibly-blocking) CV acquire therefore - # cannot strand the attempt at done=False; a joiner's bounded - # re-check then recovers the skipped notify on its own. - attempt.error = result - attempt.incomplete = result is not None or had_live - attempt.done = True + # The immutable outcome reference is the completion flag and + # result. A reader can therefore never observe completion + # without its error/incomplete payload. Publication precedes + # notification; bounded re-checks cover a skipped wakeup. + published: _CloseOutcome | None = None + incomplete = result is not None or had_live + while published is None: + try: + published = attempt.publish(result, incomplete) + except BaseException as exc: # noqa: BLE001, PERF203 + if attempt.done: + raise + if result is None: + result = exc + incomplete = True + result = published.error with self._hierarchical_start_cv: self._hierarchical_start_cv.notify_all() # Post-completion, lock-free: dropping the last refs may run a diff --git a/src/common/hierarchical/orchestrator.cpp b/src/common/hierarchical/orchestrator.cpp index 20a17493c..63b82d247 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,129 @@ 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_++; + bool map_published = false; + bool fifo_published = false; + try { + auto run = std::make_shared(run_id, *lease); + bool inserted = runs_.emplace(run_id, run).second; + if (!inserted) throw std::logic_error("Orchestrator::begin_run: duplicate run id"); + map_published = true; + if (test_hook_) test_hook_(OrchestratorTestPoint::BEGIN_RUN_MAP_PUBLISHED); + + run_fifo_.push_back(run_id); + fifo_published = true; + if (test_hook_) test_hook_(OrchestratorTestPoint::BEGIN_RUN_FIFO_PUBLISHED); + + building_run_id_ = run_id; + run->phase.store(RunPhase::BUILDING, std::memory_order_release); + } catch (...) { + if (fifo_published) run_fifo_.pop_back(); + if (map_published) runs_.erase(run_id); + pipeline_slots_.release(*lease); + runs_cv_.notify_all(); + throw; + } } - 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 +202,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 +233,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 +298,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 +351,83 @@ 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. + struct FailurePropagation { + TaskSlot slot{INVALID_SLOT}; + std::vector producers; + bool owns_releases{false}; + }; + std::vector cancelled; + cancelled.reserve(slots.size()); + for (TaskSlot slot : slots) { + TaskSlotState &state = slot_state(slot); + std::optional claimed = claim_task_failure(state, message); + if (claimed.has_value() && *claimed != TaskState::BUILDING) { + // Cancellation is the sole propagation owner for a READY/PENDING + // claim it won. Publish retryable debt before any preparation can + // throw; scheduler-owned claims deliberately expose no such debt. + state.failure_propagation_pending.store(true, std::memory_order_release); + } + // 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); + if (test_hook_) test_hook_(OrchestratorTestPoint::FAILURE_FANIN_SNAPSHOT); + FailurePropagation propagation; + propagation.slot = slot; + { + std::lock_guard lk(state.fanout_mu); + propagation.producers = state.fanin_producers; + } + cancelled.push_back(std::move(propagation)); + } + // 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. Cancellation owns every debt in this batch exclusively; the + // exchange settles each debt immediately before its non-repeatable + // releases. + for (FailurePropagation &propagation : cancelled) { + propagation.owns_releases = commit_failure_propagation(slot_state(propagation.slot)); + } + for (const FailurePropagation &propagation : cancelled) { + if (propagation.owns_releases) try_consume(propagation.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 (const FailurePropagation &propagation : cancelled) { + if (!propagation.owns_releases) continue; + for (TaskSlot producer : propagation.producers) + try_consume(producer); + } } void Orchestrator::decrement_run_tasks(RunId run_id) { @@ -295,6 +503,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,36 +601,63 @@ Tensor Orchestrator::alloc(const std::vector &shape, DataType dtype) { TaskSlotState &s = slot_state(ar.slot); s.reset(); s.run_id = run->id; - - uint64_t ptr = reinterpret_cast(ar.heap_ptr); - if (ptr != 0) { - TensorKey key = TensorKey::local_host(ptr); - tensormap_->insert(run->id, key, ar.slot); - s.output_keys.push_back(key); + s.pipeline_lease = run->lease; + // From the moment the run owns this slot until COMPLETED publication, a + // failed alloc must remain claimable by run cancellation. FREE would make + // cancellation skip the registered slot and strand active_tasks forever. + s.state.store(TaskState::BUILDING, std::memory_order_release); + try { + if (test_hook_) test_hook_(OrchestratorTestPoint::ALLOC_RUN_SLOT_REGISTERING); + register_run_slot(run, ar.slot); + } catch (...) { + // Registration itself has strong ownership semantics: push_back either + // published the slot or threw without changing the run. On the latter + // path no cancellation can discover it, so release the Ring claim here. + s.state.store(TaskState::CONSUMED, std::memory_order_release); + allocator_->release(ar.slot); + throw; } // 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 // will increment fanout_total in infer_deps. - int32_t scope_ref = (scope_->depth() > 0) ? 1 : 0; + const int32_t scope_ref = (scope_->depth() > 0) ? 1 : 0; + if (scope_ref > 0) { + // Register before charging fanout_total. If registration throws, scope + // teardown cannot release this slot, so charging first would leave a + // reference neither scope_end nor cancellation can satisfy. + scope_->register_task(ar.slot); + if (test_hook_) test_hook_(OrchestratorTestPoint::SCOPE_REGISTERED); + } { std::lock_guard lk(s.fanout_mu); s.fanout_total = scope_ref; } + + uint64_t ptr = reinterpret_cast(ar.heap_ptr); + if (ptr != 0) { + TensorKey key = TensorKey::local_host(ptr); + // Prepare the cleanup journal before publishing the TensorMap entry. + // A failed insert then leaves either no mapping or a mapping cancellation + // can erase; the reverse order can leak an unjournaled producer entry. + s.output_keys.push_back(key); + if (test_hook_) test_hook_(OrchestratorTestPoint::ALLOC_OUTPUT_KEY_PREPARED); + tensormap_->insert(run->id, key, ar.slot); + } + // Simulate the self try_consume that on_task_complete would normally // contribute for a slot that ran through the scheduler. Without this // bump, the fanout-release threshold (`>= total + 1`) would be one // short and the slot would never reach CONSUMED. + // Publish it only after every fallible BUILDING step: a failed alloc gets + // its terminal self release from cancellation instead. s.fanout_released.store(1, std::memory_order_relaxed); - if (scope_ref > 0) scope_->register_task(ar.slot); 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 +741,40 @@ 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); + try { + if (test_hook_) test_hook_(OrchestratorTestPoint::SUBMIT_RUN_SLOT_REGISTERING); + register_run_slot(run, slot); + } catch (...) { + // The run never acquired this slot, so cancellation cannot discover + // it. Return the combined slot/HeapRing reservation directly. + s.state.store(TaskState::CONSUMED, std::memory_order_release); + allocator_->release(slot); + throw; + } + // 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); + if (test_hook_) test_hook_(OrchestratorTestPoint::SCOPE_REGISTERED); + } + { + std::lock_guard lk(s.fanout_mu); + s.fanout_total = scope_ref; + } s.worker_type = worker_type; s.callable = callable; @@ -500,14 +790,28 @@ SubmitResult Orchestrator::submit_impl( // (THREAD mode) or write_blob → read_blob (PROCESS mode). The L2 ABI // ChipStorageTaskArgs conversion now runs inside ChipWorker::run // rather than at submit time. + std::vector group_member_states; + std::vector group_member_outcomes; + if (args_list.size() > 1) { + // Scheduler dispatch must not allocate after it publishes RUNNING: a + // failure there would leave neither dispatch nor cancellation owning + // the slot. Prepare both vectors while BUILDING and commit them only + // after both allocations succeed. + group_member_states.assign(args_list.size(), GroupMemberState::NOT_DISPATCHED); + if (test_hook_) test_hook_(OrchestratorTestPoint::GROUP_MEMBER_STATES_PREPARED); + group_member_outcomes.assign(args_list.size(), EndpointOutcome::SKIPPED); + } if (args_list.size() == 1) { s.is_group_ = false; s.task_args = std::move(args_list.front()); if (!remote_sidecars.empty()) s.remote_sidecar = std::move(remote_sidecars.front()); } else { + std::lock_guard lk(s.group_mu); s.is_group_ = true; s.task_args_list = std::move(args_list); s.remote_sidecars = std::move(remote_sidecars); + s.group_member_states.swap(group_member_states); + s.group_member_outcomes.swap(group_member_outcomes); } s.target_worker_ids = std::move(target_worker_ids); @@ -532,7 +836,14 @@ SubmitResult Orchestrator::submit_impl( } ps.fanout_consumers.push_back(slot); ps.fanout_total++; - s.fanin_producers.push_back(prod); + try { + if (test_hook_) test_hook_(OrchestratorTestPoint::PRODUCER_FORWARD_EDGE_PUBLISHED); + s.fanin_producers.push_back(prod); + } catch (...) { + ps.fanout_total--; + ps.fanout_consumers.pop_back(); + throw; + } if (ps_state == TaskState::FAILED) { poisoned_by_failed_producer = true; if (poison_message.empty()) poison_message = ps.failure_message; @@ -541,40 +852,80 @@ 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); - std::vector fanin_producers = s.fanin_producers; - try_consume(slot); - for (TaskSlot prod : fanin_producers) { - try_consume(prod); + // 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; + { + std::lock_guard lk(s.fanout_mu); + 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. Submission is + // still open here, so cancellation cannot take the BUILDING debt over + // concurrently. + if (commit_failure_propagation(s)) { + try_consume(slot); + for (TaskSlot prod : fanin_producers) { + try_consume(prod); + } } 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 +933,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( @@ -875,14 +1237,16 @@ void Orchestrator::infer_deps( case TensorArgType::INOUT: { TaskSlot prod = tensormap_->lookup(run_id, key); if (prod != INVALID_SLOT) add_unique_producer(prod); - tensormap_->insert(run_id, key, slot); output_keys.push_back(key); + if (test_hook_) test_hook_(OrchestratorTestPoint::SUBMIT_OUTPUT_KEY_PREPARED); + tensormap_->insert(run_id, key, slot); break; } case TensorArgType::OUTPUT: case TensorArgType::OUTPUT_EXISTING: { - tensormap_->insert(run_id, key, slot); output_keys.push_back(key); + if (test_hook_) test_hook_(OrchestratorTestPoint::SUBMIT_OUTPUT_KEY_PREPARED); + tensormap_->insert(run_id, key, slot); break; } case TensorArgType::NO_DEP: diff --git a/src/common/hierarchical/orchestrator.h b/src/common/hierarchical/orchestrator.h index 5461b6491..113ce82ce 100644 --- a/src/common/hierarchical/orchestrator.h +++ b/src/common/hierarchical/orchestrator.h @@ -33,16 +33,19 @@ #include #include #include +#include #include #include #include #include +#include #include #include "../task_interface/call_config.h" #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" @@ -63,6 +66,21 @@ struct SubmitResult { TaskSlot task_slot{INVALID_SLOT}; }; +// Deterministic seams for exception-safety unit tests. Production workers +// leave the callback unset. +enum class OrchestratorTestPoint : int32_t { + SCOPE_REGISTERED = 0, + PRODUCER_FORWARD_EDGE_PUBLISHED = 1, + FAILURE_FANIN_SNAPSHOT = 2, + GROUP_MEMBER_STATES_PREPARED = 3, + ALLOC_RUN_SLOT_REGISTERING = 4, + ALLOC_OUTPUT_KEY_PREPARED = 5, + SUBMIT_RUN_SLOT_REGISTERING = 6, + SUBMIT_OUTPUT_KEY_PREPARED = 7, + BEGIN_RUN_MAP_PUBLISHED = 8, + BEGIN_RUN_FIFO_PUBLISHED = 9, +}; + // --------------------------------------------------------------------------- // Orchestrator // --------------------------------------------------------------------------- @@ -119,6 +137,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 +146,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 @@ -167,19 +207,27 @@ class Orchestrator { // Scheduler uses the same path after releasing the final dependency. void enqueue_ready(TaskSlot slot); + void set_test_hook(std::function hook) { test_hook_ = std::move(hook); } + private: TensorMap *tensormap_ = nullptr; Ring *allocator_ = nullptr; Scope *scope_ = nullptr; WorkerManager *manager_ = nullptr; std::function ready_notify_cb_; + std::function test_hook_; ReadyQueue *ready_sub_queue_ = nullptr; 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 +239,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..2f1c3f2d8 100644 --- a/src/common/hierarchical/scheduler.cpp +++ b/src/common/hierarchical/scheduler.cpp @@ -11,6 +11,8 @@ #include "scheduler.h" +#include +#include #include #include @@ -29,10 +31,60 @@ bool is_terminal_group_state(GroupMemberState state) { state == GroupMemberState::SKIPPED; } -void reset_group_state(TaskSlotState &state, int32_t group_size, GroupMemberState initial_member_state) { +struct PreparedGroupVectors { + std::vector member_states; + std::vector member_outcomes; + int32_t terminal_count{0}; + + bool replaces_existing() const noexcept { return !member_states.empty(); } +}; + +PreparedGroupVectors +prepare_group_vectors_locked(const TaskSlotState &state, int32_t group_size, GroupMemberState initial_member_state) { + const size_t size = static_cast(group_size); + if (state.group_member_states.size() == size && state.group_member_outcomes.size() == size) return {}; + + // Prepare both allocations before changing either shared vector. Dispatch + // does this while the slot is still READY, so even a second-allocation + // failure leaves cancellation free to claim and reclaim the slot. A + // completion-side invariant repair also preserves live/terminal members; + // resetting them could consume a failed group while a peer still runs. + PreparedGroupVectors prepared; + prepared.member_states.assign(size, initial_member_state); + prepared.member_outcomes.assign(size, EndpointOutcome::SKIPPED); + std::copy_n( + state.group_member_states.begin(), std::min(size, state.group_member_states.size()), + prepared.member_states.begin() + ); + std::copy_n( + state.group_member_outcomes.begin(), std::min(size, state.group_member_outcomes.size()), + prepared.member_outcomes.begin() + ); + prepared.terminal_count = static_cast( + std::count_if(prepared.member_states.begin(), prepared.member_states.end(), is_terminal_group_state) + ); + return prepared; +} + +PreparedGroupVectors +prepare_group_vectors(TaskSlotState &state, int32_t group_size, GroupMemberState initial_member_state) { std::lock_guard lk(state.group_mu); - state.group_member_states.assign(static_cast(group_size), initial_member_state); - state.group_member_outcomes.assign(static_cast(group_size), EndpointOutcome::SKIPPED); + return prepare_group_vectors_locked(state, group_size, initial_member_state); +} + +void commit_group_vectors_locked(TaskSlotState &state, PreparedGroupVectors &prepared) noexcept { + if (!prepared.replaces_existing()) return; + state.group_member_states.swap(prepared.member_states); + state.group_member_outcomes.swap(prepared.member_outcomes); + state.group_terminal_count.store(prepared.terminal_count, std::memory_order_relaxed); +} + +void reset_group_state_locked(TaskSlotState &state, GroupMemberState initial_member_state) noexcept { + // Exact-sized storage is prepared while BUILDING and defensively repaired + // before the dispatch claim. Everything after READY -> RUNNING is now a + // no-allocation commit. + std::fill(state.group_member_states.begin(), state.group_member_states.end(), initial_member_state); + std::fill(state.group_member_outcomes.begin(), state.group_member_outcomes.end(), EndpointOutcome::SKIPPED); state.group_terminal_count.store(0, std::memory_order_relaxed); state.group_failed = false; state.group_first_failure_index = -1; @@ -82,10 +134,9 @@ void Scheduler::worker_done(WorkerCompletion completion) { { 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); - } + PreparedGroupVectors prepared = + prepare_group_vectors_locked(s, group_size, GroupMemberState::NOT_DISPATCHED); + commit_group_vectors_locked(s, prepared); bool invalid_group_index = completion.group_index < 0 || completion.group_index >= group_size; if (invalid_group_index) { terminal.outcome = EndpointOutcome::ENDPOINT_FAILURE; @@ -175,8 +226,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,18 +287,26 @@ 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 + std::string failure_message; + if (failed) failure_message = completion.error_message; + + // 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::vector producers; { std::lock_guard lk(s.fanout_mu); consumers = s.fanout_consumers; + producers = s.fanin_producers; + if (failed) s.failure_message.swap(failure_message); + s.state.store(failed ? TaskState::FAILED : TaskState::COMPLETED, std::memory_order_release); } for (TaskSlot consumer : consumers) { if (failed) { @@ -248,24 +314,18 @@ 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(); } } try_consume(slot); // Deferred release: release one fanout ref on each producer this task consumed. - std::vector producers; - { - std::lock_guard lk(s.fanout_mu); - producers = s.fanin_producers; - } for (TaskSlot prod : producers) { try_consume(prod); } @@ -273,50 +333,34 @@ 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. + // Cancellation is the only contender allowed to take over an existing + // debt, because it runs after graph submission has closed and therefore + // knows a BUILDING slot no longer has a submitting owner. + 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; + std::vector producers; { std::lock_guard lk(s.fanout_mu); consumers = s.fanout_consumers; + producers = s.fanin_producers; } for (TaskSlot consumer : consumers) { poison_task(consumer, root_message); } + // All fallible preparation, including recursive poison, finishes before + // any non-repeatable release. This claimant owns the propagation + // exclusively; cancellation does not take over scheduler-owned debt. try_consume(slot); - - std::vector producers; - { - std::lock_guard lk(s.fanout_mu); - producers = s.fanin_producers; - } for (TaskSlot prod : producers) { try_consume(prod); } @@ -349,16 +393,36 @@ void Scheduler::try_consume(TaskSlot slot) { // sched_thread_ with no surrounding handler, any throw is fatal to the whole // worker tree (std::terminate), not a per-task failure. void Scheduler::dispatch_ready() { - const std::unordered_set reserved_worker_ids = dispatch_next_level_group(); - dispatch_next_level_singles(reserved_worker_ids); - dispatch_sub_ready(); + std::optional run_snapshot; + if (cfg_.active_run_cb) { + RunId active_run = cfg_.active_run_cb(); + if (active_run == INVALID_RUN_ID) return; + run_snapshot = active_run; + } + + // Group reservations and every queue pop in one pass belong to the same + // whole-run FIFO head, even if a completion advances the head mid-pass. + const std::unordered_set reserved_worker_ids = dispatch_next_level_group(run_snapshot); + dispatch_next_level_singles(reserved_worker_ids, run_snapshot); + dispatch_sub_ready(run_snapshot); } -void Scheduler::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(const std::optional &run_snapshot) { TaskSlot slot; - while (cfg_.ready_sub_queue->try_pop(slot)) { + while (run_snapshot ? cfg_.ready_sub_queue->try_pop(*run_snapshot, 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 (run_snapshot && s.run_id != *run_snapshot) { + cfg_.enqueue_ready_cb(slot); + return; + } if (s.worker_type != WorkerType::SUB) { throw std::runtime_error("Scheduler::dispatch_sub_ready: misrouted task slot"); } @@ -369,15 +433,35 @@ 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 (s.is_group()) { - reset_group_state(s, group_size, GroupMemberState::NOT_DISPATCHED); + // The expected path is allocation-free because submit prepared + // exact-sized storage while BUILDING. Repair any violated invariant + // into local storage while the slot is still cancellation-claimable. + // Do not publish that storage unless this scheduler wins the claim: + // cancellation may have marked the shared vectors terminal meanwhile. + PreparedGroupVectors prepared = prepare_group_vectors(s, group_size, GroupMemberState::NOT_DISPATCHED); + if (cfg_.before_claim_cb) cfg_.before_claim_cb(slot); + { + // Acquire the only fallible synchronization primitive before + // publishing RUNNING; the commit after the claim only fills + // existing storage and resets scalar metadata. + std::lock_guard lk(s.group_mu); + if (!claim_for_dispatch(s)) continue; + commit_group_vectors_locked(s, prepared); + reset_group_state_locked(s, GroupMemberState::NOT_DISPATCHED); + } + } else { + if (cfg_.before_claim_cb) cfg_.before_claim_cb(slot); + if (!claim_for_dispatch(s)) continue; } for (int32_t i = 0; i < group_size; ++i) { if (s.is_group()) { @@ -391,15 +475,26 @@ void Scheduler::dispatch_sub_ready() { } } -std::unordered_set Scheduler::dispatch_next_level_group() { +std::unordered_set Scheduler::dispatch_next_level_group(const std::optional &run_snapshot) { TaskSlot slot; - while (cfg_.ready_next_level_queues->try_front_group(slot)) { + while (run_snapshot ? cfg_.ready_next_level_queues->try_front_group(*run_snapshot, 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 (run_snapshot) { + cfg_.ready_next_level_queues->try_pop_group(*run_snapshot, stale); + } else { + cfg_.ready_next_level_queues->try_pop_group(stale); + } continue; } + if (run_snapshot && s.run_id != *run_snapshot) { + TaskSlot misplaced; + cfg_.ready_next_level_queues->try_pop_group(*run_snapshot, 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 +519,30 @@ 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 = run_snapshot ? cfg_.ready_next_level_queues->try_pop_group(*run_snapshot, 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 {}; } - reset_group_state(s, group_size, GroupMemberState::RUNNING); - s.state.store(TaskState::RUNNING, std::memory_order_release); + PreparedGroupVectors prepared = prepare_group_vectors(s, group_size, GroupMemberState::RUNNING); + if (cfg_.before_claim_cb) cfg_.before_claim_cb(slot); + { + std::lock_guard lk(s.group_mu); + if (!claim_for_dispatch(s)) continue; + commit_group_vectors_locked(s, prepared); + reset_group_state_locked(s, GroupMemberState::RUNNING); + } for (int32_t i = 0; i < group_size; ++i) { workers[static_cast(i)]->dispatch(WorkerDispatch{slot, i}); } @@ -438,7 +550,9 @@ std::unordered_set Scheduler::dispatch_next_level_group() { return {}; } -void Scheduler::dispatch_next_level_singles(const std::unordered_set &reserved_worker_ids) { +void Scheduler::dispatch_next_level_singles( + const std::unordered_set &reserved_worker_ids, const std::optional &run_snapshot +) { 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 +565,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 (run_snapshot ? cfg_.ready_next_level_queues->try_pop_single(worker_id, *run_snapshot, 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 (run_snapshot && s.run_id != *run_snapshot) { + 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..4d880e9fc 100644 --- a/src/common/hierarchical/scheduler.h +++ b/src/common/hierarchical/scheduler.h @@ -37,6 +37,7 @@ #include #include #include +#include #include #include #include @@ -51,6 +52,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 +77,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); @@ -104,7 +131,9 @@ class Scheduler { void poison_task(TaskSlot slot, const std::string &root_message); void try_consume(TaskSlot slot); void dispatch_ready(); - std::unordered_set dispatch_next_level_group(); - void dispatch_next_level_singles(const std::unordered_set &reserved_worker_ids); - void dispatch_sub_ready(); + std::unordered_set dispatch_next_level_group(const std::optional &run_snapshot); + void dispatch_next_level_singles( + const std::unordered_set &reserved_worker_ids, const std::optional &run_snapshot + ); + void dispatch_sub_ready(const std::optional &run_snapshot); }; diff --git a/src/common/hierarchical/types.cpp b/src/common/hierarchical/types.cpp index 3fcebcd85..a3afc4677 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,170 @@ 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) { + // Copy before the state transition. A failed string allocation must + // leave the slot claimable rather than publish FAILED without a + // durable reason or propagation owner. + std::string durable_message(message); + if (s.state.compare_exchange_strong( + current, TaskState::FAILED, std::memory_order_acq_rel, std::memory_order_acquire + )) { + s.failure_message.swap(durable_message); + // The debt remains set while propagation takes every fallible + // snapshot for a BUILDING slot. PENDING and READY claims remain + // exclusively owned by their claimant, so cancellation cannot + // prepare them concurrently. + if (current == TaskState::BUILDING) { + s.failure_propagation_pending.store(true, std::memory_order_release); + } + return current; + } + } + return std::nullopt; +} + +bool commit_failure_propagation(TaskSlotState &s) noexcept { + return s.failure_propagation_pending.exchange(false, std::memory_order_acq_rel); +} + +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(); + const size_t size = static_cast(group_size); + const bool repair_bookkeeping = s.group_member_states.size() != size || s.group_member_outcomes.size() != size; + + std::vector member_states; + std::vector member_outcomes; + if (repair_bookkeeping) { + member_states.assign(size, GroupMemberState::NOT_DISPATCHED); + member_outcomes.assign(size, EndpointOutcome::SKIPPED); + std::copy_n(s.group_member_states.begin(), std::min(size, s.group_member_states.size()), member_states.begin()); + std::copy_n( + s.group_member_outcomes.begin(), std::min(size, s.group_member_outcomes.size()), member_outcomes.begin() + ); + } + + std::string first_failure_message; + if (!s.group_failed) { + first_failure_message = message; + } + + // Every allocation above completes before either bookkeeping vector is + // replaced, so an exception cannot expose mismatched sizes to a retry. + if (repair_bookkeeping) { + s.group_member_states.swap(member_states); + s.group_member_outcomes.swap(member_outcomes); + } + if (!s.group_failed) { + s.group_failed = true; + s.group_first_failure_index = -1; + s.group_first_failure_message.swap(first_failure_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 +252,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 +280,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..b371c33c1 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,27 @@ 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 BUILDING failure claim leaves propagation to submit, and by + // cancellation when it claims a PENDING or READY slot itself. Fallible + // preparation happens while it remains set, so cancellation can retry its + // own FAILED slot without allowing another active owner to take it over. + 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 +382,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 +452,55 @@ 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. Claims from PENDING and READY are exclusively owned by the +// claimant and may propagate immediately; they do not advertise takeover debt. +std::optional claim_task_failure(TaskSlotState &s, const std::string &message); + +// Settle a fully prepared propagation debt and claim its reference releases. +// Callers finish every fallible local snapshot first. Only BUILDING handoff and +// cancellation-owned claims use debt, so no active propagation contender may +// be preparing the same slot concurrently. +bool commit_failure_propagation(TaskSlotState &s) noexcept; + +// 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 +508,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 +535,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..5c22e0dd6 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/worker_async_fifo/test_worker_async_fifo.py @@ -0,0 +1,418 @@ +#!/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 atexit +import tempfile +import threading +import time +import uuid +from contextlib import suppress +from pathlib import Path + +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 + + +class _FileSignal: + """Pickle-safe parent/child signal for the standalone scene runner.""" + + def __init__(self, label: str): + token = uuid.uuid4().hex + self._path = Path(tempfile.gettempdir()) / f"simpler-worker-async-fifo-{label}-{token}" + + def clear(self) -> None: + with suppress(FileNotFoundError): + self._path.unlink() + + def set(self) -> None: + self._path.touch(exist_ok=True) + + def wait(self, timeout: float) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if self._path.exists(): + return True + time.sleep(0.001) + return self._path.exists() + + +_SUB_ENTERED = _FileSignal("entered") +_SUB_RELEASE = _FileSignal("release") + + +def _clear_signals() -> None: + _SUB_ENTERED.clear() + _SUB_RELEASE.clear() + + +atexit.register(_clear_signals) + + +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 _run_and_validate_l3( # noqa: PLR0913 -- mirror the scene-test runner hook + self, + worker, + compiled_callables, + sub_handles, + case, + rounds=1, + skip_golden=False, + enable_l2_swimlane=0, + enable_dump_args=False, + enable_pmu=0, + enable_dep_gen=False, + enable_scope_stats=False, + output_prefix="", + ): + """Run the custom concurrency checks from the standalone scene runner. + + Unlike a regular scene case, this test has no single orchestration + callback or generated argument set. Pytest collects the three methods + below directly; the standalone runner reaches them through this hook. + """ + del ( + rounds, + skip_golden, + enable_l2_swimlane, + enable_dump_args, + enable_pmu, + enable_dep_gen, + enable_scope_stats, + output_prefix, + ) + # The pytest fixture publishes these maps on the class. Standalone + # execution passes the same registered handles into this hook instead. + type(self)._st_chip_handles = compiled_callables + type(self)._st_sub_handles = sub_handles + platform = str(worker._config["platform"]) # noqa: SLF001 -- scene-test white-box validation + assert platform in case["platforms"] + self.test_prepared_run_device_control_waits_for_the_active_run(platform, worker) + self.test_a_run_whose_cleanup_touches_the_device_degrades_to_depth_one(platform, worker) + self.test_run(platform, worker) + + 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 = [] + tensors = [] + submitter = None + first = None + try: + 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) + tensors.clear() + tensor = None + first_a = first_b = first_out = None + 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() + tensor = None + 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..fc1a3f351 100644 --- a/tests/ut/cpp/hierarchical/test_orchestrator.cpp +++ b/tests/ut/cpp/hierarchical/test_orchestrator.cpp @@ -12,6 +12,10 @@ #include #include +#include +#include +#include +#include #include "call_config.h" #include "ring.h" @@ -42,6 +46,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 +84,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 +517,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 +694,652 @@ 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); +} + +TEST_F(OrchestratorFixture, CancellationRetriesBeforeAnyPreparedFailureReleasesCommit) { + auto producer = orch.submit_next_level(C(77), single_tensor_args(0x7700, TensorArgType::OUTPUT), cfg, 0); + TaskSlot ready; + ASSERT_TRUE(rq.try_pop(run_id, ready)); + auto consumer = orch.submit_next_level(C(78), single_tensor_args(0x7700, TensorArgType::INPUT), cfg, 0); + ASSERT_EQ(S(consumer.task_slot).state.load(std::memory_order_acquire), TaskState::PENDING); + + int32_t snapshots = 0; + orch.set_test_hook([&](OrchestratorTestPoint point) { + if (point == OrchestratorTestPoint::FAILURE_FANIN_SNAPSHOT && ++snapshots == 2) throw std::bad_alloc(); + }); + + EXPECT_THROW( + orch.fail_run_submission(run_id, std::make_exception_ptr(std::runtime_error("graph construction failed"))), + std::bad_alloc + ); + + EXPECT_TRUE(S(producer.task_slot).failure_propagation_pending.load(std::memory_order_acquire)); + EXPECT_TRUE(S(consumer.task_slot).failure_propagation_pending.load(std::memory_order_acquire)); + EXPECT_EQ(S(producer.task_slot).fanout_released.load(std::memory_order_acquire), 0); + EXPECT_EQ(S(consumer.task_slot).fanout_released.load(std::memory_order_acquire), 0) + << "one prepared slot released before every cancellation snapshot succeeded"; + + orch.set_test_hook({}); + EXPECT_NO_THROW(orch.fail_run_submission(run_id)); + EXPECT_TRUE(orch.run_done(run_id)); + EXPECT_EQ(S(producer.task_slot).state.load(std::memory_order_acquire), TaskState::CONSUMED); + EXPECT_EQ(S(consumer.task_slot).state.load(std::memory_order_acquire), TaskState::CONSUMED); + orch.release_run(run_id); +} + +TEST_F(OrchestratorFixture, FailedReverseEdgePublicationRollsBackTheProducerReference) { + auto producer = orch.submit_next_level(C(79), single_tensor_args(0x7900, TensorArgType::OUTPUT), cfg, 0); + TaskSlot ready; + ASSERT_TRUE(rq.try_pop(run_id, ready)); + + bool injected = false; + orch.set_test_hook([&](OrchestratorTestPoint point) { + if (point == OrchestratorTestPoint::PRODUCER_FORWARD_EDGE_PUBLISHED && !injected) { + injected = true; + throw std::bad_alloc(); + } + }); + EXPECT_THROW( + (void)orch.submit_next_level(C(80), single_tensor_args(0x7900, TensorArgType::INPUT), cfg, 0), std::bad_alloc + ); + ASSERT_TRUE(injected); + + { + std::lock_guard lk(S(producer.task_slot).fanout_mu); + EXPECT_TRUE(S(producer.task_slot).fanout_consumers.empty()); + EXPECT_EQ(S(producer.task_slot).fanout_total, 0); + } + + orch.set_test_hook({}); + EXPECT_NO_THROW( + orch.fail_run_submission(run_id, std::make_exception_ptr(std::runtime_error("edge publication failed"))) + ); + EXPECT_TRUE(orch.run_done(run_id)); + EXPECT_EQ(S(producer.task_slot).state.load(std::memory_order_acquire), TaskState::CONSUMED); + 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(); + } +} + +TEST(DepthOneAdmission, FailedBeginRollsBackEveryPublishedOwnerAndKeepsRunIdsMonotonic) { + for (OrchestratorTestPoint failure_point : + {OrchestratorTestPoint::BEGIN_RUN_MAP_PUBLISHED, OrchestratorTestPoint::BEGIN_RUN_FIFO_PUBLISHED}) { + 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); + + bool injected = false; + orch.set_test_hook([&](OrchestratorTestPoint point) { + if (point == failure_point && !std::exchange(injected, true)) { + throw std::runtime_error("injected begin_run publication failure"); + } + }); + EXPECT_THROW(orch.begin_run(), std::runtime_error); + ASSERT_TRUE(injected); + + orch.set_test_hook({}); + RunId recovered = orch.begin_run(); + EXPECT_EQ(recovered, 2U); + EXPECT_EQ(orch.active_run_id(), recovered); + orch.close_run_submission(recovered); + EXPECT_TRUE(orch.run_done(recovered)); + orch.release_run(recovered); + allocator.shutdown(); + } +} + +// A submit that throws after scope registration but before the matching charge +// leaves only an extra release, never a reference nothing can release. Python +// closes the scope before cancelling, which is the order used here. +TEST(SubmitFailure, ASlotWhoseSubmitThrewIsFullyReclaimedByCancellation) { + 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); + + RunId run = orch.begin_run(); + orch.scope_begin(); + bool injected = false; + orch.set_test_hook([&](OrchestratorTestPoint point) { + if (point == OrchestratorTestPoint::SCOPE_REGISTERED && !injected) { + injected = true; + throw std::bad_alloc(); + } + }); + + 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::bad_alloc); + ASSERT_TRUE(injected); + + orch.set_test_hook({}); + 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, AllocRegistrationFailureReleasesTheUnownedRingSlot) { + bool injected = false; + orch.set_test_hook([&](OrchestratorTestPoint point) { + if (point == OrchestratorTestPoint::ALLOC_RUN_SLOT_REGISTERING && !injected) { + injected = true; + throw std::bad_alloc(); + } + }); + + EXPECT_THROW((void)orch.alloc(std::vector{16}, DataType::UINT8), std::bad_alloc); + ASSERT_TRUE(injected); + ASSERT_NE(allocator.slot_state(0), nullptr); + EXPECT_EQ(allocator.slot_state(0)->state.load(std::memory_order_acquire), TaskState::CONSUMED); + EXPECT_EQ(allocator.active_count(), 0) << "a slot the run never owned was left live in the Ring"; + + orch.set_test_hook({}); + EXPECT_NO_THROW( + orch.fail_run_submission(run_id, std::make_exception_ptr(std::runtime_error("alloc registration failed"))) + ); + EXPECT_TRUE(orch.run_done(run_id)); + EXPECT_THROW(orch.wait_run(run_id), std::runtime_error); + orch.release_run(run_id); +} + +TEST_F(OrchestratorFixture, AllocFailureAfterScopeRegistrationRemainsCancellationClaimable) { + orch.scope_begin(); + bool injected = false; + orch.set_test_hook([&](OrchestratorTestPoint point) { + if (point == OrchestratorTestPoint::SCOPE_REGISTERED && !injected) { + injected = true; + throw std::bad_alloc(); + } + }); + + EXPECT_THROW((void)orch.alloc(std::vector{16}, DataType::UINT8), std::bad_alloc); + ASSERT_TRUE(injected); + ASSERT_NE(allocator.slot_state(0), nullptr); + EXPECT_EQ(allocator.slot_state(0)->state.load(std::memory_order_acquire), TaskState::BUILDING); + EXPECT_EQ(allocator.slot_state(0)->fanout_total, 0) + << "alloc charged a scope reference before scope registration completed"; + + orch.set_test_hook({}); + orch.scope_end(); + EXPECT_NO_THROW( + orch.fail_run_submission(run_id, std::make_exception_ptr(std::runtime_error("alloc scope failed"))) + ); + EXPECT_EQ(allocator.slot_state(0)->state.load(std::memory_order_acquire), TaskState::CONSUMED); + EXPECT_EQ(allocator.active_count(), 0); + EXPECT_TRUE(orch.run_done(run_id)); + EXPECT_THROW(orch.wait_run(run_id), std::runtime_error); + orch.release_run(run_id); +} + +TEST_F(OrchestratorFixture, AllocOutputPublicationFailureLeavesAReclaimableJournal) { + orch.scope_begin(); + bool injected = false; + orch.set_test_hook([&](OrchestratorTestPoint point) { + if (point == OrchestratorTestPoint::ALLOC_OUTPUT_KEY_PREPARED && !injected) { + injected = true; + throw std::bad_alloc(); + } + }); + + EXPECT_THROW((void)orch.alloc(std::vector{16}, DataType::UINT8), std::bad_alloc); + ASSERT_TRUE(injected); + ASSERT_NE(allocator.slot_state(0), nullptr); + EXPECT_EQ(allocator.slot_state(0)->state.load(std::memory_order_acquire), TaskState::BUILDING); + EXPECT_EQ(allocator.slot_state(0)->output_keys.size(), 1u); + EXPECT_EQ(tm.size(), 0) << "alloc published an output mapping before its cleanup journal"; + + orch.set_test_hook({}); + orch.scope_end(); + EXPECT_NO_THROW( + orch.fail_run_submission(run_id, std::make_exception_ptr(std::runtime_error("alloc output failed"))) + ); + EXPECT_EQ(allocator.slot_state(0)->state.load(std::memory_order_acquire), TaskState::CONSUMED); + EXPECT_EQ(allocator.active_count(), 0); + EXPECT_EQ(tm.size(), 0); + EXPECT_TRUE(orch.run_done(run_id)); + EXPECT_THROW(orch.wait_run(run_id), std::runtime_error); + orch.release_run(run_id); +} + +TEST_F(OrchestratorFixture, SubmitRegistrationFailureReleasesTheUnownedHeapRingSlot) { + bool injected = false; + orch.set_test_hook([&](OrchestratorTestPoint point) { + if (point == OrchestratorTestPoint::SUBMIT_RUN_SLOT_REGISTERING && !injected) { + injected = true; + throw std::bad_alloc(); + } + }); + + TaskArgs args; + Tensor output{}; + output.buffer.addr = 0; + output.ndims = 1; + output.shapes[0] = 16; + output.dtype = DataType::UINT8; + args.add_tensor(output, TensorArgType::OUTPUT); + EXPECT_THROW((void)orch.submit_next_level(C(90), args, cfg, 0), std::bad_alloc); + ASSERT_TRUE(injected); + ASSERT_NE(allocator.slot_state(0), nullptr); + EXPECT_EQ(allocator.slot_state(0)->state.load(std::memory_order_acquire), TaskState::CONSUMED); + EXPECT_EQ(allocator.active_count(), 0) << "a task slot absent from its run kept the HeapRing slab live"; + EXPECT_EQ(tm.size(), 0); + + orch.set_test_hook({}); + EXPECT_NO_THROW( + orch.fail_run_submission(run_id, std::make_exception_ptr(std::runtime_error("submit registration failed"))) + ); + EXPECT_TRUE(orch.run_done(run_id)); + EXPECT_THROW(orch.wait_run(run_id), std::runtime_error); + orch.release_run(run_id); +} + +TEST_F(OrchestratorFixture, SubmitOutputJournalFailurePreservesThePreviousOwnerAndReclaims) { + constexpr uint64_t first_key_addr = 0xE300; + constexpr uint64_t previous_key_addr = 0xE301; + TensorKey first_key = TensorKey::local_host(first_key_addr); + TensorKey previous_key = TensorKey::local_host(previous_key_addr); + + auto previous = + orch.submit_next_level(C(91), single_tensor_args(previous_key_addr, TensorArgType::OUTPUT_EXISTING), cfg, 0); + TaskSlot ready = INVALID_SLOT; + ASSERT_TRUE(rq.try_pop(run_id, ready)); + ASSERT_EQ(ready, previous.task_slot); + TaskState expected = TaskState::READY; + ASSERT_TRUE(S(previous.task_slot) + .state.compare_exchange_strong( + expected, TaskState::RUNNING, std::memory_order_acq_rel, std::memory_order_acquire + )); + + int32_t prepared_keys = 0; + orch.set_test_hook([&](OrchestratorTestPoint point) { + if (point == OrchestratorTestPoint::SUBMIT_OUTPUT_KEY_PREPARED && ++prepared_keys == 2) { + throw std::bad_alloc(); + } + }); + + TaskArgs replacement; + Tensor first{}; + first.buffer.addr = first_key_addr; + first.ndims = 1; + first.shapes[0] = 1; + first.dtype = DataType::UINT8; + replacement.add_tensor(first, TensorArgType::OUTPUT_EXISTING); + Tensor second = first; + second.buffer.addr = previous_key_addr; + replacement.add_tensor(second, TensorArgType::OUTPUT_EXISTING); + + EXPECT_THROW((void)orch.submit_next_level(C(92), replacement, cfg, 0), std::bad_alloc); + ASSERT_EQ(prepared_keys, 2); + TaskSlot failed_slot = allocator.next_task_id() - 1; + ASSERT_NE(failed_slot, previous.task_slot); + EXPECT_EQ(S(failed_slot).state.load(std::memory_order_acquire), TaskState::BUILDING); + ASSERT_EQ(S(failed_slot).output_keys.size(), 2u); + EXPECT_EQ(tm.lookup(run_id, first_key), failed_slot) + << "the first fully-published output is needed to exercise cancellation cleanup"; + EXPECT_EQ(tm.lookup(run_id, previous_key), previous.task_slot) + << "a failed second insert displaced the prior owner before publication"; + + orch.set_test_hook({}); + EXPECT_NO_THROW( + orch.fail_run_submission(run_id, std::make_exception_ptr(std::runtime_error("output publication failed"))) + ); + EXPECT_EQ(S(failed_slot).state.load(std::memory_order_acquire), TaskState::CONSUMED); + EXPECT_EQ(tm.lookup(run_id, first_key), INVALID_SLOT); + EXPECT_EQ(tm.lookup(run_id, previous_key), previous.task_slot) + << "failed-slot cleanup erased a TensorMap entry still owned by the running producer"; + EXPECT_FALSE(orch.run_done(run_id)); + + S(previous.task_slot).state.store(TaskState::COMPLETED, std::memory_order_release); + EXPECT_TRUE(orch.on_consumed(previous.task_slot)); + EXPECT_TRUE(orch.run_done(run_id)); + EXPECT_EQ(tm.lookup(run_id, previous_key), INVALID_SLOT); + EXPECT_EQ(allocator.active_count(), 0); + EXPECT_THROW(orch.wait_run(run_id), std::runtime_error); + orch.release_run(run_id); +} + +// Group bookkeeping is prepared by the real submit path while its slot is +// still BUILDING. Fault exactly between the two vector allocations: publishing +// RUNNING first would make this slot neither dispatchable nor cancellable, but +// BUILDING leaves it owned by the submission-failure cancellation path. +TEST_F(OrchestratorFixture, AGroupBookkeepingAllocationFailureLeavesAReclaimableSlot) { + bool injected = false; + orch.set_test_hook([&](OrchestratorTestPoint point) { + if (point == OrchestratorTestPoint::GROUP_MEMBER_STATES_PREPARED && !injected) { + injected = true; + throw std::bad_alloc(); + } + }); + + TaskArgs first = single_tensor_args(0xE200, TensorArgType::OUTPUT); + TaskArgs second = single_tensor_args(0xE201, TensorArgType::OUTPUT); + EXPECT_THROW((void)orch.submit_sub_group(C(89), {first, second}), std::bad_alloc); + ASSERT_TRUE(injected); + ASSERT_NE(allocator.slot_state(0), nullptr); + EXPECT_EQ(allocator.slot_state(0)->state.load(std::memory_order_acquire), TaskState::BUILDING) + << "group preparation published a non-cancellable state before both allocations succeeded"; + + orch.set_test_hook({}); + EXPECT_NO_THROW( + orch.fail_run_submission(run_id, std::make_exception_ptr(std::runtime_error("group preparation failed"))) + ); + EXPECT_EQ(allocator.slot_state(0)->state.load(std::memory_order_acquire), TaskState::CONSUMED); + EXPECT_TRUE(orch.run_done(run_id)); + EXPECT_THROW(orch.wait_run(run_id), std::runtime_error); + orch.release_run(run_id); +} + +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..f0036c5ce 100644 --- a/tests/ut/cpp/hierarchical/test_scheduler.cpp +++ b/tests/ut/cpp/hierarchical/test_scheduler.cpp @@ -272,6 +272,164 @@ 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"); + EXPECT_EQ(s.failure_propagation_pending.load(std::memory_order_acquire), claimable == TaskState::BUILDING) + << "only a mid-wiring claim may advertise propagation takeover debt"; + + // 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"); + + s.failure_propagation_pending.store(false, std::memory_order_release); + } +} + +TEST(MarkGroupMembersSkipped, RepairsBothBookkeepingVectorsAsOneTransaction) { + TaskSlotState s; + s.is_group_ = true; + s.task_args_list.resize(2); + s.group_member_states.assign(2, GroupMemberState::NOT_DISPATCHED); + s.group_member_outcomes.clear(); + + mark_group_members_skipped(s, "cancelled"); + + ASSERT_EQ(s.group_member_states.size(), 2u); + ASSERT_EQ(s.group_member_outcomes.size(), 2u); + EXPECT_EQ(s.group_member_states[0], GroupMemberState::SKIPPED); + EXPECT_EQ(s.group_member_states[1], GroupMemberState::SKIPPED); + EXPECT_EQ(s.group_member_outcomes[0], EndpointOutcome::SKIPPED); + EXPECT_EQ(s.group_member_outcomes[1], EndpointOutcome::SKIPPED); + EXPECT_EQ(s.group_terminal_count.load(std::memory_order_acquire), 2); +} + +// 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 +446,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 +480,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 +495,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 +559,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 +573,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 +712,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); @@ -640,6 +850,8 @@ struct GroupSchedulerFixture : public ::testing::Test { MockMailboxWorker worker_a; MockMailboxWorker worker_b; MockMailboxWorker worker_c; + MockMailboxWorker sub_worker_a; + MockMailboxWorker sub_worker_b; WorkerManager manager; Scheduler sched; CallConfig cfg; @@ -656,9 +868,13 @@ struct GroupSchedulerFixture : public ::testing::Test { worker_a.start(); worker_b.start(); worker_c.start(); + sub_worker_a.start(); + sub_worker_b.start(); manager.add_next_level(worker_a.mailbox_ptr()); manager.add_next_level(worker_b.mailbox_ptr()); manager.add_next_level(worker_c.mailbox_ptr()); + manager.add_sub(sub_worker_a.mailbox_ptr()); + manager.add_sub(sub_worker_b.mailbox_ptr()); manager.start( &allocator, [this](WorkerCompletion completion) { @@ -682,6 +898,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); @@ -735,6 +957,22 @@ TEST_F(GroupSchedulerFixture, GroupDispatchesToNWorkers) { wait_consumed(slot); } +TEST_F(GroupSchedulerFixture, SubGroupUsesTheAllocationFreeGroupCommitPath) { + TaskArgs a0 = single_tensor_args(0xA2, TensorArgType::OUTPUT); + TaskArgs a1 = single_tensor_args(0xA3, TensorArgType::OUTPUT); + auto res = orch.submit_sub_group(C(44), {a0, a1}); + + sub_worker_a.wait_running(); + sub_worker_b.wait_running(); + EXPECT_EQ(sub_worker_a.dispatched_count(), 1); + EXPECT_EQ(sub_worker_b.dispatched_count(), 1); + EXPECT_EQ(S(res.task_slot).state.load(std::memory_order_acquire), TaskState::RUNNING); + + sub_worker_a.complete(); + sub_worker_b.complete(); + wait_consumed(res.task_slot); +} + TEST_F(GroupSchedulerFixture, GroupMapsEachMemberToItsTargetWorkerIdNotIndex) { // Reversed target order: member 0 -> worker id 1 (worker_b), member 1 -> // worker id 0 (worker_a). A map-by-registration-index bug would instead @@ -842,6 +1080,120 @@ TEST_F(GroupSchedulerFixture, BlockedGroupReservesTargetsThatBecomeIdleOneAtATim wait_consumed(unrelated.task_slot); } +TEST(SchedulerDispatchPassTest, ActiveRunSwitchCannotBypassSuccessorGroupReservation) { + constexpr RunId run_a = 41; + constexpr RunId run_b = 42; + + Ring allocator; + ReadyQueue rq_sub; + NextLevelReadyQueues rq_next_level; + MockMailboxWorker worker_a; + MockMailboxWorker worker_b; + WorkerManager manager; + Scheduler sched; + + allocator.init(/*heap_bytes=*/0); + worker_a.start(); + worker_b.start(); + manager.add_next_level(worker_a.mailbox_ptr()); + manager.add_next_level(worker_b.mailbox_ptr()); + manager.start( + &allocator, + [&sched](WorkerCompletion completion) { + sched.worker_done(std::move(completion)); + }, + [](WorkerDispatch) {} + ); + rq_next_level.reset(manager.next_level_worker_ids()); + + auto allocate_slot = [&](RunId run_id, uint8_t callable_seed, int32_t worker_id, bool group) { + AllocResult allocation = allocator.alloc(/*heap_bytes=*/0, /*scope_depth=*/0); + TaskSlotState &state = *allocator.slot_state(allocation.slot); + state.reset(); + state.run_id = run_id; + state.worker_type = WorkerType::NEXT_LEVEL; + state.callable = C(callable_seed); + state.target_worker_ids.push_back(worker_id); + if (group) { + state.is_group_ = true; + state.task_args_list.push_back(single_tensor_args(callable_seed, TensorArgType::OUTPUT)); + rq_next_level.push_group(run_id, allocation.slot); + } else { + state.task_args = single_tensor_args(callable_seed, TensorArgType::OUTPUT); + rq_next_level.push_single(worker_id, run_id, allocation.slot); + } + state.state.store(TaskState::READY, std::memory_order_release); + return allocation.slot; + }; + + // Run A's group occupies worker 1. Run B's group and following single + // both target worker 0, so the group head owns that worker reservation. + allocate_slot(run_a, /*callable_seed=*/70, /*worker_id=*/1, /*group=*/true); + allocate_slot(run_b, /*callable_seed=*/71, /*worker_id=*/0, /*group=*/true); + allocate_slot(run_b, /*callable_seed=*/72, /*worker_id=*/0, /*group=*/false); + + std::atomic active_run{run_a}; + Scheduler::Config config; + config.ring = &allocator; + config.ready_sub_queue = &rq_sub; + config.ready_next_level_queues = &rq_next_level; + config.manager = &manager; + config.enqueue_ready_cb = [&](TaskSlot slot) { + TaskSlotState &state = *allocator.slot_state(slot); + if (state.is_group()) { + rq_next_level.push_group(state.run_id, slot); + } else { + rq_next_level.push_single(state.target_worker_id(0), state.run_id, slot); + } + }; + config.active_run_cb = [&] { + return active_run.load(std::memory_order_acquire); + }; + // The run switch lands after the group phase selected A and before the + // singles phase can select a queue partition. + config.before_claim_cb = [&](TaskSlot slot) { + if (allocator.slot_state(slot)->run_id == run_a) active_run.store(run_b, std::memory_order_release); + }; + config.on_consumed_cb = [&](TaskSlot slot) { + allocator.slot_state(slot)->state.store(TaskState::CONSUMED, std::memory_order_release); + allocator.release(slot); + }; + config.on_task_failed_cb = [](TaskSlot, const std::string &) {}; + sched.start(config); + + worker_a.wait_running(); + worker_b.wait_running(); + EXPECT_EQ(worker_a.dispatched_count(), 1); + EXPECT_EQ(worker_b.dispatched_count(), 1); + if (worker_a.dispatched_count() == 1) { + std::lock_guard lock(worker_a.dispatched_mu); + EXPECT_EQ(worker_a.dispatched[0].callable_hash0, 71u) + << "run B's group must dispatch before its single on the same target"; + } + if (worker_b.dispatched_count() == 1) { + std::lock_guard lock(worker_b.dispatched_mu); + EXPECT_EQ(worker_b.dispatched[0].callable_hash0, 70u); + } + + if (worker_a.is_running.load(std::memory_order_acquire)) worker_a.complete(); + auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(500); + while (worker_a.dispatched_count() < 2 && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_EQ(worker_a.dispatched_count(), 2); + worker_a.wait_running(); + if (worker_a.dispatched_count() == 2) { + std::lock_guard lock(worker_a.dispatched_mu); + EXPECT_EQ(worker_a.dispatched[1].callable_hash0, 72u); + } + + if (worker_a.is_running.load(std::memory_order_acquire)) worker_a.complete(); + if (worker_b.is_running.load(std::memory_order_acquire)) worker_b.complete(); + sched.stop(); + manager.stop(); + allocator.shutdown(); +} + TEST_F(GroupSchedulerFixture, ConsecutiveGroupsReserveOnlyBlockedHeadTargets) { SubmitResult first_group; SubmitResult second_group; @@ -986,6 +1338,42 @@ TEST_F(GroupSchedulerFixture, GroupFailureWaitsForRunningMembersThenConsumes) { EXPECT_EQ(S(slot).state.load(), TaskState::CONSUMED); } +TEST_F(GroupSchedulerFixture, CompletionRepairPreservesRunningPeersWhenOutcomesAreMissing) { + TaskArgs a0 = single_tensor_args(0xC2, TensorArgType::OUTPUT); + TaskArgs a1 = single_tensor_args(0xC3, TensorArgType::OUTPUT); + auto res = orch.submit_next_level_group(C(43), {a0, a1}, cfg, {0, 1}); + TaskSlot slot = res.task_slot; + + worker_a.wait_running(); + worker_b.wait_running(); + { + std::lock_guard lk(S(slot).group_mu); + ASSERT_EQ(S(slot).group_member_states.size(), 2u); + ASSERT_EQ(S(slot).group_member_outcomes.size(), 2u); + S(slot).group_member_outcomes.clear(); + } + + worker_a.complete_with_error("first member failed"); + auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(500); + bool repaired = false; + while (!repaired && std::chrono::steady_clock::now() < deadline) { + { + std::lock_guard lk(S(slot).group_mu); + repaired = S(slot).group_member_states.size() == 2u && S(slot).group_member_outcomes.size() == 2u && + S(slot).group_member_states[0] == GroupMemberState::FAILED && + S(slot).group_member_states[1] == GroupMemberState::RUNNING; + } + if (!repaired) std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_TRUE(repaired) << "worker_done indexed mismatched group bookkeeping instead of repairing both vectors"; + EXPECT_EQ(S(slot).state.load(std::memory_order_acquire), TaskState::RUNNING) + << "repair discarded a live peer and let group failure consume its slot early"; + + worker_b.complete(); + wait_consumed(slot); + EXPECT_EQ(S(slot).state.load(), TaskState::CONSUMED); +} + TEST_F(GroupSchedulerFixture, InvalidGroupIndexFailsAndConsumesGroup) { TaskArgs a0 = single_tensor_args(0xD0, TensorArgType::OUTPUT); TaskArgs a1 = single_tensor_args(0xD1, TensorArgType::OUTPUT); @@ -1086,6 +1474,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 +1724,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 +1794,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..832601133 100644 --- a/tests/ut/py/test_callable_identity.py +++ b/tests/ut/py/test_callable_identity.py @@ -7,6 +7,7 @@ # See LICENSE in the root of the software repository for the full text of the License. # ----------------------------------------------------------------------------------------------------------- +import contextlib import ctypes import hashlib import os @@ -14,6 +15,7 @@ import subprocess import sys import tempfile +import threading import time from multiprocessing.shared_memory import SharedMemory from typing import cast @@ -540,6 +542,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 +558,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 +603,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] @@ -848,7 +855,8 @@ def submit_sub_group(self, *args): assert not fake.called -def test_capture_remote_sidecar_refs_rolls_back_partial_acquires(): +@pytest.mark.parametrize("acquire_committed", [False, True]) +def test_remote_slot_ref_journal_cleans_both_sides_of_acquire_boundary(monkeypatch, acquire_committed): class FakeTensorSidecar: present = True @@ -860,26 +868,135 @@ def __init__(self, *handles): self.tensors = tuple(FakeTensorSidecar(handle) for handle in handles) worker = Worker(level=4, num_sub_workers=0) - first = RemoteBufferHandle._from_remote_allocation( + remote_buf = RemoteBufferHandle._from_remote_allocation( worker_id=0, buffer_id=1, generation=1, address_space=RemoteAddressSpace.REMOTE_DEVICE, nbytes=4, ) - released = RemoteBufferHandle._from_remote_allocation( - worker_id=0, - buffer_id=2, - generation=1, - address_space=RemoteAddressSpace.REMOTE_DEVICE, - nbytes=4, - released=True, - ) + resources = worker_mod._RunResources() + worker._building_run_resources = resources + real_acquire = RemoteBufferHandle._acquire_slot_ref + interrupt = SystemExit("slot-ref acquire boundary") + + def interrupt_acquire(handle, token): + if acquire_committed: + real_acquire(handle, token) + raise interrupt + + monkeypatch.setattr(RemoteBufferHandle, "_acquire_slot_ref", interrupt_acquire) + + with pytest.raises(SystemExit) as caught: + worker._adopt_remote_sidecar_refs([FakeRemoteSidecar(remote_buf)]) + + assert caught.value is interrupt + assert remote_buf._live_slot_refs == int(acquire_committed) + assert len(resources.remote_slot_refs) == 1 + assert resources.requires_ordered_cleanup + + worker._release_active_remote_slot_refs(resources) + assert remote_buf._live_slot_refs == 0 + assert resources.remote_slot_refs == [] + + +def test_remote_submit_keeps_slot_ref_after_native_commit_boundary_interrupt(): + interrupt = SystemExit("after native slot commit") + + class FakeCOrchestrator: + committed = False + + def submit_next_level(self, *args): + self.committed = True + raise interrupt + + worker = Worker(level=4, num_sub_workers=0) + resources = worker_mod._RunResources() + worker._building_run_resources = resources + try: + worker_id = worker.add_remote_worker(RemoteWorkerSpec(endpoint="127.0.0.1:19073", platform="a2a3sim")) + handle = worker.register(RemoteCallable("pkg.remote:orch"), workers=[worker_id]) + fake = FakeCOrchestrator() + orch = Orchestrator(fake, worker=worker) # type: ignore[arg-type] + remote_buf = RemoteBufferHandle._from_remote_allocation( + worker_id=worker_id, + buffer_id=1, + generation=1, + address_space=RemoteAddressSpace.REMOTE_DEVICE, + nbytes=4, + ) + args = TaskArgs() + args.add_tensor(RemoteTensorRef(remote_buf, shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT) + + with pytest.raises(SystemExit) as caught: + orch.submit_next_level(handle, args, worker=worker_id) + + assert caught.value is interrupt + assert fake.committed + assert remote_buf._live_slot_refs == 1 + assert len(resources.remote_slot_refs) == 1 + + worker._release_active_remote_slot_refs(resources) + assert remote_buf._live_slot_refs == 0 + assert resources.remote_slot_refs == [] + finally: + worker._building_run_resources = None + worker._release_active_remote_slot_refs(resources) + worker.close() + + +def test_remote_group_submit_keeps_all_slot_refs_after_native_commit_boundary_interrupt(): + interrupt = SystemExit("after native group slot commit") + + class FakeCOrchestrator: + committed = False - with pytest.raises(RuntimeError, match="already been released"): - worker._capture_remote_sidecar_refs(FakeRemoteSidecar(first, released)) + def submit_next_level_group(self, *args): + self.committed = True + raise interrupt - assert first._live_slot_refs == 0 + worker = Worker(level=4, num_sub_workers=0) + resources = worker_mod._RunResources() + worker._building_run_resources = resources + try: + worker_ids = [ + worker.add_remote_worker(RemoteWorkerSpec(endpoint=f"127.0.0.1:{19073 + i}", platform="a2a3sim")) + for i in range(2) + ] + handle = worker.register(RemoteCallable("pkg.remote:orch"), workers=worker_ids) + fake = FakeCOrchestrator() + orch = Orchestrator(fake, worker=worker) # type: ignore[arg-type] + remote_bufs = [ + RemoteBufferHandle._from_remote_allocation( + worker_id=worker_id, + buffer_id=i + 1, + generation=1, + address_space=RemoteAddressSpace.REMOTE_DEVICE, + nbytes=4, + ) + for i, worker_id in enumerate(worker_ids) + ] + args_list = [] + for remote_buf in remote_bufs: + args = TaskArgs() + args.add_tensor(RemoteTensorRef(remote_buf, shape=(4,), dtype=DataType.UINT8), TensorArgType.INPUT) + args_list.append(args) + + with pytest.raises(SystemExit) as caught: + orch.submit_next_level_group(handle, args_list, workers=worker_ids) + + assert caught.value is interrupt + assert fake.committed + assert [remote_buf._live_slot_refs for remote_buf in remote_bufs] == [1, 1] + assert len(resources.remote_slot_refs) == 2 + + worker._release_active_remote_slot_refs(resources) + assert [remote_buf._live_slot_refs for remote_buf in remote_bufs] == [0, 0] + assert resources.remote_slot_refs == [] + finally: + worker._building_run_resources = None + worker._release_active_remote_slot_refs(resources) + worker.close() @pytest.mark.parametrize( @@ -1392,6 +1509,55 @@ def parent_orch(orch, _args, cfg): daemon.stop() +def _make_remote_import_test_values(worker): + owner = RemoteBufferHandle._from_remote_allocation( + worker_id=0, + buffer_id=1, + generation=1, + address_space=RemoteAddressSpace.REMOTE_DEVICE, + nbytes=4, + ) + exported = RemoteBufferExport._from_remote_export( + owner_worker_id=0, + buffer_id=1, + generation=1, + address_space=RemoteAddressSpace.REMOTE_WINDOW, + offset=0, + nbytes=4, + export_id=1, + remote_addr=0, + rkey_or_token=1, + ub_ldst_va=0, + access_flags=3, + transport_profile="sim", + _owner_handle=owner, + worker_owner_id=worker._owner_id, + ) + return owner, exported + + +def _make_remote_import_test_handle(owner): + owner._acquire_import_ref() + return RemoteBufferHandle._from_imported_mapping( + worker_id=0, + owner_worker_id=0, + buffer_id=1, + generation=1, + import_id=7, + address_space=RemoteAddressSpace.REMOTE_WINDOW, + nbytes=4, + offset=0, + owner_handle_ref=owner, + ) + + +def _start_fake_remote_worker(worker, fake): + worker_id = worker.add_remote_worker(RemoteWorkerSpec(endpoint="127.0.0.1:19073", platform="a2a3sim")) + worker._worker = fake + worker._lifecycle = worker_mod._Lifecycle.READY + return worker_id + + def test_remote_owner_free_waits_for_import_release(): worker = Worker(level=4, num_sub_workers=0) owner = RemoteBufferHandle._from_remote_allocation( @@ -1448,7 +1614,7 @@ def remote_free(self, *args): assert worker._pending_remote_buffer_frees == [] -def test_remote_import_pins_owner_during_control_and_rolls_back_on_error(): +def test_remote_import_pins_owner_and_retains_it_after_ambiguous_transport_error(): worker = Worker(level=4, num_sub_workers=0) owner = RemoteBufferHandle._from_remote_allocation( worker_id=0, @@ -1491,7 +1657,86 @@ def remote_import(self, *args): worker.remote_import(exported, worker=worker_id) assert fake.owner_ref_seen == 1 + assert owner._live_import_refs == 1 + assert worker._ordered_cleanup_error is not None + + +def test_remote_import_interrupt_after_owner_pin_retires_the_pin(monkeypatch): + worker = Worker(level=4, num_sub_workers=0) + owner, exported = _make_remote_import_test_values(worker) + + class FakeRemoteWorker: + def __init__(self): + self.import_calls = 0 + + def remote_import(self, *args): + self.import_calls += 1 + raise AssertionError("transport must not start after the acquisition interruption") + + fake = FakeRemoteWorker() + worker_id = _start_fake_remote_worker(worker, fake) + interrupt = KeyboardInterrupt("after owner reference acquisition") + helper_calls = 0 + + def interrupt_after_acquire(items, target, **_kwargs): + nonlocal helper_calls + helper_calls += 1 + for item in items: + target(item) + if helper_calls == 1: + raise interrupt + + monkeypatch.setattr(worker_mod, "_start_and_join_threads", interrupt_after_acquire) + + with pytest.raises(KeyboardInterrupt) as caught: + worker.remote_import(exported, worker=worker_id) + + assert caught.value is interrupt + assert fake.import_calls == 0 + assert owner._live_import_refs == 0 + assert worker._ordered_cleanup_error is None + + +def test_remote_import_owner_pin_token_recovers_an_acquire_error_after_mutation(monkeypatch): + worker = Worker(level=4, num_sub_workers=0) + owner, exported = _make_remote_import_test_values(worker) + interrupt = KeyboardInterrupt("owner acquire failed after mutation") + real_acquire = RemoteBufferHandle._acquire_import_ref + + def acquire_then_interrupt(handle, token=None): + real_acquire(handle, token) + raise interrupt + + class FakeRemoteWorker: + def remote_import(self, *args): + raise AssertionError("transport must not start after the acquisition error") + + worker_id = _start_fake_remote_worker(worker, FakeRemoteWorker()) + monkeypatch.setattr(RemoteBufferHandle, "_acquire_import_ref", acquire_then_interrupt) + + with pytest.raises(KeyboardInterrupt) as caught: + worker.remote_import(exported, worker=worker_id) + + assert caught.value is interrupt assert owner._live_import_refs == 0 + assert worker._ordered_cleanup_error is None + + +def test_remote_import_completion_without_descriptor_retains_owner_and_poisons(): + worker = Worker(level=4, num_sub_workers=0) + owner, exported = _make_remote_import_test_values(worker) + + class FakeRemoteWorker: + def remote_import(self, *args): + return None + + worker_id = _start_fake_remote_worker(worker, FakeRemoteWorker()) + + with pytest.raises(RuntimeError, match="returned no import descriptor"): + worker.remote_import(exported, worker=worker_id) + + assert owner._live_import_refs == 1 + assert worker._ordered_cleanup_error is not None def test_remote_import_releases_remote_mapping_when_handle_build_fails(monkeypatch): @@ -1546,6 +1791,447 @@ def fail_from_imported_mapping(**kwargs): assert owner._live_import_refs == 0 +def test_remote_import_rolls_back_a_mapping_published_before_caller_interruption(monkeypatch): + worker = Worker(level=4, num_sub_workers=0) + owner = RemoteBufferHandle._from_remote_allocation( + worker_id=0, + buffer_id=1, + generation=1, + address_space=RemoteAddressSpace.REMOTE_DEVICE, + nbytes=4, + ) + exported = RemoteBufferExport._from_remote_export( + owner_worker_id=0, + buffer_id=1, + generation=1, + address_space=RemoteAddressSpace.REMOTE_WINDOW, + offset=0, + nbytes=4, + export_id=1, + remote_addr=0, + rkey_or_token=1, + ub_ldst_va=0, + access_flags=3, + transport_profile="sim", + _owner_handle=owner, + worker_owner_id=worker._owner_id, + ) + + class FakeRemoteWorker: + def __init__(self): + self.import_calls = 0 + self.release_calls = 0 + + def remote_import(self, *args): + self.import_calls += 1 + return (0, 0, 1, 1, 7, int(RemoteAddressSpace.REMOTE_WINDOW), 4, 0, 0, 7, 0, 3) + + def remote_release_import(self, *args): + self.release_calls += 1 + + fake = FakeRemoteWorker() + worker_id = worker.add_remote_worker(RemoteWorkerSpec(endpoint="127.0.0.1:19073", platform="a2a3sim")) + worker._worker = fake # type: ignore[assignment] + worker._lifecycle = worker_mod._Lifecycle.READY + interrupt = KeyboardInterrupt("after remote import returned") + helper_calls = 0 + + def interrupt_after_import_target(items, target, **_kwargs): + nonlocal helper_calls + helper_calls += 1 + for item in items: + target(item) + if helper_calls == 2: + raise interrupt + + monkeypatch.setattr(worker_mod, "_start_and_join_threads", interrupt_after_import_target) + + with pytest.raises(KeyboardInterrupt) as caught: + worker.remote_import(exported, worker=worker_id) + + assert caught.value is interrupt + assert fake.import_calls == 1 + assert fake.release_calls == 1 + assert owner._live_import_refs == 0 + assert worker._ordered_cleanup_error is None + + +def test_direct_import_release_publishes_completion_before_caller_interruption(monkeypatch): + worker = Worker(level=4, num_sub_workers=0) + owner = RemoteBufferHandle._from_remote_allocation( + worker_id=0, + buffer_id=1, + generation=1, + address_space=RemoteAddressSpace.REMOTE_DEVICE, + nbytes=4, + ) + owner._acquire_import_ref() + imported = RemoteBufferHandle._from_imported_mapping( + worker_id=0, + owner_worker_id=0, + buffer_id=1, + generation=1, + import_id=7, + address_space=RemoteAddressSpace.REMOTE_WINDOW, + nbytes=4, + offset=0, + owner_handle_ref=owner, + ) + + class FakeRemoteWorker: + def __init__(self): + self.release_calls = 0 + + def remote_release_import(self, *args): + self.release_calls += 1 + + fake = FakeRemoteWorker() + worker.add_remote_worker(RemoteWorkerSpec(endpoint="127.0.0.1:19073", platform="a2a3sim")) + worker._worker = fake # type: ignore[assignment] + worker._lifecycle = worker_mod._Lifecycle.READY + interrupt = KeyboardInterrupt("after remote release returned") + helper_calls = 0 + + def interrupt_after_first_target(items, target, **_kwargs): + nonlocal helper_calls + helper_calls += 1 + for item in items: + target(item) + if helper_calls == 1: + raise interrupt + + monkeypatch.setattr(worker_mod, "_start_and_join_threads", interrupt_after_first_target) + + with pytest.raises(KeyboardInterrupt) as caught: + worker.remote_release_import(imported) + + assert caught.value is interrupt + assert fake.release_calls == 1 + assert imported.released + assert imported._owner_handle_ref is None + assert owner._live_import_refs == 0 + assert worker._pending_remote_import_releases == [] + worker._flush_pending_remote_frees() + worker.remote_release_import(imported) + assert fake.release_calls == 1 + + +def test_concurrent_direct_import_releases_send_one_rpc(monkeypatch): + worker = Worker(level=4, num_sub_workers=0) + owner, _exported = _make_remote_import_test_values(worker) + imported = _make_remote_import_test_handle(owner) + + class FakeRemoteWorker: + def __init__(self): + self.release_calls = 0 + + def remote_release_import(self, *args): + self.release_calls += 1 + + fake = FakeRemoteWorker() + _start_fake_remote_worker(worker, fake) + barrier = threading.Barrier(2) + real_operation_lease = worker._operation_lease + + @contextlib.contextmanager + def synchronized_operation_lease(api): + with real_operation_lease(api): + barrier.wait(timeout=5) + yield + + monkeypatch.setattr(worker, "_operation_lease", synchronized_operation_lease) + errors = [] + + def release_import(): + try: + worker.remote_release_import(imported) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + + threads = [threading.Thread(target=release_import) for _ in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=5) + + assert all(not thread.is_alive() for thread in threads) + assert errors == [] + assert fake.release_calls == 1 + assert owner._live_import_refs == 0 + + +def test_concurrent_direct_owner_frees_send_one_rpc(monkeypatch): + worker = Worker(level=4, num_sub_workers=0) + owner, _exported = _make_remote_import_test_values(worker) + + class FakeRemoteWorker: + def __init__(self): + self.free_calls = 0 + + def remote_free(self, *args): + self.free_calls += 1 + + fake = FakeRemoteWorker() + _start_fake_remote_worker(worker, fake) + barrier = threading.Barrier(2) + real_operation_lease = worker._operation_lease + + @contextlib.contextmanager + def synchronized_operation_lease(api): + with real_operation_lease(api): + barrier.wait(timeout=5) + yield + + monkeypatch.setattr(worker, "_operation_lease", synchronized_operation_lease) + errors = [] + + def free_owner(): + try: + worker.remote_free(owner) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + + threads = [threading.Thread(target=free_owner) for _ in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=5) + + assert all(not thread.is_alive() for thread in threads) + assert errors == [] + assert fake.free_calls == 1 + + +def test_owner_free_is_durably_published_before_rpc(): + worker = Worker(level=4, num_sub_workers=0) + owner, _exported = _make_remote_import_test_values(worker) + + class FakeRemoteWorker: + def __init__(self): + self.free_calls = 0 + + def remote_free(self, *args): + assert owner.released + assert worker._pending_remote_buffer_frees == [owner] + self.free_calls += 1 + + fake = FakeRemoteWorker() + _start_fake_remote_worker(worker, fake) + + worker.remote_free(owner) + + assert fake.free_calls == 1 + assert worker._pending_remote_buffer_frees == [] + + +def test_failed_owner_free_can_be_retried_explicitly(): + worker = Worker(level=4, num_sub_workers=0) + owner, _exported = _make_remote_import_test_values(worker) + + class FakeRemoteWorker: + def __init__(self): + self.free_calls = 0 + + def remote_free(self, *args): + self.free_calls += 1 + if self.free_calls == 1: + raise RuntimeError("transient free failure") + + fake = FakeRemoteWorker() + _start_fake_remote_worker(worker, fake) + + with pytest.raises(RuntimeError, match="transient free failure"): + worker.remote_free(owner) + + assert owner.released + assert worker._pending_remote_buffer_frees == [owner] + + worker.remote_free(owner) + worker.remote_free(owner) + + assert fake.free_calls == 2 + assert worker._pending_remote_buffer_frees == [] + + +def test_owner_free_flush_cannot_miss_half_published_debt(monkeypatch): + worker = Worker(level=4, num_sub_workers=0) + owner, _exported = _make_remote_import_test_values(worker) + slot_ref_token = owner._acquire_slot_ref() + mark_entered = threading.Event() + allow_mark = threading.Event() + flush_lock_attempted = threading.Event() + + class FakeRemoteWorker: + def __init__(self): + self.free_calls = 0 + + def remote_free(self, *args): + self.free_calls += 1 + + fake = FakeRemoteWorker() + _start_fake_remote_worker(worker, fake) + real_mark_released = RemoteBufferHandle._mark_released + real_release_lock = worker._remote_import_release_mu + + class ObservedReleaseLock: + def __enter__(self): + if threading.current_thread().name == "owner-free-flusher": + flush_lock_attempted.set() + real_release_lock.acquire() + return self + + def __exit__(self, exc_type, exc_value, traceback): + real_release_lock.release() + + monkeypatch.setattr(worker, "_remote_import_release_mu", ObservedReleaseLock()) + + def blocked_mark_released(handle): + if handle is owner: + mark_entered.set() + allow_mark.wait(timeout=5) + real_mark_released(handle) + + monkeypatch.setattr(RemoteBufferHandle, "_mark_released", blocked_mark_released) + errors = [] + + def publish(): + try: + worker.remote_free(owner) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + + def retire_and_flush(): + try: + owner._release_slot_ref(slot_ref_token) + worker._flush_pending_remote_frees() + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + + publisher = threading.Thread(target=publish) + publisher.start() + assert mark_entered.wait(timeout=5) + flusher = threading.Thread(target=retire_and_flush, name="owner-free-flusher") + flusher.start() + assert flush_lock_attempted.wait(timeout=5) + assert fake.free_calls == 0 + allow_mark.set() + publisher.join(timeout=5) + flusher.join(timeout=5) + + assert not publisher.is_alive() + assert not flusher.is_alive() + assert errors == [] + assert fake.free_calls == 1 + assert worker._pending_remote_buffer_frees == [] + + +def test_import_release_flush_cannot_observe_half_published_state(monkeypatch): + worker = Worker(level=4, num_sub_workers=0) + owner, _exported = _make_remote_import_test_values(worker) + imported = _make_remote_import_test_handle(owner) + mark_entered = threading.Event() + allow_mark = threading.Event() + flush_lock_attempted = threading.Event() + rpc_started = threading.Event() + + class FakeRemoteWorker: + def __init__(self): + self.release_calls = 0 + + def remote_release_import(self, *args): + self.release_calls += 1 + rpc_started.set() + + fake = FakeRemoteWorker() + _start_fake_remote_worker(worker, fake) + real_mark_released = RemoteBufferHandle._mark_released + real_release_lock = worker._remote_import_release_mu + + class ObservedReleaseLock: + def __enter__(self): + if threading.current_thread().name == "import-release-flusher": + flush_lock_attempted.set() + real_release_lock.acquire() + return self + + def __exit__(self, exc_type, exc_value, traceback): + real_release_lock.release() + + monkeypatch.setattr(worker, "_remote_import_release_mu", ObservedReleaseLock()) + + def blocked_mark_released(handle): + if handle is imported: + mark_entered.set() + allow_mark.wait(timeout=5) + real_mark_released(handle) + + monkeypatch.setattr(RemoteBufferHandle, "_mark_released", blocked_mark_released) + errors = [] + + def publish(): + try: + worker._publish_pending_remote_import_release( + imported, + worker_mod._PendingRemoteImportReleaseState(owner_ref=owner), + ) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + + def flush(): + try: + worker._flush_pending_remote_frees() + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + + publisher = threading.Thread(target=publish) + publisher.start() + assert mark_entered.wait(timeout=5) + flusher = threading.Thread(target=flush, name="import-release-flusher") + flusher.start() + assert flush_lock_attempted.wait(timeout=5) + assert not rpc_started.is_set() + allow_mark.set() + publisher.join(timeout=5) + flusher.join(timeout=5) + + assert not publisher.is_alive() + assert not flusher.is_alive() + assert errors == [] + assert fake.release_calls == 1 + assert owner._live_import_refs == 0 + + +def test_import_release_reply_loss_is_poisoned_and_never_replayed(): + worker = Worker(level=4, num_sub_workers=0) + owner, _exported = _make_remote_import_test_values(worker) + imported = _make_remote_import_test_handle(owner) + + class FakeRemoteWorker: + def __init__(self): + self.release_calls = 0 + + def remote_release_import(self, *args): + self.release_calls += 1 + raise RuntimeError("release reply lost") + + fake = FakeRemoteWorker() + _start_fake_remote_worker(worker, fake) + + with pytest.raises(RuntimeError, match="release reply lost"): + worker.remote_release_import(imported) + + assert fake.release_calls == 1 + assert owner._live_import_refs == 1 + assert worker._ordered_cleanup_error is not None + assert worker._pending_remote_import_releases == [imported] + + with pytest.raises(RuntimeError, match="release reply lost"): + worker._flush_pending_remote_frees() + + assert fake.release_calls == 1 + assert owner._live_import_refs == 1 + + def test_remote_import_rejects_cross_worker_or_stale_export(): owner = RemoteBufferHandle._from_remote_allocation( worker_id=0, @@ -1605,7 +2291,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_child_addr_guard.py b/tests/ut/py/test_worker/test_child_addr_guard.py index 24d2a01a5..d514c85f9 100644 --- a/tests/ut/py/test_worker/test_child_addr_guard.py +++ b/tests/ut/py/test_worker/test_child_addr_guard.py @@ -586,7 +586,13 @@ def _domain_worker(self): with w._child_prov_lock: w._child_prov_record_domain(0, 0x5000, allocation_id=9, extent=64) w._child_prov_record_domain(1, 0x6000, allocation_id=9, extent=64) - handle = SimpleNamespace(name="d", workers=(0, 1), allocation_id=9) + handle = SimpleNamespace( + name="d", + workers=(0, 1), + allocation_id=9, + _domain_size=2, + _domain_ranks={0: 0, 1: 1}, + ) return w, handle def test_release_revokes_provenance_before_backend_free(self, monkeypatch): diff --git a/tests/ut/py/test_worker/test_host_worker.py b/tests/ut/py/test_worker/test_host_worker.py index ac93af168..3dcc144ec 100644 --- a/tests/ut/py/test_worker/test_host_worker.py +++ b/tests/ut/py/test_worker/test_host_worker.py @@ -12,16 +12,24 @@ Each test verifies a distinct aspect of the L3 scheduling pipeline. """ +import _thread import ctypes +import dis import gc +import inspect +import multiprocessing.shared_memory as shared_memory_mod import struct +import sys import threading 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 +138,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 +150,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 +177,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: @@ -1393,6 +1406,162 @@ def orch(o, args, cfg): class TestRunHandle: + @staticmethod + def _submission_failure_worker(failures: int): + events: list[str] = [] + + class NativeWorker: + def close(self): + events.append("close") + + class NativeOrchestrator: + def __init__(self): + self.failures_left = failures + + def _begin_run(self): + events.append("begin") + return 1 + + def _scope_begin(self): + events.append("scope_begin") + + def _scope_end(self): + events.append("scope_end") + + def _fail_run_submission(self, run_id, _error): + assert run_id == 1 + events.append("fail") + if self.failures_left: + self.failures_left -= 1 + raise RuntimeError("injected cancellation failure") + + def _wait_run(self, run_id): + assert run_id == 1 + events.append("wait") + raise RuntimeError("native graph failure") + + def _release_run(self, run_id): + assert run_id == 1 + events.append("release") + + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, NativeWorker()) + worker._orch = cast(Any, NativeOrchestrator()) + return worker, events + + def test_graph_failure_retries_native_cancellation_before_waiting(self): + worker, events = self._submission_failure_worker(failures=1) + graph_error = ValueError("bad graph") + + def bad_graph(*_args): + raise graph_error + + with pytest.raises(ValueError) as excinfo: + worker._submit_l3_locked(bad_graph, None, cast(Any, object())) + + assert excinfo.value is graph_error + assert events == ["begin", "scope_begin", "scope_end", "fail", "fail", "wait", "release"] + assert not worker._accepted_run_handles + assert worker._ordered_cleanup_error is None + + def test_unsettled_graph_cancellation_abandons_the_handle_before_close(self): + worker, events = self._submission_failure_worker(failures=2) + graph_error = ValueError("bad graph") + + def bad_graph(*_args): + raise graph_error + + with pytest.raises(ValueError) as excinfo: + worker._submit_l3_locked(bad_graph, None, cast(Any, object())) + + assert excinfo.value is graph_error + assert events == ["begin", "scope_begin", "scope_end", "fail", "fail"] + assert not worker._accepted_run_handles + assert len(worker._abandoned_run_handles) == 1 + assert worker._abandoned_run_handles[0]._keepalive is not None + assert worker._ordered_cleanup_error is not None + with pytest.raises(RuntimeError, match="no further work is admitted"): + worker._require_no_ordered_cleanup_failure("submit") + + worker.close() + assert events[-1] == "close" + assert not worker._abandoned_run_handles + + def test_unsettled_graph_cancellation_publishes_through_a_cv_enter_interrupt(self): + publication_interrupt = KeyboardInterrupt("cancellation publication") + cancellation_error = RuntimeError("injected cancellation failure") + graph_error = ValueError("bad graph") + + class ArmableCV: + def __init__(self, cv): + self.cv = cv + self.armed = False + self.interrupts = 0 + + def __enter__(self): + if self.armed: + self.armed = False + self.interrupts += 1 + raise publication_interrupt + return self.cv.__enter__() + + def __exit__(self, *exc_info): + return self.cv.__exit__(*exc_info) + + def notify_all(self): + self.cv.notify_all() + + class NativeWorker: + def close(self): + return None + + class NativeOrchestrator: + def __init__(self, lifecycle_cv): + self.lifecycle_cv = lifecycle_cv + self.cancellations = 0 + + def _begin_run(self): + return 1 + + def _scope_begin(self): + return None + + def _scope_end(self): + return None + + def _fail_run_submission(self, run_id, _error): + assert run_id == 1 + self.cancellations += 1 + if self.cancellations == worker_mod._RUN_CANCELLATION_ATTEMPTS: + self.lifecycle_cv.armed = True + raise cancellation_error + + worker = Worker(level=3, num_sub_workers=0) + lifecycle_cv = ArmableCV(worker._hierarchical_start_cv) + worker._hierarchical_start_cv = cast(Any, lifecycle_cv) + worker._worker = cast(Any, NativeWorker()) + worker._orch = cast(Any, NativeOrchestrator(lifecycle_cv)) + + def bad_graph(*_args): + raise graph_error + + with pytest.raises(ValueError) as caught: + worker._submit_l3_locked(bad_graph, None, cast(Any, object())) + + assert caught.value is graph_error + assert lifecycle_cv.interrupts == 1 + assert not worker._accepted_run_handles + assert len(worker._abandoned_run_handles) == 1 + abandoned = worker._abandoned_run_handles[0] + assert abandoned._terminal + assert isinstance(abandoned._error, RuntimeError) + assert abandoned._error is worker._ordered_cleanup_error + assert abandoned._error.__cause__ is cancellation_error + assert abandoned._keepalive is not None + + worker.close() + assert abandoned._keepalive is None + def test_submit_returns_before_completion_and_timeout_is_retryable(self): state_shm = SharedMemory(create=True, size=8) state_buf = state_shm.buf @@ -1435,6 +1604,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() @@ -1501,6 +1741,87 @@ def delayed(_args): state_shm.close() state_shm.unlink() + def test_close_uses_one_deadline_for_operations_and_run_fences_then_retries(self, monkeypatch): + class Clock: + now = 0.0 + + clock = Clock() + monkeypatch.setattr(worker_mod.time, "monotonic", lambda: clock.now) + monkeypatch.setattr(worker_mod, "_ROLLBACK_GRACEFUL_TIMEOUT_S", 10.0) + + wait_budgets: list[float] = [] + released_runs: list[int] = [] + native_closes: list[str] = [] + + class NativeWorker: + def close(self): + native_closes.append("close") + + class NativeOrchestrator: + complete = False + + def _wait_run_for(self, run_id, timeout): + assert run_id == 1 + wait_budgets.append(timeout) + if not self.complete: + clock.now += timeout + return False + return True + + def _release_run(self, run_id): + released_runs.append(run_id) + + worker = Worker(level=3, num_sub_workers=0) + native_worker = NativeWorker() + native_orch = NativeOrchestrator() + worker._worker = cast(Any, native_worker) + worker._orch = cast(Any, native_orch) + handle = RunHandle(worker, 1, ()) + worker._accepted_run_handles.add(handle) + worker._active_ops = 1 + + real_cv = worker._hierarchical_start_cv + + class AdvancingCondition: + def __enter__(self): + return real_cv.__enter__() + + def __exit__(self, *exc_info): + return real_cv.__exit__(*exc_info) + + def wait(self, timeout=None): + assert timeout == pytest.approx(10.0) + clock.now += 7.0 + worker._active_ops = 0 + return True + + def notify_all(self): + real_cv.notify_all() + + worker._hierarchical_start_cv = cast(Any, AdvancingCondition()) + + with pytest.raises(TimeoutError, match="run fence.*cleanup budget"): + worker.close() + + assert wait_budgets == [pytest.approx(3.0)] + assert handle in worker._accepted_run_handles + assert not handle._terminal + assert worker._worker is native_worker + assert not worker._teardown_attempted + assert worker._close_completion is not None and worker._close_completion.incomplete + assert native_closes == [] + + native_orch.complete = True + worker.close() + + assert wait_budgets == [pytest.approx(3.0), pytest.approx(10.0)] + assert released_runs == [1] + assert handle._terminal + assert not worker._accepted_run_handles + assert worker._teardown_attempted + assert worker._close_completion is not None and not worker._close_completion.incomplete + assert native_closes == ["close"] + def test_submit_close_race_accepts_and_drains_admitted_run(self): callback_entered = threading.Event() callback_release = threading.Event() @@ -1756,367 +2077,3450 @@ def notify_all(self): assert hw._finalize_run_handle(handle, 1, None) is interrupt assert not hw._accepted_run_handles - def test_run_finalization_releases_only_its_resources(self, monkeypatch): - class SlotRef: - def __init__(self): - self.releases = 0 + def test_finalize_drains_after_a_post_step_interrupt_and_close_completes(self): + interrupt = KeyboardInterrupt("post-step") + released_refs: list[str] = [] + released_runs: list[int] = [] + class SlotRef: def _release_slot_ref(self): - self.releases += 1 + released_refs.append("released") - class Region: - def __init__(self, region_id): - self.region_id = region_id - self._worker_id = 0 - self.mapping_closed = False - self.expired = False + class NativeOrchestrator: + def _wait_run(self, run_id): + assert run_id == 1 - def _close_l3_host_mapping(self): - self.mapping_closed = True + def _release_run(self, run_id): + released_runs.append(run_id) - def _expire(self): - self.expired = True + worker = Worker(level=3, num_sub_workers=0) + worker._orch = cast(Any, NativeOrchestrator()) + resources = worker_mod._RunResources() + resources.remote_slot_refs.append(cast(Any, SlotRef())) + handle = RunHandle(worker, 1, (object(),), resources) + worker._accepted_run_handles.add(handle) + original_finalize = worker._finalize_run_handle + interrupted = False - class NativeWorker: - def __init__(self): - self.released_regions = [] + def interrupt_after_remote_refs(step: str) -> None: + nonlocal interrupted + if step == "remote_slot_refs" and not interrupted: + interrupted = True + raise interrupt - def control_l3_l2_region_release(self, worker_id, region_id): - self.released_regions.append((worker_id, region_id)) + worker._finalize_run_handle = cast( + Any, + lambda finalized, run_id, error: original_finalize( + finalized, + run_id, + error, + _after_step=interrupt_after_remote_refs, + ), + ) + + with pytest.raises(KeyboardInterrupt) as caught: + handle.wait() + + assert caught.value is interrupt + assert released_refs == ["released"] + assert released_runs == [1] + assert handle._cleanup_published + assert handle not in worker._accepted_run_handles + assert handle._keepalive is None + worker.close() + + def test_wait_interrupted_after_election_is_re_electable(self): + interrupt = KeyboardInterrupt("after election") + nested_interrupt = SystemExit("while clearing election") + native_waits: list[int] = [] + released_runs: list[int] = [] + + class InterruptingHandle(RunHandle): + interrupt_clear = False + + def __setattr__(self, name, value): + if name == "_wait_in_progress" and value is False and self.interrupt_clear: + self.interrupt_clear = False + raise nested_interrupt + return super().__setattr__(name, value) class NativeOrchestrator: - def __init__(self): - self.released_runs = [] + def _run_done(self, run_id): + assert run_id == 1 + return False - def _release_run(self, run_id): - self.released_runs.append(run_id) + def _wait_run(self, run_id): + native_waits.append(run_id) - def domain(name, allocation_id): - return worker_mod.CommDomainHandle( - name=name, - workers=(), - contexts={}, - allocation_id=allocation_id, - _release_fn=lambda _handle: None, - ) + def _release_run(self, run_id): + released_runs.append(run_id) worker = Worker(level=3, num_sub_workers=0) - native_worker = NativeWorker() - native_orch = NativeOrchestrator() - worker._worker = cast(Any, native_worker) - worker._orch = cast(Any, native_orch) - - first = worker_mod._RunResources() - second = worker_mod._RunResources() - first_ref, second_ref = SlotRef(), SlotRef() - first_region, second_region = Region(11), Region(22) - first_live, second_live = domain("first-live", 1), domain("second-live", 2) - first_pending, second_pending = domain("first-pending", 3), domain("second-pending", 4) - first.remote_slot_refs.append(cast(Any, first_ref)) - second.remote_slot_refs.append(cast(Any, second_ref)) - first.l3_l2_regions.append(first_region) - second.l3_l2_regions.append(second_region) - first.l3_l2_orch_comm_host_buffers[0x1000] = 64 - second.l3_l2_orch_comm_host_buffers[0x2000] = 128 - first.live_domains[first_live.name] = first_live - second.live_domains[second_live.name] = second_live - first.pending_release_domains.append(first_pending) - second.pending_release_domains.append(second_pending) - worker._live_l3_l2_regions.extend([first_region, second_region]) - worker._live_domains.update({first_live.name: first_live, second_live.name: second_live}) + worker._orch = cast(Any, NativeOrchestrator()) + handle = InterruptingHandle(worker, 1, ()) + worker._accepted_run_handles.add(handle) + interrupted = False - released_domains = [] + def interrupt_after_election(phase: str) -> None: + nonlocal interrupted + if phase == "after_election" and not interrupted: + interrupted = True + raise interrupt - def release_domain_now(handle): - released_domains.append(handle) - if worker._live_domains.get(handle.name) is handle: - worker._live_domains.pop(handle.name) + handle._wait_boundary_hook = interrupt_after_election + handle.interrupt_clear = True - monkeypatch.setattr(worker, "_release_domain_now", release_domain_now) - first_handle = RunHandle(worker, 1, (), first) - second_handle = RunHandle(worker, 2, (), second) - worker._accepted_run_handles.update({first_handle, second_handle}) + with pytest.raises(KeyboardInterrupt) as caught: + handle.wait() - assert worker._finalize_run_handle(first_handle, 1, None) is None + assert caught.value is interrupt + assert not handle._wait_in_progress + assert not handle.done + assert native_waits == [] + assert released_runs == [] + assert handle in worker._accepted_run_handles - assert first_ref.releases == 1 - assert second_ref.releases == 0 - assert native_worker.released_regions == [(0, 11)] - assert first_region.mapping_closed and first_region.expired - assert not second_region.mapping_closed and not second_region.expired - assert worker._live_l3_l2_regions == [second_region] - assert first.l3_l2_orch_comm_host_buffers == {} - assert second.l3_l2_orch_comm_host_buffers == {0x2000: 128} - assert released_domains == [first_pending, first_live] - assert first_pending.freed and first_live.freed - assert not second_pending.freed and not second_live.freed - assert worker._live_domains == {second_live.name: second_live} - assert native_orch.released_runs == [1] - assert worker._accepted_run_handles == {second_handle} + handle.wait() - def test_domain_released_after_its_run_retired_is_freed_inline(self): - """A late release has no fence left to defer behind, so it frees now. + assert native_waits == [1] + assert released_runs == [1] + assert handle.done + assert handle not in worker._accepted_run_handles + worker.close() - Both deferred paths are closed to it: the run's queue is never drained - again, and _release_domain_handle has already dropped the handle from - _live_domains, so close()'s live sweep cannot reach it either. - """ - worker = Worker(level=3, num_sub_workers=0) - worker._worker = cast(Any, object()) - worker._orch = cast(Any, type("FakeOrch", (), {"_release_run": lambda self, run_id: None})()) + def test_wait_interrupted_after_finalize_publishes_terminal_once(self): + interrupt = KeyboardInterrupt("after finalize") + nested_interrupt = SystemExit("while publishing terminal") + native_waits: list[int] = [] + released_runs: list[int] = [] - freed: list[str] = [] + class InterruptingHandle(RunHandle): + interrupt_terminal = False - def release_now(handle): - freed.append(handle.name) - if worker._live_domains.get(handle.name) is handle: - worker._live_domains.pop(handle.name) + def __setattr__(self, name, value): + if name == "_terminal" and value is True and self.interrupt_terminal: + self.interrupt_terminal = False + raise nested_interrupt + return super().__setattr__(name, value) - worker._release_domain_now = cast(Any, release_now) + class NativeOrchestrator: + def _wait_run(self, run_id): + native_waits.append(run_id) - resources = worker_mod._RunResources() - late = worker_mod.CommDomainHandle( - name="late", - workers=(), - contexts={}, - allocation_id=1, - _release_fn=lambda released, owner=resources: worker._release_domain_handle(released, owner), - ) - resources.live_domains[late.name] = late - worker._live_domains[late.name] = late + def _release_run(self, run_id): + released_runs.append(run_id) - handle = RunHandle(worker, 1, (), resources) + worker = Worker(level=3, num_sub_workers=0) + worker._orch = cast(Any, NativeOrchestrator()) + handle = InterruptingHandle(worker, 1, ()) worker._accepted_run_handles.add(handle) - assert worker._finalize_run_handle(handle, 1, None) is None - assert freed == ["late"], "a domain still live at its run's fence is swept there" + interrupted = False - second = worker_mod.CommDomainHandle( - name="later", - workers=(), - contexts={}, - allocation_id=2, - _release_fn=lambda released, owner=resources: worker._release_domain_handle(released, owner), - ) - resources.live_domains[second.name] = second - worker._live_domains[second.name] = second + def interrupt_after_finalize(phase: str) -> None: + nonlocal interrupted + if phase == "after_finalize" and not interrupted: + interrupted = True + raise interrupt - second.release() + handle._wait_boundary_hook = interrupt_after_finalize + handle.interrupt_terminal = True - assert freed == ["late", "later"] - assert second.freed - assert resources.pending_release_domains == [] + with pytest.raises(KeyboardInterrupt) as caught: + handle.wait() - def test_domain_released_while_its_run_retires_is_still_freed(self): - """A release in flight across retirement must not be stranded. + assert caught.value is interrupt + assert native_waits == [1] + assert released_runs == [1] + assert handle.done + assert not handle._wait_in_progress + assert handle not in worker._accepted_run_handles + with pytest.raises(KeyboardInterrupt) as repeated: + handle.wait() + assert repeated.value is interrupt + assert native_waits == [1] + assert released_runs == [1] + worker.close() - The window the fence has to close: the deferred queue has already been - drained for the last time, so a handle appended after that is reachable - from nothing — the queue is never read again, and the release itself - popped the handle from ``_live_domains``, blinding ``close()``'s sweep. + def test_finalization_recovery_survives_an_interrupt_after_accepted_retirement(self): + finalization_interrupt = KeyboardInterrupt("finalization escaped") + recovery_interrupt = SystemExit("after accepted retirement") + recoveries: list[str] = [] - Forced deterministically by parking the fence in its live-domain sweep, - which holds no lock, while the releasing thread runs to completion. - """ - worker = Worker(level=3, num_sub_workers=0) - worker._worker = cast(Any, object()) - worker._orch = cast(Any, type("FakeOrch", (), {"_release_run": lambda self, run_id: None})()) + class NativeWorker: + def close(self): + return None + + class NativeOrchestrator: + def _wait_run(self, run_id): + assert run_id == 1 + + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, NativeWorker()) + worker._orch = cast(Any, NativeOrchestrator()) + keepalive = object() + handle = RunHandle(worker, 1, (keepalive,)) + worker._accepted_run_handles.add(handle) + worker._finalize_run_handle = cast(Any, lambda *_args: (_ for _ in ()).throw(finalization_interrupt)) + recover = worker._recover_interrupted_run_finalization + + def interrupt_after_recovery(recovering, error): + recoveries.append("recover") + recovered = recover(recovering, error) + if len(recoveries) == 1: + raise recovery_interrupt + return recovered + + worker._recover_interrupted_run_finalization = cast(Any, interrupt_after_recovery) + + with pytest.raises(KeyboardInterrupt) as caught: + handle.wait() + + assert caught.value is finalization_interrupt + assert recoveries == ["recover", "recover"] + assert handle._terminal + assert not handle._wait_in_progress + assert handle._error is finalization_interrupt + assert handle not in worker._accepted_run_handles + assert worker._abandoned_run_handles == [handle] + assert handle._keepalive == (keepalive,) + with pytest.raises(KeyboardInterrupt) as repeated: + handle.wait() + assert repeated.value is finalization_interrupt + + worker.close() + assert handle._keepalive is None + + def test_wait_timeout_includes_acceptance_owner_handoff(self): + native_waits: list[tuple[str, Any]] = [] + released_runs: list[int] = [] + + class NativeOrchestrator: + def _wait_run_for(self, run_id, timeout): + assert run_id == 1 + native_waits.append(("timed", timeout)) + return True + + def _wait_run(self, run_id): + assert run_id == 1 + native_waits.append(("untimed", None)) + + def _release_run(self, run_id): + released_runs.append(run_id) + + worker = Worker(level=3, num_sub_workers=0) + worker._orch = cast(Any, NativeOrchestrator()) + handle = RunHandle(worker, 1, ()) + handle._accept_wait_in_progress = True + worker._accepted_run_handles.add(handle) + finished = threading.Event() + errors: list[BaseException] = [] + + def wait_with_expired_deadline() -> None: + try: + handle.wait(0) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + finally: + finished.set() + + waiter = threading.Thread(target=wait_with_expired_deadline) + waiter.start() + try: + assert finished.wait(0.5), "RunHandle.wait(0) ignored its deadline during acceptance handoff" + assert len(errors) == 1 + assert isinstance(errors[0], TimeoutError) + assert handle._accept_wait_in_progress + assert not handle._wait_in_progress + assert handle in worker._accepted_run_handles + assert released_runs == [] + finally: + handle._clear_acceptance_owner() + waiter.join(5.0) + + handle.wait() + assert native_waits[0][0] == "timed" + assert native_waits[1] == ("untimed", None) + assert released_runs == [1] + worker.close() + + def test_finalizer_escape_after_retirement_preserves_cleanup_error(self): + cleanup_error = RuntimeError("cleanup failed") + boundary_error = KeyboardInterrupt("return boundary") + native_waits: list[int] = [] + released_runs: list[int] = [] + + class SlotRef: + def _release_slot_ref(self): + raise cleanup_error + + class NativeOrchestrator: + def _wait_run(self, run_id): + native_waits.append(run_id) + + def _release_run(self, run_id): + released_runs.append(run_id) + + worker = Worker(level=3, num_sub_workers=0) + worker._orch = cast(Any, NativeOrchestrator()) + resources = worker_mod._RunResources() + resources.remote_slot_refs.append(cast(Any, SlotRef())) + handle = RunHandle(worker, 1, (), resources) + worker._accepted_run_handles.add(handle) + original_finalize = worker._finalize_run_handle + + def interrupt_return(finalized, run_id, native_error): + assert original_finalize(finalized, run_id, native_error) is cleanup_error + raise boundary_error + + worker._finalize_run_handle = cast(Any, interrupt_return) + + with pytest.raises(RuntimeError) as caught: + handle.wait() + + assert caught.value is cleanup_error + assert handle._error is cleanup_error + assert worker._ordered_cleanup_error is cleanup_error + assert native_waits == [1] + assert released_runs == [1] + assert handle not in worker._accepted_run_handles + with pytest.raises(RuntimeError) as repeated: + handle.wait() + assert repeated.value is cleanup_error + + def test_acceptance_wait_interrupted_after_election_is_re_electable(self): + interrupt = KeyboardInterrupt("after acceptance election") + nested_interrupt = SystemExit("while clearing acceptance election") + native_waits: list[int] = [] + + class InterruptingHandle(RunHandle): + interrupt_clear = False + + def __setattr__(self, name, value): + if name == "_accept_wait_in_progress" and value is False and self.interrupt_clear: + self.interrupt_clear = False + raise nested_interrupt + return super().__setattr__(name, value) + + class FakeWorker: + def _wait_run_handle_accepted(self, run_id): + native_waits.append(run_id) + + handle = InterruptingHandle(cast(Worker, FakeWorker()), 1, ()) + interrupted = False + + def interrupt_after_election(phase: str) -> None: + nonlocal interrupted + if phase == "after_acceptance_election" and not interrupted: + interrupted = True + raise interrupt + + handle._wait_boundary_hook = interrupt_after_election + handle.interrupt_clear = True + + with pytest.raises(KeyboardInterrupt) as caught: + handle._wait_for_acceptance() + + assert caught.value is interrupt + assert not handle._accept_wait_in_progress + assert not handle._launch_accepted + assert native_waits == [] + + handle._wait_for_acceptance() + + assert native_waits == [1] + assert handle._launch_accepted + assert not handle._accept_wait_in_progress + + def test_acceptance_wait_interrupted_after_native_wait_stays_published(self): + interrupt = KeyboardInterrupt("after acceptance wait") + nested_interrupt = SystemExit("while publishing acceptance") + native_waits: list[int] = [] + + class InterruptingHandle(RunHandle): + interrupt_publish = False + + def __setattr__(self, name, value): + if name == "_accept_wait_in_progress" and value is False and self.interrupt_publish: + self.interrupt_publish = False + raise nested_interrupt + return super().__setattr__(name, value) + + class FakeWorker: + def _wait_run_handle_accepted(self, run_id): + native_waits.append(run_id) + + handle = InterruptingHandle(cast(Worker, FakeWorker()), 1, ()) + interrupted = False + + def interrupt_after_wait(phase: str) -> None: + nonlocal interrupted + if phase == "after_acceptance_wait" and not interrupted: + interrupted = True + raise interrupt + + handle._wait_boundary_hook = interrupt_after_wait + handle.interrupt_publish = True + + with pytest.raises(KeyboardInterrupt) as caught: + handle._wait_for_acceptance() + + assert caught.value is interrupt + assert native_waits == [1] + assert handle._launch_accepted + assert not handle._accept_wait_in_progress + + handle._wait_for_acceptance() + assert native_waits == [1] + + def test_interrupted_cleanup_is_abandoned_until_tree_teardown(self, monkeypatch): + interrupt = KeyboardInterrupt("ambiguous step boundary") + released_refs: list[str] = [] + released_runs: list[int] = [] + teardown_calls: list[str] = [] + + class SlotRef: + def __init__(self, name, *, interrupts=False): + self.name = name + self.interrupts = interrupts + + def _release_slot_ref(self): + released_refs.append(self.name) + if self.interrupts: + raise interrupt + + class NativeOrchestrator: + def _wait_run(self, run_id): + assert run_id == 1 + + def _release_run(self, run_id): + released_runs.append(run_id) + + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + worker._orch = cast(Any, NativeOrchestrator()) + resources = worker_mod._RunResources() + first_ref = SlotRef("first", interrupts=True) + second_ref = SlotRef("second") + resources.remote_slot_refs.extend([cast(Any, first_ref), cast(Any, second_ref)]) + keepalive = object() + handle = RunHandle(worker, 1, (keepalive, first_ref, second_ref), resources) + worker._accepted_run_handles.add(handle) + + def teardown_tree(): + teardown_calls.append("teardown") + worker._worker = None + worker._orch = None + + monkeypatch.setattr(worker, "_teardown_ready_tree", teardown_tree) + + with pytest.raises(KeyboardInterrupt) as caught: + handle.wait() + + assert caught.value is interrupt + assert released_refs == ["first"] + assert released_runs == [] + assert handle not in worker._accepted_run_handles + assert handle._cleanup_published + assert handle._keepalive == (keepalive, first_ref, second_ref) + assert worker._abandoned_run_handles == [handle] + assert isinstance(worker._ordered_cleanup_error, RuntimeError) + assert worker._ordered_cleanup_error.__cause__ is interrupt + + worker.close() + + assert teardown_calls == ["teardown"] + assert handle._keepalive is None + assert not worker._abandoned_run_handles + + def test_post_teardown_keepalive_drain_survives_an_interrupt(self, monkeypatch): + interrupt = KeyboardInterrupt("keepalive release") + nested_interrupt = SystemExit("keepalive traversal back-edge") + teardown_calls: list[str] = [] + + class InterruptingHandleList(list): + interrupt_bool = False + + def __bool__(self): + if self.interrupt_bool: + self.interrupt_bool = False + raise nested_interrupt + return len(self) != 0 + + retained_handles = InterruptingHandleList() + + class InterruptingRetainedHandle: + def __init__(self): + self.value = object() + self.interrupted = False + + @property + def _keepalive(self): + return self.value + + @_keepalive.setter + def _keepalive(self, value): + if not self.interrupted: + self.interrupted = True + retained_handles.interrupt_bool = True + raise interrupt + self.value = value + + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + retained = InterruptingRetainedHandle() + retained_handles.append(retained) + worker._abandoned_run_handles = cast(Any, retained_handles) + + def teardown_tree(): + teardown_calls.append("teardown") + worker._worker = None + + monkeypatch.setattr(worker, "_teardown_ready_tree", teardown_tree) + + with pytest.raises(KeyboardInterrupt) as caught: + worker.close() + + assert caught.value is interrupt + assert teardown_calls == ["teardown"] + assert retained._keepalive is None + assert not worker._abandoned_run_handles + + def test_post_teardown_keepalive_drain_precedes_the_residual_probe(self, monkeypatch): + residual_probe_interrupt = KeyboardInterrupt("residual probe") + teardown_calls: list[str] = [] + probe_calls: list[str] = [] + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + keepalive = object() + retained = RunHandle(worker, 1, (keepalive,)) + retained._finalization_abandoned = True + worker._abandoned_run_handles.append(retained) + + def teardown_tree(): + teardown_calls.append("teardown") + worker._worker = None + + def has_live_resources(): + probe_calls.append("probe") + if len(probe_calls) == 1: + return True + raise residual_probe_interrupt + + monkeypatch.setattr(worker, "_teardown_ready_tree", teardown_tree) + monkeypatch.setattr(worker, "_has_live_resources", has_live_resources) + + with pytest.raises(KeyboardInterrupt) as caught: + worker.close() + + assert caught.value is residual_probe_interrupt + assert teardown_calls == ["teardown"] + assert probe_calls == ["probe", "probe"] + assert retained._keepalive is None + assert not worker._abandoned_run_handles + + replayed: list[BaseException] = [] + + def retry_close() -> None: + try: + worker.close() + except BaseException as exc: # noqa: BLE001 + replayed.append(exc) + + retry = threading.Thread(target=retry_close) + retry.start() + retry.join(1.0) + assert not retry.is_alive() + assert replayed == [residual_probe_interrupt] + assert teardown_calls == ["teardown"] + + def test_cancellation_abandonment_survives_nested_publication_interrupts(self): + lifecycle_interrupt = KeyboardInterrupt("lifecycle publication") + terminal_interrupt = SystemExit("terminal publication") + cancellation_error = RuntimeError("cancellation did not settle") + + class OnceInterruptingCV: + def __init__(self, cv): + self.cv = cv + self.interrupted = False + + def __enter__(self): + if not self.interrupted: + self.interrupted = True + raise lifecycle_interrupt + return self.cv.__enter__() + + def __exit__(self, *exc_info): + return self.cv.__exit__(*exc_info) + + def notify_all(self): + self.cv.notify_all() + + class InterruptingHandle(RunHandle): + interrupt_terminal = False + + def __setattr__(self, name, value): + if name == "_terminal" and value is True and self.interrupt_terminal: + self.interrupt_terminal = False + raise terminal_interrupt + return super().__setattr__(name, value) + + worker = Worker(level=3, num_sub_workers=0) + handle = InterruptingHandle(worker, 1, (object(),)) + handle.interrupt_terminal = True + worker._accepted_run_handles.add(handle) + worker._hierarchical_start_cv = cast(Any, OnceInterruptingCV(worker._hierarchical_start_cv)) + + worker._abandon_unsettled_run(handle, str(cancellation_error), cancellation_error) + + assert handle._terminal + assert not handle._wait_in_progress + assert not handle._accept_wait_in_progress + assert isinstance(handle._error, RuntimeError) + assert str(handle._error) == str(cancellation_error) + assert handle._error.__cause__ is cancellation_error + assert worker._ordered_cleanup_error is handle._error + assert handle._keepalive is not None + assert handle not in worker._accepted_run_handles + assert worker._abandoned_run_handles == [handle] + + def test_run_finalization_releases_only_its_resources(self, monkeypatch): + class SlotRef: + def __init__(self): + self.releases = 0 + + def _release_slot_ref(self): + self.releases += 1 + + class Region: + def __init__(self, region_id): + self.region_id = region_id + self._worker_id = 0 + self.mapping_closed = False + self.expired = False + + def _close_l3_host_mapping(self): + self.mapping_closed = True + + def _expire(self): + self.expired = True + + class NativeWorker: + def __init__(self): + self.released_regions = [] + + def control_l3_l2_region_release(self, worker_id, region_id): + self.released_regions.append((worker_id, region_id)) + + class NativeOrchestrator: + def __init__(self): + self.released_runs = [] + + def _release_run(self, run_id): + self.released_runs.append(run_id) + + def domain(name, allocation_id): + return worker_mod.CommDomainHandle( + name=name, + workers=(), + contexts={}, + allocation_id=allocation_id, + _release_fn=lambda _handle: None, + ) + + worker = Worker(level=3, num_sub_workers=0) + native_worker = NativeWorker() + native_orch = NativeOrchestrator() + worker._worker = cast(Any, native_worker) + worker._orch = cast(Any, native_orch) + + first = worker_mod._RunResources() + second = worker_mod._RunResources() + first_ref, second_ref = SlotRef(), SlotRef() + first_region, second_region = Region(11), Region(22) + first_live, second_live = domain("first-live", 1), domain("second-live", 2) + first_pending, second_pending = domain("first-pending", 3), domain("second-pending", 4) + first.remote_slot_refs.append(cast(Any, first_ref)) + second.remote_slot_refs.append(cast(Any, second_ref)) + first.l3_l2_regions.append(first_region) + second.l3_l2_regions.append(second_region) + first.l3_l2_orch_comm_host_buffers[0x1000] = 64 + second.l3_l2_orch_comm_host_buffers[0x2000] = 128 + first.live_domains[first_live.name] = first_live + second.live_domains[second_live.name] = second_live + first.pending_release_domains.append(first_pending) + second.pending_release_domains.append(second_pending) + worker._live_l3_l2_regions.extend([first_region, second_region]) + worker._live_domains.update({first_live.name: first_live, second_live.name: second_live}) + + released_domains = [] + + def release_domain_now(handle): + released_domains.append(handle) + if worker._live_domains.get(handle.name) is handle: + worker._live_domains.pop(handle.name) + + monkeypatch.setattr(worker, "_release_domain_now", release_domain_now) + first_handle = RunHandle(worker, 1, (), first) + second_handle = RunHandle(worker, 2, (), second) + worker._accepted_run_handles.update({first_handle, second_handle}) + + assert worker._finalize_run_handle(first_handle, 1, None) is None + + assert first_ref.releases == 1 + assert second_ref.releases == 0 + assert native_worker.released_regions == [(0, 11)] + assert first_region.mapping_closed and first_region.expired + assert not second_region.mapping_closed and not second_region.expired + assert worker._live_l3_l2_regions == [second_region] + assert first.l3_l2_orch_comm_host_buffers == {} + assert second.l3_l2_orch_comm_host_buffers == {0x2000: 128} + assert released_domains == [first_pending, first_live] + assert first_pending.freed and first_live.freed + assert not second_pending.freed and not second_live.freed + assert worker._live_domains == {second_live.name: second_live} + assert native_orch.released_runs == [1] + assert worker._accepted_run_handles == {second_handle} + + def test_region_mapping_failure_still_releases_child_and_poisons_successor(self): + mapping_error = KeyboardInterrupt("mapping close") + release_error = SystemExit("child release") + + class Region: + region_id = 11 + _worker_id = 0 + + def __init__(self): + self.expired = False + + def _close_l3_host_mapping(self): + raise mapping_error + + def _expire(self): + self.expired = True + + class NativeWorker: + def __init__(self): + self.released_regions = [] + + def control_l3_l2_region_release(self, worker_id, region_id): + self.released_regions.append((worker_id, region_id)) + raise release_error + + class NativeOrchestrator: + def _release_run(self, run_id): + raise AssertionError(f"ambiguous run {run_id} must stay owned until tree teardown") + + worker = Worker(level=3, num_sub_workers=0) + native_worker = NativeWorker() + worker._worker = cast(Any, native_worker) + worker._orch = cast(Any, NativeOrchestrator()) + resources = worker_mod._RunResources() + resources.requires_ordered_cleanup = True + region = Region() + resources.l3_l2_regions.append(region) + worker._live_l3_l2_regions.append(region) + handle = RunHandle(worker, 1, (), resources) + worker._accepted_run_handles.add(handle) + + assert worker._finalize_run_handle(handle, 1, None) is mapping_error + + assert native_worker.released_regions == [(0, 11)] + assert region.expired + assert resources.l3_l2_regions == [] + assert worker._live_l3_l2_regions == [] + assert worker._ordered_cleanup_error is not None + with pytest.raises(RuntimeError, match="no further work is admitted"): + worker._require_no_ordered_cleanup_failure("submit") + + def test_region_tracking_interrupt_still_retires_every_region(self): + interrupt = KeyboardInterrupt("region tracking removal") + + class InterruptingList(list): + def __init__(self, values): + super().__init__(values) + self.interrupted = False + + def __setitem__(self, key, value): + super().__setitem__(key, value) + if isinstance(key, slice) and not self.interrupted: + self.interrupted = True + raise interrupt + + class Region: + def __init__(self, region_id): + self.region_id = region_id + self._worker_id = 0 + self.mapping_closes = 0 + self.expires = 0 + + def _close_l3_host_mapping(self): + self.mapping_closes += 1 + + def _expire(self): + self.expires += 1 + + class NativeWorker: + def __init__(self): + self.released_regions = [] + + def control_l3_l2_region_release(self, worker_id, region_id): + self.released_regions.append((worker_id, region_id)) + + first, second = Region(11), Region(22) + resources = worker_mod._RunResources() + resources.l3_l2_regions = cast(Any, InterruptingList([first, second])) + worker = Worker(level=3, num_sub_workers=0) + native_worker = NativeWorker() + worker._worker = cast(Any, native_worker) + worker._live_l3_l2_regions.extend([first, second]) + + with pytest.raises(KeyboardInterrupt) as caught: + worker._cleanup_l3_l2_regions(resources) + + assert caught.value is interrupt + assert native_worker.released_regions == [(0, 11), (0, 22)] + assert (first.mapping_closes, first.expires) == (1, 1) + assert (second.mapping_closes, second.expires) == (1, 1) + assert resources.l3_l2_regions == [] + assert worker._live_l3_l2_regions == [] + + def test_domain_released_after_its_run_retired_is_freed_inline(self): + """A late release has no fence left to defer behind, so it frees now. + + Both deferred paths are closed to it: the run's queue is never drained + again, and _release_domain_handle has already dropped the handle from + _live_domains, so close()'s live sweep cannot reach it either. + """ + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + worker._orch = cast(Any, type("FakeOrch", (), {"_release_run": lambda self, run_id: None})()) + + freed: list[str] = [] + + def release_now(handle): + freed.append(handle.name) + if worker._live_domains.get(handle.name) is handle: + worker._live_domains.pop(handle.name) + + worker._release_domain_now = cast(Any, release_now) + + resources = worker_mod._RunResources() + late = worker_mod.CommDomainHandle( + name="late", + workers=(), + contexts={}, + allocation_id=1, + _release_fn=lambda released, owner=resources: worker._release_domain_handle(released, owner), + ) + resources.live_domains[late.name] = late + worker._live_domains[late.name] = late + + handle = RunHandle(worker, 1, (), resources) + worker._accepted_run_handles.add(handle) + assert worker._finalize_run_handle(handle, 1, None) is None + assert freed == ["late"], "a domain still live at its run's fence is swept there" + + second = worker_mod.CommDomainHandle( + name="later", + workers=(), + contexts={}, + allocation_id=2, + _release_fn=lambda released, owner=resources: worker._release_domain_handle(released, owner), + ) + resources.live_domains[second.name] = second + worker._live_domains[second.name] = second + + second.release() + + assert freed == ["late", "later"] + assert second.freed + assert resources.pending_release_domains == [] + + def test_interrupted_domain_queue_publication_keeps_live_owners(self): + interrupt = KeyboardInterrupt("pending-domain publication") + + class InterruptingList(list): + def append(self, value): + raise interrupt + + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + resources = worker_mod._RunResources() + resources.pending_release_domains = cast(Any, InterruptingList()) + handle = worker_mod.CommDomainHandle( + name="owned", + workers=(), + contexts={}, + allocation_id=1, + _release_fn=lambda released, owner=resources: worker._release_domain_handle(released, owner), + ) + resources.live_domains[handle.name] = handle + worker._live_domains[handle.name] = handle + + with pytest.raises(KeyboardInterrupt) as caught: + handle.release() + + assert caught.value is interrupt + assert handle.released + assert resources.live_domains == {handle.name: handle} + assert worker._live_domains == {handle.name: handle} + assert resources.pending_release_domains == [] + + def test_interrupted_post_fence_domain_free_keeps_live_owners(self): + interrupt = KeyboardInterrupt("post-fence domain free") + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + resources = worker_mod._RunResources(retired=True) + handle = worker_mod.CommDomainHandle( + name="owned", + workers=(), + contexts={}, + allocation_id=1, + _release_fn=lambda released, owner=resources: worker._release_domain_handle(released, owner), + ) + resources.live_domains[handle.name] = handle + worker._live_domains[handle.name] = handle + worker._free_domain_after_fence = cast(Any, lambda _handle: (_ for _ in ()).throw(interrupt)) + + with pytest.raises(KeyboardInterrupt) as caught: + handle.release() + + assert caught.value is interrupt + assert handle.released + assert resources.live_domains == {handle.name: handle} + assert worker._live_domains == {handle.name: handle} + assert resources.pending_release_domains == [] + + def test_pending_domain_drain_does_not_clear_unclaimed_owners(self): + class NoClearList(list): + def clear(self): + raise AssertionError("pending owners must stay published until backend success") + + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + resources = worker_mod._RunResources() + handle = worker_mod.CommDomainHandle( + name="pending", + workers=(), + contexts={}, + allocation_id=1, + _release_fn=lambda released: None, + ) + handle._released = True + resources.pending_release_domains = cast(Any, NoClearList([handle])) + freed = [] + + def free_domain(pending): + freed.append(pending) + pending._freed = True + + worker._free_domain_after_fence = cast(Any, free_domain) + + worker._execute_pending_domain_releases(resources) + + assert freed == [handle] + assert resources.pending_release_domains == [] + + def test_failed_pending_domain_drain_retains_its_claim(self): + interrupt = KeyboardInterrupt("pending domain free") + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + resources = worker_mod._RunResources() + handle = worker_mod.CommDomainHandle( + name="pending", + workers=(), + contexts={}, + allocation_id=1, + _release_fn=lambda released: None, + ) + handle._released = True + resources.pending_release_domains.append(handle) + worker._free_domain_after_fence = cast(Any, lambda _handle: (_ for _ in ()).throw(interrupt)) + + with pytest.raises(KeyboardInterrupt) as caught: + worker._execute_pending_domain_releases(resources) + + assert caught.value is interrupt + assert resources.pending_release_domains == [handle] + + def test_domain_free_outcome_precedes_caller_interrupt(self, monkeypatch): + interrupt = KeyboardInterrupt("after isolated domain free") + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + handle = worker_mod.CommDomainHandle( + name="domain", + workers=(), + contexts={}, + allocation_id=7, + _release_fn=lambda released: None, + ) + backend_calls = [] + worker._release_domain_claimed = cast(Any, lambda claimed: backend_calls.append(claimed.allocation_id)) + + def interrupt_after_target(items, target, **kwargs): + for item in items: + target(item) + raise interrupt + + monkeypatch.setattr(worker_mod, "_start_and_join_threads", interrupt_after_target) + + with pytest.raises(KeyboardInterrupt) as caught: + worker._release_domain_now(handle) + + assert caught.value is interrupt + assert backend_calls == [7] + assert worker._domain_free_results == {7: None} + + worker._release_domain_now(handle) + assert backend_calls == [7] + + def test_domain_free_interrupted_before_isolated_admission_is_retryable(self, monkeypatch): + interrupt = KeyboardInterrupt("before isolated domain free admission") + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + handle = worker_mod.CommDomainHandle( + name="domain", + workers=(), + contexts={}, + allocation_id=7, + _release_fn=lambda released: None, + ) + backend_calls: list[int] = [] + worker._release_domain_claimed = cast(Any, lambda claimed: backend_calls.append(claimed.allocation_id)) + real_isolated_call = worker_mod._run_isolated_call + monkeypatch.setattr( + worker_mod, + "_run_isolated_call", + lambda *_args, **_kwargs: (_ for _ in ()).throw(interrupt), + ) + + with pytest.raises(KeyboardInterrupt) as caught: + worker._release_domain_now(handle) + + assert caught.value is interrupt + assert backend_calls == [] + assert 7 not in worker._domain_free_results + + monkeypatch.setattr(worker_mod, "_run_isolated_call", real_isolated_call) + worker._release_domain_now(handle) + + assert backend_calls == [7] + assert worker._domain_free_results == {7: None} + + def test_domain_free_publication_failure_never_replays_backend(self): + publication_error = MemoryError("domain outcome store") + + class InterruptingResults(dict): + def __init__(self): + super().__init__() + self.interrupted = False + + def __setitem__(self, key, value): + if value is None and not self.interrupted: + self.interrupted = True + raise publication_error + super().__setitem__(key, value) + + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + worker._domain_free_results = cast(Any, InterruptingResults()) + handle = worker_mod.CommDomainHandle( + name="domain", + workers=(), + contexts={}, + allocation_id=7, + _release_fn=lambda released: None, + ) + backend_calls = [] + worker._release_domain_claimed = cast(Any, lambda claimed: backend_calls.append(claimed.allocation_id)) + + with pytest.raises(MemoryError) as caught: + worker._release_domain_now(handle) + + assert caught.value is publication_error + assert backend_calls == [7] + assert isinstance(worker._domain_free_results[7], RuntimeError) + + with pytest.raises(RuntimeError, match="refusing to replay"): + worker._release_domain_now(handle) + assert backend_calls == [7] + + def test_domain_released_while_its_run_retires_is_still_freed(self): + """A release in flight across retirement must not be stranded. + + The window the fence has to close: the deferred queue has already been + drained for the last time, so a handle appended after that is reachable + from nothing — the queue is never read again, and the release itself + popped the handle from ``_live_domains``, blinding ``close()``'s sweep. + + Forced deterministically by parking the fence in its live-domain sweep, + which holds no lock, while the releasing thread runs to completion. + """ + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + worker._orch = cast(Any, type("FakeOrch", (), {"_release_run": lambda self, run_id: None})()) + + freed: list[str] = [] + + def release_now(handle): + freed.append(handle.name) + if worker._live_domains.get(handle.name) is handle: + worker._live_domains.pop(handle.name) + + worker._release_domain_now = cast(Any, release_now) + + resources = worker_mod._RunResources() + + def domain(name, allocation_id): + return worker_mod.CommDomainHandle( + name=name, + workers=(), + contexts={}, + allocation_id=allocation_id, + _release_fn=lambda released, owner=resources: worker._release_domain_handle(released, owner), + ) + + raced = domain("raced", 1) + # A second domain keeps the sweep branch reachable so the fence parks there. + swept = domain("swept", 2) + for handle in (raced, swept): + resources.live_domains[handle.name] = handle + worker._live_domains[handle.name] = handle + + in_sweep = threading.Event() + release_done = threading.Event() + real_sweep = worker._release_all_live_domains + + def parked_sweep(res=None): + in_sweep.set() + assert release_done.wait(5.0), "releasing thread did not finish" + real_sweep(res) + + worker._release_all_live_domains = cast(Any, parked_sweep) + + def do_release(): + assert in_sweep.wait(5.0), "fence never reached its sweep" + raced.release() + release_done.set() + + releaser = threading.Thread(target=do_release) + releaser.start() + try: + handle = RunHandle(worker, 1, (), resources) + worker._accepted_run_handles.add(handle) + assert worker._finalize_run_handle(handle, 1, None) is None + finally: + releaser.join(5.0) + + assert not releaser.is_alive() + assert raced.freed, "a release racing retirement was stranded on a drained queue" + assert swept.freed + assert resources.pending_release_domains == [] + assert freed.count("raced") == 1, f"raced domain freed more than once: {freed}" + + @staticmethod + def _gated_domain_worker(worker, target_id, outcome=None): + """Park `worker`'s backend release for `target_id`; optionally raise.""" + entered: list[int] = [] + in_backend = threading.Event() + let_go = threading.Event() + + def gated(handle): + entered.append(handle.allocation_id) + if handle.allocation_id == target_id: + in_backend.set() + let_go.wait(10.0) + if outcome is not None: + raise outcome + + worker._release_domain_claimed = cast(Any, gated) + return entered, in_backend, let_go + + def _retired_domain(self, worker, resources, name, allocation_id): + handle = worker_mod.CommDomainHandle( + name=name, + workers=(), + contexts={}, + allocation_id=allocation_id, + _release_fn=lambda released, owner=resources: worker._release_domain_handle(released, owner), + ) + worker._live_domains[name] = handle + return handle + + def test_sweep_and_post_fence_release_free_a_domain_once(self): + """Two paths reaching one handle must not both drive the backend free. + + The dangerous order is release-then-sweep: ``release()`` wins the + ``_released`` flag, so the sweep — still holding the handle in the + snapshot it took earlier — skips setting that flag and goes straight to + the backend call the release is already making. (Sweep-then-release is + already safe: the sweep sets ``_released`` before freeing, which makes + the later ``release()`` a no-op.) + """ + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + entered, in_backend, let_go = self._gated_domain_worker(worker, target_id=7) + + resources = worker_mod._RunResources() + resources.retired = True # the owning run's fence has already passed + contested = self._retired_domain(worker, resources, "contested", 7) + + releaser = threading.Thread(target=contested.release) + releaser.start() + try: + assert in_backend.wait(5.0), "release never reached the backend" + # The sweep's snapshot still holds `contested`; it must not free it + # a second time. + sweeper = threading.Thread(target=worker._release_all_live_domains) + sweeper.start() + let_go.set() + sweeper.join(5.0) + assert not sweeper.is_alive() + finally: + let_go.set() + releaser.join(5.0) + + assert entered.count(7) == 1, f"allocation 7 was released {entered.count(7)} times" + assert contested.freed + + def test_second_domain_release_waits_for_the_first(self): + """A second caller blocks until the owner's backend call returns. + + Returning early would let the caller mark the handle freed, drop it + from ``_live_domains`` and — on the ``close()`` path — tear down the + mailboxes the in-flight release is still using, which is exactly what + ``close()`` orders its domain sweep before ``_worker.close()`` to avoid. + """ + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + entered, in_backend, let_go = self._gated_domain_worker(worker, target_id=7) + + resources = worker_mod._RunResources() + contested = self._retired_domain(worker, resources, "contested", 7) + + owner = threading.Thread(target=worker._release_domain_now, args=(contested,)) + second_done = threading.Event() + + def second_caller(): + worker._release_domain_now(contested) + second_done.set() + + second = threading.Thread(target=second_caller) + owner.start() + try: + assert in_backend.wait(5.0), "owner never reached the backend" + assert not contested.freed, "freed must stay false while the backend call is in flight" + second.start() + assert not second_done.wait(0.5), "the second caller returned while the owner was still releasing" + assert not contested.freed + finally: + let_go.set() + owner.join(5.0) + second.join(5.0) + + assert second_done.is_set() + assert entered.count(7) == 1, f"allocation 7 reached the backend {entered.count(7)} times" + + def test_failed_domain_release_is_replayed_to_a_second_caller(self): + """The owner's failure reaches every later caller, so no path reports + success for an allocation whose backend release did not happen. + + ``_release_all_live_domains`` keeps an un-freed handle in + ``_live_domains`` precisely so ``close()`` reports it as a residual; a + second caller that returned success would erase that. + """ + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + boom = RuntimeError("backend release failed") + _entered, in_backend, let_go = self._gated_domain_worker(worker, target_id=7, outcome=boom) + let_go.set() + + resources = worker_mod._RunResources() + contested = self._retired_domain(worker, resources, "contested", 7) + + with pytest.raises(RuntimeError, match="backend release failed"): + worker._release_domain_now(contested) + assert in_backend.is_set() + assert not contested.freed + + # The sweep is the second caller: it must see the failure, keep the + # 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" + + def test_allocate_domain_outside_graph_construction_is_rejected(self): + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + assert worker._building_run_resources is None + + with pytest.raises(RuntimeError, match="graph is being built"): + 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_caught_submit_validation_failure_does_not_forbid_later_control(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): + with pytest.raises(TypeError, match="expects a CallableHandle"): + orch.submit_sub(object()) + orch.malloc(0, 64) + + native.submit_sub.assert_not_called() + native.malloc.assert_called_once() + + def test_native_submit_attempt_forbids_later_control_even_when_it_raises(self, monkeypatch): + worker = Worker(level=3, num_sub_workers=0) + worker._worker = cast(Any, object()) + orch, native = self._orch(worker) + monkeypatch.setattr( + orch_mod, + "_require_handle", + lambda *_args, **_kwargs: (b"digest", "python", "LOCAL_PYTHON", ()), + ) + native.submit_sub.side_effect = RuntimeError("native submit failed") + + with orch_mod._callback_run(7, worker): + with pytest.raises(RuntimeError, match="native submit failed"): + orch.submit_sub(object()) + with pytest.raises(RuntimeError, match="cannot follow a task submission"): + orch.malloc(0, 64) + + native.submit_sub.assert_called_once() + native.malloc.assert_not_called() + + 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() + + def test_interrupted_import_release_attempts_the_rest_without_losing_debt(self, monkeypatch): + worker = Worker(level=4, num_sub_workers=0) + interrupt = KeyboardInterrupt("release import") + + def imported(buffer_id, import_id): + owner = worker_mod.RemoteBufferHandle._from_remote_allocation( + worker_id=0, + buffer_id=buffer_id, + generation=1, + address_space=worker_mod.RemoteAddressSpace.REMOTE_DEVICE, + nbytes=4, + ) + owner._acquire_import_ref() + handle = worker_mod.RemoteBufferHandle._from_imported_mapping( + worker_id=1, + owner_worker_id=0, + buffer_id=buffer_id, + generation=1, + import_id=import_id, + address_space=worker_mod.RemoteAddressSpace.REMOTE_WINDOW, + nbytes=4, + offset=0, + owner_handle_ref=owner, + ) + return handle, owner + + interrupted, interrupted_owner = imported(1, 11) + released, released_owner = imported(2, 22) + blocked, blocked_owner = imported(3, 33) + blocked._acquire_slot_ref() + worker._pending_remote_import_releases.extend([interrupted, released, blocked]) + calls = [] + + def release_import(handle): + calls.append(handle.import_id) + if handle is interrupted: + raise interrupt + + monkeypatch.setattr(worker, "_send_remote_release_import", release_import) + + with pytest.raises(KeyboardInterrupt) as caught: + worker._flush_pending_remote_frees() + + assert caught.value is interrupt + assert calls == [11, 22] + assert worker._pending_remote_import_releases == [interrupted, blocked] + assert interrupted._owner_handle_ref is interrupted_owner + assert interrupted_owner._live_import_refs == 1 + assert released._owner_handle_ref is None + assert released_owner._live_import_refs == 0 + assert blocked._owner_handle_ref is blocked_owner + assert blocked_owner._live_import_refs == 1 + + def test_post_rpc_import_retirement_interrupt_never_replays_the_debt(self, monkeypatch): + worker = Worker(level=4, num_sub_workers=0) + interrupt = KeyboardInterrupt("owner reference retirement") + + class Owner: + def __init__(self, *, interrupts=False): + self.interrupts = interrupts + self.release_calls = 0 + + def _release_import_ref(self): + self.release_calls += 1 + if self.interrupts: + raise interrupt + + def imported(buffer_id, import_id, owner): + return worker_mod.RemoteBufferHandle._from_imported_mapping( + worker_id=1, + owner_worker_id=0, + buffer_id=buffer_id, + generation=1, + import_id=import_id, + address_space=worker_mod.RemoteAddressSpace.REMOTE_WINDOW, + nbytes=4, + offset=0, + owner_handle_ref=cast(Any, owner), + ) + + interrupted_owner = Owner(interrupts=True) + released_owner = Owner() + interrupted = imported(1, 11, interrupted_owner) + released = imported(2, 22, released_owner) + worker._pending_remote_import_releases.extend([interrupted, released]) + rpc_calls = [] + monkeypatch.setattr(worker, "_send_remote_release_import", lambda handle: rpc_calls.append(handle.import_id)) + + with pytest.raises(KeyboardInterrupt) as caught: + worker._flush_pending_remote_frees() + + assert caught.value is interrupt + assert rpc_calls == [11, 22] + assert interrupted_owner.release_calls == 1 + assert released_owner.release_calls == 1 + assert worker._pending_remote_import_releases == [interrupted] + + with pytest.raises(KeyboardInterrupt) as repeated: + worker._flush_pending_remote_frees() + + assert repeated.value is interrupt + assert rpc_calls == [11, 22] + assert interrupted_owner.release_calls == 1 + assert released_owner.release_calls == 1 + assert worker._pending_remote_import_releases == [interrupted] + + def test_interrupted_owner_free_attempts_the_rest_without_losing_debt(self, monkeypatch): + worker = Worker(level=4, num_sub_workers=0) + interrupt = SystemExit("remote free") + + def owner(buffer_id): + return worker_mod.RemoteBufferHandle._from_remote_allocation( + worker_id=0, + buffer_id=buffer_id, + generation=1, + address_space=worker_mod.RemoteAddressSpace.REMOTE_DEVICE, + nbytes=4, + ) + + interrupted = owner(1) + freed = owner(2) + blocked = owner(3) + blocked._acquire_import_ref() + worker._pending_remote_buffer_frees.extend([interrupted, freed, blocked]) + calls = [] + + def remote_free(handle): + calls.append(handle._buffer_id) + if handle is interrupted: + raise interrupt + + monkeypatch.setattr(worker, "_send_remote_free", remote_free) + + with pytest.raises(SystemExit) as caught: + worker._flush_pending_remote_frees() + + assert caught.value is interrupt + assert calls == [1, 2] + assert worker._pending_remote_buffer_frees == [interrupted, blocked] + + +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_keeps_its_original_release_ranks(self, monkeypatch): + """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() + worker._config = {"device_ids": [0, 1, 2]} + resources = worker_mod._RunResources() + worker._building_run_resources = resources + monkeypatch.setattr(worker, "_ensure_comm_base", lambda: None) + + def partial_failure(*, reply_shms, **_kwargs): + assert reply_shms is not None + for chip_idx in (0, 2): + reply_buf = reply_shms[chip_idx].buf + assert reply_buf is not None + struct.pack_into(" None: + if rank == 0: + rank_zero_entered.set() + assert allow_rank_zero.wait(5.0) + calls.append(rank) + if rank == 1: + rank_one_completed.set() + + def interrupt_after_first_start(rank: int) -> None: + hook_calls.append(rank) + if rank == 0: + raise first_interrupt + + def run_fanout() -> None: + try: + worker_mod._start_and_join_threads( + (0, 1), + target, + name_prefix="test_post_launch_", + _after_start=interrupt_after_first_start, + ) + except BaseException as exc: # noqa: BLE001 + raised.append(exc) + finally: + runner_completed.set() + + runner = threading.Thread(target=run_fanout) + try: + runner.start() + assert rank_zero_entered.wait(5.0) + assert rank_one_completed.wait(5.0), "the later rank was not launched after the boundary interrupt" + assert not runner_completed.wait(0.1), "fanout returned while rank zero still owned caller state" + allow_rank_zero.set() + runner.join(5.0) + + assert not runner.is_alive() + assert sorted(calls) == [0, 1] + assert hook_calls == [0, 1] + assert raised == [first_interrupt] + finally: + allow_rank_zero.set() + runner.join(5.0) + + def test_fanout_defers_a_phase_advance_interrupt_until_every_target_finishes(self): + rank_zero_entered = threading.Event() + rank_one_completed = threading.Event() + allow_rank_zero = threading.Event() + runner_completed = threading.Event() + first_interrupt = KeyboardInterrupt("phase advance") + phase_advances: list[int] = [] + raised: list[BaseException] = [] + + def target(rank: int) -> None: + if rank == 0: + rank_zero_entered.set() + assert allow_rank_zero.wait(5.0) + if rank == 1: + rank_one_completed.set() + + def interrupt_after_launch_phase(phase: int) -> None: + phase_advances.append(phase) + if phase == 0: + raise first_interrupt + + def run_fanout() -> None: + try: + worker_mod._start_and_join_threads( + (0, 1), + target, + name_prefix="test_phase_advance_", + _after_phase=interrupt_after_launch_phase, + ) + except BaseException as exc: # noqa: BLE001 + raised.append(exc) + finally: + runner_completed.set() + + runner = threading.Thread(target=run_fanout) + try: + runner.start() + assert rank_zero_entered.wait(5.0) + assert rank_one_completed.wait(5.0) + assert not runner_completed.wait(0.1), "fanout returned before the launched target released caller state" + allow_rank_zero.set() + runner.join(5.0) + + assert not runner.is_alive() + assert phase_advances == [0, 1, 2, 3] + assert len(raised) == 1 and raised[0] is first_interrupt + finally: + allow_rank_zero.set() + runner.join(5.0) + + def test_domain_fanout_cancels_and_retries_an_ambiguously_launched_thread(self, monkeypatch): + worker = self._worker() + rank_zero_entered = threading.Event() + rank_one_completed = threading.Event() + allow_rank_zero = threading.Event() + calls: list[int] = [] + + def release(chip_idx, _request_name): + if chip_idx == 0: + rank_zero_entered.set() + assert allow_rank_zero.wait(5.0) + calls.append(chip_idx) + if chip_idx == 1: + rank_one_completed.set() + + worker._worker = cast(Any, SimpleNamespace(control_release_domain=release)) + requests = {chip_idx: SharedMemory(create=True, size=worker_mod._DOMAIN_REQ_HEADER.size) for chip_idx in (0, 1)} + real_thread = threading.Thread + backing_threads: list[threading.Thread] = [] + + class AmbiguousStartInterrupt(real_thread): + armed = True + + def start(self): + if AmbiguousStartInterrupt.armed: + AmbiguousStartInterrupt.armed = False + backing = real_thread(target=self.run) + backing_threads.append(backing) + backing.start() + raise KeyboardInterrupt("ambiguous start") + return super().start() + + monkeypatch.setattr(worker_mod.threading, "Thread", AmbiguousStartInterrupt) + raised: list[BaseException] = [] + runner_completed = threading.Event() + + def run_dispatch(): + try: + worker._dispatch_control_domain( + workers=(0, 1), + request_shms=requests, + reply_shms=None, + op="release", + allocation_id=7, + ) + except BaseException as exc: # noqa: BLE001 + raised.append(exc) + finally: + runner_completed.set() + + runner = real_thread(target=run_dispatch) + try: + runner.start() + assert rank_one_completed.wait(5.0), "a later collective rank was not launched" + assert rank_zero_entered.wait(5.0), "the interrupted rank was not retried" + assert not runner_completed.wait(0.1), "fanout returned while the retried rank still owned the request shm" + allow_rank_zero.set() + runner.join(5.0) + + assert not runner.is_alive() + assert sorted(calls) == [0, 1] + assert len(raised) == 1 and isinstance(raised[0], KeyboardInterrupt) + finally: + allow_rank_zero.set() + runner.join(5.0) + for backing in backing_threads: + backing.join(5.0) + for request in requests.values(): + request.close() + request.unlink() + + def test_domain_fanout_joins_through_repeated_interruptions(self, monkeypatch): + worker = self._worker() + entered = threading.Event() + allow_finish = threading.Event() + completed: list[int] = [] + + def release(chip_idx, _request_name): + entered.set() + assert allow_finish.wait(5.0) + completed.append(chip_idx) + + worker._worker = cast(Any, SimpleNamespace(control_release_domain=release)) + request = SharedMemory(create=True, size=worker_mod._DOMAIN_REQ_HEADER.size) + real_thread = threading.Thread + instances = [] + + class RepeatedJoinInterrupt(real_thread): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.join_attempts = 0 + instances.append(self) + + def join(self, *args, **kwargs): + self.join_attempts += 1 + if self.join_attempts <= 3: + raise KeyboardInterrupt + return super().join(*args, **kwargs) + + monkeypatch.setattr(worker_mod.threading, "Thread", RepeatedJoinInterrupt) + raised: list[BaseException] = [] + + def run_dispatch(): + try: + worker._dispatch_control_domain( + workers=(0,), + request_shms={0: request}, + reply_shms=None, + op="release", + allocation_id=7, + ) + except BaseException as exc: # noqa: BLE001 + raised.append(exc) + + runner = real_thread(target=run_dispatch) + try: + runner.start() + assert entered.wait(5.0) + assert runner.is_alive(), "fanout returned while its target still owned the request shm" + allow_finish.set() + runner.join(5.0) + + assert not runner.is_alive() + assert completed == [0] + assert instances[0].join_attempts == 4 + assert len(raised) == 1 and isinstance(raised[0], KeyboardInterrupt) + finally: + allow_finish.set() + runner.join(5.0) + for instance in instances: + if instance.ident is not None: + real_thread.join(instance, 5.0) + request.close() + request.unlink() + + def test_comm_init_launches_later_ranks_after_a_start_boundary_interrupt(self, monkeypatch): + worker = self._worker() + worker._config = {"device_ids": [0, 1]} + monkeypatch.setattr(worker, "_comm_plan_rootinfo_path", lambda: "/tmp/comm-rootinfo") + calls: list[int] = [] + worker._worker = cast( + Any, + SimpleNamespace(control_comm_init=lambda chip_idx, _request_name: calls.append(chip_idx)), + ) + real_thread = threading.Thread + + class StartBoundaryInterrupt(real_thread): + armed = True + + def start(self): + super().start() + if StartBoundaryInterrupt.armed and not self.name.startswith("shm-"): + StartBoundaryInterrupt.armed = False + raise KeyboardInterrupt + + monkeypatch.setattr(worker_mod.threading, "Thread", StartBoundaryInterrupt) + + with pytest.raises(KeyboardInterrupt): + worker._ensure_comm_base() + + assert sorted(calls) == [0, 1] + assert not worker._comm_base_ready + + def test_comm_init_preserves_the_first_interrupt_through_shm_cleanup(self, monkeypatch): + worker = self._worker() + worker._config = {"device_ids": [0]} + monkeypatch.setattr(worker, "_comm_plan_rootinfo_path", lambda: "/tmp/comm-rootinfo") + worker._worker = cast(Any, SimpleNamespace(control_comm_init=lambda _chip_idx, _request_name: None)) + + created: list[SharedMemory] = [] + created_fds: list[int] = [] + real_init = SharedMemory.__init__ + + def tracked_shared_memory(shm, *args, **kwargs): + real_init(shm, *args, **kwargs) + created.append(shm) + created_fds.append(shm._fd) + + monkeypatch.setattr(SharedMemory, "__init__", tracked_shared_memory) + first_interrupt = KeyboardInterrupt("start boundary") + real_thread = threading.Thread + + class StartBoundaryInterrupt(real_thread): + armed = True + + def start(self): + super().start() + if StartBoundaryInterrupt.armed: + StartBoundaryInterrupt.armed = False + raise first_interrupt + + monkeypatch.setattr(worker_mod.threading, "Thread", StartBoundaryInterrupt) + real_close = worker_mod.os.close + close_interrupted = False + + def interrupt_first_close(fd): + nonlocal close_interrupted + if created_fds and fd == created_fds[0] and not close_interrupted: + close_interrupted = True + raise SystemExit("shm close") + return real_close(fd) + + monkeypatch.setattr(worker_mod.os, "close", interrupt_first_close) + try: + with pytest.raises(KeyboardInterrupt) as caught: + worker._ensure_comm_base() + assert caught.value is first_interrupt + assert close_interrupted + finally: + monkeypatch.setattr(worker_mod.os, "close", real_close) + for shm in created: + try: + shm.close() + shm.unlink() + except FileNotFoundError: + pass + + def test_shm_cleanup_drains_later_owners_after_a_traversal_interrupt(self, monkeypatch): + owner = worker_mod._SharedMemoryOwner(2) + shms = [owner.create(1), owner.create(1)] + names = [shm.name for shm in shms] + first_interrupt = KeyboardInterrupt("between shm owners") + real_unlink = SharedMemory.unlink + interrupted = False + + def interrupt_after_first_unlink(shm): + nonlocal interrupted + real_unlink(shm) + if shm is shms[0] and not interrupted: + interrupted = True + raise first_interrupt + + monkeypatch.setattr(SharedMemory, "unlink", interrupt_after_first_unlink) - freed: list[str] = [] + try: + cleanup_error = worker_mod._close_unlink_shms(owner) - def release_now(handle): - freed.append(handle.name) - if worker._live_domains.get(handle.name) is handle: - worker._live_domains.pop(handle.name) + assert cleanup_error is first_interrupt + for name in names: + with pytest.raises(FileNotFoundError): + SharedMemory(name=name) + finally: + for shm in shms: + try: + shm.close() + shm.unlink() + except FileNotFoundError: + pass + + def test_shm_cleanup_recovers_an_interrupt_during_cursor_setup(self, monkeypatch): + owner = worker_mod._SharedMemoryOwner(1) + shm = owner.create(1) + name = shm.name + interrupt = KeyboardInterrupt("before shm cleanup fanout") + real_cursor = worker_mod._SharedMemoryCleanupCursor + attempts = 0 + + def interrupt_first_cursor(*args, **kwargs): + nonlocal attempts + cursor = real_cursor(*args, **kwargs) + attempts += 1 + if attempts == 1: + raise interrupt + return cursor - worker._release_domain_now = cast(Any, release_now) + monkeypatch.setattr(worker_mod, "_SharedMemoryCleanupCursor", interrupt_first_cursor) + try: + cleanup_error = worker_mod._close_unlink_shms(owner) + + assert cleanup_error is interrupt + assert owner._cleanup_step == worker_mod._SHM_CLEANUP_PHASES + with pytest.raises(FileNotFoundError): + SharedMemory(name=name) + finally: + try: + shm.close() + shm.unlink() + except FileNotFoundError: + pass + @pytest.mark.parametrize("operation", ("comm_init", "domain_alloc", "domain_release")) + def test_shm_owner_retries_an_interrupt_at_cleanup_call_entry(self, monkeypatch, operation): + worker = self._worker() + worker._config = {"device_ids": [0]} + resources = worker_mod._RunResources() + worker._building_run_resources = resources + monkeypatch.setattr(worker, "_comm_plan_rootinfo_path", lambda: "/tmp/comm-rootinfo") + real_ensure_comm_base = worker._ensure_comm_base + + if operation == "comm_init": + worker._worker = cast(Any, SimpleNamespace(control_comm_init=lambda *_args: None)) + elif operation == "domain_alloc": + monkeypatch.setattr(worker, "_ensure_comm_base", lambda: None) + + def commit_domain(*, reply_shms, **_kwargs): + assert reply_shms is not None + reply = reply_shms[0] + assert reply.buf is not None + worker_mod._DOMAIN_REPLY_HEADER.pack_into(reply.buf, 0, 1, 0, 0, 0) + + monkeypatch.setattr(worker, "_dispatch_control_domain", commit_domain) + else: + monkeypatch.setattr(worker, "_dispatch_control_domain", lambda **_kwargs: None) + + real_init = SharedMemory.__init__ + created: list[SharedMemory] = [] + names: list[str] = [] + + def track_created(shm, *args, **kwargs): + real_init(shm, *args, **kwargs) + if kwargs.get("create"): + created.append(shm) + names.append(shm.name) + + monkeypatch.setattr(SharedMemory, "__init__", track_created) + real_cleanup = worker_mod._close_unlink_shms + boundary_interrupt = KeyboardInterrupt(f"{operation} cleanup call entry") + cleanup_calls = 0 + + def interrupt_first_cleanup(*args, **kwargs): + nonlocal cleanup_calls + cleanup_calls += 1 + if cleanup_calls == 1: + raise boundary_interrupt + return real_cleanup(*args, **kwargs) + + monkeypatch.setattr(worker_mod, "_close_unlink_shms", interrupt_first_cleanup) + handle = worker_mod.CommDomainHandle( + name="d", + workers=(0,), + contexts={}, + allocation_id=7, + _release_fn=lambda _released: None, + _domain_ranks={0: 0}, + ) + try: + with pytest.raises(KeyboardInterrupt) as caught: + if operation == "comm_init": + real_ensure_comm_base() + elif operation == "domain_alloc": + worker._allocate_domain(name="d", workers=(0,), window_size=64, buffers=[]) + else: + worker._release_domain_claimed(handle) + + assert caught.value is boundary_interrupt + assert cleanup_calls == 2 + assert names + monkeypatch.setattr(SharedMemory, "__init__", real_init) + for name in names: + with pytest.raises(FileNotFoundError): + SharedMemory(name=name) + finally: + monkeypatch.setattr(SharedMemory, "__init__", real_init) + for shm in created: + try: + shm.close() + shm.unlink() + except FileNotFoundError: + pass + + @pytest.mark.parametrize("operation", ("comm_init", "domain_alloc", "domain_release")) + def test_shm_lifecycle_defers_main_thread_interrupt_until_cleanup(self, monkeypatch, operation): + worker = self._worker() + worker._config = {"device_ids": [0]} resources = worker_mod._RunResources() + worker._building_run_resources = resources + monkeypatch.setattr(worker, "_comm_plan_rootinfo_path", lambda: "/tmp/comm-rootinfo") + real_ensure_comm_base = worker._ensure_comm_base + interrupt_sent = threading.Event() + + def interrupt_main() -> None: + if interrupt_sent.is_set(): + return + interrupt_sent.set() + _thread.interrupt_main() + # Keep the lifecycle active long enough for the main thread to + # observe the signal while the owned name is still live. + time.sleep(0.01) + + if operation == "comm_init": + worker._worker = cast(Any, SimpleNamespace(control_comm_init=lambda *_args: interrupt_main())) + elif operation == "domain_alloc": + monkeypatch.setattr(worker, "_ensure_comm_base", lambda: None) + + def commit_domain(*, reply_shms, **_kwargs): + assert reply_shms is not None + interrupt_main() + reply = reply_shms[0] + assert reply.buf is not None + worker_mod._DOMAIN_REPLY_HEADER.pack_into(reply.buf, 0, 1, 0, 0, 0) + + monkeypatch.setattr(worker, "_dispatch_control_domain", commit_domain) + else: + monkeypatch.setattr(worker, "_dispatch_control_domain", lambda **_kwargs: interrupt_main()) + + real_init = SharedMemory.__init__ + created: list[SharedMemory] = [] + names: list[str] = [] + + def track_created(shm, *args, **kwargs): + real_init(shm, *args, **kwargs) + if kwargs.get("create"): + created.append(shm) + names.append(shm.name) + + monkeypatch.setattr(SharedMemory, "__init__", track_created) + handle = worker_mod.CommDomainHandle( + name="d", + workers=(0,), + contexts={}, + allocation_id=7, + _release_fn=lambda _released: None, + _domain_ranks={0: 0}, + ) + try: + with pytest.raises(KeyboardInterrupt): + if operation == "comm_init": + real_ensure_comm_base() + elif operation == "domain_alloc": + worker._allocate_domain(name="d", workers=(0,), window_size=64, buffers=[]) + else: + worker._release_domain_claimed(handle) + + assert interrupt_sent.is_set() + assert names + monkeypatch.setattr(SharedMemory, "__init__", real_init) + for name in names: + with pytest.raises(FileNotFoundError): + SharedMemory(name=name) + finally: + monkeypatch.setattr(SharedMemory, "__init__", real_init) + for shm in created: + try: + shm.close() + shm.unlink() + except FileNotFoundError: + pass + + def test_shm_cleanup_drains_when_error_recording_is_interrupted(self, monkeypatch): + owner = worker_mod._SharedMemoryOwner(2) + shms = [owner.create(1), owner.create(1)] + names = [shm.name for shm in shms] + first_interrupt = KeyboardInterrupt("unlink return") + recording_interrupt = SystemExit("recording cleanup error") + real_unlink = SharedMemory.unlink + real_remember = worker_mod._remember_cleanup_error + unlink_interrupted = False + recording_interrupted = False + + def interrupt_after_unlink(shm): + nonlocal unlink_interrupted + real_unlink(shm) + if not unlink_interrupted: + unlink_interrupted = True + raise first_interrupt + + def interrupt_error_recording(first_error, cleanup_error): + nonlocal recording_interrupted + if not recording_interrupted: + recording_interrupted = True + raise recording_interrupt + return real_remember(first_error, cleanup_error) + + monkeypatch.setattr(SharedMemory, "unlink", interrupt_after_unlink) + monkeypatch.setattr(worker_mod, "_remember_cleanup_error", interrupt_error_recording) + try: + cleanup_error = worker_mod._close_unlink_shms(owner) - def domain(name, allocation_id): - return worker_mod.CommDomainHandle( - name=name, - workers=(), - contexts={}, - allocation_id=allocation_id, - _release_fn=lambda released, owner=resources: worker._release_domain_handle(released, owner), + assert cleanup_error is first_interrupt + assert recording_interrupted + for name in names: + with pytest.raises(FileNotFoundError): + SharedMemory(name=name) + finally: + monkeypatch.setattr(SharedMemory, "unlink", real_unlink) + monkeypatch.setattr(worker_mod, "_remember_cleanup_error", real_remember) + for shm in shms: + try: + shm.close() + shm.unlink() + except FileNotFoundError: + pass + + def test_shm_cleanup_replays_an_ordinary_unlink_error_after_draining(self, monkeypatch): + owner = worker_mod._SharedMemoryOwner(2) + shms = [owner.create(1), owner.create(1)] + names = [shm.name for shm in shms] + first_error = OSError("resource tracker write failed") + real_unlink = SharedMemory.unlink + injected = False + + def unlink_then_fail(shm): + nonlocal injected + real_unlink(shm) + if not injected: + injected = True + raise first_error + + monkeypatch.setattr(SharedMemory, "unlink", unlink_then_fail) + try: + cleanup_error = worker_mod._close_unlink_shms(owner) + + assert cleanup_error is first_error + for name in names: + with pytest.raises(FileNotFoundError): + SharedMemory(name=name) + finally: + monkeypatch.setattr(SharedMemory, "unlink", real_unlink) + for shm in shms: + try: + shm.close() + shm.unlink() + except FileNotFoundError: + pass + + def test_shm_cleanup_drains_through_repeated_error_and_step_boundary_interrupts(self, monkeypatch): + owner = worker_mod._SharedMemoryOwner(2) + shms = [owner.create(1), owner.create(1)] + names = [shm.name for shm in shms] + first_error = OSError("unlink completion") + real_unlink = SharedMemory.unlink + unlink_interrupted = False + error_boundary_interrupts = 0 + step_boundary_interrupts = 0 + + def unlink_then_fail(shm): + nonlocal unlink_interrupted + real_unlink(shm) + if not unlink_interrupted: + unlink_interrupted = True + raise first_error + + def interrupt_error_boundary(): + nonlocal error_boundary_interrupts + if error_boundary_interrupts < 2: + error_boundary_interrupts += 1 + raise KeyboardInterrupt(f"error boundary {error_boundary_interrupts}") + + def interrupt_step_boundary(): + nonlocal step_boundary_interrupts + if unlink_interrupted and step_boundary_interrupts < 2: + step_boundary_interrupts += 1 + raise SystemExit(f"step boundary {step_boundary_interrupts}") + + monkeypatch.setattr(SharedMemory, "unlink", unlink_then_fail) + try: + cleanup_error = worker_mod._close_unlink_shms( + owner, + _after_error=interrupt_error_boundary, + _after_step=interrupt_step_boundary, ) - raced = domain("raced", 1) - # A second domain keeps the sweep branch reachable so the fence parks there. - swept = domain("swept", 2) - for handle in (raced, swept): - resources.live_domains[handle.name] = handle - worker._live_domains[handle.name] = handle + assert cleanup_error is first_error + assert error_boundary_interrupts == 2 + assert step_boundary_interrupts == 2 + for name in names: + with pytest.raises(FileNotFoundError): + SharedMemory(name=name) + finally: + monkeypatch.setattr(SharedMemory, "unlink", real_unlink) + for shm in shms: + try: + shm.close() + shm.unlink() + except FileNotFoundError: + pass + + def test_shm_cleanup_protects_the_recursive_handoff_from_a_second_interrupt(self, monkeypatch): + owner = worker_mod._SharedMemoryOwner(2) + shms = [owner.create(1), owner.create(1)] + names = [shm.name for shm in shms] + first_interrupt = KeyboardInterrupt("step boundary") + second_interrupt = SystemExit("recursive handoff") + step_interrupted = False + handoff_interrupted = False + first_seen = threading.Event() + real_thread = threading.Thread + + class HandoffInterrupt(real_thread): + def join(self, *args, **kwargs): + nonlocal handoff_interrupted + if not handoff_interrupted: + assert first_seen.wait(5.0) + handoff_interrupted = True + raise second_interrupt + return super().join(*args, **kwargs) + + def interrupt_step_boundary(): + nonlocal step_interrupted + if not step_interrupted: + step_interrupted = True + first_seen.set() + raise first_interrupt + + monkeypatch.setattr(worker_mod.threading, "Thread", HandoffInterrupt) + try: + cleanup_error = worker_mod._close_unlink_shms(owner, _after_step=interrupt_step_boundary) - in_sweep = threading.Event() - release_done = threading.Event() - real_sweep = worker._release_all_live_domains + assert cleanup_error is first_interrupt + assert handoff_interrupted + for name in names: + with pytest.raises(FileNotFoundError): + SharedMemory(name=name) + finally: + for shm in shms: + try: + shm.close() + shm.unlink() + except FileNotFoundError: + pass + + @pytest.mark.parametrize( + ("operation", "interrupt_create_call"), + (("comm_init", 1), ("domain_alloc", 1), ("domain_alloc", 2), ("domain_release", 1)), + ) + def test_created_shm_is_owned_before_the_caller_can_register_it( + self, monkeypatch, operation, interrupt_create_call + ): + worker = self._worker() + worker._config = {"device_ids": [0]} + resources = worker_mod._RunResources() + worker._building_run_resources = resources + monkeypatch.setattr(worker, "_comm_plan_rootinfo_path", lambda: "/tmp/comm-rootinfo") + if operation == "domain_alloc": + monkeypatch.setattr(worker, "_ensure_comm_base", lambda: None) + + real_init = SharedMemory.__init__ + create_calls = 0 + created: list[SharedMemory] = [] + interrupted_name = None + first_interrupt = KeyboardInterrupt(f"{operation} create return") + + def interrupt_after_create(shm, *args, **kwargs): + nonlocal create_calls, interrupted_name + real_init(shm, *args, **kwargs) + if kwargs.get("create"): + create_calls += 1 + created.append(shm) + if create_calls == interrupt_create_call: + interrupted_name = shm.name + raise first_interrupt + + monkeypatch.setattr(SharedMemory, "__init__", interrupt_after_create) + handle = worker_mod.CommDomainHandle( + name="d", + workers=(0,), + contexts={}, + allocation_id=7, + _release_fn=lambda _released: None, + ) + try: + with pytest.raises(KeyboardInterrupt) as caught: + if operation == "comm_init": + worker._ensure_comm_base() + elif operation == "domain_alloc": + worker._allocate_domain(name="d", workers=(0,), window_size=64, buffers=[]) + else: + worker._release_domain_claimed(handle) - def parked_sweep(res=None): - in_sweep.set() - assert release_done.wait(5.0), "releasing thread did not finish" - real_sweep(res) + assert caught.value is first_interrupt + assert interrupted_name is not None + with pytest.raises(FileNotFoundError): + worker_mod.SharedMemory(name=interrupted_name) + finally: + monkeypatch.setattr(SharedMemory, "__init__", real_init) + for shm in created: + try: + shm.close() + shm.unlink() + except FileNotFoundError: + pass + + def test_shm_owner_unlinks_a_name_interrupted_before_ftruncate(self, monkeypatch): + owner = worker_mod._SharedMemoryOwner(1) + first_interrupt = KeyboardInterrupt("before ftruncate") + real_ftruncate = worker_mod.os.ftruncate + + monkeypatch.setattr( + worker_mod.os, + "ftruncate", + lambda _fd, _size: (_ for _ in ()).throw(first_interrupt), + ) + with pytest.raises(KeyboardInterrupt) as caught: + owner.create(1) + assert caught.value is first_interrupt + + shm = owner._slots[0].shm + assert shm is not None + name = shm.name + monkeypatch.setattr(worker_mod.os, "ftruncate", real_ftruncate) + + assert worker_mod._close_unlink_shms(owner) is None + with pytest.raises(FileNotFoundError): + SharedMemory(name=name) + + def test_shm_owner_unlinks_a_name_created_before_handle_publication(self, monkeypatch): + owner = worker_mod._SharedMemoryOwner(1) + first_interrupt = KeyboardInterrupt("after shm_open") + real_init = SharedMemory.__init__ + created = None + + def create_then_interrupt(shm, *args, **kwargs): + nonlocal created + real_init(shm, *args, **kwargs) + created = shm + raise first_interrupt + + monkeypatch.setattr(SharedMemory, "__init__", create_then_interrupt) + try: + with pytest.raises(KeyboardInterrupt) as caught: + owner.create(1) + assert caught.value is first_interrupt + assert created is not None + name = created.name + monkeypatch.setattr(SharedMemory, "__init__", real_init) + + assert worker_mod._close_unlink_shms(owner) is None + with pytest.raises(FileNotFoundError): + SharedMemory(name=name) + finally: + monkeypatch.setattr(SharedMemory, "__init__", real_init) + if created is not None: + created.close() + try: + created.unlink() + except FileNotFoundError: + pass + + @pytest.mark.skipif(worker_mod.os.name != "posix", reason="requires POSIX shm_open") + def test_shm_owner_isolates_the_shm_open_to_fd_publication_gap(self): + owner = worker_mod._SharedMemoryOwner(1) + instructions = list(dis.get_instructions(SharedMemory.__init__)) + fd_stores = [ + index + for index, instruction in enumerate(instructions) + if instruction.opname == "STORE_ATTR" and instruction.argval == "_fd" + ] + assert fd_stores + gap_offset = instructions[fd_stores[-1] - 1].offset + interrupted = False + first_interrupt = KeyboardInterrupt("after shm_open before fd publication") + + def interrupt_fd_publication(frame, event, _arg): + nonlocal interrupted + if frame.f_code is SharedMemory.__init__.__code__: + frame.f_trace_opcodes = True + if event == "opcode" and frame.f_lasti == gap_offset: + interrupted = True + sys.settrace(None) + raise first_interrupt + return interrupt_fd_publication + + create_error = None + leaked = False + name = None + try: + sys.settrace(interrupt_fd_publication) + try: + owner.create(1) + except BaseException as exc: # noqa: BLE001 + create_error = exc + finally: + sys.settrace(None) - worker._release_all_live_domains = cast(Any, parked_sweep) + shm = owner._slots[0].shm + assert shm is not None + name = shm.name + cleanup_error = worker_mod._close_unlink_shms(owner) + try: + getattr(shared_memory_mod, "_posixshmem").shm_unlink(name) + except FileNotFoundError: + pass + else: + leaked = True - def do_release(): - assert in_sweep.wait(5.0), "fence never reached its sweep" - raced.release() - release_done.set() + assert create_error is None + assert not interrupted + assert cleanup_error is None + assert not leaked + finally: + sys.settrace(None) + if name is not None and not leaked: + try: + getattr(shared_memory_mod, "_posixshmem").shm_unlink(name) + except FileNotFoundError: + pass + + def test_shm_owner_never_unlinks_a_colliding_foreign_name(self, monkeypatch): + owner = worker_mod._SharedMemoryOwner(1) + real_init = SharedMemory.__init__ + foreign = None + collide = True + + def collide_once(shm, *args, **kwargs): + nonlocal collide, foreign + if collide: + collide = False + foreign = SharedMemory.__new__(SharedMemory) + real_init(foreign, *args, **kwargs) + raise FileExistsError(kwargs["name"]) + real_init(shm, *args, **kwargs) + + monkeypatch.setattr(SharedMemory, "__init__", collide_once) + attached = None + try: + owned = owner.create(1) + assert foreign is not None + foreign_name = foreign.name + assert owned.name != foreign_name - releaser = threading.Thread(target=do_release) - releaser.start() + assert worker_mod._close_unlink_shms(owner) is None + attached = SharedMemory(name=foreign_name) + finally: + monkeypatch.setattr(SharedMemory, "__init__", real_init) + if attached is not None: + attached.close() + if foreign is not None: + foreign.close() + try: + foreign.unlink() + except FileNotFoundError: + pass + + def test_shm_owner_collision_stays_unowned_when_the_handler_is_interrupted(self, monkeypatch): + owner = worker_mod._SharedMemoryOwner(1) + real_init = SharedMemory.__init__ + foreign = None + first_interrupt = KeyboardInterrupt("collision handler") + + def collide(shm, *args, **kwargs): + nonlocal foreign + foreign = SharedMemory.__new__(SharedMemory) + real_init(foreign, *args, **kwargs) + raise FileExistsError(kwargs["name"]) + + source, first_line = inspect.getsourcelines(worker_mod._SharedMemoryOwner._create_in_helper) + disarm_line = next(first_line + offset for offset, line in enumerate(source) if "slot.shm = None" in line) + + def interrupt_disarm(frame, event, _arg): + if ( + frame.f_code is worker_mod._SharedMemoryOwner._create_in_helper.__code__ + and event == "line" + and frame.f_lineno == disarm_line + ): + raise first_interrupt + return interrupt_disarm + + monkeypatch.setattr(SharedMemory, "__init__", collide) + attached = None try: - handle = RunHandle(worker, 1, (), resources) - worker._accepted_run_handles.add(handle) - assert worker._finalize_run_handle(handle, 1, None) is None + threading.settrace(interrupt_disarm) + with pytest.raises(KeyboardInterrupt) as caught: + owner.create(1) + assert caught.value is first_interrupt + assert foreign is not None + foreign_name = foreign.name + + monkeypatch.setattr(SharedMemory, "__init__", real_init) + worker_mod._close_unlink_shms(owner) + attached = SharedMemory(name=foreign_name) finally: - releaser.join(5.0) + threading.settrace(cast(Any, None)) + monkeypatch.setattr(SharedMemory, "__init__", real_init) + if attached is not None: + attached.close() + if foreign is not None: + foreign.close() + try: + foreign.unlink() + except FileNotFoundError: + pass + + def test_shm_cleanup_closes_fd_after_an_interrupt_returning_from_mmap_close(self): + owner = worker_mod._SharedMemoryOwner(1) + shm = owner.create(1) + fd = shm._fd + underlying_mmap = shm._mmap + first_interrupt = KeyboardInterrupt("mmap close return") + + class InterruptingMmap: + interrupted = False + + @property + def closed(self): + return underlying_mmap.closed + + def close(self): + if underlying_mmap.closed: + raise AssertionError("closed mmap was closed twice") + underlying_mmap.close() + if not self.interrupted: + self.interrupted = True + raise first_interrupt + + shm._mmap = InterruptingMmap() + try: + cleanup_error = worker_mod._close_unlink_shms(owner) - assert not releaser.is_alive() - assert raced.freed, "a release racing retirement was stranded on a drained queue" - assert swept.freed - assert resources.pending_release_domains == [] - assert freed.count("raced") == 1, f"raced domain freed more than once: {freed}" + assert cleanup_error is first_interrupt + with pytest.raises(OSError): + worker_mod.os.fstat(fd) + finally: + if shm._buf is not None: + shm._buf.release() + shm._buf = None + if not underlying_mmap.closed: + underlying_mmap.close() + try: + worker_mod.os.close(fd) + except OSError: + pass + try: + shm.unlink() + except FileNotFoundError: + pass - @staticmethod - def _gated_domain_worker(worker, target_id, outcome=None): - """Park `worker`'s backend release for `target_id`; optionally raise.""" - entered: list[int] = [] - in_backend = threading.Event() - let_go = threading.Event() + def test_shm_cleanup_does_not_retry_close_on_a_reused_fd(self, monkeypatch): + owner = worker_mod._SharedMemoryOwner(1) + shm = owner.create(1) + fd = shm._fd + real_close = worker_mod.os.close + sentinel_fd = worker_mod.os.open("/dev/null", worker_mod.os.O_RDONLY) + first_interrupt = KeyboardInterrupt("fd close return") + injected = False + replacement_installed = False + + def close_then_reuse(close_fd): + nonlocal injected, replacement_installed + if close_fd == fd and not injected: + injected = True + real_close(close_fd) + worker_mod.os.dup2(sentinel_fd, close_fd) + replacement_installed = True + raise first_interrupt + real_close(close_fd) + + monkeypatch.setattr(worker_mod.os, "close", close_then_reuse) + try: + cleanup_error = worker_mod._close_unlink_shms(owner) - def gated(handle): - entered.append(handle.allocation_id) - if handle.allocation_id == target_id: - in_backend.set() - let_go.wait(10.0) - if outcome is not None: - raise outcome + assert cleanup_error is first_interrupt + assert replacement_installed + worker_mod.os.fstat(fd) + finally: + monkeypatch.setattr(worker_mod.os, "close", real_close) + if replacement_installed: + try: + real_close(fd) + except OSError: + pass + real_close(sentinel_fd) + try: + shm.close() + shm.unlink() + except FileNotFoundError: + pass - worker._release_domain_claimed = cast(Any, gated) - return entered, in_backend, let_go + def test_shm_cleanup_does_not_retry_close_on_a_same_inode_reused_fd(self, monkeypatch): + owner = worker_mod._SharedMemoryOwner(1) + shm = owner.create(1) + fd = shm._fd + real_close = worker_mod.os.close + duplicate_fd = worker_mod.os.dup(fd) + first_interrupt = KeyboardInterrupt("fd close return") + injected = False + replacement_installed = False + + def close_then_reuse(close_fd): + nonlocal injected, replacement_installed + if close_fd == fd and not injected: + injected = True + real_close(close_fd) + worker_mod.os.dup2(duplicate_fd, close_fd) + replacement_installed = True + raise first_interrupt + real_close(close_fd) + + monkeypatch.setattr(worker_mod.os, "close", close_then_reuse) + try: + cleanup_error = worker_mod._close_unlink_shms(owner) - def _retired_domain(self, worker, resources, name, allocation_id): + assert cleanup_error is first_interrupt + assert replacement_installed + worker_mod.os.fstat(fd) + finally: + monkeypatch.setattr(worker_mod.os, "close", real_close) + if replacement_installed: + try: + real_close(fd) + except OSError: + pass + real_close(duplicate_fd) + try: + shm.close() + shm.unlink() + except FileNotFoundError: + pass + + def test_domain_release_preserves_the_dispatch_interrupt_through_shm_cleanup(self, monkeypatch): + worker = self._worker() + first_interrupt = KeyboardInterrupt("dispatch join") + monkeypatch.setattr( + worker, + "_dispatch_control_domain", + lambda **_kwargs: (_ for _ in ()).throw(first_interrupt), + ) handle = worker_mod.CommDomainHandle( - name=name, - workers=(), + name="d", + workers=(0,), contexts={}, - allocation_id=allocation_id, - _release_fn=lambda released, owner=resources: worker._release_domain_handle(released, owner), + allocation_id=7, + _release_fn=lambda _released: None, ) - worker._live_domains[name] = handle - return handle - def test_sweep_and_post_fence_release_free_a_domain_once(self): - """Two paths reaching one handle must not both drive the backend free. + created: list[SharedMemory] = [] + created_fds: list[int] = [] + real_init = SharedMemory.__init__ - The dangerous order is release-then-sweep: ``release()`` wins the - ``_released`` flag, so the sweep — still holding the handle in the - snapshot it took earlier — skips setting that flag and goes straight to - the backend call the release is already making. (Sweep-then-release is - already safe: the sweep sets ``_released`` before freeing, which makes - the later ``release()`` a no-op.) - """ - worker = Worker(level=3, num_sub_workers=0) - worker._worker = cast(Any, object()) - entered, in_backend, let_go = self._gated_domain_worker(worker, target_id=7) + def tracked_shared_memory(shm, *args, **kwargs): + real_init(shm, *args, **kwargs) + created.append(shm) + created_fds.append(shm._fd) - resources = worker_mod._RunResources() - resources.retired = True # the owning run's fence has already passed - contested = self._retired_domain(worker, resources, "contested", 7) + monkeypatch.setattr(SharedMemory, "__init__", tracked_shared_memory) + real_close = worker_mod.os.close + close_interrupted = False - releaser = threading.Thread(target=contested.release) - releaser.start() + def interrupt_first_close(fd): + nonlocal close_interrupted + if created_fds and fd == created_fds[0] and not close_interrupted: + close_interrupted = True + raise SystemExit("shm close") + return real_close(fd) + + monkeypatch.setattr(worker_mod.os, "close", interrupt_first_close) try: - assert in_backend.wait(5.0), "release never reached the backend" - # The sweep's snapshot still holds `contested`; it must not free it - # a second time. - sweeper = threading.Thread(target=worker._release_all_live_domains) - sweeper.start() - let_go.set() - sweeper.join(5.0) - assert not sweeper.is_alive() + with pytest.raises(KeyboardInterrupt) as caught: + worker._release_domain_claimed(handle) + assert caught.value is first_interrupt + assert close_interrupted finally: - let_go.set() - releaser.join(5.0) + monkeypatch.setattr(worker_mod.os, "close", real_close) + for shm in created: + try: + shm.close() + shm.unlink() + except FileNotFoundError: + pass - assert entered.count(7) == 1, f"allocation 7 was released {entered.count(7)} times" - assert contested.freed + def test_committed_domain_is_owned_before_reply_shm_teardown(self, monkeypatch): + worker = self._worker() + worker._config = {"device_ids": [0]} + resources = worker_mod._RunResources() + worker._building_run_resources = resources + monkeypatch.setattr(worker, "_ensure_comm_base", lambda: None) + staged_shms: list[SharedMemory] = [] + + def successful_dispatch(*, request_shms, reply_shms, **_kwargs): + assert reply_shms is not None + staged_shms.extend(request_shms.values()) + staged_shms.extend(reply_shms.values()) + reply_buf = reply_shms[0].buf + assert reply_buf is not None + worker_mod._DOMAIN_REPLY_HEADER.pack_into(reply_buf, 0, 1, 0xC7, 0xB000, 0) + + monkeypatch.setattr(worker, "_dispatch_control_domain", successful_dispatch) + real_close = worker_mod.os.close + interrupted = False + target_fd = None + + def remember_target_fd(*, request_shms, reply_shms, **kwargs): + nonlocal target_fd + successful_dispatch(request_shms=request_shms, reply_shms=reply_shms, **kwargs) + target_fd = staged_shms[0]._fd + + def interrupt_first_close(fd): + nonlocal interrupted + if target_fd is not None and fd == target_fd and not interrupted: + interrupted = True + raise KeyboardInterrupt + return real_close(fd) - def test_second_domain_release_waits_for_the_first(self): - """A second caller blocks until the owner's backend call returns. + monkeypatch.setattr(worker, "_dispatch_control_domain", remember_target_fd) + monkeypatch.setattr(worker_mod.os, "close", interrupt_first_close) + try: + with pytest.raises(KeyboardInterrupt): + worker._allocate_domain(name="d", workers=(0,), window_size=64, buffers=[]) + + handle = resources.live_domains["d"] + assert worker._live_domains["d"] is handle + assert handle.workers == (0,) + assert resources.requires_ordered_cleanup + assert interrupted + finally: + monkeypatch.setattr(worker_mod.os, "close", real_close) + for shm in staged_shms: + try: + shm.close() + shm.unlink() + except FileNotFoundError: + pass - Returning early would let the caller mark the handle freed, drop it - from ``_live_domains`` and — on the ``close()`` path — tear down the - mailboxes the in-flight release is still using, which is exactly what - ``close()`` orders its domain sweep before ``_worker.close()`` to avoid. + def test_an_interrupt_through_a_region_create_refuses_further_work(self): + """The create releases the GIL, so an interrupt can land mid-flight. + + The child may still be finishing a region whose id would be written + into a reply this frame is about to unlink — something on the chip that + nothing here can name. An ordinary create failure is not that case: the + child releases its own region before reporting one. """ - worker = Worker(level=3, num_sub_workers=0) - worker._worker = cast(Any, object()) - entered, in_backend, let_go = self._gated_domain_worker(worker, target_id=7) + worker = self._worker() + worker._config = {"platform": "a2a3sim"} + worker._validate_l3_l2_worker_id = cast(Any, lambda _wid: None) - resources = worker_mod._RunResources() - contested = self._retired_domain(worker, resources, "contested", 7) + def _interrupted(*_args): + raise KeyboardInterrupt - owner = threading.Thread(target=worker._release_domain_now, args=(contested,)) - second_done = threading.Event() + worker._worker = cast(Any, SimpleNamespace(control_l3_l2_region_create=_interrupted)) + with pytest.raises(KeyboardInterrupt): + worker._create_l3_l2_region(0, 4096, 64) + assert worker._ordered_cleanup_error is not None + with pytest.raises(RuntimeError, match="no further work is admitted"): + worker._require_no_ordered_cleanup_failure("submit") - def second_caller(): - worker._release_domain_now(contested) - second_done.set() + def test_an_ordinary_region_create_failure_does_not_refuse_further_work(self): + worker = self._worker() + worker._config = {"platform": "a2a3sim"} + worker._validate_l3_l2_worker_id = cast(Any, lambda _wid: None) - second = threading.Thread(target=second_caller) - owner.start() - try: - assert in_backend.wait(5.0), "owner never reached the backend" - assert not contested.freed, "freed must stay false while the backend call is in flight" - second.start() - assert not second_done.wait(0.5), "the second caller returned while the owner was still releasing" - assert not contested.freed - finally: - let_go.set() - owner.join(5.0) - second.join(5.0) + def _failed(*_args): + raise RuntimeError("chip refused the region") - assert second_done.is_set() - assert entered.count(7) == 1, f"allocation 7 reached the backend {entered.count(7)} times" + worker._worker = cast(Any, SimpleNamespace(control_l3_l2_region_create=_failed)) + with pytest.raises(RuntimeError, match="chip refused the region"): + worker._create_l3_l2_region(0, 4096, 64) + assert worker._ordered_cleanup_error is None, "an ordinary failure must not shut the worker" - def test_failed_domain_release_is_replayed_to_a_second_caller(self): - """The owner's failure reaches every later caller, so no path reports - success for an allocation whose backend release did not happen. + def test_a_region_rollback_that_cannot_release_refuses_further_work(self): + """The id was never tracked, so no later cleanup can reclaim it. - ``_release_all_live_domains`` keeps an un-freed handle in - ``_live_domains`` precisely so ``close()`` reports it as a residual; a - second caller that returned success would erase that. + There is no handle for a fence to fail on, which is why the refusal is + recorded here rather than left to one. """ - worker = Worker(level=3, num_sub_workers=0) - worker._worker = cast(Any, object()) - boom = RuntimeError("backend release failed") - _entered, in_backend, let_go = self._gated_domain_worker(worker, target_id=7, outcome=boom) - let_go.set() + worker = self._worker() + assert worker._ordered_cleanup_error is None + + # The shape _create_l3_l2_region's rollback reaches: the region exists + # on the chip and the release for it failed. + leaked = RuntimeError("create_l3_l2_region: rollback could not release region 4") + with worker._hierarchical_start_cv: + worker._ordered_cleanup_error = leaked + + with pytest.raises(RuntimeError, match="no further work is admitted"): + worker._require_no_ordered_cleanup_failure("submit") + with pytest.raises(RuntimeError, match="no further work is admitted"): + with worker._control_reservation("Worker.malloc"): + pass - resources = worker_mod._RunResources() - contested = self._retired_domain(worker, resources, "contested", 7) - with pytest.raises(RuntimeError, match="backend release failed"): - worker._release_domain_now(contested) - assert in_backend.is_set() - assert not contested.freed +class _StubOrch: + def _release_run(self, run_id: int) -> None: + pass - # The sweep is the second caller: it must see the failure, keep the - # handle, and leave `freed` false. - 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" - def test_allocate_domain_outside_graph_construction_is_rejected(self): - worker = Worker(level=3, num_sub_workers=0) - worker._worker = cast(Any, object()) - assert worker._building_run_resources is None +class _WatchedSet(set): + """A set that reports each discard to `on_discard` before performing it.""" - with pytest.raises(RuntimeError, match="graph is being built"): - worker._allocate_domain(name="d", workers=(0,), window_size=4096, buffers=[]) + 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 # --------------------------------------------------------------------------- diff --git a/tests/ut/py/test_worker/test_l3_l2_orch_comm.py b/tests/ut/py/test_worker/test_l3_l2_orch_comm.py index 10531fb69..c25c57a60 100644 --- a/tests/ut/py/test_worker/test_l3_l2_orch_comm.py +++ b/tests/ut/py/test_worker/test_l3_l2_orch_comm.py @@ -8,7 +8,10 @@ # ----------------------------------------------------------------------------------------------------------- import ctypes +import gc import importlib +import os +import threading import time from concurrent.futures import ThreadPoolExecutor from multiprocessing.shared_memory import SharedMemory @@ -389,6 +392,102 @@ def test_onboard_direct_mapping_allows_granularity_aligned_mapping(monkeypatch): shm.unlink() +@pytest.mark.parametrize("interrupted_publication", ["worker", "run"]) +def test_direct_region_create_rolls_back_partially_published_region(monkeypatch, interrupted_publication): + class _AppendThenInterrupt(list): + def append(self, item) -> None: + super().append(item) + raise KeyboardInterrupt(f"interrupted {interrupted_publication} publication") + + worker, shm, fake_c_worker = _make_started_sim_worker() + resources = worker_module._RunResources() + worker._building_run_resources = resources + close_calls: list[int] = [] + monkeypatch.setattr(worker_module, "_l3_host_mapped_region_import_sim", lambda _token, _size: 55) + monkeypatch.setattr( + l3_l2_orch_comm, + "_l3_host_mapped_region_close", + lambda handle: close_calls.append(int(handle)), + ) + if interrupted_publication == "worker": + worker._live_l3_l2_regions = _AppendThenInterrupt() + else: + resources.l3_l2_regions = _AppendThenInterrupt() + + try: + with pytest.raises(KeyboardInterrupt, match=f"interrupted {interrupted_publication} publication"): + worker._create_l3_l2_region(0, 64, 128) + + assert worker._live_l3_l2_regions == [] + assert resources.l3_l2_regions == [] + assert resources.requires_ordered_cleanup is False + assert close_calls == [55] + assert fake_c_worker.release_calls == [(0, 1)] + finally: + worker._building_run_resources = None + worker._live_l3_l2_regions.clear() + resources.l3_l2_regions.clear() + worker._close_l3_l2_orch_comm() + shm.close() + shm.unlink() + + +def test_direct_region_create_mapping_rollback_failure_poisons_worker(monkeypatch): + class _AppendThenInterrupt(list): + def append(self, item) -> None: + super().append(item) + raise KeyboardInterrupt("interrupted publication") + + worker, shm, fake_c_worker = _make_started_sim_worker() + worker._live_l3_l2_regions = _AppendThenInterrupt() + monkeypatch.setattr(worker_module, "_l3_host_mapped_region_import_sim", lambda _token, _size: 55) + monkeypatch.setattr( + l3_l2_orch_comm, + "_l3_host_mapped_region_close", + lambda _handle: (_ for _ in ()).throw(RuntimeError("mapping close failed")), + ) + + try: + with pytest.raises(RuntimeError, match="rollback could not close the L3 Host mapping") as excinfo: + worker._create_l3_l2_region(0, 64, 128) + + assert isinstance(excinfo.value.__cause__, RuntimeError) + assert fake_c_worker.release_calls == [(0, 1)] + assert worker._live_l3_l2_regions == [] + with pytest.raises(RuntimeError, match="no further work is admitted"): + worker._require_no_ordered_cleanup_failure("submit") + finally: + worker._live_l3_l2_regions.clear() + worker._close_l3_l2_orch_comm() + shm.close() + shm.unlink() + + +def test_unadopted_native_mapping_cleanup_failure_poisons_worker(monkeypatch): + worker, shm, fake_c_worker = _make_started_sim_worker() + cleanup_errors = iter(("", "native owner cleanup failed")) + monkeypatch.setattr(worker_module, "_l3_host_mapped_region_take_cleanup_error", lambda: next(cleanup_errors)) + monkeypatch.setattr( + worker_module, + "_l3_host_mapped_region_import_sim", + lambda _token, _size: (_ for _ in ()).throw(KeyboardInterrupt("interrupted native adoption")), + ) + + try: + with pytest.raises(RuntimeError, match="rollback could not close the L3 Host mapping") as excinfo: + worker._create_l3_l2_region(0, 64, 128) + + assert isinstance(excinfo.value.__cause__, RuntimeError) + assert "native owner cleanup failed" in str(excinfo.value.__cause__) + assert fake_c_worker.release_calls == [(0, 1)] + with pytest.raises(RuntimeError, match="no further work is admitted"): + worker._require_no_ordered_cleanup_failure("submit") + finally: + worker._close_l3_l2_orch_comm() + shm.close() + shm.unlink() + + def test_onboard_region_create_handler_uses_named_export_fields(monkeypatch): req_shm = SharedMemory(create=True, size=l3_l2_orch_comm._REGION_CREATE_REQUEST_BYTES) reply_shm = SharedMemory(create=True, size=l3_l2_orch_comm._REGION_CREATE_REPLY_BYTES) @@ -487,7 +586,8 @@ def test_l3_host_mapped_counter_wait_releases_gil_for_python_notifier(): shm = SharedMemory(create=True, size=64) handle = 0 try: - handle = _task_interface_ext._l3_host_mapped_region_import_sim(shm.name, 64) + owner = _task_interface_ext._l3_host_mapped_region_import_sim(shm.name, 64) + handle = int(owner) def notify() -> None: time.sleep(0.05) @@ -512,7 +612,8 @@ def test_l3_host_mapped_sim_payload_and_counter_helpers_roundtrip(): shm = SharedMemory(create=True, size=128) handle = 0 try: - handle = _task_interface_ext._l3_host_mapped_region_import_sim(shm.name, 128) + owner = _task_interface_ext._l3_host_mapped_region_import_sim(shm.name, 128) + handle = int(owner) src_t = ctypes.c_uint8 * 8 src = src_t(*range(10, 18)) dst = src_t() @@ -547,7 +648,8 @@ def test_l3_host_mapped_region_close_makes_sim_handle_unusable(): shm = SharedMemory(create=True, size=64) handle = 0 try: - handle = _task_interface_ext._l3_host_mapped_region_import_sim(shm.name, 64) + owner = _task_interface_ext._l3_host_mapped_region_import_sim(shm.name, 64) + handle = int(owner) _task_interface_ext._l3_host_mapped_region_close(handle) with pytest.raises(RuntimeError, match="closed or unknown"): @@ -559,6 +661,103 @@ def test_l3_host_mapped_region_close_makes_sim_handle_unusable(): shm.unlink() +def test_l3_host_mapped_import_owner_closes_unadopted_mapping(): + shm = SharedMemory(create=True, size=64) + raw_handle = 0 + try: + owner = _task_interface_ext._l3_host_mapped_region_import_sim(shm.name, 64) + raw_handle = int(owner) + del owner + gc.collect() + + with pytest.raises(RuntimeError, match="closed or unknown"): + _task_interface_ext._l3_host_mapped_counter_test(raw_handle, 0, 0, int(WaitCmp.EQ)) + finally: + if raw_handle: + _task_interface_ext._l3_host_mapped_region_close(raw_handle) + shm.close() + shm.unlink() + + +def test_sim_import_registry_failure_releases_pre_registry_mapping(): + if not os.path.exists("/proc/self/maps"): + pytest.skip("requires Linux procfs resource accounting") + + shm = SharedMemory(create=True, size=64) + shm_token = shm.name.lstrip("/") + + def mapped_resource_counts() -> tuple[int, int]: + fd_count = 0 + for fd_name in os.listdir("/proc/self/fd"): + try: + target = os.readlink(f"/proc/self/fd/{fd_name}") + except OSError: + continue + fd_count += shm_token in target + with open("/proc/self/maps", encoding="utf-8") as maps_file: + map_count = sum(shm_token in line for line in maps_file) + return fd_count, map_count + + try: + baseline = mapped_resource_counts() + _task_interface_ext._l3_host_mapped_region_take_cleanup_error() + _task_interface_ext._l3_host_mapped_region_fail_next_registry_insert_for_test() + + with pytest.raises(RuntimeError, match="injected mapped-region registry insertion failure"): + _task_interface_ext._l3_host_mapped_region_import_sim(shm.name, 64) + + gc.collect() + assert mapped_resource_counts() == baseline + assert _task_interface_ext._l3_host_mapped_region_take_cleanup_error() == "" + finally: + shm.close() + shm.unlink() + + +def test_l3_host_mapped_concurrent_closes_wait_for_in_flight_counter_wait(): + shm = SharedMemory(create=True, size=64) + handle = 0 + try: + owner = _task_interface_ext._l3_host_mapped_region_import_sim(shm.name, 64) + handle = int(owner) + close_entered = [threading.Event(), threading.Event()] + close_done = [threading.Event(), threading.Event()] + + def wait_for_counter(): + return _task_interface_ext._l3_host_mapped_counter_wait(handle, 0, 1, int(WaitCmp.EQ), 1_000_000_000) + + def close_mapping(index: int) -> None: + close_entered[index].set() + _task_interface_ext._l3_host_mapped_region_close(handle) + close_done[index].set() + + with ThreadPoolExecutor(max_workers=3) as executor: + wait_future = executor.submit(wait_for_counter) + deadline = time.monotonic() + 1.0 + while _task_interface_ext._l3_host_mapped_region_active_leases(handle) != 1: + assert time.monotonic() < deadline, "counter wait never acquired its mapped-region lease" + time.sleep(0.001) + + close_futures = [executor.submit(close_mapping, index) for index in range(2)] + assert all(event.wait(1.0) for event in close_entered) + assert not any(event.wait(0.05) for event in close_done), ( + "a concurrent close returned while a native operation still held the region" + ) + + cast(memoryview, shm.buf)[:4] = b"\x01\x00\x00\x00" + assert wait_future.result(timeout=1.0) == (0, 0, 1, True, "") + for close_future in close_futures: + close_future.result(timeout=1.0) + + with pytest.raises(RuntimeError, match="closed or unknown"): + _task_interface_ext._l3_host_mapped_counter_test(handle, 0, 1, int(WaitCmp.EQ)) + finally: + if handle: + _task_interface_ext._l3_host_mapped_region_close(handle) + shm.close() + shm.unlink() + + def test_sim_direct_transfer_failure_poisons_only_region(monkeypatch): worker, shm, _fake_c_worker = _make_started_sim_worker() try: diff --git a/tests/ut/py/test_worker/test_startup_readiness.py b/tests/ut/py/test_worker/test_startup_readiness.py index 82be5ed56..57798c58b 100644 --- a/tests/ut/py/test_worker/test_startup_readiness.py +++ b/tests/ut/py/test_worker/test_startup_readiness.py @@ -925,6 +925,119 @@ def __getattr__(self, name): w._hierarchical_start_cv = real_cv assert _run_catch(w.close) is None + def test_close_outcome_retries_an_interrupt_before_publication(self, monkeypatch): + import simpler.worker as worker_mod # noqa: PLC0415 + + interrupt = SystemExit("before close outcome publication") + entered = threading.Event() + release = threading.Event() + attempt_type = worker_mod._CloseAttempt + + class _InterruptBeforeOutcome(attempt_type): + __slots__ = ("_armed",) + + def __init__(self): + super().__init__() + object.__setattr__(self, "_armed", True) + + def publish(self, error, incomplete): + if self._armed: + object.__setattr__(self, "_armed", False) + entered.set() + assert release.wait(10.0) + raise interrupt + return super().publish(error, incomplete) + + monkeypatch.setattr(worker_mod, "_CloseAttempt", _InterruptBeforeOutcome) + w = self._make_l2(monkeypatch) + owner_result: list = [] + joiner_result: list = [] + + def owner(): + w.init() + owner_result.append(_run_catch(w.close)) + + owner_thread = threading.Thread(target=owner) + joiner_thread = threading.Thread(target=lambda: joiner_result.append(_run_catch(w.close))) + owner_thread.start() + try: + assert entered.wait(5.0) + joiner_thread.start() + time.sleep(0.1) + assert joiner_thread.is_alive() + release.set() + owner_thread.join(5.0) + joiner_thread.join(5.0) + + assert not owner_thread.is_alive() + assert not joiner_thread.is_alive() + assert owner_result == [interrupt] + assert joiner_result == [interrupt] + assert w._close_completion is not None and w._close_completion.done + assert w._close_completion.error is interrupt + assert _run_catch(w.close) is interrupt + finally: + release.set() + owner_thread.join(5.0) + if joiner_thread.ident is not None: + joiner_thread.join(5.0) + + def test_close_outcome_survives_an_interrupt_after_publication(self, monkeypatch): + import simpler.worker as worker_mod # noqa: PLC0415 + + interrupt = SystemExit("after close outcome publication") + entered = threading.Event() + release = threading.Event() + attempt_type = worker_mod._CloseAttempt + + class _InterruptAfterOutcome(attempt_type): + __slots__ = ("_armed",) + + def __init__(self): + super().__init__() + object.__setattr__(self, "_armed", True) + + def __setattr__(self, name, value): + if name == "_outcome" and value is not None and self._armed: + super().__setattr__(name, value) + object.__setattr__(self, "_armed", False) + entered.set() + assert release.wait(10.0) + raise interrupt + return super().__setattr__(name, value) + + monkeypatch.setattr(worker_mod, "_CloseAttempt", _InterruptAfterOutcome) + w = self._make_l2(monkeypatch) + owner_result: list = [] + joiner_result: list = [] + + def owner(): + w.init() + owner_result.append(_run_catch(w.close)) + + owner_thread = threading.Thread(target=owner) + joiner_thread = threading.Thread(target=lambda: joiner_result.append(_run_catch(w.close))) + owner_thread.start() + try: + assert entered.wait(5.0) + assert w._close_completion is not None and w._close_completion.done + joiner_thread.start() + joiner_thread.join(5.0) + assert not joiner_thread.is_alive() + assert joiner_result == [None] + + release.set() + owner_thread.join(5.0) + assert not owner_thread.is_alive() + assert owner_result == [interrupt] + assert w._close_completion.error is None + assert _run_catch(w.close) is None + finally: + release.set() + owner_thread.join(5.0) + if joiner_thread.ident is not None: + joiner_thread.join(5.0) + def test_reap_deadline_starts_after_shutdown_broadcast(self, monkeypatch): # The child-reap grace must be measured from when SHUTDOWN is broadcast, # not from teardown entry: a slow pre-child cleanup step must not consume