diff --git a/docs/callable-identity-registration.md b/docs/callable-identity-registration.md index b2b6065530..def82b26df 100644 --- a/docs/callable-identity-registration.md +++ b/docs/callable-identity-registration.md @@ -526,10 +526,9 @@ lets both sides reject a stale activation, acceptance, or completion from an older use of the same pipeline frame. `FRAME_STAGED` confirms that this exact payload has been validated and retained; it does not resolve to a different callable identity and does not by itself prove that a native run is prepared. -An eligible HBG successor normally prepares before this publication when its -predecessor is already active. A successor staged before any active claim -publishes validation-only and may later gain a native token without another -mailbox state transition. +An eligible HBG successor prepares in its lease-selected inactive bank while +the predecessor is active. A frame published before an active claim remains +validation-only until activation or a predecessor claim makes it eligible. The target child loop owns the final execution resolve: @@ -702,11 +701,9 @@ Local mailbox task frames are hashid-based: - Chip and sub child loops resolve `hashid -> local_slot` immediately before execution. - A two-frame chip loop retains the resolved slot and immutable payload at - `FRAME_STAGED`; an eligible non-diagnostic HBG successor prepares a native - token in the lease-selected inactive bank once its predecessor owns the - active claim. A successor staged before that claim may gain the token without - another state transition. HBG tasks adjacent to diagnostics and all TMR - tasks wait for activation and native-run availability. + `FRAME_STAGED`; eligible non-diagnostic HBG successors may already own a + generation-bound native token, while other configurations defer native + prepare until activation and native-run availability. - Protocol/run/lease-generation/dispatch identity guards task-frame reuse, and a separate sticky word records only real native launch acceptance. - `ChipWorker.run(local_slot)` remains private to the child process. diff --git a/docs/dynamic-linking.md b/docs/dynamic-linking.md index 517021bb2e..cfc2e96178 100644 --- a/docs/dynamic-linking.md +++ b/docs/dynamic-linking.md @@ -291,16 +291,11 @@ ChipWorker.init(device_id, bins) # Python wrapper ctypes.CDLL(libcpu_sim_context.so, RTLD_GLOBAL) # sim only, once per process _ChipWorker.init(host_path, aicpu_path, aicore_path, device_id) # C++ dlopen(host_runtime.so, RTLD_LOCAL) - dlsym every required export declared in pto_runtime_c_api.h, including: - create_device_context, destroy_device_context, simpler_init, + dlsym: create_device_context, destroy_device_context, simpler_init, get_runtime_size, get_runtime_alignment, simpler_register_callable, simpler_prepare_run, simpler_launch_run, simpler_poll_run, simpler_wait_run, simpler_finalize_run, simpler_run, - simpler_unregister_callable, get_pipeline_contract, - supports_concurrent_native_prepare_ctx, - set_native_run_identity_ctx, set_task_accepted_state_ctx, - get_arena_bank_gm_heap_base_ctx, get_retained_temp_addr_ctx, - finalize_device + simpler_unregister_callable, finalize_device create_device_context() → DeviceContextHandle allocate zeroed, aligned, stable native-run storage per pipeline slot simpler_init(ctx, device_id, aicpu*, aicpu_size, aicore*, aicore_size) diff --git a/docs/task-flow.md b/docs/task-flow.md index 8ac8c458ab..3e31033bec 100644 --- a/docs/task-flow.md +++ b/docs/task-flow.md @@ -69,9 +69,9 @@ follows `CallConfig`, and the 32-byte digest prefixes the args blob. SUB, nested/remote L3, simulation, A5, and depth-one fallbacks carry that payload in the base compatibility frame. A direct A2/A3 onboard chip child can instead use either of two task frames after the base control frame, with -`PREPARE_READY -> FRAME_STAGED -> ACTIVATE` separating endpoint preparation -from native launch. The receiving child resolves the digest in its own address -space in both forms. +`PREPARE_READY -> FRAME_STAGED -> ACTIVATE` separating validation from native +launch. The receiving child resolves the digest in its own address space in +both forms. The proposed remote L3 path keeps the same callable identity contract, but sends it in a versioned TASK frame. The remote endpoint resolves the digest @@ -348,11 +348,16 @@ finalize`, and the existing `simpler_run` / `ChipWorker.run` surface is the blocking composition of those phases. `prepare` constructs and binds the per-run `Runtime` without crossing the device launch fence; `launch` returns after the backend has actually submitted its execution; `finalize` owns -validation, copy-back, DFX, and Runtime destruction. A backend that advertises -concurrent native preparation may own one active token and one prepared but -unlaunched and unaccepted successor token in separate lease-selected banks. -Other backends permit only one unfinished native run, including a -prepared-but-not-launched run. +validation, copy-back, DFX, and Runtime destruction. One runner still permits +one device-launch owner. A backend that explicitly advertises concurrent native +preparation may additionally own one distinct-slot, prepared-but-unlaunched +successor. Backends without that capability and diagnostic configurations keep +the depth-one native lifecycle. + +Direct L2 `Worker.submit()` composes the same phases asynchronously: it returns +after launch instead of device completion, and its `RunHandle` drives the +ordered wait/finalize handoff. At depth two the second handle may own a prepared +successor; a third submission backpressures before preparing or reusing a slot. #### Two-frame endpoint staging lane @@ -371,28 +376,27 @@ IDLE -> PREPARE_READY -> FRAME_STAGED -> ACTIVATE -> TASK_LAUNCHED -> TASK_DONE | TASK_FAILED ``` -`FRAME_STAGED` means that the child owns an immutable frame snapshot; it does -not by itself distinguish validation-only staging from completed native -preparation. For a `host_build_graph` successor whose own configuration and -active predecessor are both non-diagnostic, the child constructs a -generation-bound native run in the leased inactive arena bank while the -predecessor executes. If the predecessor is already active, that preparation -finishes before `FRAME_STAGED` publication. A successor that reaches the child -before any predecessor owns the active claim publishes validation-only, so the -parent can activate it without deadlock; native prepare follows activation or -a later predecessor claim without another mailbox state transition. HBG tasks -adjacent to diagnostic state and all -`tensormap_and_ringbuffer` tasks also use this as a validation-only state: -native prepare waits for the predecessor's complete device fence because their -shared diagnostic or device-scratch state is not safe to rewrite early. -Backend and per-run capabilities, rather than the mailbox protocol, select -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. +`FRAME_STAGED` means the child validated and owns an immutable frame snapshot; +the mailbox state does not distinguish validation-only staging from completed +native preparation. A capable backend may prepare a non-diagnostic successor +in its leased slot while the predecessor is active. Otherwise native prepare is +deferred until the predecessor has polled complete and finalized. In both cases +the successor remains unlaunched and unaccepted until FIFO activation. + +`host_build_graph` advertises this capability because each lease selects an +independent `HOST_PER_RUN` arena bank. Once the predecessor owns the active +claim, the child binds the successor into the inactive bank and provisions its +fresh AICore stream before publishing `FRAME_STAGED`. The active bank remains +immutable, and launch still waits for predecessor completion and finalization. + +`tensormap_and_ringbuffer` stages slot-private task arguments and Host +descriptors while sharing arena bank 0. Concurrent native preparation requires +the successor's resolved runtime sizing key to match the active prebuilt arena. +An incompatible successor remains staged in FIFO order and retries preparation +after the active device fence; it is not rejected. Both paths keep a fresh +AICore stream per run and persistent slot-owned AICPU streams. +Diagnostics retain validation-only staging because their collectors are +runner-global. 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 @@ -401,12 +405,6 @@ 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. -The child validates newly visible metadata from both frames before selecting -the next active `dispatch_id`. It prepares that active token first; preparation -of an ordinary HBG successor starts only after the selected predecessor owns -the native active claim. Physical frame order therefore cannot reorder native -prepare or launch. - Every task frame carries protocol, run, lease slot, generation, dispatch, and callable identity. Its sticky acceptance word is separate from the state word and is set only by the real native launch marker. `FRAME_STAGED` never satisfies @@ -414,12 +412,8 @@ the launch fence. A terminal pre-launch failure may conservatively retire the run-level acceptance waiter, but it does not set the frame's acceptance word. The parent clears that word only immediately before reusing an `IDLE` frame. Control commands continue to use a separate base frame, so they cannot -overwrite either staged task frame. Callable prepare/register/unregister -commands defer while an active or backend-prepared token owns runtime state. -The default unbounded control wait therefore follows child liveness through a -long run; a finite control timeout includes this deferral interval, and expiry -poisons the local endpoint because the pending command's completion is -uncertain. +overwrite either staged task frame. Registry mutation is deferred while an +active or backend-prepared token still owns runtime state. #### TRB temporary buffer @@ -727,8 +721,8 @@ Step-by-step (one chip worker): | 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 | -| 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` | +| 6 | chip_0 child process | validate the frame and resolve its digest, then publish `FRAME_STAGED`; a successor waits for `ACTIVATE`, while the active frame may proceed immediately | +| 7 | chip_0 native-run path | after activation, prepare and launch the native run; poll it to completion and finalize it before another staged frame may launch. 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/docs/worker-manager.md b/docs/worker-manager.md index eb6126750c..91d9065bbd 100644 --- a/docs/worker-manager.md +++ b/docs/worker-manager.md @@ -238,40 +238,31 @@ successor: IDLE -> PREPARE_READY -> FRAME_STAGED -> ACTIVATE `FRAME_STAGED` means the child validated the frame identity and arguments, resolved the callable digest, rewrote any mapped host addresses, and retained -an immutable snapshot. The state does not by itself distinguish validation-only -staging from completed native preparation. When both an HBG successor and its -already-active predecessor are non-diagnostic, the child prepares a -generation-bound native token in the successor's leased inactive bank before -publishing `FRAME_STAGED`. A successor that arrives before any active claim -publishes validation-only, allowing the parent to activate it; native prepare -follows that activation or a later predecessor claim without another mailbox -state transition. HBG tasks adjacent to diagnostic state and all TMR tasks keep -the same two-frame protocol but defer native prepare because their shared -diagnostic or device-scratch state cannot be rewritten while another run is -active. - -An HBG token remains unlaunched and unaccepted until activation, and no backend -launches a successor until the predecessor is polled and finalized. Shutdown, -stale activation, and pre-launch failure finalize any unlaunched token exactly -once. +an immutable snapshot. It does not by itself say whether the runtime-specific +native run is prepared. A backend with the explicit concurrent-prepare +capability may prepare one non-diagnostic successor in a distinct leased slot +while the predecessor is active. Unsupported backends and diagnostic runs keep +validation-only staging and defer native prepare until the predecessor is +polled and finalized. Neither path launches or accepts the successor before +FIFO activation. + +HBG opts into concurrent preparation using its two lease-selected +`HOST_PER_RUN` banks. Preparation binds the successor and creates its fresh +AICore stream in the inactive bank while leaving the active bank immutable. +A frame that arrives before any active claim publishes validation-only and may +gain its native token later without another mailbox state transition. Activation is sticky on the parent side: FIFO promotion may be observed before the child reaches `FRAME_STAGED`. The endpoint records that permission and publishes `ACTIVATE` only by a compare/exchange from the matching -`FRAME_STAGED` state. The child validates all newly visible frame metadata -before choosing the next active frame by `dispatch_id`; capable native prepare -starts for that frame first, and a successor prepares only after the selected -predecessor owns the active claim. Frame index therefore cannot let a later -dispatch bypass an earlier eligible dispatch. +`FRAME_STAGED` state. The child chooses activated frames by `dispatch_id`, so a +later frame cannot bypass an earlier eligible dispatch. Task-frame publication briefly shares the base control mutex so it has a defined order relative to a control request. Once published, ordinary controls may run while the device task is active. Registry-mutating controls are deferred until active and backend-prepared native state is finalized; final unregister also waits for every published frame using that digest to retire. -The default unbounded control wait follows child liveness through this -deferral. A finite timeout includes the deferral interval and poisons the -endpoint if it expires with completion uncertain. Each task frame is bound to its pipeline lease slot and carries a protocol trailer with `{run_id, slot_id, generation, dispatch_id}`. Parent and child diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index 01fbc7f49b..52ef84c212 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -1790,23 +1790,25 @@ NB_MODULE(_task_interface, m) { .def( "_prepare_native_run_with_pipeline_lease", [](ChipWorker &self, int32_t callable_id, TaskArgs &args, const CallConfig &config, uint32_t slot_id, - uint64_t generation) { + uint64_t generation, uint64_t run_id, uint64_t dispatch_id) { return self.prepare_native_run( - callable_id, make_view(args), config, PipelineSlotLease{slot_id, 0, generation} + callable_id, make_view(args), config, PipelineSlotLease{slot_id, 0, generation}, run_id, dispatch_id ); }, nb::arg("callable_id"), nb::arg("args"), nb::arg("config"), nb::arg("slot_id"), nb::arg("generation"), - nb::call_guard(), + nb::arg("run_id") = 0, nb::arg("dispatch_id") = 0, nb::call_guard(), "Prepare a generation-bound native run without crossing its device launch fence." ) .def( "_prepare_native_run_with_pipeline_lease", [](ChipWorker &self, int32_t callable_id, ChipStorageTaskArgs &args, const CallConfig &config, - uint32_t slot_id, uint64_t generation) { - return self.prepare_native_run(callable_id, &args, config, PipelineSlotLease{slot_id, 0, generation}); + uint32_t slot_id, uint64_t generation, uint64_t run_id, uint64_t dispatch_id) { + return self.prepare_native_run( + callable_id, &args, config, PipelineSlotLease{slot_id, 0, generation}, run_id, dispatch_id + ); }, nb::arg("callable_id"), nb::arg("args"), nb::arg("config"), nb::arg("slot_id"), nb::arg("generation"), - nb::call_guard(), + nb::arg("run_id") = 0, nb::arg("dispatch_id") = 0, nb::call_guard(), "Prepare a generation-bound native run from pre-encoded task args." ) .def( diff --git a/python/simpler/task_interface.py b/python/simpler/task_interface.py index 5133e63ac9..46442b4878 100644 --- a/python/simpler/task_interface.py +++ b/python/simpler/task_interface.py @@ -1365,7 +1365,9 @@ def _run_slot_with_pipeline_lease(self, callable_id, args, slot_id, generation, setattr(config, k, v) self._impl._run_with_pipeline_lease(int(callable_id), args, config, int(slot_id), int(generation)) - def _prepare_native_run_with_pipeline_lease(self, callable_id, args, slot_id, generation, config=None, **kwargs): + def _prepare_native_run_with_pipeline_lease( + self, callable_id, args, slot_id, generation, config=None, *, run_id=0, dispatch_id=0, **kwargs + ): """Prepare one native run without crossing its device launch fence. Private B3a seam for the hierarchical endpoint. The returned token is @@ -1379,7 +1381,7 @@ def _prepare_native_run_with_pipeline_lease(self, callable_id, args, slot_id, ge for k, v in kwargs.items(): setattr(config, k, v) return self._impl._prepare_native_run_with_pipeline_lease( - int(callable_id), args, config, int(slot_id), int(generation) + int(callable_id), args, config, int(slot_id), int(generation), int(run_id), int(dispatch_id) ) def _launch_native_run(self, run): @@ -1420,6 +1422,10 @@ def pipeline_depth(self): def runtime_slot_count(self): return self._impl.runtime_slot_count + @property + def supports_concurrent_native_prepare(self): + return bool(self._impl.supports_concurrent_native_prepare) + @property def runtime_buffer_addrs(self): """Address of each opaque host native-run storage buffer, in slot order.""" diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 9487d92882..6ac44e115c 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -2396,6 +2396,7 @@ class _StagedFrame: activated: bool native_run: Any = None published: bool = False + defer_native_prepare: bool = False supports_concurrent_native_prepare = bool(cw._impl.supports_concurrent_native_prepare) staged_frames: dict[int, _StagedFrame] = {} @@ -2457,6 +2458,7 @@ def prepare_frame_native_run(frame: _StagedFrame) -> Any: run_id, dispatch_id, ) + frame.defer_native_prepare = False return frame.native_run def finalize_frame_native_run(frame: _StagedFrame) -> None: @@ -2597,6 +2599,7 @@ def stage_frame(index: int, initial_state: int) -> _StagedFrame | None: staged.native_run is None and supports_concurrent_native_prepare and not config_has_diagnostics(staged.config) + and (not staged.defer_native_prepare or active_frame is None) and ( (active_frame is None and staged is next_active) or ( @@ -2614,6 +2617,15 @@ def stage_frame(index: int, initial_state: int) -> _StagedFrame | None: if not staged.published: publish_frame_staged(staged) except Exception as e: # noqa: BLE001 + if ( + active_frame is not None + and staged is not active_frame + and _native_prepare_requires_depth_one(e) + ): + staged.defer_native_prepare = True + if not staged.published: + publish_frame_staged(staged) + continue prepare_message = _format_exc(f"chip_process dev={device_id}: native prepare", e) finalize_failed = False try: @@ -3142,6 +3154,28 @@ class _RunResources: requires_ordered_cleanup: bool = False +@dataclass +class _L2NativeRun: + """One direct L2 submission owned by the Worker's bounded native lane.""" + + run_id: int + slot_id: int + generation: int + callable_id: int + args: Any + config: CallConfig + native_run: Any | None + permits_successor: bool + phase: str = "deferred" + error: BaseException | None = None + handle: RunHandle | None = None + + +def _native_prepare_requires_depth_one(error: BaseException) -> bool: + """Recognize the backend's explicit, correctness-preserving fallback.""" + return "native prepare requires depth-one fallback" in str(error) + + @dataclass class _PendingRemoteImportReleaseState: """Durable local phases after a deferred import-release RPC starts.""" @@ -3789,6 +3823,12 @@ def __init__( # Level-2 internals self._chip_worker: ChipWorker | None = None + self._l2_progress_mu = threading.RLock() + self._l2_progress_cv = threading.Condition(self._l2_progress_mu) + self._l2_runs: dict[int, _L2NativeRun] = {} + self._l2_fifo: list[int] = [] + self._l2_slot_generations = [0] * PTO_PIPELINE_MAX_DEPTH + self._l2_next_run_id = 1 # Level-3+ internals self._worker: _Worker | None = None @@ -8107,8 +8147,10 @@ def submit(self, callable, args=None, config=None) -> RunHandle: Dispatch: - L2: ``callable`` is a ``CallableHandle`` returned by ``Worker.register(chip_callable)``. Routes to the private slot - carried by the handle. The current L2 backend remains blocking, so - the returned handle is already complete. + carried by the handle. Submission returns after native launch, and + the handle owns completion and finalization. A capable depth-two + backend may retain one prepared successor; other configurations + deterministically wait for the active run before preparing it. - L3+: ``callable`` is a Python orch fn invoked with the ``Orchestrator`` handle. Graph construction completes synchronously; device completion is reported by the returned handle. @@ -8144,10 +8186,8 @@ def _submit_locked(self, callable, args, config) -> RunHandle: cfg = config if config is not None else CallConfig() if self.level == 2: - assert self._chip_worker is not None - state = self._resolve_handle(callable, expected_namespace="LOCAL_CHIP") - self._chip_worker._run_slot(state.slot_id, args, cfg) - return RunHandle._completed(self) + with self._submit_mu: + return self._submit_l2_locked(callable, args, cfg) with self._submit_mu: # Graph callbacks stay serialized, so a predecessor's callback has @@ -8171,6 +8211,179 @@ def _submit_locked(self, callable, args, config) -> RunHandle: self._require_no_ordered_cleanup_failure("submit") return self._submit_l3_locked(callable, args, cfg) + @staticmethod + def _l2_config_has_diagnostics(config: CallConfig) -> bool: + return bool( + config.enable_l2_swimlane + or config.enable_dump_args + or config.enable_pmu + or config.enable_dep_gen + or config.enable_scope_stats + ) + + def _l2_finish_front_locked(self, state: _L2NativeRun, error: BaseException | None) -> None: + assert self._chip_worker is not None + if state.native_run is not None: + try: + self._chip_worker._finalize_native_run(state.native_run) + except BaseException as exc: # noqa: BLE001 + if error is None: + error = exc + state.error = error + state.phase = "terminal" + if self._l2_fifo and self._l2_fifo[0] == state.run_id: + self._l2_fifo.pop(0) + self._l2_progress_cv.notify_all() + + def _l2_prepare_front_locked(self, state: _L2NativeRun) -> None: + assert self._chip_worker is not None + assert state.phase == "deferred" + try: + state.native_run = self._chip_worker._prepare_native_run_with_pipeline_lease( + state.callable_id, + state.args, + state.slot_id, + state.generation, + state.config, + run_id=state.run_id, + dispatch_id=state.run_id, + ) + except BaseException as exc: # noqa: BLE001 + self._l2_finish_front_locked(state, exc) + return + state.phase = "prepared" + self._l2_progress_cv.notify_all() + + def _l2_launch_front_locked(self, state: _L2NativeRun) -> None: + assert self._chip_worker is not None + try: + self._chip_worker._launch_native_run(state.native_run) + except BaseException as exc: # noqa: BLE001 + self._l2_finish_front_locked(state, exc) + return + state.phase = "launched" + self._l2_progress_cv.notify_all() + + def _l2_progress_locked( # noqa: PLR0912 -- one bounded loop owns all direct-L2 phases + self, target_run_id: int, deadline: float | None, *, block: bool + ) -> bool: + """Drive the direct L2 FIFO until target terminalizes or progress would block.""" + assert self._chip_worker is not None + while True: + target = self._l2_runs.get(target_run_id) + if target is None: + raise RuntimeError(f"unknown direct L2 run id {target_run_id}") + if target.phase == "terminal": + return True + if not self._l2_fifo: + raise RuntimeError("direct L2 native lane lost a nonterminal run") + + front = self._l2_runs[self._l2_fifo[0]] + if front.phase == "deferred": + self._l2_prepare_front_locked(front) + if front.phase == "terminal": + continue + if front.phase == "prepared": + self._l2_launch_front_locked(front) + if front.phase == "terminal": + continue + + completed = False + progress_error: BaseException | None = None + try: + if block and deadline is None: + self._chip_worker._wait_native_run(front.native_run) + completed = True + else: + completed = bool(self._chip_worker._poll_native_run(front.native_run)) + except BaseException as exc: # noqa: BLE001 + progress_error = exc + completed = True + + if completed: + self._l2_finish_front_locked(front, progress_error) + # Launch the successor at the same ordered handoff boundary. + if self._l2_fifo: + successor = self._l2_runs[self._l2_fifo[0]] + if successor.phase == "deferred": + self._l2_prepare_front_locked(successor) + if successor.phase == "prepared": + self._l2_launch_front_locked(successor) + continue + + if not block: + return False + if deadline is not None: + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + time.sleep(min(remaining, _RUN_HANDLE_WAIT_RECHECK_S)) + + def _submit_l2_locked(self, callable, args, cfg: CallConfig) -> RunHandle: + assert self._chip_worker is not None + callable_state = self._resolve_handle(callable, expected_namespace="LOCAL_CHIP") + permits_successor = bool( + self._chip_worker.supports_concurrent_native_prepare and not self._l2_config_has_diagnostics(cfg) + ) + + with self._l2_progress_cv: + while self._l2_fifo: + active = self._l2_runs[self._l2_fifo[0]] + capacity = 2 if permits_successor and active.permits_successor else 1 + if len(self._l2_fifo) < capacity: + break + self._l2_progress_locked(active.run_id, None, block=True) + + occupied_slots = {self._l2_runs[run_id].slot_id for run_id in self._l2_fifo} + slot_id = next( + (slot for slot in range(self._chip_worker.pipeline_depth) if slot not in occupied_slots), None + ) + if slot_id is None: + raise RuntimeError("direct L2 native lane has no free pipeline slot after admission") + generation = self._l2_slot_generations[slot_id] + 1 + if generation >= 1 << 64: + raise RuntimeError("direct L2 pipeline generation space is exhausted") + self._l2_slot_generations[slot_id] = generation + run_id = self._l2_next_run_id + self._l2_next_run_id += 1 + + state = _L2NativeRun( + run_id, + slot_id, + generation, + callable_state.slot_id, + args, + cfg, + None, + permits_successor, + ) + try: + state.native_run = self._chip_worker._prepare_native_run_with_pipeline_lease( + callable_state.slot_id, + args, + slot_id, + generation, + cfg, + run_id=run_id, + dispatch_id=run_id, + ) + except BaseException as exc: # noqa: BLE001 + if not self._l2_fifo or not _native_prepare_requires_depth_one(exc): + raise + else: + state.phase = "prepared" + handle = RunHandle(self, run_id, (callable, args, cfg)) + state.handle = handle + self._l2_runs[run_id] = state + self._l2_fifo.append(run_id) + with self._hierarchical_start_cv: + self._accepted_run_handles.add(handle) + self._hierarchical_start_cv.notify_all() + + if self._l2_fifo[0] == run_id: + self._l2_launch_front_locked(state) + return handle + def _record_unreclaimable(self, message: str, cause: BaseException | None = None) -> RuntimeError: """Refuse all further work on this worker, and return the reason. @@ -8388,10 +8601,29 @@ def _submit_l3_locked(self, callable, args, cfg: CallConfig) -> RunHandle: return handle def _run_handle_done(self, run_id: int) -> bool: + if self.level == 2: + with self._l2_progress_cv: + self._l2_progress_locked(run_id, time.monotonic(), block=False) + state = self._l2_runs.get(run_id) + if state is None: + raise RuntimeError(f"unknown direct L2 run id {run_id}") + return state.phase == "terminal" assert self._orch is not None return self._orch._run_done(run_id) def _wait_run_handle(self, run_id: int, timeout: float | None) -> bool: + if self.level == 2: + deadline = None if timeout is None else time.monotonic() + timeout + with self._l2_progress_cv: + completed = self._l2_progress_locked(run_id, deadline, block=True) + if not completed: + return False + state = self._l2_runs.get(run_id) + if state is None: + raise RuntimeError(f"unknown direct L2 run id {run_id}") + if state.error is not None: + raise state.error + return True assert self._orch is not None if timeout is None: self._orch._wait_run(run_id) @@ -8399,9 +8631,46 @@ def _wait_run_handle(self, run_id: int, timeout: float | None) -> bool: return self._orch._wait_run_for(run_id, timeout) def _wait_run_handle_accepted(self, run_id: int) -> None: + if self.level == 2: + with self._l2_progress_cv: + while True: + state = self._l2_runs.get(run_id) + if state is None: + raise RuntimeError(f"unknown direct L2 run id {run_id}") + if state.phase in ("launched", "terminal"): + if state.error is not None: + raise state.error + return + front_run_id = self._l2_fifo[0] + if front_run_id == run_id: + if state.phase == "deferred": + self._l2_prepare_front_locked(state) + if state.phase == "prepared": + self._l2_launch_front_locked(state) + continue + self._l2_progress_locked(front_run_id, None, block=True) assert self._orch is not None self._orch._wait_run_accepted(run_id) + def _finalize_l2_run_handle( + self, handle: RunHandle, run_id: int, native_error: BaseException | None + ) -> BaseException | None: + with self._l2_progress_cv: + state = self._l2_runs.get(run_id) + if state is None or state.phase != "terminal": + native_error = native_error or RuntimeError( + f"direct L2 run {run_id} reached finalization before its native fence" + ) + elif native_error is None: + native_error = state.error + self._l2_runs.pop(run_id, None) + handle._cache_finalization_error(native_error) + with self._hierarchical_start_cv: + handle._cleanup_published = True + self._accepted_run_handles.discard(handle) + self._hierarchical_start_cv.notify_all() + return handle._finalization_error + def _finalize_run_handle( self, handle: RunHandle, @@ -8411,6 +8680,9 @@ def _finalize_run_handle( _after_step: Any | None = None, ) -> BaseException | None: """Run fence-owned cleanup exactly once and return the cached result.""" + if self.level == 2: + return self._finalize_l2_run_handle(handle, run_id, native_error) + # Two different failures, deliberately not merged. A task that failed is # this run's business and says nothing about the worker; a cleanup that # failed leaves collective device state nobody can describe, and poisons diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md index 2a4cb56be4..fac5f71de3 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md @@ -110,6 +110,40 @@ The platform side is deliberately thin: `DeviceRunnerBase` only remembers a allocator; all grow/pack/slice logic lives in `runtime_maker.cpp` (`RetainedTempBump`). +Prepared whole-run admission may stage the inactive slot while the current run +uses the one shared arena bank. Before an overlapping native prepare, TMR +resolves the full per-ring sizing key and requires it to match the existing +prebuilt-arena cache. A miss never rebuilds the active run's heap, shared +memory, or runtime image. It reports a depth-one fallback to the endpoint, +which retains the successor's FIFO position and immutable frame and retries +native preparation after the predecessor fence. + +The compatibility key is the resolved four-ring `task_window`, `heap`, and +`dep_pool` layout after per-call overrides, environment values, and defaults. +It does not include the callable symbol or code image. Each compatible or +fallback run still creates a fresh run-owned AICore stream during prepare; +AICPU streams and retained argument buffers remain slot-owned. + +The overlap boundary follows the resource owner, rather than treating all TMR +initialization alike: + +- safe during predecessor execution: the successor's slot-owned TaskArgs, + retained tensor staging, Host descriptors, Runtime object, and fresh AICore + stream; a cache hit only reads the already-stable shared-arena identity and + binds its device addresses; +- device-exclusive: a cache miss or arena resize/rebuild, launch, SDMA and + completion publication, and any mutation of shared scheduler scratch; these + wait for the active device fence; +- completion-only retirement: finalizing the completed Runtime, destroying its + run-owned AICore stream, and releasing its slot generation. The slot is not + reusable until these operations finish. + +Consequently, the single-entry prebuilt-arena cache needs no concurrent writer +path: overlapping preparation is admitted only after a matching cache lookup, +and the matching bind performs another lookup without rebuilding. Cache +mutation is restricted to depth-one preparation, prewarm, and final teardown, +all of which are outside an active/prepared overlap window. + On each trb bind, `RetainedTempBump`: - packs the run's non-child, non-empty tensors to a required size, aligning diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp index 6f5cdacb8a..a185b8ac18 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp @@ -80,6 +80,13 @@ extern "C" const PipelineContract *get_pipeline_contract(void) { return &contract; } +extern "C" int concurrent_native_prepare_supported_impl(void) { + // TMR may stage slot-private arguments while the active run is on device. + // The compatibility hook below prevents a successor from replacing the + // shared device-scratch arena during that overlap. + return 1; +} + static_assert( RUNTIME_ENV_RING_COUNT == PTO2_MAX_RING_DEPTH, "RuntimeEnv ring count must match PTO2 runtime ring depth" ); @@ -502,6 +509,35 @@ static PrebuiltRuntimeArenaCacheProbe make_prebuilt_runtime_arena_cache_probe(co return probe; } +static bool resolve_arena_sizing( + const uint64_t *ring_task_window, const uint64_t *ring_heap, const uint64_t *ring_dep_pool, ArenaSizingConfig *out +); + +extern "C" int prepared_run_config_compatible_impl( + const HostApi *api, const uint64_t *ring_task_window, const uint64_t *ring_heap, const uint64_t *ring_dep_pool +) { + if (api == nullptr || api->lookup_prebuilt_runtime_arena_cache == nullptr) { + return -1; + } + + ArenaSizingConfig sizing; + if (!resolve_arena_sizing(ring_task_window, ring_heap, ring_dep_pool, &sizing)) { + return -1; + } + PrebuiltRuntimeArenaCacheProbe probe = make_prebuilt_runtime_arena_cache_probe(sizing); + void *gm_heap = nullptr; + void *gm_sm = nullptr; + void *runtime_arena = nullptr; + size_t runtime_offset = 0; + const void *image = nullptr; + size_t image_size = 0; + bool cache_hit = api->lookup_prebuilt_runtime_arena_cache( + probe.hash, probe.serialized_key.data(), probe.serialized_key.size(), &gm_heap, &gm_sm, &runtime_arena, + &runtime_offset, &image, &image_size + ); + return cache_hit ? 1 : 0; +} + // per-(cid,config): resolve the cache-key sizing knobs. Pure host parsing over // per-task overrides, PTO2_RING_* env, and compile-time defaults. Derived // allocation sizes are computed only on cache miss. diff --git a/src/common/platform/onboard/host/c_api_shared.cpp b/src/common/platform/onboard/host/c_api_shared.cpp index c6322d65b9..7c6e6956d7 100644 --- a/src/common/platform/onboard/host/c_api_shared.cpp +++ b/src/common/platform/onboard/host/c_api_shared.cpp @@ -68,6 +68,12 @@ extern "C" { int register_callable_impl(const ChipCallable *callable, uint64_t (*upload_fn)(const void *), CallableArtifacts *out); int validate_runtime_impl(Runtime *runtime, const HostApi *api, int execution_rc); __attribute__((weak)) int concurrent_native_prepare_supported_impl(void) { return 0; } +__attribute__((weak)) int prepared_run_config_compatible_impl( + const HostApi * /*api*/, const uint64_t * /*ring_task_window*/, const uint64_t * /*ring_heap*/, + const uint64_t * /*ring_dep_pool*/ +) { + return 1; +} /* =========================================================================== * Per-thread DeviceRunnerBase binding (set by simpler_register_callable / simpler_run) @@ -708,6 +714,24 @@ int simpler_prepare_run( int rc = runner->attach_current_thread(runner->device_id()); if (rc != 0) return cleanup_failed_prepare(state, rc, true); + if (overlaps_active_run) { + int compatibility_rc = 0; + { + STRACE("simpler_run.bind.compatibility"); + compatibility_rc = prepared_run_config_compatible_impl( + &g_host_api, config->runtime_env.ring_task_window, config->runtime_env.ring_heap, + config->runtime_env.ring_dep_pool + ); + } + if (compatibility_rc <= 0) { + if (compatibility_rc == 0) { + LOG_INFO("successor RuntimeEnv requires depth-one native preparation"); + compatibility_rc = PTO_RUNTIME_ERR_PREPARED_INCOMPATIBLE; + } + return cleanup_failed_prepare(state, compatibility_rc, true); + } + } + state->runner_resources_owned = true; rc = runner->provision_native_run_resources(state->pipeline_slot); if (rc != 0) return cleanup_failed_prepare(state, rc, true); @@ -716,8 +740,8 @@ int simpler_prepare_run( if (rc != 0) return cleanup_failed_prepare(state, rc, true); // Diagnostic binding reads runner-global collector configuration. It - // is depth-one, while concurrent HBG preparation must leave the active - // run's configuration untouched until launch. + // is depth-one, while concurrent successor preparation must leave the + // active run's configuration untouched until launch. if (!overlaps_active_run) runner->apply_call_config(state->config); { diff --git a/src/common/platform/onboard/host/device_runner_base.cpp b/src/common/platform/onboard/host/device_runner_base.cpp index a8f9b4c3de..e7aa2ccaaf 100644 --- a/src/common/platform/onboard/host/device_runner_base.cpp +++ b/src/common/platform/onboard/host/device_runner_base.cpp @@ -1646,8 +1646,7 @@ bool DeviceRunnerBase::try_reserve_native_run( const NativeRunReservation *existing = nullptr; for (const NativeRunReservation &reservation : native_run_reservations_) { if (reservation.owner == nullptr) continue; - if (reservation.owner == owner || reservation.pipeline_slot == pipeline_slot || - reservation.arena_bank == arena_bank) { + if (reservation.owner == owner || reservation.pipeline_slot == pipeline_slot) { return false; } ++occupied; diff --git a/src/common/platform/onboard/host/device_runner_base.h b/src/common/platform/onboard/host/device_runner_base.h index acbf910e2b..edf550eb0c 100644 --- a/src/common/platform/onboard/host/device_runner_base.h +++ b/src/common/platform/onboard/host/device_runner_base.h @@ -121,7 +121,9 @@ class DeviceRunnerBase { /** * Reserve caller-owned native-run storage before binding starts. A * concurrent reservation is admitted only while the first reservation - * owns the execution claim and selects distinct per-run resources. + * owns the execution claim and selects a distinct pipeline slot. A backend + * that shares an arena bank must reject or defer incompatible preparation + * before mutating that bank. */ bool try_reserve_native_run( const void *owner, uint32_t pipeline_slot, uint32_t arena_bank, bool allow_prepared_successor @@ -500,9 +502,9 @@ class DeviceRunnerBase { size_t host_dlopen_count() const { return host_dlopen_total_; } /** - * Number of run stream generations this runner has created. AICPU streams - * belong to pipeline slots, while an AICore stream is reused only for the - * same AICore image. Arches whose runs use the persistent pair report 0. + * Number of run-owned AICore streams this runner has created. AICPU streams + * belong to pipeline slots, while every native prepare provisions a fresh + * AICore stream. Arches whose runs use the persistent pair report 0. */ virtual size_t run_stream_set_create_count() const { return 0; } diff --git a/src/common/worker/chip_worker.cpp b/src/common/worker/chip_worker.cpp index 42add01d73..9637d6de1d 100644 --- a/src/common/worker/chip_worker.cpp +++ b/src/common/worker/chip_worker.cpp @@ -676,7 +676,11 @@ ChipWorkerNativeRun ChipWorker::prepare_native_run_on_slot( "prepare_native_run already owns a prepared successor " + format_native_run_identity(run_identity) ); } - if (!pipeline_generations_.admit(PipelineSlotLease{slot_id, 0, generation})) { + // A compatibility probe may reject this prepare while the predecessor + // is still active. Merely attempting that probe must not consume the + // lease generation, because the same lease is retried at the ordered + // depth-one handoff boundary. + if (!pipeline_generations_.is_admissible(PipelineSlotLease{slot_id, 0, generation})) { throw std::runtime_error( "native-run pipeline lease generation is stale " + format_native_run_identity(run_identity) ); @@ -710,6 +714,11 @@ ChipWorkerNativeRun ChipWorker::prepare_native_run_on_slot( if (state.run_epoch == run_epoch && state.phase == NativeRunPhase::PREPARING) { state = NativeRunSlotState{}; } + if (rc == PTO_RUNTIME_ERR_PREPARED_INCOMPATIBLE) { + throw std::runtime_error( + "native prepare requires depth-one fallback " + format_native_run_identity(run_identity) + ); + } throw std::runtime_error( "prepare_native_run failed with code " + std::to_string(rc) + " " + format_native_run_identity(run_identity) ); @@ -724,6 +733,14 @@ ChipWorkerNativeRun ChipWorker::prepare_native_run_on_slot( state = NativeRunSlotState{}; throw std::runtime_error("native-run identity changed while prepare was in progress"); } + if (!pipeline_generations_.admit(PipelineSlotLease{slot_id, 0, generation})) { + (void)finalize_run_fn_(device_ctx_, runtime_bufs_[slot_id].data()); + state = NativeRunSlotState{}; + throw std::runtime_error( + "native-run pipeline lease generation became stale while prepare was in progress " + + format_native_run_identity(run_identity) + ); + } state.phase = NativeRunPhase::PREPARED; } return run_identity; diff --git a/src/common/worker/chip_worker.h b/src/common/worker/chip_worker.h index 36ab118e38..b955fe6dc5 100644 --- a/src/common/worker/chip_worker.h +++ b/src/common/worker/chip_worker.h @@ -118,10 +118,10 @@ class ChipWorker { * the token, and the caller must still finalize it. The blocking composition * performs that cleanup internally on every exit. * - * Onboard HBG may prepare one distinct-slot successor while another run - * owns the execution claim. Diagnostics and backends without the explicit - * capability remain depth-one. The slot/lease-generation/process-unique- - * run-epoch token prevents a delayed phase call from touching reused + * A capable backend may prepare one distinct-slot successor while another + * run owns the execution claim. Diagnostics and backends without that + * explicit capability remain depth-one. The slot/lease-generation/process- + * unique-run-epoch token prevents a delayed phase call from touching reused * storage, including another run under the same pipeline lease or on * another ChipWorker. */ diff --git a/src/common/worker/pipeline_slot_pool.h b/src/common/worker/pipeline_slot_pool.h index 59c1cb08fe..4f1832f5c0 100644 --- a/src/common/worker/pipeline_slot_pool.h +++ b/src/common/worker/pipeline_slot_pool.h @@ -113,6 +113,12 @@ class PipelineSlotPool { */ class PipelineSlotGenerationFilter { public: + bool is_admissible(const PipelineSlotLease &lease) { + if (lease.slot_id >= newest_.size()) return false; + std::lock_guard lock(mu_); + return lease.generation >= newest_[lease.slot_id]; + } + bool admit(const PipelineSlotLease &lease) { if (lease.slot_id >= newest_.size()) return false; std::lock_guard lock(mu_); diff --git a/src/common/worker/pto_runtime_c_api.h b/src/common/worker/pto_runtime_c_api.h index 179d4f3222..8cdb59d363 100644 --- a/src/common/worker/pto_runtime_c_api.h +++ b/src/common/worker/pto_runtime_c_api.h @@ -76,6 +76,7 @@ typedef void *DeviceContextHandle; enum { PTO_RUNTIME_ERR_UNSUPPORTED = -2, + PTO_RUNTIME_ERR_PREPARED_INCOMPATIBLE = -3, }; /** Return values from simpler_poll_run(). */ diff --git a/tests/st/a2a3/host_build_graph/native_run_lifecycle/kernels/orchestration/long_vector_orch.cpp b/tests/st/a2a3/host_build_graph/native_run_lifecycle/kernels/orchestration/long_vector_orch.cpp index d788c355b6..0c2dc40cf6 100644 --- a/tests/st/a2a3/host_build_graph/native_run_lifecycle/kernels/orchestration/long_vector_orch.cpp +++ b/tests/st/a2a3/host_build_graph/native_run_lifecycle/kernels/orchestration/long_vector_orch.cpp @@ -17,7 +17,7 @@ namespace { constexpr uint64_t kAdd = 0; constexpr uint64_t kAddScalar = 1; -constexpr int kChainLength = 64; +constexpr int kChainLength = 512; } // namespace diff --git a/tests/st/a2a3/host_build_graph/native_run_lifecycle/test_native_run_lifecycle.py b/tests/st/a2a3/host_build_graph/native_run_lifecycle/test_native_run_lifecycle.py index 022f697623..c1d1ac9903 100644 --- a/tests/st/a2a3/host_build_graph/native_run_lifecycle/test_native_run_lifecycle.py +++ b/tests/st/a2a3/host_build_graph/native_run_lifecycle/test_native_run_lifecycle.py @@ -23,7 +23,7 @@ _SLOT = 0 _GENERATION = 1 _SIZE = 128 * 128 -_CHAIN_LENGTH = 64 +_CHAIN_LENGTH = 512 @scene_test(level=2, runtime="host_build_graph") @@ -74,7 +74,7 @@ def test_run(self, st_platform, st_worker, request, capfd): spans = list(parse_spans(capfd.readouterr().err.splitlines())) invocations = [inv for inv in group_invocations(spans) if "simpler_run" in inv.by_name()] - expected_invocations = 3 if st_platform.endswith("sim") else 7 + expected_invocations = 4 if st_platform.endswith("sim") else 8 assert len(invocations) == expected_invocations common_depths = { @@ -103,7 +103,7 @@ def test_run(self, st_platform, st_worker, request, capfd): for name in expected_depths.keys() - {"simpler_run", "simpler_run.runner_run.device_wall"}: stage = by_name[name] assert root.ts <= stage.ts <= stage.ts + stage.dur <= root_end - expected_launched = 2 if st_platform.endswith("sim") else 6 + expected_launched = 3 if st_platform.endswith("sim") else 7 assert launched_count == expected_launched if not st_platform.endswith("sim"): root_attrs = [inv.by_name()["simpler_run"].attrs for inv in invocations] @@ -133,8 +133,10 @@ def _run_and_validate_l2( # noqa: PLR0913, PLR0915 -- lifecycle contract is int config = self._build_config(case["config"]) chip_worker = worker._chip_worker assert chip_worker is not None - supports_concurrent_prepare = bool(chip_worker._impl.supports_concurrent_native_prepare) + supports_concurrent_prepare = chip_worker.supports_concurrent_native_prepare chip_worker._register_callable_at_slot(_SLOT, callable_obj) + private_slot_registered = True + public_handle = None native_run = None successor_run = None try: @@ -148,8 +150,7 @@ def _run_and_validate_l2( # noqa: PLR0913, PLR0915 -- lifecycle contract is int _SLOT, chip_args, _SLOT, _GENERATION, config=config ) first_run = native_run - expected_stream_count = stream_count_before_prepare + int(supports_concurrent_prepare) - assert chip_worker.run_stream_set_create_count == expected_stream_count + assert chip_worker.run_stream_set_create_count == stream_count_before_prepare + 1 assert torch.count_nonzero(test_args.out) == 0, "prepare crossed the device launch fence" with pytest.raises(RuntimeError, match="unfinished native run|owns the runner|active predecessor"): chip_worker._prepare_native_run_with_pipeline_lease(_SLOT, chip_args, 1, _GENERATION, config=config) @@ -208,15 +209,11 @@ def build_run_args(): self.compute_golden(run_golden, case["params"]) return run_args, run_chip_args, run_output_names, run_golden - # The successor owns a distinct bank and fresh stream while A - # still owns the execution claim. A failed early launch must - # leave B prepared so the same token can launch after A's - # complete fence and finalization. active_args, active_chip_args, active_outputs, active_golden = build_run_args() successor_args, successor_chip_args, successor_outputs, successor_golden = build_run_args() stream_count = chip_worker.run_stream_set_create_count native_run = chip_worker._prepare_native_run_with_pipeline_lease( - _SLOT, active_chip_args, 0, _GENERATION + 1, config=config + _SLOT, active_chip_args, 0, _GENERATION, config=config ) assert chip_worker.run_stream_set_create_count == stream_count + 1 chip_worker._launch_native_run(native_run) @@ -253,37 +250,49 @@ def build_run_args(): diagnostic_config.enable_dep_gen = True diagnostic_config.output_prefix = output_dir - # A diagnostic successor cannot overlap an ordinary active - # run, even though the predecessor otherwise permits one. active_args, active_chip_args, active_outputs, active_golden = build_run_args() native_run = chip_worker._prepare_native_run_with_pipeline_lease( - _SLOT, active_chip_args, 0, _GENERATION + 2, config=config + _SLOT, active_chip_args, 0, _GENERATION, config=config ) chip_worker._launch_native_run(native_run) with pytest.raises(RuntimeError, match="active predecessor"): chip_worker._prepare_native_run_with_pipeline_lease( - _SLOT, successor_chip_args, 1, _GENERATION + 1, config=diagnostic_config + _SLOT, successor_chip_args, 1, _GENERATION, config=diagnostic_config ) chip_worker._wait_native_run(native_run) chip_worker._finalize_native_run(native_run) native_run = None _compare_outputs(active_args, active_golden, active_outputs, self.RTOL, self.ATOL) - # A diagnostic predecessor also cannot admit an ordinary - # successor while it owns the execution claim. active_args, active_chip_args, active_outputs, active_golden = build_run_args() native_run = chip_worker._prepare_native_run_with_pipeline_lease( - _SLOT, active_chip_args, 0, _GENERATION + 3, config=diagnostic_config + _SLOT, active_chip_args, 0, _GENERATION, config=diagnostic_config ) chip_worker._launch_native_run(native_run) with pytest.raises(RuntimeError, match="active predecessor"): chip_worker._prepare_native_run_with_pipeline_lease( - _SLOT, successor_chip_args, 1, _GENERATION + 1, config=config + _SLOT, successor_chip_args, 1, _GENERATION, config=config ) chip_worker._wait_native_run(native_run) chip_worker._finalize_native_run(native_run) native_run = None _compare_outputs(active_args, active_golden, active_outputs, self.RTOL, self.ATOL) + + chip_worker._unregister_slot(_SLOT) + private_slot_registered = False + public_handle = worker.register(callable_obj) + async_args = self.generate_args(case["params"]) + async_chip_args, async_output_names = _build_chip_task_args( + async_args, self.CALLABLE["orchestration"]["signature"] + ) + async_golden = async_args.clone() + self.compute_golden(async_golden, case["params"]) + run_handle = worker.submit(public_handle, args=async_chip_args, config=config) + assert not run_handle._terminal, "direct L2 submit returned a pre-completed compatibility handle" + run_handle.wait(30.0) + _compare_outputs(async_args, async_golden, async_output_names, self.RTOL, self.ATOL) + worker.unregister(public_handle) + public_handle = None finally: for unfinished_run in (successor_run, native_run): if unfinished_run is None: @@ -292,7 +301,10 @@ def build_run_args(): chip_worker._finalize_native_run(unfinished_run) except Exception: pass - chip_worker._unregister_slot(_SLOT) + if public_handle is not None: + worker.unregister(public_handle) + elif private_slot_registered: + chip_worker._unregister_slot(_SLOT) if __name__ == "__main__": diff --git a/tests/st/a2a3/host_build_graph/run_stream_reuse/test_run_stream_reuse.py b/tests/st/a2a3/host_build_graph/run_stream_reuse/test_run_stream_reuse.py index 07ef9875d7..7f7daa4183 100644 --- a/tests/st/a2a3/host_build_graph/run_stream_reuse/test_run_stream_reuse.py +++ b/tests/st/a2a3/host_build_graph/run_stream_reuse/test_run_stream_reuse.py @@ -139,7 +139,7 @@ def test_every_run_creates_its_own_aicore_stream(self, st_platform, st_worker): f"{after_first} -> {st_worker.run_stream_set_create_count}" ) - def _run_registered(self, worker, handle, *, subtract): + def _run_registered(self, worker, handle, *, subtract, unleased=False): params = self.CASES[0]["params"] test_args = self.generate_args(params) chip_args, output_names = _build_chip_task_args(test_args, self.CALLABLE["orchestration"]["signature"]) @@ -147,7 +147,12 @@ def _run_registered(self, worker, handle, *, subtract): a, b = golden_args.a, golden_args.b base = a - b if subtract else a + b golden_args.f[:] = (base + 1) * (base + 2) - worker.run(handle, chip_args, config=self._build_config(self.CASES[0]["config"])) + config = self._build_config(self.CASES[0]["config"]) + if unleased: + state = worker._resolve_handle(handle) + worker._chip_worker._run_slot(state.slot_id, chip_args, config=config) + else: + worker.run(handle, chip_args, config=config) _compare_outputs(test_args, golden_args, output_names, self.RTOL, self.ATOL) # The st_worker fixture is shared by every test in this class, and so is the @@ -258,7 +263,7 @@ def test_unleased_runs_do_not_consume_lease_generations(self, st_platform, st_wo handle = other.register(self.build_callable(st_platform)) try: for _ in range(3): - self._run_registered(other, handle, subtract=False) + self._run_registered(other, handle, subtract=False, unleased=True) self._run_registered_with_lease(other, handle, slot_id=0, generation=1) finally: other.unregister(handle) @@ -285,3 +290,25 @@ def test_depth_two_slot_is_generation_safe(self, st_platform, st_worker): assert st_worker.run_stream_set_create_count == stream_sets finally: st_worker.unregister(add_handle) + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestRunStreamFreshTmr(SceneTestCase): + """TMR preparation keeps the same fresh run-owned AICore stream rule.""" + + CALLABLE = TestRunStreamReuseHbg.CALLABLE + CASES = TestRunStreamReuseHbg.CASES + + generate_args = TestRunStreamReuseHbg.generate_args + compute_golden = TestRunStreamReuseHbg.compute_golden + + def test_every_run_creates_its_own_aicore_stream(self, st_platform, st_worker): + if st_platform != "a2a3": + pytest.skip("run stream sets are an a2a3 onboard resource") + + callable_obj = self.build_callable(st_platform) + self._run_and_validate_l2(st_worker, callable_obj, self.CASES[0], rounds=1) + after_first = st_worker.run_stream_set_create_count + rounds = _REPEATED_RUNS - 1 + self._run_and_validate_l2(st_worker, callable_obj, self.CASES[0], rounds=rounds) + assert st_worker.run_stream_set_create_count == after_first + rounds diff --git a/tests/st/a2a3/host_build_graph/worker_async_endpoint/test_worker_async_endpoint.py b/tests/st/a2a3/host_build_graph/worker_async_endpoint/test_worker_async_endpoint.py index b3fc24d060..a06ab66479 100644 --- a/tests/st/a2a3/host_build_graph/worker_async_endpoint/test_worker_async_endpoint.py +++ b/tests/st/a2a3/host_build_graph/worker_async_endpoint/test_worker_async_endpoint.py @@ -210,10 +210,10 @@ def second_graph(orch, _args, _cfg): assert accepted == 0, "the successor crossed its launch fence before activation" saw_staged = True break - assert accepted == 0, "the successor crossed FIFO activation while its predecessor run remained open" + assert accepted == 0, "the successor launched while its predecessor was still active" time.sleep(0.001) - assert saw_staged, "the successor did not reach FRAME_STAGED while its predecessor run remained open" + assert saw_staged, "the successor did not reach FRAME_STAGED behind its predecessor" assert not first.done, "the predecessor escaped its SubTask fence" _SUB_RELEASE.set() diff --git a/tests/st/a2a3/host_build_graph/worker_async_fifo/test_worker_async_fifo.py b/tests/st/a2a3/host_build_graph/worker_async_fifo/test_worker_async_fifo.py index ad3b8f5c3e..6f2bf05e87 100644 --- a/tests/st/a2a3/host_build_graph/worker_async_fifo/test_worker_async_fifo.py +++ b/tests/st/a2a3/host_build_graph/worker_async_fifo/test_worker_async_fifo.py @@ -479,5 +479,84 @@ def submit_third(): st_worker.free_host_buffer(buffer) +@scene_test(level=3, runtime="tensormap_and_ringbuffer") +class TestWorkerAsyncWholeRunFifoTmr(TestWorkerAsyncWholeRunFifo): + """TMR uses the common FIFO with shared-arena compatibility fallback.""" + + def test_incompatible_runtime_env_falls_back_to_depth_one(self, st_platform, st_worker): + if st_platform != "a2a3": + pytest.skip("TMR prepared-state validation requires an a2a3 onboard worker") + + _SUB_ENTERED.clear() + _SUB_RELEASE.clear() + buffers = [] + tensors = [] + first = None + second = None + try: + for value in (2.0, 3.0, 0.0, 5.0, 7.0, 0.0): + buffer, tensor = self._tensor_from_host_buffer(st_worker, value) + buffers.append(buffer) + tensors.append(tensor) + first_a, first_b, first_out, second_a, second_b, second_out = tensors + vector_handle = type(self)._st_chip_handles["vector"] + vector_signature = type(self)._st_chip_handles["vector_sig"] + sub_handle = type(self)._st_sub_handles["wait_for_release"] + + def submit_vector(orch, a, b, out, config, *, spin_iters=0, hold_open=False): + builder = TaskArgsBuilder( + Tensor("a", a), Tensor("b", b), Tensor("f", out), Scalar("spin_iters", spin_iters) + ) + chip_args, _ = _build_l3_task_args(builder, vector_signature) + orch.submit_next_level(vector_handle, chip_args, config, worker=0) + if hold_open: + orch.submit_sub(sub_handle) + + default_config = self._build_config(self.CASES[0]["config"]) + incompatible_config = self._build_config( + { + **self.CASES[0]["config"], + "runtime_env": {"ring_heap": [128 * 1024 * 1024] * 4}, + } + ) + first = st_worker.submit( + lambda orch, _args, _cfg: submit_vector( + orch, + first_a, + first_b, + first_out, + default_config, + spin_iters=_DEVICE_SPIN_ITERS, + hold_open=True, + ) + ) + _wait_for_active_device_run(st_worker, 10.0) + assert _SUB_ENTERED.wait(10.0) + + second = st_worker.submit( + lambda orch, _args, _cfg: submit_vector(orch, second_a, second_b, second_out, incompatible_config) + ) + _wait_for_backend_prepared_successor(st_worker, 10.0) + assert torch.count_nonzero(second_out).item() == 0 + + _SUB_RELEASE.set() + first.wait(30.0) + second.wait(30.0) + assert torch.allclose(first_out, first_a + first_b + _CHAIN_LENGTH) + assert torch.allclose(second_out, second_a + second_b + _CHAIN_LENGTH) + finally: + _SUB_RELEASE.set() + for handle in (first, second): + if handle is not None: + with suppress(Exception): + handle.wait(30.0) + tensors.clear() + tensor = None + first_a = first_b = first_out = second_a = second_b = second_out = None + if all(handle is None or handle.done for handle in (first, second)): + for buffer in buffers: + st_worker.free_host_buffer(buffer) + + if __name__ == "__main__": SceneTestCase.run_module(__name__) diff --git a/tests/ut/cpp/common/test_trb_runtime_temp_buffer.cpp b/tests/ut/cpp/common/test_trb_runtime_temp_buffer.cpp index 30dba36523..fbac01f6ce 100644 --- a/tests/ut/cpp/common/test_trb_runtime_temp_buffer.cpp +++ b/tests/ut/cpp/common/test_trb_runtime_temp_buffer.cpp @@ -41,6 +41,10 @@ extern "C" int bind_callable_to_runtime_impl( const uint64_t *ring_dep_pool ); extern "C" int validate_runtime_impl(Runtime *runtime, const HostApi *api, int execution_rc); +extern "C" int concurrent_native_prepare_supported_impl(void); +extern "C" int prepared_run_config_compatible_impl( + const HostApi *api, const uint64_t *ring_task_window, const uint64_t *ring_heap, const uint64_t *ring_dep_pool +); namespace { @@ -66,6 +70,11 @@ struct FakeHostApi { std::vector gm_heap; std::vector gm_sm; std::vector runtime_arena; + bool compatibility_key_valid = false; + uint64_t compatibility_hash = 0; + std::vector compatibility_key; + uint64_t observed_hash = 0; + std::vector observed_key; ~FakeHostApi() { release_all(); } @@ -168,6 +177,25 @@ bool fake_lookup_prebuilt_runtime_arena_cache( ) { return false; } +bool fake_lookup_compatible_runtime_arena_cache( + uint64_t hash, const void *key_data, size_t key_size, void **gm_heap_base, void **sm_base, + void **runtime_arena_base, size_t *runtime_off, const void **image_data, size_t *image_size +) { + const auto *key = static_cast(key_data); + g_fake->observed_hash = hash; + g_fake->observed_key.assign(key, key + key_size); + bool hit = g_fake->compatibility_key_valid && hash == g_fake->compatibility_hash && + g_fake->observed_key == g_fake->compatibility_key; + if (hit) { + *gm_heap_base = reinterpret_cast(1); + *sm_base = reinterpret_cast(2); + *runtime_arena_base = reinterpret_cast(3); + *runtime_off = 4; + *image_data = reinterpret_cast(5); + *image_size = 6; + } + return hit; +} void fake_mark_prebuilt_runtime_arena_cached( uint64_t /* hash */, const void * /* key_data */, size_t /* key_size */, void * /* gm_heap_base */, void * /* sm_base */, void * /* runtime_arena_base */, size_t /* runtime_off */, const void * /* image_data */, @@ -198,6 +226,12 @@ HostApi make_host_api(bool with_temporary_buffer = true) { }; } +HostApi make_compatibility_host_api() { + HostApi api = make_host_api(); + api.lookup_prebuilt_runtime_arena_cache = fake_lookup_compatible_runtime_arena_cache; + return api; +} + Tensor make_tensor(std::vector &storage, bool child_memory = false) { Tensor tensor; uint32_t shape[1] = {static_cast(storage.size())}; @@ -465,3 +499,22 @@ TEST_F(TrbRuntimeTempBufferTest, FailedCopyOnTemporaryPathDoesNotFreeRetainedBuf EXPECT_NE(fake_.retained_addr, nullptr); EXPECT_TRUE(runtime.tensor_leases_.empty()); } + +TEST_F(TrbRuntimeTempBufferTest, PreparedRuntimeEnvRequiresTheActiveArenaKey) { + fake_.reset(); + HostApi compatibility_api = make_compatibility_host_api(); + uint64_t task_window[PTO2_MAX_RING_DEPTH] = {4, 4, 4, 4}; + uint64_t heap[PTO2_MAX_RING_DEPTH] = {1024, 1024, 1024, 1024}; + uint64_t dep_pool[PTO2_MAX_RING_DEPTH] = {4, 4, 4, 4}; + + EXPECT_EQ(concurrent_native_prepare_supported_impl(), 1); + EXPECT_EQ(prepared_run_config_compatible_impl(&compatibility_api, task_window, heap, dep_pool), 0); + fake_.compatibility_key_valid = true; + fake_.compatibility_hash = fake_.observed_hash; + fake_.compatibility_key = fake_.observed_key; + + EXPECT_EQ(prepared_run_config_compatible_impl(&compatibility_api, task_window, heap, dep_pool), 1); + heap[2] = 2048; + EXPECT_EQ(prepared_run_config_compatible_impl(&compatibility_api, task_window, heap, dep_pool), 0); + EXPECT_NE(fake_.observed_key, fake_.compatibility_key); +} diff --git a/tests/ut/cpp/hierarchical/test_pipeline_contract.cpp b/tests/ut/cpp/hierarchical/test_pipeline_contract.cpp index 1a01c91b90..98b01cf07c 100644 --- a/tests/ut/cpp/hierarchical/test_pipeline_contract.cpp +++ b/tests/ut/cpp/hierarchical/test_pipeline_contract.cpp @@ -35,7 +35,7 @@ TEST(PipelineContract, AcceptsADeclarationThisBuildCanHonor) { EXPECT_TRUE(is_valid_pipeline_contract(&c)); } -// Runtime loading requires a contract; this helper still rejects malformed null values. +// A runtime that exports no contract is handled by the caller, not here. TEST(PipelineContract, RejectsNull) { EXPECT_FALSE(is_valid_pipeline_contract(nullptr)); } TEST(PipelineContract, AcceptsAnEmptyResourceList) { @@ -242,4 +242,22 @@ TEST(PipelineSlotPool, StaleGenerationCannotAccessOrReleaseAReusedSlot) { EXPECT_TRUE(pool.release(replacement)); } +TEST(PipelineSlotGenerationFilter, CompatibilityPreviewDoesNotConsumeGeneration) { + PipelineSlotGenerationFilter filter; + const PipelineSlotLease retried{1, 0, 2}; + + EXPECT_TRUE(filter.is_admissible(retried)); + EXPECT_TRUE(filter.admit(PipelineSlotLease{1, 0, 1})); + EXPECT_TRUE(filter.admit(retried)); +} + +TEST(PipelineSlotGenerationFilter, CommitRechecksGenerationAfterCompatibilityPreview) { + PipelineSlotGenerationFilter filter; + const PipelineSlotLease delayed{0, 0, 1}; + + EXPECT_TRUE(filter.is_admissible(delayed)); + EXPECT_TRUE(filter.admit(PipelineSlotLease{0, 0, 2})); + EXPECT_FALSE(filter.admit(delayed)); +} + } // namespace diff --git a/tests/ut/py/test_worker/test_host_worker.py b/tests/ut/py/test_worker/test_host_worker.py index 97b156df1c..0b7b66d372 100644 --- a/tests/ut/py/test_worker/test_host_worker.py +++ b/tests/ut/py/test_worker/test_host_worker.py @@ -519,7 +519,7 @@ def test_two_frame_stages_b_without_native_prepare_until_a_finalizes(): harness.close() -def test_two_frame_hbg_prepares_b_while_a_runs_but_accepts_only_after_launch(): +def test_two_frame_capable_backend_prepares_b_while_a_runs_but_accepts_only_after_launch(): harness = _TwoFrameLoopHarness( supports_concurrent_native_prepare=True, chip_runtime="host_build_graph", @@ -560,7 +560,7 @@ def test_two_frame_hbg_prepares_b_while_a_runs_but_accepts_only_after_launch(): harness.close() -def test_two_frame_hbg_publishes_failure_instead_of_staged_when_prepare_fails(): +def test_two_frame_capable_backend_publishes_failure_instead_of_staged_when_prepare_fails(): harness = _TwoFrameLoopHarness( supports_concurrent_native_prepare=True, chip_runtime="host_build_graph", @@ -577,7 +577,45 @@ def test_two_frame_hbg_publishes_failure_instead_of_staged_when_prepare_fails(): harness.close() -def test_two_frame_hbg_waits_for_first_token_to_launch_before_preparing_second(): +def test_two_frame_incompatible_successor_stays_staged_and_prepares_after_predecessor(): + harness = _TwoFrameLoopHarness( + supports_concurrent_native_prepare=True, + chip_runtime="tensormap_and_ringbuffer", + ) + try: + harness.publish(0, 1) + harness.start() + assert harness.cw._impl.launched[0].wait(5.0) + + fallback = RuntimeError("native prepare requires depth-one fallback run_id=2 slot=1") + harness.cw._impl.prepare_errors[(1, 11)] = fallback + harness.publish(1, 2, state=worker_mod._PREPARE_READY) + harness.wait_state(1, worker_mod._FRAME_STAGED) + assert not harness.cw._impl.prepared[1].is_set() + assert not harness.cw._impl.finalized[0].is_set() + + del harness.cw._impl.prepare_errors[(1, 11)] + _mailbox_store_i32(harness.state_addr(1), worker_mod._ACTIVATE) + harness.cw._impl.completed[0].set() + assert harness.cw._impl.finalized[0].wait(5.0) + assert harness.cw._impl.launched[1].wait(5.0) + harness.cw._impl.completed[1].set() + harness.wait_state(1, worker_mod._TASK_DONE) + + lifecycle = [event[:2] for event in harness.cw._impl.events if event[0] in {"prepare", "launch", "finalize"}] + assert lifecycle == [ + ("prepare", 0), + ("launch", 0), + ("finalize", 0), + ("prepare", 1), + ("launch", 1), + ("finalize", 1), + ] + finally: + harness.close() + + +def test_two_frame_capable_backend_waits_for_first_token_to_launch_before_preparing_second(): harness = _TwoFrameLoopHarness( supports_concurrent_native_prepare=True, chip_runtime="host_build_graph", @@ -606,7 +644,7 @@ def test_two_frame_hbg_waits_for_first_token_to_launch_before_preparing_second() harness.close() -def test_two_frame_hbg_prepares_and_launches_reverse_ready_frames_by_dispatch_id(): +def test_two_frame_capable_backend_prepares_and_launches_reverse_ready_frames_by_dispatch_id(): harness = _TwoFrameLoopHarness( supports_concurrent_native_prepare=True, chip_runtime="host_build_graph", @@ -637,7 +675,7 @@ def test_two_frame_hbg_prepares_and_launches_reverse_ready_frames_by_dispatch_id harness.close() -def test_two_frame_hbg_does_not_prepare_high_dispatch_successor_before_active_frame(): +def test_two_frame_capable_backend_does_not_prepare_high_dispatch_successor_before_active_frame(): harness = _TwoFrameLoopHarness( supports_concurrent_native_prepare=True, chip_runtime="host_build_graph", @@ -667,7 +705,7 @@ def test_two_frame_hbg_does_not_prepare_high_dispatch_successor_before_active_fr @pytest.mark.parametrize("diagnostic_frame", ["active", "successor"]) -def test_two_frame_hbg_defers_diagnostic_native_prepare_until_predecessor_finalizes(diagnostic_frame): +def test_two_frame_capable_backend_defers_diagnostic_prepare_until_predecessor_finalizes(diagnostic_frame): harness = _TwoFrameLoopHarness( supports_concurrent_native_prepare=True, chip_runtime="host_build_graph", @@ -740,7 +778,7 @@ def test_two_frame_prepare_ready_waits_for_sticky_activation(): harness.close() -def test_two_frame_hbg_lone_prepare_ready_stages_before_native_prepare(): +def test_two_frame_capable_backend_lone_prepare_ready_stages_before_native_prepare(): harness = _TwoFrameLoopHarness( supports_concurrent_native_prepare=True, chip_runtime="host_build_graph", @@ -760,7 +798,7 @@ def test_two_frame_hbg_lone_prepare_ready_stages_before_native_prepare(): harness.close() -def test_two_frame_hbg_prepares_already_staged_successor_after_active_claim(): +def test_two_frame_capable_backend_prepares_already_staged_successor_after_active_claim(): harness = _TwoFrameLoopHarness( supports_concurrent_native_prepare=True, chip_runtime="host_build_graph", @@ -2363,6 +2401,175 @@ def orch(o, args, cfg): counter_shm.unlink() +class TestDirectL2RunHandle: + class FakeChipWorker: + pipeline_depth = 2 + + def __init__(self, *, concurrent: bool) -> None: + self.supports_concurrent_native_prepare = concurrent + self.events: list[tuple] = [] + self.completed = [threading.Event(), threading.Event()] + self.wait_entered = threading.Event() + self.depth_one_once: set[int] = set() + + def _prepare_native_run_with_pipeline_lease( + self, callable_id, args, slot_id, generation, config, *, run_id, dispatch_id + ): + del args, config + if run_id in self.depth_one_once: + self.depth_one_once.remove(run_id) + raise RuntimeError(f"native prepare requires depth-one fallback run_id={run_id}") + token = SimpleNamespace( + callable_id=callable_id, + slot_id=slot_id, + generation=generation, + run_id=run_id, + dispatch_id=dispatch_id, + ) + self.events.append(("prepare", run_id, slot_id, generation, dispatch_id)) + return token + + def _launch_native_run(self, token) -> None: + self.events.append(("launch", token.run_id, token.slot_id)) + + def _poll_native_run(self, token) -> bool: + self.events.append(("poll", token.run_id)) + return self.completed[token.slot_id].is_set() + + def _wait_native_run(self, token) -> None: + self.events.append(("wait", token.run_id)) + self.wait_entered.set() + assert self.completed[token.slot_id].wait(5.0) + + def _finalize_native_run(self, token) -> None: + self.events.append(("finalize", token.run_id, token.slot_id)) + + @staticmethod + def _worker(*, concurrent: bool): + worker = Worker(level=2, device_id=0, platform="a2a3", runtime="host_build_graph") + callable_handle = worker.register(_unique_chip_callable(31)) + fake = TestDirectL2RunHandle.FakeChipWorker(concurrent=concurrent) + worker._chip_worker = cast(Any, fake) + return worker, callable_handle, fake + + @staticmethod + def _config(): + return SimpleNamespace( + enable_l2_swimlane=0, + enable_dump_args=0, + enable_pmu=0, + enable_dep_gen=0, + enable_scope_stats=0, + ) + + def test_submit_returns_after_launch_before_native_completion(self): + worker, callable_handle, fake = self._worker(concurrent=False) + + handle = worker._submit_l2_locked(callable_handle, object(), cast(Any, self._config())) + + assert isinstance(handle, RunHandle) + assert fake.events[:2] == [("prepare", 1, 0, 1, 1), ("launch", 1, 0)] + assert not handle.done + fake.completed[0].set() + handle.wait(5.0) + assert handle.done + assert ("finalize", 1, 0) in fake.events + assert not worker._accepted_run_handles + + def test_capable_lane_keeps_one_launched_and_one_prepared_and_backpressures_third(self): + worker, callable_handle, fake = self._worker(concurrent=True) + first = worker._submit_l2_locked(callable_handle, object(), cast(Any, self._config())) + second = worker._submit_l2_locked(callable_handle, object(), cast(Any, self._config())) + + assert fake.events[:3] == [ + ("prepare", 1, 0, 1, 1), + ("launch", 1, 0), + ("prepare", 2, 1, 1, 2), + ] + assert worker._l2_runs[1].phase == "launched" + assert worker._l2_runs[2].phase == "prepared" + + third_result: dict[str, RunHandle] = {} + submitter = threading.Thread( + target=lambda: third_result.setdefault( + "handle", worker._submit_l2_locked(callable_handle, object(), cast(Any, self._config())) + ) + ) + submitter.start() + assert fake.wait_entered.wait(1.0) + assert submitter.is_alive() + assert all(event[:2] != ("prepare", 3) for event in fake.events) + + fake.completed[0].set() + submitter.join(5.0) + assert not submitter.is_alive() + third = third_result["handle"] + assert ("finalize", 1, 0) in fake.events + assert ("launch", 2, 1) in fake.events + assert ("prepare", 3, 0, 2, 3) in fake.events + assert worker._l2_runs[2].phase == "launched" + assert worker._l2_runs[3].phase == "prepared" + + fake.completed[1].set() + fake.completed[0].set() + third.wait(5.0) + first.wait(5.0) + second.wait(5.0) + assert not worker._accepted_run_handles + assert [event[1] for event in fake.events if event[0] == "launch"] == [1, 2, 3] + + def test_successor_acceptance_wait_returns_after_launch_before_completion(self): + worker, callable_handle, fake = self._worker(concurrent=True) + first = worker._submit_l2_locked(callable_handle, object(), cast(Any, self._config())) + second = worker._submit_l2_locked(callable_handle, object(), cast(Any, self._config())) + accepted = threading.Event() + + waiter = threading.Thread(target=lambda: (second._wait_for_acceptance(), accepted.set())) + waiter.start() + assert fake.wait_entered.wait(1.0) + assert not accepted.is_set() + + fake.completed[0].set() + waiter.join(5.0) + assert accepted.is_set() + assert worker._l2_runs[2].phase == "launched" + assert not fake.completed[1].is_set() + + fake.completed[1].set() + first.wait(5.0) + second.wait(5.0) + + def test_incompatible_successor_defers_prepare_without_failing_its_handle(self): + worker, callable_handle, fake = self._worker(concurrent=True) + first = worker._submit_l2_locked(callable_handle, object(), cast(Any, self._config())) + fake.depth_one_once.add(2) + + second = worker._submit_l2_locked(callable_handle, object(), cast(Any, self._config())) + assert worker._l2_runs[2].phase == "deferred" + assert not second.done + assert all(event[:2] != ("prepare", 2) for event in fake.events) + + accepted = threading.Event() + waiter = threading.Thread(target=lambda: (second._wait_for_acceptance(), accepted.set())) + waiter.start() + assert fake.wait_entered.wait(1.0) + assert not accepted.is_set() + + fake.completed[0].set() + waiter.join(5.0) + assert accepted.is_set() + assert ("finalize", 1, 0) in fake.events + assert ("prepare", 2, 1, 1, 2) in fake.events + assert ("launch", 2, 1) in fake.events + assert not fake.completed[1].is_set() + + fake.completed[1].set() + first.wait(5.0) + second.wait(5.0) + assert second.done + assert not worker._accepted_run_handles + + class TestRunHandle: @staticmethod def _submission_failure_worker(failures: int): diff --git a/tests/ut/py/test_worker/test_startup_readiness.py b/tests/ut/py/test_worker/test_startup_readiness.py index 57798c58bb..60a85a481c 100644 --- a/tests/ut/py/test_worker/test_startup_readiness.py +++ b/tests/ut/py/test_worker/test_startup_readiness.py @@ -689,6 +689,13 @@ def owner_body(): class _FakeChipOk: """Stand-in for a ChipWorker whose init succeeds — no NPU touched.""" + pipeline_depth = 1 + supports_concurrent_native_prepare = False + + def __init__(self): + self.events = [] + self.complete = False + def init(self, *_a, **_k): pass @@ -698,6 +705,28 @@ def _register_callable_at_slot(self, *_a, **_k): # pragma: no cover def _run_slot(self, *_a, **_k): pass + def _prepare_native_run_with_pipeline_lease( + self, callable_id, args, slot_id, generation, config, *, run_id, dispatch_id + ): + del args, config + token = (callable_id, slot_id, generation, run_id, dispatch_id) + self.events.append(("prepare", token)) + return token + + def _launch_native_run(self, token): + self.events.append(("launch", token)) + + def _poll_native_run(self, token): + self.events.append(("poll", token)) + return self.complete + + def _wait_native_run(self, token): + self.events.append(("wait", token)) + self.complete = True + + def _finalize_native_run(self, token): + self.events.append(("finalize", token)) + def finalize(self): pass @@ -761,7 +790,7 @@ def test_l2_init_then_close_does_not_hang(self, monkeypatch): # A second close is a clean no-op (does not re-block on the epoch cv). w.close() - def test_l2_submit_returns_completed_handle(self, monkeypatch): + def test_l2_submit_returns_live_handle_and_wait_finalizes_it(self, monkeypatch): from simpler.task_interface import ChipCallable # noqa: PLC0415 w = self._make_l2(monkeypatch) @@ -771,8 +800,18 @@ def test_l2_submit_returns_completed_handle(self, monkeypatch): w.init() run_handle = w.submit(callable_handle) assert isinstance(run_handle, RunHandle) - assert run_handle.done + assert not run_handle.done assert run_handle.wait() is None + assert run_handle.done + chip_worker = w._chip_worker + assert chip_worker is not None + assert [event[0] for event in chip_worker.events] == [ + "prepare", + "launch", + "poll", + "wait", + "finalize", + ] w.close() def test_l2_close_without_init_is_noop(self, monkeypatch):