From 6b3904277b9b58bf1a0853483b76292379ab0fab Mon Sep 17 00:00:00 2001 From: Crane-Liu Date: Mon, 27 Jul 2026 22:55:58 +0800 Subject: [PATCH 1/3] Feat: add bounded whole-run FIFO admission --- docs/task-flow.md | 44 ++++- python/bindings/worker_bind.h | 7 +- python/simpler/worker.py | 42 ++-- src/common/hierarchical/orchestrator.cpp | 186 +++++++++++++++--- src/common/hierarchical/orchestrator.h | 18 +- src/common/hierarchical/scheduler.cpp | 51 ++++- src/common/hierarchical/scheduler.h | 3 + src/common/hierarchical/types.cpp | 90 ++++++++- src/common/hierarchical/types.h | 47 +++-- src/common/hierarchical/worker.cpp | 3 + src/common/hierarchical/worker.h | 5 + src/common/hierarchical/worker_manager.cpp | 1 + src/common/hierarchical/worker_manager.h | 11 +- src/common/worker/pipeline_slot_pool.h | 9 +- .../test_worker_async_fifo.py | 186 ++++++++++++++++++ .../ut/cpp/hierarchical/test_orchestrator.cpp | 141 ++++++++++++- .../hierarchical/test_pipeline_contract.cpp | 10 + tests/ut/cpp/hierarchical/test_scheduler.cpp | 35 ++++ tests/ut/py/test_callable_identity.py | 5 + tests/ut/py/test_worker/test_host_worker.py | 76 +++++++ 20 files changed, 890 insertions(+), 80 deletions(-) create mode 100644 tests/st/a2a3/host_build_graph/worker_async_fifo/test_worker_async_fifo.py diff --git a/docs/task-flow.md b/docs/task-flow.md index f9b109ea9..a370afc89 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,6 +293,39 @@ 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. +#### Whole-run FIFO admission + +L3 graph callbacks remain synchronous and serialized, but native admission +allows two live run reservations: + +```text +run N: EXECUTING +run N+1: PREPARED +run N+2: blocked in begin_run before its graph callback +``` + +`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 a blocked third submission, 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. + 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 @@ -382,6 +416,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 +427,8 @@ fanout wiring), see [orchestrator.md](orchestrator.md). For local endpoints, after the Scheduler resolves the submitted NEXT_LEVEL target (or chooses an idle SUB worker), `LocalMailboxEndpoint` encodes -`(callable digest, CallConfig, TaskArgs)` into the per-worker shm mailbox and +`(callable digest, CallConfig, PipelineSlotLease, TaskArgs)` into the +per-worker shm mailbox and the forked child decodes it. Remote NEXT_LEVEL dispatch through `RemoteL3Endpoint` serializes the same logical payload into a framed TASK request instead. diff --git a/python/bindings/worker_bind.h b/python/bindings/worker_bind.h index f110c9a16..10e206b5c 100644 --- a/python/bindings/worker_bind.h +++ b/python/bindings/worker_bind.h @@ -350,7 +350,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", @@ -410,6 +410,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) { @@ -766,6 +770,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/worker.py b/python/simpler/worker.py index 5535007aa..dd20b491f 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -82,6 +82,7 @@ def my_l4_orch(orch, args, config): import cloudpickle from _task_interface import ( # pyright: ignore[reportMissingImports] MAX_REGISTERED_CALLABLE_IDS, + PTO_PIPELINE_MAX_DEPTH, RUNTIME_ENV_RING_COUNT, TENSOR_CHILD_MEMORY_OFFSET, WorkerType, @@ -176,10 +177,13 @@ def my_l4_orch(orch, args, config): # travels separately via ChipWorker.init(log_level) — not on per-task wire. _RUNTIME_ENV_UINT64_FIELD_COUNT = 3 * RUNTIME_ENV_RING_COUNT _CFG_FMT = struct.Struct("=iiiiii" + ("Q" * _RUNTIME_ENV_UINT64_FIELD_COUNT) + "1024s") -# Args region starts after CONFIG, rounded up to 8 bytes so the first +# The generation-safe pipeline lease follows CONFIG. Args start after the +# lease, rounded up to 8 bytes so the first # Tensor.data (uint64_t at OFF_ARGS+8) is 8-byte aligned, avoiding # SIGBUS on strict-alignment platforms (aarch64 atomics, some ARM cores). -_OFF_ARGS = (_OFF_CONFIG + _CFG_FMT.size + 7) & ~7 +_PIPELINE_LEASE_FMT = struct.Struct("=IIQ") +_OFF_PIPELINE_LEASE = (_OFF_CONFIG + _CFG_FMT.size + 7) & ~7 +_OFF_ARGS = (_OFF_PIPELINE_LEASE + _PIPELINE_LEASE_FMT.size + 7) & ~7 assert _OFF_ARGS % 8 == 0, "_OFF_ARGS must be 8-aligned for Tensor.data" _OFF_TASK_CALLABLE_HASH = _OFF_ARGS _OFF_TASK_ARGS_BLOB = _OFF_TASK_CALLABLE_HASH + CALLABLE_HASH_DIGEST_BYTES @@ -1558,9 +1562,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( @@ -1570,6 +1579,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 @@ -1767,6 +1778,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() @@ -4489,6 +4501,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 @@ -4594,6 +4607,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 @@ -4657,6 +4678,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 @@ -6147,9 +6169,9 @@ 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. Native admission permits one + active and one prepared run; a third submission blocks before invoking + its graph callback. Completion and cleanup stay attached to each handle. """ with self._operation_lease("submit"): return self._submit_locked(callable, args, config) @@ -6176,12 +6198,8 @@ 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. The native begin_run admission + # boundary owns depth-two backpressure and slot generations. return self._submit_l3_locked(callable, args, cfg) def _submit_l3_locked(self, callable, args, cfg: CallConfig) -> RunHandle: diff --git a/src/common/hierarchical/orchestrator.cpp b/src/common/hierarchical/orchestrator.cpp index 649a25bc7..04b58554e 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,106 @@ std::shared_ptr Orchestrator::current_building_run() const { } RunId Orchestrator::begin_run() { - std::lock_guard lk(runs_mu_); - if (building_run_id_ != INVALID_RUN_ID) { - throw std::logic_error("Orchestrator::begin_run: another run is still building"); - } - if (next_run_id_ == INVALID_RUN_ID || next_run_id_ == std::numeric_limits::max()) { - throw std::overflow_error("Orchestrator::begin_run: run id space exhausted"); + RunId run_id = INVALID_RUN_ID; + { + std::unique_lock lk(runs_mu_); + if (building_run_id_ != INVALID_RUN_ID) { + throw std::logic_error("Orchestrator::begin_run: another run is still building"); + } + std::optional lease; + runs_cv_.wait(lk, [this, &lease] { + lease = pipeline_slots_.try_acquire(admission_depth_); + return lease.has_value(); + }); + if (next_run_id_ == INVALID_RUN_ID || next_run_id_ == std::numeric_limits::max()) { + pipeline_slots_.release(*lease); + throw std::overflow_error("Orchestrator::begin_run: run id space exhausted"); + } + run_id = next_run_id_++; + auto run = std::make_shared(run_id, *lease); + runs_.emplace(run_id, run); + run_fifo_.push_back(run_id); + building_run_id_ = run_id; + run->phase.store(RunPhase::BUILDING, std::memory_order_release); } - RunId run_id = next_run_id_++; - runs_.emplace(run_id, std::make_shared(run_id)); - building_run_id_ = run_id; + + // The FIFO head may execute while its graph callback is still building. + // Existing orchestration callbacks use this path to submit device work and + // wait for L2 communication before returning. A successor remains gated by + // active_run_id_ until the prior run reaches its terminal fence. + activate_fifo_head(); return run_id; } +void Orchestrator::configure_pipeline_depth(uint32_t depth) { + if (depth == 0 || depth > PTO_PIPELINE_MAX_DEPTH) { + throw std::invalid_argument("Orchestrator: pipeline depth is outside the supported range"); + } + std::lock_guard lk(runs_mu_); + if (!runs_.empty() || building_run_id_ != INVALID_RUN_ID || active_run_id_ != INVALID_RUN_ID) { + throw std::logic_error("Orchestrator: pipeline depth cannot change after admission starts"); + } + admission_depth_ = depth; +} + void Orchestrator::finish_run_if_ready(const std::shared_ptr &run) { bool notify = false; { std::lock_guard lk(run->completion_mu); - if (!run->submission_closed || run->active_tasks.load(std::memory_order_acquire) != 0 || - is_terminal(run->phase.load(std::memory_order_acquire))) { + RunPhase phase = run->phase.load(std::memory_order_acquire); + if (!run->submission_closed || run->active_tasks.load(std::memory_order_acquire) != 0 || is_terminal(phase) || + (phase != RunPhase::EXECUTING && !run->submission_failed)) { return; } bool failed = run->submission_failed || static_cast(run->first_error); run->phase.store(failed ? RunPhase::FAILED : RunPhase::COMPLETED, std::memory_order_release); notify = true; } - if (notify) run->completion_cv.notify_all(); + if (!notify) return; + run->completion_cv.notify_all(); + retire_terminal_run(run); +} + +void Orchestrator::clear_run_ready_queues(RunId run_id) { + if (ready_sub_queue_ != nullptr) ready_sub_queue_->erase_run(run_id); + if (ready_next_level_queues_ != nullptr) ready_next_level_queues_->erase_run(run_id); +} + +void Orchestrator::retire_terminal_run(const std::shared_ptr &run) { + PipelineSlotLease lease{}; + bool release_lease = false; + { + 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); + if (!run->lease_released) { + run->lease_released = true; + lease = run->lease; + release_lease = true; + } + } + clear_run_ready_queues(run->id); + if (release_lease) pipeline_slots_.release(lease); + 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); + } + if (ready_notify_cb_) ready_notify_cb_(); + finish_run_if_ready(run); } void Orchestrator::close_run_submission(RunId run_id) { @@ -104,16 +179,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 +210,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 +275,19 @@ bool Orchestrator::run_failed(RunId run_id) const { return run->submission_failed || static_cast(run->first_error); } +bool Orchestrator::can_dispatch_run(RunId run_id) const { + std::lock_guard lk(runs_mu_); + 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); +} + +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 +316,44 @@ 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()); + } + + for (TaskSlot slot : slots) { + TaskSlotState &state = slot_state(slot); + TaskState current = state.state.load(std::memory_order_acquire); + if (current == TaskState::PENDING || current == TaskState::READY) { + state.failure_message = message; + (void)state.state.compare_exchange_strong( + current, TaskState::FAILED, std::memory_order_release, std::memory_order_acquire + ); + } + } + for (TaskSlot slot : slots) { + TaskSlotState &state = slot_state(slot); + TaskState current = state.state.load(std::memory_order_acquire); + if (current == TaskState::FAILED) try_consume(slot); + } + for (TaskSlot slot : slots) { + TaskSlotState &state = slot_state(slot); + std::vector producers; + { + std::lock_guard lk(state.fanout_mu); + producers = state.fanin_producers; + } + for (TaskSlot producer : producers) + try_consume(producer); + } } void Orchestrator::decrement_run_tasks(RunId run_id) { @@ -366,6 +500,8 @@ Tensor Orchestrator::alloc(const std::vector &shape, DataType dtype) { TaskSlotState &s = slot_state(ar.slot); s.reset(); s.run_id = run->id; + s.pipeline_lease = run->lease; + register_run_slot(run, ar.slot); uint64_t ptr = reinterpret_cast(ar.heap_ptr); if (ptr != 0) { @@ -394,8 +530,6 @@ Tensor Orchestrator::alloc(const std::vector &shape, DataType dtype) { s.state.store(TaskState::COMPLETED, std::memory_order_release); - increment_run_tasks(run->id); - // Build a contiguous external Tensor over the allocated buffer. ptr may be // 0 for a 0-byte request (a shape with a zero dim), in which case // init_external sets buffer.addr == 0 — the "no tensor" sentinel honored by @@ -479,6 +613,9 @@ SubmitResult Orchestrator::submit_impl( TaskSlotState &s = slot_state(slot); s.reset(); s.run_id = run->id; + s.pipeline_lease = run->lease; + s.state.store(TaskState::PENDING, std::memory_order_release); + register_run_slot(run, slot); s.worker_type = worker_type; s.callable = callable; @@ -546,7 +683,6 @@ 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); increment_run_accepts(run->id, s.group_size()); if (poisoned_by_failed_producer) { @@ -580,13 +716,13 @@ void Orchestrator::enqueue_ready(TaskSlot slot) { if (ready_next_level_queues_ == nullptr) throw std::runtime_error("Orchestrator::enqueue_ready: NEXT_LEVEL queues are not initialized"); if (s.is_group()) { - ready_next_level_queues_->push_group(slot); + ready_next_level_queues_->push_group(s.run_id, slot); } else { - ready_next_level_queues_->push_single(s.target_worker_id(0), slot); + ready_next_level_queues_->push_single(s.target_worker_id(0), s.run_id, slot); } return; } - ready_sub_queue_->push(slot); + ready_sub_queue_->push(s.run_id, slot); } void Orchestrator::validate_worker_eligibility( diff --git a/src/common/hierarchical/orchestrator.h b/src/common/hierarchical/orchestrator.h index d131bc0dd..591e27e6e 100644 --- a/src/common/hierarchical/orchestrator.h +++ b/src/common/hierarchical/orchestrator.h @@ -33,6 +33,7 @@ #include #include #include +#include #include #include #include @@ -43,6 +44,7 @@ #include "../task_interface/data_type.h" #include "../task_interface/task_args.h" #include "../task_interface/tensor.h" +#include "../worker/pipeline_slot_pool.h" #include "ring.h" #include "scope.h" #include "tensormap.h" @@ -118,6 +120,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); @@ -126,6 +129,8 @@ 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; + RunId active_run_id() const; void release_run(RunId run_id); // Open a nested scope. Every task submitted between this call and the @@ -176,9 +181,14 @@ class Orchestrator { NextLevelReadyQueues *ready_next_level_queues_ = nullptr; mutable std::mutex runs_mu_; + std::condition_variable runs_cv_; std::unordered_map> runs_; + std::deque run_fifo_; + PipelineSlotPool pipeline_slots_{PTO_PIPELINE_MAX_DEPTH}; + uint32_t admission_depth_{PTO_PIPELINE_MAX_DEPTH}; RunId next_run_id_{1}; RunId building_run_id_{INVALID_RUN_ID}; + RunId active_run_id_{INVALID_RUN_ID}; // Scheduler's loop mutex (not owned). Held across optional quiescent // compaction so the scheduler cannot retain a slot pointer being removed. @@ -190,13 +200,17 @@ 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; + 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 acd6e2242..e819454f1 100644 --- a/src/common/hierarchical/scheduler.cpp +++ b/src/common/hierarchical/scheduler.cpp @@ -176,8 +176,15 @@ void Scheduler::run() { { std::unique_lock lk(completion_mu_); completion_cv_.wait(lk, [this] { - return !completion_queue_.empty() || !cfg_.ready_next_level_queues->empty() || - !cfg_.ready_sub_queue->empty() || stop_requested_.load(std::memory_order_acquire); + bool ready = false; + if (cfg_.active_run_cb) { + RunId active = cfg_.active_run_cb(); + ready = active != INVALID_RUN_ID && + (!cfg_.ready_next_level_queues->empty(active) || !cfg_.ready_sub_queue->empty(active)); + } else { + ready = !cfg_.ready_next_level_queues->empty() || !cfg_.ready_sub_queue->empty(); + } + return !completion_queue_.empty() || ready || stop_requested_.load(std::memory_order_acquire); }); } @@ -356,10 +363,16 @@ void Scheduler::dispatch_ready() { } void Scheduler::dispatch_sub_ready() { + RunId active_run = cfg_.active_run_cb ? cfg_.active_run_cb() : INVALID_RUN_ID; + if (cfg_.active_run_cb && active_run == INVALID_RUN_ID) return; TaskSlot slot; - while (cfg_.ready_sub_queue->try_pop(slot)) { + while (cfg_.active_run_cb ? cfg_.ready_sub_queue->try_pop(active_run, slot) : cfg_.ready_sub_queue->try_pop(slot)) { TaskSlotState &s = *cfg_.ring->slot_state(slot); if (s.state.load(std::memory_order_acquire) != TaskState::READY) continue; + if (cfg_.active_run_cb && s.run_id != active_run) { + cfg_.enqueue_ready_cb(slot); + return; + } if (s.worker_type != WorkerType::SUB) { throw std::runtime_error("Scheduler::dispatch_sub_ready: misrouted task slot"); } @@ -370,7 +383,7 @@ 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); + cfg_.ready_sub_queue->push(active_run, slot); return; } workers.push_back(worker); @@ -393,14 +406,27 @@ void Scheduler::dispatch_sub_ready() { } void Scheduler::dispatch_next_level_group() { + RunId active_run = cfg_.active_run_cb ? cfg_.active_run_cb() : INVALID_RUN_ID; + if (cfg_.active_run_cb && active_run == INVALID_RUN_ID) return; TaskSlot slot; - while (cfg_.ready_next_level_queues->try_front_group(slot)) { + while (cfg_.active_run_cb ? cfg_.ready_next_level_queues->try_front_group(active_run, slot) : + cfg_.ready_next_level_queues->try_front_group(slot)) { TaskSlotState &s = *cfg_.ring->slot_state(slot); if (s.state.load(std::memory_order_acquire) != TaskState::READY) { TaskSlot stale; - cfg_.ready_next_level_queues->try_pop_group(stale); + if (cfg_.active_run_cb) { + cfg_.ready_next_level_queues->try_pop_group(active_run, stale); + } else { + cfg_.ready_next_level_queues->try_pop_group(stale); + } continue; } + if (cfg_.active_run_cb && s.run_id != active_run) { + TaskSlot misplaced; + cfg_.ready_next_level_queues->try_pop_group(active_run, misplaced); + cfg_.enqueue_ready_cb(slot); + return; + } if (s.worker_type != WorkerType::NEXT_LEVEL || !s.is_group()) { throw std::runtime_error("Scheduler::dispatch_next_level_group: misrouted task slot"); } @@ -422,7 +448,9 @@ void Scheduler::dispatch_next_level_group() { } TaskSlot popped; - if (!cfg_.ready_next_level_queues->try_pop_group(popped) || popped != slot) { + bool popped_ok = cfg_.active_run_cb ? cfg_.ready_next_level_queues->try_pop_group(active_run, popped) : + cfg_.ready_next_level_queues->try_pop_group(popped); + if (!popped_ok || popped != slot) { throw std::runtime_error("Scheduler::dispatch_next_level_group: group queue changed unexpectedly"); } @@ -435,6 +463,8 @@ void Scheduler::dispatch_next_level_group() { } void Scheduler::dispatch_next_level_singles() { + RunId active_run = cfg_.active_run_cb ? cfg_.active_run_cb() : INVALID_RUN_ID; + if (cfg_.active_run_cb && active_run == INVALID_RUN_ID) return; for (int32_t worker_id : cfg_.ready_next_level_queues->worker_ids()) { WorkerThread *worker = cfg_.manager->get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); if (worker == nullptr) { @@ -445,9 +475,14 @@ void Scheduler::dispatch_next_level_singles() { if (!worker->idle()) continue; TaskSlot slot; - while (cfg_.ready_next_level_queues->try_pop_single(worker_id, slot)) { + while (cfg_.active_run_cb ? cfg_.ready_next_level_queues->try_pop_single(worker_id, active_run, slot) : + cfg_.ready_next_level_queues->try_pop_single(worker_id, slot)) { TaskSlotState &s = *cfg_.ring->slot_state(slot); if (s.state.load(std::memory_order_acquire) != TaskState::READY) continue; + if (cfg_.active_run_cb && s.run_id != active_run) { + cfg_.enqueue_ready_cb(slot); + break; + } if (s.worker_type != WorkerType::NEXT_LEVEL || s.is_group() || s.target_worker_id(0) != worker_id) { throw std::runtime_error("Scheduler::dispatch_next_level_singles: misrouted task slot"); } diff --git a/src/common/hierarchical/scheduler.h b/src/common/hierarchical/scheduler.h index 00c1abfd2..95be5dc77 100644 --- a/src/common/hierarchical/scheduler.h +++ b/src/common/hierarchical/scheduler.h @@ -59,6 +59,9 @@ 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 diff --git a/src/common/hierarchical/types.cpp b/src/common/hierarchical/types.cpp index 3fcebcd85..288f507d5 100644 --- a/src/common/hierarchical/types.cpp +++ b/src/common/hierarchical/types.cpp @@ -11,6 +11,7 @@ #include "types.h" +#include #include // ============================================================================= @@ -20,6 +21,7 @@ void TaskSlotState::reset() { state.store(TaskState::FREE, std::memory_order_relaxed); run_id = INVALID_RUN_ID; + pipeline_lease = PipelineSlotLease{}; fanin_count = 0; fanin_released.store(0, std::memory_order_relaxed); { @@ -59,31 +61,78 @@ void TaskSlotState::reset() { // 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 +163,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 +191,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..2c9083399 100644 --- a/src/common/hierarchical/types.h +++ b/src/common/hierarchical/types.h @@ -32,16 +32,19 @@ #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 +126,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}; }; // ============================================================================= @@ -322,6 +330,7 @@ struct WorkerCompletion { struct TaskSlotState { std::atomic state{TaskState::FREE}; RunId run_id{INVALID_RUN_ID}; + PipelineSlotLease pipeline_lease{}; // --- Fanin (orch writes once; scheduler reads atomically) --- int32_t fanin_count{0}; @@ -423,17 +432,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 +459,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..751eeddc4 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_.active_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 ad2466985..22dc9f01a 100644 --- a/src/common/hierarchical/worker_manager.cpp +++ b/src/common/hierarchical/worker_manager.cpp @@ -383,6 +383,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 6d00a9551..257b39b28 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..4f1890a8e --- /dev/null +++ b/tests/st/a2a3/host_build_graph/worker_async_fifo/test_worker_async_fifo.py @@ -0,0 +1,186 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Onboard validation for bounded whole-run FIFO admission. + +The first run completes real NPU work but remains active behind a SubTask +fence. The second run builds its graph into the other pipeline slot and must +not dispatch until the first run becomes terminal. A third submission must +block before its graph callback while both slots are admitted. +""" + +import multiprocessing +import threading +import time +from contextlib import suppress + +import pytest +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler_setup.scene_test import _build_l3_task_args + +_VECTOR_KERNELS = "../vector_example/kernels" +_SIZE = 128 * 128 +_SUB_ENTERED = multiprocessing.Event() +_SUB_RELEASE = multiprocessing.Event() + + +def _wait_for_release(_args): + _SUB_ENTERED.set() + if not _SUB_RELEASE.wait(30.0): + raise RuntimeError("whole-run FIFO test timed out waiting for the release fence") + + +@scene_test(level=3, runtime="host_build_graph") +class TestWorkerAsyncWholeRunFifo(SceneTestCase): + """A prepared run may build ahead but cannot dispatch ahead.""" + + CALLABLE = { + "callables": [ + { + "name": "vector", + "orchestration": { + "source": f"{_VECTOR_KERNELS}/orchestration/example_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": f"{_VECTOR_KERNELS}/aiv/kernel_add.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "source": f"{_VECTOR_KERNELS}/aiv/kernel_add_scalar.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT], + }, + { + "func_id": 2, + "source": f"{_VECTOR_KERNELS}/aiv/kernel_mul.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + ], + }, + {"name": "wait_for_release", "callable": _wait_for_release}, + ], + } + + CASES = [ + { + "name": "whole_run_fifo", + "platforms": ["a2a3"], + "config": {"device_count": 1, "num_sub_workers": 1, "aicpu_thread_num": 4}, + "params": {}, + }, + ] + + @staticmethod + def _tensor_from_host_buffer(worker, value): + buffer = worker.create_host_buffer(_SIZE * torch.float32.itemsize) + tensor = torch.frombuffer(buffer.buffer, dtype=torch.float32, count=_SIZE) + tensor.fill_(value) + return buffer, tensor + + def test_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" + + _SUB_RELEASE.set() + first.wait(10.0) + assert third_callback.wait(10.0), "the third submission did not enter after the first run freed its slot" + second.wait(10.0) + second_expected = (second_a + second_b + 1) * (second_a + second_b + 2) + assert torch.allclose(second_out, second_expected), "the prepared run did not execute correctly on the NPU" + + submitter.join(10.0) + assert not submitter.is_alive() + third_result["handle"].wait(10.0) + finally: + _SUB_RELEASE.set() + if submitter is not None: + submitter.join(10.0) + handles = [first, second, third_result.get("handle")] + for handle in handles: + if handle is not None: + with suppress(Exception): + handle.wait(10.0) + tensors.clear() + first_a = first_b = first_out = second_a = second_b = second_out = None + if all(handle is None or handle.done for handle in handles): + for buffer in buffers: + st_worker.free_host_buffer(buffer) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/ut/cpp/hierarchical/test_orchestrator.cpp b/tests/ut/cpp/hierarchical/test_orchestrator.cpp index 8dd4791c1..e4e8c0e44 100644 --- a/tests/ut/cpp/hierarchical/test_orchestrator.cpp +++ b/tests/ut/cpp/hierarchical/test_orchestrator.cpp @@ -12,6 +12,8 @@ #include #include +#include +#include #include "call_config.h" #include "ring.h" @@ -42,6 +44,7 @@ struct OrchestratorFixture : public ::testing::Test { struct WorkerQueueView { NextLevelReadyQueues *queues; bool try_pop(TaskSlot &out) { return queues->try_pop_single(0, out); } + bool try_pop(RunId run_id, TaskSlot &out) { return queues->try_pop_single(0, run_id, out); } } rq{&rq_next_level}; void SetUp() override { @@ -79,6 +82,24 @@ struct OrchestratorFixture : public ::testing::Test { // Tests // --------------------------------------------------------------------------- +TEST(ReadyQueueTest, UnscopedAccessPreservesRunInsertionOrder) { + ReadyQueue queue; + queue.push(/*run_id=*/7, /*slot=*/70); + queue.push(/*run_id=*/8, /*slot=*/80); + queue.push(/*run_id=*/7, /*slot=*/71); + + TaskSlot slot = INVALID_SLOT; + ASSERT_TRUE(queue.try_front(slot)); + EXPECT_EQ(slot, 70); + ASSERT_TRUE(queue.try_pop(slot)); + EXPECT_EQ(slot, 70); + ASSERT_TRUE(queue.try_pop(slot)); + EXPECT_EQ(slot, 71); + ASSERT_TRUE(queue.try_pop(slot)); + EXPECT_EQ(slot, 80); + EXPECT_TRUE(queue.empty()); +} + TEST_F(OrchestratorFixture, IndependentTaskIsImmediatelyReady) { auto a = single_tensor_args(0xCAFE, TensorArgType::OUTPUT); auto res = orch.submit_next_level(C(42), a, cfg, 0); @@ -494,11 +515,7 @@ TEST_F(OrchestratorFixture, RunAcceptanceWaitsForEveryDispatchedGroupMember) { TEST_F(OrchestratorFixture, TerminalFailureUnblocksRunAcceptance) { auto result = orch.submit_next_level(C(80), single_tensor_args(0x8030, TensorArgType::OUTPUT), cfg, 0); orch.fail_run_submission(run_id, std::make_exception_ptr(std::runtime_error("graph build failed"))); - EXPECT_FALSE(orch.run_accepted(run_id)); - - S(result.task_slot).state.store(TaskState::FAILED, std::memory_order_release); - EXPECT_TRUE(orch.on_consumed(result.task_slot)); - + EXPECT_EQ(S(result.task_slot).state.load(std::memory_order_acquire), TaskState::CONSUMED); EXPECT_TRUE(orch.run_accepted(run_id)); EXPECT_NO_THROW(orch.wait_run_accepted(run_id)); } @@ -675,3 +692,117 @@ TEST_F(OrchestratorFixture, FailedSubmissionCarriesItsMessageToTheFence) { } 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); +} + +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 31ca70b43..61a154658 100644 --- a/tests/ut/cpp/hierarchical/test_scheduler.cpp +++ b/tests/ut/cpp/hierarchical/test_scheduler.cpp @@ -318,6 +318,9 @@ struct SchedulerFixture : public ::testing::Test { c.enqueue_ready_cb = [this](TaskSlot slot) { orch.enqueue_ready(slot); }; + c.active_run_cb = [this] { + return orch.active_run_id(); + }; c.on_consumed_cb = [this](TaskSlot s) { orch.on_consumed(s); std::lock_guard lk(consumed_mu); @@ -387,6 +390,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 +404,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) {} @@ -1180,6 +1190,9 @@ struct MixedTypeSchedulerFixture : public ::testing::Test { c.enqueue_ready_cb = [this](TaskSlot slot) { orch.enqueue_ready(slot); }; + c.active_run_cb = [this] { + return orch.active_run_id(); + }; c.on_consumed_cb = [this](TaskSlot s) { orch.on_consumed(s); std::lock_guard lk(consumed_mu); @@ -1243,6 +1256,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..ac01dce0c 100644 --- a/tests/ut/py/test_callable_identity.py +++ b/tests/ut/py/test_callable_identity.py @@ -540,6 +540,7 @@ class FakeCWorker: def __init__(self, *args): self.remote_worker_ids = [] self.closed = False + self.pipeline_depth = None def add_remote_l3_socket(self, worker_id, *args): self.remote_worker_ids.append(worker_id) @@ -555,6 +556,9 @@ def add_next_level_worker(self, *args): def add_next_level_worker_at(self, *args): pass + def configure_pipeline_depth(self, depth): + self.pipeline_depth = depth + def init(self): pass @@ -597,6 +601,7 @@ def fake_open_remote_session(self, *, spec, worker_id, session_id, deadline): assert worker._resolve_handle(handle).eligible_worker_ids == (remote_worker_id,) worker.init() + assert fake_c_worker.pipeline_depth == 2 assert opened_worker_ids == [remote_worker_id] assert fake_c_worker.remote_worker_ids == [remote_worker_id] diff --git a/tests/ut/py/test_worker/test_host_worker.py b/tests/ut/py/test_worker/test_host_worker.py index ac93af168..a3cf24399 100644 --- a/tests/ut/py/test_worker/test_host_worker.py +++ b/tests/ut/py/test_worker/test_host_worker.py @@ -130,8 +130,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 +142,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 +169,7 @@ def fake_run_chip_main_loop(cw, *_args, chip_platform, chip_runtime, prepared=No assert events[1][0] == "main_loop" assert events[1][2:] == ("a2a3", "tensormap_and_ringbuffer") assert events[2] == ("finalize",) + assert published_depths == [2] def _chip_digest(callable_obj: ChipCallable, *, platform: str = "", runtime: str = "") -> bytes: @@ -1435,6 +1440,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() From d6dfaa90ea144fc7b23937d16476d13f58f0b3e1 Mon Sep 17 00:00:00 2001 From: Crane-Liu Date: Wed, 29 Jul 2026 15:27:38 +0800 Subject: [PATCH 2/3] Add: generation-safe two-frame local endpoint - Separate control traffic from two identified task frames - Accept a queued frame while one sequential child executor is active - Bound parent dispatch capacity and fill every available frame - Retire publish ordering across every pre-publish failure path - Reject oversized argument blobs before publishing mailbox state - Preserve single-frame fallbacks on unsupported local backends - Keep failure-path tests and host-buffer cleanup terminal-safe - Release the GIL around native execution so admission remains live --- docs/task-flow.md | 10 +- python/bindings/task_interface.cpp | 1 + python/bindings/worker_bind.h | 17 +- python/simpler/task_interface.py | 2 + python/simpler/worker.py | 166 ++++++++- src/common/hierarchical/scheduler.cpp | 6 +- src/common/hierarchical/types.h | 5 +- src/common/hierarchical/worker.cpp | 8 +- src/common/hierarchical/worker.h | 4 +- src/common/hierarchical/worker_manager.cpp | 322 ++++++++++++++++-- src/common/hierarchical/worker_manager.h | 63 +++- .../orchestration/long_vector_orch.cpp | 66 ++++ .../test_worker_async_endpoint.py | 137 ++++++++ .../ut/cpp/hierarchical/test_orchestrator.cpp | 34 +- tests/ut/cpp/hierarchical/test_scheduler.cpp | 125 ++++++- tests/ut/py/test_worker/test_host_worker.py | 115 ++++++- 16 files changed, 997 insertions(+), 84 deletions(-) create mode 100644 tests/st/a2a3/host_build_graph/worker_async_endpoint/kernels/orchestration/long_vector_orch.cpp create mode 100644 tests/st/a2a3/host_build_graph/worker_async_endpoint/test_worker_async_endpoint.py diff --git a/docs/task-flow.md b/docs/task-flow.md index a370afc89..66cfd8024 100644 --- a/docs/task-flow.md +++ b/docs/task-flow.md @@ -326,13 +326,6 @@ 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. -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. - 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 buffer per slot, and its single-entry prebuilt-arena cache stays owned by @@ -447,7 +440,8 @@ Local mailbox path: ```text slot.callable.digest ─┐ slot.config ─┼─► memcpy into shm mailbox ─► child resolves digest -slot.task_args ─┘ (dispatch_process) and runs local slot +slot.pipeline_lease ─┤ (dispatch_process) and runs local slot +slot.task_args ─┘ ``` For SUB children the same mailbox layout is reused; the Python child diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index 4d588551a..b7c07cc2f 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -1456,6 +1456,7 @@ NB_MODULE(_task_interface, m) { [](ChipWorker &self, int32_t callable_id, uint64_t args_blob_ptr, size_t blob_capacity, const CallConfig &config, uint64_t accepted_state_addr, int32_t accepted_value, uint32_t pipeline_slot, uint64_t pipeline_generation) { + nb::gil_scoped_release release; // The mailbox region is the on-wire format `write_blob` produced; // `read_blob` is the matching reader that returns a zero-copy // TaskArgsView into the caller-owned bytes. Forwards to the diff --git a/python/bindings/worker_bind.h b/python/bindings/worker_bind.h index 10e206b5c..0610a0ddf 100644 --- a/python/bindings/worker_bind.h +++ b/python/bindings/worker_bind.h @@ -393,10 +393,12 @@ inline void bind_worker(nb::module_ &m) { .def( "add_next_level_worker", - [](Worker &self, uint64_t mailbox_ptr, int child_pid) { - self.add_worker(WorkerType::NEXT_LEVEL, reinterpret_cast(mailbox_ptr), child_pid); + [](Worker &self, uint64_t mailbox_ptr, int child_pid, uint32_t task_frame_count) { + self.add_worker( + WorkerType::NEXT_LEVEL, reinterpret_cast(mailbox_ptr), child_pid, task_frame_count + ); }, - nb::arg("mailbox_ptr"), nb::arg("child_pid") = -1, + nb::arg("mailbox_ptr"), nb::arg("child_pid") = -1, nb::arg("task_frame_count") = 1, "Add a NEXT_LEVEL sub-worker. `mailbox_ptr` is the address of a " "MAILBOX_SIZE-byte MAP_SHARED region; the child process loop is " "Python-managed (fork + _chip_process_loop). `child_pid` is that " @@ -404,10 +406,12 @@ inline void bind_worker(nb::module_ &m) { ) .def( "add_next_level_worker_at", - [](Worker &self, int32_t worker_id, uint64_t mailbox_ptr, int child_pid) { - self.add_next_level_worker(worker_id, reinterpret_cast(mailbox_ptr), child_pid); + [](Worker &self, int32_t worker_id, uint64_t mailbox_ptr, int child_pid, uint32_t task_frame_count) { + self.add_next_level_worker( + worker_id, reinterpret_cast(mailbox_ptr), child_pid, task_frame_count + ); }, - nb::arg("worker_id"), nb::arg("mailbox_ptr"), nb::arg("child_pid") = -1, + nb::arg("worker_id"), nb::arg("mailbox_ptr"), nb::arg("child_pid") = -1, nb::arg("task_frame_count") = 1, "Add a NEXT_LEVEL sub-worker with an explicit worker id." ) .def( @@ -768,6 +772,7 @@ inline void bind_worker(nb::module_ &m) { m.attr("DEFAULT_HEAP_RING_SIZE") = static_cast(DEFAULT_HEAP_RING_SIZE); m.attr("MAILBOX_SIZE") = static_cast(MAILBOX_SIZE); + m.attr("MAILBOX_FRAME_SIZE") = static_cast(MAILBOX_FRAME_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); diff --git a/python/simpler/task_interface.py b/python/simpler/task_interface.py index a7a3260d8..1e1b96e3c 100644 --- a/python/simpler/task_interface.py +++ b/python/simpler/task_interface.py @@ -41,6 +41,7 @@ import _task_interface as _ti_module # pyright: ignore[reportMissingImports] from _task_interface import ( # pyright: ignore[reportMissingImports] MAILBOX_ERROR_MSG_SIZE, + MAILBOX_FRAME_SIZE, MAILBOX_OFF_ERROR_MSG, MAILBOX_SIZE, MAX_REGISTERED_CALLABLE_IDS, @@ -156,6 +157,7 @@ def _assert_bindings_match_source_tree() -> None: "TaskState", "_Worker", "MAILBOX_SIZE", + "MAILBOX_FRAME_SIZE", "MAILBOX_OFF_ERROR_MSG", "MAILBOX_ERROR_MSG_SIZE", "read_args_from_blob", diff --git a/python/simpler/worker.py b/python/simpler/worker.py index dd20b491f..15483ae4a 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -129,6 +129,7 @@ def my_l4_orch(orch, args, config): from .orchestrator import Orchestrator from .task_interface import ( MAILBOX_ERROR_MSG_SIZE, + MAILBOX_FRAME_SIZE, MAILBOX_OFF_ERROR_MSG, MAILBOX_SIZE, CallConfig, @@ -194,9 +195,15 @@ def my_l4_orch(orch, args, config): # sticky word rather than a MailboxState, because a state carrying it is lost # whenever the child reaches TASK_DONE between two parent polls. The parent # clears it when it publishes the next TASK_READY. -_OFF_ACCEPTED = MAILBOX_SIZE - MAILBOX_ERROR_MSG_SIZE - 8 +_OFF_ACCEPTED = MAILBOX_FRAME_SIZE - MAILBOX_ERROR_MSG_SIZE - 8 _TASK_ACCEPTED = 1 -_MAILBOX_ARGS_CAPACITY = MAILBOX_SIZE - _OFF_TASK_ARGS_BLOB - MAILBOX_ERROR_MSG_SIZE - 8 +_OFF_FRAME_PROTOCOL = _OFF_ACCEPTED - 40 +_OFF_FRAME_RUN_ID = _OFF_ACCEPTED - 32 +_OFF_FRAME_SLOT_ID = _OFF_ACCEPTED - 24 +_OFF_FRAME_GENERATION = _OFF_ACCEPTED - 16 +_OFF_FRAME_DISPATCH_ID = _OFF_ACCEPTED - 8 +_TASK_PROTOCOL_VERSION = 1 +_MAILBOX_ARGS_CAPACITY = _OFF_FRAME_PROTOCOL - _OFF_TASK_ARGS_BLOB _OFF_CONTROL_CALLABLE_HASH = _OFF_ARGS + 32 # MAILBOX_OFF_ERROR_MSG / MAILBOX_ERROR_MSG_SIZE come from the C++ # nanobind module so the two sides cannot drift. @@ -216,6 +223,17 @@ def my_l4_orch(orch, args, config): # deadline aborts startup with a bounded error instead of an unbounded spin. _INIT_READY = 6 _INIT_FAILED = 7 +_TASK_ACCEPTED_STATE = 8 +_TASK_ACTIVE = 9 +_TASK_FAILED = 10 +_TASK_FRAME_COUNT = 2 + + +def _local_task_frame_count(platform: str, runtime: str, pipeline_depth: int) -> int: + if platform == "a2a3" and runtime == "host_build_graph" and pipeline_depth >= 2: + return _TASK_FRAME_COUNT + return 1 + # Startup readiness bound. A child that neither reports INIT_READY/INIT_FAILED # nor exits within this window is treated as hung and startup is aborted. @@ -1512,6 +1530,7 @@ def _run_chip_main_loop( # noqa: PLR0913, PLR0915 -- fork-child entry: every de chip_runtime: str = "", on_task_done_success=None, prepared: set[int] | None = None, + task_frame_count: int = 1, ) -> None: """Chip-process handlers for `_run_mailbox_loop`. @@ -1537,10 +1556,15 @@ def _run_chip_main_loop( # noqa: PLR0913, PLR0915 -- fork-child entry: every de host_buf_table: dict[int, tuple[SharedMemory, int, int, int]] = {} # token -> (shm, lo, hi, child_base) host_buf_ranges: list[tuple[int, int, int]] = [] # (parent_lo, parent_hi, child_base) - def handle_task() -> tuple[int, str]: - digest = _read_task_digest(buf) + def handle_task( + task_buf: memoryview = buf, + task_mailbox_addr: int = mailbox_addr, + *, + publish_native_acceptance: bool = True, + ) -> tuple[int, str]: + digest = _read_task_digest(task_buf) cid = identity_table.get(digest) - cfg = _read_config_from_mailbox(buf) + cfg = _read_config_from_mailbox(task_buf) code = 0 msg = "" @@ -1561,9 +1585,9 @@ def handle_task() -> tuple[int, str]: # blob to this child's own mapping before the runtime reads it. # No-op when nothing is registered. if host_buf_ranges: - _rewrite_blob_host_addrs(buf, _OFF_TASK_ARGS_BLOB, host_buf_ranges) + _rewrite_blob_host_addrs(task_buf, _OFF_TASK_ARGS_BLOB, host_buf_ranges) pipeline_slot, pipeline_reserved, pipeline_generation = _PIPELINE_LEASE_FMT.unpack_from( - buf, _OFF_PIPELINE_LEASE + task_buf, _OFF_PIPELINE_LEASE ) if pipeline_reserved != 0: raise RuntimeError(f"chip_process dev={device_id}: invalid pipeline lease reserved field") @@ -1574,11 +1598,11 @@ def handle_task() -> tuple[int, str]: # is the source of truth. cw._impl.run_from_blob( cid, - mailbox_addr + _OFF_TASK_ARGS_BLOB, + task_mailbox_addr + _OFF_TASK_ARGS_BLOB, _MAILBOX_ARGS_CAPACITY, cfg, - mailbox_addr + _OFF_ACCEPTED, - _TASK_ACCEPTED, + task_mailbox_addr + _OFF_ACCEPTED if publish_native_acceptance else 0, + _TASK_ACCEPTED if publish_native_acceptance else 0, pipeline_slot, pipeline_generation, ) @@ -1701,8 +1725,113 @@ def handle_control(sub_cmd: int) -> tuple[int, str]: # noqa: PLR0912 -- one bra msg = _format_exc(f"chip_process dev={device_id} ctrl={int(sub_cmd)}", e) return code, msg + def run_two_frame_loop() -> None: + action_cv = threading.Condition() + actions: list[tuple[str, int, tuple[int, int, int, int, int] | None]] = [] + stop_admission = threading.Event() + control_queued = False + frame_bufs = [ + buf[(1 + index) * MAILBOX_FRAME_SIZE : (2 + index) * MAILBOX_FRAME_SIZE] + for index in range(_TASK_FRAME_COUNT) + ] + frame_addrs = [mailbox_addr + (1 + index) * MAILBOX_FRAME_SIZE for index in range(_TASK_FRAME_COUNT)] + + def read_identity(frame_buf: memoryview) -> tuple[int, int, int, int, int]: + return ( + struct.unpack_from("=Q", frame_buf, _OFF_FRAME_PROTOCOL)[0], + struct.unpack_from("=Q", frame_buf, _OFF_FRAME_RUN_ID)[0], + struct.unpack_from("=Q", frame_buf, _OFF_FRAME_SLOT_ID)[0], + struct.unpack_from("=Q", frame_buf, _OFF_FRAME_GENERATION)[0], + struct.unpack_from("=Q", frame_buf, _OFF_FRAME_DISPATCH_ID)[0], + ) + + def admission_loop() -> None: + nonlocal control_queued + while not stop_admission.is_set(): + control_state = _mailbox_load_i32(state_addr) + with action_cv: + if control_state == _SHUTDOWN: + actions.append(("shutdown", -1, None)) + action_cv.notify() + return + if control_state == _CONTROL_REQUEST and not control_queued: + control_queued = True + actions.append(("control", -1, None)) + action_cv.notify() + + for index, (frame_buf, frame_addr) in enumerate(zip(frame_bufs, frame_addrs, strict=True)): + frame_state_addr = frame_addr + _OFF_STATE + if _mailbox_load_i32(frame_state_addr) != _TASK_READY: + continue + identity = read_identity(frame_buf) + protocol, run_id, slot_id, generation, dispatch_id = identity + if ( + protocol != _TASK_PROTOCOL_VERSION + or run_id == 0 + or slot_id >= _TASK_FRAME_COUNT + or generation == 0 + or dispatch_id == 0 + ): + _write_error( + frame_buf, + 1, + f"chip_process dev={device_id}: invalid task frame identity {identity}", + ) + _mailbox_store_i32(frame_state_addr, _TASK_FAILED) + continue + _mailbox_store_i32(frame_state_addr, _TASK_ACCEPTED_STATE) + with action_cv: + actions.append(("task", index, identity)) + action_cv.notify() + + admission_thread = threading.Thread( + target=admission_loop, + name=f"simpler-chip-admission-{device_id}", + daemon=True, + ) + admission_thread.start() + try: + while True: + with action_cv: + action_cv.wait_for(lambda: bool(actions)) + kind, index, identity = actions.pop(0) + if kind == "shutdown": + return + if kind == "control": + sub_cmd = struct.unpack_from("Q", buf, _OFF_CALLABLE)[0] + code, msg = handle_control(int(sub_cmd)) + _write_error(buf, code, msg) + _mailbox_store_i32(state_addr, _CONTROL_DONE) + with action_cv: + control_queued = False + continue + + frame_buf = frame_bufs[index] + frame_addr = frame_addrs[index] + frame_state_addr = frame_addr + _OFF_STATE + if identity is None or read_identity(frame_buf) != identity: + _write_error(frame_buf, 1, f"chip_process dev={device_id}: stale task frame identity") + _mailbox_store_i32(frame_state_addr, _TASK_FAILED) + continue + _mailbox_store_i32(frame_state_addr, _TASK_ACTIVE) + code, msg = handle_task( + frame_buf, + frame_addr, + publish_native_acceptance=False, + ) + _write_error(frame_buf, code, msg) + _mailbox_store_i32(frame_state_addr, _TASK_DONE if code == 0 else _TASK_FAILED) + finally: + stop_admission.set() + admission_thread.join() + for frame_buf in frame_bufs: + frame_buf.release() + try: - _run_mailbox_loop(buf, state_addr, handle_task=handle_task, handle_control=handle_control) + if task_frame_count >= 2: + run_two_frame_loop() + else: + _run_mailbox_loop(buf, state_addr, handle_task=handle_task, handle_control=handle_control) finally: _sweep_l2_host_l3_l2_regions(l3_l2_region_store) for host_shm, _lo, _hi, _base in host_buf_table.values(): @@ -1778,6 +1907,9 @@ 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. + # Before the first task, the lease word is startup metadata: slot_id carries + # the backend's supported admission depth. Dispatches later overwrite the + # same fixed wire region with the run-owned slot/generation lease. _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") @@ -1796,6 +1928,7 @@ def _chip_process_loop( # noqa: PLR0913 -- fork-child entry: all context (bins, chip_platform=platform, chip_runtime=runtime, prepared=prepared, + task_frame_count=_local_task_frame_count(platform, runtime, int(cw.pipeline_depth)), ) finally: cw.finalize() @@ -4611,6 +4744,8 @@ def _setup(): for shm in self._chip_shms: buf = shm.buf assert buf is not None + # INIT_READY repurposes the lease slot_id as the child's depth + # advertisement; task dispatch restores normal lease semantics. 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}") @@ -4685,8 +4820,11 @@ def _setup(inner=inner_worker): # mailbox that can no longer be completed. if device_ids: _require_matching_pids(self._chip_shms, self._chip_pids, "chip") - for shm, pid in zip(self._chip_shms, self._chip_pids): - dw.add_next_level_worker(_mailbox_addr(shm), pid) + task_frame_count = _local_task_frame_count( + str(self._config["platform"]), str(self._config["runtime"]), direct_chip_pipeline_depth + ) + for shm, pid in zip(self._chip_shms, self._chip_pids, strict=True): + dw.add_next_level_worker(_mailbox_addr(shm), pid, task_frame_count) # Register Worker children as NEXT_LEVEL (L4+) if self._next_level_shms and not hasattr(dw, "add_next_level_worker_at"): @@ -4697,7 +4835,7 @@ def _setup(inner=inner_worker): dw.add_next_level_worker_at(worker_id, _mailbox_addr(shm), self._next_level_pids[idx]) _require_matching_pids(self._sub_shms, self._sub_pids, "sub") - for shm, pid in zip(self._sub_shms, self._sub_pids): + for shm, pid in zip(self._sub_shms, self._sub_pids, strict=True): dw.add_sub_worker(_mailbox_addr(shm), pid) # Start Scheduler + WorkerThreads (C++ threads start here, after fork) diff --git a/src/common/hierarchical/scheduler.cpp b/src/common/hierarchical/scheduler.cpp index e819454f1..3c4cb7351 100644 --- a/src/common/hierarchical/scheduler.cpp +++ b/src/common/hierarchical/scheduler.cpp @@ -443,7 +443,7 @@ void Scheduler::dispatch_next_level_group() { if (std::find(workers.begin(), workers.end(), worker) != workers.end()) { throw std::runtime_error("Scheduler::dispatch_next_level_group: duplicate target worker"); } - if (!worker->idle()) return; + if (!worker->has_capacity()) return; workers.push_back(worker); } @@ -472,7 +472,7 @@ void Scheduler::dispatch_next_level_singles() { "Scheduler::dispatch_next_level_singles: unknown worker id " + std::to_string(worker_id) ); } - if (!worker->idle()) continue; + if (!worker->has_capacity()) continue; TaskSlot slot; while (cfg_.active_run_cb ? cfg_.ready_next_level_queues->try_pop_single(worker_id, active_run, slot) : @@ -488,7 +488,7 @@ void Scheduler::dispatch_next_level_singles() { } s.state.store(TaskState::RUNNING, std::memory_order_release); worker->dispatch(WorkerDispatch{slot, 0}); - break; + if (!worker->has_capacity()) break; } } } diff --git a/src/common/hierarchical/types.h b/src/common/hierarchical/types.h index 2c9083399..84b5dde7c 100644 --- a/src/common/hierarchical/types.h +++ b/src/common/hierarchical/types.h @@ -125,6 +125,7 @@ static constexpr int32_t INVALID_SLOT = -1; using RunId = uint64_t; static constexpr RunId INVALID_RUN_ID = 0; +using TaskSlot = int32_t; // Admission reserves a generation-safe pipeline slot before graph construction. // Closing the graph publishes PREPARED; only the whole-run FIFO head may enter @@ -151,7 +152,7 @@ struct RunState { mutable std::mutex completion_mu; std::condition_variable completion_cv; std::exception_ptr first_error; - std::vector task_slots; + std::vector task_slots; bool submission_closed{false}; bool submission_failed{false}; bool lease_released{false}; @@ -161,8 +162,6 @@ struct RunState { // Task slot index type // ============================================================================= -using TaskSlot = int32_t; - static constexpr size_t CALLABLE_HASH_DIGEST_SIZE = 32; enum class CallableKind : int32_t { diff --git a/src/common/hierarchical/worker.cpp b/src/common/hierarchical/worker.cpp index 751eeddc4..0d9e9aa3f 100644 --- a/src/common/hierarchical/worker.cpp +++ b/src/common/hierarchical/worker.cpp @@ -68,15 +68,15 @@ Worker::~Worker() { if (initialized_) close(); } -void Worker::add_worker(WorkerType type, void *mailbox, int child_pid) { +void Worker::add_worker(WorkerType type, void *mailbox, int child_pid, uint32_t task_frame_count) { if (initialized_) throw std::runtime_error("Worker: add_worker after init"); - if (type == WorkerType::NEXT_LEVEL) manager_.add_next_level(mailbox, child_pid); + if (type == WorkerType::NEXT_LEVEL) manager_.add_next_level(mailbox, child_pid, task_frame_count); else manager_.add_sub(mailbox, child_pid); } -void Worker::add_next_level_worker(int32_t worker_id, void *mailbox, int child_pid) { +void Worker::add_next_level_worker(int32_t worker_id, void *mailbox, int child_pid, uint32_t task_frame_count) { if (initialized_) throw std::runtime_error("Worker: add_next_level_worker after init"); - manager_.add_next_level_at(worker_id, mailbox, child_pid); + manager_.add_next_level_at(worker_id, mailbox, child_pid, task_frame_count); } void Worker::add_remote_l3_socket( diff --git a/src/common/hierarchical/worker.h b/src/common/hierarchical/worker.h index e6969f01c..af4a5b3b8 100644 --- a/src/common/hierarchical/worker.h +++ b/src/common/hierarchical/worker.h @@ -76,8 +76,8 @@ class Worker { // child and consumes the mailbox via the Python child loop. // `child_pid` is the forked child servicing `mailbox`, or -1 when the // caller owns no waitable child. - void add_worker(WorkerType type, void *mailbox, int child_pid = -1); - void add_next_level_worker(int32_t worker_id, void *mailbox, int child_pid = -1); + void add_worker(WorkerType type, void *mailbox, int child_pid = -1, uint32_t task_frame_count = 1); + void add_next_level_worker(int32_t worker_id, void *mailbox, int child_pid = -1, uint32_t task_frame_count = 1); // Register a REMOTE_L3 endpoint only after its session runner completed // prestart and reported HELLO READY on the command lane. diff --git a/src/common/hierarchical/worker_manager.cpp b/src/common/hierarchical/worker_manager.cpp index 22dc9f01a..970927dd4 100644 --- a/src/common/hierarchical/worker_manager.cpp +++ b/src/common/hierarchical/worker_manager.cpp @@ -150,14 +150,20 @@ WorkerEndpoint::run_with_accept(Ring *ring, const WorkerDispatch &dispatch, cons // LocalMailboxEndpoint — mailbox helpers // ============================================================================= -LocalMailboxEndpoint::LocalMailboxEndpoint(int32_t worker_id, void *mailbox, int child_pid) : +LocalMailboxEndpoint::LocalMailboxEndpoint(int32_t worker_id, void *mailbox, int child_pid, uint32_t task_frame_count) : mailbox_(mailbox), + task_frame_count_(task_frame_count), child_pid_(child_pid) { if (mailbox == nullptr) throw std::invalid_argument("LocalMailboxEndpoint: null mailbox"); + if (task_frame_count == 0 || task_frame_count > MAILBOX_TASK_FRAME_COUNT) { + throw std::invalid_argument("LocalMailboxEndpoint: invalid task frame count"); + } caps_.worker_id = worker_id; + caps_.max_inflight_tasks = task_frame_count; } std::string LocalMailboxEndpoint::check_child_death() { + std::lock_guard child_lk(child_mu_); if (child_dead_) return child_death_reason_; if (child_pid_ <= 0) return {}; @@ -185,8 +191,9 @@ std::string LocalMailboxEndpoint::check_child_death() { return child_death_reason_; } -MailboxState LocalMailboxEndpoint::read_mailbox_state() const { - volatile int32_t *ptr = reinterpret_cast(mbox() + MAILBOX_OFF_STATE); +MailboxState LocalMailboxEndpoint::read_mailbox_state(const char *frame) const { + const char *base = frame == nullptr ? mbox() : frame; + volatile int32_t *ptr = reinterpret_cast(const_cast(base) + MAILBOX_OFF_STATE); int32_t v; #if defined(__aarch64__) __asm__ volatile("ldar %w0, [%1]" : "=r"(v) : "r"(ptr) : "memory"); @@ -199,8 +206,9 @@ MailboxState LocalMailboxEndpoint::read_mailbox_state() const { return static_cast(v); } -void LocalMailboxEndpoint::write_mailbox_state(MailboxState s) { - volatile int32_t *ptr = reinterpret_cast(mbox() + MAILBOX_OFF_STATE); +void LocalMailboxEndpoint::write_mailbox_state(MailboxState s, char *frame) { + char *base = frame == nullptr ? mbox() : frame; + volatile int32_t *ptr = reinterpret_cast(base + MAILBOX_OFF_STATE); int32_t v = static_cast(s); #if defined(__aarch64__) __asm__ volatile("stlr %w0, [%1]" : : "r"(v), "r"(ptr) : "memory"); @@ -212,21 +220,28 @@ void LocalMailboxEndpoint::write_mailbox_state(MailboxState s) { #endif } -bool LocalMailboxEndpoint::read_task_accepted() const { - const int32_t *ptr = reinterpret_cast(mbox() + MAILBOX_OFF_ACCEPTED); +bool LocalMailboxEndpoint::read_task_accepted(const char *frame) const { + const char *base = frame == nullptr ? mbox() : frame; + const int32_t *ptr = reinterpret_cast(base + MAILBOX_OFF_ACCEPTED); int32_t v = 0; __atomic_load(ptr, &v, __ATOMIC_ACQUIRE); return v == MAILBOX_TASK_ACCEPTED; } -void LocalMailboxEndpoint::clear_task_accepted() { - int32_t *ptr = reinterpret_cast(mbox() + MAILBOX_OFF_ACCEPTED); +void LocalMailboxEndpoint::clear_task_accepted(char *frame) { + char *base = frame == nullptr ? mbox() : frame; + int32_t *ptr = reinterpret_cast(base + MAILBOX_OFF_ACCEPTED); int32_t v = 0; __atomic_store(ptr, &v, __ATOMIC_RELEASE); } void LocalMailboxEndpoint::shutdown_child() { write_mailbox_state(MailboxState::SHUTDOWN); } +char *LocalMailboxEndpoint::task_frame(size_t index) const { + return mbox() + + static_cast(MAILBOX_FIRST_TASK_FRAME + index) * static_cast(MAILBOX_FRAME_SIZE); +} + // ============================================================================= // WorkerThread — lifecycle // ============================================================================= @@ -241,12 +256,22 @@ void WorkerThread::start( on_accept_ = on_accept; endpoint_ = std::move(endpoint); shutdown_ = false; - idle_.store(true, std::memory_order_relaxed); - thread_ = std::thread(&WorkerThread::loop, this); + capacity_ = endpoint_->caps().max_inflight_tasks; + if (capacity_ == 0) throw std::invalid_argument("WorkerThread::start: endpoint capacity is zero"); + inflight_.store(0, std::memory_order_relaxed); + threads_.reserve(capacity_); + for (uint32_t i = 0; i < capacity_; ++i) { + threads_.emplace_back(&WorkerThread::loop, this); + } } void WorkerThread::dispatch(WorkerDispatch d) { - idle_.store(false, std::memory_order_release); + uint32_t previous = inflight_.fetch_add(1, std::memory_order_acq_rel); + if (previous >= capacity_) { + inflight_.fetch_sub(1, std::memory_order_acq_rel); + throw std::logic_error("WorkerThread::dispatch: endpoint capacity exceeded"); + } + d.dispatch_id = next_dispatch_id_.fetch_add(1, std::memory_order_relaxed); std::lock_guard lk(mu_); queue_.push(d); cv_.notify_one(); @@ -258,7 +283,10 @@ void WorkerThread::stop() { shutdown_ = true; } cv_.notify_all(); - if (thread_.joinable()) thread_.join(); + for (auto &thread : threads_) { + if (thread.joinable()) thread.join(); + } + threads_.clear(); } void WorkerThread::shutdown_child() { @@ -328,7 +356,7 @@ void WorkerThread::loop() { } } - idle_.store(true, std::memory_order_release); + inflight_.fetch_sub(1, std::memory_order_acq_rel); on_complete_(std::move(completion)); } } @@ -345,6 +373,7 @@ WorkerCompletion LocalMailboxEndpoint::run(Ring *ring, const WorkerDispatch &dis WorkerCompletion LocalMailboxEndpoint::run_with_accept( Ring *ring, const WorkerDispatch &dispatch, const std::function &on_accept ) { + if (task_frame_count_ > 1) return run_two_frame(ring, dispatch, on_accept); if (ring == nullptr) throw std::invalid_argument("LocalMailboxEndpoint::run: null ring"); TaskSlotState &s = *ring->slot_state(dispatch.task_slot); int32_t group_index = dispatch.group_index; @@ -360,7 +389,6 @@ WorkerCompletion LocalMailboxEndpoint::run_with_accept( WorkerCompletion completion; completion.task_slot = dispatch.task_slot; completion.group_index = group_index; - // Hold mailbox_mu_ for the entire round trip (write payload + state + // spin-poll TASK_DONE + reset to IDLE). Any control_* request from the // orch thread waits for the dispatch to finish before claiming the @@ -472,17 +500,261 @@ WorkerCompletion LocalMailboxEndpoint::run_with_accept( return completion; } +WorkerCompletion LocalMailboxEndpoint::run_two_frame( + Ring *ring, const WorkerDispatch &dispatch, const std::function &on_accept +) { + WorkerCompletion completion; + completion.task_slot = dispatch.task_slot; + completion.group_index = dispatch.group_index; + + bool publish_sequence_retired = false; + bool task_published = false; + auto retire_publish_sequence = [&]() { + if (publish_sequence_retired) return; + std::unique_lock publish_lk(publish_mu_); + publish_cv_.wait(publish_lk, [&] { + return dispatch.dispatch_id == next_publish_id_ || endpoint_poisoned_.load(std::memory_order_acquire); + }); + if (dispatch.dispatch_id == next_publish_id_) ++next_publish_id_; + publish_sequence_retired = true; + publish_lk.unlock(); + publish_cv_.notify_all(); + }; + + size_t frame_index = MAILBOX_TASK_FRAME_COUNT; + bool frame_is_claimed = false; + auto release_claim = [&]() { + if (!frame_is_claimed) return; + std::lock_guard claim_lk(frame_claim_mu_); + frame_claimed_[frame_index] = false; + frame_is_claimed = false; + }; + + try { + if (ring == nullptr) throw std::invalid_argument("LocalMailboxEndpoint::run_two_frame: null ring"); + TaskSlotState *slot_state = ring->slot_state(dispatch.task_slot); + if (slot_state == nullptr) { + throw std::out_of_range("LocalMailboxEndpoint::run_two_frame: invalid task slot"); + } + TaskSlotState &s = *slot_state; + const int32_t group_index = dispatch.group_index; + const TaskArgsView view = s.args_view(group_index); + + if (!s.remote_sidecar_for(group_index).empty()) { + retire_publish_sequence(); + completion.outcome = EndpointOutcome::ENDPOINT_FAILURE; + completion.error_message = + "LocalMailboxEndpoint::run_two_frame: remote task sidecar is not supported by local mailbox"; + return completion; + } + + { + std::lock_guard claim_lk(frame_claim_mu_); + for (size_t i = 0; i < task_frame_count_; ++i) { + if (!frame_claimed_[i]) { + frame_claimed_[i] = true; + frame_index = i; + frame_is_claimed = true; + break; + } + } + } + if (!frame_is_claimed) { + retire_publish_sequence(); + completion.outcome = EndpointOutcome::ENDPOINT_FAILURE; + completion.error_message = "LocalMailboxEndpoint::run_two_frame: no free task frame"; + return completion; + } + + std::lock_guard frame_lk(frame_mu_[frame_index]); + char *frame = task_frame(frame_index); + if (endpoint_poisoned_.load(std::memory_order_acquire)) { + retire_publish_sequence(); + completion.outcome = EndpointOutcome::ENDPOINT_FAILURE; + completion.error_message = "LocalMailboxEndpoint::run_two_frame: endpoint is poisoned"; + release_claim(); + return completion; + } + if (read_mailbox_state(frame) != MailboxState::IDLE) { + endpoint_poisoned_.store(true, std::memory_order_release); + publish_cv_.notify_all(); + retire_publish_sequence(); + completion.outcome = EndpointOutcome::ENDPOINT_FAILURE; + completion.error_message = "LocalMailboxEndpoint::run_two_frame: claimed frame is not IDLE"; + release_claim(); + return completion; + } + + int32_t zero_err = 0; + std::memcpy(frame + MAILBOX_OFF_ERROR, &zero_err, sizeof(zero_err)); + std::memset(frame + MAILBOX_OFF_ERROR_MSG, 0, MAILBOX_ERROR_MSG_SIZE); + clear_task_accepted(frame); + + uint64_t reserved_callable = 0; + std::memcpy(frame + MAILBOX_OFF_CALLABLE, &reserved_callable, sizeof(reserved_callable)); + std::memcpy(frame + MAILBOX_OFF_CONFIG, &s.config, sizeof(CallConfig)); + std::memcpy(frame + MAILBOX_OFF_PIPELINE_LEASE, &s.pipeline_lease, sizeof(PipelineSlotLease)); + std::memcpy( + frame + MAILBOX_OFF_TASK_CALLABLE_HASH, s.callable.digest.data(), + static_cast(CALLABLE_HASH_DIGEST_SIZE) + ); + + const size_t blob_bytes = TASK_ARGS_BLOB_HEADER_SIZE + static_cast(view.tensor_count) * sizeof(Tensor) + + static_cast(view.scalar_count) * sizeof(uint64_t); + if (blob_bytes > MAILBOX_ARGS_CAPACITY) { + retire_publish_sequence(); + completion.outcome = EndpointOutcome::ENDPOINT_FAILURE; + completion.error_message = + "LocalMailboxEndpoint::run_two_frame: args blob exceeds mailbox capacity: need " + + std::to_string(blob_bytes) + ", capacity " + std::to_string(MAILBOX_ARGS_CAPACITY); + release_claim(); + return completion; + } + uint8_t *blob = reinterpret_cast(frame + MAILBOX_OFF_TASK_ARGS_BLOB); + std::memcpy(blob, &view.tensor_count, sizeof(int32_t)); + std::memcpy(blob + 4, &view.scalar_count, sizeof(int32_t)); + if (view.tensor_count > 0) { + std::memcpy( + blob + TASK_ARGS_BLOB_HEADER_SIZE, view.tensor_bytes, + static_cast(view.tensor_count) * sizeof(Tensor) + ); + } + if (view.scalar_count > 0) { + std::memcpy( + blob + TASK_ARGS_BLOB_HEADER_SIZE + static_cast(view.tensor_count) * sizeof(Tensor), + view.scalars, static_cast(view.scalar_count) * sizeof(uint64_t) + ); + } + + const uint64_t protocol = MAILBOX_TASK_PROTOCOL_VERSION; + const uint64_t slot_id = s.pipeline_lease.slot_id; + std::memcpy(frame + MAILBOX_OFF_FRAME_PROTOCOL, &protocol, sizeof(protocol)); + std::memcpy(frame + MAILBOX_OFF_FRAME_RUN_ID, &s.run_id, sizeof(s.run_id)); + std::memcpy(frame + MAILBOX_OFF_FRAME_SLOT_ID, &slot_id, sizeof(slot_id)); + std::memcpy(frame + MAILBOX_OFF_FRAME_GENERATION, &s.pipeline_lease.generation, sizeof(uint64_t)); + std::memcpy(frame + MAILBOX_OFF_FRAME_DISPATCH_ID, &dispatch.dispatch_id, sizeof(dispatch.dispatch_id)); + + { + std::unique_lock publish_lk(publish_mu_); + publish_cv_.wait(publish_lk, [&] { + return dispatch.dispatch_id == next_publish_id_ || endpoint_poisoned_.load(std::memory_order_acquire); + }); + if (endpoint_poisoned_.load(std::memory_order_acquire)) { + publish_sequence_retired = true; + completion.outcome = EndpointOutcome::ENDPOINT_FAILURE; + completion.error_message = "LocalMailboxEndpoint::run_two_frame: endpoint poisoned before publish"; + release_claim(); + return completion; + } + write_mailbox_state(MailboxState::TASK_READY, frame); + ++next_publish_id_; + publish_sequence_retired = true; + task_published = true; + } + publish_cv_.notify_all(); + + const auto identity_matches = [&]() { + uint64_t observed_protocol = 0; + RunId observed_run = INVALID_RUN_ID; + uint64_t observed_slot = 0; + uint64_t observed_generation = 0; + uint64_t observed_dispatch = 0; + std::memcpy(&observed_protocol, frame + MAILBOX_OFF_FRAME_PROTOCOL, sizeof(observed_protocol)); + std::memcpy(&observed_run, frame + MAILBOX_OFF_FRAME_RUN_ID, sizeof(observed_run)); + std::memcpy(&observed_slot, frame + MAILBOX_OFF_FRAME_SLOT_ID, sizeof(observed_slot)); + std::memcpy(&observed_generation, frame + MAILBOX_OFF_FRAME_GENERATION, sizeof(observed_generation)); + std::memcpy(&observed_dispatch, frame + MAILBOX_OFF_FRAME_DISPATCH_ID, sizeof(observed_dispatch)); + return observed_protocol == protocol && observed_run == s.run_id && observed_slot == slot_id && + observed_generation == s.pipeline_lease.generation && observed_dispatch == dispatch.dispatch_id; + }; + + auto next_liveness_check = std::chrono::steady_clock::now() + kChildLivenessPollPeriod; + bool acceptance_observed = false; + MailboxState terminal_state = MailboxState::IDLE; + while (true) { + MailboxState state = read_mailbox_state(frame); + if (!acceptance_observed && (state == MailboxState::TASK_ACCEPTED || state == MailboxState::TASK_ACTIVE || + state == MailboxState::TASK_DONE || state == MailboxState::TASK_FAILED)) { + if (!identity_matches()) { + endpoint_poisoned_.store(true, std::memory_order_release); + publish_cv_.notify_all(); + completion.outcome = EndpointOutcome::ENDPOINT_FAILURE; + completion.error_message = "LocalMailboxEndpoint::run_two_frame: stale frame acceptance identity"; + release_claim(); + return completion; + } + acceptance_observed = true; + if (on_accept) on_accept(); + } + if (state == MailboxState::TASK_DONE || state == MailboxState::TASK_FAILED) { + terminal_state = state; + break; + } + auto now = std::chrono::steady_clock::now(); + if (now >= next_liveness_check) { + next_liveness_check = now + kChildLivenessPollPeriod; + std::string death = check_child_death(); + if (!death.empty()) { + endpoint_poisoned_.store(true, std::memory_order_release); + publish_cv_.notify_all(); + completion.outcome = EndpointOutcome::ENDPOINT_FAILURE; + completion.error_message = "LocalMailboxEndpoint::run_two_frame: " + death; + release_claim(); + return completion; + } + } + } + + if (!identity_matches()) { + endpoint_poisoned_.store(true, std::memory_order_release); + publish_cv_.notify_all(); + completion.outcome = EndpointOutcome::ENDPOINT_FAILURE; + completion.error_message = "LocalMailboxEndpoint::run_two_frame: stale frame terminal identity"; + release_claim(); + return completion; + } + + int32_t error_code = 0; + std::memcpy(&error_code, frame + MAILBOX_OFF_ERROR, sizeof(error_code)); + if (terminal_state == MailboxState::TASK_FAILED || error_code != 0) { + completion.outcome = EndpointOutcome::TASK_FAILURE; + completion.error_message = + "LocalMailboxEndpoint::run_two_frame: child failed (worker_id=" + std::to_string(caps_.worker_id) + + ", code=" + std::to_string(error_code) + "): " + read_error_msg(frame); + } else { + completion.outcome = EndpointOutcome::SUCCESS; + } + write_mailbox_state(MailboxState::IDLE, frame); + release_claim(); + return completion; + } catch (...) { + if (task_published) { + endpoint_poisoned_.store(true, std::memory_order_release); + publish_cv_.notify_all(); + } else if (!publish_sequence_retired) { + try { + retire_publish_sequence(); + } catch (...) { + endpoint_poisoned_.store(true, std::memory_order_release); + publish_cv_.notify_all(); + } + } + release_claim(); + throw; + } +} + // ============================================================================= // WorkerManager // ============================================================================= -void WorkerManager::add_next_level(void *mailbox, int child_pid) { - add_next_level_at(static_cast(next_level_entries_.size()), mailbox, child_pid); +void WorkerManager::add_next_level(void *mailbox, int child_pid, uint32_t task_frame_count) { + add_next_level_at(static_cast(next_level_entries_.size()), mailbox, child_pid, task_frame_count); } -void WorkerManager::add_next_level_at(int32_t worker_id, void *mailbox, int child_pid) { +void WorkerManager::add_next_level_at(int32_t worker_id, void *mailbox, int child_pid, uint32_t task_frame_count) { if (worker_id < 0) throw std::invalid_argument("WorkerManager::add_next_level_at: negative worker_id"); - next_level_entries_.push_back(LocalNextLevelEntry{worker_id, mailbox, child_pid}); + next_level_entries_.push_back(LocalNextLevelEntry{worker_id, mailbox, child_pid, task_frame_count}); } void WorkerManager::add_next_level_endpoint(std::unique_ptr endpoint) { @@ -519,7 +791,9 @@ void WorkerManager::start(Ring *ring, const OnCompleteFn &on_complete, const OnA auto make_next_level_threads = [&]() { for (const auto &entry : next_level_entries_) { auto wt = std::make_unique(); - auto endpoint = std::make_unique(entry.worker_id, entry.mailbox, entry.child_pid); + auto endpoint = std::make_unique( + entry.worker_id, entry.mailbox, entry.child_pid, entry.task_frame_count + ); wt->start(ring, on_complete, on_accept, std::move(endpoint)); next_level_threads_.push_back(std::move(wt)); } @@ -634,7 +908,9 @@ void LocalMailboxEndpoint::run_control_command(const char *op_name, double timeo while (read_mailbox_state() != MailboxState::CONTROL_DONE) { auto now = std::chrono::steady_clock::now(); if (now >= deadline) { - mailbox_control_timed_out_ = true; + mailbox_control_timed_out_.store(true, std::memory_order_release); + endpoint_poisoned_.store(true, std::memory_order_release); + publish_cv_.notify_all(); throw std::runtime_error(std::string(op_name) + " timed out waiting for CONTROL_DONE"); } if (now >= next_liveness_check) { @@ -644,7 +920,9 @@ void LocalMailboxEndpoint::run_control_command(const char *op_name, double timeo // The mailbox is poisoned rather than reset to IDLE: with the // child gone no later command can complete, so admitting one // would restore the hang this check exists to break. - mailbox_control_timed_out_ = true; + mailbox_control_timed_out_.store(true, std::memory_order_release); + endpoint_poisoned_.store(true, std::memory_order_release); + publish_cv_.notify_all(); throw std::runtime_error(std::string(op_name) + ": " + death); } } diff --git a/src/common/hierarchical/worker_manager.h b/src/common/hierarchical/worker_manager.h index 257b39b28..189807da7 100644 --- a/src/common/hierarchical/worker_manager.h +++ b/src/common/hierarchical/worker_manager.h @@ -27,6 +27,7 @@ #pragma once #include +#include #include #include #include @@ -72,6 +73,9 @@ enum class MailboxState : int32_t { // PLATFORM_STREAM_SYNC_TIMEOUT_MS budget (issue #897). INIT_READY = 6, INIT_FAILED = 7, + TASK_ACCEPTED = 8, + TASK_ACTIVE = 9, + TASK_FAILED = 10, }; // Sized so the args region can hold any TaskArgs the runtime itself accepts @@ -82,7 +86,12 @@ enum class MailboxState : int32_t { // to the unified 128 B Tensor: the worst-case blob (CHIP_MAX_TENSOR_ARGS tensors) // grew ~3x, and 128*128 B = 16 KB alone exceeded the old mailbox (see the // capacity static_assert after MAILBOX_ARGS_CAPACITY). -static constexpr size_t MAILBOX_SIZE = 32768; +static constexpr size_t MAILBOX_FRAME_SIZE = 32768; +static constexpr size_t MAILBOX_TASK_FRAME_COUNT = 2; +static constexpr size_t MAILBOX_CONTROL_FRAME = 0; +static constexpr size_t MAILBOX_FIRST_TASK_FRAME = 1; +static constexpr size_t MAILBOX_SIZE = MAILBOX_FRAME_SIZE * (1 + MAILBOX_TASK_FRAME_COUNT); +static constexpr uint32_t MAILBOX_TASK_PROTOCOL_VERSION = 1; // Error message region lives at the mailbox tail. 256 B of headroom is // enough for `: ` produced by the child-side @@ -117,7 +126,7 @@ static_assert( "PipelineSlotLease overflows reserved lease region" ); static constexpr ptrdiff_t MAILBOX_OFF_ERROR_MSG = - static_cast(MAILBOX_SIZE) - static_cast(MAILBOX_ERROR_MSG_SIZE); + static_cast(MAILBOX_FRAME_SIZE) - static_cast(MAILBOX_ERROR_MSG_SIZE); // Launch acceptance is sticky, not a MailboxState: a state word carrying it // loses the ACK whenever the child reaches TASK_DONE between two parent polls, // which is exactly the short-task case the fence exists to pipeline. The child @@ -125,14 +134,23 @@ static constexpr ptrdiff_t MAILBOX_OFF_ERROR_MSG = // TASK_READY, so nothing else can overwrite it in between. static constexpr int32_t MAILBOX_TASK_ACCEPTED = 1; static constexpr ptrdiff_t MAILBOX_OFF_ACCEPTED = MAILBOX_OFF_ERROR_MSG - 8; +static constexpr ptrdiff_t MAILBOX_OFF_FRAME_PROTOCOL = MAILBOX_OFF_ACCEPTED - 40; +static constexpr ptrdiff_t MAILBOX_OFF_FRAME_RUN_ID = MAILBOX_OFF_ACCEPTED - 32; +static constexpr ptrdiff_t MAILBOX_OFF_FRAME_SLOT_ID = MAILBOX_OFF_ACCEPTED - 24; +static constexpr ptrdiff_t MAILBOX_OFF_FRAME_GENERATION = MAILBOX_OFF_ACCEPTED - 16; +static constexpr ptrdiff_t MAILBOX_OFF_FRAME_DISPATCH_ID = MAILBOX_OFF_ACCEPTED - 8; static constexpr ptrdiff_t MAILBOX_OFF_TASK_CALLABLE_HASH = MAILBOX_OFF_ARGS; static constexpr ptrdiff_t MAILBOX_OFF_TASK_ARGS_BLOB = MAILBOX_OFF_TASK_CALLABLE_HASH + static_cast(CALLABLE_HASH_DIGEST_SIZE); static constexpr size_t CTRL_SHM_NAME_BYTES = 32; static constexpr ptrdiff_t MAILBOX_OFF_CONTROL_CALLABLE_HASH = MAILBOX_OFF_ARGS + static_cast(CTRL_SHM_NAME_BYTES); +static_assert( + MAILBOX_OFF_TASK_ARGS_BLOB < MAILBOX_OFF_FRAME_PROTOCOL, + "mailbox task-args region must precede the frame protocol trailer" +); static constexpr size_t MAILBOX_ARGS_CAPACITY = - MAILBOX_SIZE - static_cast(MAILBOX_OFF_TASK_ARGS_BLOB) - MAILBOX_ERROR_MSG_SIZE - 8; + static_cast(MAILBOX_OFF_FRAME_PROTOCOL) - static_cast(MAILBOX_OFF_TASK_ARGS_BLOB); static_assert( MAILBOX_ARGS_CAPACITY >= TASK_ARGS_BLOB_HEADER_SIZE + static_cast(CHIP_MAX_TENSOR_ARGS) * sizeof(Tensor) + static_cast(CHIP_MAX_SCALAR_ARGS) * sizeof(uint64_t), @@ -207,6 +225,7 @@ struct WorkerEndpointCaps { bool remote{false}; bool supports_task_dispatch{true}; bool supports_control{true}; + uint32_t max_inflight_tasks{1}; std::string transport{"local-mailbox"}; }; @@ -271,7 +290,7 @@ class LocalMailboxEndpoint : public WorkerEndpoint { // caller does not own a waitable child. A valid pid enables liveness // checks that turn a child that dies before publishing TASK_DONE / // CONTROL_DONE into a reported failure instead of an endless spin-poll. - LocalMailboxEndpoint(int32_t worker_id, void *mailbox, int child_pid = -1); + LocalMailboxEndpoint(int32_t worker_id, void *mailbox, int child_pid = -1, uint32_t task_frame_count = 1); const WorkerEndpointCaps &caps() const override { return caps_; } WorkerCompletion run(Ring *ring, const WorkerDispatch &dispatch) override; @@ -325,7 +344,16 @@ class LocalMailboxEndpoint : public WorkerEndpoint { WorkerEndpointCaps caps_; void *mailbox_{nullptr}; std::mutex mailbox_mu_; - bool mailbox_control_timed_out_{false}; + std::mutex child_mu_; + std::mutex frame_claim_mu_; + std::mutex publish_mu_; + std::condition_variable publish_cv_; + std::array frame_mu_; + std::array frame_claimed_{}; + uint32_t task_frame_count_{1}; + std::atomic endpoint_poisoned_{false}; + uint64_t next_publish_id_{1}; + std::atomic mailbox_control_timed_out_{false}; int child_pid_{-1}; // Set once the child has been reaped or observed unwaitable; the exit // status is only available from the reaping waitpid(), so it is retained @@ -334,12 +362,14 @@ class LocalMailboxEndpoint : public WorkerEndpoint { std::string child_death_reason_; char *mbox() const { return static_cast(mailbox_); } - MailboxState read_mailbox_state() const; - void write_mailbox_state(MailboxState s); + MailboxState read_mailbox_state(const char *frame = nullptr) const; + void write_mailbox_state(MailboxState s, char *frame = nullptr); // Sticky launch acceptance, cleared only when this endpoint publishes the // next TASK_READY. See MAILBOX_OFF_ACCEPTED. - bool read_task_accepted() const; - void clear_task_accepted(); + bool read_task_accepted(const char *frame = nullptr) const; + void clear_task_accepted(char *frame = nullptr); + WorkerCompletion run_two_frame(Ring *ring, const WorkerDispatch &dispatch, const std::function &on_accept); + char *task_frame(size_t index) const; void run_control_command(const char *op_name, double timeout_s = -1.0); // Returns a description of the child's death, or an empty string while it @@ -359,6 +389,7 @@ class LocalMailboxEndpoint : public WorkerEndpoint { struct WorkerDispatch { TaskSlot task_slot{INVALID_SLOT}; int32_t group_index{0}; + uint64_t dispatch_id{0}; }; // ============================================================================= @@ -393,7 +424,8 @@ class WorkerThread { void dispatch(WorkerDispatch d); // True if the worker has no active task. - bool idle() const { return idle_.load(std::memory_order_acquire); } + bool idle() const { return inflight_.load(std::memory_order_acquire) == 0; } + bool has_capacity() const { return inflight_.load(std::memory_order_acquire) < capacity_; } const WorkerEndpointCaps &caps() const; int32_t worker_id() const; @@ -478,12 +510,14 @@ class WorkerThread { std::function on_complete_; std::function on_accept_; - std::thread thread_; + std::vector threads_; std::queue queue_; std::mutex mu_; std::condition_variable cv_; bool shutdown_{false}; - std::atomic idle_{true}; + uint32_t capacity_{1}; + std::atomic inflight_{0}; + std::atomic next_dispatch_id_{1}; void loop(); WorkerCompletion dispatch_process(WorkerDispatch d, const std::function &on_accept); @@ -503,8 +537,8 @@ class WorkerManager { // callable for SUB) lives in the forked child. // `child_pid` is the forked child servicing `mailbox`; pass -1 when the // caller owns no waitable child. - void add_next_level(void *mailbox, int child_pid = -1); - void add_next_level_at(int32_t worker_id, void *mailbox, int child_pid = -1); + void add_next_level(void *mailbox, int child_pid = -1, uint32_t task_frame_count = 1); + void add_next_level_at(int32_t worker_id, void *mailbox, int child_pid = -1, uint32_t task_frame_count = 1); void add_next_level_endpoint(std::unique_ptr endpoint); void add_sub(void *mailbox, int child_pid = -1); @@ -590,6 +624,7 @@ class WorkerManager { int32_t worker_id{-1}; void *mailbox{nullptr}; int child_pid{-1}; + uint32_t task_frame_count{1}; }; struct LocalSubEntry { void *mailbox{nullptr}; diff --git a/tests/st/a2a3/host_build_graph/worker_async_endpoint/kernels/orchestration/long_vector_orch.cpp b/tests/st/a2a3/host_build_graph/worker_async_endpoint/kernels/orchestration/long_vector_orch.cpp new file mode 100644 index 000000000..ef460db82 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/worker_async_endpoint/kernels/orchestration/long_vector_orch.cpp @@ -0,0 +1,66 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include + +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) + +namespace { + +constexpr uint64_t kAdd = 0; +constexpr uint64_t kAddScalar = 1; +constexpr int kChainLength = 64; + +} // namespace + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &args) { + (void)args; + return PTO2OrchestrationConfig{.expected_arg_count = 3}; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &args) { + const Tensor &a = args.tensor(0).ref(); + const Tensor &b = args.tensor(1).ref(); + const Tensor &out = args.tensor(2).ref(); + uint32_t shape[1] = {a.shapes[0]}; + TensorCreateInfo temporary(shape, 1, DataType::FLOAT32); + + L0TaskArgs add_args; + add_args.add_input(a); + add_args.add_input(b); + add_args.add_output(temporary); + TaskOutputTensors add_outputs = rt_submit_aiv_task(kAdd, add_args); + Tensor current = add_outputs.get_ref(0); + + union { + float f32; + uint64_t u64; + } scalar{}; + scalar.f32 = 1.0F; + for (int i = 0; i < kChainLength; ++i) { + L0TaskArgs step_args; + step_args.add_input(current); + if (i + 1 == kChainLength) { + step_args.add_output(out); + step_args.add_scalar(scalar.u64); + rt_submit_aiv_task(kAddScalar, step_args); + } else { + step_args.add_output(temporary); + step_args.add_scalar(scalar.u64); + TaskOutputTensors step_outputs = rt_submit_aiv_task(kAddScalar, step_args); + current = step_outputs.get_ref(0); + } + } +} + +} // extern "C" diff --git a/tests/st/a2a3/host_build_graph/worker_async_endpoint/test_worker_async_endpoint.py b/tests/st/a2a3/host_build_graph/worker_async_endpoint/test_worker_async_endpoint.py new file mode 100644 index 000000000..bd4fe56f7 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/worker_async_endpoint/test_worker_async_endpoint.py @@ -0,0 +1,137 @@ +#!/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 the generation-safe two-frame local endpoint.""" + +import ctypes +import time +from contextlib import suppress + +import pytest +import torch +from simpler.task_interface import ArgDirection as D +from simpler.worker import ( + _OFF_STATE, + _TASK_ACCEPTED_STATE, + _TASK_ACTIVE, + MAILBOX_FRAME_SIZE, + _mailbox_load_i32, +) + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler_setup.scene_test import _build_l3_task_args + +_VECTOR_KERNELS = "../vector_example/kernels/aiv" +_SIZE = 128 * 128 +_CHAIN_LENGTH = 64 + + +@scene_test(level=3, runtime="host_build_graph") +class TestWorkerAsyncEndpoint(SceneTestCase): + CALLABLE = { + "callables": [ + { + "name": "long_vector", + "orchestration": { + "source": "kernels/orchestration/long_vector_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "source": f"{_VECTOR_KERNELS}/kernel_add.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "source": f"{_VECTOR_KERNELS}/kernel_add_scalar.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT], + }, + ], + }, + ], + } + + CASES = [ + { + "name": "two_frame_sequential_execution", + "platforms": ["a2a3"], + "config": {"device_count": 1, "aicpu_thread_num": 4}, + "params": {}, + }, + ] + + @staticmethod + def _tensor_from_host_buffer(worker, value): + buffer = worker.create_host_buffer(_SIZE * torch.float32.itemsize) + tensor = torch.frombuffer(buffer.buffer, dtype=torch.float32, count=_SIZE) + tensor.fill_(value) + return buffer, tensor + + def test_run(self, st_platform, st_worker): + if st_platform != "a2a3": + pytest.skip("the two-frame local endpoint is enabled on a2a3 onboard workers") + + buffers = [] + tensors = [] + run = 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 + handle = type(self)._st_chip_handles["long_vector"] + signature = type(self)._st_chip_handles["long_vector_sig"] + cfg = self._build_config(self.CASES[0]["config"]) + + def submit_task(orch, a, b, out): + builder = TaskArgsBuilder(Tensor("a", a), Tensor("b", b), Tensor("out", out)) + chip_args, _ = _build_l3_task_args(builder, signature) + orch.submit_next_level(handle, chip_args, cfg, worker=0) + + run = st_worker.submit( + lambda orch, _args, _cfg: ( + submit_task(orch, first_a, first_b, first_out), + submit_task(orch, second_a, second_b, second_out), + ) + ) + + shm_buf = st_worker._chip_shms[0].buf + assert shm_buf is not None + mailbox_addr = ctypes.addressof(ctypes.c_char.from_buffer(shm_buf)) + frame_state_addrs = [mailbox_addr + (index + 1) * MAILBOX_FRAME_SIZE + _OFF_STATE for index in range(2)] + saw_active_and_accepted = False + deadline = time.monotonic() + 10.0 + while not run.done and time.monotonic() < deadline: + states = [_mailbox_load_i32(addr) for addr in frame_state_addrs] + if _TASK_ACTIVE in states and _TASK_ACCEPTED_STATE in states: + saw_active_and_accepted = True + break + + run.wait(20.0) + assert saw_active_and_accepted, "frame B was not accepted while frame A was active" + assert torch.allclose(first_out, first_a + first_b + _CHAIN_LENGTH) + assert torch.allclose(second_out, second_a + second_b + _CHAIN_LENGTH) + finally: + if run is not None and not run.done: + with suppress(Exception): + run.wait(20.0) + tensors.clear() + first_a = first_b = first_out = second_a = second_b = second_out = None + if run is None or run.done: + 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 e4e8c0e44..6beca4b60 100644 --- a/tests/ut/cpp/hierarchical/test_orchestrator.cpp +++ b/tests/ut/cpp/hierarchical/test_orchestrator.cpp @@ -758,13 +758,25 @@ TEST_F(OrchestratorFixture, PreparedRunWaitsForActiveRunAndThirdAdmissionBlocks) S(first.task_slot).state.store(TaskState::COMPLETED, std::memory_order_release); ASSERT_TRUE(orch.on_consumed(first.task_slot)); + bool prepared_consumed_early = false; + const std::future_status third_status = third.wait_for(std::chrono::seconds(1)); + EXPECT_EQ(third_status, std::future_status::ready); + if (third_status != std::future_status::ready) { + // Returning the remaining lease guarantees the async begin_run is + // joinable even when this admission invariant fails. + S(prepared.task_slot).state.store(TaskState::COMPLETED, std::memory_order_release); + EXPECT_TRUE(orch.on_consumed(prepared.task_slot)); + prepared_consumed_early = true; + } 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); + if (!prepared_consumed_early) { + 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)); + 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)); @@ -793,11 +805,21 @@ TEST_F(OrchestratorFixture, FailedPreparedConstructionReturnsAdmissionWithoutDis auto replacement = std::async(std::launch::async, [this] { return orch.begin_run(); }); + bool active_consumed_early = false; + const std::future_status replacement_status = replacement.wait_for(std::chrono::seconds(1)); + EXPECT_EQ(replacement_status, std::future_status::ready); + if (replacement_status != std::future_status::ready) { + S(active.task_slot).state.store(TaskState::COMPLETED, std::memory_order_release); + EXPECT_TRUE(orch.on_consumed(active.task_slot)); + active_consumed_early = true; + } 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)); + if (!active_consumed_early) { + 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); diff --git a/tests/ut/cpp/hierarchical/test_scheduler.cpp b/tests/ut/cpp/hierarchical/test_scheduler.cpp index 61a154658..80bbea093 100644 --- a/tests/ut/cpp/hierarchical/test_scheduler.cpp +++ b/tests/ut/cpp/hierarchical/test_scheduler.cpp @@ -247,6 +247,47 @@ class FakeEndpoint final : public WorkerEndpoint { std::atomic *prepare_count_{nullptr}; }; +class BlockingTwoFrameEndpoint final : public WorkerEndpoint { +public: + BlockingTwoFrameEndpoint() { + caps_.worker_id = 0; + caps_.max_inflight_tasks = 2; + } + + const WorkerEndpointCaps &caps() const override { return caps_; } + + WorkerCompletion run(Ring *, const WorkerDispatch &dispatch) override { + int current = active_.fetch_add(1, std::memory_order_acq_rel) + 1; + int observed = max_active_.load(std::memory_order_acquire); + while (current > observed && !max_active_.compare_exchange_weak(observed, current, std::memory_order_acq_rel)) { + } + entered_.fetch_add(1, std::memory_order_release); + std::unique_lock release_lk(release_mu_); + release_cv_.wait(release_lk, [this] { + return release_; + }); + active_.fetch_sub(1, std::memory_order_acq_rel); + return WorkerCompletion{dispatch.task_slot, dispatch.group_index, EndpointOutcome::SUCCESS, {}}; + } + + int entered() const { return entered_.load(std::memory_order_acquire); } + int max_active() const { return max_active_.load(std::memory_order_acquire); } + void release() { + std::lock_guard lk(release_mu_); + release_ = true; + release_cv_.notify_all(); + } + +private: + WorkerEndpointCaps caps_; + std::atomic active_{0}; + std::atomic max_active_{0}; + std::atomic entered_{0}; + std::mutex release_mu_; + std::condition_variable release_cv_; + bool release_{false}; +}; + // --------------------------------------------------------------------------- // Helper: build a TaskArgs whose only tensor has the given (data, tag). // --------------------------------------------------------------------------- @@ -378,6 +419,86 @@ TEST(WorkerManagerTest, StartRejectsDuplicateNextLevelWorkerId) { EXPECT_TRUE(threw); } +TEST(WorkerManagerTest, WorkerThreadUsesBoundedEndpointCapacity) { + Ring allocator; + allocator.init(/*heap_bytes=*/0); + WorkerThread worker; + auto endpoint = std::make_unique(); + BlockingTwoFrameEndpoint *endpoint_ptr = endpoint.get(); + std::atomic completions{0}; + worker.start( + &allocator, + [&](WorkerCompletion) { + completions.fetch_add(1, std::memory_order_release); + }, + {}, std::move(endpoint) + ); + + worker.dispatch(WorkerDispatch{1, 0}); + worker.dispatch(WorkerDispatch{2, 0}); + auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(3); + while (endpoint_ptr->entered() != 2 && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_EQ(endpoint_ptr->entered(), 2); + EXPECT_EQ(endpoint_ptr->max_active(), 2); + EXPECT_FALSE(worker.has_capacity()); + EXPECT_FALSE(worker.idle()); + EXPECT_THROW(worker.dispatch(WorkerDispatch{3, 0}), std::logic_error); + + endpoint_ptr->release(); + deadline = std::chrono::steady_clock::now() + std::chrono::seconds(3); + while (completions.load(std::memory_order_acquire) != 2 && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_EQ(completions.load(std::memory_order_acquire), 2); + EXPECT_TRUE(worker.has_capacity()); + EXPECT_TRUE(worker.idle()); + worker.stop(); + allocator.shutdown(); +} + +TEST(WorkerManagerTest, TwoFramePrePublishExceptionRetiresDispatchSequence) { + alignas(8) std::array mailbox{}; + Ring allocator; + allocator.init(/*heap_bytes=*/0); + AllocResult ar = allocator.alloc(/*heap_bytes=*/0, /*depth=*/0); + ASSERT_NE(ar.slot, INVALID_SLOT); + TaskSlotState *slot = allocator.slot_state(ar.slot); + ASSERT_NE(slot, nullptr); + slot->reset(); + slot->remote_sidecar.inline_payload.push_back(1); + + LocalMailboxEndpoint endpoint( + /*worker_id=*/0, mailbox.data(), /*child_pid=*/-1, /*task_frame_count=*/2 + ); + EXPECT_THROW( + (void) + endpoint.run(&allocator, WorkerDispatch{/*task_slot=*/INVALID_SLOT, /*group_index=*/0, /*dispatch_id=*/1}), + std::out_of_range + ); + + auto second = std::async(std::launch::async, [&] { + return endpoint.run(&allocator, WorkerDispatch{ar.slot, /*group_index=*/0, /*dispatch_id=*/2}); + }); + const std::future_status second_status = second.wait_for(std::chrono::seconds(1)); + EXPECT_EQ(second_status, std::future_status::ready) + << "dispatch 2 remained blocked behind dispatch 1's pre-publish exception"; + if (second_status != std::future_status::ready) { + AllocResult recovery = allocator.alloc(/*heap_bytes=*/0, /*depth=*/0); + ASSERT_NE(recovery.slot, INVALID_SLOT); + allocator.slot_state(recovery.slot)->reset(); + int32_t active = static_cast(MailboxState::TASK_ACTIVE); + std::memcpy(mailbox.data() + MAILBOX_FRAME_SIZE + MAILBOX_OFF_STATE, &active, sizeof(active)); + (void)endpoint.run(&allocator, WorkerDispatch{recovery.slot, /*group_index=*/0, /*dispatch_id=*/3}); + } + EXPECT_EQ(second.wait_for(std::chrono::seconds(1)), std::future_status::ready); + if (second.wait_for(std::chrono::seconds(0)) == std::future_status::ready) { + EXPECT_EQ(second.get().outcome, EndpointOutcome::ENDPOINT_FAILURE); + } + allocator.shutdown(); +} + TEST(WorkerManagerTest, LocalMailboxPublishesAcceptanceBeforeCompletion) { MockMailboxWorker child; child.start(); @@ -412,7 +533,9 @@ TEST(WorkerManagerTest, LocalMailboxPublishesAcceptanceBeforeCompletion) { 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) {} + while (!accepted.load(std::memory_order_acquire) && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } EXPECT_TRUE(accepted.load(std::memory_order_acquire)); EXPECT_EQ(done.wait_for(std::chrono::milliseconds(0)), std::future_status::timeout); diff --git a/tests/ut/py/test_worker/test_host_worker.py b/tests/ut/py/test_worker/test_host_worker.py index a3cf24399..d208f2c30 100644 --- a/tests/ut/py/test_worker/test_host_worker.py +++ b/tests/ut/py/test_worker/test_host_worker.py @@ -131,6 +131,7 @@ 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] = [] + published_frame_counts: list[int] = [] class FakeChipWorker: pipeline_depth = 2 @@ -141,8 +142,9 @@ def init(self, device_id, bins, *, log_level, prewarm_config=None, enable_sdma=F def finalize(self) -> None: events.append(("finalize",)) - def fake_run_chip_main_loop(cw, *_args, chip_platform, chip_runtime, prepared=None): + def fake_run_chip_main_loop(cw, *_args, chip_platform, chip_runtime, prepared=None, task_frame_count=1): published_depths.append(worker_mod._PIPELINE_LEASE_FMT.unpack_from(_args[0], worker_mod._OFF_PIPELINE_LEASE)[0]) + published_frame_counts.append(task_frame_count) events.append(("main_loop", cw, chip_platform, chip_runtime)) monkeypatch.setattr(worker_mod, "ChipWorker", FakeChipWorker) @@ -170,6 +172,117 @@ def fake_run_chip_main_loop(cw, *_args, chip_platform, chip_runtime, prepared=No assert events[1][2:] == ("a2a3", "tensormap_and_ringbuffer") assert events[2] == ("finalize",) assert published_depths == [2] + assert published_frame_counts == [1] + + +@pytest.mark.parametrize( + ("platform", "runtime", "depth", "expected"), + [ + ("a2a3", "host_build_graph", 2, 2), + ("a2a3", "host_build_graph", 1, 1), + ("a2a3", "tensormap_and_ringbuffer", 2, 1), + ("a5", "host_build_graph", 2, 1), + ], +) +def test_local_task_frame_count_requires_supported_backend(platform, runtime, depth, expected): + assert worker_mod._local_task_frame_count(platform, runtime, depth) == expected + + +def test_two_frame_chip_loop_accepts_b_while_a_executes_sequentially(): + entered = [threading.Event(), threading.Event()] + release = [threading.Event(), threading.Event()] + active = 0 + max_active = 0 + calls = 0 + call_lock = threading.Lock() + + class FakeImpl: + def run_from_blob(self, *_args): + nonlocal active, max_active, calls + with call_lock: + index = calls + calls += 1 + active += 1 + max_active = max(max_active, active) + entered[index].set() + assert release[index].wait(5.0) + with call_lock: + active -= 1 + + class FakeChipWorker: + pipeline_depth = 2 + _impl = FakeImpl() + + shm = SharedMemory(create=True, size=MAILBOX_SIZE) + loop_thread = None + frame_bufs = [] + try: + buf = shm.buf + assert buf is not None + mailbox_addr = _mailbox_addr(shm) + digest = bytes([0x42]) * worker_mod.CALLABLE_HASH_DIGEST_BYTES + loop_thread = threading.Thread( + target=worker_mod._run_chip_main_loop, + args=(FakeChipWorker(), buf, mailbox_addr, mailbox_addr + _OFF_STATE, 0, {7: object()}, {digest: 7}, {}), + kwargs={ + "chip_platform": "a2a3", + "prepared": {7}, + "task_frame_count": 2, + }, + ) + loop_thread.start() + + for index in range(2): + frame = buf[(1 + index) * worker_mod.MAILBOX_FRAME_SIZE : (2 + index) * worker_mod.MAILBOX_FRAME_SIZE] + frame_bufs.append(frame) + frame[worker_mod._OFF_TASK_CALLABLE_HASH : worker_mod._OFF_TASK_ARGS_BLOB] = digest + struct.pack_into("=ii", frame, worker_mod._OFF_TASK_ARGS_BLOB, 0, 0) + cfg_values = [0] * (6 + 3 * worker_mod.RUNTIME_ENV_RING_COUNT) + worker_mod._CFG_FMT.pack_into(frame, worker_mod._OFF_CONFIG, *cfg_values, b"") + worker_mod._PIPELINE_LEASE_FMT.pack_into(frame, worker_mod._OFF_PIPELINE_LEASE, index, 0, 11) + struct.pack_into("=Q", frame, worker_mod._OFF_FRAME_PROTOCOL, worker_mod._TASK_PROTOCOL_VERSION) + struct.pack_into("=Q", frame, worker_mod._OFF_FRAME_RUN_ID, 5) + struct.pack_into("=Q", frame, worker_mod._OFF_FRAME_SLOT_ID, index) + struct.pack_into("=Q", frame, worker_mod._OFF_FRAME_GENERATION, 11) + struct.pack_into("=Q", frame, worker_mod._OFF_FRAME_DISPATCH_ID, index + 1) + + frame0_addr = mailbox_addr + worker_mod.MAILBOX_FRAME_SIZE + frame1_addr = mailbox_addr + 2 * worker_mod.MAILBOX_FRAME_SIZE + _mailbox_store_i32(frame0_addr + _OFF_STATE, worker_mod._TASK_READY) + assert entered[0].wait(5.0) + assert _mailbox_load_i32(frame0_addr + _OFF_STATE) == worker_mod._TASK_ACTIVE + + _mailbox_store_i32(frame1_addr + _OFF_STATE, worker_mod._TASK_READY) + deadline = time.monotonic() + 5.0 + while _mailbox_load_i32(frame1_addr + _OFF_STATE) != worker_mod._TASK_ACCEPTED_STATE: + assert time.monotonic() < deadline + time.sleep(0.001) + assert calls == 1 + + release[0].set() + assert entered[1].wait(5.0) + assert _mailbox_load_i32(frame0_addr + _OFF_STATE) == worker_mod._TASK_DONE + assert _mailbox_load_i32(frame1_addr + _OFF_STATE) == worker_mod._TASK_ACTIVE + release[1].set() + deadline = time.monotonic() + 5.0 + while _mailbox_load_i32(frame1_addr + _OFF_STATE) != worker_mod._TASK_DONE: + assert time.monotonic() < deadline + time.sleep(0.001) + assert max_active == 1 + + _mailbox_store_i32(mailbox_addr + _OFF_STATE, worker_mod._SHUTDOWN) + loop_thread.join(5.0) + assert not loop_thread.is_alive() + finally: + for event in release: + event.set() + if loop_thread is not None and loop_thread.is_alive(): + _mailbox_store_i32(_mailbox_addr(shm) + _OFF_STATE, worker_mod._SHUTDOWN) + loop_thread.join(5.0) + for frame in frame_bufs: + frame.release() + shm.close() + shm.unlink() def _chip_digest(callable_obj: ChipCallable, *, platform: str = "", runtime: str = "") -> bytes: From a264f8e530c147a67023a82cd67675cf844d3359 Mon Sep 17 00:00:00 2001 From: Crane-Liu Date: Wed, 29 Jul 2026 17:46:48 +0800 Subject: [PATCH 3/3] Add: generation-safe prepare activation contract --- docs/task-flow.md | 28 +++ python/bindings/worker_bind.h | 15 +- python/simpler/worker.py | 5 +- src/common/hierarchical/orchestrator.cpp | 13 ++ src/common/hierarchical/orchestrator.h | 1 + src/common/hierarchical/scheduler.cpp | 67 +++++- src/common/hierarchical/scheduler.h | 3 + src/common/hierarchical/worker.cpp | 18 +- src/common/hierarchical/worker.h | 10 +- src/common/hierarchical/worker_manager.cpp | 189 ++++++++++++++++- src/common/hierarchical/worker_manager.h | 53 ++++- tests/ut/cpp/hierarchical/test_scheduler.cpp | 205 +++++++++++++++++++ 12 files changed, 578 insertions(+), 29 deletions(-) diff --git a/docs/task-flow.md b/docs/task-flow.md index 66cfd8024..b79006fa6 100644 --- a/docs/task-flow.md +++ b/docs/task-flow.md @@ -326,6 +326,34 @@ 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. +#### Prepare/activate endpoint lane + +A local two-frame endpoint may separately advertise +`supports_prepare_activate`. The scheduler then stages at most one ready +NEXT_LEVEL task slot from the prepared FIFO successor. A group is staged as one +task slot across all of its target workers; an endpoint set that cannot stage +the whole group leaves it on the normal ready queue. + +The prepared frame uses a distinct protocol path: + +```text +PREPARE_READY -> BACKEND_READY -> ACTIVATE -> TASK_ACTIVE -> TASK_DONE | TASK_FAILED +``` + +`BACKEND_READY` reports only that unpublished backend state exists. It does not +satisfy the launch-acceptance fence and does not permit device work. The +WorkerThread waits on an activation condition variable while retaining the +frame and its generation-safe identity. Whole-run FIFO promotion latches +`ACTIVATE`; the latch may arrive before backend preparation finishes, but the +endpoint cannot publish it to the child until `BACKEND_READY` is observed for +the same run, slot, generation, and dispatch id. + +Only the first eligible dispatch in a run uses this lane. Other tasks from the +same run retain the existing sequential compatibility path. Remote, SUB, A5, +single-frame, and endpoints without the capability remain unchanged. The +common contract is dormant until a backend explicitly advertises support; the +HBG and TMR prepared-state implementations provide that support separately. + 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 buffer per slot, and its single-entry prebuilt-arena cache stays owned by diff --git a/python/bindings/worker_bind.h b/python/bindings/worker_bind.h index 0610a0ddf..58756f001 100644 --- a/python/bindings/worker_bind.h +++ b/python/bindings/worker_bind.h @@ -393,12 +393,15 @@ inline void bind_worker(nb::module_ &m) { .def( "add_next_level_worker", - [](Worker &self, uint64_t mailbox_ptr, int child_pid, uint32_t task_frame_count) { + [](Worker &self, uint64_t mailbox_ptr, int child_pid, uint32_t task_frame_count, + bool supports_prepare_activate) { self.add_worker( - WorkerType::NEXT_LEVEL, reinterpret_cast(mailbox_ptr), child_pid, task_frame_count + WorkerType::NEXT_LEVEL, reinterpret_cast(mailbox_ptr), child_pid, task_frame_count, + supports_prepare_activate ); }, nb::arg("mailbox_ptr"), nb::arg("child_pid") = -1, nb::arg("task_frame_count") = 1, + nb::arg("supports_prepare_activate") = false, "Add a NEXT_LEVEL sub-worker. `mailbox_ptr` is the address of a " "MAILBOX_SIZE-byte MAP_SHARED region; the child process loop is " "Python-managed (fork + _chip_process_loop). `child_pid` is that " @@ -406,13 +409,15 @@ inline void bind_worker(nb::module_ &m) { ) .def( "add_next_level_worker_at", - [](Worker &self, int32_t worker_id, uint64_t mailbox_ptr, int child_pid, uint32_t task_frame_count) { + [](Worker &self, int32_t worker_id, uint64_t mailbox_ptr, int child_pid, uint32_t task_frame_count, + bool supports_prepare_activate) { self.add_next_level_worker( - worker_id, reinterpret_cast(mailbox_ptr), child_pid, task_frame_count + worker_id, reinterpret_cast(mailbox_ptr), child_pid, task_frame_count, + supports_prepare_activate ); }, nb::arg("worker_id"), nb::arg("mailbox_ptr"), nb::arg("child_pid") = -1, nb::arg("task_frame_count") = 1, - "Add a NEXT_LEVEL sub-worker with an explicit worker id." + nb::arg("supports_prepare_activate") = false, "Add a NEXT_LEVEL sub-worker with an explicit worker id." ) .def( "configure_pipeline_depth", &Worker::configure_pipeline_depth, nb::arg("depth"), diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 15483ae4a..207050e36 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -202,7 +202,7 @@ def my_l4_orch(orch, args, config): _OFF_FRAME_SLOT_ID = _OFF_ACCEPTED - 24 _OFF_FRAME_GENERATION = _OFF_ACCEPTED - 16 _OFF_FRAME_DISPATCH_ID = _OFF_ACCEPTED - 8 -_TASK_PROTOCOL_VERSION = 1 +_TASK_PROTOCOL_VERSION = 2 _MAILBOX_ARGS_CAPACITY = _OFF_FRAME_PROTOCOL - _OFF_TASK_ARGS_BLOB _OFF_CONTROL_CALLABLE_HASH = _OFF_ARGS + 32 # MAILBOX_OFF_ERROR_MSG / MAILBOX_ERROR_MSG_SIZE come from the C++ @@ -226,6 +226,9 @@ def my_l4_orch(orch, args, config): _TASK_ACCEPTED_STATE = 8 _TASK_ACTIVE = 9 _TASK_FAILED = 10 +_BACKEND_READY = 11 +_ACTIVATE = 12 +_PREPARE_READY = 13 _TASK_FRAME_COUNT = 2 diff --git a/src/common/hierarchical/orchestrator.cpp b/src/common/hierarchical/orchestrator.cpp index 04b58554e..38d8cb824 100644 --- a/src/common/hierarchical/orchestrator.cpp +++ b/src/common/hierarchical/orchestrator.cpp @@ -184,6 +184,7 @@ void Orchestrator::close_run_submission(RunId run_id) { } } run->completion_cv.notify_all(); + if (ready_notify_cb_) ready_notify_cb_(); activate_fifo_head(); finish_run_if_ready(run); } @@ -288,6 +289,18 @@ RunId Orchestrator::active_run_id() const { return active_run_id_; } +RunId Orchestrator::preparable_run_id() const { + std::lock_guard lk(runs_mu_); + if (active_run_id_ == INVALID_RUN_ID || run_fifo_.size() < 2 || run_fifo_.front() != active_run_id_) { + return INVALID_RUN_ID; + } + auto it = runs_.find(run_fifo_[1]); + if (it == runs_.end() || it->second->phase.load(std::memory_order_acquire) != RunPhase::PREPARED) { + return INVALID_RUN_ID; + } + return it->first; +} + bool Orchestrator::quiescent_locked() const { return runs_.empty() && building_run_id_ == INVALID_RUN_ID; } void Orchestrator::compact_if_quiescent() { diff --git a/src/common/hierarchical/orchestrator.h b/src/common/hierarchical/orchestrator.h index 591e27e6e..a817a88b5 100644 --- a/src/common/hierarchical/orchestrator.h +++ b/src/common/hierarchical/orchestrator.h @@ -131,6 +131,7 @@ class Orchestrator { bool run_failed(RunId run_id) const; bool can_dispatch_run(RunId run_id) const; RunId active_run_id() const; + RunId preparable_run_id() const; void release_run(RunId run_id); // Open a nested scope. Every task submitted between this call and the diff --git a/src/common/hierarchical/scheduler.cpp b/src/common/hierarchical/scheduler.cpp index 3c4cb7351..fa7681332 100644 --- a/src/common/hierarchical/scheduler.cpp +++ b/src/common/hierarchical/scheduler.cpp @@ -180,7 +180,13 @@ void Scheduler::run() { 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)); + (!cfg_.ready_next_level_queues->empty(active) || !cfg_.ready_sub_queue->empty(active) || + cfg_.manager->needs_activation(active)); + if (!ready && cfg_.preparable_run_cb) { + RunId preparable = cfg_.preparable_run_cb(); + ready = preparable != INVALID_RUN_ID && !cfg_.manager->has_staged_run(preparable) && + !cfg_.ready_next_level_queues->empty(preparable); + } } else { ready = !cfg_.ready_next_level_queues->empty() || !cfg_.ready_sub_queue->empty(); } @@ -357,11 +363,70 @@ 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() { + RunId active_run = cfg_.active_run_cb ? cfg_.active_run_cb() : INVALID_RUN_ID; + if (active_run != INVALID_RUN_ID) cfg_.manager->activate_prepared_run(active_run); + dispatch_preparable_next_level_group(); + dispatch_preparable_next_level_singles(); dispatch_next_level_group(); dispatch_next_level_singles(); dispatch_sub_ready(); } +void Scheduler::dispatch_preparable_next_level_group() { + if (!cfg_.preparable_run_cb) return; + RunId run_id = cfg_.preparable_run_cb(); + if (run_id == INVALID_RUN_ID || cfg_.manager->has_staged_run(run_id)) return; + + TaskSlot slot; + if (!cfg_.ready_next_level_queues->try_front_group(run_id, slot)) return; + TaskSlotState &state = *cfg_.ring->slot_state(slot); + if (state.state.load(std::memory_order_acquire) != TaskState::READY || state.run_id != run_id || + state.worker_type != WorkerType::NEXT_LEVEL || !state.is_group()) { + return; + } + + const int32_t group_size = state.group_size(); + std::vector workers; + workers.reserve(static_cast(group_size)); + for (int32_t index = 0; index < group_size; ++index) { + WorkerThread *worker = cfg_.manager->get_worker_by_id(WorkerType::NEXT_LEVEL, state.target_worker_id(index)); + if (worker == nullptr || !worker->caps().supports_prepare_activate || !worker->has_capacity()) return; + workers.push_back(worker); + } + + TaskSlot popped; + if (!cfg_.ready_next_level_queues->try_pop_group(run_id, popped) || popped != slot) { + throw std::runtime_error("Scheduler::dispatch_preparable_next_level_group: group queue changed unexpectedly"); + } + reset_group_state(state, group_size, GroupMemberState::RUNNING); + state.state.store(TaskState::RUNNING, std::memory_order_release); + for (int32_t index = 0; index < group_size; ++index) { + workers[static_cast(index)]->dispatch_prepared(WorkerDispatch{slot, index}); + } +} + +void Scheduler::dispatch_preparable_next_level_singles() { + if (!cfg_.preparable_run_cb) return; + RunId run_id = cfg_.preparable_run_cb(); + if (run_id == INVALID_RUN_ID || cfg_.manager->has_staged_run(run_id)) return; + + for (int32_t worker_id : cfg_.ready_next_level_queues->worker_ids()) { + WorkerThread *worker = cfg_.manager->get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); + if (worker == nullptr || !worker->caps().supports_prepare_activate || !worker->has_capacity()) continue; + TaskSlot slot; + if (!cfg_.ready_next_level_queues->try_pop_single(worker_id, run_id, slot)) continue; + TaskSlotState &state = *cfg_.ring->slot_state(slot); + if (state.state.load(std::memory_order_acquire) != TaskState::READY || state.run_id != run_id || + state.worker_type != WorkerType::NEXT_LEVEL || state.is_group() || state.target_worker_id(0) != worker_id) { + cfg_.enqueue_ready_cb(slot); + continue; + } + state.state.store(TaskState::RUNNING, std::memory_order_release); + worker->dispatch_prepared(WorkerDispatch{slot, 0}); + return; + } +} + void Scheduler::dispatch_sub_ready() { RunId active_run = cfg_.active_run_cb ? cfg_.active_run_cb() : INVALID_RUN_ID; if (cfg_.active_run_cb && active_run == INVALID_RUN_ID) return; diff --git a/src/common/hierarchical/scheduler.h b/src/common/hierarchical/scheduler.h index 95be5dc77..3e4e0085a 100644 --- a/src/common/hierarchical/scheduler.h +++ b/src/common/hierarchical/scheduler.h @@ -62,6 +62,7 @@ class Scheduler { // 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; + std::function preparable_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 @@ -106,6 +107,8 @@ class Scheduler { void poison_task(TaskSlot slot, const std::string &root_message); void try_consume(TaskSlot slot); void dispatch_ready(); + void dispatch_preparable_next_level_group(); + void dispatch_preparable_next_level_singles(); void dispatch_next_level_group(); void dispatch_next_level_singles(); void dispatch_sub_ready(); diff --git a/src/common/hierarchical/worker.cpp b/src/common/hierarchical/worker.cpp index 0d9e9aa3f..79cb885af 100644 --- a/src/common/hierarchical/worker.cpp +++ b/src/common/hierarchical/worker.cpp @@ -68,15 +68,20 @@ Worker::~Worker() { if (initialized_) close(); } -void Worker::add_worker(WorkerType type, void *mailbox, int child_pid, uint32_t task_frame_count) { +void Worker::add_worker( + WorkerType type, void *mailbox, int child_pid, uint32_t task_frame_count, bool supports_prepare_activate +) { if (initialized_) throw std::runtime_error("Worker: add_worker after init"); - if (type == WorkerType::NEXT_LEVEL) manager_.add_next_level(mailbox, child_pid, task_frame_count); - else manager_.add_sub(mailbox, child_pid); + if (type == WorkerType::NEXT_LEVEL) { + manager_.add_next_level(mailbox, child_pid, task_frame_count, supports_prepare_activate); + } else manager_.add_sub(mailbox, child_pid); } -void Worker::add_next_level_worker(int32_t worker_id, void *mailbox, int child_pid, uint32_t task_frame_count) { +void Worker::add_next_level_worker( + int32_t worker_id, void *mailbox, int child_pid, uint32_t task_frame_count, bool supports_prepare_activate +) { if (initialized_) throw std::runtime_error("Worker: add_next_level_worker after init"); - manager_.add_next_level_at(worker_id, mailbox, child_pid, task_frame_count); + manager_.add_next_level_at(worker_id, mailbox, child_pid, task_frame_count, supports_prepare_activate); } void Worker::add_remote_l3_socket( @@ -125,6 +130,9 @@ void Worker::init() { cfg.active_run_cb = [this] { return orchestrator_.active_run_id(); }; + cfg.preparable_run_cb = [this] { + return orchestrator_.preparable_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 af4a5b3b8..289f4e923 100644 --- a/src/common/hierarchical/worker.h +++ b/src/common/hierarchical/worker.h @@ -76,8 +76,14 @@ class Worker { // child and consumes the mailbox via the Python child loop. // `child_pid` is the forked child servicing `mailbox`, or -1 when the // caller owns no waitable child. - void add_worker(WorkerType type, void *mailbox, int child_pid = -1, uint32_t task_frame_count = 1); - void add_next_level_worker(int32_t worker_id, void *mailbox, int child_pid = -1, uint32_t task_frame_count = 1); + void add_worker( + WorkerType type, void *mailbox, int child_pid = -1, uint32_t task_frame_count = 1, + bool supports_prepare_activate = false + ); + void add_next_level_worker( + int32_t worker_id, void *mailbox, int child_pid = -1, uint32_t task_frame_count = 1, + bool supports_prepare_activate = false + ); // Register a REMOTE_L3 endpoint only after its session runner completed // prestart and reported HELLO READY on the command lane. diff --git a/src/common/hierarchical/worker_manager.cpp b/src/common/hierarchical/worker_manager.cpp index 970927dd4..d05f42ffc 100644 --- a/src/common/hierarchical/worker_manager.cpp +++ b/src/common/hierarchical/worker_manager.cpp @@ -146,11 +146,20 @@ WorkerEndpoint::run_with_accept(Ring *ring, const WorkerDispatch &dispatch, cons return completion; } +WorkerCompletion WorkerEndpoint::run_prepared_with_activation( + Ring *, const WorkerDispatch &, const std::function &, const std::function &, + const std::function & +) { + throw std::runtime_error("prepared activation is not supported by this WorkerEndpoint"); +} + // ============================================================================= // LocalMailboxEndpoint — mailbox helpers // ============================================================================= -LocalMailboxEndpoint::LocalMailboxEndpoint(int32_t worker_id, void *mailbox, int child_pid, uint32_t task_frame_count) : +LocalMailboxEndpoint::LocalMailboxEndpoint( + int32_t worker_id, void *mailbox, int child_pid, uint32_t task_frame_count, bool supports_prepare_activate +) : mailbox_(mailbox), task_frame_count_(task_frame_count), child_pid_(child_pid) { @@ -160,6 +169,7 @@ LocalMailboxEndpoint::LocalMailboxEndpoint(int32_t worker_id, void *mailbox, int } caps_.worker_id = worker_id; caps_.max_inflight_tasks = task_frame_count; + caps_.supports_prepare_activate = supports_prepare_activate; } std::string LocalMailboxEndpoint::check_child_death() { @@ -266,6 +276,33 @@ void WorkerThread::start( } void WorkerThread::dispatch(WorkerDispatch d) { + d.prepare_only = false; + enqueue_dispatch(d); +} + +void WorkerThread::dispatch_prepared(WorkerDispatch d) { + if (!caps().supports_prepare_activate) { + throw std::logic_error("WorkerThread::dispatch_prepared: endpoint does not support prepare/activate"); + } + if (ring_ == nullptr) throw std::logic_error("WorkerThread::dispatch_prepared: null ring"); + TaskSlotState *slot = ring_->slot_state(d.task_slot); + if (slot == nullptr || slot->run_id == INVALID_RUN_ID) { + throw std::logic_error("WorkerThread::dispatch_prepared: dispatch has no run identity"); + } + RunId expected = INVALID_RUN_ID; + if (!staged_run_id_.compare_exchange_strong(expected, slot->run_id, std::memory_order_acq_rel)) { + throw std::logic_error("WorkerThread::dispatch_prepared: worker already owns a staged run"); + } + d.prepare_only = true; + try { + enqueue_dispatch(d); + } catch (...) { + staged_run_id_.store(INVALID_RUN_ID, std::memory_order_release); + throw; + } +} + +void WorkerThread::enqueue_dispatch(WorkerDispatch d) { uint32_t previous = inflight_.fetch_add(1, std::memory_order_acq_rel); if (previous >= capacity_) { inflight_.fetch_sub(1, std::memory_order_acq_rel); @@ -277,12 +314,24 @@ void WorkerThread::dispatch(WorkerDispatch d) { cv_.notify_one(); } +bool WorkerThread::activate_prepared(RunId run_id) { + if (run_id == INVALID_RUN_ID || !needs_activation(run_id)) return false; + { + std::lock_guard lk(activation_mu_); + activation_requested_run_id_ = run_id; + activated_run_id_.store(run_id, std::memory_order_release); + } + activation_cv_.notify_all(); + return true; +} + void WorkerThread::stop() { { std::lock_guard lk(mu_); shutdown_ = true; } cv_.notify_all(); + activation_cv_.notify_all(); for (auto &thread : threads_) { if (thread.joinable()) thread.join(); } @@ -310,7 +359,7 @@ void WorkerThread::loop() { { std::unique_lock lk(mu_); cv_.wait(lk, [this] { - return !queue_.empty() || shutdown_; + return !queue_.empty() || shutdown_.load(std::memory_order_acquire); }); if (queue_.empty()) break; d = queue_.front(); @@ -363,6 +412,44 @@ void WorkerThread::loop() { WorkerCompletion WorkerThread::dispatch_process(WorkerDispatch d, const std::function &on_accept) { if (!endpoint_) throw std::runtime_error("WorkerThread::dispatch_process: null endpoint"); + if (d.prepare_only) { + TaskSlotState *slot = ring_ == nullptr ? nullptr : ring_->slot_state(d.task_slot); + if (slot == nullptr) throw std::runtime_error("WorkerThread::dispatch_process: invalid prepared task slot"); + const RunId run_id = slot->run_id; + auto clear_staged = [&]() { + { + std::lock_guard lk(activation_mu_); + if (activation_requested_run_id_ == run_id) activation_requested_run_id_ = INVALID_RUN_ID; + } + backend_ready_run_id_.store(INVALID_RUN_ID, std::memory_order_release); + activated_run_id_.store(INVALID_RUN_ID, std::memory_order_release); + staged_run_id_.store(INVALID_RUN_ID, std::memory_order_release); + }; + try { + WorkerCompletion completion = endpoint_->run_prepared_with_activation( + ring_, d, + [this, run_id]() { + backend_ready_run_id_.store(run_id, std::memory_order_release); + }, + [this, run_id]() { + std::unique_lock lk(activation_mu_); + activation_cv_.wait(lk, [this, run_id] { + return activation_requested_run_id_ == run_id || shutdown_.load(std::memory_order_acquire); + }); + if (shutdown_.load(std::memory_order_acquire)) { + throw std::runtime_error("prepared activation interrupted by worker shutdown"); + } + activation_requested_run_id_ = INVALID_RUN_ID; + }, + on_accept + ); + clear_staged(); + return completion; + } catch (...) { + clear_staged(); + throw; + } + } return endpoint_->run_with_accept(ring_, d, on_accept); } @@ -500,8 +587,19 @@ WorkerCompletion LocalMailboxEndpoint::run_with_accept( return completion; } +WorkerCompletion LocalMailboxEndpoint::run_prepared_with_activation( + Ring *ring, const WorkerDispatch &dispatch, const std::function &on_backend_ready, + const std::function &wait_for_activation, const std::function &on_accept +) { + if (!caps_.supports_prepare_activate || task_frame_count_ < 2) { + throw std::runtime_error("LocalMailboxEndpoint: prepared activation requires an enabled two-frame endpoint"); + } + return run_two_frame(ring, dispatch, on_accept, on_backend_ready, wait_for_activation); +} + WorkerCompletion LocalMailboxEndpoint::run_two_frame( - Ring *ring, const WorkerDispatch &dispatch, const std::function &on_accept + Ring *ring, const WorkerDispatch &dispatch, const std::function &on_accept, + const std::function &on_backend_ready, const std::function &wait_for_activation ) { WorkerCompletion completion; completion.task_slot = dispatch.task_slot; @@ -646,7 +744,7 @@ WorkerCompletion LocalMailboxEndpoint::run_two_frame( release_claim(); return completion; } - write_mailbox_state(MailboxState::TASK_READY, frame); + write_mailbox_state(wait_for_activation ? MailboxState::PREPARE_READY : MailboxState::TASK_READY, frame); ++next_publish_id_; publish_sequence_retired = true; task_published = true; @@ -671,6 +769,51 @@ WorkerCompletion LocalMailboxEndpoint::run_two_frame( auto next_liveness_check = std::chrono::steady_clock::now() + kChildLivenessPollPeriod; bool acceptance_observed = false; MailboxState terminal_state = MailboxState::IDLE; + if (wait_for_activation) { + while (true) { + MailboxState state = read_mailbox_state(frame); + if (state == MailboxState::BACKEND_READY || state == MailboxState::TASK_FAILED) { + terminal_state = state; + break; + } + auto now = std::chrono::steady_clock::now(); + if (now >= next_liveness_check) { + next_liveness_check = now + kChildLivenessPollPeriod; + std::string death = check_child_death(); + if (!death.empty()) { + endpoint_poisoned_.store(true, std::memory_order_release); + publish_cv_.notify_all(); + completion.outcome = EndpointOutcome::ENDPOINT_FAILURE; + completion.error_message = "LocalMailboxEndpoint::run_two_frame: " + death; + release_claim(); + return completion; + } + } + } + if (!identity_matches()) { + endpoint_poisoned_.store(true, std::memory_order_release); + publish_cv_.notify_all(); + completion.outcome = EndpointOutcome::ENDPOINT_FAILURE; + completion.error_message = "LocalMailboxEndpoint::run_two_frame: stale backend-ready identity"; + release_claim(); + return completion; + } + if (terminal_state == MailboxState::BACKEND_READY) { + if (on_backend_ready) on_backend_ready(); + wait_for_activation(); + if (!identity_matches() || read_mailbox_state(frame) != MailboxState::BACKEND_READY) { + endpoint_poisoned_.store(true, std::memory_order_release); + publish_cv_.notify_all(); + completion.outcome = EndpointOutcome::ENDPOINT_FAILURE; + completion.error_message = + "LocalMailboxEndpoint::run_two_frame: prepared frame changed before activation"; + release_claim(); + return completion; + } + write_mailbox_state(MailboxState::ACTIVATE, frame); + terminal_state = MailboxState::IDLE; + } + } while (true) { MailboxState state = read_mailbox_state(frame); if (!acceptance_observed && (state == MailboxState::TASK_ACCEPTED || state == MailboxState::TASK_ACTIVE || @@ -748,13 +891,22 @@ WorkerCompletion LocalMailboxEndpoint::run_two_frame( // WorkerManager // ============================================================================= -void WorkerManager::add_next_level(void *mailbox, int child_pid, uint32_t task_frame_count) { - add_next_level_at(static_cast(next_level_entries_.size()), mailbox, child_pid, task_frame_count); +void WorkerManager::add_next_level( + void *mailbox, int child_pid, uint32_t task_frame_count, bool supports_prepare_activate +) { + add_next_level_at( + static_cast(next_level_entries_.size()), mailbox, child_pid, task_frame_count, + supports_prepare_activate + ); } -void WorkerManager::add_next_level_at(int32_t worker_id, void *mailbox, int child_pid, uint32_t task_frame_count) { +void WorkerManager::add_next_level_at( + int32_t worker_id, void *mailbox, int child_pid, uint32_t task_frame_count, bool supports_prepare_activate +) { if (worker_id < 0) throw std::invalid_argument("WorkerManager::add_next_level_at: negative worker_id"); - next_level_entries_.push_back(LocalNextLevelEntry{worker_id, mailbox, child_pid, task_frame_count}); + next_level_entries_.push_back( + LocalNextLevelEntry{worker_id, mailbox, child_pid, task_frame_count, supports_prepare_activate} + ); } void WorkerManager::add_next_level_endpoint(std::unique_ptr endpoint) { @@ -792,7 +944,7 @@ void WorkerManager::start(Ring *ring, const OnCompleteFn &on_complete, const OnA for (const auto &entry : next_level_entries_) { auto wt = std::make_unique(); auto endpoint = std::make_unique( - entry.worker_id, entry.mailbox, entry.child_pid, entry.task_frame_count + entry.worker_id, entry.mailbox, entry.child_pid, entry.task_frame_count, entry.supports_prepare_activate ); wt->start(ring, on_complete, on_accept, std::move(endpoint)); next_level_threads_.push_back(std::move(wt)); @@ -1279,6 +1431,25 @@ bool WorkerManager::any_busy() const { return false; } +bool WorkerManager::has_staged_run(RunId run_id) const { + for (const auto &worker : next_level_threads_) + if (worker->has_staged_run(run_id)) return true; + return false; +} + +bool WorkerManager::needs_activation(RunId run_id) const { + for (const auto &worker : next_level_threads_) + if (worker->needs_activation(run_id)) return true; + return false; +} + +bool WorkerManager::activate_prepared_run(RunId run_id) { + bool activated = false; + for (const auto &worker : next_level_threads_) + activated = worker->activate_prepared(run_id) || activated; + return activated; +} + // ============================================================================= // Dynamic register/unregister broadcast (POSIX shm staging + parallel fan-out) // ============================================================================= diff --git a/src/common/hierarchical/worker_manager.h b/src/common/hierarchical/worker_manager.h index 189807da7..f4230fad5 100644 --- a/src/common/hierarchical/worker_manager.h +++ b/src/common/hierarchical/worker_manager.h @@ -76,6 +76,9 @@ enum class MailboxState : int32_t { TASK_ACCEPTED = 8, TASK_ACTIVE = 9, TASK_FAILED = 10, + BACKEND_READY = 11, + ACTIVATE = 12, + PREPARE_READY = 13, }; // Sized so the args region can hold any TaskArgs the runtime itself accepts @@ -91,7 +94,7 @@ static constexpr size_t MAILBOX_TASK_FRAME_COUNT = 2; static constexpr size_t MAILBOX_CONTROL_FRAME = 0; static constexpr size_t MAILBOX_FIRST_TASK_FRAME = 1; static constexpr size_t MAILBOX_SIZE = MAILBOX_FRAME_SIZE * (1 + MAILBOX_TASK_FRAME_COUNT); -static constexpr uint32_t MAILBOX_TASK_PROTOCOL_VERSION = 1; +static constexpr uint32_t MAILBOX_TASK_PROTOCOL_VERSION = 2; // Error message region lives at the mailbox tail. 256 B of headroom is // enough for `: ` produced by the child-side @@ -226,6 +229,7 @@ struct WorkerEndpointCaps { bool supports_task_dispatch{true}; bool supports_control{true}; uint32_t max_inflight_tasks{1}; + bool supports_prepare_activate{false}; std::string transport{"local-mailbox"}; }; @@ -239,6 +243,10 @@ class WorkerEndpoint { // conservative completion-time acceptance for remote, SUB and A5 paths. virtual WorkerCompletion run_with_accept(Ring *ring, const WorkerDispatch &dispatch, const std::function &on_accept); + virtual WorkerCompletion run_prepared_with_activation( + Ring *ring, const WorkerDispatch &dispatch, const std::function &on_backend_ready, + const std::function &wait_for_activation, const std::function &on_accept + ); virtual void shutdown_child() {} virtual uint64_t control_malloc(size_t size); @@ -290,12 +298,19 @@ class LocalMailboxEndpoint : public WorkerEndpoint { // caller does not own a waitable child. A valid pid enables liveness // checks that turn a child that dies before publishing TASK_DONE / // CONTROL_DONE into a reported failure instead of an endless spin-poll. - LocalMailboxEndpoint(int32_t worker_id, void *mailbox, int child_pid = -1, uint32_t task_frame_count = 1); + LocalMailboxEndpoint( + int32_t worker_id, void *mailbox, int child_pid = -1, uint32_t task_frame_count = 1, + bool supports_prepare_activate = false + ); const WorkerEndpointCaps &caps() const override { return caps_; } WorkerCompletion run(Ring *ring, const WorkerDispatch &dispatch) override; WorkerCompletion run_with_accept(Ring *ring, const WorkerDispatch &dispatch, const std::function &on_accept) override; + WorkerCompletion run_prepared_with_activation( + Ring *ring, const WorkerDispatch &dispatch, const std::function &on_backend_ready, + const std::function &wait_for_activation, const std::function &on_accept + ) override; void shutdown_child() override; uint64_t control_malloc(size_t size) override; @@ -368,7 +383,10 @@ class LocalMailboxEndpoint : public WorkerEndpoint { // next TASK_READY. See MAILBOX_OFF_ACCEPTED. bool read_task_accepted(const char *frame = nullptr) const; void clear_task_accepted(char *frame = nullptr); - WorkerCompletion run_two_frame(Ring *ring, const WorkerDispatch &dispatch, const std::function &on_accept); + WorkerCompletion run_two_frame( + Ring *ring, const WorkerDispatch &dispatch, const std::function &on_accept, + const std::function &on_backend_ready = {}, const std::function &wait_for_activation = {} + ); char *task_frame(size_t index) const; void run_control_command(const char *op_name, double timeout_s = -1.0); @@ -390,6 +408,7 @@ struct WorkerDispatch { TaskSlot task_slot{INVALID_SLOT}; int32_t group_index{0}; uint64_t dispatch_id{0}; + bool prepare_only{false}; }; // ============================================================================= @@ -422,6 +441,12 @@ class WorkerThread { // Enqueue a dispatch for the worker. Non-blocking. void dispatch(WorkerDispatch d); + void dispatch_prepared(WorkerDispatch d); + bool has_staged_run(RunId run_id) const { return staged_run_id_.load(std::memory_order_acquire) == run_id; } + bool needs_activation(RunId run_id) const { + return has_staged_run(run_id) && activated_run_id_.load(std::memory_order_acquire) != run_id; + } + bool activate_prepared(RunId run_id); // True if the worker has no active task. bool idle() const { return inflight_.load(std::memory_order_acquire) == 0; } @@ -514,13 +539,20 @@ class WorkerThread { std::queue queue_; std::mutex mu_; std::condition_variable cv_; - bool shutdown_{false}; + std::atomic shutdown_{false}; uint32_t capacity_{1}; std::atomic inflight_{0}; std::atomic next_dispatch_id_{1}; + std::atomic staged_run_id_{INVALID_RUN_ID}; + std::atomic backend_ready_run_id_{INVALID_RUN_ID}; + std::atomic activated_run_id_{INVALID_RUN_ID}; + std::mutex activation_mu_; + std::condition_variable activation_cv_; + RunId activation_requested_run_id_{INVALID_RUN_ID}; void loop(); WorkerCompletion dispatch_process(WorkerDispatch d, const std::function &on_accept); + void enqueue_dispatch(WorkerDispatch d); }; // ============================================================================= @@ -537,8 +569,13 @@ class WorkerManager { // callable for SUB) lives in the forked child. // `child_pid` is the forked child servicing `mailbox`; pass -1 when the // caller owns no waitable child. - void add_next_level(void *mailbox, int child_pid = -1, uint32_t task_frame_count = 1); - void add_next_level_at(int32_t worker_id, void *mailbox, int child_pid = -1, uint32_t task_frame_count = 1); + void add_next_level( + void *mailbox, int child_pid = -1, uint32_t task_frame_count = 1, bool supports_prepare_activate = false + ); + void add_next_level_at( + int32_t worker_id, void *mailbox, int child_pid = -1, uint32_t task_frame_count = 1, + bool supports_prepare_activate = false + ); void add_next_level_endpoint(std::unique_ptr endpoint); void add_sub(void *mailbox, int child_pid = -1); @@ -556,6 +593,9 @@ class WorkerManager { WorkerThread *pick_idle_sub_excluding(const std::vector &exclude) const; bool any_busy() const; + bool has_staged_run(RunId run_id) const; + bool needs_activation(RunId run_id) const; + bool activate_prepared_run(RunId run_id); // Forward CTRL_PREPARE to a specific NEXT_LEVEL worker. Thin wrapper // over WorkerThread::control_prepare; exposed at manager level so the @@ -625,6 +665,7 @@ class WorkerManager { void *mailbox{nullptr}; int child_pid{-1}; uint32_t task_frame_count{1}; + bool supports_prepare_activate{false}; }; struct LocalSubEntry { void *mailbox{nullptr}; diff --git a/tests/ut/cpp/hierarchical/test_scheduler.cpp b/tests/ut/cpp/hierarchical/test_scheduler.cpp index 80bbea093..6a7511376 100644 --- a/tests/ut/cpp/hierarchical/test_scheduler.cpp +++ b/tests/ut/cpp/hierarchical/test_scheduler.cpp @@ -288,6 +288,57 @@ class BlockingTwoFrameEndpoint final : public WorkerEndpoint { bool release_{false}; }; +class PreparedActivationEndpoint final : public WorkerEndpoint { +public: + PreparedActivationEndpoint() { + caps_.worker_id = 0; + caps_.max_inflight_tasks = 2; + caps_.supports_prepare_activate = true; + } + + const WorkerEndpointCaps &caps() const override { return caps_; } + + WorkerCompletion run(Ring *, const WorkerDispatch &dispatch) override { + normal_running_.store(true, std::memory_order_release); + std::unique_lock lk(mu_); + cv_.wait(lk, [this] { + return release_normal_; + }); + normal_running_.store(false, std::memory_order_release); + return WorkerCompletion{dispatch.task_slot, dispatch.group_index, EndpointOutcome::SUCCESS, {}}; + } + + WorkerCompletion run_prepared_with_activation( + Ring *, const WorkerDispatch &dispatch, const std::function &on_backend_ready, + const std::function &wait_for_activation, const std::function &on_accept + ) override { + backend_ready_.store(true, std::memory_order_release); + if (on_backend_ready) on_backend_ready(); + wait_for_activation(); + activated_.store(true, std::memory_order_release); + if (on_accept) on_accept(); + return WorkerCompletion{dispatch.task_slot, dispatch.group_index, EndpointOutcome::SUCCESS, {}}; + } + + bool normal_running() const { return normal_running_.load(std::memory_order_acquire); } + bool backend_ready() const { return backend_ready_.load(std::memory_order_acquire); } + bool activated() const { return activated_.load(std::memory_order_acquire); } + void release_normal() { + std::lock_guard lk(mu_); + release_normal_ = true; + cv_.notify_all(); + } + +private: + WorkerEndpointCaps caps_; + std::atomic normal_running_{false}; + std::atomic backend_ready_{false}; + std::atomic activated_{false}; + std::mutex mu_; + std::condition_variable cv_; + bool release_normal_{false}; +}; + // --------------------------------------------------------------------------- // Helper: build a TaskArgs whose only tensor has the given (data, tag). // --------------------------------------------------------------------------- @@ -499,6 +550,160 @@ TEST(WorkerManagerTest, TwoFramePrePublishExceptionRetiresDispatchSequence) { allocator.shutdown(); } +TEST(WorkerManagerTest, PreparedFrameRequiresExplicitActivation) { + alignas(8) std::array mailbox{}; + Ring allocator; + allocator.init(/*heap_bytes=*/0); + AllocResult allocation = allocator.alloc(/*heap_bytes=*/0, /*depth=*/0); + ASSERT_NE(allocation.slot, INVALID_SLOT); + TaskSlotState *slot = allocator.slot_state(allocation.slot); + ASSERT_NE(slot, nullptr); + slot->reset(); + slot->run_id = 41; + slot->pipeline_lease = PipelineSlotLease{1, 0, 9}; + + LocalMailboxEndpoint endpoint( + /*worker_id=*/0, mailbox.data(), /*child_pid=*/-1, /*task_frame_count=*/2, + /*supports_prepare_activate=*/true + ); + std::atomic allow_activation{false}; + std::atomic backend_ready{false}; + std::atomic accepted{false}; + std::promise result; + auto done = result.get_future(); + + std::thread child([&] { + char *frame = mailbox.data() + MAILBOX_FRAME_SIZE; + auto *state = reinterpret_cast(frame + MAILBOX_OFF_STATE); + while (__atomic_load_n(state, __ATOMIC_ACQUIRE) != static_cast(MailboxState::PREPARE_READY)) { + std::this_thread::sleep_for(std::chrono::microseconds(50)); + } + __atomic_store_n(state, static_cast(MailboxState::BACKEND_READY), __ATOMIC_RELEASE); + while (__atomic_load_n(state, __ATOMIC_ACQUIRE) != static_cast(MailboxState::ACTIVATE)) { + std::this_thread::sleep_for(std::chrono::microseconds(50)); + } + __atomic_store_n(state, static_cast(MailboxState::TASK_ACTIVE), __ATOMIC_RELEASE); + __atomic_store_n(state, static_cast(MailboxState::TASK_DONE), __ATOMIC_RELEASE); + }); + + std::thread caller([&] { + result.set_value(endpoint.run_prepared_with_activation( + &allocator, WorkerDispatch{allocation.slot, 0, 1, true}, + [&] { + backend_ready.store(true, std::memory_order_release); + }, + [&] { + while (!allow_activation.load(std::memory_order_acquire)) { + std::this_thread::sleep_for(std::chrono::microseconds(50)); + } + }, + [&] { + accepted.store(true, std::memory_order_release); + } + )); + }); + + auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(3); + while (!backend_ready.load(std::memory_order_acquire) && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_TRUE(backend_ready.load(std::memory_order_acquire)); + EXPECT_FALSE(accepted.load(std::memory_order_acquire)); + EXPECT_EQ(done.wait_for(std::chrono::milliseconds(0)), std::future_status::timeout); + + allow_activation.store(true, std::memory_order_release); + EXPECT_EQ(done.wait_for(std::chrono::seconds(3)), std::future_status::ready); + if (done.valid() && done.wait_for(std::chrono::seconds(0)) == std::future_status::ready) { + EXPECT_EQ(done.get().outcome, EndpointOutcome::SUCCESS); + } + EXPECT_TRUE(accepted.load(std::memory_order_acquire)); + caller.join(); + child.join(); + allocator.shutdown(); +} + +TEST(SchedulerPreparedRunTest, SuccessorPreparesButActivatesOnlyAfterFifoPromotion) { + TensorMap tensor_map; + Ring allocator; + Scope scope; + ReadyQueue ready_sub; + NextLevelReadyQueues ready_next; + Orchestrator orchestrator; + WorkerManager manager; + Scheduler scheduler; + CallConfig config; + allocator.init(/*heap_bytes=*/1ULL << 20); + + auto endpoint = std::make_unique(); + PreparedActivationEndpoint *endpoint_ptr = endpoint.get(); + manager.add_next_level_endpoint(std::move(endpoint)); + manager.start( + &allocator, + [&](WorkerCompletion completion) { + scheduler.worker_done(std::move(completion)); + }, + [&](WorkerDispatch dispatch) { + orchestrator.mark_task_accepted(dispatch.task_slot); + } + ); + ready_next.reset(manager.next_level_worker_ids()); + orchestrator.init(&tensor_map, &allocator, &scope, &ready_sub, &ready_next, &manager, [&] { + scheduler.notify_ready(); + }); + orchestrator.configure_pipeline_depth(2); + + Scheduler::Config scheduler_config; + scheduler_config.ring = &allocator; + scheduler_config.ready_sub_queue = &ready_sub; + scheduler_config.ready_next_level_queues = &ready_next; + scheduler_config.manager = &manager; + scheduler_config.enqueue_ready_cb = [&](TaskSlot task_slot) { + orchestrator.enqueue_ready(task_slot); + }; + scheduler_config.active_run_cb = [&] { + return orchestrator.active_run_id(); + }; + scheduler_config.preparable_run_cb = [&] { + return orchestrator.preparable_run_id(); + }; + scheduler_config.on_consumed_cb = [&](TaskSlot task_slot) { + orchestrator.on_consumed(task_slot); + }; + scheduler_config.on_task_failed_cb = [&](TaskSlot task_slot, const std::string &message) { + orchestrator.report_task_error(task_slot, message); + }; + scheduler.start(scheduler_config); + + RunId first_run = orchestrator.begin_run(); + orchestrator.submit_next_level(C(1), single_tensor_args(0x1000, TensorArgType::OUTPUT), config, 0); + orchestrator.close_run_submission(first_run); + RunId second_run = orchestrator.begin_run(); + orchestrator.submit_next_level(C(2), single_tensor_args(0x2000, TensorArgType::OUTPUT), config, 0); + orchestrator.close_run_submission(second_run); + + auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(3); + while ((!endpoint_ptr->normal_running() || !endpoint_ptr->backend_ready()) && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_TRUE(endpoint_ptr->normal_running()); + EXPECT_TRUE(endpoint_ptr->backend_ready()); + EXPECT_FALSE(endpoint_ptr->activated()); + EXPECT_EQ(orchestrator.active_run_id(), first_run); + EXPECT_EQ(orchestrator.preparable_run_id(), second_run); + + endpoint_ptr->release_normal(); + orchestrator.wait_run(first_run); + orchestrator.wait_run(second_run); + EXPECT_TRUE(endpoint_ptr->activated()); + + orchestrator.release_run(first_run); + orchestrator.release_run(second_run); + scheduler.stop(); + manager.stop(); + allocator.shutdown(); +} + TEST(WorkerManagerTest, LocalMailboxPublishesAcceptanceBeforeCompletion) { MockMailboxWorker child; child.start();