diff --git a/docs/task-flow.md b/docs/task-flow.md index f9b109ea9..b79006fa6 100644 --- a/docs/task-flow.md +++ b/docs/task-flow.md @@ -64,8 +64,9 @@ to C++: | `w3.submit_sub(handle, …)` dispatched to a SUB child | `LOCAL_PYTHON` | child resolves digest to a Python callable and calls `fn(args)` | All three paths share one mailbox wire format: `MAILBOX_OFF_CALLABLE` is -reserved, and the 32-byte digest prefixes the args blob. The receiving child -does the digest-to-slot resolve in its own address space. +reserved, the run's generation-safe pipeline lease follows `CallConfig`, and +the 32-byte digest prefixes the args blob. The receiving child does the +digest-to-slot resolve in its own address space. The proposed remote L3 path keeps the same callable identity contract, but sends it in a versioned TASK frame. The remote endpoint resolves the digest @@ -292,12 +293,66 @@ own AICore stream and retires it on every exit path, and no record of which image a stream last ran is load-bearing. The AICPU stream carries no such state and stays with its slot. -This is dormant capacity at this layer. The ordinary synchronous entry point -continues to use slot 0, and the chip child's mailbox loop passes no lease, so -every production run is unleased. This contract does not enable a second -mailbox frame, a second device execution, or cross-run publication overlap. -Carrying a lease across the mailbox, and deciding when slot 1 may be leased at -all, belong to whole-run admission. +#### Whole-run FIFO admission + +L3 graph callbacks remain synchronous and serialized, 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. + +#### 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 @@ -382,6 +437,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 +448,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. @@ -409,7 +468,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 f110c9a16..58756f001 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", @@ -393,10 +393,15 @@ 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, + bool supports_prepare_activate) { + self.add_worker( + 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("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 " @@ -404,11 +409,19 @@ 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, + bool supports_prepare_activate) { + self.add_next_level_worker( + 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, - "Add a NEXT_LEVEL sub-worker with an explicit worker id." + nb::arg("worker_id"), 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 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", @@ -764,8 +777,10 @@ 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); 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/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 5535007aa..207050e36 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, @@ -128,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, @@ -176,10 +178,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 @@ -190,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 = 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++ # nanobind module so the two sides cannot drift. @@ -212,6 +223,20 @@ 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 +_BACKEND_READY = 11 +_ACTIVATE = 12 +_PREPARE_READY = 13 +_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. @@ -1508,6 +1533,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`. @@ -1533,10 +1559,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 = "" @@ -1557,19 +1588,26 @@ 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( + task_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( 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, ) except Exception as e: # noqa: BLE001 code = 1 @@ -1690,8 +1728,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(): @@ -1767,6 +1910,10 @@ 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") sys.stderr.flush() @@ -1784,6 +1931,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() @@ -4489,6 +4637,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 +4743,16 @@ 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 + # 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}") + 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,14 +4816,18 @@ 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 # 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"): @@ -4675,7 +4838,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) @@ -6147,9 +6310,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 +6339,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..38d8cb824 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,26 @@ 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(); + if (ready_notify_cb_) ready_notify_cb_(); + 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 +211,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 +276,31 @@ 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_; +} + +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() { @@ -218,8 +329,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 +513,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 +543,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 +626,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 +696,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 +729,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..a817a88b5 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,9 @@ 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; + RunId preparable_run_id() const; void release_run(RunId run_id); // Open a nested scope. Every task submitted between this call and the @@ -176,9 +182,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 +201,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..fa7681332 100644 --- a/src/common/hierarchical/scheduler.cpp +++ b/src/common/hierarchical/scheduler.cpp @@ -176,8 +176,21 @@ 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) || + 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(); + } + return !completion_queue_.empty() || ready || stop_requested_.load(std::memory_order_acquire); }); } @@ -350,16 +363,81 @@ 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; 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 +448,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 +471,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"); } @@ -417,12 +508,14 @@ 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); } 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 +528,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) { @@ -442,18 +537,23 @@ 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_.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"); } 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/scheduler.h b/src/common/hierarchical/scheduler.h index 00c1abfd2..3e4e0085a 100644 --- a/src/common/hierarchical/scheduler.h +++ b/src/common/hierarchical/scheduler.h @@ -59,6 +59,10 @@ 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; + 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 @@ -103,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/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..84b5dde7c 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 @@ -122,39 +125,43 @@ static constexpr int32_t INVALID_SLOT = -1; using RunId = uint64_t; static constexpr RunId INVALID_RUN_ID = 0; +using TaskSlot = int32_t; -// 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}; }; // ============================================================================= // Task slot index type // ============================================================================= -using TaskSlot = int32_t; - static constexpr size_t CALLABLE_HASH_DIGEST_SIZE = 32; enum class CallableKind : int32_t { @@ -322,6 +329,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 +431,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 +458,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..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) { +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); - 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) { +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); + manager_.add_next_level_at(worker_id, mailbox, child_pid, task_frame_count, supports_prepare_activate); } void Worker::add_remote_l3_socket( @@ -122,6 +127,12 @@ 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.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 f8e2e82b0..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); - 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, + 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. @@ -92,6 +98,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..d05f42ffc 100644 --- a/src/common/hierarchical/worker_manager.cpp +++ b/src/common/hierarchical/worker_manager.cpp @@ -146,18 +146,34 @@ 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) : +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) { 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; + caps_.supports_prepare_activate = supports_prepare_activate; } 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 +201,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 +216,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 +230,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,24 +266,76 @@ 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); + 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); + 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(); } +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(); - if (thread_.joinable()) thread_.join(); + activation_cv_.notify_all(); + for (auto &thread : threads_) { + if (thread.joinable()) thread.join(); + } + threads_.clear(); } void WorkerThread::shutdown_child() { @@ -282,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(); @@ -328,13 +405,51 @@ void WorkerThread::loop() { } } - idle_.store(true, std::memory_order_release); + inflight_.fetch_sub(1, std::memory_order_acq_rel); on_complete_(std::move(completion)); } } 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); } @@ -345,6 +460,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 +476,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 @@ -383,6 +498,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) + @@ -471,17 +587,326 @@ 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, + const std::function &on_backend_ready, const std::function &wait_for_activation +) { + 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(wait_for_activation ? MailboxState::PREPARE_READY : 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; + 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 || + 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, 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) { +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}); + 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) { @@ -518,7 +943,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, entry.supports_prepare_activate + ); wt->start(ring, on_complete, on_accept, std::move(endpoint)); next_level_threads_.push_back(std::move(wt)); } @@ -633,7 +1060,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) { @@ -643,7 +1072,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); } } @@ -1000,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 6d00a9551..f4230fad5 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,12 @@ 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, + BACKEND_READY = 11, + ACTIVATE = 12, + PREPARE_READY = 13, }; // Sized so the args region can hold any TaskArgs the runtime itself accepts @@ -82,7 +89,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 = 2; // Error message region lives at the mailbox tail. 256 B of headroom is // enough for `: ` produced by the child-side @@ -92,7 +104,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,15 +115,21 @@ 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); + 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 @@ -118,14 +137,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), @@ -200,6 +228,8 @@ struct WorkerEndpointCaps { bool remote{false}; 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"}; }; @@ -213,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); @@ -264,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); + 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; @@ -318,7 +359,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 @@ -327,12 +377,17 @@ 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, + 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); // Returns a description of the child's death, or an empty string while it @@ -352,6 +407,8 @@ class LocalMailboxEndpoint : public WorkerEndpoint { struct WorkerDispatch { TaskSlot task_slot{INVALID_SLOT}; int32_t group_index{0}; + uint64_t dispatch_id{0}; + bool prepare_only{false}; }; // ============================================================================= @@ -384,9 +441,16 @@ 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 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; @@ -471,15 +535,24 @@ 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}; + 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); }; // ============================================================================= @@ -496,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); - 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, 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); @@ -515,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 @@ -583,6 +664,8 @@ class WorkerManager { int32_t worker_id{-1}; 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/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_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/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..6beca4b60 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,139 @@ 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)); + + 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(); + 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)); + } + 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(); + }); + 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(); + + 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); + 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..6a7511376 100644 --- a/tests/ut/cpp/hierarchical/test_scheduler.cpp +++ b/tests/ut/cpp/hierarchical/test_scheduler.cpp @@ -247,6 +247,98 @@ 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}; +}; + +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). // --------------------------------------------------------------------------- @@ -318,6 +410,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); @@ -375,6 +470,240 @@ 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, 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(); @@ -387,6 +716,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,9 +730,17 @@ 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) {} + 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); @@ -1180,6 +1518,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 +1584,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..d208f2c30 100644 --- a/tests/ut/py/test_worker/test_host_worker.py +++ b/tests/ut/py/test_worker/test_host_worker.py @@ -130,15 +130,21 @@ 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 + 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)) 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) @@ -165,6 +171,118 @@ 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] + 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: @@ -1435,6 +1553,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()