diff --git a/docs/chip-level-arch.md b/docs/chip-level-arch.md index a0e76f0f5d..1735d001f4 100644 --- a/docs/chip-level-arch.md +++ b/docs/chip-level-arch.md @@ -109,7 +109,7 @@ runner.set_executors(aicpu_binary, aicore_binary); // once, at init time std::unique_ptr prepared; runner.prepare_execution(runtime, config, pipeline_slot, identity, &prepared); auto launched = runner.launch_execution(std::move(prepared), std::move(permit)); -runner.drain_execution(*launched.active); // child progress path owns progress +runner.drain_execution(*launched.active); // resident lane lifecycle owner waits/finalizes runner.finalize(); ``` diff --git a/docs/task-flow.md b/docs/task-flow.md index eb8200c1f0..3d8be1a9cb 100644 --- a/docs/task-flow.md +++ b/docs/task-flow.md @@ -379,10 +379,10 @@ prepared-but-not-launched run. A direct A2/A3 chip endpoint with a negotiated depth of at least two uses two task frames and advertises `supports_frame_staging`. One `WorkerThread` owns -both frames and drives them through a non-blocking progress interface; the -child process likewise has one loop that services control traffic, both task -frames, and the bounded active/prepared native lifecycles. There is no thread -per frame. +both frames and drives them through a non-blocking progress interface. In the +child process, the mailbox loop services control traffic and both task frames, +while one resident C++ lifecycle thread per chip lane blocks on the launched +native token's device completion. There is no thread per frame or per run. The active and successor paths are: @@ -390,6 +390,9 @@ The active and successor paths are: IDLE -> TASK_READY -> FRAME_STAGED -> TASK_LAUNCHED -> TASK_DONE | TASK_FAILED IDLE -> PREPARE_READY -> FRAME_STAGED -> ACTIVATE -> TASK_LAUNCHED -> TASK_DONE | TASK_FAILED +IDLE -> NATIVE_PREPARE_READY -> FRAME_STAGED -> ACTIVATE -> TASK_LAUNCHED + -> TASK_DONE | TASK_FAILED + -> ABANDON -> TASK_FAILED ``` `FRAME_STAGED` means that the child owns an immutable frame snapshot; it does @@ -411,16 +414,29 @@ between these meanings. An HBG successor's prepared token remains unlaunched and unaccepted until `ACTIVATE`, and activation still cannot launch it until the predecessor has -polled complete and finalized. The sticky acceptance word therefore remains -zero throughout preparation. Shutdown, stale activation, and pre-launch -failure finalize the token exactly once before the frame becomes terminal. +completed and finalized. The resident lifecycle thread owns the launched +native token's blocking wait and finalizes outside the lane mutex; it takes the +mutex only to publish terminal state and launch the successor. The mailbox +loop uses a short, bounded wait for launched handles so the binding releases +the GIL and does not starve that lifecycle thread, while unlaunched handles +remain non-blocking status probes. The sticky acceptance word therefore +remains zero throughout preparation. Shutdown, stale activation, and +pre-launch failure finalize the token exactly once before the frame becomes +terminal. The scheduler stages only the first eligible single NEXT_LEVEL task from the -prepared FIFO successor. Tasks from the active run use only the active lane, so -the second frame cannot create same-device execution overlap. Prepared groups -remain on their normal queue and dispatch synchronously after FIFO promotion. -Remote, SUB, A5, simulation, nested-worker, and single-frame endpoints retain -the blocking compatibility path. +prepared FIFO successor. A prepared group never crosses the whole-run FIFO; +the second frame therefore cannot create same-device execution overlap. A +non-diagnostic NEXT_LEVEL group whose members all target local two-frame endpoints uses a +group barrier: the scheduler publishes every member as +`NATIVE_PREPARE_READY`, keeps all target workers reserved, and publishes no +`ACTIVATE` until every member has completed native bind/prepare and then +reported `FRAME_STAGED` with disposition `NATIVE_PREPARED`. It activates every +member in the same scheduler progress pass. If preparation of one member fails, staged peers receive +`ABANDON`; the child releases each unlaunched native run and acknowledges +`TASK_FAILED`, so a partial group can never launch. Diagnostic groups and +remote, SUB, A5, simulation, nested-worker, and single-frame endpoints retain +the compatibility path. The child validates newly visible metadata from both frames before selecting the next active `dispatch_id`. It prepares that active token first; preparation @@ -750,9 +766,9 @@ Step-by-step (one chip worker): | 2 | `Worker::run` | `scope_begin` → call `my_orch(&orch_, args.view(), cfg)` | | 3 | `Orchestrator::submit_next_level` | `slot = ring.alloc()`; move `chip_args` into `slot.task_args`; walk tags → `tensormap.lookup(a.data)`, `tensormap.lookup(b.data)`, `tensormap.insert(c.data, slot)`; push ready | | 4 | Scheduler thread | pop `slot` from worker 0's FIFO; resolve stable worker ID 0 to WT_chip_0; dispatch | -| 5 | WT_chip_0 parent side | encode one leased task frame: write `config`, digest prefix, and the args blob; publish `TASK_READY` for the active lane or `PREPARE_READY` for a staged successor | +| 5 | WT_chip_0 parent side | encode one leased task frame: write `config`, digest prefix, and the args blob; publish `TASK_READY` for the active lane, `PREPARE_READY` for a staged successor, or `NATIVE_PREPARE_READY` for a local group barrier member | | 6 | chip_0 child process | validate the frame and resolve its digest; ordinary HBG with an active predecessor also prepares the leased inactive arena bank before publishing `FRAME_STAGED`, while a frame with no active predecessor, diagnostic HBG, and TMR publish after validation and defer native prepare | -| 7 | chip_0 native-run path | after activation and the predecessor's finalization fence, launch an already-prepared HBG run or finish deferred native preparation and then launch; poll it to completion and finalize it before another staged frame may launch. Compatibility endpoints perform the equivalent operation through blocking `ChipWorker::run` | +| 7 | chip_0 native-run path | after activation and the predecessor's finalization fence, launch an already-prepared HBG run or finish deferred native preparation and then launch; the resident lifecycle owner blocks on device completion, finalizes outside the lane mutex, publishes terminal state, and launches the staged successor. Compatibility endpoints perform the equivalent operation through blocking `ChipWorker::run` | | 8 | runtime.so | translate host ptrs → device ptrs; dispatch AICPU / AICore; write output into `c`'s shm | | 9 | chip_0 child | native finalization returns; write `TASK_DONE` | | 10 | WT_chip_0 parent | observe `TASK_DONE`; push success completion | diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index 484c817640..215ddb228b 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -2759,6 +2759,7 @@ NB_MODULE(_task_interface, m) { nb::class_(m, "_ChipRun") .def("done", &ChipRun::done) + .def("prepare", &ChipRun::prepare, nb::call_guard()) .def("activate", &ChipRun::activate) .def("abandon", &ChipRun::abandon, nb::call_guard()) .def_prop_ro("launched", &ChipRun::launched) diff --git a/python/bindings/worker_bind.h b/python/bindings/worker_bind.h index 6c148ac321..414faadd2d 100644 --- a/python/bindings/worker_bind.h +++ b/python/bindings/worker_bind.h @@ -234,6 +234,12 @@ inline void mailbox_store_i32(uint64_t addr, int32_t v) { #endif } +inline void mailbox_notify_i32(uint64_t addr) { + auto *ptr = reinterpret_cast(addr); + (void)__atomic_add_fetch(ptr, 1, __ATOMIC_RELEASE); + mpi_group_mailbox::wake_word(ptr); +} + inline void bind_worker(nb::module_ &m) { // --- WorkerType --- nb::enum_(m, "WorkerType").value("NEXT_LEVEL", WorkerType::NEXT_LEVEL).value("SUB", WorkerType::SUB); @@ -859,6 +865,7 @@ inline void bind_worker(nb::module_ &m) { 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_OFF_NOTIFICATION") = static_cast(MAILBOX_OFF_NOTIFICATION); m.attr("MAILBOX_ERROR_MSG_SIZE") = static_cast(MAILBOX_ERROR_MSG_SIZE); // The MailboxState values as the C++ side defines them, keyed by // enumerator name. They are a cross-process wire contract: the word at @@ -881,6 +888,8 @@ inline void bind_worker(nb::module_ &m) { mailbox_states["TASK_FAILED"] = static_cast(MailboxState::TASK_FAILED); mailbox_states["ACTIVATE"] = static_cast(MailboxState::ACTIVATE); mailbox_states["PREPARE_READY"] = static_cast(MailboxState::PREPARE_READY); + mailbox_states["ABANDON"] = static_cast(MailboxState::ABANDON); + mailbox_states["NATIVE_PREPARE_READY"] = static_cast(MailboxState::NATIVE_PREPARE_READY); m.attr("MAILBOX_STATE_VALUES") = mailbox_states; nb::dict mailbox_dispositions; mailbox_dispositions["NONE"] = static_cast(MailboxPreparationDisposition::NONE); @@ -908,6 +917,13 @@ inline void bind_worker(nb::module_ &m) { }, nb::arg("addr"), nb::arg("value"), "Release-store a 32-bit mailbox word at `addr`." ); + m.def( + "_mailbox_notify_i32", + [](uint64_t addr) { + mailbox_notify_i32(addr); + }, + nb::arg("addr"), "Increment and wake a 32-bit mailbox notification generation at `addr`." + ); m.def( "_mailbox_wait_i32", [](uint64_t addr, int32_t expected, double timeout_s) { diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 45777523b1..3de91f1726 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -88,13 +88,16 @@ def my_l4_orch(orch, args, config): import cloudpickle from _task_interface import ( # pyright: ignore[reportMissingImports] HOST_STRACE_ENABLED, + MAILBOX_OFF_NOTIFICATION, MAX_REGISTERED_CALLABLE_IDS, PTO_PIPELINE_MAX_DEPTH, RUNTIME_ENV_RING_COUNT, WorkerType, _emit_host_span, _mailbox_load_i32, + _mailbox_notify_i32, _mailbox_store_i32, + _mailbox_wait_i32, _read_control_copy_request, _set_host_span_level_prefix, _worker_host_mapped_region_ack_cleanup_error, @@ -335,7 +338,7 @@ def _host_spans_active() -> bool: _OFF_FRAME_SLOT_ID = _OFF_ACCEPTED - 24 _OFF_FRAME_GENERATION = _OFF_ACCEPTED - 16 _OFF_FRAME_DISPATCH_ID = _OFF_ACCEPTED - 8 -_TASK_PROTOCOL_VERSION = 3 +_TASK_PROTOCOL_VERSION = 5 # Mirrors MAILBOX_OFF_SHUTDOWN / MAILBOX_SHUTDOWN_REQUESTED: termination is a # sticky one-way word on the control frame, not a MailboxState. _OFF_STATE has # three writers (parent CONTROL_REQUEST, child CONTROL_DONE, C++ @@ -369,6 +372,8 @@ def _host_spans_active() -> bool: _TASK_FAILED = 10 _ACTIVATE = 11 _PREPARE_READY = 12 +_ABANDON = 13 +_NATIVE_PREPARE_READY = 14 _TASK_FRAME_COUNT = 2 @@ -397,6 +402,8 @@ def _assert_mailbox_wire_constants() -> None: "TASK_FAILED": _TASK_FAILED, "ACTIVATE": _ACTIVATE, "PREPARE_READY": _PREPARE_READY, + "ABANDON": _ABANDON, + "NATIVE_PREPARE_READY": _NATIVE_PREPARE_READY, } dispositions = { "NONE": _DISPOSITION_NONE, @@ -2048,6 +2055,7 @@ def _request_child_shutdown(buf) -> None: """ _mailbox_store_i32(_buffer_field_addr(buf, _OFF_SHUTDOWN), _SHUTDOWN_REQUESTED) _mailbox_store_i32(_buffer_field_addr(buf, _OFF_STATE), _SHUTDOWN) + _mailbox_notify_i32(_buffer_field_addr(buf, MAILBOX_OFF_NOTIFICATION)) def _write_error(buf, code: int, msg: str = "") -> None: @@ -2112,6 +2120,7 @@ def _reexport_args_from_mailbox(buf, worker: Worker) -> TaskArgs: # orphan is reaped before it is noticeable, cheap enough to be lost in the # noise of the poll itself. _PARENT_LIVENESS_POLL_INTERVAL = 1000 +_CHIP_RUN_PROGRESS_WAIT_S = 0.001 def _run_mailbox_loop( @@ -2970,6 +2979,7 @@ class _StagedFrame: cid: int config: CallConfig activated: bool + require_native_prepare: bool chip_run: Any = None launched_published: bool = False @@ -2985,7 +2995,15 @@ def read_identity(frame_buf: memoryview) -> tuple[int, int, int, int, int]: ) def task_frame_references_digest(digest: bytes) -> bool: - live_states = (_TASK_READY, _PREPARE_READY, _ACTIVATE, _FRAME_STAGED, _TASK_LAUNCHED) + live_states = ( + _TASK_READY, + _PREPARE_READY, + _NATIVE_PREPARE_READY, + _ACTIVATE, + _FRAME_STAGED, + _TASK_LAUNCHED, + _ABANDON, + ) for index, frame_buf in enumerate(frame_bufs): if _mailbox_load_i32(frame_addrs[index] + _OFF_STATE) not in live_states: continue @@ -3037,6 +3055,7 @@ def stage_frame(index: int, initial_state: int) -> _StagedFrame | None: cid=int(cid), config=_read_config_from_mailbox(frame_buf), activated=initial_state in (_TASK_READY, _ACTIVATE), + require_native_prepare=initial_state == _NATIVE_PREPARE_READY, ) except Exception as e: # noqa: BLE001 _write_error(frame_buf, 1, _format_exc(f"chip_process dev={device_id} frame={index}", e)) @@ -3066,6 +3085,8 @@ def submit_frame(frame: _StagedFrame) -> None: _TASK_ACCEPTED, False, ) + if frame.require_native_prepare: + frame.chip_run.prepare() raw_disposition = frame.chip_run.preparation_disposition disposition = int(getattr(raw_disposition, "value", raw_disposition)) if disposition not in (_VALIDATED_ONLY, _NATIVE_PREPARED): @@ -3080,8 +3101,10 @@ def submit_frame(frame: _StagedFrame) -> None: liveness_countdown = _PARENT_LIVENESS_POLL_INTERVAL shutdown_message = f"chip_process dev={device_id}: task loop shut down" shutdown_addr = _buffer_field_addr(buf, _OFF_SHUTDOWN) + notification_addr = _buffer_field_addr(buf, MAILBOX_OFF_NOTIFICATION) try: while True: + notification_snapshot = _mailbox_load_i32(notification_addr) control_state = _mailbox_load_i32(state_addr) if control_state == _SHUTDOWN or _mailbox_load_i32(shutdown_addr) == _SHUTDOWN_REQUESTED: break @@ -3101,7 +3124,12 @@ def submit_frame(frame: _StagedFrame) -> None: frame_state = _mailbox_load_i32(frame_addrs[index] + _OFF_STATE) staged = staged_frames.get(index) if staged is None: - if frame_state in (_TASK_READY, _PREPARE_READY, _ACTIVATE): + if frame_state in ( + _TASK_READY, + _PREPARE_READY, + _NATIVE_PREPARE_READY, + _ACTIVATE, + ): staged = stage_frame(index, frame_state) if staged is not None: new_frames.append(staged) @@ -3124,6 +3152,15 @@ def submit_frame(frame: _StagedFrame) -> None: except Exception as e: # noqa: BLE001 shutdown_message = _format_exc(f"chip_process dev={device_id}: native activation", e) break + if frame_state == _ABANDON and not staged.activated: + try: + staged.chip_run.abandon() + except Exception as e: # noqa: BLE001 + shutdown_message = _format_exc(f"chip_process dev={device_id}: native abandonment", e) + break + _mailbox_store_i32(staged.frame_addr + _OFF_STATE, _TASK_FAILED) + staged_frames.pop(index, None) + continue else: for staged in sorted(new_frames, key=lambda frame: frame.identity[4]): try: @@ -3140,7 +3177,15 @@ def submit_frame(frame: _StagedFrame) -> None: _mailbox_store_i32(staged.frame_addr + _OFF_STATE, _TASK_LAUNCHED) staged.launched_published = True continue - run_complete = bool(staged.chip_run.done()) + # A launched run has a resident C++ completion + # owner. Its bounded wait releases the GIL while + # preserving mailbox responsiveness. Unlaunched + # staged work remains a nonblocking status probe. + run_complete = bool( + staged.chip_run.wait(_CHIP_RUN_PROGRESS_WAIT_S) + if staged.chip_run.launched + else staged.chip_run.done() + ) if not run_complete and staged.chip_run.launched and not staged.launched_published: _mailbox_store_i32(staged.frame_addr + _OFF_STATE, _TASK_LAUNCHED) staged.launched_published = True @@ -3181,6 +3226,14 @@ def submit_frame(frame: _StagedFrame) -> None: if os.getppid() != parent_pid: shutdown_message = f"chip_process dev={device_id}: parent exited" break + # With no launched run, every interesting transition is + # a parent mailbox publication. Park on the shared + # generation instead of burning a core in Python. A + # publication racing this scan changes the expected + # value, so FUTEX_WAIT returns immediately rather than + # losing the wake. + if not any(frame.chip_run.launched for frame in staged_frames.values()): + _mailbox_wait_i32(notification_addr, notification_snapshot, 0.01) continue break finally: @@ -3195,8 +3248,10 @@ def submit_frame(frame: _StagedFrame) -> None: if _mailbox_load_i32(frame_state_addr) in ( _TASK_READY, _PREPARE_READY, + _NATIVE_PREPARE_READY, _ACTIVATE, _FRAME_STAGED, + _ABANDON, _TASK_LAUNCHED, ): _write_error(frame_buf, 1, shutdown_message) diff --git a/src/common/hierarchical/scheduler.cpp b/src/common/hierarchical/scheduler.cpp index fd347cb32c..673218c1e3 100644 --- a/src/common/hierarchical/scheduler.cpp +++ b/src/common/hierarchical/scheduler.cpp @@ -113,6 +113,7 @@ void Scheduler::start(const Config &cfg) { ++wake_generation_; } dispatch_round_count_.store(0, std::memory_order_relaxed); + pending_group_barriers_.clear(); reservation_stall_episode_.reset(); stop_requested_.store(false, std::memory_order_relaxed); running_.store(true, std::memory_order_release); @@ -276,6 +277,7 @@ void Scheduler::run() { std::lock_guard loop_lk(loop_mu_); cfg_.manager->progress(); + progress_group_barriers(); // Phase 1: drain completions while (true) { @@ -432,7 +434,12 @@ void Scheduler::dispatch_ready() { RunId active_run = cfg_.active_run_cb(); if (active_run == INVALID_RUN_ID) return; run_snapshot = active_run; - cfg_.manager->activate_prepared_run(active_run); + const bool group_barrier_pending = + std::any_of(pending_group_barriers_.begin(), pending_group_barriers_.end(), [&](TaskSlot slot) { + const TaskSlotState *state = cfg_.ring->slot_state(slot); + return state != nullptr && state->run_id == active_run; + }); + if (!group_barrier_pending) cfg_.manager->activate_prepared_run(active_run); } dispatch_preparable_next_level_singles(); @@ -441,16 +448,90 @@ void Scheduler::dispatch_ready() { // whole-run FIFO head, even if a completion advances the head mid-pass. bool group_arrived_between_phases = false; do { - const NextLevelGroupDispatchResult group_result = dispatch_next_level_group(run_snapshot); + NextLevelGroupDispatchResult group_result = dispatch_next_level_group(run_snapshot); update_reservation_stall(group_result); if (cfg_.after_group_phase_cb) cfg_.after_group_phase_cb(); + // A prepared group owns every target until the all-member barrier has + // activated it. The workers' active lanes are still empty during this + // interval, so reserve them explicitly; otherwise a following single + // from the same run can occupy an active lane and split the barrier. + for (TaskSlot pending_slot : pending_group_barriers_) { + const TaskSlotState *pending = cfg_.ring->slot_state(pending_slot); + if (pending == nullptr) continue; + for (int32_t i = 0; i < pending->group_size(); ++i) { + group_result.reserved_worker_ids.insert(pending->target_worker_id(i)); + } + } group_arrived_between_phases = dispatch_next_level_singles( - group_result.reserved_worker_ids, run_snapshot, group_result.blocked_group_slot == INVALID_SLOT + group_result.reserved_worker_ids, run_snapshot, + group_result.blocked_group_slot == INVALID_SLOT && pending_group_barriers_.empty() ); } while (group_arrived_between_phases); dispatch_sub_ready(run_snapshot); } +void Scheduler::progress_group_barriers() { + const RunId active_run = cfg_.active_run_cb ? cfg_.active_run_cb() : INVALID_RUN_ID; + for (auto it = pending_group_barriers_.begin(); it != pending_group_barriers_.end();) { + const TaskSlot slot = *it; + TaskSlotState &state = *cfg_.ring->slot_state(slot); + if (state.state.load(std::memory_order_acquire) != TaskState::RUNNING || !state.is_group()) { + it = pending_group_barriers_.erase(it); + continue; + } + + bool failed = false; + std::string failure_message; + { + std::lock_guard lk(state.group_mu); + failed = state.group_failed; + failure_message = state.group_first_failure_message; + } + + const int32_t group_size = state.group_size(); + std::vector workers; + workers.reserve(static_cast(group_size)); + // Activation is legal only for the current whole-run FIFO head and + // after every target endpoint has retired its predecessor active lane. + bool all_ready = !failed && (active_run == INVALID_RUN_ID || state.run_id == active_run); + for (int32_t i = 0; i < group_size; ++i) { + WorkerThread *worker = cfg_.manager->get_worker_by_id(WorkerType::NEXT_LEVEL, state.target_worker_id(i)); + if (worker == nullptr) { + failed = true; + all_ready = false; + failure_message = "group barrier lost a target worker before activation"; + break; + } + workers.push_back(worker); + if (!worker->prepared_ready(state.run_id) || !worker->idle()) all_ready = false; + } + + if (failed) { + if (failure_message.empty()) failure_message = "group member failed before the activation barrier"; + for (WorkerThread *worker : workers) + (void)worker->cancel_prepared(state.run_id, failure_message); + it = pending_group_barriers_.erase(it); + continue; + } + if (!all_ready) { + ++it; + continue; + } + // Every target owns a fully prepared frame before any activation is + // requested. activate_prepared only moves endpoint-lane metadata; the + // following manager progress pass publishes all ACTIVATE stores. + for (WorkerThread *worker : workers) { + if (!worker->activate_prepared(state.run_id)) { + // Readiness was checked for every member above and the + // Scheduler is the sole lane owner, so this is an invariant + // failure rather than a recoverable capacity race. + throw std::runtime_error("group barrier lost a prepared member during activation"); + } + } + it = pending_group_barriers_.erase(it); + } +} + bool claim_for_dispatch(TaskSlotState &s) { TaskState expected = TaskState::READY; return s.state.compare_exchange_strong( @@ -482,12 +563,16 @@ void Scheduler::dispatch_preparable_next_level_singles() { WorkerThread *worker = cfg_.manager->get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); if (worker == nullptr || !worker->can_stage()) continue; TaskSlot slot; - if (!cfg_.ready_next_level_queues->try_pop_single(worker_id, run_id, slot)) continue; + if (!cfg_.ready_next_level_queues->try_front_single(worker_id, run_id, slot)) continue; + TaskSlotState &state = *cfg_.ring->slot_state(slot); + if (!worker->can_stage(state.pipeline_lease.slot_id)) continue; + TaskSlot popped = INVALID_SLOT; + if (!cfg_.ready_next_level_queues->try_pop_single(worker_id, run_id, popped)) continue; + if (popped != slot) throw std::runtime_error("prepared single queue head changed during dispatch"); if (!cfg_.ready_next_level_queues->groups_empty(run_id)) { cfg_.enqueue_ready_cb(slot); return; } - TaskSlotState &state = *cfg_.ring->slot_state(slot); if (state.state.load(std::memory_order_acquire) != TaskState::READY) continue; if (state.run_id != run_id || state.worker_type != WorkerType::NEXT_LEVEL || state.is_group() || state.target_worker_id(0) != worker_id) { @@ -600,6 +685,7 @@ Scheduler::NextLevelGroupDispatchResult Scheduler::dispatch_next_level_group(con workers.reserve(static_cast(group_size)); result.reserved_worker_ids.reserve(static_cast(group_size)); bool all_workers_idle = true; + bool supports_prepared_barrier = !s.config.diagnostics_any(); for (int32_t i = 0; i < group_size; ++i) { const int32_t worker_id = s.target_worker_id(i); WorkerThread *worker = cfg_.manager->get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); @@ -609,6 +695,9 @@ Scheduler::NextLevelGroupDispatchResult Scheduler::dispatch_next_level_group(con if (!result.reserved_worker_ids.insert(worker_id).second) { throw std::runtime_error("Scheduler::dispatch_next_level_group: duplicate target worker"); } + const WorkerEndpointCaps &caps = worker->caps(); + supports_prepared_barrier = supports_prepared_barrier && caps.kind == WorkerEndpointKind::LOCAL_MAILBOX && + caps.supports_frame_staging; const bool worker_idle = worker->idle(); if (!worker_idle) { all_workers_idle = false; @@ -616,6 +705,22 @@ Scheduler::NextLevelGroupDispatchResult Scheduler::dispatch_next_level_group(con } workers.push_back(worker); } + if (supports_prepared_barrier) { + // A local full-rank group may occupy the staged lane, but it is + // not activated until every member reports native preparation + // complete and the target active lanes are idle. + all_workers_idle = std::all_of(workers.begin(), workers.end(), [&](WorkerThread *worker) { + return worker->can_stage(s.pipeline_lease.slot_id); + }); + if (!all_workers_idle) { + result.busy_target_worker_ids.clear(); + for (WorkerThread *worker : workers) { + if (!worker->can_stage(s.pipeline_lease.slot_id)) { + result.busy_target_worker_ids.push_back(worker->worker_id()); + } + } + } + } if (!all_workers_idle) { for (size_t i = 0; i < workers.size(); ++i) { TaskSlot single_head = INVALID_SLOT; @@ -658,8 +763,11 @@ Scheduler::NextLevelGroupDispatchResult Scheduler::dispatch_next_level_group(con commit_group_vectors_locked(s, prepared); reset_group_state_locked(s, GroupMemberState::RUNNING); } + if (supports_prepared_barrier) pending_group_barriers_.insert(slot); for (int32_t i = 0; i < group_size; ++i) { - dispatch_claimed(workers[static_cast(i)], WorkerDispatch{slot, i}, /*prepared=*/false); + WorkerDispatch dispatch{slot, i}; + dispatch.require_native_prepare = supports_prepared_barrier; + dispatch_claimed(workers[static_cast(i)], dispatch, /*prepared=*/supports_prepared_barrier); } } return {}; diff --git a/src/common/hierarchical/scheduler.h b/src/common/hierarchical/scheduler.h index 529f292b31..8daab8ae49 100644 --- a/src/common/hierarchical/scheduler.h +++ b/src/common/hierarchical/scheduler.h @@ -173,6 +173,9 @@ class Scheduler { // resets it before the thread exists. Any reader added off sched_thread_ // needs a lock the two paths do not currently share. std::optional reservation_stall_episode_; + // Full-rank local groups are staged behind one all-member native-prepare + // barrier and activated together. Groups do not cross the whole-run FIFO. + std::unordered_set pending_group_barriers_; void run(); void on_task_complete(const WorkerCompletion &completion); @@ -180,6 +183,7 @@ class Scheduler { void try_consume(TaskSlot slot); void dispatch_ready(); + void progress_group_barriers(); void dispatch_claimed(WorkerThread *worker, WorkerDispatch dispatch, bool prepared); void dispatch_preparable_next_level_singles(); NextLevelGroupDispatchResult dispatch_next_level_group(const std::optional &run_snapshot); diff --git a/src/common/hierarchical/worker_manager.cpp b/src/common/hierarchical/worker_manager.cpp index 11fd94dae8..3a4e646c75 100644 --- a/src/common/hierarchical/worker_manager.cpp +++ b/src/common/hierarchical/worker_manager.cpp @@ -32,6 +32,7 @@ #include "common/host_span_names.h" #include "common/host_span_scope.h" +#include "mpi_group_mailbox.h" #include "ring.h" namespace { @@ -91,7 +92,8 @@ trace_dispatch_attrs(RunId run_id, const WorkerDispatch &dispatch, const WorkerE attrs << "run_id=" << run_id << " task_slot=" << dispatch.task_slot << " group_index=" << dispatch.group_index << " worker_id=" << caps.worker_id << " dispatch_id=" << dispatch.dispatch_id << " endpoint_kind=" << endpoint_kind_name(caps.kind) - << " prepare_only=" << static_cast(dispatch.prepare_only) << " role=" << role; + << " prepare_only=" << static_cast(dispatch.prepare_only) + << " require_native_prepare=" << static_cast(dispatch.require_native_prepare) << " role=" << role; return attrs.str(); } @@ -206,6 +208,8 @@ void WorkerEndpoint::submit_progress(Ring *, const WorkerDispatch &) { } bool WorkerEndpoint::poll_progress(WorkerEndpointProgress &) { return false; } bool WorkerEndpoint::activate_progress(RunId) { return false; } + +bool WorkerEndpoint::cancel_progress(RunId, const std::string &) { return false; } void WorkerEndpoint::request_progress_stop() noexcept {} void WorkerEndpoint::report_progress_error(const std::string &) { request_progress_stop(); } bool WorkerEndpoint::report_submission_error(const WorkerDispatch &, const std::string &reason) { @@ -287,6 +291,13 @@ void LocalMailboxEndpoint::write_mailbox_state(MailboxState s, char *frame) { #else __atomic_store(ptr, &v, __ATOMIC_RELEASE); #endif + notify_child(); +} + +void LocalMailboxEndpoint::notify_child() noexcept { + auto *notification = reinterpret_cast(mbox() + MAILBOX_OFF_NOTIFICATION); + (void)__atomic_add_fetch(notification, 1, __ATOMIC_RELEASE); + mpi_group_mailbox::wake_word(notification); } bool mailbox_compare_exchange_state(char *frame, MailboxState expected, MailboxState desired) noexcept { @@ -402,19 +413,21 @@ void WorkerThread::dispatch_prepared(WorkerDispatch d) { complete_unpublished(d, "WorkerThread::dispatch_prepared: dispatch has no run identity"); return; } - bool staged_lane_occupied = false; + bool staging_unavailable = false; { std::lock_guard lane_lk(lane_mu_); + const LaneState &active = lane(LaneKind::ACTIVE); LaneState &staged = lane(LaneKind::STAGED); - if (staged.occupied) { - staged_lane_occupied = true; + if (staged.occupied || (active.occupied && active.pipeline_slot_id == slot->pipeline_lease.slot_id)) { + staging_unavailable = true; } else { staged.occupied = true; staged.run_id = slot->run_id; + staged.pipeline_slot_id = slot->pipeline_lease.slot_id; } } - if (staged_lane_occupied) { - complete_unpublished(d, "WorkerThread::dispatch_prepared: worker already owns a staged run"); + if (staging_unavailable) { + complete_unpublished(d, "WorkerThread::dispatch_prepared: worker cannot stage on the requested pipeline slot"); return; } SubmitDispatchResult result; @@ -462,6 +475,8 @@ WorkerThread::submit_dispatch(WorkerDispatch d, LaneKind lane_kind, RunId expect return SubmitDispatchResult::STAGED_IDENTITY_CHANGED; } dispatch_lane.dispatch_id = d.dispatch_id; + const TaskSlotState *state = ring_ == nullptr ? nullptr : ring_->slot_state(d.task_slot); + if (state != nullptr) dispatch_lane.pipeline_slot_id = state->pipeline_lease.slot_id; } ++next_dispatch_id_; inflight_.fetch_add(1, std::memory_order_release); @@ -504,26 +519,58 @@ bool WorkerThread::activate_prepared(RunId run_id) { std::lock_guard lane_lk(lane_mu_); LaneState &active = lane(LaneKind::ACTIVE); LaneState &staged = lane(LaneKind::STAGED); - if (!staged.occupied || staged.run_id != run_id || staged.dispatch_id == 0 || active.occupied) { + if (!staged.occupied || staged.run_id != run_id || staged.dispatch_id == 0 || staged.cancellation_requested) { return false; } + // A host-staged successor must not be activated while the predecessor + // still owns this endpoint. Local chip completion is earlier than the + // whole L3 run's retire fence; arming here lets the successor enter the + // native runtime before parent resources are safe to reuse. + if (active.occupied) return false; active = staged; active.activation_requested = true; staged = {}; return true; } +bool WorkerThread::cancel_prepared(RunId run_id, const std::string &reason) { + if (run_id == INVALID_RUN_ID) return false; + std::lock_guard admission_lk(admission_mu_); + if (shutdown_.load(std::memory_order_acquire)) return false; + { + std::lock_guard lane_lk(lane_mu_); + LaneState &staged = lane(LaneKind::STAGED); + if (!staged.occupied || staged.run_id != run_id || staged.dispatch_id == 0) return false; + staged.cancellation_requested = true; + } + return endpoint_->cancel_progress(run_id, reason); +} + bool WorkerThread::has_staged_run(RunId run_id) const { std::lock_guard lane_lk(lane_mu_); const LaneState &staged = lane(LaneKind::STAGED); return staged.occupied && staged.run_id == run_id; } +bool WorkerThread::prepared_ready(RunId run_id) const { + std::lock_guard lane_lk(lane_mu_); + const LaneState &staged = lane(LaneKind::STAGED); + return staged.occupied && staged.run_id == run_id && staged.dispatch_id != 0 && staged.staged_ready && + !staged.cancellation_requested; +} + bool WorkerThread::can_stage() const { std::lock_guard lane_lk(lane_mu_); return caps().supports_frame_staging && !lane(LaneKind::STAGED).occupied; } +bool WorkerThread::can_stage(uint32_t pipeline_slot_id) const { + std::lock_guard lane_lk(lane_mu_); + const LaneState &active = lane(LaneKind::ACTIVE); + return caps().supports_frame_staging && !lane(LaneKind::STAGED).occupied && + (!active.occupied || active.pipeline_slot_id != pipeline_slot_id); +} + bool WorkerThread::idle() const { std::lock_guard lane_lk(lane_mu_); return !lane(LaneKind::ACTIVE).occupied; @@ -585,14 +632,23 @@ void WorkerThread::progress() { { std::lock_guard lane_lk(lane_mu_); const LaneState &active = lane(LaneKind::ACTIVE); - if (active.occupied && active.activation_requested) activated = active.run_id; + const LaneState &staged = lane(LaneKind::STAGED); + if (active.occupied && active.activation_requested) { + activated = active.run_id; + } else if (staged.occupied && staged.activation_requested) { + activated = staged.run_id; + } } if (activated != INVALID_RUN_ID) { try { if (endpoint_->activate_progress(activated)) { std::lock_guard lane_lk(lane_mu_); - LaneState &active = lane(LaneKind::ACTIVE); - if (active.occupied && active.run_id == activated) active.activation_requested = false; + for (LaneState &dispatch_lane : lanes_) { + if (dispatch_lane.occupied && dispatch_lane.run_id == activated) { + dispatch_lane.activation_requested = false; + break; + } + } } } catch (const std::exception &e) { fail_progress_driver(std::string("activate_progress failed: ") + e.what()); @@ -636,9 +692,12 @@ void WorkerThread::fail_submission(const WorkerDispatch &dispatch, const std::st void WorkerThread::finish_progress_dispatch(const WorkerEndpointProgress &progress) { const WorkerDispatch &dispatch = progress.dispatch; if (progress.kind == WorkerProgressKind::FRAME_STAGED) { - // The endpoint already owns the prepared frame. This cursor-only event - // keeps the progress poll moving; acceptance, completion, and inflight - // ownership intentionally remain unchanged until activation/terminal. + // Publish readiness to the group barrier. Acceptance, completion, and + // inflight ownership intentionally remain unchanged until + // activation/terminal. + std::lock_guard lane_lk(lane_mu_); + LaneState &staged = lane(LaneKind::STAGED); + if (staged.occupied && staged.dispatch_id == progress.dispatch.dispatch_id) staged.staged_ready = true; return; } @@ -700,11 +759,12 @@ void WorkerThread::finish_progress_dispatch(const WorkerEndpointProgress &progre on_complete_(std::move(completion)); { std::lock_guard lane_lk(lane_mu_); - for (LaneState &dispatch_lane : lanes_) { - if (dispatch_lane.occupied && dispatch_lane.dispatch_id == dispatch.dispatch_id) { - dispatch_lane = {}; - break; - } + LaneState &active = lane(LaneKind::ACTIVE); + LaneState &staged = lane(LaneKind::STAGED); + if (active.occupied && active.dispatch_id == dispatch.dispatch_id) { + active = {}; + } else if (staged.occupied && staged.dispatch_id == dispatch.dispatch_id) { + staged = {}; } } inflight_.fetch_sub(1, std::memory_order_acq_rel); @@ -721,6 +781,9 @@ void WorkerThread::fail_progress_driver(const std::string &reason) noexcept { void LocalMailboxEndpoint::submit_progress(Ring *ring, const WorkerDispatch &dispatch) { if (ring == nullptr) throw std::invalid_argument("LocalMailboxEndpoint::submit_progress: null ring"); + if (dispatch.require_native_prepare && !dispatch.prepare_only) { + throw std::invalid_argument("LocalMailboxEndpoint::submit_progress: prepared mode requires staged dispatch"); + } TaskSlotState *slot_state = ring->slot_state(dispatch.task_slot); if (slot_state == nullptr) throw std::out_of_range("LocalMailboxEndpoint::submit_progress: invalid task slot"); TaskSlotState &state = *slot_state; @@ -813,7 +876,11 @@ void LocalMailboxEndpoint::submit_progress(Ring *ring, const WorkerDispatch &dis record.run_id = state.run_id; record.slot_id = slot_id; record.generation = state.pipeline_lease.generation; - write_mailbox_state(dispatch.prepare_only ? MailboxState::PREPARE_READY : MailboxState::TASK_READY, frame); + const MailboxState ready_state = + dispatch.require_native_prepare ? + MailboxState::NATIVE_PREPARE_READY : + (dispatch.prepare_only ? MailboxState::PREPARE_READY : MailboxState::TASK_READY); + write_mailbox_state(ready_state, frame); } bool LocalMailboxEndpoint::frame_identity_matches(const FrameRecord &record, const char *frame) const { @@ -842,10 +909,31 @@ bool LocalMailboxEndpoint::try_publish_activation(FrameRecord &record, char *fra } if (mailbox_compare_exchange_state(frame, MailboxState::FRAME_STAGED, MailboxState::ACTIVATE)) { record.activation_published = true; + notify_child(); } return record.activation_published; } +bool LocalMailboxEndpoint::try_publish_cancellation(FrameRecord &record, char *frame) { + if (!record.cancellation_requested || record.cancellation_published) return record.cancellation_published; + if (!frame_identity_matches(record, frame)) { + poison_progress("stale staged frame identity before cancellation"); + return false; + } + int32_t error_code = -1; + std::memcpy(frame + MAILBOX_OFF_ERROR, &error_code, sizeof(error_code)); + std::memset(frame + MAILBOX_OFF_ERROR_MSG, 0, MAILBOX_ERROR_MSG_SIZE); + if (!record.cancellation_reason.empty()) { + const size_t n = std::min(record.cancellation_reason.size(), MAILBOX_ERROR_MSG_SIZE - 1); + std::memcpy(frame + MAILBOX_OFF_ERROR_MSG, record.cancellation_reason.data(), n); + } + if (mailbox_compare_exchange_state(frame, MailboxState::FRAME_STAGED, MailboxState::ABANDON)) { + record.cancellation_published = true; + notify_child(); + } + return record.cancellation_published; +} + void LocalMailboxEndpoint::poison_progress(const std::string &reason) { if (endpoint_poisoned_) return; endpoint_poisoned_ = true; @@ -932,7 +1020,10 @@ bool LocalMailboxEndpoint::poll_progress(WorkerEndpointProgress &progress) { poison_progress("stale frame identity at endpoint staging"); break; } - if (record.activation_requested) { + if (record.cancellation_requested) { + (void)try_publish_cancellation(record, frame); + if (endpoint_poisoned_) break; + } else if (record.activation_requested) { (void)try_publish_activation(record, frame); if (endpoint_poisoned_) break; } @@ -981,7 +1072,8 @@ bool LocalMailboxEndpoint::poll_progress(WorkerEndpointProgress &progress) { } if (state != MailboxState::TASK_READY && state != MailboxState::PREPARE_READY && - state != MailboxState::ACTIVATE && state != MailboxState::TASK_LAUNCHED) { + state != MailboxState::NATIVE_PREPARE_READY && state != MailboxState::ACTIVATE && + state != MailboxState::TASK_LAUNCHED && state != MailboxState::ABANDON) { poison_progress("task frame entered an invalid state " + std::to_string(static_cast(state))); break; } @@ -1032,6 +1124,23 @@ bool LocalMailboxEndpoint::activate_progress(RunId run_id) { return false; } +bool LocalMailboxEndpoint::cancel_progress(RunId run_id, const std::string &reason) { + std::lock_guard lk(progress_mu_); + if (endpoint_poisoned_) return false; + for (size_t index = 0; index < task_frame_count_; ++index) { + FrameRecord &record = frames_[index]; + if (!record.occupied || !record.dispatch.prepare_only || record.run_id != run_id) continue; + record.cancellation_requested = true; + record.cancellation_reason = reason; + char *frame = task_frame(index); + if (read_mailbox_state(frame) == MailboxState::FRAME_STAGED) { + (void)try_publish_cancellation(record, frame); + } + return true; + } + return false; +} + void LocalMailboxEndpoint::request_progress_stop() noexcept { shutdown_child(); } void LocalMailboxEndpoint::report_progress_error(const std::string &reason) { diff --git a/src/common/hierarchical/worker_manager.h b/src/common/hierarchical/worker_manager.h index 4bb25cf416..07e3ea7923 100644 --- a/src/common/hierarchical/worker_manager.h +++ b/src/common/hierarchical/worker_manager.h @@ -79,6 +79,8 @@ enum class MailboxState : int32_t { TASK_FAILED = 10, ACTIVATE = 11, PREPARE_READY = 12, + ABANDON = 13, + NATIVE_PREPARE_READY = 14, }; enum class MailboxPreparationDisposition : int32_t { @@ -98,7 +100,7 @@ static constexpr size_t MAILBOX_TASK_FRAME_COUNT = 2; static constexpr size_t MAILBOX_CONTROL_FRAME = 0; static constexpr size_t MAILBOX_FIRST_TASK_FRAME = 1; static constexpr size_t MAILBOX_SIZE = MAILBOX_FRAME_SIZE * (1 + MAILBOX_TASK_FRAME_COUNT); -static constexpr uint32_t MAILBOX_TASK_PROTOCOL_VERSION = 3; +static constexpr uint32_t MAILBOX_TASK_PROTOCOL_VERSION = 5; // Error message region lives at the mailbox tail. 256 B of headroom is // enough for `: ` produced by the child-side @@ -156,6 +158,16 @@ static constexpr ptrdiff_t MAILBOX_OFF_FRAME_DISPATCH_ID = MAILBOX_OFF_ACCEPTED // matter what state word a concurrent control command leaves behind. static constexpr ptrdiff_t MAILBOX_OFF_SHUTDOWN = MAILBOX_OFF_FRAME_PROTOCOL - 8; static constexpr int32_t MAILBOX_SHUTDOWN_REQUESTED = 1; +// Mailbox-wide parent-to-child event generation. A child cannot futex-wait +// on both task-frame state words (and the control state) at once, so every +// parent publication increments this shared word and wakes it. Keeping a +// generation, rather than wake-only notification, closes the lost-wakeup +// race between the child's final frame scan and FUTEX_WAIT. +static constexpr ptrdiff_t MAILBOX_OFF_NOTIFICATION = MAILBOX_OFF_SHUTDOWN + 4; +static_assert( + MAILBOX_OFF_NOTIFICATION + static_cast(sizeof(int32_t)) <= MAILBOX_OFF_FRAME_PROTOCOL, + "mailbox notification word must fit in the trailer before frame protocol" +); 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); @@ -285,6 +297,7 @@ struct WorkerDispatch { int32_t group_index{0}; uint64_t dispatch_id{0}; bool prepare_only{false}; + bool require_native_prepare{false}; }; enum class WorkerProgressKind : int32_t { @@ -330,6 +343,7 @@ class WorkerEndpoint { virtual void submit_progress(Ring *ring, const WorkerDispatch &dispatch); virtual bool poll_progress(WorkerEndpointProgress &progress); virtual bool activate_progress(RunId run_id); + virtual bool cancel_progress(RunId run_id, const std::string &reason); virtual void request_progress_stop() noexcept; // Called when the owning progress driver catches an endpoint exception. // Implementations must turn already-published work into terminal progress @@ -400,6 +414,7 @@ class LocalMailboxEndpoint : public WorkerEndpoint { void submit_progress(Ring *ring, const WorkerDispatch &dispatch) override; bool poll_progress(WorkerEndpointProgress &progress) override; bool activate_progress(RunId run_id) override; + bool cancel_progress(RunId run_id, const std::string &reason) override; void request_progress_stop() noexcept override; void report_progress_error(const std::string &reason) override; bool report_submission_error(const WorkerDispatch &dispatch, const std::string &reason) override; @@ -470,6 +485,7 @@ class LocalMailboxEndpoint : public WorkerEndpoint { char *mbox() const { return static_cast(mailbox_); } MailboxState read_mailbox_state(const char *frame = nullptr) const; void write_mailbox_state(MailboxState s, char *frame = nullptr); + void notify_child() noexcept; // Sticky launch acceptance, cleared before this endpoint reuses a task // frame. See MAILBOX_OFF_ACCEPTED. bool read_task_accepted(const char *frame = nullptr) const; @@ -483,6 +499,9 @@ class LocalMailboxEndpoint : public WorkerEndpoint { bool accepted_reported{false}; bool activation_requested{false}; bool activation_published{false}; + bool cancellation_requested{false}; + bool cancellation_published{false}; + std::string cancellation_reason; WorkerDispatch dispatch{}; RunId run_id{INVALID_RUN_ID}; uint64_t slot_id{0}; @@ -492,6 +511,7 @@ class LocalMailboxEndpoint : public WorkerEndpoint { bool frame_identity_matches(const FrameRecord &record, const char *frame) const; bool try_publish_activation(FrameRecord &record, char *frame); + bool try_publish_cancellation(FrameRecord &record, char *frame); void poison_progress(const std::string &reason); bool poisoned_progress_quiesced(); WorkerCompletion poisoned_completion(const FrameRecord &record) const; @@ -537,7 +557,9 @@ class WorkerThread { // still needs its conservative terminal fallback. void complete_unpublished(WorkerDispatch d, const std::string &error_message); bool has_staged_run(RunId run_id) const; + bool prepared_ready(RunId run_id) const; bool activate_prepared(RunId run_id); + bool cancel_prepared(RunId run_id, const std::string &reason); void progress(); // The active lane and staged-successor lane are intentionally distinct. @@ -545,6 +567,7 @@ class WorkerThread { // dispatchable on the same device. bool idle() const; bool can_stage() const; + bool can_stage(uint32_t pipeline_slot_id) const; bool busy() const { return inflight_.load(std::memory_order_acquire) != 0; } const WorkerEndpointCaps &caps() const; int32_t worker_id() const; @@ -640,8 +663,11 @@ class WorkerThread { struct LaneState { bool occupied{false}; bool activation_requested{false}; + bool staged_ready{false}; + bool cancellation_requested{false}; RunId run_id{INVALID_RUN_ID}; uint64_t dispatch_id{0}; + uint32_t pipeline_slot_id{UINT32_MAX}; }; Ring *ring_{nullptr}; diff --git a/src/common/worker/chip_run_lane.cpp b/src/common/worker/chip_run_lane.cpp index ed20b669a3..44c76d6f2b 100644 --- a/src/common/worker/chip_run_lane.cpp +++ b/src/common/worker/chip_run_lane.cpp @@ -14,12 +14,14 @@ #include "chip_worker.h" #include +#include #include #include #include #include #include #include +#include #include struct ChipRunState { @@ -40,13 +42,26 @@ struct ChipRunState { bool activated{false}; bool crossed_launch_fence{false}; bool depth_one_fallback{false}; + bool wait_in_progress{false}; std::exception_ptr error; }; struct ChipRunLaneState { explicit ChipRunLaneState(ChipWorker &worker) : worker(&worker), - generations(worker.pipeline_depth(), 0) {} + generations(worker.pipeline_depth(), 0), + progress_worker([this]() { + progress_loop(); + }) {} + + ~ChipRunLaneState() { + { + std::lock_guard lock(mu); + stopping = true; + } + cv.notify_all(); + if (progress_worker.joinable()) progress_worker.join(); + } void require_usable() const { if (closed) throw std::runtime_error("chip run lane is closed"); @@ -70,6 +85,16 @@ struct ChipRunLaneState { !successor_config.diagnostics_any() && predecessor.phase == ChipRunState::Phase::LAUNCHED; } + static bool permits_resident_completion(const ChipRunState &run) noexcept { + // Some diagnostics, notably host_build_graph dep_gen, keep capture + // state on the submitter thread and export it during finalize. Keep + // those depth-one runs on that same thread. Direct runs also retain + // their caller-owned lifecycle: their second submission must be able + // to stage against an unfinalized predecessor. The resident owner is + // for scheduler-leased asynchronous serving runs only. + return run.pipeline_leased && !run.config.diagnostics_any(); + } + bool permits_native_successor(const ChipRunState &predecessor, const ChipRunState &successor) const { return permits_native_successor(predecessor, successor.config); } @@ -172,6 +197,7 @@ struct ChipRunLaneState { return true; } if (target->phase != ChipRunState::Phase::LAUNCHED) return false; + if (target->wait_in_progress) return false; try { if (!worker->poll_native_run(target->native_run)) return false; @@ -188,6 +214,64 @@ struct ChipRunLaneState { return true; } + void progress_loop() noexcept { + std::unique_lock lock(mu); + while (true) { + cv.wait(lock, [this]() { + return stopping || (!fifo.empty() && fifo.front()->phase == ChipRunState::Phase::LAUNCHED && + !fifo.front()->wait_in_progress && permits_resident_completion(*fifo.front())); + }); + if (stopping) return; + const auto target = fifo.front(); + target->wait_in_progress = true; + lock.unlock(); + std::exception_ptr native_error; + try { + worker->wait_native_run(target->native_run); + } catch (...) { + native_error = std::current_exception(); + } + // This thread owns the native token until terminal publication. + // Native wait and finalization stay outside the lane mutex; the + // mutex protects terminal publication and successor launch. + try { + worker->finalize_native_run(target->native_run); + } catch (...) { + if (native_error == nullptr) native_error = std::current_exception(); + } + lock.lock(); + target->wait_in_progress = false; + if (target->phase == ChipRunState::Phase::LAUNCHED) { + if (native_error != nullptr) { + target->error = native_error; + poison_with(native_error); + } + target->phase = ChipRunState::Phase::TERMINAL; + if (!fifo.empty() && fifo.front() == target) fifo.pop_front(); + launch_front(); + } + cv.notify_all(); + } + } + + bool block_front_on_caller() noexcept { + if (fifo.empty()) return false; + const auto run = fifo.front(); + if (run->phase != ChipRunState::Phase::LAUNCHED || run->wait_in_progress) { + return false; + } + try { + worker->wait_native_run(run->native_run); + } catch (...) { + run->error = std::current_exception(); + poison_with(run->error); + } + finish(run); + launch_front(); + cv.notify_all(); + return true; + } + void drain_front() noexcept { if (fifo.empty()) return; const auto run = fifo.front(); @@ -220,35 +304,16 @@ struct ChipRunLaneState { } } - // Block on the device for the launched front, for waiters with no deadline - // to bound them. Only the front can be LAUNCHED, so its completion is what - // lets any waiter in the FIFO advance. Re-polling instead would hold a core - // for the whole run — the case codestyle rule 5 sends to a wakeup primitive - // rather than a busy loop. Reports whether it actually blocked, so a caller - // that cannot be unblocked this way does not spin on it. - bool block_on_front() noexcept { - if (fifo.empty()) return false; - const auto front = fifo.front(); - if (front->phase != ChipRunState::Phase::LAUNCHED) return false; - try { - worker->wait_native_run(front->native_run); - } catch (...) { - const std::exception_ptr wait_error = std::current_exception(); - if (front->error == nullptr) front->error = wait_error; - poison_with(front->error); - } - finish(front); - launch_front(); - return true; - } - ChipWorker *worker; mutable std::mutex mu; + std::condition_variable cv; std::deque> fifo; std::vector generations; uint64_t direct_generation{0}; std::exception_ptr poison; bool closed{false}; + bool stopping{false}; + std::thread progress_worker; }; ChipRun::ChipRun(std::shared_ptr lane, std::shared_ptr run) : @@ -258,26 +323,69 @@ ChipRun::ChipRun(std::shared_ptr lane, std::shared_ptr lk(lane_->mu); + lane_->cv.notify_all(); return lane_->progress(run_); } bool ChipRun::wait_until(Deadline deadline) { if (lane_ == nullptr || run_ == nullptr) throw std::runtime_error("empty ChipRun handle"); const bool unbounded = deadline == Deadline::max(); - while (true) { - { - std::lock_guard lk(lane_->mu); - if (lane_->progress(run_)) { - ChipRunLaneState::rethrow_run_error(run_); - return true; + if (run_->config.diagnostics_any() || !run_->pipeline_leased) { + // Diagnostic finalizers can consume thread-affine capture state. Match + // the pre-resident direct lifecycle: poll (for a bounded wait), or + // block and finalize on the submitter/waiter thread for an unbounded + // wait. This also preserves direct depth-two successor staging. + while (true) { + { + std::lock_guard lk(lane_->mu); + if (lane_->progress(run_)) { + ChipRunLaneState::rethrow_run_error(run_); + return true; + } + if (unbounded && lane_->block_front_on_caller()) continue; } - // An unbounded waiter has no deadline to end its loop, so polling - // here would spin for the whole run. Block on the device instead; - // if nothing is blockable yet the poll loop below still applies. - if (unbounded && lane_->block_on_front()) continue; + if (Clock::now() >= deadline) return false; + std::this_thread::yield(); } - if (Clock::now() >= deadline) return false; } + std::unique_lock lk(lane_->mu); + lane_->cv.notify_all(); + const auto terminal = [this]() { + return run_->phase == ChipRunState::Phase::TERMINAL; + }; + if (unbounded) { + lane_->cv.wait(lk, terminal); + } else if (!lane_->cv.wait_until(lk, deadline, terminal)) { + return false; + } + ChipRunLaneState::rethrow_run_error(run_); + return true; +} + +void ChipRun::prepare() { + if (lane_ == nullptr || run_ == nullptr) throw std::runtime_error("empty ChipRun handle"); + std::lock_guard lk(lane_->mu); + if (run_->phase == ChipRunState::Phase::TERMINAL) { + ChipRunLaneState::rethrow_run_error(run_); + return; + } + if (run_->activated || run_->phase == ChipRunState::Phase::LAUNCHED) { + throw std::logic_error("cannot prepare a ChipRun after activation"); + } + if (run_->phase == ChipRunState::Phase::PREPARED) return; + if (lane_->fifo.empty() || lane_->fifo.front() != run_) { + throw std::logic_error("only the front ChipRun can be prepared without activation"); + } + try { + lane_->prepare(run_); + } catch (...) { + run_->error = std::current_exception(); + run_->phase = ChipRunState::Phase::TERMINAL; + lane_->fifo.pop_front(); + lane_->cv.notify_all(); + throw; + } + lane_->cv.notify_all(); } void ChipRun::activate() { @@ -292,6 +400,8 @@ void ChipRun::activate() { if (lane_->fifo.size() == 2 && lane_->fifo.front() == run_) { lane_->prepare_successor_if_eligible(lane_->fifo.back()); } + lane_->cv.notify_all(); + ChipRunLaneState::rethrow_run_error(run_); } void ChipRun::abandon() { @@ -404,6 +514,7 @@ ChipRun ChipRunLane::submit( auto it = std::find(state_->fifo.begin(), state_->fifo.end(), run); if (it != state_->fifo.end()) state_->fifo.erase(it); } + state_->cv.notify_all(); return ChipRun(state_, std::move(run)); } @@ -411,7 +522,7 @@ ChipRun ChipRunLane::submit( int32_t callable_id, const ChipStorageTaskArgs &args, const CallConfig &config, volatile int32_t *accepted_state, int32_t accepted_value ) { - std::lock_guard lk(state_->mu); + std::unique_lock lk(state_->mu); state_->require_usable(); if (state_->generations.empty()) throw std::runtime_error("chip run lane has no runtime slots"); @@ -424,6 +535,16 @@ ChipRun ChipRunLane::submit( const bool has_successor_capacity = state_->fifo.size() == 1 && state_->permits_native_successor(*state_->fifo.front(), config); if (has_successor_capacity) break; + if (state_->fifo.front()->wait_in_progress) { + state_->cv.wait(lk, [this, &config]() { + if (state_->fifo.empty()) return true; + if (state_->fifo.size() == 1 && state_->permits_native_successor(*state_->fifo.front(), config)) { + return true; + } + return !state_->fifo.front()->wait_in_progress; + }); + continue; + } state_->drain_front(); state_->require_usable(); state_->launch_front(); @@ -461,25 +582,45 @@ ChipRun ChipRunLane::submit( } else { state_->prepare_successor_if_eligible(run); } + state_->cv.notify_all(); return ChipRun(state_, std::move(run)); } void ChipRunLane::drain() { - std::lock_guard lk(state_->mu); - while (!state_->fifo.empty()) + std::unique_lock lk(state_->mu); + while (!state_->fifo.empty()) { + if (state_->fifo.front()->wait_in_progress) { + state_->cv.wait(lk, [this]() { + return state_->fifo.empty() || !state_->fifo.front()->wait_in_progress; + }); + continue; + } state_->drain_front(); + } if (state_->poison != nullptr) std::rethrow_exception(state_->poison); } void ChipRunLane::close() { - std::lock_guard lk(state_->mu); + std::unique_lock lk(state_->mu); if (state_->closed) { if (state_->poison != nullptr) std::rethrow_exception(state_->poison); return; } - while (!state_->fifo.empty()) + while (!state_->fifo.empty()) { + if (state_->fifo.front()->wait_in_progress) { + state_->cv.wait(lk, [this]() { + return state_->fifo.empty() || !state_->fifo.front()->wait_in_progress; + }); + continue; + } state_->drain_front(); + } state_->closed = true; + state_->stopping = true; + state_->cv.notify_all(); + lk.unlock(); + if (state_->progress_worker.joinable()) state_->progress_worker.join(); + lk.lock(); if (state_->poison != nullptr) std::rethrow_exception(state_->poison); } diff --git a/src/common/worker/chip_run_lane.h b/src/common/worker/chip_run_lane.h index 5464c452cb..b26242c749 100644 --- a/src/common/worker/chip_run_lane.h +++ b/src/common/worker/chip_run_lane.h @@ -37,6 +37,7 @@ class ChipRun { bool done(); bool wait_until(Deadline deadline); + void prepare(); void activate(); void abandon(); diff --git a/tests/ut/cpp/hierarchical/test_chip_run_lane.cpp b/tests/ut/cpp/hierarchical/test_chip_run_lane.cpp index 47c9d6737a..d3d0a0f826 100644 --- a/tests/ut/cpp/hierarchical/test_chip_run_lane.cpp +++ b/tests/ut/cpp/hierarchical/test_chip_run_lane.cpp @@ -11,6 +11,7 @@ #include #include +#include #include #include #include @@ -39,7 +40,7 @@ namespace { std::unordered_map g_slots; -std::array g_complete{}; +std::array, 2> g_complete{}; std::array g_prepare_rc{}; std::array g_launch_rc{}; std::array g_poll_rc{}; @@ -56,9 +57,16 @@ bool g_supports_successor{true}; std::mutex g_wait_mu; std::condition_variable g_wait_cv; bool g_wait_entered{false}; -bool g_release_wait{true}; +bool g_release_wait{false}; +std::thread::id g_prepare_thread; +std::thread::id g_finalize_thread; std::vector g_events; +void record_event(std::string event) { + std::lock_guard lock(g_wait_mu); + g_events.push_back(std::move(event)); +} + uint32_t slot_of(void *runtime) { return g_slots.at(runtime); } int prepare_run( @@ -66,7 +74,8 @@ int prepare_run( ) { EXPECT_EQ(slot_of(runtime), descriptor->pipeline_slot); g_complete[descriptor->pipeline_slot] = false; - g_events.push_back("prepare" + std::to_string(descriptor->pipeline_slot)); + g_prepare_thread = std::this_thread::get_id(); + record_event("prepare" + std::to_string(descriptor->pipeline_slot)); ++g_prepare_count[descriptor->pipeline_slot]; if (g_reject_first_prepare[descriptor->pipeline_slot] && g_prepare_count[descriptor->pipeline_slot] == 1) { return PTO_RUNTIME_ERR_PREPARED_INCOMPATIBLE; @@ -76,7 +85,7 @@ int prepare_run( int launch_run(void *, void *runtime) { const uint32_t slot = slot_of(runtime); - g_events.push_back("launch" + std::to_string(slot)); + record_event("launch" + std::to_string(slot)); return g_launch_rc[slot]; } @@ -90,13 +99,13 @@ int poll_run(void *, void *runtime) { int wait_run(void *, void *runtime) { const uint32_t slot = slot_of(runtime); - g_events.push_back("wait" + std::to_string(slot)); + record_event("wait" + std::to_string(slot)); { std::unique_lock lk(g_wait_mu); g_wait_entered = true; g_wait_cv.notify_all(); - g_wait_cv.wait(lk, [] { - return g_release_wait; + g_wait_cv.wait(lk, [slot] { + return g_release_wait || g_complete[slot].load(); }); } g_complete[slot] = true; @@ -105,7 +114,8 @@ int wait_run(void *, void *runtime) { int finalize_run(void *, void *runtime) { const uint32_t slot = slot_of(runtime); - g_events.push_back("finalize" + std::to_string(slot)); + g_finalize_thread = std::this_thread::get_id(); + record_event("finalize" + std::to_string(slot)); return g_finalize_rc[slot]; } @@ -113,7 +123,8 @@ int supports_successor(void *) { return g_supports_successor ? 1 : 0; } void prime_worker(ChipWorker &worker) { g_slots.clear(); - g_complete = {}; + for (auto &complete : g_complete) + complete.store(false); g_prepare_rc = {}; g_launch_rc = {}; g_poll_rc = {}; @@ -124,12 +135,17 @@ void prime_worker(ChipWorker &worker) { g_poll_count = 0; g_poll_completes_after = 0; g_supports_successor = true; + g_prepare_thread = {}; + g_finalize_thread = {}; { std::lock_guard lk(g_wait_mu); g_wait_entered = false; - g_release_wait = true; + g_release_wait = false; + } + { + std::lock_guard lock(g_wait_mu); + g_events.clear(); } - g_events.clear(); worker.initialized_ = true; worker.pipeline_contract_ = {PTO_PIPELINE_CONTRACT_ABI_VERSION, 0, 2, {}}; worker.runtime_bufs_.emplace_back(64, alignof(std::max_align_t)); @@ -150,6 +166,41 @@ ChipRun submit(ChipRunLane &lane, uint64_t run_id, uint32_t slot, bool activate return lane.submit(1, args, CallConfig{}, PipelineSlotLease{slot, 0, run_id}, run_id, run_id, nullptr, 0, activate); } +void complete(uint32_t slot) { + // Publish the predicate while holding the same mutex used by wait_run. + // Atomic storage alone does not prevent a notification from landing + // between the predicate check and the condition-variable wait. + { + std::lock_guard lock(g_wait_mu); + g_complete[slot] = true; + } + g_wait_cv.notify_all(); +} + +bool await_done(ChipRun &run) { return run.wait_until(std::chrono::steady_clock::now() + std::chrono::seconds(5)); } + +bool await_terminal(ChipRun &run) { + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(1); + while (!run.done() && std::chrono::steady_clock::now() < deadline) { + std::this_thread::yield(); + } + return run.done(); +} + +std::vector events_without_wait() { + std::lock_guard lock(g_wait_mu); + std::vector events; + std::copy_if(g_events.begin(), g_events.end(), std::back_inserter(events), [](const std::string &event) { + return event.rfind("wait", 0) != 0; + }); + return events; +} + +bool has_event(const std::string &event) { + std::lock_guard lock(g_wait_mu); + return std::find(g_events.begin(), g_events.end(), event) != g_events.end(); +} + } // namespace TEST(ChipRunLaneTest, OwnsFifoPreparationAndLaunch) { @@ -160,18 +211,43 @@ TEST(ChipRunLaneTest, OwnsFifoPreparationAndLaunch) { ChipRun first = submit(lane, 101, 0); ChipRun second = submit(lane, 102, 1, false); EXPECT_EQ(second.preparation_disposition(), ChipRunPreparationDisposition::NATIVE_PREPARED); - EXPECT_EQ(g_events, (std::vector{"prepare0", "launch0", "prepare1"})); + EXPECT_EQ(events_without_wait(), (std::vector{"prepare0", "launch0", "prepare1"})); second.activate(); EXPECT_FALSE(second.done()); - g_complete[0] = true; - EXPECT_TRUE(first.done()); + complete(0); + EXPECT_TRUE(await_done(first)); EXPECT_FALSE(second.done()); EXPECT_TRUE(second.launched()); - EXPECT_EQ(g_events, (std::vector{"prepare0", "launch0", "prepare1", "finalize0", "launch1"})); + EXPECT_EQ( + events_without_wait(), (std::vector{"prepare0", "launch0", "prepare1", "finalize0", "launch1"}) + ); - g_complete[1] = true; - EXPECT_TRUE(second.done()); + complete(1); + EXPECT_TRUE(await_done(second)); + lane.close(); + worker.finalize(); +} + +TEST(ChipRunLaneTest, ExplicitPrepareDoesNotLaunchBeforeActivation) { + ChipWorker worker; + prime_worker(worker); + ChipRunLane lane(worker); + + ChipRun run = submit(lane, 101, 0, false); + EXPECT_EQ(run.preparation_disposition(), ChipRunPreparationDisposition::VALIDATED_ONLY); + EXPECT_FALSE(run.launched()); + + run.prepare(); + EXPECT_EQ(run.preparation_disposition(), ChipRunPreparationDisposition::NATIVE_PREPARED); + EXPECT_FALSE(run.launched()); + EXPECT_EQ(events_without_wait(), (std::vector{"prepare0"})); + + run.activate(); + EXPECT_TRUE(run.launched()); + EXPECT_EQ(events_without_wait(), (std::vector{"prepare0", "launch0"})); + complete(0); + EXPECT_TRUE(await_done(run)); lane.close(); worker.finalize(); } @@ -188,13 +264,17 @@ TEST(ChipRunLaneTest, ValidationOnlySuccessorPreparesAfterPromotion) { ChipRun first = submit(lane, 101, 0); ChipRun second = lane.submit(1, args, diagnostic, PipelineSlotLease{1, 0, 102}, 102, 102, nullptr, 0, false); EXPECT_EQ(second.preparation_disposition(), ChipRunPreparationDisposition::VALIDATED_ONLY); - EXPECT_EQ(g_events, (std::vector{"prepare0", "launch0"})); + EXPECT_EQ(events_without_wait(), (std::vector{"prepare0", "launch0"})); second.activate(); - g_complete[0] = true; - EXPECT_TRUE(first.done()); + complete(0); + EXPECT_TRUE(await_done(first)); EXPECT_FALSE(second.done()); - EXPECT_EQ(g_events, (std::vector{"prepare0", "launch0", "finalize0", "prepare1", "launch1"})); + EXPECT_EQ( + events_without_wait(), (std::vector{"prepare0", "launch0", "finalize0", "prepare1", "launch1"}) + ); + complete(1); + EXPECT_TRUE(await_done(second)); lane.close(); worker.finalize(); } @@ -208,18 +288,19 @@ TEST(ChipRunLaneTest, IncompatibleSuccessorRetriesAfterPredecessorFence) { ChipRun first = submit(lane, 101, 0); ChipRun second = submit(lane, 102, 1, false); EXPECT_EQ(second.preparation_disposition(), ChipRunPreparationDisposition::VALIDATED_ONLY); - EXPECT_EQ(g_events, (std::vector{"prepare0", "launch0", "prepare1"})); + EXPECT_EQ(events_without_wait(), (std::vector{"prepare0", "launch0", "prepare1"})); second.activate(); - g_complete[0] = true; - EXPECT_TRUE(first.done()); + complete(0); + EXPECT_TRUE(await_done(first)); EXPECT_TRUE(second.launched()); EXPECT_EQ( - g_events, (std::vector{"prepare0", "launch0", "prepare1", "finalize0", "prepare1", "launch1"}) + events_without_wait(), + (std::vector{"prepare0", "launch0", "prepare1", "finalize0", "prepare1", "launch1"}) ); - g_complete[1] = true; - EXPECT_TRUE(second.done()); + complete(1); + EXPECT_TRUE(await_done(second)); lane.close(); worker.finalize(); } @@ -234,14 +315,14 @@ TEST(ChipRunLaneTest, EarlierActiveDispatchOrdersBeforeAnAlreadyStagedSuccessor) ChipRun active = submit(lane, 101, 0, true); EXPECT_TRUE(active.launched()); EXPECT_EQ(successor.preparation_disposition(), ChipRunPreparationDisposition::NATIVE_PREPARED); - EXPECT_EQ(g_events, (std::vector{"prepare0", "launch0", "prepare1"})); + EXPECT_EQ(events_without_wait(), (std::vector{"prepare0", "launch0", "prepare1"})); successor.activate(); - g_complete[0] = true; - EXPECT_TRUE(active.done()); + complete(0); + EXPECT_TRUE(await_done(active)); EXPECT_TRUE(successor.launched()); - g_complete[1] = true; - EXPECT_TRUE(successor.done()); + complete(1); + EXPECT_TRUE(await_done(successor)); lane.close(); worker.finalize(); } @@ -251,13 +332,13 @@ TEST(ChipRunLaneTest, AcceptsCurrentLeaseGenerationAndRejectsOlderOne) { prime_worker(worker); ChipRunLane lane(worker); ChipRun first = submit(lane, 101, 0); - g_complete[0] = true; - ASSERT_TRUE(first.done()); + complete(0); + ASSERT_TRUE(await_done(first)); ChipStorageTaskArgs args{}; ChipRun repeated = lane.submit(1, args, CallConfig{}, PipelineSlotLease{0, 0, 101}, 102, 102, nullptr, 0, true); - g_complete[0] = true; - ASSERT_TRUE(repeated.done()); + complete(0); + ASSERT_TRUE(await_done(repeated)); EXPECT_THROW( lane.submit(1, args, CallConfig{}, PipelineSlotLease{0, 0, 100}, 103, 103, nullptr, 0, true), std::runtime_error ); @@ -274,9 +355,61 @@ TEST(ChipRunLaneTest, ActivateStopsAtTheLaunchFenceBeforePolling) { run.activate(); EXPECT_TRUE(run.launched()); EXPECT_EQ(g_poll_count, 0); - g_complete[0] = true; - EXPECT_TRUE(run.done()); - EXPECT_EQ(g_poll_count, 1); + complete(0); + EXPECT_TRUE(await_done(run)); + EXPECT_LE(g_poll_count, 1u); + lane.close(); + worker.finalize(); +} + +TEST(ChipRunLaneTest, ResidentLifecycleLaunchesSuccessorWithoutCallerPolling) { + ChipWorker worker; + prime_worker(worker); + ChipRunLane lane(worker); + + ChipRun first = submit(lane, 101, 0); + ChipRun second = submit(lane, 102, 1, false); + second.activate(); + complete(0); + + const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(100); + while (!second.launched() && std::chrono::steady_clock::now() < deadline) { + std::this_thread::yield(); + } + EXPECT_TRUE(second.launched()); + + complete(1); + EXPECT_TRUE(second.wait_until(std::chrono::steady_clock::now() + std::chrono::seconds(1))); + EXPECT_TRUE(first.done()); + lane.close(); + worker.finalize(); +} + +TEST(ChipRunLaneTest, DiagnosticFinalizeStaysOnCallingThread) { + ChipWorker worker; + prime_worker(worker); + ChipRunLane lane(worker); + ChipStorageTaskArgs args{}; + CallConfig config{}; + config.enable_dep_gen = 1; + config.output_prefix[0] = 'x'; + + const std::thread::id caller = std::this_thread::get_id(); + ChipRun run = lane.submit(1, args, config, PipelineSlotLease{0, 0, 101}, 101, 101, nullptr, 0, true); + std::thread completer([] { + { + std::unique_lock lk(g_wait_mu); + (void)g_wait_cv.wait_for(lk, std::chrono::seconds(1), [] { + return g_wait_entered; + }); + } + complete(0); + }); + + EXPECT_TRUE(run.wait_until(ChipRunLane::Deadline::max())); + completer.join(); + EXPECT_EQ(g_prepare_thread, caller); + EXPECT_EQ(g_finalize_thread, caller); lane.close(); worker.finalize(); } @@ -288,13 +421,13 @@ TEST(ChipRunLaneTest, DirectGenerationDoesNotConsumeTheFirstPipelineLease) { ChipStorageTaskArgs args{}; ChipRun direct = lane.submit(1, args, CallConfig{}); - g_complete[0] = true; - ASSERT_TRUE(direct.done()); + complete(0); + ASSERT_TRUE(await_done(direct)); ChipRun leased = submit(lane, 1, 0); EXPECT_TRUE(leased.launched()); - g_complete[0] = true; - EXPECT_TRUE(leased.done()); + complete(0); + EXPECT_TRUE(await_done(leased)); lane.close(); worker.finalize(); } @@ -310,7 +443,7 @@ TEST(ChipRunLaneTest, DirectCapacityTwoPreparesSuccessorAndBackpressuresThird) { EXPECT_TRUE(first.launched()); EXPECT_FALSE(second.launched()); EXPECT_EQ(second.preparation_disposition(), ChipRunPreparationDisposition::NATIVE_PREPARED); - EXPECT_EQ(g_events, (std::vector{"prepare0", "launch0", "prepare1"})); + EXPECT_EQ(events_without_wait(), (std::vector{"prepare0", "launch0", "prepare1"})); { std::lock_guard lk(g_wait_mu); @@ -336,11 +469,7 @@ TEST(ChipRunLaneTest, DirectCapacityTwoPreparesSuccessorAndBackpressuresThird) { EXPECT_TRUE(entered_wait); EXPECT_EQ(g_prepare_count[0], 1u) << "third submit prepared before capacity released"; EXPECT_EQ(g_prepare_count[1], 1u); - { - std::lock_guard lk(g_wait_mu); - g_release_wait = true; - } - g_wait_cv.notify_all(); + complete(0); submitter.join(); ASSERT_EQ(submit_error, nullptr); @@ -350,11 +479,11 @@ TEST(ChipRunLaneTest, DirectCapacityTwoPreparesSuccessorAndBackpressuresThird) { EXPECT_EQ(third->preparation_disposition(), ChipRunPreparationDisposition::NATIVE_PREPARED); EXPECT_EQ(g_prepare_count[0], 2u); - g_complete[1] = true; - EXPECT_TRUE(second.done()); + complete(1); + EXPECT_TRUE(await_done(second)); EXPECT_TRUE(third->launched()); - g_complete[0] = true; - EXPECT_TRUE(third->done()); + complete(0); + EXPECT_TRUE(await_done(*third)); EXPECT_TRUE(first.done()); lane.close(); worker.finalize(); @@ -372,13 +501,13 @@ TEST(ChipRunLaneTest, DirectIncompatibleSuccessorRetriesAfterPredecessorFence) { EXPECT_EQ(second.preparation_disposition(), ChipRunPreparationDisposition::VALIDATED_ONLY); EXPECT_EQ(g_prepare_count[1], 1u); - g_complete[0] = true; + complete(0); EXPECT_FALSE(second.done()); - EXPECT_TRUE(first.done()); + EXPECT_TRUE(await_done(first)); EXPECT_TRUE(second.launched()); EXPECT_EQ(g_prepare_count[1], 2u); - g_complete[1] = true; - EXPECT_TRUE(second.done()); + complete(1); + EXPECT_TRUE(await_done(second)); lane.close(); worker.finalize(); } @@ -409,8 +538,8 @@ TEST(ChipRunLaneTest, LatchedFlowControlDeadlockIsTerminalRatherThanRetried) { EXPECT_EQ(g_prepare_count[1], 1u) << "latched deadlock was retried as an incompatible prepare"; EXPECT_FALSE(lane.poisoned()); - g_complete[0] = true; - EXPECT_TRUE(first.done()); + complete(0); + EXPECT_TRUE(await_done(first)); lane.close(); worker.finalize(); } @@ -423,14 +552,15 @@ TEST(ChipRunLaneTest, DirectRuntimeWithoutConcurrentPrepareRetainsDepthOne) { ChipStorageTaskArgs args{}; ChipRun first = lane.submit(1, args, CallConfig{}); + complete(0); ChipRun second = lane.submit(1, args, CallConfig{}); EXPECT_TRUE(first.done()); EXPECT_TRUE(second.launched()); EXPECT_EQ(g_prepare_count[0], 2u); EXPECT_EQ(g_prepare_count[1], 0u); - g_complete[0] = true; - EXPECT_TRUE(second.done()); + complete(0); + EXPECT_TRUE(await_done(second)); lane.close(); worker.finalize(); } @@ -448,8 +578,8 @@ TEST(ChipRunLaneTest, PrepareFailureIsTerminalWithoutPoisoningTheLane) { ChipRun successor = submit(lane, 102, 1); EXPECT_TRUE(successor.launched()); - g_complete[1] = true; - EXPECT_TRUE(successor.done()); + complete(1); + EXPECT_TRUE(await_done(successor)); lane.close(); worker.finalize(); } @@ -467,8 +597,8 @@ TEST(ChipRunLaneTest, ReclaimedLaunchFailureDoesNotPoisonTheLane) { ChipRun successor = submit(lane, 102, 1); EXPECT_TRUE(successor.launched()); - g_complete[1] = true; - EXPECT_TRUE(successor.done()); + complete(1); + EXPECT_TRUE(await_done(successor)); lane.close(); worker.finalize(); } @@ -481,7 +611,7 @@ TEST(ChipRunLaneTest, ExpiredWaitLeavesTheRunLive) { EXPECT_FALSE(run.wait_until(ChipRunLane::Clock::now())); EXPECT_FALSE(run.done()); - g_complete[0] = true; + complete(0); EXPECT_TRUE(run.wait_until(ChipRunLane::Deadline::max())); lane.close(); worker.finalize(); @@ -497,18 +627,24 @@ TEST(ChipRunLaneTest, UnboundedWaitBlocksInsteadOfPolling) { ChipRunLane lane(worker); ChipRun run = submit(lane, 101, 0); - // Nothing completes this run except wait_run, so a re-polling implementation - // would loop forever. The escape hatch keeps that a fast failure with a - // readable message rather than a wedged suite: past the bound the stub - // completes the run itself, and the assertions below report the poll count. + // Complete only after the resident waiter has entered the blocking native + // wait. This keeps the assertion deterministic and proves that no polling + // loop is required to make progress. ASSERT_FALSE(g_complete[0]); g_poll_count = 0; - g_poll_completes_after = 64; + std::thread completer([] { + { + std::unique_lock lk(g_wait_mu); + (void)g_wait_cv.wait_for(lk, std::chrono::seconds(1), [] { + return g_wait_entered; + }); + } + complete(0); + }); EXPECT_TRUE(run.wait_until(ChipRunLane::Deadline::max())); - g_poll_completes_after = 0; + completer.join(); - EXPECT_NE(std::find(g_events.begin(), g_events.end(), "wait0"), g_events.end()) - << "unbounded wait never reached the device's blocking wait"; + EXPECT_TRUE(has_event("wait0")) << "unbounded wait never reached the device's blocking wait"; EXPECT_LE(g_poll_count, 1u) << "unbounded wait polled " << g_poll_count << " times instead of blocking"; lane.close(); worker.finalize(); @@ -519,9 +655,9 @@ TEST(ChipRunLaneTest, FinalizeFailurePoisonsAdmissionAndCloseReportsIt) { prime_worker(worker); ChipRunLane lane(worker); ChipRun first = submit(lane, 101, 0); - g_complete[0] = true; g_finalize_rc[0] = -7; - EXPECT_TRUE(first.done()); + complete(0); + EXPECT_TRUE(await_terminal(first)); EXPECT_THROW(first.wait_until(ChipRunLane::Clock::time_point::max()), std::runtime_error); EXPECT_TRUE(lane.poisoned()); @@ -533,15 +669,21 @@ TEST(ChipRunLaneTest, FinalizeFailurePoisonsAdmissionAndCloseReportsIt) { worker.finalize(); } -TEST(ChipRunLaneTest, PollFailureReportsTheTerminalNativeError) { +TEST(ChipRunLaneTest, WaitFailureReportsTheTerminalNativeError) { ChipWorker worker; prime_worker(worker); ChipRunLane lane(worker); - ChipRun run = submit(lane, 101, 0); - g_poll_rc[0] = SIMPLER_NATIVE_RUN_POLL_ERROR; g_wait_rc[0] = 507015; + ChipRun run = submit(lane, 101, 0); + { + std::unique_lock lk(g_wait_mu); + ASSERT_TRUE(g_wait_cv.wait_for(lk, std::chrono::seconds(1), [] { + return g_wait_entered; + })); + } + complete(0); - EXPECT_TRUE(run.done()); + EXPECT_TRUE(await_terminal(run)); EXPECT_THROW(run.wait_until(ChipRunLane::Deadline::max()), std::runtime_error); try { run.wait_until(ChipRunLane::Deadline::max()); @@ -559,6 +701,7 @@ TEST(ChipRunLaneTest, CloseDrainsAndRejectsNewSubmissions) { prime_worker(worker); ChipRunLane lane(worker); ChipRun first = submit(lane, 101, 0); + complete(0); lane.close(); EXPECT_TRUE(first.done()); diff --git a/tests/ut/cpp/hierarchical/test_scheduler.cpp b/tests/ut/cpp/hierarchical/test_scheduler.cpp index 85fd48b597..8b7a95f257 100644 --- a/tests/ut/cpp/hierarchical/test_scheduler.cpp +++ b/tests/ut/cpp/hierarchical/test_scheduler.cpp @@ -438,6 +438,26 @@ class DeterministicProgressEndpoint final : public WorkerEndpoint { return false; } + bool cancel_progress(RunId run_id, const std::string &reason) override { + ProgressCall call(*this); + std::lock_guard lk(mu_); + for (const auto &[dispatch_id, outstanding] : outstanding_) { + (void)dispatch_id; + if (!outstanding.dispatch.prepare_only || outstanding.run_id != run_id) continue; + cancelled_runs_.insert(run_id); + WorkerEndpointProgress progress; + progress.kind = WorkerProgressKind::COMPLETED; + progress.dispatch = outstanding.dispatch; + progress.completion = WorkerCompletion{ + outstanding.dispatch.task_slot, outstanding.dispatch.group_index, EndpointOutcome::TASK_FAILURE, reason + }; + events_.push_back(std::move(progress)); + cv_.notify_all(); + return true; + } + return false; + } + void request_progress_stop() noexcept override { ProgressCall call(*this); std::lock_guard lk(mu_); @@ -479,6 +499,13 @@ class DeterministicProgressEndpoint final : public WorkerEndpoint { }); } + bool wait_cancelled(RunId run_id, std::chrono::milliseconds timeout = std::chrono::seconds(3)) { + std::unique_lock lk(mu_); + return cv_.wait_for(lk, timeout, [this, run_id] { + return cancelled_runs_.count(run_id) != 0; + }); + } + bool wait_stop_requested(std::chrono::milliseconds timeout = std::chrono::seconds(3)) { std::unique_lock lk(mu_); return cv_.wait_for(lk, timeout, [this] { @@ -562,6 +589,17 @@ class DeterministicProgressEndpoint final : public WorkerEndpoint { cv_.notify_all(); } + void emit_failure(const WorkerDispatch &dispatch, const std::string &message) { + std::lock_guard lk(mu_); + WorkerEndpointProgress progress; + progress.kind = WorkerProgressKind::COMPLETED; + progress.dispatch = dispatch; + progress.completion = + WorkerCompletion{dispatch.task_slot, dispatch.group_index, EndpointOutcome::TASK_FAILURE, message}; + events_.push_back(std::move(progress)); + cv_.notify_all(); + } + int max_concurrent_progress_calls() const { return max_concurrent_calls_.load(std::memory_order_acquire); } bool progress_owner_changed() const { std::lock_guard lk(owner_mu_); @@ -629,6 +667,7 @@ class DeterministicProgressEndpoint final : public WorkerEndpoint { std::unordered_map outstanding_; std::deque events_; std::set activated_runs_; + std::set cancelled_runs_; bool stop_requested_{false}; bool stop_terminalized_{false}; size_t stop_request_count_{0}; @@ -1162,6 +1201,8 @@ TEST(WorkerManagerTest, WorkerThreadUsesOneProgressOwnerForActiveAndStagedLanes) ); worker.dispatch(WorkerDispatch{active_slot, 0}); + EXPECT_FALSE(worker.can_stage(/*pipeline_slot_id=*/0)); + EXPECT_TRUE(worker.can_stage(/*pipeline_slot_id=*/1)); worker.dispatch_prepared(WorkerDispatch{staged_slot, 0}); EXPECT_TRUE(endpoint_ptr->wait_submitted(2)); std::vector submitted = endpoint_ptr->submitted(); @@ -1536,6 +1577,58 @@ TEST(WorkerManagerTest, TwoFrameLeaseSlotsDoNotDefineFifoOrAcceptance) { allocator.shutdown(); } +TEST(WorkerManagerTest, PreparedLocalMailboxDispatchCanBeCancelledBeforeActivation) { + alignas(8) std::array mailbox{}; + Ring allocator; + allocator.init(/*heap_bytes=*/0); + constexpr RunId run_id = 23; + TaskSlot task_slot = make_progress_slot(allocator, run_id, /*pipeline_slot=*/0, /*generation=*/1); + ASSERT_NE(task_slot, INVALID_SLOT); + + LocalMailboxEndpoint endpoint(/*worker_id=*/0, mailbox.data(), /*child_pid=*/-1, /*task_frame_count=*/2); + WorkerDispatch dispatch{task_slot, 0, /*dispatch_id=*/43, /*prepare_only=*/true}; + endpoint.submit_progress(&allocator, dispatch); + + char *frame = test_task_frame(mailbox, 0); + EXPECT_EQ(test_frame_state(frame), MailboxState::PREPARE_READY); + const int32_t validated_only = static_cast(MailboxPreparationDisposition::VALIDATED_ONLY); + std::memcpy(frame + MAILBOX_OFF_PREPARATION_DISPOSITION, &validated_only, sizeof(validated_only)); + set_test_frame_state(frame, MailboxState::FRAME_STAGED); + + WorkerEndpointProgress progress; + ASSERT_TRUE(endpoint.poll_progress(progress)); + ASSERT_EQ(progress.kind, WorkerProgressKind::FRAME_STAGED); + ASSERT_TRUE(endpoint.cancel_progress(run_id, "peer preparation failed")); + EXPECT_EQ(test_frame_state(frame), MailboxState::ABANDON); + EXPECT_STREQ(frame + MAILBOX_OFF_ERROR_MSG, "peer preparation failed"); + + // The Python child acknowledges ABANDON by abandoning the native staged + // run and publishing TASK_FAILED; the parent then completes the dispatch + // through the ordinary terminal-progress path. + set_test_frame_state(frame, MailboxState::TASK_FAILED); + ASSERT_TRUE(endpoint.poll_progress(progress)); + EXPECT_EQ(progress.kind, WorkerProgressKind::COMPLETED); + EXPECT_EQ(progress.completion.outcome, EndpointOutcome::TASK_FAILURE); + EXPECT_NE(progress.completion.error_message.find("peer preparation failed"), std::string::npos); + allocator.shutdown(); +} + +TEST(WorkerManagerTest, NativePreparedBarrierUsesDedicatedMailboxState) { + alignas(8) std::array mailbox{}; + Ring allocator; + allocator.init(/*heap_bytes=*/0); + TaskSlot task_slot = make_progress_slot(allocator, /*run_id=*/24, /*pipeline_slot=*/0, /*generation=*/1); + ASSERT_NE(task_slot, INVALID_SLOT); + + LocalMailboxEndpoint endpoint(/*worker_id=*/0, mailbox.data(), /*child_pid=*/-1, /*task_frame_count=*/2); + WorkerDispatch dispatch{task_slot, 0, /*dispatch_id=*/44, /*prepare_only=*/true}; + dispatch.require_native_prepare = true; + endpoint.submit_progress(&allocator, dispatch); + + EXPECT_EQ(test_frame_state(test_task_frame(mailbox, 0)), MailboxState::NATIVE_PREPARE_READY); + allocator.shutdown(); +} + TEST(WorkerManagerTest, CapacityOneMailboxUsesTheProgressTaskFrame) { alignas(8) std::array mailbox{}; Ring allocator; @@ -2108,6 +2201,64 @@ TEST_F(ProgressSchedulerFixture, GroupSubmitReportsNoSingleWorkerAndNoSingleInde ); } +TEST_F(ProgressSchedulerFixture, GroupWaitsForEveryPreparedMemberBeforeActivation) { + RunId run = orchestrator.begin_run(); + SubmitResult group = orchestrator.submit_next_level_group( + C(10), {single_tensor_args(0x9100, TensorArgType::OUTPUT), single_tensor_args(0xA100, TensorArgType::OUTPUT)}, + config, {0, 1} + ); + orchestrator.close_run_submission(run); + + ASSERT_TRUE(endpoint0->wait_submitted(1)); + ASSERT_TRUE(endpoint1->wait_submitted(1)); + WorkerDispatch first = endpoint0->submitted().front(); + WorkerDispatch second = endpoint1->submitted().front(); + EXPECT_TRUE(first.prepare_only); + EXPECT_TRUE(second.prepare_only); + EXPECT_TRUE(first.require_native_prepare); + EXPECT_TRUE(second.require_native_prepare); + EXPECT_EQ(first.task_slot, group.task_slot); + EXPECT_EQ(second.task_slot, group.task_slot); + + endpoint0->emit(WorkerProgressKind::FRAME_STAGED, first); + EXPECT_FALSE(endpoint0->wait_activated(run, std::chrono::milliseconds(20))); + EXPECT_FALSE(endpoint1->wait_activated(run, std::chrono::milliseconds(20))); + + endpoint1->emit(WorkerProgressKind::FRAME_STAGED, second); + EXPECT_TRUE(endpoint0->wait_activated(run)); + EXPECT_TRUE(endpoint1->wait_activated(run)); + + endpoint0->emit(WorkerProgressKind::ACCEPTED, first); + endpoint1->emit(WorkerProgressKind::ACCEPTED, second); + endpoint0->emit(WorkerProgressKind::COMPLETED, first); + endpoint1->emit(WorkerProgressKind::COMPLETED, second); + EXPECT_TRUE(orchestrator.wait_run_for(run, 3.0)); + if (orchestrator.run_done(run)) orchestrator.release_run(run); +} + +TEST_F(ProgressSchedulerFixture, GroupPrepareFailureCancelsStagedPeersWithoutActivation) { + RunId run = orchestrator.begin_run(); + (void)orchestrator.submit_next_level_group( + C(11), {single_tensor_args(0x9200, TensorArgType::OUTPUT), single_tensor_args(0xA200, TensorArgType::OUTPUT)}, + config, {0, 1} + ); + orchestrator.close_run_submission(run); + + ASSERT_TRUE(endpoint0->wait_submitted(1)); + ASSERT_TRUE(endpoint1->wait_submitted(1)); + WorkerDispatch first = endpoint0->submitted().front(); + WorkerDispatch second = endpoint1->submitted().front(); + endpoint1->emit(WorkerProgressKind::FRAME_STAGED, second); + endpoint0->emit_failure(first, "injected group prepare failure"); + + EXPECT_TRUE(endpoint1->wait_cancelled(run)); + EXPECT_FALSE(endpoint0->wait_activated(run, std::chrono::milliseconds(20))); + EXPECT_FALSE(endpoint1->wait_activated(run, std::chrono::milliseconds(20))); + EXPECT_THROW((void)orchestrator.wait_run_for(run, 3.0), std::runtime_error); + EXPECT_TRUE(orchestrator.run_failed(run)); + if (orchestrator.run_done(run)) orchestrator.release_run(run); +} + TEST_F(ProgressSchedulerFixture, SuccessorStagesButActivatesOnlyAfterFifoPromotion) { RunId first_run = orchestrator.begin_run(); SubmitResult first = @@ -2277,48 +2428,6 @@ TEST_F(CapacityOneProgressSchedulerFixture, PublicationFailureCompletesTheClaime if (orchestrator.run_done(second_run)) orchestrator.release_run(second_run); } -TEST_F(ProgressSchedulerFixture, PreparedSuccessorGroupRemainsQueuedUntilPromotion) { - RunId first_run = orchestrator.begin_run(); - orchestrator.submit_next_level(C(3), single_tensor_args(0x3000, TensorArgType::OUTPUT), config, 0); - orchestrator.close_run_submission(first_run); - RunId second_run = orchestrator.begin_run(); - orchestrator.submit_next_level_group( - C(4), {single_tensor_args(0x4000, TensorArgType::OUTPUT), single_tensor_args(0x5000, TensorArgType::OUTPUT)}, - config, {0, 1} - ); - orchestrator.close_run_submission(second_run); - - EXPECT_TRUE(endpoint0->wait_submitted(1)); - EXPECT_TRUE(ready_next.singles_empty(second_run)); - EXPECT_FALSE(ready_next.groups_empty(second_run)); - EXPECT_FALSE(manager.has_staged_run(second_run)); - EXPECT_TRUE(endpoint1->submitted().empty()); - std::vector first_submissions = endpoint0->submitted(); - ASSERT_EQ(first_submissions.size(), 1u); - - endpoint0->emit(WorkerProgressKind::ACCEPTED, first_submissions[0]); - endpoint0->emit(WorkerProgressKind::COMPLETED, first_submissions[0]); - EXPECT_TRUE(endpoint0->wait_submitted(2)); - EXPECT_TRUE(endpoint1->wait_submitted(1)); - std::vector worker0_submissions = endpoint0->submitted(); - std::vector worker1_submissions = endpoint1->submitted(); - ASSERT_EQ(worker0_submissions.size(), 2u); - ASSERT_EQ(worker1_submissions.size(), 1u); - EXPECT_FALSE(worker0_submissions[1].prepare_only); - EXPECT_FALSE(worker1_submissions[0].prepare_only); - EXPECT_EQ(orchestrator.active_run_id(), second_run); - EXPECT_EQ(orchestrator.preparable_run_id(), INVALID_RUN_ID); - - endpoint0->emit(WorkerProgressKind::ACCEPTED, worker0_submissions[1]); - endpoint1->emit(WorkerProgressKind::ACCEPTED, worker1_submissions[0]); - endpoint0->emit(WorkerProgressKind::COMPLETED, worker0_submissions[1]); - endpoint1->emit(WorkerProgressKind::COMPLETED, worker1_submissions[0]); - EXPECT_TRUE(orchestrator.wait_run_for(first_run, 3.0)); - EXPECT_TRUE(orchestrator.wait_run_for(second_run, 3.0)); - if (orchestrator.run_done(first_run)) orchestrator.release_run(first_run); - if (orchestrator.run_done(second_run)) orchestrator.release_run(second_run); -} - TEST_F(ProgressSchedulerFixture, PreparedSuccessorSingleCannotBypassItsReadyGroup) { RunId first_run = orchestrator.begin_run(); SubmitResult first = @@ -2350,8 +2459,13 @@ TEST_F(ProgressSchedulerFixture, PreparedSuccessorSingleCannotBypassItsReadyGrou ASSERT_EQ(worker1_submissions.size(), 1u); EXPECT_EQ(worker0_submissions[1].task_slot, group.task_slot); EXPECT_EQ(worker1_submissions[0].task_slot, group.task_slot); - EXPECT_FALSE(worker0_submissions[1].prepare_only); - EXPECT_FALSE(worker1_submissions[0].prepare_only); + EXPECT_TRUE(worker0_submissions[1].prepare_only); + EXPECT_TRUE(worker1_submissions[0].prepare_only); + + endpoint0->emit(WorkerProgressKind::FRAME_STAGED, worker0_submissions[1]); + endpoint1->emit(WorkerProgressKind::FRAME_STAGED, worker1_submissions[0]); + EXPECT_TRUE(endpoint0->wait_activated(second_run)); + EXPECT_TRUE(endpoint1->wait_activated(second_run)); endpoint0->emit(WorkerProgressKind::ACCEPTED, worker0_submissions[1]); endpoint1->emit(WorkerProgressKind::ACCEPTED, worker1_submissions[0]); diff --git a/tests/ut/py/test_worker/test_host_worker.py b/tests/ut/py/test_worker/test_host_worker.py index 11bc91c9a1..3216241480 100644 --- a/tests/ut/py/test_worker/test_host_worker.py +++ b/tests/ut/py/test_worker/test_host_worker.py @@ -647,6 +647,7 @@ def __init__(self, lane, submission) -> None: self.activated = False self._launched = False self.terminal = False + self.wait_calls = 0 self.error: Optional[BaseException] = None self._disposition = worker_mod._VALIDATED_ONLY @@ -667,6 +668,17 @@ def activate(self) -> None: self._lane._launch_front() self._lane._prepare_successor() + def prepare(self) -> None: + if self.activated or self._launched: + raise RuntimeError("cannot prepare an activated fake ChipRun") + if self.terminal: + self._raise_if_failed() + return + if not self._lane._runs or self._lane._runs[0] is not self: + raise RuntimeError("only the front fake ChipRun can be prepared") + if self.token is None: + self._lane._prepare(self) + def abandon(self) -> None: if self._launched: raise RuntimeError("cannot abandon a launched fake ChipRun") @@ -675,6 +687,11 @@ def abandon(self) -> None: def done(self) -> bool: return self._lane._progress(self) + def wait(self, timeout: float = -1.0) -> bool: + del timeout + self.wait_calls += 1 + return self.done() + def _raise_if_failed(self) -> None: assert self.terminal if self.error is not None: @@ -1112,6 +1129,7 @@ def test_two_frame_stages_b_without_native_prepare_until_a_finalizes(): ("launch", 1), ("finalize", 1), ] + assert all(run.wait_calls > 0 for run in harness.cw._impl._runs) launch_entries = [event for event in harness.cw._impl.events if event[0] == "launch_enter"] assert launch_entries == [ ("launch_enter", 0, 0, worker_mod._FRAME_STAGED), @@ -1357,6 +1375,40 @@ def test_two_frame_prepare_ready_waits_for_sticky_activation(): harness.close() +def test_two_frame_prepare_ready_can_be_abandoned_before_activation(): + harness = _TwoFrameLoopHarness() + try: + harness.publish(0, 1, state=worker_mod._PREPARE_READY) + harness.start() + harness.wait_state(0, worker_mod._FRAME_STAGED) + assert not harness.cw._impl.launched[0].is_set() + + _mailbox_store_i32(harness.state_addr(0), worker_mod._ABANDON) + harness.wait_state(0, worker_mod._TASK_FAILED) + assert not harness.cw._impl.launched[0].is_set() + assert not harness.cw._impl._runs + finally: + harness.close() + + +def test_two_frame_native_prepare_ready_prepares_before_group_activation(): + harness = _TwoFrameLoopHarness() + try: + harness.publish(0, 1, state=worker_mod._NATIVE_PREPARE_READY) + harness.start() + harness.wait_state(0, worker_mod._FRAME_STAGED) + assert harness.cw._impl.prepared[0].is_set() + assert harness.preparation_disposition(0) == worker_mod._NATIVE_PREPARED + assert not harness.cw._impl.launched[0].is_set() + + _mailbox_store_i32(harness.state_addr(0), worker_mod._ACTIVATE) + assert harness.cw._impl.launched[0].wait(timeout=1.0) + harness.cw._impl.completed[0].set() + harness.wait_state(0, worker_mod._TASK_DONE) + finally: + harness.close() + + def test_two_frame_hbg_lone_prepare_ready_stages_before_native_prepare(): harness = _TwoFrameLoopHarness( supports_concurrent_native_prepare=True, diff --git a/tests/ut/py/test_worker/test_mailbox_atomics.py b/tests/ut/py/test_worker/test_mailbox_atomics.py index bf2e18c3af..1c6b5d2b6f 100644 --- a/tests/ut/py/test_worker/test_mailbox_atomics.py +++ b/tests/ut/py/test_worker/test_mailbox_atomics.py @@ -33,7 +33,12 @@ from multiprocessing.shared_memory import SharedMemory import pytest -from _task_interface import _mailbox_load_i32, _mailbox_store_i32 # pyright: ignore[reportMissingImports] +from _task_interface import ( # pyright: ignore[reportMissingImports] + _mailbox_load_i32, + _mailbox_notify_i32, + _mailbox_store_i32, + _mailbox_wait_i32, +) # --------------------------------------------------------------------------- # Helpers @@ -88,6 +93,29 @@ def test_offset(self, shm): class TestCrossProcess: + def test_notification_wakes_cross_process_waiter(self, shm): + state_addr = _addr(shm.buf, 0) + result_addr = _addr(shm.buf, 8) + _mailbox_store_i32(state_addr, 0) + _mailbox_store_i32(result_addr, 0) + + pid = os.fork() + if pid == 0: + try: + deadline = time.monotonic() + 5.0 + while _mailbox_load_i32(state_addr) == 0 and time.monotonic() < deadline: + _mailbox_wait_i32(state_addr, 0, 5.0) + _mailbox_store_i32(result_addr, int(_mailbox_load_i32(state_addr) == 1)) + finally: + os._exit(0) + + time.sleep(0.02) + started = time.monotonic() + _mailbox_notify_i32(state_addr) + os.waitpid(pid, 0) + assert _mailbox_load_i32(result_addr) == 1 + assert time.monotonic() - started < 1.0 + def test_child_transitions_visible_in_parent(self, shm): """Child cycles state 0→1→2→3→0; parent must at least see the final 0. @@ -245,6 +273,8 @@ def test_every_declared_state_matches_the_native_enum(self): "TASK_FAILED": worker_mod._TASK_FAILED, "ACTIVATE": worker_mod._ACTIVATE, "PREPARE_READY": worker_mod._PREPARE_READY, + "ABANDON": worker_mod._ABANDON, + "NATIVE_PREPARE_READY": worker_mod._NATIVE_PREPARE_READY, } assert declared == dict(MAILBOX_STATE_VALUES) @@ -278,7 +308,7 @@ def test_the_guard_rejects_an_undeclared_enumerator(self): from simpler import worker as worker_mod # noqa: PLC0415 original = worker_mod.MAILBOX_STATE_VALUES - worker_mod.MAILBOX_STATE_VALUES = {**dict(original), "FUTURE_STATE": 13} + worker_mod.MAILBOX_STATE_VALUES = {**dict(original), "FUTURE_STATE": 15} try: with pytest.raises(RuntimeError, match="does not declare"): worker_mod._assert_mailbox_wire_constants()