From 6b3904277b9b58bf1a0853483b76292379ab0fab Mon Sep 17 00:00:00 2001 From: Crane-Liu Date: Mon, 27 Jul 2026 22:55:58 +0800 Subject: [PATCH 1/5] 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/5] 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/5] 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/5] 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)) From 3a6d649475ba4709f215607a40a1c0fcb568b3b0 Mon Sep 17 00:00:00 2001 From: Crane-Liu Date: Wed, 29 Jul 2026 20:57:13 +0800 Subject: [PATCH 5/5] Add: TMR prepared state pipeline --- docs/task-flow.md | 39 ++++++---- python/bindings/task_interface.cpp | 10 +-- python/simpler/worker.py | 26 +++++-- .../platform/onboard/host/device_runner.cpp | 35 ++++----- .../platform/onboard/host/device_runner.h | 11 ++- .../docs/RUNTIME_LOGIC.md | 19 +++++ .../host/runtime_maker.cpp | 32 +++++++++ src/common/hierarchical/scheduler.cpp | 64 ++++++++++++++--- src/common/hierarchical/scheduler.h | 2 + src/common/hierarchical/types.cpp | 4 ++ src/common/hierarchical/types.h | 1 + src/common/hierarchical/worker.cpp | 1 + src/common/hierarchical/worker_manager.cpp | 22 ++++-- src/common/hierarchical/worker_manager.h | 3 + .../platform/include/host/run_stream_slots.h | 71 ++++++++++++++----- .../platform/onboard/host/c_api_shared.cpp | 21 ++++++ .../onboard/host/device_runner_base.cpp | 5 ++ .../onboard/host/device_runner_base.h | 1 + src/common/worker/chip_worker.cpp | 21 +++++- src/common/worker/chip_worker.h | 8 +-- src/common/worker/pto_runtime_c_api.h | 8 ++- .../run_stream_reuse/test_run_stream_reuse.py | 51 ++++++------- .../orchestration/pipelined_vector_orch.cpp | 2 +- .../test_worker_async_fifo.py | 34 +++++---- tests/ut/cpp/CMakeLists.txt | 6 +- .../common/test_trb_runtime_temp_buffer.cpp | 53 ++++++++++++++ .../hierarchical/test_run_stream_slots.cpp | 54 ++++++++------ tests/ut/cpp/hierarchical/test_scheduler.cpp | 69 ++++++++++++++---- tests/ut/py/test_worker/test_host_worker.py | 22 ++++-- 29 files changed, 520 insertions(+), 175 deletions(-) diff --git a/docs/task-flow.md b/docs/task-flow.md index 7a250ea78..34d554ebd 100644 --- a/docs/task-flow.md +++ b/docs/task-flow.md @@ -284,14 +284,13 @@ that window needs the admission layer to gate dispatch on `pool.owns(lease)` before handing work down, which is whole-run admission's job, not this layer's. -AICore streams are outside that lease. The instruction cache belongs to the -cores, every slot publishes its image to the same GM code address, and the -platform offers no cache invalidation for code replaced there. Creating a -stream is the only operation known to leave a core free of the previous image's -instructions; selecting an already-existing one is not. So each run creates its -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. +AICore streams carry a hardware-generation key. The instruction cache belongs +to the cores and the platform offers no cache invalidation for code replaced at +a reused GM address, so a stream is reused only when the next run has the same +`aicore_image_hash`. An image transition destroys the prior stream and creates +a fresh one before launch. A destroy failure poisons that slot and preserves +the handle for finalize to retry; an error run retires its stream +conservatively. The AICPU stream carries no such state and stays with its slot. #### Whole-run FIFO admission @@ -350,14 +349,28 @@ 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. +The scheduler also waits for the active run's native sticky launch ACK before +staging its successor. `TASK_ACCEPTED` and `TASK_ACTIVE` in the frame state only +report child-executor admission; neither proves that the AICore and AICPU +kernels were enqueued. The native runner writes the separate sticky ACK after +both launches, for normal and prepared execution, and only that signal opens +the next prepare window. Terminal fallback still drains the public run fence, +but does not masquerade as a chip-launch ACK. + 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`. +collectors cannot overlap Host preparation. Remote, SUB, A5, sim, single-frame, +and endpoints without the capability remain unchanged. + +Onboard a2a3 HBG builds into the lease-selected arena bank, reports +`BACKEND_READY` without publishing device work, and launches only after the +matching `ACTIVATE`. Onboard a2a3 TMR stages the next run's slot-private tensor +arguments and Host descriptor while sharing arena bank 0. Its prepare path must +hit the prebuilt-arena cache for the fully resolved `RuntimeEnv` sizing key; a +miss returns `PTO_RUNTIME_ERR_PREPARED_INCOMPATIBLE` instead of rebuilding the +active run's shared heap, SM, or runtime image. Activation then publishes the +already-compatible descriptor after the current run's completion fence. 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 cae811638..61adbb689 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -1504,15 +1504,17 @@ NB_MODULE(_task_interface, m) { .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) { + uint64_t pipeline_generation, uint64_t run_id, uint64_t dispatch_id, uint64_t accepted_state_addr, + int32_t accepted_value) { nb::gil_scoped_release release; self.execute_prepared( - callable_id, config, PipelineSlotLease{pipeline_slot, 0, pipeline_generation}, run_id, dispatch_id + callable_id, config, PipelineSlotLease{pipeline_slot, 0, pipeline_generation}, run_id, dispatch_id, + reinterpret_cast(accepted_state_addr), accepted_value ); }, 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." + nb::arg("run_id"), nb::arg("dispatch_id"), nb::arg("accepted_state_addr") = 0, + nb::arg("accepted_value") = 0, "Publish and complete a matching prepared run on the sequential executor." ) .def( "abort_prepared", diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 296fa9fb4..54b2818d2 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -231,10 +231,11 @@ def my_l4_orch(orch, args, config): _PREPARE_READY = 13 _TASK_FRAME_COUNT = 2 _PIPELINE_CAP_PREPARE_ACTIVATE = 1 << 0 +_A2A3_PREPARED_RUNTIMES = frozenset(("host_build_graph", "tensormap_and_ringbuffer")) 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: + if platform == "a2a3" and runtime in _A2A3_PREPARED_RUNTIMES and pipeline_depth >= 2: return _TASK_FRAME_COUNT return 1 @@ -1807,6 +1808,7 @@ def resolve_prepared_frame( def execute_prepared_frame( frame_buf: memoryview, + frame_addr: int, identity: tuple[int, int, int, int, int], ) -> tuple[int, str]: _protocol, run_id, slot_id, generation, dispatch_id = identity @@ -1817,7 +1819,14 @@ def execute_prepared_frame( 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) + int(cid), + cfg, + int(slot_id), + int(generation), + int(run_id), + int(dispatch_id), + frame_addr + _OFF_ACCEPTED, + _TASK_ACCEPTED, ) except Exception as e: # noqa: BLE001 return 1, _format_exc(f"chip_process dev={device_id} execute prepared", e) @@ -1958,7 +1967,7 @@ def admission_loop() -> None: _mailbox_store_i32(frame_state_addr, _TASK_ACTIVE) if kind == "prepared": prepared_frames.pop(index, None) - code, msg = execute_prepared_frame(frame_buf, identity) + code, msg = execute_prepared_frame(frame_buf, frame_addr, identity) else: code, msg = handle_task( frame_buf, @@ -2018,6 +2027,11 @@ def _chip_process_loop( # noqa: PLR0913 -- fork-child entry: all context (bins, import traceback as _tb # noqa: PLC0415 try: + if prewarm_config is None and platform == "a2a3" and runtime == "tensormap_and_ringbuffer": + # A prepared successor may reach this child before the first active + # run has built the shared arena. Establish the default key before + # READY so prepare is read-only from the first dispatch onward. + prewarm_config = CallConfig() cw = ChipWorker() cw.init( device_id, @@ -4545,8 +4559,10 @@ def init(self, prewarm_config: CallConfig | None = None, *, _startup_deadline: f same sizing skips the (~800ms) cold prebuilt runtime-arena build. An L2 worker prewarms here; an L3+ worker prewarms each chip child during hierarchy startup, before it publishes INIT_READY. A no-op - for runtimes without a prebuilt arena (host_build_graph). ``None`` - (default) disables prewarm. + for runtimes without a prebuilt arena (host_build_graph). For an + L3+ A2A3 TMR worker, ``None`` prewarms the default ``CallConfig`` + because prepared dispatch requires the shared arena to exist + before the first active run; elsewhere ``None`` disables prewarm. _startup_deadline: Internal. Absolute ``time.monotonic()`` deadline inherited from a parent's startup epoch so a recursive descendant consumes the parent's remaining budget instead of restarting the diff --git a/src/a2a3/platform/onboard/host/device_runner.cpp b/src/a2a3/platform/onboard/host/device_runner.cpp index 981a21358..e66a87631 100644 --- a/src/a2a3/platform/onboard/host/device_runner.cpp +++ b/src/a2a3/platform/onboard/host/device_runner.cpp @@ -248,22 +248,17 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { return rc; } - rc = ensure_run_stream_set(selected_pipeline_slot); + const uint64_t aicore_image_hash = callable_aicore_image_hash(runtime.get_active_callable_id()); + rc = ensure_run_stream_set(selected_pipeline_slot, aicore_image_hash); if (rc != 0) { - LOG_ERROR("ensure_run_stream_set(%u) failed: %d", selected_pipeline_slot, rc); + LOG_ERROR("ensure_run_stream_set(%u, image=0x%lx) failed: %d", selected_pipeline_slot, aicore_image_hash, rc); return rc; } - // The AICore stream is this run's alone: creation is the only operation - // this platform offers that is known to leave a core free of a previous - // image's cached instructions, and there is no evidence that selecting an - // already-existing stream does anything to the instruction cache. Retiring - // it on every exit path is what keeps that guarantee from depending on - // which image the next run happens to publish. - // On an early return the guard retires it; the success path below retires - // it explicitly so a failed destroy is reported instead of discarded. - bool aicore_stream_retired = false; - auto aicore_stream_retire = RAIIScopeGuard([this, selected_pipeline_slot, &aicore_stream_retired]() { - if (!aicore_stream_retired) (void)retire_run_aicore_stream(selected_pipeline_slot); + // A successfully reaped stream remains reusable for this exact image. Any + // error path retires it conservatively; a failed destroy poisons the slot. + bool run_reaped = false; + auto aicore_stream_retire = RAIIScopeGuard([this, selected_pipeline_slot, &run_reaped]() { + if (!run_reaped) (void)retire_run_aicore_stream(selected_pipeline_slot); }); ensure_device_wall_buffer(); @@ -485,13 +480,7 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { rc = reap_run(selected_pipeline_slot); if (rc != 0) return rc; - - // The run owns its AICore stream, so a destroy this run cannot complete is - // this run's failure: reporting success would leave the caller believing a - // slot is reusable that ensure_run_stream_set will now refuse. - aicore_stream_retired = true; - rc = retire_run_aicore_stream(selected_pipeline_slot); - if (rc != 0) return rc; + run_reaped = true; // Print handshake results (reads from device memory, must be before free) print_handshake_results(); @@ -499,10 +488,10 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { return 0; } -int DeviceRunner::ensure_run_stream_set(unsigned slot) { - int rc = run_stream_slots_.acquire(slot); +int DeviceRunner::ensure_run_stream_set(unsigned slot, uint64_t aicore_image_hash) { + int rc = run_stream_slots_.acquire(slot, aicore_image_hash); if (rc != 0) { - LOG_ERROR("ensure_run_stream_set(%u) failed: %d", slot, rc); + LOG_ERROR("ensure_run_stream_set(%u, image=0x%lx) failed: %d", slot, aicore_image_hash, rc); ACL_LOG_ERROR_DETAIL(rc); } return rc; diff --git a/src/a2a3/platform/onboard/host/device_runner.h b/src/a2a3/platform/onboard/host/device_runner.h index e45ea0745..cc1f4d545 100644 --- a/src/a2a3/platform/onboard/host/device_runner.h +++ b/src/a2a3/platform/onboard/host/device_runner.h @@ -238,12 +238,9 @@ class DeviceRunner : public DeviceRunnerBase { static constexpr unsigned kRunStreamSetCount = PTO_PIPELINE_MAX_DEPTH; // AICPU streams belong to slots for the worker's lifetime. AICore streams - // belong to a single run: the platform offers no instruction-cache - // invalidation for code replaced at a reused GM address, and creating the - // stream is the only operation known to leave a core free of the previous - // image's instructions. Selecting an existing stream is not. - // Stream lifetimes live in RunStreamSlots so the slot state machine — - // fresh AICore stream per run, handle kept when a destroy fails — is + // are reusable only for the same aicore_image_hash; an image transition + // destroys and recreates the stream before launch. Stream lifetimes live in + // RunStreamSlots so generation matching and poison-on-destroy-failure are // testable without a device. RunStreamSlots run_stream_slots_{ [](void **out) { @@ -253,7 +250,7 @@ class DeviceRunner : public DeviceRunnerBase { return rtStreamDestroy(static_cast(stream)); } }; - int ensure_run_stream_set(unsigned slot); + int ensure_run_stream_set(unsigned slot, uint64_t aicore_image_hash); // Destroys this run's AICore stream. Returns the driver's error and KEEPS // the handle when the destroy fails: the stream may still hold the previous // image's instructions, so the slot must refuse the next run until finalize diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md index 168ec2ad5..5624a8e7b 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md @@ -110,6 +110,25 @@ The platform side is deliberately thin: `DeviceRunnerBase` only remembers a allocator; all grow/pack/slice logic lives in `runtime_maker.cpp` (`RetainedTempBump`). +Prepared whole-run admission may stage the inactive slot while the current run +uses the one shared arena bank. Before staging, the runtime resolves the full +per-ring sizing key and requires it to hit the existing prebuilt-arena cache. +An incompatible key is rejected; preparation never grows or replaces the +shared heap, SM, or runtime image underneath an active run. + +That arena key is the resolved four-ring `task_window`, `heap`, and `dep_pool` +layout after per-call overrides, environment values, and defaults are applied. +The callable's AICore image has a separate generation key on each run-stream +slot, so an image transition recreates the AICore stream without changing the +arena key. SDMA workspaces are provisioned once when the Worker is initialized +and cannot vary between adjacent runs. + +An L3 A2A3 TMR Worker prewarms the default key before its chip child publishes +READY, so the first successor cannot race the first active run's cold arena +build. Callers using non-default ring sizing pass that `CallConfig` to +`Worker.init(prewarm_config=...)`; a prepared transition to any other resolved +key is rejected explicitly. + On each trb bind, `RetainedTempBump`: - packs the run's non-child, non-empty tensors to a required size, aligning diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp index 6f5cdacb8..79afa2a09 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp @@ -80,6 +80,13 @@ extern "C" const PipelineContract *get_pipeline_contract(void) { return &contract; } +extern "C" int prepared_run_supported_impl(void) { + // TRB preparation may stage slot-private arguments while the active run is + // on device. The compatibility hook below prevents a prepared cache miss + // from rebuilding the one shared device-scratch arena. + return 1; +} + static_assert( RUNTIME_ENV_RING_COUNT == PTO2_MAX_RING_DEPTH, "RuntimeEnv ring count must match PTO2 runtime ring depth" ); @@ -525,6 +532,31 @@ static bool resolve_arena_sizing( return true; } +extern "C" int prepared_run_config_compatible_impl( + const HostApi *api, const uint64_t *ring_task_window, const uint64_t *ring_heap, const uint64_t *ring_dep_pool +) { + if (api == nullptr || api->lookup_prebuilt_runtime_arena_cache == nullptr) { + return -1; + } + + ArenaSizingConfig sizing; + if (!resolve_arena_sizing(ring_task_window, ring_heap, ring_dep_pool, &sizing)) { + return -1; + } + PrebuiltRuntimeArenaCacheProbe probe = make_prebuilt_runtime_arena_cache_probe(sizing); + void *gm_heap = nullptr; + void *gm_sm = nullptr; + void *runtime_arena = nullptr; + size_t runtime_offset = 0; + const void *image = nullptr; + size_t image_size = 0; + bool cache_hit = api->lookup_prebuilt_runtime_arena_cache( + probe.hash, probe.serialized_key.data(), probe.serialized_key.size(), &gm_heap, &gm_sm, &runtime_arena, + &runtime_offset, &image, &image_size + ); + return cache_hit ? 1 : 0; +} + static bool derive_arena_static_sizes(const ArenaSizingConfig &sizing, ArenaStaticSizes *out) { out->total_heap = 0; for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { diff --git a/src/common/hierarchical/scheduler.cpp b/src/common/hierarchical/scheduler.cpp index 1d94168cd..0da370d55 100644 --- a/src/common/hierarchical/scheduler.cpp +++ b/src/common/hierarchical/scheduler.cpp @@ -182,11 +182,7 @@ void Scheduler::run() { ready = active != INVALID_RUN_ID && (!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); - } + if (!ready && cfg_.preparable_run_cb) ready = preparable_work_ready(); } else { ready = !cfg_.ready_next_level_queues->empty() || !cfg_.ready_sub_queue->empty(); } @@ -362,6 +358,49 @@ void Scheduler::try_consume(TaskSlot slot) { // helpers below are therefore unreachable invariant checks; because they run on // sched_thread_ with no surrounding handler, any throw is fatal to the whole // worker tree (std::terminate), not a per-task failure. +bool Scheduler::active_run_chip_dispatch_accepted(RunId active_run) const { + return active_run != INVALID_RUN_ID && cfg_.manager->any_next_level_worker_has_accepted_run(active_run); +} + +bool Scheduler::preparable_work_ready() const { + if (!cfg_.active_run_cb || !cfg_.preparable_run_cb) return false; + RunId active_run = cfg_.active_run_cb(); + RunId preparable_run = cfg_.preparable_run_cb(); + if (active_run == INVALID_RUN_ID || preparable_run == INVALID_RUN_ID || + cfg_.manager->has_staged_run(preparable_run)) { + return false; + } + + TaskSlot slot; + if (cfg_.ready_next_level_queues->try_front_group(preparable_run, slot)) { + TaskSlotState &state = *cfg_.ring->slot_state(slot); + bool ready = state.state.load(std::memory_order_acquire) == TaskState::READY && + state.run_id == preparable_run && state.worker_type == WorkerType::NEXT_LEVEL && + state.is_group() && !state.config.diagnostics_any() && + active_run_chip_dispatch_accepted(active_run); + for (int32_t index = 0; ready && index < state.group_size(); ++index) { + WorkerThread *worker = + cfg_.manager->get_worker_by_id(WorkerType::NEXT_LEVEL, state.target_worker_id(index)); + ready = worker != nullptr && worker->caps().supports_prepare_activate && worker->has_capacity(); + } + if (ready) return true; + } + + for (int32_t worker_id : cfg_.ready_next_level_queues->worker_ids()) { + if (!cfg_.ready_next_level_queues->try_front_single(worker_id, preparable_run, slot)) continue; + TaskSlotState &state = *cfg_.ring->slot_state(slot); + WorkerThread *worker = cfg_.manager->get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); + if (state.state.load(std::memory_order_acquire) == TaskState::READY && state.run_id == preparable_run && + state.worker_type == WorkerType::NEXT_LEVEL && !state.is_group() && + state.target_worker_id(0) == worker_id && !state.config.diagnostics_any() && worker != nullptr && + worker->caps().supports_prepare_activate && worker->has_capacity() && + active_run_chip_dispatch_accepted(active_run)) { + return true; + } + } + return false; +} + 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); @@ -376,6 +415,7 @@ 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; + RunId active_run = cfg_.active_run_cb ? cfg_.active_run_cb() : INVALID_RUN_ID; TaskSlot slot; if (!cfg_.ready_next_level_queues->try_front_group(run_id, slot)) return; @@ -384,6 +424,7 @@ void Scheduler::dispatch_preparable_next_level_group() { state.worker_type != WorkerType::NEXT_LEVEL || !state.is_group() || state.config.diagnostics_any()) { return; } + if (!active_run_chip_dispatch_accepted(active_run)) return; const int32_t group_size = state.group_size(); std::vector workers; @@ -409,21 +450,24 @@ 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; + RunId active_run = cfg_.active_run_cb ? cfg_.active_run_cb() : INVALID_RUN_ID; 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; + if (!cfg_.ready_next_level_queues->try_front_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; } - if (state.config.diagnostics_any()) { - cfg_.enqueue_ready_cb(slot); - continue; + if (state.config.diagnostics_any() || !active_run_chip_dispatch_accepted(active_run)) continue; + TaskSlot popped; + if (!cfg_.ready_next_level_queues->try_pop_single(worker_id, run_id, popped) || popped != slot) { + throw std::runtime_error( + "Scheduler::dispatch_preparable_next_level_singles: worker queue changed unexpectedly" + ); } state.state.store(TaskState::RUNNING, std::memory_order_release); worker->dispatch_prepared(WorkerDispatch{slot, 0}); diff --git a/src/common/hierarchical/scheduler.h b/src/common/hierarchical/scheduler.h index 3e4e0085a..d3e26461f 100644 --- a/src/common/hierarchical/scheduler.h +++ b/src/common/hierarchical/scheduler.h @@ -106,6 +106,8 @@ class Scheduler { void on_task_complete(const WorkerCompletion &completion); void poison_task(TaskSlot slot, const std::string &root_message); void try_consume(TaskSlot slot); + bool preparable_work_ready() const; + bool active_run_chip_dispatch_accepted(RunId active_run) const; void dispatch_ready(); void dispatch_preparable_next_level_group(); void dispatch_preparable_next_level_singles(); diff --git a/src/common/hierarchical/types.cpp b/src/common/hierarchical/types.cpp index 288f507d5..e83327adf 100644 --- a/src/common/hierarchical/types.cpp +++ b/src/common/hierarchical/types.cpp @@ -175,6 +175,10 @@ bool NextLevelReadyQueues::try_pop_single(int32_t worker_id, RunId run_id, TaskS return queues_[index_for(worker_id)]->try_pop(run_id, out); } +bool NextLevelReadyQueues::try_front_single(int32_t worker_id, RunId run_id, TaskSlot &out) { + return queues_[index_for(worker_id)]->try_front(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); } diff --git a/src/common/hierarchical/types.h b/src/common/hierarchical/types.h index 84b5dde7c..7a9c0e09e 100644 --- a/src/common/hierarchical/types.h +++ b/src/common/hierarchical/types.h @@ -461,6 +461,7 @@ class NextLevelReadyQueues { 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); + bool try_front_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); diff --git a/src/common/hierarchical/worker.cpp b/src/common/hierarchical/worker.cpp index 79cb885af..429f244a2 100644 --- a/src/common/hierarchical/worker.cpp +++ b/src/common/hierarchical/worker.cpp @@ -110,6 +110,7 @@ void Worker::init() { }, [this](WorkerDispatch dispatch) { orchestrator_.mark_task_accepted(dispatch.task_slot); + scheduler_.notify_ready(); } ); ready_next_level_queues_.reset(manager_.next_level_worker_ids()); diff --git a/src/common/hierarchical/worker_manager.cpp b/src/common/hierarchical/worker_manager.cpp index d05f42ffc..45b7f4548 100644 --- a/src/common/hierarchical/worker_manager.cpp +++ b/src/common/hierarchical/worker_manager.cpp @@ -368,13 +368,19 @@ void WorkerThread::loop() { WorkerCompletion completion; bool accepted = false; - auto accept_once = [&]() { + auto accept_once = [&](bool chip_dispatch_accepted) { if (accepted) return; accepted = true; + TaskSlotState *slot = ring_ == nullptr ? nullptr : ring_->slot_state(d.task_slot); + if (chip_dispatch_accepted && slot != nullptr) { + accepted_run_id_.store(slot->run_id, std::memory_order_release); + } if (on_accept_) on_accept_(d); }; try { - completion = dispatch_process(d, accept_once); + completion = dispatch_process(d, [&] { + accept_once(true); + }); } catch (const std::exception &e) { completion.task_slot = d.task_slot; completion.group_index = d.group_index; @@ -392,7 +398,7 @@ void WorkerThread::loop() { // process. mark_task_accepted is non-throwing by construction; this // keeps that a local property rather than a remote assumption. try { - accept_once(); + accept_once(false); } catch (const std::exception &e) { if (completion.outcome == EndpointOutcome::SUCCESS) { completion.outcome = EndpointOutcome::ENDPOINT_FAILURE; @@ -816,8 +822,7 @@ WorkerCompletion LocalMailboxEndpoint::run_two_frame( } 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 (!acceptance_observed && read_task_accepted(frame)) { if (!identity_matches()) { endpoint_poisoned_.store(true, std::memory_order_release); publish_cv_.notify_all(); @@ -1437,6 +1442,13 @@ bool WorkerManager::has_staged_run(RunId run_id) const { return false; } +bool WorkerManager::any_next_level_worker_has_accepted_run(RunId run_id) const { + for (const auto &worker : next_level_threads_) { + if (worker->has_accepted_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; diff --git a/src/common/hierarchical/worker_manager.h b/src/common/hierarchical/worker_manager.h index f4230fad5..f124cb5d2 100644 --- a/src/common/hierarchical/worker_manager.h +++ b/src/common/hierarchical/worker_manager.h @@ -443,6 +443,7 @@ class WorkerThread { 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 has_accepted_run(RunId run_id) const { return accepted_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; } @@ -544,6 +545,7 @@ class WorkerThread { std::atomic inflight_{0}; std::atomic next_dispatch_id_{1}; std::atomic staged_run_id_{INVALID_RUN_ID}; + std::atomic accepted_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_; @@ -594,6 +596,7 @@ class WorkerManager { bool any_busy() const; bool has_staged_run(RunId run_id) const; + bool any_next_level_worker_has_accepted_run(RunId run_id) const; bool needs_activation(RunId run_id) const; bool activate_prepared_run(RunId run_id); diff --git a/src/common/platform/include/host/run_stream_slots.h b/src/common/platform/include/host/run_stream_slots.h index 194a16474..d70824ef8 100644 --- a/src/common/platform/include/host/run_stream_slots.h +++ b/src/common/platform/include/host/run_stream_slots.h @@ -22,11 +22,11 @@ * Per-slot ownership of the two streams a run submits on. * * The AICPU stream carries no instruction-cache state, so it belongs to its - * slot for the runner's lifetime. The AICore stream belongs to a single run: - * the platform offers no instruction-cache invalidation for code replaced at a - * reused GM address, and creating the stream is the only operation known to - * leave a core free of the previous image's instructions — selecting an - * existing one is not. + * slot for the runner's lifetime. A successfully reaped AICore stream may be + * reused only for the same loaded code image. A different image first destroys + * the old stream, because creating a stream is the only operation known to + * leave a core free of instructions cached from code replaced at a reused GM + * address. * * A destroy that fails keeps its handle. Such a stream may still hold the * previous image's instructions, so the slot must refuse the next run rather @@ -44,13 +44,14 @@ class RunStreamSlots { destroy_(std::move(destroy)) {} /** - * Ready `slot` for a run: its AICPU stream on first use, and always a fresh - * AICore stream. Fails when the slot still holds an AICore stream a prior - * run could not retire. + * Ready `slot` for a run. The AICPU stream is created on first use. The + * AICore stream is reused for `aicore_image_hash`, or replaced before a + * different image runs. Fails when a prior stream could not be destroyed. */ - int acquire(unsigned slot) { + int acquire(unsigned slot, uint64_t aicore_image_hash) { if (slot >= slots_.size()) return -1; Slot &s = slots_[slot]; + if (s.poisoned) return -1; if (s.aicpu == nullptr) { int rc = create_(&s.aicpu); if (rc != 0) { @@ -58,12 +59,25 @@ class RunStreamSlots { return rc; } } - if (s.aicore != nullptr) return -1; + if (s.aicore != nullptr && s.image_valid && s.aicore_image_hash == aicore_image_hash) { + return 0; + } + if (s.aicore != nullptr) { + int rc = destroy_(s.aicore); + if (rc != 0) { + s.poisoned = true; + return rc; + } + s.aicore = nullptr; + s.image_valid = false; + } int rc = create_(&s.aicore); if (rc != 0) { s.aicore = nullptr; return rc; } + s.aicore_image_hash = aicore_image_hash; + s.image_valid = true; ++created_count_; return 0; } @@ -74,8 +88,14 @@ class RunStreamSlots { Slot &s = slots_[slot]; if (s.aicore == nullptr) return 0; int rc = destroy_(s.aicore); - if (rc != 0) return rc; + if (rc != 0) { + s.poisoned = true; + return rc; + } s.aicore = nullptr; + s.aicore_image_hash = 0; + s.image_valid = false; + s.poisoned = false; return 0; } @@ -83,14 +103,25 @@ class RunStreamSlots { int destroy_all() { int first_error = 0; for (Slot &s : slots_) { - for (void **stream : {&s.aicpu, &s.aicore}) { - if (*stream == nullptr) continue; - int rc = destroy_(*stream); + if (s.aicpu != nullptr) { + int rc = destroy_(s.aicpu); if (rc != 0) { if (first_error == 0) first_error = rc; - continue; + } else { + s.aicpu = nullptr; + } + } + if (s.aicore != nullptr) { + int rc = destroy_(s.aicore); + if (rc != 0) { + if (first_error == 0) first_error = rc; + s.poisoned = true; + } else { + s.aicore = nullptr; + s.aicore_image_hash = 0; + s.image_valid = false; + s.poisoned = false; } - *stream = nullptr; } } return first_error; @@ -98,7 +129,10 @@ class RunStreamSlots { void *aicpu(unsigned slot) const { return slot < slots_.size() ? slots_[slot].aicpu : nullptr; } void *aicore(unsigned slot) const { return slot < slots_.size() ? slots_[slot].aicore : nullptr; } - bool ready(unsigned slot) const { return aicpu(slot) != nullptr && aicore(slot) != nullptr; } + bool ready(unsigned slot) const { + return slot < slots_.size() && !slots_[slot].poisoned && slots_[slot].aicpu != nullptr && + slots_[slot].aicore != nullptr && slots_[slot].image_valid; + } size_t created_count() const { return created_count_; } static constexpr size_t capacity() { return PTO_PIPELINE_MAX_DEPTH; } @@ -106,6 +140,9 @@ class RunStreamSlots { struct Slot { void *aicpu{nullptr}; void *aicore{nullptr}; + uint64_t aicore_image_hash{0}; + bool image_valid{false}; + bool poisoned{false}; }; CreateFn create_; diff --git a/src/common/platform/onboard/host/c_api_shared.cpp b/src/common/platform/onboard/host/c_api_shared.cpp index 670452106..008d3aca9 100644 --- a/src/common/platform/onboard/host/c_api_shared.cpp +++ b/src/common/platform/onboard/host/c_api_shared.cpp @@ -62,6 +62,12 @@ 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; } +__attribute__((weak)) int prepared_run_config_compatible_impl( + const HostApi * /*api*/, const uint64_t * /*ring_task_window*/, const uint64_t * /*ring_heap*/, + const uint64_t * /*ring_dep_pool*/ +) { + return 1; +} /* =========================================================================== * Per-thread DeviceRunnerBase binding (set by simpler_register_callable / simpler_run) @@ -716,6 +722,21 @@ int simpler_prepare_run( try { int rc = runner->attach_current_thread(runner->device_id()); if (rc != 0) return rc; + int compatibility_rc = 0; + { + STRACE_A("simpler_prepare.compatibility", attrs); + compatibility_rc = prepared_run_config_compatible_impl( + &g_host_api, config->runtime_env.ring_task_window, config->runtime_env.ring_heap, + config->runtime_env.ring_dep_pool + ); + } + if (compatibility_rc <= 0) { + if (compatibility_rc == 0) { + LOG_ERROR("prepared run RuntimeEnv is incompatible with the active device-scratch image"); + return PTO_RUNTIME_ERR_PREPARED_INCOMPATIBLE; + } + return compatibility_rc; + } Runtime *r = new (runtime) Runtime(); bool cleanup_validated = false; auto runtime_guard = RAIIScopeGuard([&]() { diff --git a/src/common/platform/onboard/host/device_runner_base.cpp b/src/common/platform/onboard/host/device_runner_base.cpp index 04315b9a4..cc0f809bc 100644 --- a/src/common/platform/onboard/host/device_runner_base.cpp +++ b/src/common/platform/onboard/host/device_runner_base.cpp @@ -960,6 +960,11 @@ uint64_t DeviceRunnerBase::callable_hash(int32_t callable_id) const { return it == callables_.end() ? 0 : it->second.hash; } +uint64_t DeviceRunnerBase::callable_aicore_image_hash(int32_t callable_id) const { + auto it = callables_.find(callable_id); + return it == callables_.end() ? 0 : it->second.aicore_image_hash; +} + // Per-run binding half, defined in each runtime's runtime_maker.cpp and linked // into this same host_runtime.so. Declared here (rather than only in // c_api_shared.cpp) so bind_callable_to_runtime can call it directly, keeping diff --git a/src/common/platform/onboard/host/device_runner_base.h b/src/common/platform/onboard/host/device_runner_base.h index b9df44636..bef1f4db7 100644 --- a/src/common/platform/onboard/host/device_runner_base.h +++ b/src/common/platform/onboard/host/device_runner_base.h @@ -394,6 +394,7 @@ class DeviceRunnerBase { * per-stage timing to a specific callable. */ uint64_t callable_hash(int32_t callable_id) const; + uint64_t callable_aicore_image_hash(int32_t callable_id) const; /** * Publish this run's core geometry onto `Runtime` before the graph is diff --git a/src/common/worker/chip_worker.cpp b/src/common/worker/chip_worker.cpp index c53ed815d..88a6aaf75 100644 --- a/src/common/worker/chip_worker.cpp +++ b/src/common/worker/chip_worker.cpp @@ -583,6 +583,9 @@ void ChipWorker::prepare_with_lease( if (rc != 0) { std::scoped_lock lk(prepared_slots_mu_); prepared_slots_[lease.slot_id] = PreparedSlot{}; + if (rc == PTO_RUNTIME_ERR_PREPARED_INCOMPATIBLE) { + throw std::runtime_error("prepared run RuntimeEnv is incompatible with the active device-scratch image"); + } throw std::runtime_error("prepare run failed with code " + std::to_string(rc)); } @@ -605,7 +608,8 @@ void ChipWorker::prepare_with_lease( } void ChipWorker::execute_prepared( - int32_t callable_id, const CallConfig &config, const PipelineSlotLease &lease, uint64_t run_id, uint64_t dispatch_id + int32_t callable_id, const CallConfig &config, const PipelineSlotLease &lease, uint64_t run_id, + uint64_t dispatch_id, volatile int32_t *accepted_state, int32_t accepted_value ) { config.validate(); if (config.diagnostics_any()) { @@ -630,14 +634,29 @@ void ChipWorker::execute_prepared( } int rc = -1; + if (accepted_state != nullptr && set_task_accepted_state_fn_ != nullptr) { + int bind_rc = set_task_accepted_state_fn_(device_ctx_, accepted_state, accepted_value); + if (bind_rc != 0) { + std::scoped_lock lk(prepared_slots_mu_); + prepared_slots_[lease.slot_id] = PreparedSlot{}; + throw std::runtime_error("set_task_accepted_state_ctx failed with code " + std::to_string(bind_rc)); + } + } + auto clear_accepted_state = [&]() { + if (accepted_state != nullptr && set_task_accepted_state_fn_ != nullptr) { + (void)set_task_accepted_state_fn_(device_ctx_, nullptr, 0); + } + }; 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 (...) { + clear_accepted_state(); std::scoped_lock lk(prepared_slots_mu_); prepared_slots_[lease.slot_id] = PreparedSlot{}; throw; } + clear_accepted_state(); { std::scoped_lock lk(prepared_slots_mu_); prepared_slots_[lease.slot_id] = PreparedSlot{}; diff --git a/src/common/worker/chip_worker.h b/src/common/worker/chip_worker.h index 1b82e0808..390aa3a5a 100644 --- a/src/common/worker/chip_worker.h +++ b/src/common/worker/chip_worker.h @@ -94,7 +94,7 @@ class ChipWorker { /// 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 + uint64_t dispatch_id, volatile int32_t *accepted_state = nullptr, int32_t accepted_value = 0 ); /// Reclaim an unpublished prepared run after a local protocol failure. void abort_prepared(const PipelineSlotLease &lease, uint64_t run_id, uint64_t dispatch_id); @@ -121,9 +121,9 @@ class ChipWorker { size_t host_dlopen_count() const; /// Number of AICore run streams the bound runner has created. AICPU streams - /// belong to slots for the worker's lifetime; each run gets a freshly - /// created AICore stream, so this advances once per run. Platforms using - /// the persistent bootstrap pair report 0. + /// belong to slots for the worker's lifetime; each slot retains its AICore + /// stream while the loaded image is unchanged and recreates it on an image + /// transition. Platforms using the persistent bootstrap pair report 0. size_t run_stream_set_create_count() const; uint64_t malloc(size_t size); diff --git a/src/common/worker/pto_runtime_c_api.h b/src/common/worker/pto_runtime_c_api.h index b77586831..349730f80 100644 --- a/src/common/worker/pto_runtime_c_api.h +++ b/src/common/worker/pto_runtime_c_api.h @@ -62,6 +62,7 @@ typedef void *DeviceContextHandle; enum { PTO_RUNTIME_ERR_UNSUPPORTED = -2, + PTO_RUNTIME_ERR_PREPARED_INCOMPATIBLE = -3, }; enum { @@ -383,9 +384,10 @@ size_t get_host_dlopen_count(DeviceContextHandle ctx); /** * Number of AICore run streams the runner bound to `ctx` has created. AICPU - * streams belong to pipeline slots for the worker's lifetime; each run gets a - * freshly created AICore stream, so this advances once per run. Returns 0 on - * platforms whose runs use the persistent bootstrap pair. + * streams belong to pipeline slots for the worker's lifetime; each slot retains + * its AICore stream while the loaded image is unchanged and recreates it on an + * image transition. Returns 0 on platforms whose runs use the persistent + * bootstrap pair. */ size_t get_run_stream_set_create_count(DeviceContextHandle ctx); diff --git a/tests/st/a2a3/host_build_graph/run_stream_reuse/test_run_stream_reuse.py b/tests/st/a2a3/host_build_graph/run_stream_reuse/test_run_stream_reuse.py index f828ebac9..474e90bd5 100644 --- a/tests/st/a2a3/host_build_graph/run_stream_reuse/test_run_stream_reuse.py +++ b/tests/st/a2a3/host_build_graph/run_stream_reuse/test_run_stream_reuse.py @@ -7,19 +7,18 @@ # 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. # ----------------------------------------------------------------------------------------------------------- -"""A2A3 gives every run a freshly created AICore stream. +"""A2A3 reuses an AICore stream only for the same loaded code image. A2A3 submits each run's AICore and AICPU kernels on the stream set of the selected pipeline slot rather than on the persistent bootstrap pair. The AICPU stream carries no instruction-cache state and belongs to the slot for the -worker's lifetime. The AICore stream belongs to one run: the platform offers no -instruction-cache invalidation for code replaced at a reused GM address, and -creating the stream is the only operation known to leave a core free of the -previous image's instructions — selecting an already-existing stream is not, so -reuse cannot be made safe by tracking which image a stream last ran. - -`Worker.run_stream_set_create_count` counts those creations, so it advances once -per run. +worker's lifetime. The slot also retains its AICore stream while the loaded code +image is unchanged. A different image destroys and recreates that stream because +the platform offers no instruction-cache invalidation for code replaced at a +reused GM address. + +`Worker.run_stream_set_create_count` counts AICore stream creations, so repeated +runs of one image do not advance it while image transitions do. """ import itertools @@ -69,7 +68,7 @@ class _SubtractCallable(SceneTestCase): @scene_test(level=2, runtime="host_build_graph") class TestRunStreamReuseHbg(SceneTestCase): - """Each run gets its own AICore stream; its pipeline slot owns the rest.""" + """Each slot retains one AICore stream for its currently loaded image.""" RTOL = 1e-5 ATOL = 1e-5 @@ -123,8 +122,8 @@ def compute_golden(self, args, params): a, b = args.a, args.b args.f[:] = (a + b + 1) * (a + b + 2) - def test_every_run_creates_its_own_aicore_stream(self, st_platform, st_worker): - """N runs create N AICore streams, and every result is right.""" + def test_same_image_reuses_aicore_stream(self, st_platform, st_worker): + """Repeated runs of one image reuse the slot's AICore stream.""" if st_platform != "a2a3": pytest.skip("run stream sets are an a2a3 onboard resource") @@ -134,8 +133,8 @@ def test_every_run_creates_its_own_aicore_stream(self, st_platform, st_worker): rounds = _REPEATED_RUNS - 1 self._run_and_validate_l2(st_worker, callable_obj, self.CASES[0], rounds=rounds) - assert st_worker.run_stream_set_create_count == after_first + rounds, ( - f"{rounds} runs did not each create an AICore stream: " + assert st_worker.run_stream_set_create_count == after_first, ( + f"{rounds} same-image runs recreated the AICore stream: " f"{after_first} -> {st_worker.run_stream_set_create_count}" ) @@ -176,7 +175,7 @@ def _run_registered_with_lease(self, worker, handle, *, slot_id, generation): _compare_outputs(test_args, golden_args, output_names, self.RTOL, self.ATOL) def test_alternating_code_images_never_execute_stale_instructions(self, st_platform, st_worker): - """A→B→A→B at one reused GM code address always runs the selected image. + """A-A-B-B-A reuses matching streams and replaces mismatched streams. Each result is compared against the golden for the image that run asked for, so an AICore that executed the previous image's instructions shows @@ -188,15 +187,19 @@ def test_alternating_code_images_never_execute_stale_instructions(self, st_platf add_handle = st_worker.register(self.build_callable(st_platform)) sub_handle = st_worker.register(_SubtractCallable.compile_chip_callable(st_platform)) try: - for handle, subtract in ( - (add_handle, False), - (sub_handle, True), - (add_handle, False), - (sub_handle, True), - ): - before = st_worker.run_stream_set_create_count - self._run_registered(st_worker, handle, subtract=subtract) - assert st_worker.run_stream_set_create_count == before + 1 + self._run_registered(st_worker, add_handle, subtract=False) + after_add = st_worker.run_stream_set_create_count + self._run_registered(st_worker, add_handle, subtract=False) + assert st_worker.run_stream_set_create_count == after_add + + self._run_registered(st_worker, sub_handle, subtract=True) + after_sub = st_worker.run_stream_set_create_count + assert after_sub == after_add + 1 + self._run_registered(st_worker, sub_handle, subtract=True) + assert st_worker.run_stream_set_create_count == after_sub + + self._run_registered(st_worker, add_handle, subtract=False) + assert st_worker.run_stream_set_create_count == after_sub + 1 finally: st_worker.unregister(sub_handle) st_worker.unregister(add_handle) 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 index 1345b754e..898f983f7 100644 --- 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 @@ -17,7 +17,7 @@ namespace { constexpr uint64_t kAdd = 0; constexpr uint64_t kAddScalar = 1; -constexpr int kChainLength = 512; +constexpr int kChainLength = 4096; } // namespace 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 c76689d8c..b6695cdfa 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,7 +15,6 @@ block before its graph callback while both slots are admitted. """ -import ctypes import multiprocessing import threading import time @@ -24,14 +23,13 @@ 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 +_CHAIN_LENGTH = 4096 _SUB_ENTERED = multiprocessing.Event() _SUB_RELEASE = multiprocessing.Event() @@ -103,6 +101,7 @@ def test_run(self, st_platform, st_worker): submitter = None first = None second = None + repeat = 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) @@ -126,17 +125,6 @@ def submit_vector(orch, a, b, out, *, hold_open=False): ) assert _SUB_ENTERED.wait(10.0), "the first run's SubTask did not start" - 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 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) - second_graph_done = threading.Event() def second_graph(orch, _args, _cfg): @@ -174,11 +162,22 @@ def third_graph(_orch, _args, _cfg): submitter.join(10.0) assert not submitter.is_alive() third_result["handle"].wait(10.0) + + stream_count = st_worker.run_stream_set_create_count + first_out.zero_() + repeat = st_worker.submit( + lambda orch, _args, _cfg: submit_vector(orch, first_a, first_b, first_out) + ) + repeat.wait(10.0) + assert torch.allclose(first_out, first_expected), "the repeated image produced the wrong result" + assert st_worker.run_stream_set_create_count == stream_count, ( + "a same-image run recreated a slot's AICore stream" + ) finally: _SUB_RELEASE.set() if submitter is not None: submitter.join(10.0) - handles = [first, second, third_result.get("handle")] + handles = [first, second, third_result.get("handle"), repeat] for handle in handles: if handle is not None: with suppress(Exception): @@ -190,5 +189,10 @@ def third_graph(_orch, _args, _cfg): st_worker.free_host_buffer(buffer) +@scene_test(level=3, runtime="tensormap_and_ringbuffer") +class TestWorkerAsyncWholeRunFifoTmr(TestWorkerAsyncWholeRunFifo): + """The same FIFO drives TMR's slot-private prepare and shared-image gate.""" + + if __name__ == "__main__": SceneTestCase.run_module(__name__) diff --git a/tests/ut/cpp/CMakeLists.txt b/tests/ut/cpp/CMakeLists.txt index cc75282a0..43d3647a7 100644 --- a/tests/ut/cpp/CMakeLists.txt +++ b/tests/ut/cpp/CMakeLists.txt @@ -368,9 +368,9 @@ add_hierarchical_test(test_remote_wire hierarchical/test_remote_wire.cpp) add_hierarchical_test(test_remote_endpoint hierarchical/test_remote_endpoint.cpp) add_hierarchical_test(test_pipeline_contract hierarchical/test_pipeline_contract.cpp) -# Run-stream slot state machine: fresh AICore stream per run, and the handle a -# failed destroy must keep. Header-only with injected create/destroy, so it -# needs no runtime objects and no device. +# Run-stream slot state machine: same-image reuse, image-transition replacement, +# and the handle a failed destroy must keep. Header-only with injected +# create/destroy, so it needs no runtime objects and no device. add_executable(test_run_stream_slots hierarchical/test_run_stream_slots.cpp) target_include_directories(test_run_stream_slots PRIVATE ${GTEST_INCLUDE_DIRS} diff --git a/tests/ut/cpp/common/test_trb_runtime_temp_buffer.cpp b/tests/ut/cpp/common/test_trb_runtime_temp_buffer.cpp index 30dba3652..0eae85426 100644 --- a/tests/ut/cpp/common/test_trb_runtime_temp_buffer.cpp +++ b/tests/ut/cpp/common/test_trb_runtime_temp_buffer.cpp @@ -41,6 +41,10 @@ extern "C" int bind_callable_to_runtime_impl( const uint64_t *ring_dep_pool ); extern "C" int validate_runtime_impl(Runtime *runtime, const HostApi *api, int execution_rc); +extern "C" int prepared_run_supported_impl(void); +extern "C" int prepared_run_config_compatible_impl( + const HostApi *api, const uint64_t *ring_task_window, const uint64_t *ring_heap, const uint64_t *ring_dep_pool +); namespace { @@ -66,6 +70,11 @@ struct FakeHostApi { std::vector gm_heap; std::vector gm_sm; std::vector runtime_arena; + bool compatibility_key_valid = false; + uint64_t compatibility_hash = 0; + std::vector compatibility_key; + uint64_t observed_hash = 0; + std::vector observed_key; ~FakeHostApi() { release_all(); } @@ -168,6 +177,25 @@ bool fake_lookup_prebuilt_runtime_arena_cache( ) { return false; } +bool fake_lookup_compatible_runtime_arena_cache( + uint64_t hash, const void *key_data, size_t key_size, void **gm_heap_base, void **sm_base, + void **runtime_arena_base, size_t *runtime_off, const void **image_data, size_t *image_size +) { + const auto *key = static_cast(key_data); + g_fake->observed_hash = hash; + g_fake->observed_key.assign(key, key + key_size); + bool hit = g_fake->compatibility_key_valid && hash == g_fake->compatibility_hash && + g_fake->observed_key == g_fake->compatibility_key; + if (hit) { + *gm_heap_base = reinterpret_cast(1); + *sm_base = reinterpret_cast(2); + *runtime_arena_base = reinterpret_cast(3); + *runtime_off = 4; + *image_data = reinterpret_cast(5); + *image_size = 6; + } + return hit; +} void fake_mark_prebuilt_runtime_arena_cached( uint64_t /* hash */, const void * /* key_data */, size_t /* key_size */, void * /* gm_heap_base */, void * /* sm_base */, void * /* runtime_arena_base */, size_t /* runtime_off */, const void * /* image_data */, @@ -198,6 +226,12 @@ HostApi make_host_api(bool with_temporary_buffer = true) { }; } +HostApi make_compatibility_host_api() { + HostApi api = make_host_api(); + api.lookup_prebuilt_runtime_arena_cache = fake_lookup_compatible_runtime_arena_cache; + return api; +} + Tensor make_tensor(std::vector &storage, bool child_memory = false) { Tensor tensor; uint32_t shape[1] = {static_cast(storage.size())}; @@ -465,3 +499,22 @@ TEST_F(TrbRuntimeTempBufferTest, FailedCopyOnTemporaryPathDoesNotFreeRetainedBuf EXPECT_NE(fake_.retained_addr, nullptr); EXPECT_TRUE(runtime.tensor_leases_.empty()); } + +TEST_F(TrbRuntimeTempBufferTest, PreparedRuntimeEnvRequiresTheActiveArenaKey) { + fake_.reset(); + HostApi compatibility_api = make_compatibility_host_api(); + uint64_t task_window[PTO2_MAX_RING_DEPTH] = {4, 4, 4, 4}; + uint64_t heap[PTO2_MAX_RING_DEPTH] = {1024, 1024, 1024, 1024}; + uint64_t dep_pool[PTO2_MAX_RING_DEPTH] = {4, 4, 4, 4}; + + EXPECT_EQ(prepared_run_supported_impl(), 1); + EXPECT_EQ(prepared_run_config_compatible_impl(&compatibility_api, task_window, heap, dep_pool), 0); + fake_.compatibility_key_valid = true; + fake_.compatibility_hash = fake_.observed_hash; + fake_.compatibility_key = fake_.observed_key; + + EXPECT_EQ(prepared_run_config_compatible_impl(&compatibility_api, task_window, heap, dep_pool), 1); + heap[2] = 2048; + EXPECT_EQ(prepared_run_config_compatible_impl(&compatibility_api, task_window, heap, dep_pool), 0); + EXPECT_NE(fake_.observed_key, fake_.compatibility_key); +} diff --git a/tests/ut/cpp/hierarchical/test_run_stream_slots.cpp b/tests/ut/cpp/hierarchical/test_run_stream_slots.cpp index 250b64566..7c1917bb4 100644 --- a/tests/ut/cpp/hierarchical/test_run_stream_slots.cpp +++ b/tests/ut/cpp/hierarchical/test_run_stream_slots.cpp @@ -69,25 +69,35 @@ RunStreamSlots make_slots(FakeStreams &fake) { ); } -// The AICPU stream is the slot's for the runner's lifetime; the AICore stream -// belongs to one run, so the count advances once per acquire. -TEST(RunStreamSlots, EveryAcquireCreatesAnAicoreStreamAndKeepsTheAicpuOne) { +// Both streams remain reusable while the loaded AICore image is unchanged. +TEST(RunStreamSlots, SameImageReusesBothStreams) { FakeStreams fake; RunStreamSlots slots = make_slots(fake); - ASSERT_EQ(slots.acquire(0), 0); + ASSERT_EQ(slots.acquire(0, 0xabc), 0); void *aicpu = slots.aicpu(0); void *first_aicore = slots.aicore(0); ASSERT_NE(aicpu, nullptr); ASSERT_NE(first_aicore, nullptr); EXPECT_EQ(slots.created_count(), 1u); - ASSERT_EQ(slots.retire_aicore(0), 0); - EXPECT_EQ(slots.aicore(0), nullptr); - - ASSERT_EQ(slots.acquire(0), 0); + ASSERT_EQ(slots.acquire(0, 0xabc), 0); EXPECT_EQ(slots.aicpu(0), aicpu) << "the AICPU stream must not be recreated"; - EXPECT_NE(slots.aicore(0), first_aicore) << "the AICore stream must be a new one"; + EXPECT_EQ(slots.aicore(0), first_aicore) << "the same AICore image must keep its stream"; + EXPECT_EQ(slots.created_count(), 1u); +} + +TEST(RunStreamSlots, DifferentImageRecreatesOnlyTheAicoreStream) { + FakeStreams fake; + RunStreamSlots slots = make_slots(fake); + + ASSERT_EQ(slots.acquire(0, 0xabc), 0); + void *aicpu = slots.aicpu(0); + void *first_aicore = slots.aicore(0); + + ASSERT_EQ(slots.acquire(0, 0xdef), 0); + EXPECT_EQ(slots.aicpu(0), aicpu); + EXPECT_NE(slots.aicore(0), first_aicore); EXPECT_EQ(slots.created_count(), 2u); } @@ -96,18 +106,18 @@ TEST(RunStreamSlots, AFailedDestroyReportsKeepsTheHandleAndLocksTheSlot) { FakeStreams fake; RunStreamSlots slots = make_slots(fake); - ASSERT_EQ(slots.acquire(0), 0); + ASSERT_EQ(slots.acquire(0, 0xabc), 0); void *stranded = slots.aicore(0); - // 1. the error reaches the caller + // 1. the image transition's destroy error reaches the caller fake.fail_next_destroys(1); - EXPECT_EQ(slots.retire_aicore(0), -13); + EXPECT_EQ(slots.acquire(0, 0xdef), -13); // 2. the handle survives, so teardown still has something to reclaim EXPECT_EQ(slots.aicore(0), stranded); // 3. the next run is refused rather than handed a second live stream - EXPECT_NE(slots.acquire(0), 0); + EXPECT_NE(slots.acquire(0, 0xabc), 0); EXPECT_EQ(slots.aicore(0), stranded); EXPECT_EQ(slots.created_count(), 1u) << "a refused acquire must not create anything"; @@ -122,11 +132,11 @@ TEST(RunStreamSlots, AStrandedSlotDoesNotBlockItsPeer) { FakeStreams fake; RunStreamSlots slots = make_slots(fake); - ASSERT_EQ(slots.acquire(0), 0); + ASSERT_EQ(slots.acquire(0, 0xabc), 0); fake.fail_next_destroys(1); - ASSERT_NE(slots.retire_aicore(0), 0); + ASSERT_NE(slots.acquire(0, 0xdef), 0); - ASSERT_EQ(slots.acquire(1), 0); + ASSERT_EQ(slots.acquire(1, 0xdef), 0); EXPECT_NE(slots.aicore(1), nullptr); EXPECT_EQ(slots.retire_aicore(1), 0); } @@ -136,8 +146,8 @@ TEST(RunStreamSlots, AStrandedSlotDoesNotBlockItsPeer) { TEST(RunStreamSlots, DestroyAllReportsFailureAndRetriesWhatSurvived) { FakeStreams fake; RunStreamSlots slots = make_slots(fake); - ASSERT_EQ(slots.acquire(0), 0); - ASSERT_EQ(slots.acquire(1), 0); + ASSERT_EQ(slots.acquire(0, 0xabc), 0); + ASSERT_EQ(slots.acquire(1, 0xdef), 0); fake.fail_next_destroys(1); EXPECT_EQ(slots.destroy_all(), -13); @@ -153,7 +163,7 @@ TEST(RunStreamSlots, AFailedCreateLeavesTheSlotEmpty) { RunStreamSlots slots = make_slots(fake); fake.fail_next_creates(1); - EXPECT_NE(slots.acquire(0), 0); + EXPECT_NE(slots.acquire(0, 0xabc), 0); EXPECT_EQ(slots.aicpu(0), nullptr); EXPECT_EQ(slots.aicore(0), nullptr); EXPECT_EQ(slots.created_count(), 0u); @@ -161,10 +171,10 @@ TEST(RunStreamSlots, AFailedCreateLeavesTheSlotEmpty) { // The AICPU stream lands, the AICore one fails: the slot stays unusable // rather than half-ready. fake.fail_next_creates(0); - ASSERT_EQ(slots.acquire(0), 0); + ASSERT_EQ(slots.acquire(0, 0xabc), 0); ASSERT_EQ(slots.retire_aicore(0), 0); fake.fail_next_creates(1); - EXPECT_NE(slots.acquire(0), 0); + EXPECT_NE(slots.acquire(0, 0xabc), 0); EXPECT_FALSE(slots.ready(0)); } @@ -172,7 +182,7 @@ TEST(RunStreamSlots, RejectsASlotOutsideTheContractDepth) { FakeStreams fake; RunStreamSlots slots = make_slots(fake); const unsigned past_end = static_cast(RunStreamSlots::capacity()); - EXPECT_NE(slots.acquire(past_end), 0); + EXPECT_NE(slots.acquire(past_end, 0xabc), 0); EXPECT_NE(slots.retire_aicore(past_end), 0); EXPECT_EQ(slots.aicore(past_end), nullptr); } diff --git a/tests/ut/cpp/hierarchical/test_scheduler.cpp b/tests/ut/cpp/hierarchical/test_scheduler.cpp index 738ae87b1..73013b549 100644 --- a/tests/ut/cpp/hierarchical/test_scheduler.cpp +++ b/tests/ut/cpp/hierarchical/test_scheduler.cpp @@ -290,18 +290,28 @@ class BlockingTwoFrameEndpoint final : public WorkerEndpoint { class PreparedActivationEndpoint final : public WorkerEndpoint { public: - PreparedActivationEndpoint() { - caps_.worker_id = 0; + explicit PreparedActivationEndpoint(int32_t worker_id = 0, bool gate_normal_acceptance = false) : + gate_normal_acceptance_(gate_normal_acceptance) { + caps_.worker_id = worker_id; caps_.max_inflight_tasks = 2; caps_.supports_prepare_activate = true; } const WorkerEndpointCaps &caps() const override { return caps_; } - WorkerCompletion run(Ring *, const WorkerDispatch &dispatch) override { + WorkerCompletion run(Ring *ring, const WorkerDispatch &dispatch) override { + return run_with_accept(ring, dispatch, {}); + } + + WorkerCompletion + run_with_accept(Ring *, const WorkerDispatch &dispatch, const std::function &on_accept) 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] { + return !gate_normal_acceptance_ || accept_normal_; + }); + if (on_accept) on_accept(); cv_.wait(lk, [this] { return release_normal_; }); @@ -325,6 +335,11 @@ class PreparedActivationEndpoint final : public WorkerEndpoint { 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 accept_normal() { + std::lock_guard lk(mu_); + accept_normal_ = true; + cv_.notify_all(); + } void release_normal() { std::lock_guard lk(mu_); release_normal_ = true; @@ -340,6 +355,8 @@ class PreparedActivationEndpoint final : public WorkerEndpoint { std::mutex mu_; std::condition_variable cv_; bool release_normal_{false}; + bool gate_normal_acceptance_{false}; + bool accept_normal_{false}; }; // --------------------------------------------------------------------------- @@ -570,6 +587,8 @@ TEST(WorkerManagerTest, PreparedFrameRequiresExplicitActivation) { /*supports_prepare_activate=*/true ); std::atomic allow_activation{false}; + std::atomic task_active{false}; + std::atomic publish_acceptance{false}; std::atomic backend_ready{false}; std::atomic accepted{false}; std::promise result; @@ -586,6 +605,12 @@ TEST(WorkerManagerTest, PreparedFrameRequiresExplicitActivation) { std::this_thread::sleep_for(std::chrono::microseconds(50)); } __atomic_store_n(state, static_cast(MailboxState::TASK_ACTIVE), __ATOMIC_RELEASE); + task_active.store(true, std::memory_order_release); + while (!publish_acceptance.load(std::memory_order_acquire)) { + std::this_thread::sleep_for(std::chrono::microseconds(50)); + } + auto *accepted_word = reinterpret_cast(frame + MAILBOX_OFF_ACCEPTED); + __atomic_store_n(accepted_word, MAILBOX_TASK_ACCEPTED, __ATOMIC_RELEASE); __atomic_store_n(state, static_cast(MailboxState::TASK_DONE), __ATOMIC_RELEASE); }); @@ -615,6 +640,13 @@ TEST(WorkerManagerTest, PreparedFrameRequiresExplicitActivation) { EXPECT_EQ(done.wait_for(std::chrono::milliseconds(0)), std::future_status::timeout); allow_activation.store(true, std::memory_order_release); + while (!task_active.load(std::memory_order_acquire) && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_TRUE(task_active.load(std::memory_order_acquire)); + EXPECT_FALSE(accepted.load(std::memory_order_acquire)) + << "TASK_ACTIVE is admission state, not native launch acceptance"; + publish_acceptance.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); @@ -637,9 +669,12 @@ TEST(SchedulerPreparedRunTest, SuccessorPreparesButActivatesOnlyAfterFifoPromoti 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)); + auto active_endpoint = std::make_unique(0, true); + PreparedActivationEndpoint *active_endpoint_ptr = active_endpoint.get(); + auto successor_endpoint = std::make_unique(1); + PreparedActivationEndpoint *successor_endpoint_ptr = successor_endpoint.get(); + manager.add_next_level_endpoint(std::move(active_endpoint)); + manager.add_next_level_endpoint(std::move(successor_endpoint)); manager.start( &allocator, [&](WorkerCompletion completion) { @@ -647,6 +682,7 @@ TEST(SchedulerPreparedRunTest, SuccessorPreparesButActivatesOnlyAfterFifoPromoti }, [&](WorkerDispatch dispatch) { orchestrator.mark_task_accepted(dispatch.task_slot); + scheduler.notify_ready(); } ); ready_next.reset(manager.next_level_worker_ids()); @@ -681,24 +717,29 @@ TEST(SchedulerPreparedRunTest, SuccessorPreparesButActivatesOnlyAfterFifoPromoti 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.submit_next_level(C(2), single_tensor_args(0x2000, TensorArgType::OUTPUT), config, 1); 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) { + while (!active_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()); - EXPECT_TRUE(endpoint_ptr->backend_ready()); - EXPECT_FALSE(endpoint_ptr->activated()); + EXPECT_TRUE(active_endpoint_ptr->normal_running()); + EXPECT_FALSE(successor_endpoint_ptr->backend_ready()); + + active_endpoint_ptr->accept_normal(); + while (!successor_endpoint_ptr->backend_ready() && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_TRUE(successor_endpoint_ptr->backend_ready()); + EXPECT_FALSE(successor_endpoint_ptr->activated()); EXPECT_EQ(orchestrator.active_run_id(), first_run); EXPECT_EQ(orchestrator.preparable_run_id(), second_run); - endpoint_ptr->release_normal(); + active_endpoint_ptr->release_normal(); orchestrator.wait_run(first_run); orchestrator.wait_run(second_run); - EXPECT_TRUE(endpoint_ptr->activated()); + EXPECT_TRUE(successor_endpoint_ptr->activated()); orchestrator.release_run(first_run); orchestrator.release_run(second_run); diff --git a/tests/ut/py/test_worker/test_host_worker.py b/tests/ut/py/test_worker/test_host_worker.py index 1f38d1edf..b278f7235 100644 --- a/tests/ut/py/test_worker/test_host_worker.py +++ b/tests/ut/py/test_worker/test_host_worker.py @@ -131,10 +131,15 @@ 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_caps: list[int] = [] published_frame_counts: list[int] = [] + class FakeImpl: + supports_prepared_runs = True + class FakeChipWorker: pipeline_depth = 2 + _impl = FakeImpl() 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)) @@ -143,7 +148,12 @@ def finalize(self) -> None: events.append(("finalize",)) 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]) + depth, caps, generation = worker_mod._PIPELINE_LEASE_FMT.unpack_from( + _args[0], worker_mod._OFF_PIPELINE_LEASE + ) + published_depths.append(depth) + published_caps.append(caps) + assert generation == 0 published_frame_counts.append(task_frame_count) events.append(("main_loop", cw, chip_platform, chip_runtime)) @@ -167,12 +177,15 @@ def fake_run_chip_main_loop(cw, *_args, chip_platform, chip_runtime, prepared=No shm.close() shm.unlink() - assert events[0] == ("init", 7, "bins", 25, None, False) + assert events[0][:4] == ("init", 7, "bins", 25) + assert isinstance(events[0][4], worker_mod.CallConfig) + assert events[0][5] is False assert events[1][0] == "main_loop" assert events[1][2:] == ("a2a3", "tensormap_and_ringbuffer") assert events[2] == ("finalize",) assert published_depths == [2] - assert published_frame_counts == [1] + assert published_caps == [worker_mod._PIPELINE_CAP_PREPARE_ACTIVATE] + assert published_frame_counts == [2] @pytest.mark.parametrize( @@ -180,7 +193,8 @@ def fake_run_chip_main_loop(cw, *_args, chip_platform, chip_runtime, prepared=No [ ("a2a3", "host_build_graph", 2, 2), ("a2a3", "host_build_graph", 1, 1), - ("a2a3", "tensormap_and_ringbuffer", 2, 1), + ("a2a3", "tensormap_and_ringbuffer", 2, 2), + ("a2a3", "tensormap_and_ringbuffer", 1, 1), ("a5", "host_build_graph", 2, 1), ], )