From 6b3904277b9b58bf1a0853483b76292379ab0fab Mon Sep 17 00:00:00 2001 From: Crane-Liu Date: Mon, 27 Jul 2026 22:55:58 +0800 Subject: [PATCH 1/4] 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/4] 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/4] 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(); From fb31c0541f637ee736ef6b9eea8cf3973451030c Mon Sep 17 00:00:00 2001 From: Crane-Liu Date: Wed, 29 Jul 2026 19:21:51 +0800 Subject: [PATCH 4/4] Add: HBG prepared epoch pipeline --- docs/task-flow.md | 23 +- python/bindings/task_interface.cpp | 41 ++++ python/simpler/worker.py | 224 +++++++++++++++--- .../platform/onboard/host/device_runner.h | 7 +- .../host_build_graph/host/runtime_maker.cpp | 7 + src/common/hierarchical/scheduler.cpp | 6 +- .../platform/onboard/host/c_api_shared.cpp | 193 +++++++++++++++ .../onboard/host/device_runner_base.cpp | 79 ++++-- .../onboard/host/device_runner_base.h | 20 +- src/common/worker/chip_worker.cpp | 181 ++++++++++++++ src/common/worker/chip_worker.h | 41 +++- src/common/worker/pto_runtime_c_api.h | 33 +++ .../orchestration/pipelined_vector_orch.cpp | 66 ++++++ .../test_worker_async_fifo.py | 30 ++- tests/ut/cpp/hierarchical/test_scheduler.cpp | 87 +++++++ tests/ut/py/test_worker/test_host_worker.py | 112 ++++++++- 16 files changed, 1054 insertions(+), 96 deletions(-) create mode 100644 tests/st/a2a3/host_build_graph/worker_async_fifo/kernels/orchestration/pipelined_vector_orch.cpp diff --git a/docs/task-flow.md b/docs/task-flow.md index b79006fa6..7a250ea78 100644 --- a/docs/task-flow.md +++ b/docs/task-flow.md @@ -321,10 +321,12 @@ 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. +Each direct chip child publishes its runtime contract's `pipeline_depth` and +endpoint capability bits in the startup mailbox before `INIT_READY`. The parent +configures admission to the minimum published depth and enables prepare/activate +per endpoint only when the loaded backend reports the complete prepared-run ABI. +Backends without a depth-two contract therefore keep depth-one serial behavior +instead of receiving an invalid slot-1 lease. #### Prepare/activate endpoint lane @@ -348,11 +350,14 @@ frame and its generation-safe identity. Whole-run FIFO promotion latches 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. +Only the first eligible, non-diagnostic dispatch in a run uses this lane. Other +tasks from the same run retain the existing sequential compatibility path. +Diagnostic execution takes an exclusive backend gate so its runner-global +collectors cannot overlap Host preparation. Remote, SUB, A5, TMR, sim, +single-frame, and endpoints without the capability remain unchanged. Onboard +a2a3 HBG is the first backend to advertise support: it builds into the +lease-selected arena bank, reports `BACKEND_READY` without publishing device +work, and launches only after the matching `ACTIVATE`. 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 diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index b7c07cc2f..cae811638 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -1484,6 +1484,46 @@ NB_MODULE(_task_interface, m) { "of the per-task tensor/scalar layout). The blob must be in the format " "produced by `write_blob`; read_blob enforces capacity bounds against shm corruption." ) + .def( + "prepare_from_blob", + [](ChipWorker &self, int32_t callable_id, uint64_t args_blob_ptr, size_t blob_capacity, + const CallConfig &config, uint32_t pipeline_slot, uint64_t pipeline_generation, uint64_t run_id, + uint64_t dispatch_id) { + nb::gil_scoped_release release; + TaskArgsView view = read_blob(reinterpret_cast(args_blob_ptr), blob_capacity); + ChipStorageTaskArgs storage = view_to_chip_storage(view); + self.prepare_with_lease( + callable_id, &storage, config, PipelineSlotLease{pipeline_slot, 0, pipeline_generation}, run_id, + dispatch_id + ); + }, + nb::arg("callable_id"), nb::arg("args_blob_ptr"), nb::arg("blob_capacity"), nb::arg("config"), + nb::arg("pipeline_slot"), nb::arg("pipeline_generation"), nb::arg("run_id"), nb::arg("dispatch_id"), + "Build one unpublished prepared run in a generation-safe pipeline slot." + ) + .def( + "execute_prepared", + [](ChipWorker &self, int32_t callable_id, const CallConfig &config, uint32_t pipeline_slot, + uint64_t pipeline_generation, uint64_t run_id, uint64_t dispatch_id) { + nb::gil_scoped_release release; + self.execute_prepared( + callable_id, config, PipelineSlotLease{pipeline_slot, 0, pipeline_generation}, run_id, dispatch_id + ); + }, + nb::arg("callable_id"), nb::arg("config"), nb::arg("pipeline_slot"), nb::arg("pipeline_generation"), + nb::arg("run_id"), nb::arg("dispatch_id"), + "Publish and complete a matching prepared run on the sequential executor." + ) + .def( + "abort_prepared", + [](ChipWorker &self, uint32_t pipeline_slot, uint64_t pipeline_generation, uint64_t run_id, + uint64_t dispatch_id) { + nb::gil_scoped_release release; + self.abort_prepared(PipelineSlotLease{pipeline_slot, 0, pipeline_generation}, run_id, dispatch_id); + }, + nb::arg("pipeline_slot"), nb::arg("pipeline_generation"), nb::arg("run_id"), nb::arg("dispatch_id"), + "Reclaim an unpublished prepared run after a local protocol failure." + ) .def( "unregister_callable", [](ChipWorker &self, int32_t callable_id) { @@ -1498,6 +1538,7 @@ NB_MODULE(_task_interface, m) { .def_prop_ro("initialized", &ChipWorker::initialized) .def_prop_ro("pipeline_depth", &ChipWorker::pipeline_depth) .def_prop_ro("runtime_slot_count", &ChipWorker::runtime_slot_count) + .def_prop_ro("supports_prepared_runs", &ChipWorker::supports_prepared_runs) .def_prop_ro( "runtime_buffer_addrs", &ChipWorker::runtime_buffer_addrs, "Host Runtime staging buffer address of every copy the runtime's " diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 207050e36..296fa9fb4 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -230,6 +230,7 @@ def my_l4_orch(orch, args, config): _ACTIVATE = 12 _PREPARE_READY = 13 _TASK_FRAME_COUNT = 2 +_PIPELINE_CAP_PREPARE_ACTIVATE = 1 << 0 def _local_task_frame_count(platform: str, runtime: str, pipeline_depth: int) -> int: @@ -1728,11 +1729,12 @@ 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: + def run_two_frame_loop() -> None: # noqa: PLR0915 -- two-frame mailbox state machine action_cv = threading.Condition() actions: list[tuple[str, int, tuple[int, int, int, int, int] | None]] = [] stop_admission = threading.Event() control_queued = False + prepared_frames: dict[int, tuple[int, int, int, int, int]] = {} frame_bufs = [ buf[(1 + index) * MAILBOX_FRAME_SIZE : (2 + index) * MAILBOX_FRAME_SIZE] for index in range(_TASK_FRAME_COUNT) @@ -1748,6 +1750,89 @@ def read_identity(frame_buf: memoryview) -> tuple[int, int, int, int, int]: struct.unpack_from("=Q", frame_buf, _OFF_FRAME_DISPATCH_ID)[0], ) + def validate_identity(identity: tuple[int, int, int, int, int]) -> bool: + protocol, run_id, slot_id, generation, dispatch_id = identity + return ( + protocol == _TASK_PROTOCOL_VERSION + and run_id != 0 + and slot_id < task_frame_count + and generation != 0 + and dispatch_id != 0 + ) + + def resolve_prepared_frame( + frame_buf: memoryview, + frame_addr: int, + identity: tuple[int, int, int, int, int], + ) -> tuple[int, CallConfig]: + _protocol, run_id, slot_id, generation, dispatch_id = identity + digest = _read_task_digest(frame_buf) + cid = identity_table.get(digest) + if cid is None: + raise RuntimeError(f"callable hash {_format_digest(digest)} not registered") + if cid not in prepared: + raise RuntimeError( + f"chip_process dev={device_id}: cid {cid} not prepared before PREPARE_READY " + f"(register via _CTRL_PREPARE first)" + ) + if host_buf_ranges: + _rewrite_blob_host_addrs(frame_buf, _OFF_TASK_ARGS_BLOB, host_buf_ranges) + pipeline_slot, pipeline_reserved, pipeline_generation = _PIPELINE_LEASE_FMT.unpack_from( + frame_buf, _OFF_PIPELINE_LEASE + ) + if pipeline_reserved != 0 or pipeline_slot != slot_id or pipeline_generation != generation: + raise RuntimeError(f"chip_process dev={device_id}: task frame and pipeline lease disagree") + cfg = _read_config_from_mailbox(frame_buf) + if any( + ( + cfg.enable_l2_swimlane, + cfg.enable_dump_args, + cfg.enable_pmu, + cfg.enable_dep_gen, + cfg.enable_scope_stats, + ) + ): + raise RuntimeError("prepared execution does not support per-run diagnostics") + cw._impl.prepare_from_blob( + int(cid), + frame_addr + _OFF_TASK_ARGS_BLOB, + _MAILBOX_ARGS_CAPACITY, + cfg, + int(slot_id), + int(generation), + int(run_id), + int(dispatch_id), + ) + return int(cid), cfg + + def execute_prepared_frame( + frame_buf: memoryview, + identity: tuple[int, int, int, int, int], + ) -> tuple[int, str]: + _protocol, run_id, slot_id, generation, dispatch_id = identity + digest = _read_task_digest(frame_buf) + cid = identity_table.get(digest) + cfg = _read_config_from_mailbox(frame_buf) + try: + if cid is None: + raise RuntimeError(f"callable hash {_format_digest(digest)} not registered") + cw._impl.execute_prepared( + int(cid), cfg, int(slot_id), int(generation), int(run_id), int(dispatch_id) + ) + except Exception as e: # noqa: BLE001 + return 1, _format_exc(f"chip_process dev={device_id} execute prepared", e) + if on_task_done_success is not None: + return on_task_done_success() + return 0, "" + + def abort_prepared(identity: tuple[int, int, int, int, int]) -> None: + _protocol, run_id, slot_id, generation, dispatch_id = identity + cw._impl.abort_prepared(int(slot_id), int(generation), int(run_id), int(dispatch_id)) + + def fail_frame(frame_buf: memoryview, frame_state_addr: int, message: str) -> None: + _write_error(frame_buf, 1, message) + _mailbox_store_i32(frame_state_addr, _TASK_FAILED) + def admission_loop() -> None: nonlocal control_queued while not stop_admission.is_set(): @@ -1764,23 +1849,71 @@ def admission_loop() -> None: 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: + frame_state = _mailbox_load_i32(frame_state_addr) + if frame_state not in (_TASK_READY, _PREPARE_READY, _ACTIVATE): 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( + if not validate_identity(identity): + fail_frame( frame_buf, - 1, + frame_state_addr, f"chip_process dev={device_id}: invalid task frame identity {identity}", ) - _mailbox_store_i32(frame_state_addr, _TASK_FAILED) + continue + + if frame_state == _PREPARE_READY: + if index in prepared_frames: + fail_frame( + frame_buf, + frame_state_addr, + f"chip_process dev={device_id}: pipeline slot {index} already prepared", + ) + continue + try: + resolve_prepared_frame(frame_buf, frame_addr, identity) + except Exception as e: # noqa: BLE001 + fail_frame( + frame_buf, + frame_state_addr, + _format_exc(f"chip_process dev={device_id} prepare", e), + ) + continue + prepared_frames[index] = identity + _mailbox_store_i32(frame_state_addr, _BACKEND_READY) + continue + + if frame_state == _ACTIVATE: + prepared_identity = prepared_frames.get(index) + if prepared_identity != identity: + if prepared_identity is not None: + try: + abort_prepared(prepared_identity) + except Exception: # noqa: BLE001 + pass + prepared_frames.pop(index, None) + fail_frame( + frame_buf, + frame_state_addr, + f"chip_process dev={device_id}: ACTIVATE identity does not match prepared epoch", + ) + continue + _mailbox_store_i32(frame_state_addr, _TASK_ACCEPTED_STATE) + with action_cv: + actions.append(("prepared", index, identity)) + action_cv.notify() + continue + + if index in prepared_frames: + prepared_identity = prepared_frames.pop(index) + try: + abort_prepared(prepared_identity) + except Exception: # noqa: BLE001 + pass + fail_frame( + frame_buf, + frame_state_addr, + f"chip_process dev={device_id}: TASK_READY attempted to overwrite prepared epoch", + ) continue _mailbox_store_i32(frame_state_addr, _TASK_ACCEPTED_STATE) with action_cv: @@ -1813,20 +1946,36 @@ def admission_loop() -> None: frame_addr = frame_addrs[index] frame_state_addr = frame_addr + _OFF_STATE if identity is None or read_identity(frame_buf) != identity: + if kind == "prepared" and identity is not None: + try: + abort_prepared(identity) + except Exception: # noqa: BLE001 + pass + prepared_frames.pop(index, None) _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, - ) + if kind == "prepared": + prepared_frames.pop(index, None) + code, msg = execute_prepared_frame(frame_buf, identity) + else: + code, msg = handle_task( + frame_buf, + frame_addr, + publish_native_acceptance=True, + ) _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 identity in tuple(prepared_frames.values()): + try: + abort_prepared(identity) + except Exception: # noqa: BLE001 + pass + prepared_frames.clear() for frame_buf in frame_bufs: frame_buf.release() @@ -1911,9 +2060,15 @@ def _chip_process_loop( # noqa: PLR0913 -- fork-child entry: all context (bins, # 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) + # the backend's supported admission depth and reserved carries endpoint + # capability bits. Dispatches later overwrite the same fixed wire region + # with the run-owned slot/generation lease, whose reserved field stays zero. + local_frame_count = _local_task_frame_count(platform, runtime, int(cw.pipeline_depth)) + supports_prepare_activate = local_frame_count >= 2 and bool( + getattr(getattr(cw, "_impl", None), "supports_prepared_runs", False) + ) + startup_caps = _PIPELINE_CAP_PREPARE_ACTIVATE if supports_prepare_activate else 0 + _PIPELINE_LEASE_FMT.pack_into(buf, _OFF_PIPELINE_LEASE, int(cw.pipeline_depth), startup_caps, 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() @@ -1931,7 +2086,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)), + task_frame_count=_TASK_FRAME_COUNT if supports_prepare_activate else 1, ) finally: cw.finalize() @@ -4638,6 +4793,7 @@ def _start_hierarchical(self) -> None: # noqa: PLR0912 -- three parallel fork l n_sub = self._config.get("num_sub_workers", 0) deadline = self._startup_deadline direct_chip_pipeline_depth = PTO_PIPELINE_MAX_DEPTH + direct_chip_prepare_caps: list[bool] = [] # Freeze the startup registry snapshot. init() already holds the epoch in # the INITIALIZING state, so a concurrent register/unregister is blocked @@ -4748,8 +4904,16 @@ def _setup(): 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]) + # advertisement and reserved as capability bits; task dispatch + # restores normal lease semantics with reserved == 0. + depth, caps, startup_generation = _PIPELINE_LEASE_FMT.unpack_from(buf, _OFF_PIPELINE_LEASE) + if startup_generation != 0 or caps & ~_PIPELINE_CAP_PREPARE_ACTIVATE: + raise RuntimeError( + f"chip worker published invalid pipeline startup metadata: " + f"depth={depth}, caps={caps}, generation={startup_generation}" + ) + chip_depths.append(depth) + direct_chip_prepare_caps.append(bool(caps & _PIPELINE_CAP_PREPARE_ACTIVATE)) 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) @@ -4823,11 +4987,13 @@ 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") - 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) + for shm, pid, supports_prepare_activate in zip( + self._chip_shms, self._chip_pids, direct_chip_prepare_caps, strict=True + ): + task_frame_count = _TASK_FRAME_COUNT if supports_prepare_activate else 1 + dw.add_next_level_worker( + _mailbox_addr(shm), pid, task_frame_count, supports_prepare_activate + ) # Register Worker children as NEXT_LEVEL (L4+) if self._next_level_shms and not hasattr(dw, "add_next_level_worker_at"): diff --git a/src/a2a3/platform/onboard/host/device_runner.h b/src/a2a3/platform/onboard/host/device_runner.h index 62f818375..e45ea0745 100644 --- a/src/a2a3/platform/onboard/host/device_runner.h +++ b/src/a2a3/platform/onboard/host/device_runner.h @@ -24,6 +24,7 @@ #include +#include #include #include #include @@ -115,7 +116,7 @@ class DeviceRunner : public DeviceRunnerBase { * and read off DeviceRunner state / HostLogger here — no per-run args. */ int run(Runtime &runtime, const CallConfig &config) override; - bool can_accept_run() const override { return !device_unusable_; } + bool can_accept_run() const override { return !device_unusable_.load(std::memory_order_acquire); } // Map/unmap a device buffer into host address space via // halHostRegister(DEV_SVM_MAP_HOST) / halHostUnregister. The returned host @@ -224,7 +225,9 @@ class DeviceRunner : public DeviceRunnerBase { // this path so the next Worker re-inits clean in the same process (see // force_reset_device()). This flag fails run() fast and drives that // recovery. See run() and recover_device_or_mark_unusable(). - bool device_unusable_{false}; + // The prepared-run admission thread reads this while the sole device + // executor may poison the runner after a failed launch. + std::atomic device_unusable_{false}; struct RunStreamSet { rtStream_t aicpu{nullptr}; diff --git a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp index df21926bf..821c2763a 100644 --- a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp @@ -82,6 +82,13 @@ extern "C" const PipelineContract *get_pipeline_contract(void) { return &contract; } +extern "C" int prepared_run_supported_impl(void) { + // HBG can materialize a complete graph into the lease-selected unpublished + // arena bank. The common C API keeps collector-bearing configurations on + // the sequential path until their state is per-epoch. + return 1; +} + // RuntimeEnv (call_config.h) is the cross-runtime ABI for per-ring config and // carries RUNTIME_ENV_RING_COUNT slots, shared with tensormap_and_ringbuffer. // host_build_graph is single-ring (PTO2_MAX_RING_DEPTH == 1) and reads only the diff --git a/src/common/hierarchical/scheduler.cpp b/src/common/hierarchical/scheduler.cpp index fa7681332..1d94168cd 100644 --- a/src/common/hierarchical/scheduler.cpp +++ b/src/common/hierarchical/scheduler.cpp @@ -381,7 +381,7 @@ void Scheduler::dispatch_preparable_next_level_group() { 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()) { + state.worker_type != WorkerType::NEXT_LEVEL || !state.is_group() || state.config.diagnostics_any()) { return; } @@ -421,6 +421,10 @@ void Scheduler::dispatch_preparable_next_level_singles() { cfg_.enqueue_ready_cb(slot); continue; } + if (state.config.diagnostics_any()) { + cfg_.enqueue_ready_cb(slot); + continue; + } state.state.store(TaskState::RUNNING, std::memory_order_release); worker->dispatch_prepared(WorkerDispatch{slot, 0}); return; diff --git a/src/common/platform/onboard/host/c_api_shared.cpp b/src/common/platform/onboard/host/c_api_shared.cpp index 716a41738..670452106 100644 --- a/src/common/platform/onboard/host/c_api_shared.cpp +++ b/src/common/platform/onboard/host/c_api_shared.cpp @@ -35,7 +35,10 @@ #include #include +#include #include +#include +#include #include #include @@ -58,6 +61,7 @@ extern "C" { * =========================================================================== */ int register_callable_impl(const ChipCallable *callable, uint64_t (*upload_fn)(const void *), CallableArtifacts *out); int validate_runtime_impl(Runtime *runtime, const HostApi *api, int execution_rc); +__attribute__((weak)) int prepared_run_supported_impl(void) { return 0; } /* =========================================================================== * Per-thread DeviceRunnerBase binding (set by simpler_register_callable / simpler_run) @@ -67,6 +71,12 @@ static pthread_key_t g_runner_key; static pthread_once_t g_runner_key_once = PTHREAD_ONCE_INIT; static void create_runner_key() { pthread_key_create(&g_runner_key, nullptr); } +// Host graph preparation may overlap ordinary device execution, but the +// collector configuration is runner-global. Diagnostic runs therefore take +// this gate exclusively while non-diagnostic execution and preparation share +// it. Each loaded runtime SO has its own gate and one DeviceRunner per child. +static std::shared_mutex g_prepared_run_gate; + static DeviceRunnerBase *current_runner() { return static_cast(pthread_getspecific(g_runner_key)); } /* =========================================================================== @@ -562,6 +572,14 @@ int simpler_run( return -1; } + std::shared_lock run_lock(g_prepared_run_gate, std::defer_lock); + std::unique_lock diagnostic_lock(g_prepared_run_gate, std::defer_lock); + if (config->diagnostics_any()) { + diagnostic_lock.lock(); + } else { + run_lock.lock(); + } + pthread_once(&g_runner_key_once, create_runner_key); pthread_setspecific(g_runner_key, ctx); auto tsd_guard = RAIIScopeGuard([]() { @@ -622,6 +640,7 @@ int simpler_run( return validation_rc != 0 ? validation_rc : rc; } + runner->activate_launch_shape(*r); { STRACE("simpler_run.runner_run"); // run() latches the diagnostic enables from config via @@ -648,6 +667,180 @@ int simpler_run( } } +static bool valid_prepared_identity(const PreparedRunIdentity *identity) { + return identity != nullptr && identity->run_id != 0 && identity->generation != 0 && identity->dispatch_id != 0 && + identity->reserved == 0 && identity->slot_id < PTO_PIPELINE_MAX_DEPTH; +} + +static void format_prepared_attrs(const PreparedRunIdentity &identity, char *attrs, size_t attrs_size) { + std::snprintf( + attrs, attrs_size, "request=%llu run=%llu epoch=%llu slot=%u generation=%llu dispatch=%llu", + static_cast(identity.run_id), static_cast(identity.run_id), + static_cast(identity.generation), identity.slot_id, + static_cast(identity.generation), static_cast(identity.dispatch_id) + ); +} + +int supports_prepared_run_ctx(DeviceContextHandle ctx) { + return ctx != NULL && prepared_run_supported_impl() != 0 ? 1 : 0; +} + +int simpler_prepare_run( + DeviceContextHandle ctx, RuntimeHandle runtime, int32_t callable_id, const void *args, const CallConfig *config, + const PreparedRunIdentity *identity +) { + if (ctx == NULL || runtime == NULL || args == NULL || config == NULL || !valid_prepared_identity(identity) || + prepared_run_supported_impl() == 0) { + return PTO_RUNTIME_ERR_UNSUPPORTED; + } + DeviceRunnerBase *runner = static_cast(ctx); + if (!runner->has_callable(callable_id) || !runner->can_accept_run() || config->diagnostics_any() || + runner->pipeline_slot() != identity->slot_id) { + return -1; + } + + std::shared_lock prepare_lock(g_prepared_run_gate); + + pthread_once(&g_runner_key_once, create_runner_key); + pthread_setspecific(g_runner_key, ctx); + auto tsd_guard = RAIIScopeGuard([]() { + pthread_setspecific(g_runner_key, nullptr); + }); + + char attrs[160]; + format_prepared_attrs(*identity, attrs, sizeof(attrs)); + STRACE_NEW_INV(); + STRACE_SET_HID(runner->callable_hash(callable_id)); + STRACE_A("simpler_prepare", attrs); + + try { + int rc = runner->attach_current_thread(runner->device_id()); + if (rc != 0) return rc; + Runtime *r = new (runtime) Runtime(); + bool cleanup_validated = false; + auto runtime_guard = RAIIScopeGuard([&]() { + if (!cleanup_validated) { + r->set_gm_sm_ptr(nullptr); + try { + (void)validate_runtime_impl(r, &g_host_api, -1); + } catch (...) {} + } + r->~Runtime(); + }); + rc = runner->prepare_launch_shape(*r, *config); + if (rc == 0) { + STRACE_A("simpler_prepare.bind", attrs); + rc = runner->bind_callable_to_runtime( + *r, callable_id, &g_host_api, args, config->runtime_env.ring_task_window, config->runtime_env.ring_heap, + config->runtime_env.ring_dep_pool + ); + } + if (rc != 0) { + r->set_gm_sm_ptr(nullptr); + int validation_rc = validate_runtime_impl(r, &g_host_api, rc); + cleanup_validated = true; + return validation_rc != 0 ? validation_rc : rc; + } + runtime_guard.dismiss(); + return 0; + } catch (...) { + return -1; + } +} + +int simpler_execute_prepared( + DeviceContextHandle ctx, RuntimeHandle runtime, int32_t callable_id, const CallConfig *config, + const PreparedRunIdentity *identity +) { + if (ctx == NULL || runtime == NULL || config == NULL || !valid_prepared_identity(identity) || + prepared_run_supported_impl() == 0) { + return PTO_RUNTIME_ERR_UNSUPPORTED; + } + DeviceRunnerBase *runner = static_cast(ctx); + if (!runner->has_callable(callable_id) || config->diagnostics_any() || + runner->pipeline_slot() != identity->slot_id) { + return -1; + } + + std::shared_lock execute_lock(g_prepared_run_gate); + + pthread_once(&g_runner_key_once, create_runner_key); + pthread_setspecific(g_runner_key, ctx); + auto tsd_guard = RAIIScopeGuard([]() { + pthread_setspecific(g_runner_key, nullptr); + }); + + char attrs[160]; + format_prepared_attrs(*identity, attrs, sizeof(attrs)); + STRACE_NEW_INV(); + STRACE_SET_HID(runner->callable_hash(callable_id)); + STRACE_A("simpler_execute_prepared", attrs); + + Runtime *r = static_cast(runtime); + try { + bool cleanup_validated = false; + auto runtime_guard = RAIIScopeGuard([&]() { + if (!cleanup_validated) { + r->set_gm_sm_ptr(nullptr); + try { + (void)validate_runtime_impl(r, &g_host_api, -1); + } catch (...) {} + } + r->~Runtime(); + }); + int rc = runner->attach_current_thread(runner->device_id()); + if (rc == 0) { + runner->activate_launch_shape(*r); + STRACE_A("simpler_execute_prepared.runner_run", attrs); + rc = runner->run(*r, *config); + } + int validation_rc = 0; + { + STRACE_A("simpler_execute_prepared.validate", attrs); + validation_rc = validate_runtime_impl(r, &g_host_api, rc); + cleanup_validated = true; + } + emit_device_phase_markers(runner); + return validation_rc != 0 ? validation_rc : rc; + } catch (...) { + return -1; + } +} + +int simpler_abort_prepared(DeviceContextHandle ctx, RuntimeHandle runtime, const PreparedRunIdentity *identity) { + if (ctx == NULL || runtime == NULL || !valid_prepared_identity(identity) || prepared_run_supported_impl() == 0) { + return PTO_RUNTIME_ERR_UNSUPPORTED; + } + std::shared_lock abort_lock(g_prepared_run_gate); + pthread_once(&g_runner_key_once, create_runner_key); + pthread_setspecific(g_runner_key, ctx); + auto tsd_guard = RAIIScopeGuard([]() { + pthread_setspecific(g_runner_key, nullptr); + }); + DeviceRunnerBase *runner = static_cast(ctx); + Runtime *r = static_cast(runtime); + try { + bool cleanup_validated = false; + auto runtime_guard = RAIIScopeGuard([&]() { + if (!cleanup_validated) { + r->set_gm_sm_ptr(nullptr); + try { + (void)validate_runtime_impl(r, &g_host_api, -1); + } catch (...) {} + } + r->~Runtime(); + }); + int rc = runner->attach_current_thread(runner->device_id()); + if (rc != 0) return rc; + r->set_gm_sm_ptr(nullptr); + rc = validate_runtime_impl(r, &g_host_api, -1); + cleanup_validated = true; + return rc; + } catch (...) { + return -1; + } +} + int set_task_accepted_state_ctx(DeviceContextHandle ctx, volatile int32_t *state, int32_t accepted_value) { if (ctx == NULL) return -1; try { diff --git a/src/common/platform/onboard/host/device_runner_base.cpp b/src/common/platform/onboard/host/device_runner_base.cpp index 1f92a8984..04315b9a4 100644 --- a/src/common/platform/onboard/host/device_runner_base.cpp +++ b/src/common/platform/onboard/host/device_runner_base.cpp @@ -27,6 +27,7 @@ #include #include #include +#include #include #include @@ -63,6 +64,32 @@ extern "C" const char *const *runtime_extra_aicpu_symbols(size_t *count); namespace { +struct ThreadRunSelection { + uint32_t pipeline_slot{0}; + uint32_t arena_bank{0}; + volatile int32_t *accepted_state{nullptr}; + int32_t accepted_value{0}; +}; + +pthread_key_t g_run_selection_key; +pthread_once_t g_run_selection_once = PTHREAD_ONCE_INIT; + +void create_run_selection_key() { + pthread_key_create(&g_run_selection_key, [](void *p) { + delete static_cast(p); + }); +} + +ThreadRunSelection &run_selection() { + pthread_once(&g_run_selection_once, create_run_selection_key); + auto *selection = static_cast(pthread_getspecific(g_run_selection_key)); + if (selection == nullptr) { + selection = new ThreadRunSelection{}; + pthread_setspecific(g_run_selection_key, selection); + } + return *selection; +} + HostRuntimeTimeoutConfig resolve_onboard_timeout_config() { RuntimeTimeoutConfig order_defaults{ PLATFORM_OP_EXECUTE_TIMEOUT_US, PLATFORM_STREAM_SYNC_TIMEOUT_MS, PLATFORM_ONBOARD_SCHEDULER_TIMEOUT_MS @@ -131,7 +158,7 @@ int DeviceRunnerBase::select_pipeline_slot(uint32_t slot_id) { LOG_ERROR("pipeline slot %u is outside [0, %u)", slot_id, PTO_PIPELINE_MAX_DEPTH); return -1; } - pipeline_slot_ = slot_id; + run_selection().pipeline_slot = slot_id; return 0; } @@ -140,11 +167,13 @@ int DeviceRunnerBase::select_arena_bank(uint32_t bank_id) { LOG_ERROR("arena bank %u is outside [0, %u)", bank_id, PTO_PIPELINE_MAX_DEPTH); return -1; } - arena_bank_ = bank_id; + run_selection().arena_bank = bank_id; return 0; } -uint32_t DeviceRunnerBase::pipeline_slot() const { return pipeline_slot_; } +uint32_t DeviceRunnerBase::pipeline_slot() const { return run_selection().pipeline_slot; } + +uint32_t DeviceRunnerBase::selected_arena_bank() const { return run_selection().arena_bank; } uint64_t DeviceRunnerBase::arena_bank_gm_heap_base(uint32_t bank_id) const { if (bank_id >= arena_banks_.size()) return 0; @@ -178,13 +207,13 @@ int DeviceRunnerBase::device_memset(void *dev_ptr, int value, std::size_t bytes) } void DeviceRunnerBase::get_retained_temp_buffer(void **addr, size_t *size) { - if (addr != nullptr) *addr = retained_temp_addrs_[pipeline_slot_]; - if (size != nullptr) *size = retained_temp_sizes_[pipeline_slot_]; + if (addr != nullptr) *addr = retained_temp_addrs_[pipeline_slot()]; + if (size != nullptr) *size = retained_temp_sizes_[pipeline_slot()]; } void DeviceRunnerBase::set_retained_temp_buffer(void *addr, size_t size) { - retained_temp_addrs_[pipeline_slot_] = addr; - retained_temp_sizes_[pipeline_slot_] = size; + retained_temp_addrs_[pipeline_slot()] = addr; + retained_temp_sizes_[pipeline_slot()] = size; } void DeviceRunnerBase::clear_temporary_buffer() { @@ -222,7 +251,7 @@ bool DeviceRunnerBase::lookup_prebuilt_runtime_arena_cache( ) const { // The cache holds one entry and its bases point into bank 0, so any other // bank must rebuild rather than be handed a region it does not own. - if (arena_bank_ != 0) return false; + if (selected_arena_bank() != 0) return false; if (!prebuilt_runtime_arena_cache_valid_ || prebuilt_runtime_arena_cache_hash_ != hash || prebuilt_runtime_arena_cache_key_.size() != key_size || key_data == nullptr || gm_heap_base == nullptr || sm_base == nullptr || runtime_arena_base == nullptr || runtime_off == nullptr || image_data == nullptr || @@ -246,7 +275,7 @@ void DeviceRunnerBase::mark_prebuilt_runtime_arena_cached( size_t runtime_off, const void *image_data, size_t image_size ) { // Single-entry cache owned by bank 0; see lookup_prebuilt_runtime_arena_cache. - if (arena_bank_ != 0) return; + if (selected_arena_bank() != 0) return; prebuilt_runtime_arena_cache_valid_ = false; prebuilt_runtime_arena_cache_hash_ = hash; prebuilt_runtime_arena_cache_key_.assign( @@ -322,7 +351,7 @@ int DeviceRunnerBase::setup_static_arena(size_t gm_heap_size, size_t gm_sm_size, bank.cached_gm_heap_size = 0; bank.cached_gm_sm_size = 0; bank.cached_runtime_arena_size = 0; - if (arena_bank_ == 0) { + if (selected_arena_bank() == 0) { prebuilt_runtime_arena_cache_valid_ = false; prebuilt_runtime_arena_cache_key_.clear(); prebuilt_runtime_arena_cache_gm_heap_base_ = nullptr; @@ -332,7 +361,7 @@ int DeviceRunnerBase::setup_static_arena(size_t gm_heap_size, size_t gm_sm_size, } return -1; } - if (arena_changed && arena_bank_ == 0) { + if (arena_changed && selected_arena_bank() == 0) { prebuilt_runtime_arena_cache_valid_ = false; prebuilt_runtime_arena_cache_key_.clear(); prebuilt_runtime_arena_cache_gm_heap_base_ = nullptr; @@ -372,12 +401,14 @@ int DeviceRunnerBase::attach_current_thread(int device_id) { return rc; } + // simpler_init performs the only lifetime write. Prepared-run admission + // and execution subsequently attach different host threads, so repeated + // same-value writes here would still be a C++ data race. if (device_id_ == -1) { timeout_config_ = resolve_onboard_timeout_config(); configure_aicore_op_timeout(); + device_id_ = device_id; } - - device_id_ = device_id; return 0; } @@ -1145,8 +1176,6 @@ int DeviceRunnerBase::finalize_common() { bank->cached_gm_sm_size = 0; bank->cached_runtime_arena_size = 0; } - pipeline_slot_ = 0; - arena_bank_ = 0; return rc; } @@ -1258,9 +1287,8 @@ int DeviceRunnerBase::resolve_block_dim() { ); return -1; } - block_dim_ = max_block_dim_; - LOG_INFO("block_dim resolved to %d (cube=%u, vector=%u)", block_dim_, max_cube_cores_, max_vector_cores_); - return block_dim_; + LOG_INFO("block_dim resolved to %d (cube=%u, vector=%u)", max_block_dim_, max_cube_cores_, max_vector_cores_); + return max_block_dim_; } int DeviceRunnerBase::prepare_launch_shape(Runtime &runtime, const CallConfig &config) { @@ -1279,7 +1307,6 @@ int DeviceRunnerBase::prepare_launch_shape(Runtime &runtime, const CallConfig &c } runtime.set_worker_count(num_aicore); - worker_count_ = num_aicore; // Stored for print_handshake_results in destructor runtime.set_aicpu_thread_num(config.aicpu_thread_num); // First `block_dim` cores are AIC; remaining ~2/3 are AIV. @@ -1294,6 +1321,11 @@ int DeviceRunnerBase::prepare_launch_shape(Runtime &runtime, const CallConfig &c return 0; } +void DeviceRunnerBase::activate_launch_shape(const Runtime &runtime) { + worker_count_ = runtime.get_worker_count(); + block_dim_ = worker_count_ / cores_per_blockdim_; +} + void DeviceRunnerBase::resolve_task_binary_addrs(Runtime &runtime) { // Runtime::func_id_to_addr_[] stores a CoreCallable device address; the // binary code address is one compile-time offset further in. The dispatch @@ -1482,13 +1514,14 @@ void DeviceRunnerBase::teardown_shared_collectors_after_run() { } int DeviceRunnerBase::set_task_accepted_state(volatile int32_t *state, int32_t accepted_value) { - task_accepted_state_ = state; - task_accepted_value_ = accepted_value; + run_selection().accepted_state = state; + run_selection().accepted_value = accepted_value; return 0; } void DeviceRunnerBase::publish_task_accepted() const { - if (task_accepted_state_ != nullptr) { - __atomic_store_n(task_accepted_state_, task_accepted_value_, __ATOMIC_RELEASE); + ThreadRunSelection &selection = run_selection(); + if (selection.accepted_state != nullptr) { + __atomic_store_n(selection.accepted_state, selection.accepted_value, __ATOMIC_RELEASE); } } diff --git a/src/common/platform/onboard/host/device_runner_base.h b/src/common/platform/onboard/host/device_runner_base.h index d34b97593..b9df44636 100644 --- a/src/common/platform/onboard/host/device_runner_base.h +++ b/src/common/platform/onboard/host/device_runner_base.h @@ -411,6 +411,8 @@ class DeviceRunnerBase { * Returns 0 on success, -1 on a bad `block_dim` / `aicpu_thread_num`. */ int prepare_launch_shape(Runtime &runtime, const CallConfig &config); + /** Latch a prepared Runtime's geometry immediately before Device S. */ + void activate_launch_shape(const Runtime &runtime); /** * Replay a previously-registered callable's state onto a fresh Runtime and @@ -871,16 +873,11 @@ class DeviceRunnerBase { // Same re-register semantics as `aicpu_dlopen_total_`, but for hbg // variants. size_t host_dlopen_total_{0}; - volatile int32_t *task_accepted_state_{nullptr}; - int32_t task_accepted_value_{0}; - // ---- State shared by both a2a3 and a5 --------------------------------- // // `device_id_` is set once in `attach_current_thread()` (called from - // simpler_init during ChipWorker::init) and read on every subsequent - // op. All ChipWorker callers run on the same thread that called - // init, so plain int + the init→user happens-before edge is - // sufficient. + // simpler_init during ChipWorker::init) and is immutable while prepared- + // run admission and execution attach their respective host threads. int device_id_{-1}; int block_dim_{0}; int cores_per_blockdim_{PLATFORM_CORES_PER_BLOCKDIM}; @@ -958,13 +955,8 @@ class DeviceRunnerBase { // Held by pointer because DeviceArena is non-copyable and non-movable, so // the array cannot be brace-initialised without naming every bank. std::array, PTO_PIPELINE_MAX_DEPTH> arena_banks_; - ArenaBank &arena_bank() { return *arena_banks_[arena_bank_]; } - - // Selection carried by the lease of the run currently being served. Both - // default to zero, which is what an unleased run and every depth-one - // runtime use. - uint32_t pipeline_slot_{0}; - uint32_t arena_bank_{0}; + ArenaBank &arena_bank() { return *arena_banks_[selected_arena_bank()]; } + uint32_t selected_arena_bank() const; bool prebuilt_runtime_arena_cache_valid_{false}; uint64_t prebuilt_runtime_arena_cache_hash_{0}; diff --git a/src/common/worker/chip_worker.cpp b/src/common/worker/chip_worker.cpp index 1136deaf8..c53ed815d 100644 --- a/src/common/worker/chip_worker.cpp +++ b/src/common/worker/chip_worker.cpp @@ -117,6 +117,13 @@ void ChipWorker::init( simpler_init_fn_ = load_symbol(handle, "simpler_init"); register_callable_fn_ = load_symbol(handle, "simpler_register_callable"); run_fn_ = load_symbol(handle, "simpler_run"); + // Prepared execution is an optional runtime extension. Keeping all + // four probes optional lets runtimes without the extension remain on + // the existing synchronous ABI without requiring compatibility stubs. + supports_prepared_run_fn_ = load_optional_symbol(handle, "supports_prepared_run_ctx"); + prepare_run_fn_ = load_optional_symbol(handle, "simpler_prepare_run"); + execute_prepared_fn_ = load_optional_symbol(handle, "simpler_execute_prepared"); + abort_prepared_fn_ = load_optional_symbol(handle, "simpler_abort_prepared"); select_pipeline_slot_fn_ = load_symbol(handle, "select_pipeline_slot_ctx"); select_arena_bank_fn_ = load_symbol(handle, "select_arena_bank_ctx"); get_arena_bank_gm_heap_base_fn_ = @@ -187,6 +194,7 @@ void ChipWorker::init( resolved_contract.pipeline_depth, std::vector(get_runtime_size_fn_()) ); runtime_bufs_.swap(runtime_bufs); + prepared_slots_.assign(resolved_contract.pipeline_depth, PreparedSlot{}); } catch (...) { destroy_device_context_fn_(device_ctx_); device_ctx_ = nullptr; @@ -241,6 +249,10 @@ void ChipWorker::init( simpler_init_fn_ = nullptr; register_callable_fn_ = nullptr; run_fn_ = nullptr; + supports_prepared_run_fn_ = nullptr; + prepare_run_fn_ = nullptr; + execute_prepared_fn_ = nullptr; + abort_prepared_fn_ = nullptr; select_pipeline_slot_fn_ = nullptr; select_arena_bank_fn_ = nullptr; get_arena_bank_gm_heap_base_fn_ = nullptr; @@ -264,6 +276,7 @@ void ChipWorker::init( comm_barrier_fn_ = nullptr; comm_destroy_fn_ = nullptr; runtime_bufs_.clear(); + prepared_slots_.clear(); throw; } if (init_rc != 0) { @@ -286,6 +299,10 @@ void ChipWorker::init( simpler_init_fn_ = nullptr; register_callable_fn_ = nullptr; run_fn_ = nullptr; + supports_prepared_run_fn_ = nullptr; + prepare_run_fn_ = nullptr; + execute_prepared_fn_ = nullptr; + abort_prepared_fn_ = nullptr; select_pipeline_slot_fn_ = nullptr; select_arena_bank_fn_ = nullptr; get_arena_bank_gm_heap_base_fn_ = nullptr; @@ -310,6 +327,7 @@ void ChipWorker::init( comm_barrier_fn_ = nullptr; comm_destroy_fn_ = nullptr; runtime_bufs_.clear(); + prepared_slots_.clear(); throw std::runtime_error("simpler_init failed with code " + std::to_string(init_rc)); } @@ -337,6 +355,28 @@ void ChipWorker::init( } void ChipWorker::finalize() { + // A shutdown/control failure may leave a built epoch unpublished. Destroy + // those Runtime objects before the backend and their arena banks go away. + std::vector prepared_identities; + { + std::scoped_lock lk(prepared_slots_mu_); + for (const PreparedSlot &slot : prepared_slots_) { + if (slot.state == PreparedSlotState::PREPARED) { + prepared_identities.push_back(slot.identity); + } + } + } + for (const PreparedRunIdentity &identity : prepared_identities) { + try { + abort_prepared( + PipelineSlotLease{identity.slot_id, 0, identity.generation}, identity.run_id, identity.dispatch_id + ); + } catch (...) { + // Finalization remains best-effort so the owning device context is + // torn down even if a backend abort reports failure. + } + } + // Defensive: if the user never called comm_destroy, reclaim all owned // communicator handles and streams before tearing down the device context. clear_comm_sessions(); @@ -361,6 +401,10 @@ void ChipWorker::finalize() { get_runtime_size_fn_ = nullptr; register_callable_fn_ = nullptr; run_fn_ = nullptr; + supports_prepared_run_fn_ = nullptr; + prepare_run_fn_ = nullptr; + execute_prepared_fn_ = nullptr; + abort_prepared_fn_ = nullptr; select_pipeline_slot_fn_ = nullptr; select_arena_bank_fn_ = nullptr; get_arena_bank_gm_heap_base_fn_ = nullptr; @@ -385,6 +429,7 @@ void ChipWorker::finalize() { comm_barrier_fn_ = nullptr; comm_destroy_fn_ = nullptr; runtime_bufs_.clear(); + prepared_slots_.clear(); pipeline_generations_.reset(); pipeline_contract_ = {PTO_PIPELINE_CONTRACT_ABI_VERSION, 0, 1, {}}; initialized_ = false; @@ -488,6 +533,142 @@ void ChipWorker::run_with_lease( run_on_slot(callable_id, args, config, lease.slot_id, accepted_state, accepted_value); } +bool ChipWorker::supports_prepared_runs() const { + return initialized_ && supports_prepared_run_fn_ != nullptr && prepare_run_fn_ != nullptr && + execute_prepared_fn_ != nullptr && abort_prepared_fn_ != nullptr && + supports_prepared_run_fn_(device_ctx_) > 0; +} + +void ChipWorker::prepare_with_lease( + int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, const PipelineSlotLease &lease, + uint64_t run_id, uint64_t dispatch_id +) { + config.validate(); + if (!supports_prepared_runs()) { + throw std::runtime_error("bound runtime does not support prepared runs"); + } + if (config.diagnostics_any()) { + throw std::runtime_error("prepared runs with per-run diagnostics use the sequential compatibility path"); + } + if (lease.reserved != 0 || lease.generation == 0 || lease.slot_id >= pipeline_contract_.pipeline_depth || + run_id == 0 || dispatch_id == 0) { + throw std::runtime_error("prepared run identity is outside the runtime PipelineContract"); + } + if (!pipeline_generations_.admit(lease)) { + throw std::runtime_error("prepared run pipeline lease generation is stale"); + } + + const PreparedRunIdentity identity{run_id, lease.generation, dispatch_id, lease.slot_id, 0}; + { + std::scoped_lock lk(prepared_slots_mu_); + PreparedSlot &slot = prepared_slots_.at(lease.slot_id); + if (slot.state != PreparedSlotState::FREE) { + throw std::runtime_error("prepared run slot is not free"); + } + slot.state = PreparedSlotState::BUILDING; + slot.identity = identity; + slot.callable_id = callable_id; + slot.config = config; + } + + int rc = -1; + try { + (void)select_slot_resources(lease.slot_id); + rc = prepare_run_fn_(device_ctx_, runtime_bufs_[lease.slot_id].data(), callable_id, args, &config, &identity); + } catch (...) { + std::scoped_lock lk(prepared_slots_mu_); + prepared_slots_[lease.slot_id] = PreparedSlot{}; + throw; + } + if (rc != 0) { + std::scoped_lock lk(prepared_slots_mu_); + prepared_slots_[lease.slot_id] = PreparedSlot{}; + throw std::runtime_error("prepare run failed with code " + std::to_string(rc)); + } + + bool identity_changed = false; + { + std::scoped_lock lk(prepared_slots_mu_); + PreparedSlot &slot = prepared_slots_[lease.slot_id]; + identity_changed = slot.state != PreparedSlotState::BUILDING || slot.identity.run_id != run_id || + slot.identity.generation != lease.generation || slot.identity.dispatch_id != dispatch_id; + if (identity_changed) { + slot = PreparedSlot{}; + } else { + slot.state = PreparedSlotState::PREPARED; + } + } + if (identity_changed) { + (void)abort_prepared_fn_(device_ctx_, runtime_bufs_[lease.slot_id].data(), &identity); + throw std::runtime_error("prepared run identity changed while building"); + } +} + +void ChipWorker::execute_prepared( + int32_t callable_id, const CallConfig &config, const PipelineSlotLease &lease, uint64_t run_id, uint64_t dispatch_id +) { + config.validate(); + if (config.diagnostics_any()) { + throw std::runtime_error("prepared runs with per-run diagnostics use the sequential compatibility path"); + } + const PreparedRunIdentity identity{run_id, lease.generation, dispatch_id, lease.slot_id, 0}; + { + std::scoped_lock lk(prepared_slots_mu_); + if (lease.slot_id >= prepared_slots_.size()) { + throw std::runtime_error("prepared run slot is outside the runtime PipelineContract"); + } + PreparedSlot &slot = prepared_slots_[lease.slot_id]; + if (slot.state != PreparedSlotState::PREPARED || slot.callable_id != callable_id || + slot.identity.run_id != run_id || slot.identity.generation != lease.generation || + slot.identity.dispatch_id != dispatch_id) { + throw std::runtime_error("prepared run identity does not match the published slot"); + } + if (std::memcmp(&slot.config, &config, sizeof(CallConfig)) != 0) { + throw std::runtime_error("prepared run config does not match the published slot"); + } + slot.state = PreparedSlotState::ACTIVE; + } + + int rc = -1; + try { + (void)select_slot_resources(lease.slot_id); + rc = execute_prepared_fn_(device_ctx_, runtime_bufs_[lease.slot_id].data(), callable_id, &config, &identity); + } catch (...) { + std::scoped_lock lk(prepared_slots_mu_); + prepared_slots_[lease.slot_id] = PreparedSlot{}; + throw; + } + { + std::scoped_lock lk(prepared_slots_mu_); + prepared_slots_[lease.slot_id] = PreparedSlot{}; + } + if (rc != 0) { + throw std::runtime_error("execute prepared run failed with code " + std::to_string(rc)); + } +} + +void ChipWorker::abort_prepared(const PipelineSlotLease &lease, uint64_t run_id, uint64_t dispatch_id) { + const PreparedRunIdentity identity{run_id, lease.generation, dispatch_id, lease.slot_id, 0}; + { + std::scoped_lock lk(prepared_slots_mu_); + if (lease.slot_id >= prepared_slots_.size()) return; + const PreparedSlot &slot = prepared_slots_[lease.slot_id]; + if (slot.state != PreparedSlotState::PREPARED || slot.identity.run_id != run_id || + slot.identity.generation != lease.generation || slot.identity.dispatch_id != dispatch_id) { + return; + } + } + (void)select_slot_resources(lease.slot_id); + int rc = abort_prepared_fn_(device_ctx_, runtime_bufs_[lease.slot_id].data(), &identity); + { + std::scoped_lock lk(prepared_slots_mu_); + prepared_slots_[lease.slot_id] = PreparedSlot{}; + } + if (rc != 0) { + throw std::runtime_error("abort prepared run failed with code " + std::to_string(rc)); + } +} + void ChipWorker::run_on_slot( int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, uint32_t slot_id, volatile int32_t *accepted_state, int32_t accepted_value diff --git a/src/common/worker/chip_worker.h b/src/common/worker/chip_worker.h index f248fb25f..1b82e0808 100644 --- a/src/common/worker/chip_worker.h +++ b/src/common/worker/chip_worker.h @@ -13,6 +13,7 @@ #define SRC_COMMON_WORKER_CHIP_WORKER_H_ #include +#include #include #include #include @@ -84,6 +85,20 @@ class ChipWorker { int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, const PipelineSlotLease &lease, volatile int32_t *accepted_state = nullptr, int32_t accepted_value = 0 ); + + /// Build one unpublished HBG run into the lease-selected slot/bank. + void prepare_with_lease( + int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, const PipelineSlotLease &lease, + uint64_t run_id, uint64_t dispatch_id + ); + /// Publish and complete a run previously prepared under the same identity. + void execute_prepared( + int32_t callable_id, const CallConfig &config, const PipelineSlotLease &lease, uint64_t run_id, + uint64_t dispatch_id + ); + /// Reclaim an unpublished prepared run after a local protocol failure. + void abort_prepared(const PipelineSlotLease &lease, uint64_t run_id, uint64_t dispatch_id); + bool supports_prepared_runs() const; void run_with_lease( int32_t callable_id, TaskArgsView args, const CallConfig &config, const PipelineSlotLease &lease, volatile int32_t *accepted_state = nullptr, int32_t accepted_value = 0 @@ -193,6 +208,11 @@ class ChipWorker { ); using SimplerRegisterCallableFn = int (*)(void *, int32_t, const void *); using SimplerRunFn = int (*)(void *, void *, int32_t, const void *, const CallConfig *); + using SupportsPreparedRunFn = int (*)(void *); + using SimplerPrepareRunFn = + int (*)(void *, void *, int32_t, const void *, const CallConfig *, const PreparedRunIdentity *); + using SimplerExecutePreparedFn = int (*)(void *, void *, int32_t, const CallConfig *, const PreparedRunIdentity *); + using SimplerAbortPreparedFn = int (*)(void *, void *, const PreparedRunIdentity *); using SetTaskAcceptedStateFn = int (*)(void *, volatile int32_t *, int32_t); using SelectPipelineSlotFn = int (*)(void *, uint32_t); using SelectArenaBankFn = int (*)(void *, uint32_t); @@ -245,6 +265,10 @@ class ChipWorker { SimplerInitFn simpler_init_fn_ = nullptr; SimplerRegisterCallableFn register_callable_fn_ = nullptr; SimplerRunFn run_fn_ = nullptr; + SupportsPreparedRunFn supports_prepared_run_fn_ = nullptr; + SimplerPrepareRunFn prepare_run_fn_ = nullptr; + SimplerExecutePreparedFn execute_prepared_fn_ = nullptr; + SimplerAbortPreparedFn abort_prepared_fn_ = nullptr; SetTaskAcceptedStateFn set_task_accepted_state_fn_ = nullptr; SelectPipelineSlotFn select_pipeline_slot_fn_ = nullptr; SelectArenaBankFn select_arena_bank_fn_ = nullptr; @@ -284,13 +308,20 @@ class ChipWorker { uint32_t select_slot_resources(uint32_t slot_id); std::vector> runtime_bufs_; + enum class PreparedSlotState : uint8_t { FREE, BUILDING, PREPARED, ACTIVE }; + struct PreparedSlot { + PreparedSlotState state{PreparedSlotState::FREE}; + PreparedRunIdentity identity{}; + int32_t callable_id{-1}; + CallConfig config{}; + }; + std::mutex prepared_slots_mu_; + std::vector prepared_slots_; PipelineSlotGenerationFilter pipeline_generations_; PipelineContract pipeline_contract_{PTO_PIPELINE_CONTRACT_ABI_VERSION, 0, 1, {}}; - // device_id_ is set once in init() and never modified afterward. All - // ChipWorker callers run on the thread that called init() (the same - // thread is the only one that subsequently calls malloc / copy_to / - // run / finalize), so plain `int` is sufficient — no cross-thread - // synchronization required. + // device_id_ is set once in init() and never modified afterward. The + // prepared-run admission and executor threads only read it after startup; + // lifecycle and memory/control operations remain owner-thread-only. int device_id_ = -1; bool initialized_ = false; bool finalized_ = false; diff --git a/src/common/worker/pto_runtime_c_api.h b/src/common/worker/pto_runtime_c_api.h index 8bf3544fa..b77586831 100644 --- a/src/common/worker/pto_runtime_c_api.h +++ b/src/common/worker/pto_runtime_c_api.h @@ -137,6 +137,15 @@ typedef struct PipelineSlotLease { uint64_t generation; } PipelineSlotLease; +/** Stable identity carried across the Host-prepare and Device-execute halves. */ +typedef struct PreparedRunIdentity { + uint64_t run_id; + uint64_t generation; + uint64_t dispatch_id; + uint32_t slot_id; + uint32_t reserved; +} PreparedRunIdentity; + /* Per-stage run timing is no longer returned. The platform emits it as * `[STRACE]` log markers (host stages + the AICPU device-phase breakdown, * gated on SIMPLER_HOST_STRACE) — parse with simpler_setup.tools.strace_timing. @@ -286,6 +295,30 @@ int simpler_run( DeviceContextHandle ctx, RuntimeHandle runtime, int32_t callable_id, const void *args, const CallConfig *config ); +/** Whether this runtime implements the split prepare/execute contract. */ +int supports_prepared_run_ctx(DeviceContextHandle ctx); + +/** + * Construct and fully bind one unpublished run in `runtime`. + * + * On success the Runtime object remains live until exactly one matching + * `simpler_execute_prepared` or `simpler_abort_prepared` call. No device + * scheduler work is published by this entry point. + */ +int simpler_prepare_run( + DeviceContextHandle ctx, RuntimeHandle runtime, int32_t callable_id, const void *args, const CallConfig *config, + const PreparedRunIdentity *identity +); + +/** Publish, execute, validate, and destroy a successfully prepared Runtime. */ +int simpler_execute_prepared( + DeviceContextHandle ctx, RuntimeHandle runtime, int32_t callable_id, const CallConfig *config, + const PreparedRunIdentity *identity +); + +/** Reclaim an unpublished prepared Runtime without launching device work. */ +int simpler_abort_prepared(DeviceContextHandle ctx, RuntimeHandle runtime, const PreparedRunIdentity *identity); + /** Select the per-run/exec-handle slot used by the next synchronous run. */ int select_pipeline_slot_ctx(DeviceContextHandle ctx, uint32_t slot_id); diff --git a/tests/st/a2a3/host_build_graph/worker_async_fifo/kernels/orchestration/pipelined_vector_orch.cpp b/tests/st/a2a3/host_build_graph/worker_async_fifo/kernels/orchestration/pipelined_vector_orch.cpp new file mode 100644 index 000000000..1345b754e --- /dev/null +++ b/tests/st/a2a3/host_build_graph/worker_async_fifo/kernels/orchestration/pipelined_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 = 512; + +} // 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_fifo/test_worker_async_fifo.py b/tests/st/a2a3/host_build_graph/worker_async_fifo/test_worker_async_fifo.py index 4f1890a8e..c76689d8c 100644 --- 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 @@ -15,6 +15,7 @@ block before its graph callback while both slots are admitted. """ +import ctypes import multiprocessing import threading import time @@ -23,12 +24,14 @@ import pytest import torch from simpler.task_interface import ArgDirection as D +from simpler.worker import _OFF_ACCEPTED, _TASK_ACCEPTED, 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" _SIZE = 128 * 128 +_CHAIN_LENGTH = 512 _SUB_ENTERED = multiprocessing.Event() _SUB_RELEASE = multiprocessing.Event() @@ -48,7 +51,7 @@ class TestWorkerAsyncWholeRunFifo(SceneTestCase): { "name": "vector", "orchestration": { - "source": f"{_VECTOR_KERNELS}/orchestration/example_orch.cpp", + "source": "kernels/orchestration/pipelined_vector_orch.cpp", "function_name": "aicpu_orchestration_entry", "signature": [D.IN, D.IN, D.OUT], }, @@ -65,12 +68,6 @@ class TestWorkerAsyncWholeRunFifo(SceneTestCase): "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}, @@ -129,11 +126,16 @@ def submit_vector(orch, a, b, out, *, hold_open=False): ) 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) + 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)) + accepted_addrs = [ + mailbox_addr + (index + 1) * MAILBOX_FRAME_SIZE + _OFF_ACCEPTED for index in range(2) + ] deadline = time.monotonic() + 10.0 - while not torch.allclose(first_out, first_expected) and time.monotonic() < deadline: + while not any(_mailbox_load_i32(addr) == _TASK_ACCEPTED for addr in accepted_addrs): + assert time.monotonic() < deadline, "the first run was not published to the NPU" 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() @@ -147,6 +149,12 @@ def second_graph(orch, _args, _cfg): "the prepared run dispatched before the active run ended" ) + first_expected = first_a + first_b + _CHAIN_LENGTH + 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" + def third_graph(_orch, _args, _cfg): third_callback.set() @@ -160,7 +168,7 @@ def third_graph(_orch, _args, _cfg): 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) + second_expected = second_a + second_b + _CHAIN_LENGTH assert torch.allclose(second_out, second_expected), "the prepared run did not execute correctly on the NPU" submitter.join(10.0) diff --git a/tests/ut/cpp/hierarchical/test_scheduler.cpp b/tests/ut/cpp/hierarchical/test_scheduler.cpp index 6a7511376..738ae87b1 100644 --- a/tests/ut/cpp/hierarchical/test_scheduler.cpp +++ b/tests/ut/cpp/hierarchical/test_scheduler.cpp @@ -299,6 +299,7 @@ class PreparedActivationEndpoint final : public WorkerEndpoint { const WorkerEndpointCaps &caps() const override { return caps_; } WorkerCompletion run(Ring *, const WorkerDispatch &dispatch) override { + normal_runs_.fetch_add(1, std::memory_order_relaxed); normal_running_.store(true, std::memory_order_release); std::unique_lock lk(mu_); cv_.wait(lk, [this] { @@ -321,6 +322,7 @@ class PreparedActivationEndpoint final : public WorkerEndpoint { } bool normal_running() const { return normal_running_.load(std::memory_order_acquire); } + int normal_runs() const { return normal_runs_.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() { @@ -331,6 +333,7 @@ class PreparedActivationEndpoint final : public WorkerEndpoint { private: WorkerEndpointCaps caps_; + std::atomic normal_runs_{0}; std::atomic normal_running_{false}; std::atomic backend_ready_{false}; std::atomic activated_{false}; @@ -704,6 +707,90 @@ TEST(SchedulerPreparedRunTest, SuccessorPreparesButActivatesOnlyAfterFifoPromoti allocator.shutdown(); } +TEST(SchedulerPreparedRunTest, DiagnosticSuccessorUsesSequentialCompatibilityPath) { + TensorMap tensor_map; + Ring allocator; + Scope scope; + ReadyQueue ready_sub; + NextLevelReadyQueues ready_next; + Orchestrator orchestrator; + WorkerManager manager; + Scheduler scheduler; + CallConfig normal_config; + CallConfig diagnostic_config; + diagnostic_config.enable_pmu = 1; + std::strcpy(diagnostic_config.output_prefix, "/tmp/simpler-diagnostic-successor"); + 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), normal_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), diagnostic_config, 0); + orchestrator.close_run_submission(second_run); + + auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(3); + while (!endpoint_ptr->normal_running() && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_TRUE(endpoint_ptr->normal_running()); + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + EXPECT_FALSE(endpoint_ptr->backend_ready()); + + endpoint_ptr->release_normal(); + orchestrator.wait_run(first_run); + orchestrator.wait_run(second_run); + EXPECT_EQ(endpoint_ptr->normal_runs(), 2); + EXPECT_FALSE(endpoint_ptr->backend_ready()); + EXPECT_FALSE(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(); diff --git a/tests/ut/py/test_worker/test_host_worker.py b/tests/ut/py/test_worker/test_host_worker.py index d208f2c30..1f38d1edf 100644 --- a/tests/ut/py/test_worker/test_host_worker.py +++ b/tests/ut/py/test_worker/test_host_worker.py @@ -194,11 +194,13 @@ def test_two_frame_chip_loop_accepts_b_while_a_executes_sequentially(): active = 0 max_active = 0 calls = 0 + native_acceptance_args = [] call_lock = threading.Lock() class FakeImpl: def run_from_blob(self, *_args): nonlocal active, max_active, calls + native_acceptance_args.append((_args[4], _args[5])) with call_lock: index = calls calls += 1 @@ -239,10 +241,10 @@ class FakeChipWorker: 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) + worker_mod._PIPELINE_LEASE_FMT.pack_into(frame, worker_mod._OFF_PIPELINE_LEASE, 0, 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_SLOT_ID, 0) struct.pack_into("=Q", frame, worker_mod._OFF_FRAME_GENERATION, 11) struct.pack_into("=Q", frame, worker_mod._OFF_FRAME_DISPATCH_ID, index + 1) @@ -269,6 +271,7 @@ class FakeChipWorker: assert time.monotonic() < deadline time.sleep(0.001) assert max_active == 1 + assert all(ptr != 0 and value == worker_mod._TASK_ACCEPTED for ptr, value in native_acceptance_args) _mailbox_store_i32(mailbox_addr + _OFF_STATE, worker_mod._SHUTDOWN) loop_thread.join(5.0) @@ -285,6 +288,111 @@ class FakeChipWorker: shm.unlink() +def test_two_frame_chip_loop_prepares_b_while_a_executes_and_waits_for_activation(): + run_entered = threading.Event() + release_run = threading.Event() + prepare_called = threading.Event() + execute_called = threading.Event() + abort_called = threading.Event() + + class FakeImpl: + supports_prepared_runs = True + + def run_from_blob(self, *_args): + run_entered.set() + assert release_run.wait(5.0) + + def prepare_from_blob(self, *_args): + assert run_entered.is_set() + assert not release_run.is_set() + prepare_called.set() + + def execute_prepared(self, *_args): + execute_called.set() + + def abort_prepared(self, *_args): + abort_called.set() + + 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([0x43]) * 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", + "chip_runtime": "host_build_graph", + "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 + index) + 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 run_entered.wait(5.0) + + _mailbox_store_i32(frame1_addr + _OFF_STATE, worker_mod._PREPARE_READY) + assert prepare_called.wait(5.0) + deadline = time.monotonic() + 5.0 + while _mailbox_load_i32(frame1_addr + _OFF_STATE) != worker_mod._BACKEND_READY: + assert time.monotonic() < deadline + time.sleep(0.001) + assert not execute_called.is_set() + + release_run.set() + deadline = time.monotonic() + 5.0 + while _mailbox_load_i32(frame0_addr + _OFF_STATE) != worker_mod._TASK_DONE: + assert time.monotonic() < deadline + time.sleep(0.001) + assert not execute_called.is_set() + + _mailbox_store_i32(frame1_addr + _OFF_STATE, worker_mod._ACTIVATE) + assert execute_called.wait(5.0) + 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 not abort_called.is_set() + + _mailbox_store_i32(mailbox_addr + _OFF_STATE, worker_mod._SHUTDOWN) + loop_thread.join(5.0) + assert not loop_thread.is_alive() + finally: + release_run.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: descriptor = build_chip_callable_descriptor(target=callable_obj, platform=platform, runtime=runtime) return hashid_to_digest(compute_callable_hashid(descriptor))