From 7b7b206d82971d26b07f674f543d4c39e38e7be3 Mon Sep 17 00:00:00 2001 From: Crane-Liu Date: Fri, 24 Jul 2026 02:12:50 -0700 Subject: [PATCH] Refactor: isolate hierarchical worker run state - Track task completion and first errors on per-run fences - Namespace TensorMap entries by run and protect concurrent access - Keep Worker.run blocking while removing global drain dependencies - Cover run identity, failure isolation, and resource lifetime - Keep run bookkeeping non-throwing on the scheduler and worker threads, where an escaping exception would terminate the process - Carry the orchestration failure message onto the run fence so a submission failure is retrievable from the run, not just the caller - Hold runs_mu_ across both the quiescence test and Ring compaction: begin_run takes only that mutex and no slot is allocated without a building run, so a run registered mid-compaction keeps its first slot - Close the run inside the block that releases it, so a failure there cannot strand the run id or skip run-owned resource cleanup --- docs/comm-domain.md | 26 +- docs/hierarchical_level_runtime.md | 8 +- docs/orchestrator.md | 93 +++--- docs/remote-l3-worker-design.md | 19 +- docs/scheduler.md | 7 +- docs/task-flow.md | 12 +- python/bindings/worker_bind.h | 23 +- python/simpler/orchestrator.py | 20 +- python/simpler/task_interface.py | 6 +- python/simpler/worker.py | 90 +++--- src/common/hierarchical/orchestrator.cpp | 273 +++++++++++++----- src/common/hierarchical/orchestrator.h | 61 ++-- src/common/hierarchical/ring.cpp | 2 +- src/common/hierarchical/ring.h | 3 +- src/common/hierarchical/scheduler.cpp | 14 +- src/common/hierarchical/scheduler.h | 9 +- src/common/hierarchical/tensormap.cpp | 20 +- src/common/hierarchical/tensormap.h | 28 +- src/common/hierarchical/types.cpp | 1 + src/common/hierarchical/types.h | 33 ++- src/common/hierarchical/worker.cpp | 6 +- src/common/hierarchical/worker_manager.cpp | 33 +-- src/common/hierarchical/worker_manager.h | 21 +- .../ut/cpp/hierarchical/test_orchestrator.cpp | 138 ++++++++- tests/ut/cpp/hierarchical/test_scheduler.cpp | 25 +- tests/ut/cpp/hierarchical/test_tensormap.cpp | 80 +++-- .../py/test_worker/test_error_propagation.py | 32 +- .../ut/py/test_worker/test_l3_l2_orch_comm.py | 16 +- 28 files changed, 749 insertions(+), 350 deletions(-) diff --git a/docs/comm-domain.md b/docs/comm-domain.md index 0b9df0b9fb..00a9efa0bd 100644 --- a/docs/comm-domain.md +++ b/docs/comm-domain.md @@ -58,26 +58,26 @@ The handle is a context manager. Its lifecycle has **two distinct states**: exits). Further indexing (`handle[i]`) raises. This is the *user-visible* state: "do not hand this domain to any new `submit_*`." - **`freed`** — the backend `comm_release_domain_windows` has actually run and - the device memory is gone. This happens **after** `Worker.run` drains the - DAG, never inside the `with` block. + the device memory is gone. This happens **after** the owning run's completion + fence, never inside the `with` block. This split exists because `submit_next_level()` only *enqueues* DAG work; -`Worker.run()` does not drain until the orch function returns. If `release()` -freed memory immediately on `with`-exit, a still-queued task that captured the -domain's `device_ctx` / `buffer_ptrs` would read freed memory. So **release is -deferred**: `release()` flips `released` and queues the backend free; the real -free runs after drain, when every task that could reference the window has -completed. +`Worker.run()` does not wait for completion until the orch function returns. +If `release()` freed memory immediately on `with`-exit, a still-queued task that +captured the domain's `device_ctx` / `buffer_ptrs` would read freed memory. So +**release is deferred**: `release()` flips `released` and queues the backend +free; the real free runs after the run fence, when every task that could +reference the window has completed. Mental model: like `with open(f) as fh: ...` — the user-visible close is lexical (end of block), the physical teardown is managed for you. Use `handle.released` to guard against accidental reuse; use `handle.freed` only if you must assert physical teardown. -Cleanup is **drain-safe**: even if a chip task fails and `drain()` re-raises, -`Worker.run` still executes the pending releases and sweeps any live domains the -orch fn forgot to release (LIFO), so a failed run cannot strand backend -allocations into the next run. +Cleanup is **failure-safe**: even if a chip task fails and the run wait +re-raises, `Worker.run` still executes pending releases and sweeps any live +domains the orch fn forgot to release (LIFO), so a failed run cannot strand +backend allocations into the next run. --- @@ -94,7 +94,7 @@ called many times: - the **base communicator is created once** and reused — it is *not* rebuilt per `run` or per domain; -- only the **per-domain windows** are allocated (and freed after drain) on each +- only the **per-domain windows** are allocated (and freed after the run fence) on each `allocate_domain` / `run`. Each allocation gets a fresh `allocation_id` so concurrent or sequential domains never collide on IPC handshake / barrier names. diff --git a/docs/hierarchical_level_runtime.md b/docs/hierarchical_level_runtime.md index adbd5823bc..1cdf5ceb36 100644 --- a/docs/hierarchical_level_runtime.md +++ b/docs/hierarchical_level_runtime.md @@ -157,7 +157,7 @@ what flows through `ChipWorker::run`. │ │ success → COMPLETED │ │ failure → FAILED + poison downstream │ │ try_consume → ring release - │ drain() ◄── notify when all done │ + │ wait_run(id) ◄── run becomes terminal ``` Communication channels: @@ -212,9 +212,9 @@ When L4's `WorkerThread` writes a task frame to the L3 child's mailbox, the frame carries the callable hash digest plus `config` and `args_blob`. The child loop reads the digest, resolves it through its local identity table to a private orch-function slot, and calls `inner_worker.run(orch_fn, args, cfg)`. The inner -Worker opens its own scope, executes the orch function with its own -`Orchestrator`, and drains. Each level's orch fn receives its own Orchestrator -— recursion is symmetric. +Worker opens its own run and scope, executes the orch function with its own +`Orchestrator`, closes submission, and waits for that run. Each level's orch fn +receives its own Orchestrator — recursion is symmetric. **Nested fork ordering**: L3's own children (sub/chip) are forked **inside** the L4 child process, in L3's eager `init()` (which the L4 child runs before it diff --git a/docs/orchestrator.md b/docs/orchestrator.md index 4c8f46e384..a189c15842 100644 --- a/docs/orchestrator.md +++ b/docs/orchestrator.md @@ -7,9 +7,9 @@ internals, not public submit arguments. See [callable-identity-registration.md](callable-identity-registration.md). The Orchestrator is the **DAG builder**. It runs single-threaded on the user's -thread (inside `Worker::run` between `scope_begin` and `drain`) and owns the -three data structures that turn a sequence of `submit_*` calls into a scheduled -DAG: `Ring`, `TensorMap`, and `Scope`. +thread while a run is open for submission and owns the three data structures +that turn a sequence of `submit_*` calls into a scheduled DAG: `Ring`, +`TensorMap`, and `Scope`. For the high-level role of the Orchestrator among the three engine components, see [hierarchical_level_runtime.md](hierarchical_level_runtime.md). For what @@ -51,14 +51,18 @@ public: // --- Intermediate-buffer allocation (runtime-owned lifetime) --- Tensor alloc(const std::vector &shape, DataType dtype); - // --- Internal lifecycle (invoked by Worker::run only, bound as _scope_begin - // / _scope_end / _drain in the Python facade) --- + // --- Internal lifecycle (invoked by Worker::run only) --- + RunId begin_run(); + void close_run_submission(RunId run_id); + void fail_run_submission(RunId run_id, std::exception_ptr error); + void wait_run(RunId run_id); + bool run_done(RunId run_id) const; + void release_run(RunId run_id); void scope_begin(); void scope_end(); - void drain(); private: - // ... components: Ring, TensorMap, Scope, slot pool, active_tasks_ counter + // ... components: Ring, TensorMap, Scope, and the RunState registry }; struct SubmitResult { TaskSlot task_slot; }; // internal only; not bound to Python @@ -67,9 +71,10 @@ struct SubmitResult { TaskSlot task_slot; }; // internal only; not bound to Pyt **Status**: `submit_sub` takes only `(CallableIdentity, args)` — no `config`, since SUB has no per-call config. -`scope_begin` / `scope_end` / `drain` are invoked from Python `Worker.run` via -`_scope_begin` / `_scope_end` / `_drain` bindings. They are not part of the -user-facing orch-fn API. +The run lifecycle and outer `scope_begin` / `scope_end` are invoked from Python +`Worker.run` through private bindings. They are not part of the user-facing +orch-fn API. `Worker.run` remains blocking: it closes submission, waits for the +matching run fence, performs run-owned cleanup, and returns `None`. Remote L3 submit adds two hidden pieces of metadata: final eligible worker-id sets and optional `RemoteTaskArgsSidecar` entries aligned by tensor index. @@ -161,9 +166,9 @@ small POD copied by value. `callable` is a `uint64_t` opaque handle (see **Step 3 — tag walk**: The only place tags are consumed. After this step tags are never inspected again; they are not carried into the slot's stored `task_args` value during dispatch (see [task-flow.md](task-flow.md) §3). -Local tensors key TensorMap by `(LOCAL_HOST, ptr)` or -`(LOCAL_CHILD, worker, ptr)`. Remote tensors with sidecars key by -`(address_kind, owner_worker_id, buffer_id, generation, offset)`. +Every TensorMap key starts with the current `RunId`. Local tensor identity is +then `(LOCAL_HOST, ptr)` or `(LOCAL_CHILD, worker, ptr)`. Remote tensors with +sidecars use `(address_kind, owner_worker_id, buffer_id, generation, offset)`. | Tag | `tensormap.lookup` | `tensormap.insert` | | --- | ------------------ | ------------------ | @@ -276,8 +281,7 @@ SubmitResult Orchestrator::submit_sub(const CallableIdentity &callable, TaskArgs state lives in parent-process heap (never crossed into child workers), so the ring-index addressing scheme L2 needs for shmem descriptors buys us nothing here. A monotonic `int32_t` gives ~2 billion ids per - `reset_to_empty()` interval, reset to 0 at the end of every - `Worker.run()`. + globally quiescent compaction interval. 2. **`MAX_RING_DEPTH = 4` independent shared-memory heap slabs** (Strict-1; matches L2's `PTO2_MAX_RING_DEPTH`). Each slab has its own `mmap(MAP_SHARED | MAP_ANONYMOUS)` region, bump cursor, FIFO @@ -299,7 +303,7 @@ SubmitResult Orchestrator::submit_sub(const CallableIdentity &callable, TaskArgs records its `ring_idx` and `ring_slot_idx` (position within that ring's FIFO order). `std::deque::push_back` never invalidates pointers to existing elements, so the pointer returned by `slot_state(id)` - stays valid until `reset_to_empty()` drops the whole deque. + stays valid until globally quiescent `reset_to_empty()` drops the deque. ```cpp struct AllocResult { @@ -351,12 +355,13 @@ next-oldest in-ring slot is released, walking the ring's `heap_tail` forward. Rings never touch each other — inner-scope tasks reclaim without waiting for an outer-scope task to finish. -**End-of-run reset**: `Orchestrator::drain()` waits for -`active_tasks_` to hit 0, then calls `ring.reset_to_empty()` which -drops the whole slot-state deque *and* rewinds every ring's cursors / -`released[]` / `slot_heap_end[]` back to 0. Memory per `Worker.run()` -is bounded by that run's peak alive task count; nothing accumulates -across runs. +**Run completion and compaction**: every committed slot increments its owning +`RunState.active_tasks`; `on_consumed` decrements that same run exactly once. A +run becomes terminal only after submission is closed and its count reaches +zero. Slot and heap reclamation still happens incrementally through +`ring.release(slot)`. `release_run()` may call `ring.reset_to_empty()` only +when no registered runs or live slots remain, so one run never resets storage +owned by another. **Locking**: each ring has its own `mu` / `cv`; the shared `next_task_id_` and slot deque are guarded by a separate `slots_mu_`. @@ -457,29 +462,28 @@ threshold transitions to CONSUMED inline; others stay COMPLETED or PENDING until the scheduler and consumers finish their own releases. This mirrors L2's `pto2_scope_end`. -Users who need a synchronous wait for *all* in-flight tasks must call -`drain()` (or let `Worker::run` finish — its outer `scope_end` is followed -by `drain()` before the call returns). There is deliberately no -per-scope drain primitive: the extra machinery (per-scope active counter -and cv) would only pay for itself in patterns we do not have yet. +The internal run fence, not `scope_end`, provides synchronous completion. +`Worker.run` closes its outer scope, closes submission, and waits for that +run's active count to reach zero before returning. There is deliberately no +per-scope wait primitive. --- ## 7. TensorMap -The TensorMap maps `tensor_base_ptr → current_producer_slot`. It drives -automatic dependency inference. +The TensorMap maps `(RunId, TensorKey) → current_producer_slot`. It drives +automatic dependency inference without resolving a producer from another run. ```cpp class TensorMap { public: - TaskSlot lookup(uint64_t base_ptr) const; // returns INVALID if absent - void insert(uint64_t base_ptr, TaskSlot sid); // overwrites; previous - // producer remains wire-referenced - void erase(uint64_t base_ptr); // called when producer - // reaches CONSUMED + TaskSlot lookup(RunId run_id, TensorKey key) const; + void insert(RunId run_id, TensorKey key, TaskSlot sid); + void erase_task_outputs(RunId run_id, + const std::vector &keys); private: - std::unordered_map map_; + std::mutex mu_; + std::unordered_map map_; }; ``` @@ -503,11 +507,9 @@ private: ### Thread safety -TensorMap is written only by the Orch thread (in `submit_*`) and modified by -the Scheduler thread via `erase` (on CONSUMED). Since `submit_*` and `erase` -for different entries are non-overlapping in practice, a single mutex guards -the map in the current implementation. If contention becomes a concern, a -concurrent hash map can replace it. +TensorMap is written by the Orch thread in `submit_*` and modified by the +Scheduler thread when a slot becomes CONSUMED. A mutex serializes lookup, +insert, erase, and size operations across those threads. --- @@ -577,7 +579,8 @@ Tensor Orchestrator::alloc(const std::vector &shape, DataType dtype) { // 2. Register as this slot's output so downstream tensors with the same // data pointer look up this slot as producer. uint64_t key = reinterpret_cast(ar.heap_ptr); - tensormap_.insert(key, ar.slot); + s.run_id = current_run_id; + tensormap_.insert(current_run_id, key, ar.slot); s.output_keys.push_back(key); // 3. No fanin — alloc has no work to wait on. s.fanin_count = 0; @@ -591,7 +594,7 @@ Tensor Orchestrator::alloc(const std::vector &shape, DataType dtype) { s.fanout_released = 1; // 6. Straight to COMPLETED — no dispatch needed. s.state = TaskState::COMPLETED; - active_tasks_++; + current_run.active_tasks++; return Tensor{key, shape, dtype}; } ``` @@ -695,9 +698,9 @@ instead of stalling forever. Default timeout: 10 s. ## 9. Invariants -1. **Orch is single-threaded**: only one thread ever calls `submit_*` or holds - the `Orchestrator`. No locking is needed on TensorMap, Scope, or Ring-head - for self-writes. +1. **Orch is single-threaded**: only one thread builds a run and calls + `submit_*` at a time. TensorMap still uses a mutex because Scheduler-driven + consumption erases entries concurrently. 2. **Tags are consumed at submit**: `task_args.tag(i)` is read only inside `submit_*`. Phases after submit (slot storage, dispatch, execution) do not see tags. diff --git a/docs/remote-l3-worker-design.md b/docs/remote-l3-worker-design.md index 66e774a110..607a6f9bc4 100644 --- a/docs/remote-l3-worker-design.md +++ b/docs/remote-l3-worker-design.md @@ -143,7 +143,7 @@ Relevant code paths: - Tags are stripped after submit. - `docs/comm-domain.md` - Dynamic communication domains already model deferred release after - `drain()`. + their owning run's completion fence. The Scheduler should not inspect transport details, but it does need enough metadata to avoid dispatching a task to an endpoint that cannot run it. Remote @@ -178,11 +178,10 @@ the endpoint, reports endpoint errors, and notifies the Scheduler with an explicit success/failure outcome. Directed ready queues, exact group dispatch, fanin/fanout, and ring release -remain in the common runtime. The first-error-wins policy remains only as the error -reporting policy for choosing which root error `drain()` raises. The important -change is that completion is no longer implicitly success; every endpoint, -including `LocalMailboxEndpoint`, must report an explicit success/failure -outcome. +remain in the common runtime. First-error-wins is scoped to each `RunState` and +chooses which root error that run's wait raises. Completion is not implicitly +success; every endpoint, including `LocalMailboxEndpoint`, reports an explicit +success/failure outcome. ## Fork-Safe Remote Process Model @@ -517,17 +516,17 @@ Required parent-side behavior: `task_failure` instead of reporting a successful completion. - Non-zero task or endpoint errors become candidates for the worker's first reported error. -- The worker still notifies the Scheduler so `drain()` cannot hang. +- The worker still notifies the Scheduler so the originating run can finish. - The notification carries an outcome: success, task failure, or endpoint failure. - Failed slots transition to a failed/poisoned state rather than successful `COMPLETED`. - Downstream consumers of a failed producer are marked failed/skipped and are not dispatched. -- `drain()` waits for bookkeeping and cleanup, then raises the first root - error with remote host, worker id, hashid, and sequence in the message. +- The run fence waits for bookkeeping and cleanup, then raises that run's first + root error with remote host, worker id, hashid, and sequence in the message. -Local mailbox dispatch keeps first-error-wins only for final error reporting. +Local mailbox dispatch keeps per-run first-error-wins only for final reporting. It must not mark a failed child dispatch as successful `COMPLETED`. The remote buffer path and the local adapter both use the same poisoned dependency propagation before dependent tasks are exposed to failed producer outputs. diff --git a/docs/scheduler.md b/docs/scheduler.md index ad119c33ce..d064f0a4ac 100644 --- a/docs/scheduler.md +++ b/docs/scheduler.md @@ -66,9 +66,10 @@ while (true) { } ``` -`loop_mutex` covers completion and dispatch slot access. `Orchestrator::drain` -uses the same mutex while releasing/resetting slots, preventing slot reuse -while the Scheduler still holds a reference. +`loop_mutex` covers completion and dispatch slot access. +`Orchestrator::release_run` uses the same mutex during optional globally +quiescent compaction, preventing slot removal while the Scheduler holds a +reference. ## 3. Directed NEXT_LEVEL dispatch diff --git a/docs/task-flow.md b/docs/task-flow.md index 64509c2433..6eabd6ef3e 100644 --- a/docs/task-flow.md +++ b/docs/task-flow.md @@ -369,7 +369,7 @@ The mailbox layout, fork ordering, and child loop are in | Region | Lives in | Used by | Lifetime | | ------ | -------- | ------- | -------- | -| `Ring` slot-state pool (`std::deque>`) | parent heap | Orchestrator, Scheduler, WorkerThread parent side | monotonic task-id; reset at `Worker.run` drain | +| `Ring` slot-state pool (`std::deque>`) | parent heap | Orchestrator, Scheduler, WorkerThread parent side | monotonic task-id; compacted only when globally quiescent | | `slot.task_args` (single) or `task_args_list[N]` (group, vector-backed) | parent heap | same | until slot reaches CONSUMED | | per-WT mailbox | shm MAP_SHARED | parent WorkerThread writes, child reads | lifetime of WorkerThread | | **HeapRing[0..3]** (user OUTPUT auto-alloc + `orch.alloc`) | **4 separate shm MAP_SHARED mmaps**, one per scope-layer ring | output to user code; inherited by forked children | per-ring FIFO via `rings_[r].last_alive`; scope depth picks the ring | @@ -378,9 +378,9 @@ The mailbox layout, fork ordering, and child loop are in Slot state lives inside `Ring` as `std::deque>` so `push_back` never invalidates pointers to live slots. -`ring.slot_state(id)` hands out a stable pointer for every live slot; -`drain()` calls `ring.reset_to_empty()` to drop all slot state at the -end of each `Worker.run`, bounding per-run memory. +`ring.slot_state(id)` hands out a stable pointer for every live slot. Each slot +is reclaimed individually when it reaches CONSUMED. The deque is reset only +when the worker has no registered runs or live slots. The HeapRing is **partitioned into `MAX_RING_DEPTH = 4` independent rings** (Strict-1; matches L2's `PTO2_MAX_RING_DEPTH`). Each ring is its @@ -502,10 +502,10 @@ L4 parent process | 6 | L3 child | `inner_worker.run(my_l3_orch, args, cfg)` → `scope_begin` → `my_l3_orch(orch3, ...)` | | 7 | L3 `Orchestrator.submit_sub` | `l3_sub_handle` digest dispatched to L3's own sub worker child | | 8 | L3 sub child | child resolves digest to its local Python callable and executes `verify_result()` | -| 9 | L3 drain | all L3 tasks complete; `scope_end` + `drain` return | +| 9 | L3 run fence | all L3 tasks complete; `scope_end` + `wait_run` return | | 10 | L3 child | `inner_worker.run()` returns; `_child_worker_loop` writes `TASK_DONE` | | 11 | L4 LocalMailboxEndpoint | sees `TASK_DONE`; returns success completion | -| 12 | L4 drain | L4 scope_end + drain; `w4.run()` returns | +| 12 | L4 run fence | L4 `scope_end` + `wait_run`; `w4.run()` returns | Each level's orch fn receives **its own** `Orchestrator` — the recursion is symmetric. `Worker` code does not branch on `level`; the level is only a diff --git a/python/bindings/worker_bind.h b/python/bindings/worker_bind.h index a74ee3dd84..9901979699 100644 --- a/python/bindings/worker_bind.h +++ b/python/bindings/worker_bind.h @@ -350,14 +350,25 @@ 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( - "_drain", &Orchestrator::drain, nb::call_guard(), - "Block until all submitted tasks are CONSUMED (releases GIL). " - "Rethrows the first dispatch failure seen in this run, if any." + .def("_begin_run", &Orchestrator::begin_run) + .def("_close_run_submission", &Orchestrator::close_run_submission, nb::arg("run_id")) + .def( + "_fail_run_submission", + [](Orchestrator &self, RunId run_id, const std::string &message) { + self.fail_run_submission( + run_id, + message.empty() ? std::exception_ptr{} : std::make_exception_ptr(std::runtime_error(message)) + ); + }, + nb::arg("run_id"), nb::arg("message") = std::string(), + "Close a run whose submission failed. A non-empty message becomes the run's first error." ) .def( - "_clear_error", &Orchestrator::clear_error, "Clear any stored dispatch error so the next run can proceed." - ); + "_wait_run", &Orchestrator::wait_run, nb::arg("run_id"), nb::call_guard(), + "Block until one run is terminal and raise only that run's error." + ) + .def("_run_done", &Orchestrator::run_done, nb::arg("run_id")) + .def("_release_run", &Orchestrator::release_run, nb::arg("run_id")); // --- Worker --- // Bound as `_Worker` because the Python user-facing `Worker` factory diff --git a/python/simpler/orchestrator.py b/python/simpler/orchestrator.py index 1d8ced9e32..29e3537df0 100644 --- a/python/simpler/orchestrator.py +++ b/python/simpler/orchestrator.py @@ -552,8 +552,20 @@ def _scope_begin(self) -> None: def _scope_end(self) -> None: self._o._scope_end() - def _drain(self) -> None: - self._o._drain() + def _begin_run(self) -> int: + return int(self._o._begin_run()) - def _clear_error(self) -> None: - self._o._clear_error() + def _close_run_submission(self, run_id: int) -> None: + self._o._close_run_submission(run_id) + + def _fail_run_submission(self, run_id: int, message: str = "") -> None: + self._o._fail_run_submission(run_id, message) + + def _wait_run(self, run_id: int) -> None: + self._o._wait_run(run_id) + + def _run_done(self, run_id: int) -> bool: + return bool(self._o._run_done(run_id)) + + def _release_run(self, run_id: int) -> None: + self._o._release_run(run_id) diff --git a/python/simpler/task_interface.py b/python/simpler/task_interface.py index 5c2190471a..21f3ff986a 100644 --- a/python/simpler/task_interface.py +++ b/python/simpler/task_interface.py @@ -854,7 +854,7 @@ def released(self) -> bool: def freed(self) -> bool: """True once the backend ``comm_release_domain_windows`` has executed. - Only flips after the owning ``Worker.run`` drains and processes the + Only flips after the owning ``Worker.run`` completes and processes the pending-release queue. An ``orch_fn`` will never observe ``True`` for a handle it released within the same ``run`` call. """ @@ -865,7 +865,7 @@ def release(self) -> None: Inside an orch function, this is a non-blocking mark — the actual backend ``comm_release_domain_windows`` runs after - ``Worker.run.drain()`` so that any tasks already submitted with + the owning run's completion wait so that tasks already submitted with this domain's ``device_ctx`` see live memory through execution. After this returns, the handle is treated as released for the @@ -877,7 +877,7 @@ def release(self) -> None: return self._released = True # _release_fn is owned by Worker; it queues the actual backend - # release and runs it after drain. Worker also flips _freed. + # release and runs it after the owning run completes. Worker also flips _freed. self._release_fn(self) def __enter__(self) -> CommDomainHandle: diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 3236185cd7..cfd7550113 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -2069,7 +2069,7 @@ def __init__( # ``release()`` removes them and queues a deferred backend free. self._live_domains: dict[str, CommDomainHandle] = {} # Handles whose `release()` has been called inside an orch function. - # The backend free is deferred until after Worker.run.drain() so that + # The backend free is deferred until the current Worker.run completes so that # tasks already submitted with this domain's device_ctx / buffer_ptrs # see live memory through execution. self._pending_release_domains: list[CommDomainHandle] = [] @@ -4956,7 +4956,7 @@ def _release_domain_handle(self, handle: CommDomainHandle) -> None: to have already submitted DAG tasks that capture the handle's ``device_ctx`` / ``buffer_ptrs``. Those tasks must see live memory through execution; ``Worker.run`` calls - ``_execute_pending_domain_releases`` only after ``drain()``. + ``_execute_pending_domain_releases`` only after the run-specific wait. """ if self._worker is None: return @@ -4968,7 +4968,7 @@ def _release_domain_handle(self, handle: CommDomainHandle) -> None: def _execute_pending_domain_releases(self) -> None: """Drive CTRL_RELEASE_DOMAIN for every queued handle. Must run - after ``self._orch._drain()`` so chip-side tasks have completed + after ``self._orch._wait_run()`` so chip-side tasks have completed their use of the domain memory. """ if not self._pending_release_domains: @@ -5648,48 +5648,62 @@ def _run_locked(self, callable, args, config) -> None: assert self._orch is not None assert self._worker is not None - # Drop any error stashed by a previous run() so this call starts - # clean. drain() rethrows on the way out; every successful run() - # leaves the error slot empty, but an unrelated caller may have - # poked it. - self._orch._clear_error() - self._orch._scope_begin() + run_id = self._orch._begin_run() + scope_open = False + submission_failed = True + submission_error: BaseException | None = None try: + self._orch._scope_begin() + scope_open = True callable(self._orch, args, cfg) + submission_failed = False + except BaseException as e: + submission_error = e + raise finally: - # Always release scope refs and drain so ring slots aren't - # stranded when the orch fn raises mid-DAG. drain() also - # rethrows the first dispatch failure for this run — that - # is how child-task exceptions surface to the caller of - # Worker.run(). scope_end deliberately does NOT throw: if - # it did, released refs would be incomplete and drain - # would hang on in-flight tasks. - self._orch._scope_end() - # ORDER MATTERS: drain() must complete first so any in-flight - # task that captured a now-pending handle's device_ctx / - # buffer_ptrs sees live memory. THEN execute the pending - # backend releases. Last, sweep any handles that the orch - # function neither released nor passed out (covers exception - # unwind and "forgot to release" — auto-release in LIFO). - # drain() rethrows the first chip-task/dispatch failure, so the - # cleanup lives in a finally: a failed task must not strand - # backend domain allocations into the next run. try: - try: - self._orch._drain() - except Exception as e: - self._poison_l3_l2_region_from_endpoint_error(e) - raise + if scope_open: + self._orch._scope_end() + except BaseException as e: + submission_failed = True + if submission_error is None: + submission_error = e + raise finally: - self._release_active_remote_slot_refs() - self._flush_pending_remote_frees() try: - self._cleanup_l3_l2_regions() + # Closing the run must not be able to strand `run_id` in the + # engine's registry: `_release_run` and the run-owned + # resource cleanup below both live under this `try`. + if submission_failed: + self._orch._fail_run_submission( + run_id, + "" if submission_error is None else _format_exc("orchestration", submission_error), + ) + else: + self._orch._close_run_submission(run_id) + try: + self._orch._wait_run(run_id) + except Exception as e: + self._poison_l3_l2_region_from_endpoint_error(e) + # Preserve a synchronous graph-construction exception + # already propagating from the orchestration callback. + if not submission_failed: + raise finally: - self._l3_l2_orch_comm_host_buffers.clear() - self._execute_pending_domain_releases() - if self._live_domains: - self._release_all_live_domains() + try: + # Run-owned resources remain live until the per-run + # completion fence fires, even on failure. + self._release_active_remote_slot_refs() + self._flush_pending_remote_frees() + try: + self._cleanup_l3_l2_regions() + finally: + self._l3_l2_orch_comm_host_buffers.clear() + self._execute_pending_domain_releases() + if self._live_domains: + self._release_all_live_domains() + finally: + self._orch._release_run(run_id) # L3+ returns None like every other worker level; per-L2-child timing # is emitted as `[STRACE]` markers from each simpler_run. return None diff --git a/src/common/hierarchical/orchestrator.cpp b/src/common/hierarchical/orchestrator.cpp index dda1db8107..c462998fbd 100644 --- a/src/common/hierarchical/orchestrator.cpp +++ b/src/common/hierarchical/orchestrator.cpp @@ -12,6 +12,7 @@ #include "orchestrator.h" #include +#include #include #include @@ -28,7 +29,187 @@ void Orchestrator::init( ready_next_level_queues_ = ready_next_level_queues; manager_ = manager; ready_notify_cb_ = std::move(ready_notify_cb); - active_tasks_.store(0, std::memory_order_relaxed); +} + +bool Orchestrator::is_terminal(RunPhase phase) { return phase == RunPhase::COMPLETED || phase == RunPhase::FAILED; } + +std::shared_ptr Orchestrator::find_run(RunId run_id) const { + std::lock_guard lk(runs_mu_); + auto it = runs_.find(run_id); + return it == runs_.end() ? nullptr : it->second; +} + +std::shared_ptr Orchestrator::get_run(RunId run_id) const { + auto run = find_run(run_id); + if (run == nullptr) throw std::invalid_argument("Orchestrator: unknown run id"); + return run; +} + +std::shared_ptr Orchestrator::current_building_run() const { + std::lock_guard lk(runs_mu_); + auto it = runs_.find(building_run_id_); + if (building_run_id_ == INVALID_RUN_ID || it == runs_.end()) { + throw std::logic_error("Orchestrator: task submission requires an active run"); + } + return it->second; +} + +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 = next_run_id_++; + runs_.emplace(run_id, std::make_shared(run_id)); + building_run_id_ = run_id; + return run_id; +} + +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))) { + 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(); +} + +void Orchestrator::close_run_submission(RunId run_id) { + auto run = get_run(run_id); + { + std::lock_guard runs_lk(runs_mu_); + if (building_run_id_ != run_id) { + throw std::logic_error("Orchestrator::close_run_submission: run is not building"); + } + building_run_id_ = INVALID_RUN_ID; + } + { + std::lock_guard lk(run->completion_mu); + if (run->submission_closed) { + 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); + } + } + finish_run_if_ready(run); +} + +void Orchestrator::fail_run_submission(RunId run_id, std::exception_ptr error) { + auto run = get_run(run_id); + if (error) record_run_error(run_id, std::move(error)); + { + std::lock_guard runs_lk(runs_mu_); + if (building_run_id_ != run_id) { + throw std::logic_error("Orchestrator::fail_run_submission: run is not building"); + } + building_run_id_ = INVALID_RUN_ID; + } + { + 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); + } + } + finish_run_if_ready(run); +} + +void Orchestrator::wait_run(RunId run_id) { + auto run = get_run(run_id); + std::exception_ptr error; + { + std::unique_lock lk(run->completion_mu); + run->completion_cv.wait(lk, [&run] { + return is_terminal(run->phase.load(std::memory_order_acquire)); + }); + error = run->first_error; + } + if (error) std::rethrow_exception(error); +} + +bool Orchestrator::run_done(RunId run_id) const { + return is_terminal(get_run(run_id)->phase.load(std::memory_order_acquire)); +} + +bool Orchestrator::run_failed(RunId run_id) const { + auto run = get_run(run_id); + std::lock_guard lk(run->completion_mu); + return run->submission_failed || static_cast(run->first_error); +} + +bool Orchestrator::quiescent_locked() const { return runs_.empty() && building_run_id_ == INVALID_RUN_ID; } + +void Orchestrator::compact_if_quiescent() { + // `begin_run` takes only runs_mu_, and no slot is allocated without a + // building run, so holding runs_mu_ across both the test and the reset is + // what stops a run registered mid-compaction from losing its first slot. + // Lock order is the scheduler loop mutex then runs_mu_, matching the + // scheduler's on_consumed path. + std::unique_lock sched_lk; + if (sched_loop_mu_ != nullptr) sched_lk = std::unique_lock(*sched_loop_mu_); + std::lock_guard runs_lk(runs_mu_); + if (quiescent_locked() && allocator_->active_count() == 0) allocator_->reset_to_empty(); +} + +void Orchestrator::release_run(RunId run_id) { + { + std::lock_guard lk(runs_mu_); + auto it = runs_.find(run_id); + if (it == runs_.end()) throw std::invalid_argument("Orchestrator::release_run: unknown run id"); + if (!is_terminal(it->second->phase.load(std::memory_order_acquire))) { + throw std::logic_error("Orchestrator::release_run: run is not terminal"); + } + runs_.erase(it); + } + + 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::decrement_run_tasks(RunId run_id) { + auto run = find_run(run_id); + if (run == nullptr) return; + int32_t remaining = run->active_tasks.fetch_sub(1, std::memory_order_acq_rel) - 1; + if (remaining < 0) { + // Reaching the fence is what lets the caller observe the failure, so a + // count mismatch is surfaced as the run's error rather than thrown: + // this runs on the scheduler thread, where an escaping exception is + // fatal to the process. + run->active_tasks.store(0, std::memory_order_release); + record_run_error(run, std::make_exception_ptr(std::logic_error("Orchestrator: run task count underflow"))); + } + if (remaining <= 0) finish_run_if_ready(run); +} + +void Orchestrator::record_run_error(const std::shared_ptr &run, std::exception_ptr error) { + if (run == nullptr || !error) return; + std::lock_guard lk(run->completion_mu); + if (!run->first_error) run->first_error = std::move(error); +} + +void Orchestrator::record_run_error(RunId run_id, std::exception_ptr error) { + if (!error) return; + record_run_error(find_run(run_id), std::move(error)); +} + +void Orchestrator::report_task_error(TaskSlot slot, const std::string &message) { + TaskSlotState &task = slot_state(slot); + record_run_error(task.run_id, std::make_exception_ptr(std::runtime_error(message))); } uint64_t Orchestrator::malloc(int worker_id, size_t size) { @@ -68,6 +249,7 @@ TaskSlotState &Orchestrator::slot_state(TaskSlot s) { uint64_t Orchestrator::output_alloc_bytes(const Tensor &t) { return align_up(t.nbytes(), HEAP_ALIGN); } Tensor Orchestrator::alloc(const std::vector &shape, DataType dtype) { + auto run = current_building_run(); if (shape.empty()) { // Rank-0 tensors are not supported across the ABI (Tensor enforces // ndims > 0). Reject here so we never allocate + register a buffer in @@ -100,11 +282,12 @@ 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(key, ar.slot); + tensormap_->insert(run->id, key, ar.slot); s.output_keys.push_back(key); } @@ -128,7 +311,7 @@ Tensor Orchestrator::alloc(const std::vector &shape, DataType dtype) { s.state.store(TaskState::COMPLETED, std::memory_order_release); - active_tasks_.fetch_add(1, std::memory_order_relaxed); + 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 @@ -189,38 +372,30 @@ SubmitResult Orchestrator::submit_impl( std::vector target_worker_ids, std::vector> eligible_worker_ids, std::vector remote_sidecars ) { + auto run = current_building_run(); if (args_list.empty()) throw std::invalid_argument("Orchestrator: args_list must not be empty"); config.validate(); validate_worker_eligibility(worker_type, args_list.size(), target_worker_ids, eligible_worker_ids); validate_remote_sidecars(args_list, remote_sidecars, eligible_worker_ids); - // Fail-fast: if a previously-dispatched task has already failed, abort - // this submit before any bookkeeping so the orch fn unwinds promptly - // and no further work is queued. Tasks already in flight run to - // completion; drain() picks up any remaining bookkeeping and rethrows - // at the finally: _drain() site in Worker.run. - if (manager_ && manager_->has_error()) { - std::rethrow_exception(manager_->take_error()); + { + std::lock_guard lk(run->completion_mu); + if (run->first_error) std::rethrow_exception(run->first_error); } - // Track this submission for drain() before any allocations so the count - // is incremented exactly once per submitted DAG node, regardless of the - // group_size N. - active_tasks_.fetch_add(1, std::memory_order_relaxed); - // --- Step 1: Atomically claim slot + auto-alloc any OUTPUT tensors that // arrived with a null data pointer. Both resources come from the same // merged allocator (Strict-2) so there is no partial-failure rollback // path. AllocResult ar = reserve_outputs_and_slot(args_list, remote_sidecars); if (ar.slot == INVALID_SLOT) { - active_tasks_.fetch_sub(1, std::memory_order_relaxed); throw std::runtime_error("Orchestrator: allocator shutdown"); } TaskSlot slot = ar.slot; TaskSlotState &s = slot_state(slot); s.reset(); + s.run_id = run->id; s.worker_type = worker_type; s.callable = callable; @@ -288,10 +463,12 @@ SubmitResult Orchestrator::submit_impl( s.fanout_released.store(0, std::memory_order_relaxed); if (scope_ref > 0) scope_->register_task(slot); + increment_run_tasks(run->id); 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); @@ -541,6 +718,7 @@ void Orchestrator::infer_deps( const std::vector &remote_sidecars, std::vector &producers, std::vector &output_keys ) { + RunId run_id = slot_state(slot).run_id; std::unordered_set producer_seen; size_t tensor_count_hint = 0; for (const TaskArgs &args : args_list) { @@ -600,20 +778,20 @@ void Orchestrator::infer_deps( TensorArgType tag = a.tag(i); switch (tag) { case TensorArgType::INPUT: { - TaskSlot prod = tensormap_->lookup(key); + TaskSlot prod = tensormap_->lookup(run_id, key); if (prod != INVALID_SLOT) add_unique_producer(prod); break; } case TensorArgType::INOUT: { - TaskSlot prod = tensormap_->lookup(key); + TaskSlot prod = tensormap_->lookup(run_id, key); if (prod != INVALID_SLOT) add_unique_producer(prod); - tensormap_->insert(key, slot); + tensormap_->insert(run_id, key, slot); output_keys.push_back(key); break; } case TensorArgType::OUTPUT: case TensorArgType::OUTPUT_EXISTING: { - tensormap_->insert(key, slot); + tensormap_->insert(run_id, key, slot); output_keys.push_back(key); break; } @@ -629,7 +807,10 @@ void Orchestrator::infer_deps( // Scope // ============================================================================= -void Orchestrator::scope_begin() { scope_->scope_begin(); } +void Orchestrator::scope_begin() { + (void)current_building_run(); + scope_->scope_begin(); +} void Orchestrator::scope_end() { scope_->scope_end([this](TaskSlot slot) { @@ -680,57 +861,13 @@ bool Orchestrator::on_consumed(TaskSlot slot) { } } - tensormap_->erase_task_outputs(s.output_keys); + RunId run_id = s.run_id; + tensormap_->erase_task_outputs(run_id, s.output_keys); // HeapRing-owned OUTPUT slabs are reclaimed implicitly when the allocator // advances last_alive past this slot — no per-slot munmap needed. allocator_->release(slot); - // Decrement active-task counter so drain() observes completion. Gated - // on the CAS win so both consume paths — release_ref (Orch thread, - // scope_end) and try_consume (scheduler thread, consumer's deferred - // release) — decrement exactly once. Notify drain_cv when the count - // hits zero. - int32_t remaining = active_tasks_.fetch_sub(1, std::memory_order_acq_rel) - 1; - if (remaining == 0) { - std::lock_guard lk(drain_mu_); - drain_cv_.notify_all(); - } + decrement_run_tasks(run_id); return true; } - -void Orchestrator::drain() { - { - std::unique_lock lk(drain_mu_); - drain_cv_.wait(lk, [this] { - return active_tasks_.load(std::memory_order_acquire) == 0; - }); - } - // Every slot is CONSUMED (active_tasks_ == 0 ⇒ allocator last_alive_ == - // next_task_id_). Drop all per-slot state so the next Worker.run() - // starts from task_id = 0 with no accumulated memory. - // - // Hold the scheduler's loop mutex across the reset: active_tasks_ can reach - // zero while the scheduler thread is still inside on_task_complete (it reads - // the slot after the consume that drives the count to 0). Freeing the slots - // here without this guard is a heap-use-after-free. The scheduler releases - // loop_mu_ only between iterations, and with active_tasks_ == 0 it has no - // further slots to touch, so this blocks at most one in-flight iteration. - if (sched_loop_mu_ != nullptr) { - std::lock_guard sched_lk(*sched_loop_mu_); - allocator_->reset_to_empty(); - } else { - allocator_->reset_to_empty(); - } - - // Rethrow the first dispatch failure seen during this run. Deferred to - // after allocator reset so the next Worker.run() can proceed cleanly - // once clear_error() is called. - if (manager_ && manager_->has_error()) { - std::rethrow_exception(manager_->take_error()); - } -} - -void Orchestrator::clear_error() { - if (manager_) manager_->clear_error(); -} diff --git a/src/common/hierarchical/orchestrator.h b/src/common/hierarchical/orchestrator.h index cd6d266065..e27d1ef9e6 100644 --- a/src/common/hierarchical/orchestrator.h +++ b/src/common/hierarchical/orchestrator.h @@ -36,6 +36,7 @@ #include #include #include +#include #include #include "../task_interface/call_config.h" @@ -115,6 +116,15 @@ class Orchestrator { // Submit a group of SUB tasks: N args -> N workers, 1 DAG node. SubmitResult submit_sub_group(const CallableIdentity &callable, const std::vector &args_list); + // Only the calling orchestration thread builds a run at a time. + RunId begin_run(); + void close_run_submission(RunId run_id); + void fail_run_submission(RunId run_id, std::exception_ptr error = nullptr); + void wait_run(RunId run_id); + bool run_done(RunId run_id) const; + bool run_failed(RunId run_id) const; + void release_run(RunId run_id); + // Open a nested scope. Every task submitted between this call and the // matching `scope_end()` picks a heap ring based on the current scope // depth (`min(depth, MAX_RING_DEPTH - 1)`) so its slab reclaims @@ -125,27 +135,17 @@ class Orchestrator { // Non-blocking: `scope_end` walks the scope's tasks and releases one // ref per task, returning immediately. Actual CONSUMED transitions // happen asynchronously as each task's consumer count reaches - // threshold (mirrors L2's `pto2_scope_end`). Callers that need a - // synchronous wait must call `drain()` separately. + // threshold (mirrors L2's `pto2_scope_end`). The owning run fence + // provides the synchronous completion boundary. void scope_begin(); void scope_end(); - // Block until every submitted task has reached CONSUMED. Invoked by - // Worker::run after scope_end; not part of the user-facing orch-fn API. - // Rethrows the first exception reported by any WorkerThread during - // dispatch (fail-fast): the wait itself is unaffected — every in-flight - // task is allowed to finish so ring slots aren't leaked. - void drain(); - - // Wire the Scheduler's loop mutex so drain() can hold it across - // reset_to_empty(), preventing the scheduler thread from touching slot - // state mid-teardown (heap-use-after-free). Set once by Worker after the - // Scheduler is constructed. + // Wire the Scheduler's loop mutex so release_run() can safely perform an + // optional allocator compaction when the whole worker is quiescent. void set_scheduler_loop_mutex(std::mutex *m) { sched_loop_mu_ = m; } - // Clear any stored dispatch error so the next Worker::run() starts - // from a clean slate. Called by Worker::run before scope_begin. - void clear_error(); + // Attach a scheduler/endpoint failure to the task's originating run. + void report_task_error(TaskSlot slot, const std::string &message); // Called by Scheduler (via Worker) when a task becomes CONSUMED: // erases TensorMap entries, releases the allocator slot (and implicitly @@ -168,14 +168,31 @@ class Orchestrator { ReadyQueue *ready_sub_queue_ = nullptr; NextLevelReadyQueues *ready_next_level_queues_ = nullptr; - // --- Drain support (owned here, not on Worker) --- - std::atomic active_tasks_{0}; - std::mutex drain_mu_; - std::condition_variable drain_cv_; - // Scheduler's loop mutex (not owned). Held across reset_to_empty() in - // drain() so the scheduler can't be mid-on_task_complete during teardown. + mutable std::mutex runs_mu_; + std::unordered_map> runs_; + RunId next_run_id_{1}; + RunId building_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. std::mutex *sched_loop_mu_{nullptr}; + // Returns nullptr for an unknown id. Bookkeeping reached from the + // scheduler and worker threads uses this instead of `get_run` — an + // exception escaping those threads terminates the process. + 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); + static bool is_terminal(RunPhase phase); + // Callers hold runs_mu_. + bool quiescent_locked() const; + void compact_if_quiescent(); + void increment_run_tasks(RunId run_id); + void decrement_run_tasks(RunId run_id); + static void record_run_error(const std::shared_ptr &run, std::exception_ptr error); + void record_run_error(RunId run_id, std::exception_ptr error); + // Slot state lives in the Ring; the pointer stays stable for the // slot's lifetime. Throws if the id is out of range — callers that // hold a recently-allocated slot id should always get a valid pointer. diff --git a/src/common/hierarchical/ring.cpp b/src/common/hierarchical/ring.cpp index e8ac6e5b2e..5b31b769fc 100644 --- a/src/common/hierarchical/ring.cpp +++ b/src/common/hierarchical/ring.cpp @@ -201,7 +201,7 @@ void Ring::reset_to_empty() { if (r.last_alive != static_cast(r.released.size())) { throw std::logic_error( "Ring::reset_to_empty: tasks still live on at least one ring. " - "Did drain() complete?" + "Did every registered run complete?" ); } } diff --git a/src/common/hierarchical/ring.h b/src/common/hierarchical/ring.h index 0e0b9da8b8..497f77a583 100644 --- a/src/common/hierarchical/ring.h +++ b/src/common/hierarchical/ring.h @@ -130,8 +130,7 @@ class Ring { // Rewind every ring's cursors + released/slot_heap_end vectors and drop // all slot states. Requires that no slots are currently live - // (`active_count() == 0`) — typically called by `Orchestrator::drain` - // right after the active count hits zero. + // (`active_count() == 0`) and no registered runs remain. void reset_to_empty(); int32_t active_count() const; diff --git a/src/common/hierarchical/scheduler.cpp b/src/common/hierarchical/scheduler.cpp index fcba6e7a65..acd6e22420 100644 --- a/src/common/hierarchical/scheduler.cpp +++ b/src/common/hierarchical/scheduler.cpp @@ -72,6 +72,10 @@ void Scheduler::stop() { void Scheduler::worker_done(WorkerCompletion completion) { TaskSlotState &s = *cfg_.ring->slot_state(completion.task_slot); + bool failure_reported = is_failure(completion.outcome); + if (failure_reported && cfg_.on_task_failed_cb) { + cfg_.on_task_failed_cb(completion.task_slot, completion.error_message); + } // Group aggregation: only push to completion queue when ALL workers done if (s.is_group()) { @@ -146,6 +150,10 @@ void Scheduler::worker_done(WorkerCompletion completion) { completion = std::move(terminal); } + if (!failure_reported && is_failure(completion.outcome) && cfg_.on_task_failed_cb) { + cfg_.on_task_failed_cb(completion.task_slot, completion.error_message); + } + { std::lock_guard lk(completion_mu_); completion_queue_.push(std::move(completion)); @@ -173,9 +181,9 @@ void Scheduler::run() { }); } - // Hold loop_mu_ across the entire slot-touching body so drain()'s - // reset_to_empty() cannot free TaskSlotStates while on_task_complete / - // dispatch_ready are still reading them (heap-use-after-free). + // Hold loop_mu_ across the entire slot-touching body so quiescent + // compaction cannot free TaskSlotStates while on_task_complete or + // dispatch_ready is still reading them. std::lock_guard loop_lk(loop_mu_); // Phase 1: drain completions diff --git a/src/common/hierarchical/scheduler.h b/src/common/hierarchical/scheduler.h index 328dd9a57c..00c1abfd27 100644 --- a/src/common/hierarchical/scheduler.h +++ b/src/common/hierarchical/scheduler.h @@ -61,6 +61,9 @@ class Scheduler { std::function enqueue_ready_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; }; void start(const Config &cfg); @@ -77,9 +80,9 @@ class Scheduler { void notify_ready(); // Mutex held by run() across each loop iteration's slot-touching body - // (completion processing + dispatch). Orchestrator::drain() acquires it - // before Ring::reset_to_empty() so the ring can't be torn down while the - // scheduler thread is mid-on_task_complete (heap-use-after-free). + // (completion processing + dispatch). Orchestrator::release_run() acquires + // it before optional Ring::reset_to_empty() compaction so the ring cannot + // be torn down while the scheduler thread is mid-on_task_complete. std::mutex &loop_mutex() { return loop_mu_; } private: diff --git a/src/common/hierarchical/tensormap.cpp b/src/common/hierarchical/tensormap.cpp index 3e7fa9e4e8..b090a65553 100644 --- a/src/common/hierarchical/tensormap.cpp +++ b/src/common/hierarchical/tensormap.cpp @@ -11,17 +11,25 @@ #include "tensormap.h" -TaskSlot TensorMap::lookup(TensorKey key) const { - auto it = map_.find(key); +TaskSlot TensorMap::lookup(RunId run_id, TensorKey key) const { + std::lock_guard lk(mu_); + auto it = map_.find(RunTensorKey{run_id, key}); if (it == map_.end()) return INVALID_SLOT; return it->second; } -void TensorMap::insert(TensorKey key, TaskSlot producer) { map_[key] = producer; } +void TensorMap::insert(RunId run_id, TensorKey key, TaskSlot producer) { + std::lock_guard lk(mu_); + map_[RunTensorKey{run_id, key}] = producer; +} -void TensorMap::erase_task_outputs(const std::vector &keys) { +void TensorMap::erase_task_outputs(RunId run_id, const std::vector &keys) { + std::lock_guard lk(mu_); for (const auto &key : keys) - map_.erase(key); + map_.erase(RunTensorKey{run_id, key}); } -int32_t TensorMap::size() const { return static_cast(map_.size()); } +int32_t TensorMap::size() const { + std::lock_guard lk(mu_); + return static_cast(map_.size()); +} diff --git a/src/common/hierarchical/tensormap.h b/src/common/hierarchical/tensormap.h index 326970e269..f74a018b4e 100644 --- a/src/common/hierarchical/tensormap.h +++ b/src/common/hierarchical/tensormap.h @@ -10,7 +10,7 @@ */ /** - * TensorMap — TensorKey → producer task slot mapping. + * TensorMap — (RunId, TensorKey) → producer task slot mapping. * * At the hierarchical host level, every tensor is identified by a TensorKey * consisting of (ptr, worker_id). Host tensors (HeapRing) use @@ -28,6 +28,7 @@ #pragma once +#include #include #include @@ -37,19 +38,36 @@ class TensorMap { public: // Look up the producer for a tensor key. // Returns INVALID_SLOT when not found. - TaskSlot lookup(TensorKey key) const; + TaskSlot lookup(RunId run_id, TensorKey key) const; // Register key → producer mapping. // Overwrites any existing entry (re-use of the same buffer by a new producer). - void insert(TensorKey key, TaskSlot producer); + void insert(RunId run_id, TensorKey key, TaskSlot producer); // Remove all entries whose key appears in 'keys'. // Called when a producer task transitions to CONSUMED. - void erase_task_outputs(const std::vector &keys); + void erase_task_outputs(RunId run_id, const std::vector &keys); // Number of entries currently tracked. int32_t size() const; private: - std::unordered_map map_; + struct RunTensorKey { + RunId run_id{INVALID_RUN_ID}; + TensorKey tensor{}; + + bool operator==(const RunTensorKey &other) const { return run_id == other.run_id && tensor == other.tensor; } + }; + + struct RunTensorKeyHash { + size_t operator()(const RunTensorKey &key) const { + size_t h = std::hash{}(key.run_id); + size_t tensor_hash = TensorKeyHash{}(key.tensor); + h ^= tensor_hash + 0x9e3779b9 + (h << 6) + (h >> 2); + return h; + } + }; + + mutable std::mutex mu_; + std::unordered_map map_; }; diff --git a/src/common/hierarchical/types.cpp b/src/common/hierarchical/types.cpp index 7d060c5e02..3fcebcd85a 100644 --- a/src/common/hierarchical/types.cpp +++ b/src/common/hierarchical/types.cpp @@ -19,6 +19,7 @@ void TaskSlotState::reset() { state.store(TaskState::FREE, std::memory_order_relaxed); + run_id = INVALID_RUN_ID; fanin_count = 0; fanin_released.store(0, std::memory_order_relaxed); { diff --git a/src/common/hierarchical/types.h b/src/common/hierarchical/types.h index 346757c561..5819f88ae3 100644 --- a/src/common/hierarchical/types.h +++ b/src/common/hierarchical/types.h @@ -30,7 +30,9 @@ #include #include +#include #include +#include #include #include #include @@ -118,6 +120,34 @@ static constexpr int32_t MAX_RING_DEPTH = 4; 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. +enum class RunPhase : int32_t { + BUILDING = 0, + PREPARED = 1, + EXECUTING = 2, + COMPLETED = 3, + FAILED = 4, +}; + +struct RunState { + explicit RunState(RunId run_id) : + id(run_id) {} + + RunId id{INVALID_RUN_ID}; + std::atomic phase{RunPhase::BUILDING}; + std::atomic active_tasks{0}; + mutable std::mutex completion_mu; + std::condition_variable completion_cv; + std::exception_ptr first_error; + bool submission_closed{false}; + bool submission_failed{false}; +}; + // ============================================================================= // Task slot index type // ============================================================================= @@ -290,6 +320,7 @@ 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}; @@ -318,7 +349,7 @@ struct TaskSlotState { // --- Failure state --- // Root worker failures and downstream poison both land here. The - // WorkerManager still owns first-error-wins reporting for drain(). + // originating RunState owns first-error-wins reporting. std::string failure_message; // --- Task data (stored on parent heap, lives until slot CONSUMED) --- diff --git a/src/common/hierarchical/worker.cpp b/src/common/hierarchical/worker.cpp index e061248eea..f5f52b01ee 100644 --- a/src/common/hierarchical/worker.cpp +++ b/src/common/hierarchical/worker.cpp @@ -119,10 +119,12 @@ void Worker::init() { cfg.on_consumed_cb = [this](TaskSlot slot) { orchestrator_.on_consumed(slot); }; + cfg.on_task_failed_cb = [this](TaskSlot slot, const std::string &message) { + orchestrator_.report_task_error(slot, message); + }; scheduler_.start(cfg); - // Let drain() hold the scheduler's loop mutex across ring teardown so slots - // aren't freed while the scheduler thread is mid-on_task_complete. + // Allocator compaction and scheduler slot access share this mutex. orchestrator_.set_scheduler_loop_mutex(&scheduler_.loop_mutex()); initialized_ = true; } diff --git a/src/common/hierarchical/worker_manager.cpp b/src/common/hierarchical/worker_manager.cpp index 5044b37f6a..ecef136aed 100644 --- a/src/common/hierarchical/worker_manager.cpp +++ b/src/common/hierarchical/worker_manager.cpp @@ -166,12 +166,10 @@ void LocalMailboxEndpoint::shutdown_child() { write_mailbox_state(MailboxState:: // ============================================================================= void WorkerThread::start( - Ring *ring, WorkerManager *manager, const std::function &on_complete, - std::unique_ptr endpoint + Ring *ring, const std::function &on_complete, std::unique_ptr endpoint ) { if (!endpoint) throw std::invalid_argument("WorkerThread::start: null endpoint"); ring_ = ring; - manager_ = manager; on_complete_ = on_complete; endpoint_ = std::move(endpoint); shutdown_ = false; @@ -238,10 +236,6 @@ void WorkerThread::loop() { completion.error_message = "WorkerThread endpoint failed with unknown exception"; } - if (completion.outcome != EndpointOutcome::SUCCESS && manager_) { - manager_->report_error(std::make_exception_ptr(std::runtime_error(completion.error_message))); - } - idle_.store(true, std::memory_order_release); on_complete_(std::move(completion)); } @@ -398,7 +392,7 @@ void WorkerManager::start(Ring *ring, const OnCompleteFn &on_complete) { for (const auto &entry : next_level_entries_) { auto wt = std::make_unique(); auto endpoint = std::make_unique(entry.worker_id, entry.mailbox); - wt->start(ring, this, on_complete, std::move(endpoint)); + wt->start(ring, on_complete, std::move(endpoint)); next_level_threads_.push_back(std::move(wt)); } }; @@ -407,39 +401,20 @@ void WorkerManager::start(Ring *ring, const OnCompleteFn &on_complete) { for (size_t i = 0; i < entries.size(); ++i) { auto wt = std::make_unique(); auto endpoint = std::make_unique(static_cast(i), entries[i]); - wt->start(ring, this, on_complete, std::move(endpoint)); + wt->start(ring, on_complete, std::move(endpoint)); threads.push_back(std::move(wt)); } }; make_next_level_threads(); for (auto &endpoint : next_level_endpoint_entries_) { auto wt = std::make_unique(); - wt->start(ring, this, on_complete, std::move(endpoint)); + wt->start(ring, on_complete, std::move(endpoint)); next_level_threads_.push_back(std::move(wt)); } next_level_endpoint_entries_.clear(); make_sub_threads(sub_entries_, sub_threads_); } -void WorkerManager::report_error(std::exception_ptr e) { - if (!e) return; - std::lock_guard lk(err_mu_); - if (first_error_) return; // first-error-wins - first_error_ = std::move(e); - has_error_.store(true, std::memory_order_release); -} - -std::exception_ptr WorkerManager::take_error() { - std::lock_guard lk(err_mu_); - return first_error_; -} - -void WorkerManager::clear_error() { - std::lock_guard lk(err_mu_); - first_error_ = nullptr; - has_error_.store(false, std::memory_order_release); -} - void WorkerManager::stop() { for (auto &wt : next_level_threads_) wt->stop(); diff --git a/src/common/hierarchical/worker_manager.h b/src/common/hierarchical/worker_manager.h index 7d4cfe1fd8..49567f079c 100644 --- a/src/common/hierarchical/worker_manager.h +++ b/src/common/hierarchical/worker_manager.h @@ -345,11 +345,8 @@ class WorkerThread { // `ring->slot_state(task_slot)` on each dispatch. // on_complete(completion) is called (in the WorkerThread) after each // endpoint run(). - // `manager` is a borrowed pointer used to report dispatch failures - // (exception_ptr routed out of the worker thread to the orch thread). void start( - Ring *ring, WorkerManager *manager, const std::function &on_complete, - std::unique_ptr endpoint + Ring *ring, const std::function &on_complete, std::unique_ptr endpoint ); // Enqueue a dispatch for the worker. Non-blocking. @@ -437,7 +434,6 @@ class WorkerThread { private: Ring *ring_{nullptr}; - WorkerManager *manager_{nullptr}; std::unique_ptr endpoint_; std::function on_complete_; @@ -542,14 +538,6 @@ class WorkerManager { double timeout_s ); - // Error propagation: first dispatch failure from any WorkerThread wins. - // The orch thread inspects via `has_error()` / `take_error()` and - // clears between Worker.run() invocations via `clear_error()`. - void report_error(std::exception_ptr e); - bool has_error() const { return has_error_.load(std::memory_order_acquire); } - std::exception_ptr take_error(); - void clear_error(); - private: struct LocalNextLevelEntry { int32_t worker_id{-1}; @@ -561,11 +549,4 @@ class WorkerManager { std::vector> next_level_threads_; std::vector> sub_threads_; - - // First-error-wins exception slot. Written under err_mu_ by - // WorkerThread::loop() catch handlers; read by the orch thread at - // submit_*/drain boundaries. - std::atomic has_error_{false}; - mutable std::mutex err_mu_; - std::exception_ptr first_error_; }; diff --git a/tests/ut/cpp/hierarchical/test_orchestrator.cpp b/tests/ut/cpp/hierarchical/test_orchestrator.cpp index f456d18e58..428596ddbd 100644 --- a/tests/ut/cpp/hierarchical/test_orchestrator.cpp +++ b/tests/ut/cpp/hierarchical/test_orchestrator.cpp @@ -33,6 +33,7 @@ struct OrchestratorFixture : public ::testing::Test { ReadyQueue rq_sub; Orchestrator orch; CallConfig cfg; + RunId run_id{INVALID_RUN_ID}; // Tests in this file only submit NEXT_LEVEL tasks, so `rq` is a // convenience alias for the next-level queue. Kept public so existing @@ -47,6 +48,7 @@ struct OrchestratorFixture : public ::testing::Test { allocator.init(/*heap_bytes=*/1ULL << 20); rq_next_level.reset({0, 1, 3}); orch.init(&tm, &allocator, &scope, &rq_sub, &rq_next_level); + run_id = orch.begin_run(); } void TearDown() override { allocator.shutdown(); } @@ -166,7 +168,7 @@ TEST_F(OrchestratorFixture, TensorMapTracksProducer) { TaskSlot drain_slot; rq.try_pop(drain_slot); - EXPECT_EQ(tm.lookup(TensorKey{0x1234, -1}), a.task_slot); + EXPECT_EQ(tm.lookup(run_id, TensorKey{0x1234, -1}), a.task_slot); } TEST_F(OrchestratorFixture, OnConsumedCleansUpTensorMap) { @@ -175,12 +177,12 @@ TEST_F(OrchestratorFixture, OnConsumedCleansUpTensorMap) { TaskSlot slot; rq.try_pop(slot); - EXPECT_EQ(tm.lookup(TensorKey{0x42, -1}), slot); + EXPECT_EQ(tm.lookup(run_id, TensorKey{0x42, -1}), slot); S(slot).state.store(TaskState::COMPLETED, std::memory_order_relaxed); orch.on_consumed(slot); - EXPECT_EQ(tm.lookup(TensorKey{0x42, -1}), INVALID_SLOT); + EXPECT_EQ(tm.lookup(run_id, TensorKey{0x42, -1}), INVALID_SLOT); EXPECT_EQ(S(slot).state.load(), TaskState::CONSUMED); } @@ -237,8 +239,8 @@ TEST_F(OrchestratorFixture, GroupTaskStoresArgsListPerMember) { EXPECT_EQ(S(res.task_slot).args_view(1).tensors(0).buffer.addr, 0xA1u); // Both keys registered as producers for the group slot. - EXPECT_EQ(tm.lookup(TensorKey{0xA0, -1}), res.task_slot); - EXPECT_EQ(tm.lookup(TensorKey{0xA1, -1}), res.task_slot); + EXPECT_EQ(tm.lookup(run_id, TensorKey{0xA0, -1}), res.task_slot); + EXPECT_EQ(tm.lookup(run_id, TensorKey{0xA1, -1}), res.task_slot); } TEST_F(OrchestratorFixture, SingleTaskStoresTaskArgsDirectly) { @@ -274,7 +276,7 @@ TEST_F(OrchestratorFixture, OutputAutoAllocsFromHeapRing) { EXPECT_LT(data, base + allocator.heap_size(0)); EXPECT_EQ(data % HEAP_ALIGN, 0u); - EXPECT_EQ(tm.lookup(TensorKey{data, -1}), res.task_slot); + EXPECT_EQ(tm.lookup(run_id, TensorKey{data, -1}), res.task_slot); } TEST_F(OrchestratorFixture, RemoteOutputSidecarSkipsLocalAutoAllocAndRegistersRemoteKey) { @@ -301,7 +303,7 @@ TEST_F(OrchestratorFixture, RemoteOutputSidecarSkipsLocalAutoAllocAndRegistersRe EXPECT_EQ(S(res.task_slot).task_args.tensor(0).buffer.addr, 0u); TensorKey key = TensorKey::remote_buffer(TensorAddressKind::REMOTE_BUFFER, 3, 9, 2, 64); - EXPECT_EQ(tm.lookup(key), res.task_slot); + EXPECT_EQ(tm.lookup(run_id, key), res.task_slot); } TEST_F(OrchestratorFixture, RemoteBarePayloadFailsBeforeSlotCommit) { @@ -388,7 +390,7 @@ TEST_F(OrchestratorFixture, InoutWiresCreatorAsFanin) { rq.try_pop(writer_slot); // TensorMap now points at the new writer. - EXPECT_EQ(tm.lookup(TensorKey{0xFEED, -1}), writer.task_slot); + EXPECT_EQ(tm.lookup(run_id, TensorKey{0xFEED, -1}), writer.task_slot); // Writer has the creator recorded as a fanin producer (via INOUT // lookup) but no *live* fanin since the creator is already COMPLETED. EXPECT_EQ(S(writer.task_slot).fanin_count, 0); @@ -422,7 +424,7 @@ TEST_F(OrchestratorFixture, OutputAndOutputExistingAreInsertOnly) { auto writer_args = single_tensor_args(c.key, c.tag); auto writer = orch.submit_next_level(C(42), writer_args, cfg, 0); - EXPECT_EQ(tm.lookup(TensorKey{c.key, -1}), writer.task_slot); + EXPECT_EQ(tm.lookup(run_id, TensorKey{c.key, -1}), writer.task_slot); EXPECT_EQ(S(writer.task_slot).fanin_count, 0); EXPECT_TRUE(S(writer.task_slot).fanin_producers.empty()) << "tag=" << static_cast(c.tag); { @@ -431,3 +433,121 @@ TEST_F(OrchestratorFixture, OutputAndOutputExistingAreInsertOnly) { } } } + +TEST_F(OrchestratorFixture, EmptyRunCompletesWhenSubmissionCloses) { + EXPECT_FALSE(orch.run_done(run_id)); + orch.close_run_submission(run_id); + EXPECT_TRUE(orch.run_done(run_id)); + EXPECT_FALSE(orch.run_failed(run_id)); + EXPECT_NO_THROW(orch.wait_run(run_id)); + orch.release_run(run_id); +} + +TEST_F(OrchestratorFixture, OneTaskRunCompletesAfterConsumption) { + auto result = orch.submit_next_level(C(80), single_tensor_args(0x8000, TensorArgType::OUTPUT), cfg, 0); + EXPECT_EQ(S(result.task_slot).run_id, run_id); + + orch.close_run_submission(run_id); + EXPECT_FALSE(orch.run_done(run_id)); + + S(result.task_slot).state.store(TaskState::COMPLETED, std::memory_order_release); + EXPECT_TRUE(orch.on_consumed(result.task_slot)); + EXPECT_TRUE(orch.run_done(run_id)); + EXPECT_NO_THROW(orch.wait_run(run_id)); + orch.release_run(run_id); +} + +TEST_F(OrchestratorFixture, SequentialRunsHaveDistinctIdsAndErrorsDoNotLeak) { + auto failed_task = orch.submit_next_level(C(81), single_tensor_args(0x8100, TensorArgType::OUTPUT), cfg, 0); + orch.report_task_error(failed_task.task_slot, "run one failed"); + S(failed_task.task_slot).failure_message = "run one failed"; + S(failed_task.task_slot).state.store(TaskState::FAILED, std::memory_order_release); + EXPECT_TRUE(orch.on_consumed(failed_task.task_slot)); + orch.close_run_submission(run_id); + EXPECT_TRUE(orch.run_failed(run_id)); + EXPECT_THROW(orch.wait_run(run_id), std::runtime_error); + orch.release_run(run_id); + + RunId second = orch.begin_run(); + EXPECT_NE(second, run_id); + orch.close_run_submission(second); + EXPECT_FALSE(orch.run_failed(second)); + EXPECT_NO_THROW(orch.wait_run(second)); + orch.release_run(second); +} + +TEST_F(OrchestratorFixture, FailedSubmissionCompletesWithoutEnteringDeviceWork) { + orch.fail_run_submission(run_id, std::make_exception_ptr(std::runtime_error("graph build failed"))); + EXPECT_TRUE(orch.run_done(run_id)); + EXPECT_TRUE(orch.run_failed(run_id)); + EXPECT_THROW(orch.wait_run(run_id), std::runtime_error); + orch.release_run(run_id); +} + +TEST_F(OrchestratorFixture, ReleasingRunDoesNotResetSlotsOwnedByAnotherRegisteredRun) { + auto first_task = orch.submit_next_level(C(82), single_tensor_args(0x8200, TensorArgType::OUTPUT), cfg, 0); + S(first_task.task_slot).state.store(TaskState::COMPLETED, std::memory_order_release); + EXPECT_TRUE(orch.on_consumed(first_task.task_slot)); + orch.close_run_submission(run_id); + orch.wait_run(run_id); + + RunId second = orch.begin_run(); + EXPECT_EQ(allocator.next_task_id(), 1); + orch.release_run(run_id); + EXPECT_EQ(allocator.next_task_id(), 1); + + orch.close_run_submission(second); + orch.wait_run(second); + orch.release_run(second); + EXPECT_EQ(allocator.next_task_id(), 0); +} + +// report_task_error runs on a worker thread, where an escaping exception would +// terminate the process, so a slot whose run is already gone is a no-op. +TEST_F(OrchestratorFixture, ReportingErrorForAReleasedRunIsIgnored) { + orch.close_run_submission(run_id); + orch.wait_run(run_id); + RunId released = run_id; + orch.release_run(released); + + run_id = orch.begin_run(); + auto live = orch.submit_next_level(C(83), single_tensor_args(0x8300, TensorArgType::OUTPUT), cfg, 0); + S(live.task_slot).run_id = released; + + EXPECT_NO_THROW(orch.report_task_error(live.task_slot, "late endpoint failure")); + EXPECT_FALSE(orch.run_failed(run_id)); +} + +// on_consumed runs on the scheduler thread and reaches the same accounting. +TEST_F(OrchestratorFixture, ConsumingASlotOwnedByAReleasedRunIsIgnored) { + auto task = orch.submit_next_level(C(84), single_tensor_args(0x8400, TensorArgType::OUTPUT), cfg, 0); + S(task.task_slot).state.store(TaskState::COMPLETED, std::memory_order_release); + EXPECT_TRUE(orch.on_consumed(task.task_slot)); + orch.close_run_submission(run_id); + orch.wait_run(run_id); + RunId released = run_id; + orch.release_run(released); + + run_id = orch.begin_run(); + auto live = orch.submit_next_level(C(85), single_tensor_args(0x8500, TensorArgType::OUTPUT), cfg, 0); + orch.close_run_submission(run_id); + + S(live.task_slot).run_id = released; + S(live.task_slot).state.store(TaskState::COMPLETED, std::memory_order_release); + EXPECT_NO_THROW((void)orch.on_consumed(live.task_slot)); + + // The decrement landed on no run at all, so the live run still owes a task. + EXPECT_FALSE(orch.run_done(run_id)); +} + +TEST_F(OrchestratorFixture, FailedSubmissionCarriesItsMessageToTheFence) { + orch.fail_run_submission(run_id, std::make_exception_ptr(std::runtime_error("orchestration: ValueError: bad arg"))); + EXPECT_TRUE(orch.run_failed(run_id)); + try { + orch.wait_run(run_id); + FAIL() << "wait_run must rethrow the submission error"; + } catch (const std::runtime_error &e) { + EXPECT_STREQ(e.what(), "orchestration: ValueError: bad arg"); + } + orch.release_run(run_id); +} diff --git a/tests/ut/cpp/hierarchical/test_scheduler.cpp b/tests/ut/cpp/hierarchical/test_scheduler.cpp index 7383744692..fd0cf6ed70 100644 --- a/tests/ut/cpp/hierarchical/test_scheduler.cpp +++ b/tests/ut/cpp/hierarchical/test_scheduler.cpp @@ -273,6 +273,7 @@ struct SchedulerFixture : public ::testing::Test { WorkerManager manager; Scheduler sched; CallConfig cfg; + RunId run_id{INVALID_RUN_ID}; std::vector consumed_slots; std::mutex consumed_mu; @@ -291,6 +292,7 @@ struct SchedulerFixture : public ::testing::Test { orch.init(&tm, &allocator, &scope, &rq_sub, &rq_next_level, &manager, [this] { sched.notify_ready(); }); + run_id = orch.begin_run(); Scheduler::Config c; c.ring = &allocator; @@ -305,6 +307,9 @@ struct SchedulerFixture : public ::testing::Test { std::lock_guard lk(consumed_mu); consumed_slots.push_back(s); }; + c.on_task_failed_cb = [this](TaskSlot s, const std::string &message) { + orch.report_task_error(s, message); + }; sched.start(c); } @@ -461,7 +466,7 @@ TEST_F(SchedulerFixture, FailedProducerPoisonsDependentTask) { wait_consumed(a.task_slot); wait_consumed(b.task_slot); - EXPECT_TRUE(manager.has_error()); + EXPECT_TRUE(orch.run_failed(run_id)); EXPECT_EQ(mock_worker.dispatched_count(), 1) << "poisoned consumer must not dispatch"; EXPECT_EQ(S(a.task_slot).state.load(), TaskState::CONSUMED); EXPECT_EQ(S(b.task_slot).state.load(), TaskState::CONSUMED); @@ -483,6 +488,7 @@ struct GroupSchedulerFixture : public ::testing::Test { WorkerManager manager; Scheduler sched; CallConfig cfg; + RunId run_id{INVALID_RUN_ID}; std::vector consumed_slots; std::mutex consumed_mu; @@ -503,6 +509,7 @@ struct GroupSchedulerFixture : public ::testing::Test { orch.init(&tm, &allocator, &scope, &rq_sub, &rq_next_level, &manager, [this] { sched.notify_ready(); }); + run_id = orch.begin_run(); Scheduler::Config c; c.ring = &allocator; @@ -517,6 +524,9 @@ struct GroupSchedulerFixture : public ::testing::Test { std::lock_guard lk(consumed_mu); consumed_slots.push_back(s); }; + c.on_task_failed_cb = [this](TaskSlot s, const std::string &message) { + orch.report_task_error(s, message); + }; sched.start(c); } @@ -712,10 +722,10 @@ TEST_F(GroupSchedulerFixture, GroupFailureWaitsForRunningMembersThenConsumes) { worker_a.complete_with_error("member boom"); auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(300); - while (!manager.has_error() && std::chrono::steady_clock::now() < deadline) { + while (!orch.run_failed(run_id) && std::chrono::steady_clock::now() < deadline) { std::this_thread::sleep_for(std::chrono::milliseconds(1)); } - EXPECT_TRUE(manager.has_error()); + EXPECT_TRUE(orch.run_failed(run_id)); EXPECT_EQ(S(slot).state.load(), TaskState::RUNNING); worker_b.complete(); @@ -860,6 +870,7 @@ TEST(SchedulerWorkerTargetTest, NextLevelTargetUsesWorkerIdNotVectorIndex) { orch.init(&tm, &allocator, &scope, &rq_sub, &rq_next_level, &manager, [&sched] { sched.notify_ready(); }); + (void)orch.begin_run(); Scheduler::Config c; c.ring = &allocator; @@ -874,6 +885,9 @@ TEST(SchedulerWorkerTargetTest, NextLevelTargetUsesWorkerIdNotVectorIndex) { std::lock_guard lk(consumed_mu); consumed_slots.push_back(s); }; + c.on_task_failed_cb = [&orch](TaskSlot s, const std::string &message) { + orch.report_task_error(s, message); + }; sched.start(c); auto wait_consumed_slot = [&consumed_slots, &consumed_mu](TaskSlot slot) { @@ -963,6 +977,7 @@ struct MixedTypeSchedulerFixture : public ::testing::Test { WorkerManager manager; Scheduler sched; CallConfig cfg; + RunId run_id{INVALID_RUN_ID}; std::vector consumed_slots; std::mutex consumed_mu; @@ -983,6 +998,7 @@ struct MixedTypeSchedulerFixture : public ::testing::Test { orch.init(&tm, &allocator, &scope, &rq_sub, &rq_next_level, &manager, [this] { sched.notify_ready(); }); + run_id = orch.begin_run(); Scheduler::Config c; c.ring = &allocator; @@ -997,6 +1013,9 @@ struct MixedTypeSchedulerFixture : public ::testing::Test { std::lock_guard lk(consumed_mu); consumed_slots.push_back(s); }; + c.on_task_failed_cb = [this](TaskSlot s, const std::string &message) { + orch.report_task_error(s, message); + }; sched.start(c); } diff --git a/tests/ut/cpp/hierarchical/test_tensormap.cpp b/tests/ut/cpp/hierarchical/test_tensormap.cpp index 89c47b3cb9..1556ad5851 100644 --- a/tests/ut/cpp/hierarchical/test_tensormap.cpp +++ b/tests/ut/cpp/hierarchical/test_tensormap.cpp @@ -18,84 +18,100 @@ static TensorKey hk(uint64_t ptr) { return TensorKey::local_host(ptr); } // Helper: child key scoped by NEXT_LEVEL worker id. static TensorKey ck(uint64_t ptr, int32_t worker_id) { return TensorKey::local_child(ptr, worker_id); } +static constexpr RunId RUN1 = 1; +static constexpr RunId RUN2 = 2; TEST(TensorMap, LookupEmptyReturnsInvalid) { TensorMap tm; - EXPECT_EQ(tm.lookup(hk(0xDEADBEEF)), INVALID_SLOT); + EXPECT_EQ(tm.lookup(RUN1, hk(0xDEADBEEF)), INVALID_SLOT); } TEST(TensorMap, InsertAndLookup) { TensorMap tm; - tm.insert(hk(0x1000), 5); - EXPECT_EQ(tm.lookup(hk(0x1000)), 5); - EXPECT_EQ(tm.lookup(hk(0x2000)), INVALID_SLOT); + tm.insert(RUN1, hk(0x1000), 5); + EXPECT_EQ(tm.lookup(RUN1, hk(0x1000)), 5); + EXPECT_EQ(tm.lookup(RUN1, hk(0x2000)), INVALID_SLOT); EXPECT_EQ(tm.size(), 1); } TEST(TensorMap, OverwriteExistingEntry) { TensorMap tm; - tm.insert(hk(0x1000), 3); - tm.insert(hk(0x1000), 7); // new producer reuses same buffer - EXPECT_EQ(tm.lookup(hk(0x1000)), 7); + tm.insert(RUN1, hk(0x1000), 3); + tm.insert(RUN1, hk(0x1000), 7); // new producer reuses same buffer + EXPECT_EQ(tm.lookup(RUN1, hk(0x1000)), 7); EXPECT_EQ(tm.size(), 1); } TEST(TensorMap, EraseTaskOutputs) { TensorMap tm; - tm.insert(hk(0x1000), 0); - tm.insert(hk(0x2000), 0); - tm.insert(hk(0x3000), 1); + tm.insert(RUN1, hk(0x1000), 0); + tm.insert(RUN1, hk(0x2000), 0); + tm.insert(RUN1, hk(0x3000), 1); - tm.erase_task_outputs({hk(0x1000), hk(0x2000)}); + tm.erase_task_outputs(RUN1, {hk(0x1000), hk(0x2000)}); - EXPECT_EQ(tm.lookup(hk(0x1000)), INVALID_SLOT); - EXPECT_EQ(tm.lookup(hk(0x2000)), INVALID_SLOT); - EXPECT_EQ(tm.lookup(hk(0x3000)), 1); + EXPECT_EQ(tm.lookup(RUN1, hk(0x1000)), INVALID_SLOT); + EXPECT_EQ(tm.lookup(RUN1, hk(0x2000)), INVALID_SLOT); + EXPECT_EQ(tm.lookup(RUN1, hk(0x3000)), 1); EXPECT_EQ(tm.size(), 1); } TEST(TensorMap, EraseWithEmptyKeyList) { TensorMap tm; - tm.insert(hk(0x1000), 2); - tm.erase_task_outputs({}); - EXPECT_EQ(tm.lookup(hk(0x1000)), 2); + tm.insert(RUN1, hk(0x1000), 2); + tm.erase_task_outputs(RUN1, {}); + EXPECT_EQ(tm.lookup(RUN1, hk(0x1000)), 2); } TEST(TensorMap, MultipleEntries) { TensorMap tm; for (int i = 0; i < 100; ++i) - tm.insert(hk(static_cast(i) * 0x1000), i % 16); + tm.insert(RUN1, hk(static_cast(i) * 0x1000), i % 16); EXPECT_EQ(tm.size(), 100); for (int i = 0; i < 100; ++i) - EXPECT_EQ(tm.lookup(hk(static_cast(i) * 0x1000)), i % 16); + EXPECT_EQ(tm.lookup(RUN1, hk(static_cast(i) * 0x1000)), i % 16); } // --- TensorKey compound key tests --- TEST(TensorMap, SamePtrDifferentEndpointAreDistinct) { TensorMap tm; - tm.insert(ck(0xABC, 0), 10); - tm.insert(ck(0xABC, 1), 20); - EXPECT_EQ(tm.lookup(ck(0xABC, 0)), 10); - EXPECT_EQ(tm.lookup(ck(0xABC, 1)), 20); + tm.insert(RUN1, ck(0xABC, 0), 10); + tm.insert(RUN1, ck(0xABC, 1), 20); + EXPECT_EQ(tm.lookup(RUN1, ck(0xABC, 0)), 10); + EXPECT_EQ(tm.lookup(RUN1, ck(0xABC, 1)), 20); EXPECT_EQ(tm.size(), 2); } TEST(TensorMap, HostAndChildKeyAreDistinct) { TensorMap tm; - tm.insert(hk(0x1000), 5); - tm.insert(ck(0x1000, 0), 6); - EXPECT_EQ(tm.lookup(hk(0x1000)), 5); - EXPECT_EQ(tm.lookup(ck(0x1000, 0)), 6); + tm.insert(RUN1, hk(0x1000), 5); + tm.insert(RUN1, ck(0x1000, 0), 6); + EXPECT_EQ(tm.lookup(RUN1, hk(0x1000)), 5); + EXPECT_EQ(tm.lookup(RUN1, ck(0x1000, 0)), 6); EXPECT_EQ(tm.size(), 2); } TEST(TensorMap, EraseChildKeyLeavesHostKey) { TensorMap tm; - tm.insert(hk(0x1000), 5); - tm.insert(ck(0x1000, 0), 6); - tm.erase_task_outputs({ck(0x1000, 0)}); - EXPECT_EQ(tm.lookup(hk(0x1000)), 5); - EXPECT_EQ(tm.lookup(ck(0x1000, 0)), INVALID_SLOT); + tm.insert(RUN1, hk(0x1000), 5); + tm.insert(RUN1, ck(0x1000, 0), 6); + tm.erase_task_outputs(RUN1, {ck(0x1000, 0)}); + EXPECT_EQ(tm.lookup(RUN1, hk(0x1000)), 5); + EXPECT_EQ(tm.lookup(RUN1, ck(0x1000, 0)), INVALID_SLOT); EXPECT_EQ(tm.size(), 1); } + +TEST(TensorMap, SameAddressIsNamespacedByRun) { + TensorMap tm; + tm.insert(RUN1, hk(0x1000), 5); + tm.insert(RUN2, hk(0x1000), 9); + + EXPECT_EQ(tm.lookup(RUN1, hk(0x1000)), 5); + EXPECT_EQ(tm.lookup(RUN2, hk(0x1000)), 9); + EXPECT_EQ(tm.size(), 2); + + tm.erase_task_outputs(RUN1, {hk(0x1000)}); + EXPECT_EQ(tm.lookup(RUN1, hk(0x1000)), INVALID_SLOT); + EXPECT_EQ(tm.lookup(RUN2, hk(0x1000)), 9); +} diff --git a/tests/ut/py/test_worker/test_error_propagation.py b/tests/ut/py/test_worker/test_error_propagation.py index 85bbed400d..7f523465e6 100644 --- a/tests/ut/py/test_worker/test_error_propagation.py +++ b/tests/ut/py/test_worker/test_error_propagation.py @@ -13,8 +13,8 @@ 1. SubWorker callable raises → Worker.run(orch) re-raises with original Python exception type + message in the text. -2. Scope mid-failure → Nth submit sees has_error, next Worker.run - is not wedged (error state cleared). +2. Scope mid-failure → Nth submit sees its run error; the next + Worker.run uses an independent error state. 3. L4 → L3 → SubWorker chain → Exception raised in bottom sub surfaces at L4 Worker.run() with an identifiable chain (sub_worker / child_worker prefixes). @@ -100,19 +100,33 @@ def orch(o, args, cfg): # --------------------------------------------------------------------------- -# 2. Scope mid-failure — drain rethrows once, next run clean +# 2. Scope mid-failure — per-run wait rethrows once, next run clean # --------------------------------------------------------------------------- class TestScopeMidFailure: + def test_orchestration_exception_is_synchronous(self): + hw = Worker(level=3, num_sub_workers=1) + hw.init() + try: + + def graph_boom(orch, args, cfg): + raise SentinelError("graph-build-failure") + + with pytest.raises(SentinelError, match="graph-build-failure"): + hw.run(graph_boom) + + assert hw.run(lambda orch, args, cfg: None) is None + finally: + hw.close() + def test_failure_does_not_wedge_worker(self): """After a run() fails, the next run() using a clean orch succeeds. Register two callables — one that always raises, one that increments a shared counter. Fire the failing one first (should raise) then the - succeeding one (should run to completion). Proves ``_clear_error`` at - the top of ``Worker.run`` resets the error slot so the next run is - not permanently poisoned. + succeeding one (should run to completion). Proves the next run owns a + separate error state and is not permanently poisoned. Two callables rather than a shared-state toggle: the sub-worker is a forked child that inherits closure state copy-on-write; updates @@ -155,8 +169,8 @@ def test_subsequent_submit_raises_after_failure(self): """A second submit_sub after a failed first one must not swallow the error. With one sub worker we submit two tasks sequentially in the orch fn; - the first one fails, so the second submit sees has_error and re-raises - immediately (fail-fast). The exception reaching Worker.run is the + the first one fails, so the second submit sees the current run's error + and re-raises immediately (fail-fast). The exception reaching Worker.run is the original child failure — not some generic "orchestrator poisoned" wrapper. """ @@ -197,7 +211,7 @@ def test_bottom_sub_failure_surfaces_at_l4(self): The error first bubbles from the SubWorker through the parent's dispatch_process as `std::runtime_error("sub_worker: ... SentinelError: ...")`. - Inside the L3 orch fn this reaches _drain which rethrows. The L3 + Inside the L3 orch fn this reaches the per-run wait, which rethrows. The L3 child process catches that in _child_worker_loop and rewrites it as `child_worker level=3: RuntimeError: `. The L4 parent's dispatch_process rethrows again. The final string visible at the L4 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 e681453af4..10531fb690 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 @@ -110,8 +110,8 @@ def control_l3_l2_region_release(self, worker_id: int, region_id: int) -> None: class _EndpointFailingOrch: - def _clear_error(self) -> None: - pass + def _begin_run(self) -> int: + return 1 def _scope_begin(self) -> None: pass @@ -119,12 +119,22 @@ def _scope_begin(self) -> None: def _scope_end(self) -> None: pass - def _drain(self) -> None: + def _close_run_submission(self, run_id: int) -> None: + assert run_id == 1 + + def _fail_run_submission(self, run_id: int) -> None: + assert run_id == 1 + + def _wait_run(self, run_id: int) -> None: + assert run_id == 1 raise RuntimeError( "child failed: L3-L2 endpoint error op=signal_wait kind=3 region=2 " "counter_addr=0x200000 counter_operand=7 observed_counter=0 msg=wait timed out" ) + def _release_run(self, run_id: int) -> None: + assert run_id == 1 + class _NamedOnboardRegionExport: def __init__(self) -> None: