From eba8aacac2c5a154c776a177fcb82f5277e56a2c Mon Sep 17 00:00:00 2001 From: Chao Wang <26245345+ChaoWao@users.noreply.github.com> Date: Mon, 3 Aug 2026 06:13:18 -0700 Subject: [PATCH] Fix: fence register/unregister and make shutdown sticky close() publishes CLOSED and then drains _active_ops before it touches the tree, so a registry mutation made outside an operation lease is invisible to that drain. register() published the handle before taking its lease on the L3+ post-init ChipCallable path and on the L2 pre-warm path, and unregister() held no lease at all, so a close() could pass admission mid-transaction and tear the tree down underneath it. - register(): the post-READY L3+ and L2 paths take the lease before publication, so publication and the child-facing call are one transaction. Pre-start (NEW) and INITIALIZING registrations keep the epoch linearization they had, since the lease admits only READY. - unregister(): the live-target paths (L3+ child broadcast, L2 device slot release) take a symmetric lease. Whether the target is reachable is decided once by the caller under that lease instead of being re-read from the lifecycle across the broadcast. A registry-only decrement keeps the no-lease path so it still works before init(). - The rollback on a failed broadcast stays. The lease covers a racing close(); it says nothing about a broadcast that fails on its own. - _post_start_register_l2 rolls back on a failed pre-warm, matching _post_start_register_chip: if _register_callable_at_slot raises, the caller never receives the handle, so the published cid leaks against MAX_REGISTERED_CALLABLE_IDS otherwise. Termination had the mirror-image problem on the wire. The mailbox state word has three writers, so the CONTROL_DONE a child publishes for an in-flight control command overwrites a concurrent SHUTDOWN store; the child then polls a mailbox whose request has been erased until the parent's reap deadline expires and reports it as a survivor. MAILBOX_OFF_SHUTDOWN is a sticky word written only by a terminating parent, 0 -> 1, never cleared, and both serve loops read it at the top of every iteration alongside the state word. It is reserved on every frame by shrinking MAILBOX_ARGS_CAPACITY by 8 bytes, so no existing offset moves and run_control_command is unchanged. tests/ut/py/test_worker/test_admission_fence.py covers all three: the lease is held at publication and at the unregister broadcast, close() drains a transaction in flight, and a shutdown that races an in-flight control command still ends the child. Four of the five fail against the previous code. The host-buffer registration test scripted the serve loop's poll sequence, so its fake now answers the shutdown address. --- docs/worker-manager.md | 22 +- python/simpler/worker.py | 166 +++++++--- src/common/hierarchical/worker_manager.cpp | 9 +- src/common/hierarchical/worker_manager.h | 16 +- .../ut/py/test_worker/test_admission_fence.py | 287 ++++++++++++++++++ .../test_host_buffer_registration.py | 7 +- 6 files changed, 454 insertions(+), 53 deletions(-) create mode 100644 tests/ut/py/test_worker/test_admission_fence.py diff --git a/docs/worker-manager.md b/docs/worker-manager.md index eb6126750c..2a90793bc0 100644 --- a/docs/worker-manager.md +++ b/docs/worker-manager.md @@ -316,6 +316,7 @@ MAILBOX_OFF_PIPELINE_LEASE: PipelineSlotLease MAILBOX_OFF_TASK_CALLABLE_HASH: uint8[32] callable digest MAILBOX_OFF_TASK_ARGS_BLOB: bytes [int32 T][int32 S] [Tensor x T][uint64_t x S] +MAILBOX_OFF_SHUTDOWN: int32 sticky shutdown request task-frame trailer: protocol, run id, lease slot, lease generation, dispatch id MAILBOX_OFF_ACCEPTED: int32 sticky native-launch marker @@ -324,17 +325,28 @@ frame tail: fixed-size NUL-terminated error message The C++ `MAILBOX_FRAME_SIZE` and `MAILBOX_SIZE` constants are exported through the nanobind module. Python derives frame slicing and offsets from those -bindings where possible. `MAILBOX_ARGS_CAPACITY` ends before the protocol -trailer, acceptance word, and error-message tail. +bindings where possible. `MAILBOX_ARGS_CAPACITY` ends before the shutdown word, +protocol trailer, acceptance word, and error-message tail. + +`MAILBOX_OFF_SHUTDOWN` carries the termination request on the control base +frame. The state word has three writers — the parent's `CONTROL_REQUEST`, the +child's `CONTROL_DONE`, and the endpoint's return to `IDLE` — so a `SHUTDOWN` +state store is erasable. Only a terminating parent writes the shutdown word, +0 -> 1, and nothing clears it, so a child's serve loop still observes the +request after an in-flight control command completes over it. Both sides write +it before the `SHUTDOWN` state store and read it at the top of every serve-loop +iteration alongside the state word. ### 3.5 Stop and child shutdown `Worker::close()` first asks the Scheduler to stop admitting dispatch, then calls `WorkerManager::stop_workers()` while Scheduler callbacks and worker pool entries are still valid. A progressable `WorkerThread` repeatedly publishes -`SHUTDOWN` on the control base until its frames terminalize; repetition prevents -a concurrently finishing control handler's `CONTROL_DONE` from losing the stop -request. The child finalizes any active native run and marks every +`SHUTDOWN` on the control base until its frames terminalize; the sticky +shutdown word (section 3.4) is what makes the request itself unloseable, while +the repetition keeps the *state* word from resting on a `CONTROL_DONE` a +concurrently finishing control handler published. The child finalizes any +active native run and marks every active or staged task frame failed before leaving; the parent drains those terminal events, brings its in-flight count to zero, and joins the one progress thread. The Scheduler is joined only after worker threads can no longer invoke diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 9487d92882..34f3107474 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -207,7 +207,15 @@ def my_l4_orch(orch, args, config): _OFF_FRAME_GENERATION = _OFF_ACCEPTED - 16 _OFF_FRAME_DISPATCH_ID = _OFF_ACCEPTED - 8 _TASK_PROTOCOL_VERSION = 2 -_MAILBOX_ARGS_CAPACITY = _OFF_FRAME_PROTOCOL - _OFF_TASK_ARGS_BLOB +# 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++ +# return-to-IDLE), any of which overwrites a _SHUTDOWN store; only a +# terminating parent writes this word, 0 -> 1, and nothing clears it. The word +# is reserved on every frame so a task-args blob can never reach it. +_OFF_SHUTDOWN = _OFF_FRAME_PROTOCOL - 8 +_SHUTDOWN_REQUESTED = 1 +_MAILBOX_ARGS_CAPACITY = _OFF_SHUTDOWN - _OFF_TASK_ARGS_BLOB _OFF_CONTROL_CALLABLE_HASH = _OFF_ARGS + 32 # MAILBOX_OFF_ERROR_MSG / MAILBOX_ERROR_MSG_SIZE come from the C++ # nanobind module so the two sides cannot drift. @@ -1655,6 +1663,18 @@ def _buffer_field_addr(buf, offset: int) -> int: return ctypes.addressof(ctypes.c_char.from_buffer(buf)) + offset +def _request_child_shutdown(buf) -> None: + """Ask the child owning this mailbox to leave its serve loop. + + Sole writer of the termination request. The sticky ``_OFF_SHUTDOWN`` word + goes first so a child sampling it between the two stores already leaves by + the shutdown path; the ``_SHUTDOWN`` state word follows for a child parked + on the state word alone. + """ + _mailbox_store_i32(_buffer_field_addr(buf, _OFF_SHUTDOWN), _SHUTDOWN_REQUESTED) + _mailbox_store_i32(_buffer_field_addr(buf, _OFF_STATE), _SHUTDOWN) + + def _write_error(buf, code: int, msg: str = "") -> None: """Write an (error code, message) tuple into the mailbox error region. @@ -1742,11 +1762,21 @@ def _run_mailbox_loop( kill, a cancelled CI job) would otherwise leave this loop polling a mailbox nobody writes to, for the lifetime of the machine. The loop therefore samples its own parent and leaves by the SHUTDOWN path once it changes. + + Termination is read from the sticky ``_OFF_SHUTDOWN`` word as well as the + state word: the ``_CONTROL_DONE`` this loop publishes for an in-flight + control command overwrites a concurrent ``_SHUTDOWN`` store, and only the + sticky word survives that. """ parent_pid = os.getppid() liveness_countdown = _PARENT_LIVENESS_POLL_INTERVAL + shutdown_addr = _buffer_field_addr(buf, _OFF_SHUTDOWN) while True: state = _mailbox_load_i32(state_addr) + if state == _SHUTDOWN or _mailbox_load_i32(shutdown_addr) == _SHUTDOWN_REQUESTED: + if on_shutdown is not None: + on_shutdown() + break if state == _TASK_READY: code, msg = handle_task() _write_error(buf, code, msg) @@ -1756,10 +1786,6 @@ def _run_mailbox_loop( code, msg = handle_control(int(sub_cmd)) _write_error(buf, code, msg) _mailbox_store_i32(state_addr, _CONTROL_DONE) - elif state == _SHUTDOWN: - if on_shutdown is not None: - on_shutdown() - break else: liveness_countdown -= 1 if liveness_countdown <= 0: @@ -2534,10 +2560,14 @@ def stage_frame(index: int, initial_state: int) -> _StagedFrame | None: parent_pid = os.getppid() liveness_countdown = _PARENT_LIVENESS_POLL_INTERVAL shutdown_message = f"chip_process dev={device_id}: task loop shut down" + # The sticky shutdown word outlives the _CONTROL_DONE this loop + # publishes for an in-flight control command, which overwrites a + # concurrent _SHUTDOWN store on the state word. + shutdown_addr = _buffer_field_addr(buf, _OFF_SHUTDOWN) try: while True: control_state = _mailbox_load_i32(state_addr) - if control_state == _SHUTDOWN: + if control_state == _SHUTDOWN or _mailbox_load_i32(shutdown_addr) == _SHUTDOWN_REQUESTED: break if control_state == _CONTROL_REQUEST: sub_cmd = struct.unpack_from("Q", buf, _OFF_CALLABLE)[0] @@ -5150,41 +5180,69 @@ def register(self, target, *, workers: list[int] | None = None) -> CallableHandl handle = self._register_into_snapshot_or_wait(reg) if handle is not None: return handle - if not isinstance(target, ChipCallable): - with self._operation_lease("register"): + # Post-start publication touches the live tree; hold a lease across + # the whole transaction, publication included, so close() drains it + # before teardown (re-checks READY, closing the gate-then-teardown + # race). + with self._operation_lease("register"): + if not isinstance(target, ChipCallable): return self._post_start_register_python(reg) - else: - # L2 has no pre-start snapshot, but still linearizes against the - # epoch: reject a terminal (CLOSED/FAILED) worker and wait out an - # in-progress init so the callable is installed and its device slot - # prepared after READY — never left registered-but-not-prepared, and - # never accepted onto a closed worker as an inert handle. - with self._hierarchical_start_cv: - self._wait_out_init_locked("register") + return self._post_start_register_chip(reg, target) + # L2 has no pre-start snapshot, but still linearizes against the + # epoch: reject a terminal (CLOSED/FAILED) worker and wait out an + # in-progress init so the callable is installed and its device slot + # prepared after READY — never left registered-but-not-prepared, and + # never accepted onto a closed worker as an inert handle. + with self._hierarchical_start_cv: + self._wait_out_init_locked("register") + if self.level == 2 and self._initialized: + with self._operation_lease("register"): + return self._post_start_register_l2(reg, target) + with self._registry_lock: + handle, _is_new = self._install_registration_locked(reg) + return handle + + def _post_start_register_chip(self, reg: _CallableRegistration, target: ChipCallable) -> CallableHandle: + """Publish a post-READY L3+ ChipCallable and broadcast it to the chip / + next-level children via C++ after Host-side slot allocation. + + Caller holds an ``_operation_lease``, so publication and broadcast are + one transaction: a close() either drains the whole thing or is refused + admission before anything is published. The slot is target-private; task + dispatches carry only ``handle.digest``. + """ with self._registry_lock: handle, is_new = self._install_registration_locked(reg) + try: + self._post_init_register(target, handle.digest, is_new=is_new) + except Exception: + with self._registry_lock: + self._rollback_handle_locked(handle) + raise + return handle - # L3+ post-init ChipCallable: broadcast to chip / next-level children - # via C++ after L3 Host-side slot allocation is complete. The slot is - # target-private; task dispatches carry only handle.digest. - if self.level >= 3 and self._initialized and isinstance(target, ChipCallable): - try: - with self._operation_lease("register"): - self._post_init_register(target, handle.digest, is_new=is_new) - except Exception: - with self._registry_lock: - self._rollback_handle_locked(handle) - raise + def _post_start_register_l2(self, reg: _CallableRegistration, target: ChipCallable) -> CallableHandle: + """Publish a post-READY L2 registration and pre-warm its device slot, so + the very first ``run(handle, …)`` is a clean cache hit. - # L2 post-init: pre-warm immediately so the very first run(handle, …) - # is a clean cache hit. - if self.level == 2 and self._initialized and isinstance(target, ChipCallable) and is_new: - assert self._chip_worker is not None + Caller holds an ``_operation_lease``. L2 has no child subtree to + broadcast to, so the lease covers only publication and the local + pre-warm. + """ + with self._registry_lock: + handle, is_new = self._install_registration_locked(reg) + if not is_new: + return handle + assert self._chip_worker is not None + with self._registry_lock: + slot_id = self._identity_registry[handle.digest].slot_id + try: + self._chip_worker._register_callable_at_slot(slot_id, target) + except Exception: with self._registry_lock: - slot_id = self._identity_registry[handle.digest].slot_id - with self._operation_lease("register"): - self._chip_worker._register_callable_at_slot(slot_id, target) + self._rollback_handle_locked(handle) + raise return handle def _python_worker_types(self) -> list[WorkerType]: @@ -5694,6 +5752,26 @@ def unregister(self, handle_or_slot) -> None: return if self._pre_start_unregister_if_needed(handle_or_slot): return + # Symmetric with register: a path that drives the live tree — the L3+ + # child broadcast, and the L2 device-slot release — makes its registry + # mutation and its child-facing call one transaction under a lease, so + # a concurrent close() drains it before teardown instead of racing it. + # A registry-only decrement takes no lease: it must still work on a + # worker that was never initialized, which the lease rejects outright. + if self._initialized and self.level >= 2: + with self._operation_lease("unregister"): + self._unregister_handle(handle_or_slot, live=True) + return + self._unregister_handle(handle_or_slot, live=False) + + def _unregister_handle(self, handle_or_slot, *, live: bool) -> None: + """Pop a handle from the registry and propagate cleanup to its target. + + ``live`` says whether the target is reachable — an L3+ child subtree or + an L2 device slot. It is decided once by the caller, under the lease + that pins the worker in READY for the whole transaction, rather than + re-read from the lifecycle across the broadcast. + """ target = None digest = b"" cid = -1 @@ -5706,9 +5784,7 @@ def unregister(self, handle_or_slot) -> None: raise KeyError("UNREGISTER_TOMBSTONE_ACTIVE: callable handle already pending unregister") self._live_handles.pop(handle_id, None) state.ref_count -= 1 - should_broadcast_decrement = ( - self.level >= 3 and self._initialized and getattr(self, "_hierarchical_started", False) - ) + should_broadcast_decrement = live and self.level >= 3 if state.ref_count > 0 and not should_broadcast_decrement: return target = self._callable_registry[cid] @@ -5717,7 +5793,7 @@ def unregister(self, handle_or_slot) -> None: self._pending_unregister_cids.add(cid) if state.ref_count > 0: remove_target = False - elif self.level == 2 and self._initialized: + elif live and self.level == 2: assert self._chip_worker is not None self._chip_worker._unregister_slot(cid) self._callable_registry.pop(cid, None) @@ -6572,7 +6648,7 @@ def _abort_hierarchical(self, deadline: float | None = None) -> None: # noqa: P buf = shms_list[idx].buf if idx < len(shms_list) else None if buf is None: continue - _mailbox_store_i32(_buffer_field_addr(buf, _OFF_STATE), _SHUTDOWN) + _request_child_shutdown(buf) graceful.append(pid) # Phase 1b: mid-init next-level children get a cooperative cancel so they @@ -8839,15 +8915,19 @@ def close(self) -> None: # noqa: PLR0912, PLR0915 -- lifecycle linearization: r @staticmethod def _broadcast_child_shutdown(shms: list[SharedMemory]) -> None: - """Store _SHUTDOWN into every child mailbox in one group (next-level - children trigger ``inner_worker.close()``; chip/sub children exit their - serve loop). The first store error is raised after all are attempted.""" + """Store the shutdown request into every child mailbox in one group + (next-level children trigger ``inner_worker.close()``; chip/sub children + exit their serve loop). The first store error is raised after all are + attempted. + + The request is a sticky word plus the state word (see + ``_request_child_shutdown``).""" errors: list[BaseException] = [] for shm in shms: try: buf = shm.buf if buf is not None: - _mailbox_store_i32(_buffer_field_addr(buf, _OFF_STATE), _SHUTDOWN) + _request_child_shutdown(buf) except BaseException as exc: # noqa: BLE001 errors.append(exc) if errors: diff --git a/src/common/hierarchical/worker_manager.cpp b/src/common/hierarchical/worker_manager.cpp index 8fa9c193da..92a082a1b8 100644 --- a/src/common/hierarchical/worker_manager.cpp +++ b/src/common/hierarchical/worker_manager.cpp @@ -256,7 +256,14 @@ void LocalMailboxEndpoint::clear_task_accepted(char *frame) { __atomic_store(ptr, &v, __ATOMIC_RELEASE); } -void LocalMailboxEndpoint::shutdown_child() { write_mailbox_state(MailboxState::SHUTDOWN); } +void LocalMailboxEndpoint::shutdown_child() { + // Sticky word first: a child that samples it between the two stores leaves + // by the shutdown path even though the state word still reads IDLE. + int32_t *ptr = reinterpret_cast(mbox() + MAILBOX_OFF_SHUTDOWN); + int32_t requested = MAILBOX_SHUTDOWN_REQUESTED; + __atomic_store(ptr, &requested, __ATOMIC_RELEASE); + write_mailbox_state(MailboxState::SHUTDOWN); +} char *LocalMailboxEndpoint::task_frame(size_t index) const { return mbox() + diff --git a/src/common/hierarchical/worker_manager.h b/src/common/hierarchical/worker_manager.h index 8d051b20c4..4e94228101 100644 --- a/src/common/hierarchical/worker_manager.h +++ b/src/common/hierarchical/worker_manager.h @@ -141,6 +141,14 @@ static constexpr ptrdiff_t MAILBOX_OFF_FRAME_RUN_ID = MAILBOX_OFF_ACCEPTED - 32; static constexpr ptrdiff_t MAILBOX_OFF_FRAME_SLOT_ID = MAILBOX_OFF_ACCEPTED - 24; static constexpr ptrdiff_t MAILBOX_OFF_FRAME_GENERATION = MAILBOX_OFF_ACCEPTED - 16; static constexpr ptrdiff_t MAILBOX_OFF_FRAME_DISPATCH_ID = MAILBOX_OFF_ACCEPTED - 8; +// Termination is a sticky one-way word on the control frame, not a MailboxState: +// MAILBOX_OFF_STATE has three writers (the parent's CONTROL_REQUEST, the child's +// CONTROL_DONE, and this endpoint's return-to-IDLE), any of which overwrites a +// SHUTDOWN store. Only a terminating parent writes this word, 0 -> 1, and +// nothing ever clears it, so a child that observes it exits its serve loop no +// 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; 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); @@ -148,11 +156,13 @@ static constexpr size_t CTRL_SHM_NAME_BYTES = 32; static constexpr ptrdiff_t MAILBOX_OFF_CONTROL_CALLABLE_HASH = MAILBOX_OFF_ARGS + static_cast(CTRL_SHM_NAME_BYTES); static_assert( - MAILBOX_OFF_TASK_ARGS_BLOB < MAILBOX_OFF_FRAME_PROTOCOL, - "mailbox task-args region must precede the frame protocol trailer" + MAILBOX_OFF_TASK_ARGS_BLOB < MAILBOX_OFF_SHUTDOWN, + "mailbox task-args region must precede the shutdown word and the frame protocol trailer" ); +// The shutdown word is reserved on every frame, not just the control frame, so +// the args region a task frame accepts can never reach it. static constexpr size_t MAILBOX_ARGS_CAPACITY = - static_cast(MAILBOX_OFF_FRAME_PROTOCOL) - static_cast(MAILBOX_OFF_TASK_ARGS_BLOB); + static_cast(MAILBOX_OFF_SHUTDOWN) - static_cast(MAILBOX_OFF_TASK_ARGS_BLOB); static_assert( MAILBOX_ARGS_CAPACITY >= TASK_ARGS_BLOB_HEADER_SIZE + static_cast(CHIP_MAX_TENSOR_ARGS) * sizeof(Tensor) + static_cast(CHIP_MAX_SCALAR_ARGS) * sizeof(uint64_t), diff --git a/tests/ut/py/test_worker/test_admission_fence.py b/tests/ut/py/test_worker/test_admission_fence.py new file mode 100644 index 0000000000..8fe8ad562b --- /dev/null +++ b/tests/ut/py/test_worker/test_admission_fence.py @@ -0,0 +1,287 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Admission fencing for register / unregister, and the one-way shutdown word. + +``Worker.close()`` publishes CLOSED and then drains ``_active_ops`` before it +touches the tree, so every API that mutates the registry or drives a child must +hold an ``_operation_lease`` for its WHOLE transaction — publication included. +A registry mutation made outside the lease is invisible to that drain, which is +how a registration lands on a worker whose teardown is already running. + +Termination has the mirror-image problem on the wire: the mailbox state word has +three writers, and the ``CONTROL_DONE`` a child publishes for an in-flight +control command overwrites a concurrent ``SHUTDOWN`` store, leaving the child +polling a mailbox whose request has been erased. The sticky ``_OFF_SHUTDOWN`` +word is written only by a terminating parent and never cleared, so it survives +that overwrite. + +Every test here is device-free (an L3 worker with one SUB child and no chips) +and carries its own hard timeout, so a regression that reintroduces an unbounded +wait fails promptly instead of hanging CI. +""" + +from __future__ import annotations + +import contextlib +import os +import signal +import threading +import time +from multiprocessing.shared_memory import SharedMemory + +import pytest +import simpler.worker as worker_mod +from _task_interface import ChipCallable # pyright: ignore[reportMissingImports] +from simpler.worker import Worker + +# Comfortably above every wait these tests actually take, well under any hang. +_TEST_WALL_BUDGET_S = 30.0 +# How long close() must be observed NOT to return while a fenced operation is +# parked mid-transaction. Only has to outlast the scheduling noise of handing +# the GIL to the close() thread. +_FENCE_OBSERVATION_S = 1.0 + + +def _chip_callable(func_name: str = "admission_fence") -> ChipCallable: + return ChipCallable.build(signature=[], func_name=func_name, binary=b"\x00", children=[]) + + +@pytest.fixture +def ready_worker(): + """A READY L3 worker with one SUB child and no chips — no NPU required.""" + w = Worker(level=3, num_sub_workers=1, startup_timeout_s=_TEST_WALL_BUDGET_S) + w.init() + try: + yield w + finally: + with contextlib.suppress(BaseException): # a test may have closed it already + w.close() + + +class _ParkedTransaction: + """Park a worker method mid-transaction so close() can be raced against it. + + ``entered`` fires once the patched method has run the real implementation; + the method then blocks until ``release`` is set (bounded), which is the + window a close() must not be able to slip through. + """ + + def __init__(self, obj, name: str): + self.entered = threading.Event() + self.release = threading.Event() + self.lease_held: bool | None = None + self._obj = obj + self._name = name + self._original = getattr(obj, name) + + def __enter__(self): + def patched(*args, **kwargs): + result = self._original(*args, **kwargs) + self.lease_held = threading.get_ident() in self._obj._lease_depth + self.entered.set() + self.release.wait(timeout=_TEST_WALL_BUDGET_S) + return result + + setattr(self._obj, self._name, patched) + return self + + def __exit__(self, *exc_info): + self.release.set() + try: + delattr(self._obj, self._name) + except AttributeError: + pass + return False + + +def _close_while_parked(worker: Worker, parked: _ParkedTransaction) -> float: + """close() the worker from THIS thread while ``parked`` holds a transaction + open, and return how long it took. + + close() runs on the caller's thread because a worker with a live native tree + may only be closed on the thread that init()'d it. A timer releases the + parked transaction, so a close() that is properly fenced behind the lease + takes at least ``_FENCE_OBSERVATION_S`` and an unfenced one returns at once. + """ + releaser = threading.Timer(_FENCE_OBSERVATION_S, parked.release.set) + releaser.start() + started = time.monotonic() + try: + worker.close() + finally: + releaser.cancel() + parked.release.set() + return time.monotonic() - started + + +def _run_in_thread(fn): + box: dict[str, BaseException] = {} + + def target(): + try: + fn() + except BaseException as exc: # noqa: BLE001 -- reported to the test thread + box["error"] = exc + + thread = threading.Thread(target=target, daemon=True) + thread.start() + return thread, box + + +class TestRegisterAdmissionFence: + """register() must publish inside the lease, not before it.""" + + def test_publication_holds_the_lease(self, ready_worker): + with _ParkedTransaction(ready_worker, "_install_registration_locked") as parked: + parked.release.set() + handle = ready_worker.register(_chip_callable()) + assert parked.lease_held is True + ready_worker.unregister(handle) + + def test_close_drains_a_register_in_flight(self, ready_worker): + # Parks at the broadcast, not at publication: publication runs under + # ``_registry_lock``, which close()'s registry detach also takes, so a + # park there would delay close() even with no lease at all. + with _ParkedTransaction(ready_worker, "_post_init_register") as parked: + register_thread, register_box = _run_in_thread(lambda: ready_worker.register(_chip_callable())) + assert parked.entered.wait(timeout=_TEST_WALL_BUDGET_S), "register never reached its broadcast" + elapsed = _close_while_parked(ready_worker, parked) + register_thread.join(timeout=_TEST_WALL_BUDGET_S) + + assert elapsed >= _FENCE_OBSERVATION_S, "close() tore the worker down while a register was mid-broadcast" + assert not register_thread.is_alive() + assert "error" not in register_box, f"register() failed: {register_box.get('error')}" + # Nothing survives a terminal close: the registration either completed + # inside the drained lease and was detached with the rest of the + # registry, or never landed. + assert ready_worker._identity_registry == {} + assert ready_worker._live_handles == {} + + +class TestUnregisterAdmissionFence: + """unregister() is symmetric to register() — it held no lease at all.""" + + def test_broadcast_holds_the_lease(self, ready_worker): + handle = ready_worker.register(_chip_callable()) + with _ParkedTransaction(ready_worker, "_broadcast_unregister") as parked: + parked.release.set() + ready_worker.unregister(handle) + assert parked.lease_held is True + + def test_close_cannot_slip_into_the_unregister_broadcast(self, ready_worker): + handle = ready_worker.register(_chip_callable()) + with _ParkedTransaction(ready_worker, "_broadcast_unregister") as parked: + unregister_thread, unregister_box = _run_in_thread(lambda: ready_worker.unregister(handle)) + assert parked.entered.wait(timeout=_TEST_WALL_BUDGET_S), "unregister never reached its broadcast" + elapsed = _close_while_parked(ready_worker, parked) + unregister_thread.join(timeout=_TEST_WALL_BUDGET_S) + + assert elapsed >= _FENCE_OBSERVATION_S, "close() tore the worker down while an unregister was mid-transaction" + assert not unregister_thread.is_alive() + assert "error" not in unregister_box, f"unregister() failed: {unregister_box.get('error')}" + assert ready_worker._identity_registry == {} + assert ready_worker._live_handles == {} + + +# --------------------------------------------------------------------------- +# One-way shutdown word +# --------------------------------------------------------------------------- + +_GATE_CHILD_IN_CONTROL = 0 +_GATE_PARENT_RELEASED = 1 + + +def _serve_until_shutdown(mailbox: SharedMemory, gate: SharedMemory) -> None: + """Child half of the shutdown race: a serve loop whose control handler is + held open long enough for the parent to request shutdown underneath it.""" + mailbox_buf = mailbox.buf + gate_buf = gate.buf + assert mailbox_buf is not None + assert gate_buf is not None + state_addr = worker_mod._buffer_field_addr(mailbox_buf, worker_mod._OFF_STATE) + + def handle_task(): + return 0, "" + + def handle_control(_sub_cmd): + gate_buf[_GATE_CHILD_IN_CONTROL] = 1 + deadline = time.monotonic() + _TEST_WALL_BUDGET_S + while not gate_buf[_GATE_PARENT_RELEASED] and time.monotonic() < deadline: + pass + return 0, "" + + worker_mod._run_mailbox_loop( + mailbox_buf, + state_addr, + handle_task=handle_task, + handle_control=handle_control, + ) + + +def _wait_for_exit(pid: int, budget_s: float) -> int | None: + deadline = time.monotonic() + budget_s + while time.monotonic() < deadline: + reaped, status = os.waitpid(pid, os.WNOHANG) + if reaped == pid: + return status + time.sleep(0.01) + return None + + +class TestOneWayShutdown: + def test_shutdown_survives_an_in_flight_control_command(self): + """A shutdown requested while a control command is in flight must still + end the child, even though the CONTROL_DONE it publishes overwrites the + SHUTDOWN state word.""" + mailbox = SharedMemory(create=True, size=worker_mod.MAILBOX_FRAME_SIZE) + gate = SharedMemory(create=True, size=8) + mailbox_buf = mailbox.buf + gate_buf = gate.buf + assert mailbox_buf is not None + assert gate_buf is not None + pid = None + try: + pid = os.fork() + if pid == 0: # pragma: no cover -- child process + try: + _serve_until_shutdown(mailbox, gate) + finally: + os._exit(0) + + state_addr = worker_mod._buffer_field_addr(mailbox_buf, worker_mod._OFF_STATE) + worker_mod._mailbox_store_i32(state_addr, worker_mod._CONTROL_REQUEST) + + deadline = time.monotonic() + _TEST_WALL_BUDGET_S + while not gate_buf[_GATE_CHILD_IN_CONTROL] and time.monotonic() < deadline: + time.sleep(0.01) + assert gate_buf[_GATE_CHILD_IN_CONTROL], "child never entered its control handler" + + # The request lands while the control command is still in flight, so + # the child's CONTROL_DONE store comes after it. + worker_mod._request_child_shutdown(mailbox_buf) + gate_buf[_GATE_PARENT_RELEASED] = 1 + + status = _wait_for_exit(pid, _TEST_WALL_BUDGET_S) + assert status is not None, "child did not exit after SHUTDOWN raced an in-flight control command" + assert os.WIFEXITED(status) and os.WEXITSTATUS(status) == 0 + pid = None + assert worker_mod._mailbox_load_i32(state_addr) == worker_mod._CONTROL_DONE, ( + "expected the child's CONTROL_DONE to have overwritten the SHUTDOWN state word" + ) + finally: + if pid is not None: + with contextlib.suppress(ProcessLookupError, ChildProcessError): + os.kill(pid, signal.SIGKILL) + os.waitpid(pid, 0) + mailbox_buf.release() + gate_buf.release() + mailbox.close() + mailbox.unlink() + gate.close() + gate.unlink() diff --git a/tests/ut/py/test_worker/test_host_buffer_registration.py b/tests/ut/py/test_worker/test_host_buffer_registration.py index 4027243ba2..21a7a64c80 100644 --- a/tests/ut/py/test_worker/test_host_buffer_registration.py +++ b/tests/ut/py/test_worker/test_host_buffer_registration.py @@ -177,8 +177,13 @@ def test_l3_sub_worker_maps_rewrites_and_unmaps_host_buffer(monkeypatch): ) ) called = [] + # The loop polls two words per iteration; only the state word is + # scripted, so the sticky shutdown word always reads "not requested". + shutdown_addr = worker_mod._buffer_field_addr(mailbox_buf, worker_mod._OFF_SHUTDOWN) - def load_state(_state_addr): + def load_state(state_addr): + if state_addr == shutdown_addr: + return 0 state = next(states) if state == worker_mod._TASK_READY: start = worker_mod._OFF_TASK_CALLABLE_HASH