diff --git a/python/simpler/worker.py b/python/simpler/worker.py index d77c1e4d6a..74dd5f1348 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -63,6 +63,7 @@ def my_l4_orch(orch, args, config): import ctypes import importlib import json +import math import os import re import signal @@ -70,6 +71,7 @@ def my_l4_orch(orch, args, config): import struct import sys import threading +import time import uuid from dataclasses import dataclass, field from multiprocessing.shared_memory import SharedMemory @@ -192,11 +194,31 @@ def my_l4_orch(orch, args, config): _SHUTDOWN = 3 _CONTROL_REQUEST = 4 _CONTROL_DONE = 5 -# Child writes this after its expensive init (ChipWorker.init) completes. -# Parent's _start_hierarchical spin-waits for every chip child to reach -# INIT_DONE before allowing any dispatch — keeps cross-rank init skew out -# of the per-rank host-side stream sync budget (issue #897). -_INIT_DONE = 6 +# Startup readiness handshake. A child writes INIT_READY after its own init +# (ChipWorker.init / inner Worker.init) succeeds, or INIT_FAILED after it fails, +# leaving the cause in the mailbox error region. The parent's readiness barrier +# (_await_children_ready) blocks on every child reaching INIT_READY before any +# dispatch, which also keeps cross-rank init skew out of the per-rank host-side +# stream sync budget (issue #897); INIT_FAILED, a dead child, or a blown +# deadline aborts startup with a bounded error instead of an unbounded spin. +_INIT_READY = 6 +_INIT_FAILED = 7 + +# Startup readiness bound. A child that neither reports INIT_READY/INIT_FAILED +# nor exits within this window is treated as hung and startup is aborted. +# Generous by default so a legitimately slow device/runtime init (large +# PTO2_RING_HEAP, cold arena build) is never falsely reaped; override per Worker +# via the `startup_timeout_s` config kwarg. The point is to bound *hangs*, not +# to police slow-but-progressing init. +_STARTUP_TIMEOUT_S = 300.0 +# Parent poll granularity while waiting for children to become ready. Cheap +# shared-memory reads dominate; the sleep only caps waitpid/deadline syscall +# frequency and is far below any real init-skew alignment concern. +_STARTUP_POLL_INTERVAL_S = 0.001 +# On startup rollback, a next-level child that reached its serve loop is asked +# to close gracefully (so it unlinks the nested mailbox shms only it knows the +# names of) before being SIGKILLed. This bounds that graceful wait. +_ROLLBACK_GRACEFUL_TIMEOUT_S = 10.0 # Control sub-commands (written at _OFF_CALLABLE as uint64) _CTRL_MALLOC = 0 @@ -1547,20 +1569,20 @@ def _chip_process_loop( # noqa: PLR0913 -- fork-child entry: all context (bins, cw.init(device_id, bins, log_level=log_level, log_info_v=log_info_v, prewarm_config=prewarm_config) except Exception as e: _tb.print_exc() - # Write the message so any parent reader that *does* inspect this - # path sees the real cause. State handshake for this init-time - # failure is broken — see KNOWN_ISSUES.md — and that is not part - # of the L4 scope. + # Publish the cause into the mailbox and flag INIT_FAILED so the + # parent's readiness barrier returns a bounded error instead of + # spinning forever on a child that will never reach INIT_READY. _write_error(buf, 1, _format_exc(f"chip_process dev={device_id} init", e)) + _mailbox_store_i32(_buffer_field_addr(buf, _OFF_STATE), _INIT_FAILED) return mailbox_addr = ctypes.addressof(ctypes.c_char.from_buffer(buf)) state_addr = mailbox_addr + _OFF_STATE - # Signal init complete. Parent's _start_hierarchical spin-waits for - # every chip child to reach _INIT_DONE before dispatching the first - # task, so the per-rank host-side stream sync budget only covers - # actual op execution rather than absorbing peer-rank init skew. - _mailbox_store_i32(state_addr, _INIT_DONE) + # Signal init success. The parent's readiness barrier waits for every chip + # child to reach _INIT_READY before dispatching the first task, so the + # per-rank host-side stream sync budget only covers actual op execution + # rather than absorbing peer-rank init skew. + _mailbox_store_i32(state_addr, _INIT_READY) sys.stderr.write(f"[chip_process pid={os.getpid()} dev={device_id}] ready\n") sys.stderr.flush() @@ -1714,6 +1736,39 @@ def _child_worker_loop( break +def _forked_child_main(buf: memoryview, label: str, setup, serve) -> None: + """Run a forked child to completion, always terminating via ``os._exit``. + + ``setup()`` runs the child's fallible init and returns an opaque context; + any failure there publishes INIT_FAILED with the cause and exits. On + success the child publishes INIT_READY (the parent's readiness barrier + unblocks) and ``serve(ctx)`` runs the mailbox loop. + + Load-bearing invariant: a forked child must NEVER let an exception unwind + back into the forked copy of the parent's ``_start_hierarchical`` frames. + Those frames carry the parent's inherited child-PID lists, so an unwind + into the startup rollback path would SIGKILL this child's *siblings* (real + processes at those PIDs). Catch everything and exit instead. + """ + import traceback as _tb # noqa: PLC0415 + + state_addr = _buffer_field_addr(buf, _OFF_STATE) + try: + ctx = setup() + except BaseException as e: # noqa: BLE001 + _tb.print_exc() + _write_error(buf, 1, _format_exc(f"{label} init", e)) + _mailbox_store_i32(state_addr, _INIT_FAILED) + os._exit(1) + _mailbox_store_i32(state_addr, _INIT_READY) + try: + serve(ctx) + except BaseException: # noqa: BLE001 + _tb.print_exc() + os._exit(1) + os._exit(0) + + # --------------------------------------------------------------------------- # Worker factory # --------------------------------------------------------------------------- @@ -1754,6 +1809,22 @@ def __init__( self._pending_unregister_cids: set[int] = set() self._pending_remote_unregister_hashids: set[bytes] = set() self._py_control_timeout_s = float(config.get("py_control_timeout_s", _PY_CONTROL_TIMEOUT_S)) + # Upper bound on how long the readiness barrier waits for a forked child + # to report INIT_READY/INIT_FAILED before treating it as hung. Must be + # finite, else the deadline can never trip and the "bounded startup" + # guarantee is void (NaN compares false against every deadline). + self._startup_timeout_s = float(config.get("startup_timeout_s", _STARTUP_TIMEOUT_S)) + if not (self._startup_timeout_s > 0 and math.isfinite(self._startup_timeout_s)): + raise ValueError("Worker startup_timeout_s must be a positive finite number of seconds") + # Per-startup bookkeeping consumed by the rollback path: PIDs the barrier + # already reaped (must not be re-SIGKILLed — the PID may be reused) and + # PIDs that reached their serve loop (READY → asked to close gracefully + # so they unlink their own nested shms). Reset at each _start_hierarchical. + self._startup_reaped_pids: set[int] = set() + self._startup_ready_pids: set[int] = set() + # Disposition of the last rollback (graceful vs. killed PIDs); diagnostics + # and tests read it to confirm READY children were closed, not killed. + self._last_rollback: dict[str, list[int]] | None = None self._hierarchical_start_state = "not_started" self._hierarchical_start_mu = threading.Lock() self._hierarchical_start_cv = threading.Condition(self._hierarchical_start_mu) @@ -3387,6 +3458,11 @@ def _start_hierarchical(self) -> None: # noqa: PLR0912 -- three parallel fork l device_ids = self._config.get("device_ids", []) n_sub = self._config.get("num_sub_workers", 0) + # Until the C++ scheduler is running (dw.init), every forked child is + # ours alone to reap: a readiness-barrier failure rolls the whole epoch + # back. After dw.init the scheduler owns the mailboxes, so we leave the + # existing (state="failed") handling and let close() tear down. + scheduler_started = False try: # Fork children from an immutable snapshot. The state transition # and snapshot are one gate, so dynamic register/unregister callers @@ -3406,22 +3482,33 @@ def _start_hierarchical(self) -> None: # noqa: PLR0912 -- three parallel fork l ] self._hierarchical_start_cv.notify_all() + self._startup_reaped_pids = set() + self._startup_ready_pids = set() + # Fork SubWorker processes (MUST be before any C++ threads) for i in range(n_sub): pid = os.fork() if pid == 0: buf = self._sub_shms[i].buf assert buf is not None - registry, identity_table, identity_refs = _make_local_identity_tables( - identity_snapshot, - callable_kind=("PYTHON_SERIALIZED", "PYTHON_IMPORT"), - target_namespace="LOCAL_PYTHON", - ) - _sub_worker_loop(buf, registry, identity_table, identity_refs) - os._exit(0) + + def _setup(): + return _make_local_identity_tables( + identity_snapshot, + callable_kind=("PYTHON_SERIALIZED", "PYTHON_IMPORT"), + target_namespace="LOCAL_PYTHON", + ) + + _forked_child_main(buf, f"sub worker {i}", _setup, lambda t, b=buf: _sub_worker_loop(b, *t)) else: self._sub_pids.append(pid) + # SUB children have no fallible device/runtime init, but they join + # the same readiness contract so a child that dies before entering + # its loop aborts startup rather than surfacing later as a hung + # submit_sub. + self._await_children_ready(self._sub_shms, self._sub_pids, "sub") + # Fork ChipWorker processes (L3 with device_ids). Always use the # plain task-loop variant; the base communicator is established # lazily on first ``orch.allocate_domain`` via CTRL_COMM_INIT. @@ -3432,21 +3519,36 @@ def _start_hierarchical(self) -> None: # noqa: PLR0912 -- three parallel fork l if pid == 0: buf = self._chip_shms[idx].buf assert buf is not None - _chip_process_loop( - buf, - self._l3_bins, - dev_id, - *_make_local_identity_tables( - identity_snapshot, - callable_kind="CHIP_CALLABLE", - target_namespace="LOCAL_CHIP", - ), - log_level=chip_log_level, - log_info_v=chip_log_info_v, - platform=str(self._config["platform"]), - runtime=str(self._config["runtime"]), - prewarm_config=self._prewarm_config, - ) + # _chip_process_loop publishes INIT_READY/INIT_FAILED + # itself (around cw.init). This guard only ensures the + # child exits rather than unwinding into the parent's + # startup frames (see _forked_child_main). A throw before + # cw.init (e.g. identity-table build) leaves the mailbox + # IDLE, so publish INIT_FAILED for a bounded parent error. + try: + _chip_process_loop( + buf, + self._l3_bins, + dev_id, + *_make_local_identity_tables( + identity_snapshot, + callable_kind="CHIP_CALLABLE", + target_namespace="LOCAL_CHIP", + ), + log_level=chip_log_level, + log_info_v=chip_log_info_v, + platform=str(self._config["platform"]), + runtime=str(self._config["runtime"]), + prewarm_config=self._prewarm_config, + ) + except BaseException as e: # noqa: BLE001 + import traceback as _tb # noqa: PLC0415 + + _tb.print_exc() + if _mailbox_load_i32(_buffer_field_addr(buf, _OFF_STATE)) == _IDLE: + _write_error(buf, 1, _format_exc(f"chip worker {idx} dev={dev_id} init", e)) + _mailbox_store_i32(_buffer_field_addr(buf, _OFF_STATE), _INIT_FAILED) + os._exit(1) os._exit(0) else: self._chip_pids.append(pid) @@ -3459,15 +3561,9 @@ def _start_hierarchical(self) -> None: # noqa: PLR0912 -- three parallel fork l # point, and any cross-rank wait inside the op (HCCL notify, # etc.) charges the slow peer's remaining init time against # the fast peer's PLATFORM_STREAM_SYNC_TIMEOUT_MS budget — - # the cascade documented in issue #897. Reset each child to - # _IDLE once observed so the standard dispatch state machine - # resumes from the canonical "ready for work" state. - for shm in self._chip_shms: - assert shm.buf is not None - addr = _buffer_field_addr(shm.buf, _OFF_STATE) - while _mailbox_load_i32(addr) != _INIT_DONE: - pass - _mailbox_store_i32(addr, _IDLE) + # the cascade documented in issue #897. A chip that fails or + # dies during init raises here rather than spinning forever. + self._await_children_ready(self._chip_shms, self._chip_pids, "chip") # Fork next-level Worker children (L4+ with Worker children). # Each child process: init the inner Worker (which mmaps its own @@ -3480,20 +3576,37 @@ def _start_hierarchical(self) -> None: # noqa: PLR0912 -- three parallel fork l if pid == 0: buf = self._next_level_shms[idx].buf assert buf is not None - # Propagate the fork-constant prewarm sizing down so L4+ trees - # prewarm their L3 chips too (each inner Worker prewarms on its - # own first run). Closes the prior L4+ silent no-op. - inner_worker.init(prewarm_config=self._prewarm_config) - registry, identity_table, identity_refs = _make_local_identity_tables( - identity_snapshot, - callable_kind=("PYTHON_SERIALIZED", "PYTHON_IMPORT"), - target_namespace="LOCAL_PYTHON", + + def _setup(inner=inner_worker): + # Propagate the fork-constant prewarm sizing down so L4+ + # trees prewarm their L3 chips too (each inner Worker + # prewarms on its own first run). INIT_READY is published + # only after BOTH the inner init and the identity-table + # build succeed, so the parent never observes READY for a + # child that then dies in fallible post-init setup. + inner.init(prewarm_config=self._prewarm_config) + return _make_local_identity_tables( + identity_snapshot, + callable_kind=("PYTHON_SERIALIZED", "PYTHON_IMPORT"), + target_namespace="LOCAL_PYTHON", + ) + + _forked_child_main( + buf, + f"next_level worker {idx}", + _setup, + lambda tables, b=buf, inner=inner_worker: _child_worker_loop(b, *tables, inner), ) - _child_worker_loop(buf, registry, identity_table, identity_refs, inner_worker) - os._exit(0) else: self._next_level_pids.append(pid) + # A next-level child publishes INIT_READY only after its own init + # succeeds; a failure, exit, or hang aborts startup here. This is + # the recursive readiness edge: once eager init lands, the child's + # init itself blocks on its descendants, so INIT_READY propagates up + # only after the whole subtree is ready. + self._await_children_ready(self._next_level_shms, self._next_level_pids, "next_level") + # _Worker was constructed in _init_hierarchical (pre-fork) so # children inherit the HeapRing MAP_SHARED mmap. Register PROCESS-mode # workers via the unified mailbox. @@ -3517,6 +3630,7 @@ def _start_hierarchical(self) -> None: # noqa: PLR0912 -- three parallel fork l # Start Scheduler + WorkerThreads (C++ threads start here, after fork) dw.init() + scheduler_started = True self._orch = Orchestrator(dw.get_orchestrator(), self) @@ -3539,38 +3653,152 @@ def _start_hierarchical(self) -> None: # noqa: PLR0912 -- three parallel fork l with self._hierarchical_start_cv: self._hierarchical_start_state = "started" self._hierarchical_start_cv.notify_all() - except Exception: + except BaseException: + # Pre-scheduler failures (a readiness-barrier abort, or a + # KeyboardInterrupt/SystemExit during the wait) roll the whole epoch + # back: gracefully close READY children so they unlink their own + # nested shms, SIGKILL the rest, reap all, and free the mailboxes — + # so a failed startup leaks no process or shm. After dw.init the + # scheduler owns the mailboxes, so leave teardown to close(). + if not scheduler_started: + self._abort_hierarchical() with self._hierarchical_start_cv: self._hierarchical_start_state = "failed" self._hierarchical_start_cv.notify_all() raise + def _await_children_ready(self, shms, pids, kind: str) -> None: + """Block until every forked child reports INIT_READY, or abort. + + Polls each child's mailbox: INIT_READY resets the slot to _IDLE (so the + C++ dispatch state machine resumes from the canonical "ready for work" + state), records the pid as having reached its serve loop, and retires + it; INIT_FAILED surfaces the child's own error; ``waitpid(WNOHANG)`` + catches a child that died before signalling (recording the reaped pid so + rollback never re-SIGKILLs a possibly-reused PID); the per-Worker + ``startup_timeout_s`` deadline catches a child that hangs. A failure + raises ``RuntimeError`` — the caller rolls back the whole startup epoch. + """ + deadline = time.monotonic() + self._startup_timeout_s + pending = list(range(len(shms))) + while pending: + still_pending = [] + for i in pending: + buf = shms[i].buf + assert buf is not None + addr = _buffer_field_addr(buf, _OFF_STATE) + state = _mailbox_load_i32(addr) + if state == _INIT_READY: + _mailbox_store_i32(addr, _IDLE) + self._startup_ready_pids.add(pids[i]) + continue + if state == _INIT_FAILED: + raise RuntimeError(f"{kind} worker {i} (pid {pids[i]}) failed during init: {_read_error_msg(buf)}") + try: + wpid, status = os.waitpid(pids[i], os.WNOHANG) + except ChildProcessError: + self._startup_reaped_pids.add(pids[i]) + raise RuntimeError( + f"{kind} worker {i} (pid {pids[i]}) exited during init before signalling ready" + ) from None + if wpid != 0: + self._startup_reaped_pids.add(pids[i]) + raise RuntimeError( + f"{kind} worker {i} (pid {pids[i]}) exited during init " + f"before signalling ready (wait status {status})" + ) + still_pending.append(i) + pending = still_pending + if pending: + if time.monotonic() > deadline: + raise RuntimeError( + f"{kind} worker(s) {pending} did not become ready within " + f"{self._startup_timeout_s}s (startup deadline exceeded)" + ) + time.sleep(_STARTUP_POLL_INTERVAL_S) + # ------------------------------------------------------------------ # Hierarchical abort # ------------------------------------------------------------------ - def _abort_hierarchical(self) -> None: + def _abort_hierarchical(self) -> None: # noqa: PLR0912 -- graceful-close-then-kill rollback: classify next-level children, bounded-wait, SIGKILL the rest, reap, free shms """Tear down all forked children + shms after a bootstrap failure. - Best-effort: SIGKILL every child we spawned, reap them, then close - and unlink every mailbox. Called only from the init() failure path, - so `dw.init()` has not run and the C++ scheduler is not holding any - mailbox references. + Called from the init() failure path and from a pre-scheduler + ``_start_hierarchical`` abort (readiness-barrier failure, or an + interrupt during the wait) — in both cases `dw.init()` has not run, so + the C++ scheduler is not holding any mailbox references. + + A next-level child that already reached its serve loop owns nested + mailbox shms that only *it* can unlink (the parent never learns their + names), so it is asked to close gracefully (SHUTDOWN, bounded wait) + before any SIGKILL. Every other child — and any graceful child that + overran the deadline — is SIGKILLed. PIDs the barrier already reaped are + excluded from the kill so a reused PID is never signalled. + + Residual: a next-level child still inside ``inner.init()`` when a sibling + fails has not reached its serve loop, so it cannot service SHUTDOWN and + is SIGKILLed; any nested shms it had already created are reclaimed by the + multiprocessing resource_tracker at interpreter exit rather than here. + Closing that race needs the eager-init handshake (follow-up PR). """ + reaped = set(self._startup_reaped_pids) + graceful: list[int] = [] + killed: list[int] = [] + + # Ask serve-loop-reached next-level children to close (unlinks their + # own nested shms), then bounded-wait for them to exit. + for idx, pid in enumerate(self._next_level_pids): + if pid in reaped or pid not in self._startup_ready_pids: + continue + buf = self._next_level_shms[idx].buf if idx < len(self._next_level_shms) else None + if buf is None: + continue + _mailbox_store_i32(_buffer_field_addr(buf, _OFF_STATE), _SHUTDOWN) + graceful.append(pid) + if graceful: + gdeadline = time.monotonic() + _ROLLBACK_GRACEFUL_TIMEOUT_S + waiting = set(graceful) + while waiting and time.monotonic() <= gdeadline: + for pid in list(waiting): + try: + wpid, _status = os.waitpid(pid, os.WNOHANG) + except ChildProcessError: + waiting.discard(pid) + reaped.add(pid) + continue + if wpid != 0: + waiting.discard(pid) + reaped.add(pid) + if waiting: + time.sleep(_STARTUP_POLL_INTERVAL_S) + # Any graceful child still alive overran the deadline; fall through + # to SIGKILL below. + pids = list(self._chip_pids) + list(self._sub_pids) + list(self._next_level_pids) for pid in pids: + if pid in reaped: + continue try: os.kill(pid, signal.SIGKILL) + killed.append(pid) except ProcessLookupError: pass except OSError: pass for pid in pids: + if pid in reaped and pid not in killed: + continue try: os.waitpid(pid, 0) except ChildProcessError: pass + self._last_rollback = { + "graceful": [p for p in graceful if p not in killed], + "killed": list(killed), + } + for shm in self._sub_shms + self._chip_shms + self._next_level_shms: try: shm.close() diff --git a/src/common/hierarchical/worker_manager.h b/src/common/hierarchical/worker_manager.h index c6aba8cf8a..c4668a0dde 100644 --- a/src/common/hierarchical/worker_manager.h +++ b/src/common/hierarchical/worker_manager.h @@ -62,13 +62,16 @@ enum class MailboxState : int32_t { SHUTDOWN = 3, CONTROL_REQUEST = 4, CONTROL_DONE = 5, - // Child writes this after its expensive init (ChipWorker::init / inner - // Worker::init) completes. Parent's _start_hierarchical spin-waits for - // EVERY chip child to reach INIT_DONE before any dispatch (CTRL_PREPARE - // or TASK_READY) goes out. This aligns the host-side stream-sync windows - // across distributed ranks so cross-rank init skew never charges against - // the per-rank PLATFORM_STREAM_SYNC_TIMEOUT_MS budget (issue #897). - INIT_DONE = 6, + // Startup readiness handshake, driven entirely by the Python facade + // (_await_children_ready). A child writes INIT_READY after its expensive + // init (ChipWorker::init / inner Worker::init) succeeds, or INIT_FAILED + // after it fails. The parent barrier waits for EVERY child to reach + // INIT_READY before any dispatch (CTRL_PREPARE or TASK_READY) goes out, + // which also aligns the host-side stream-sync windows across distributed + // ranks so cross-rank init skew never charges against the per-rank + // PLATFORM_STREAM_SYNC_TIMEOUT_MS budget (issue #897). + INIT_READY = 6, + INIT_FAILED = 7, }; // Sized so the args region can hold any TaskArgs the runtime itself accepts diff --git a/tests/ut/py/test_worker/test_startup_readiness.py b/tests/ut/py/test_worker/test_startup_readiness.py new file mode 100644 index 0000000000..a9c51608f8 --- /dev/null +++ b/tests/ut/py/test_worker/test_startup_readiness.py @@ -0,0 +1,459 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""Recursive worker-startup readiness protocol. + +A hierarchical Worker's startup is a strong READY boundary: every child +process (sub, chip, and next-level) must either publish INIT_READY after its +own init succeeds, or publish INIT_FAILED with a bounded error. The parent +waits for each child with a deadline and a ``waitpid(WNOHANG)`` liveness +check, so a child that crashes, exits, or hangs during init surfaces as a +prompt ``RuntimeError`` instead of an unbounded parent spin (the #1003 / #980 +hang). On failure the parent rolls the whole startup epoch back: children that +reached their serve loop are closed gracefully so they unlink their own nested +shms, the rest are SIGKILLed, and every child is reaped. + +Most tests inject failures at the L4 -> L3 (next-level) edge, which needs no +NPU device: the child runs ``inner_worker.init()`` before entering its serve +loop. The chip (L2) edge shares the same parent-side barrier; its device-free +failure path is covered by ``TestChipStartupFailure`` with a faked +``ChipWorker`` on the ``a2a3sim`` platform (no silicon). + +Every failure test is wrapped in a hard SIGALRM timeout so a protocol +regression that reintroduces an unbounded spin fails the suite promptly +instead of hanging CI. +""" + +import os +import signal +import struct +import time +from contextlib import contextmanager +from multiprocessing.shared_memory import SharedMemory + +import pytest +from simpler.task_interface import CallConfig, TaskArgs +from simpler.worker import Worker + +# Hard wall budget for a single failure scenario — comfortably above the +# injected startup_timeout_s values below, well under any real hang. +_TEST_WALL_BUDGET_S = 30.0 + + +@contextmanager +def _hard_timeout(seconds: float, msg: str = "startup did not return within the hard test budget"): + """Abort the body with TimeoutError if it runs longer than ``seconds``. + + A backstop against a regression that reintroduces an unbounded startup + spin: instead of hanging CI, the test fails with TimeoutError. Uses SIGALRM + (pytest runs on the main thread), which interrupts the barrier's + ``time.sleep`` poll. + """ + + def _handler(_signum, _frame): + raise TimeoutError(msg) + + old = signal.signal(signal.SIGALRM, _handler) + signal.setitimer(signal.ITIMER_REAL, seconds) + try: + yield + finally: + signal.setitimer(signal.ITIMER_REAL, 0) + signal.signal(signal.SIGALRM, old) + + +def _make_shared_counter(): + shm = SharedMemory(create=True, size=4) + buf = shm.buf + assert buf is not None + struct.pack_into("i", buf, 0, 0) + return shm, buf + + +def _read_counter(buf) -> int: + return struct.unpack_from("i", buf, 0)[0] + + +def _increment_counter(buf) -> None: + v = struct.unpack_from("i", buf, 0)[0] + struct.pack_into("i", buf, 0, v + 1) + + +# Injected inner-worker init failures. Defined at module scope so the forked +# next-level child inherits them (copy-on-write) and calls the replacement in +# place of the real Worker.init. + + +def _init_raises(*_a, **_k): + raise RuntimeError("injected inner init failure") + + +def _init_slow_raises(*_a, **_k): + # Delay so a healthy sibling reliably reaches READY before this one fails, + # exercising the graceful-close rollback path deterministically. + time.sleep(0.5) + raise RuntimeError("injected slow inner init failure") + + +def _init_hard_exits(*_a, **_k): + os._exit(42) + + +def _init_hangs(*_a, **_k): + time.sleep(3600) + + +def _l3_child(sub_fn=None, num_sub_workers=1): + l3 = Worker(level=3, num_sub_workers=num_sub_workers) + l3.register(sub_fn if sub_fn is not None else (lambda args: None)) + return l3 + + +def _trivial_orch(orch, args, config): + return None + + +class TestNextLevelStartupFailure: + def test_inner_init_failure_raises_bounded_error(self): + """A next-level child whose init raises surfaces its error at startup.""" + l3 = _l3_child() + l3.init = _init_raises # noqa: SLF001 -- test injection inherited across fork + + w4 = Worker(level=4, num_sub_workers=0, startup_timeout_s=10.0) + w4.register(_trivial_orch) + w4.add_worker(l3) + w4.init() + start = time.monotonic() + try: + with _hard_timeout(_TEST_WALL_BUDGET_S): + with pytest.raises(RuntimeError, match="injected inner init failure"): + w4.run(_trivial_orch) + assert time.monotonic() - start < _TEST_WALL_BUDGET_S + finally: + w4.close() + + def test_inner_exit_before_ready_raises(self): + """A child that exits during init (before READY) is detected via waitpid.""" + l3 = _l3_child() + l3.init = _init_hard_exits # noqa: SLF001 + + w4 = Worker(level=4, num_sub_workers=0, startup_timeout_s=10.0) + w4.register(_trivial_orch) + w4.add_worker(l3) + w4.init() + try: + with _hard_timeout(_TEST_WALL_BUDGET_S): + with pytest.raises(RuntimeError, match="exited during init"): + w4.run(_trivial_orch) + finally: + w4.close() + + def test_startup_deadline_fires_on_hung_child(self): + """A child that hangs in init trips the startup deadline, not an infinite spin.""" + l3 = _l3_child() + l3.init = _init_hangs # noqa: SLF001 + + w4 = Worker(level=4, num_sub_workers=0, startup_timeout_s=1.5) + w4.register(_trivial_orch) + w4.add_worker(l3) + w4.init() + start = time.monotonic() + try: + with _hard_timeout(_TEST_WALL_BUDGET_S): + with pytest.raises(RuntimeError, match="deadline"): + w4.run(_trivial_orch) + elapsed = time.monotonic() - start + assert 1.5 <= elapsed < _TEST_WALL_BUDGET_S + finally: + w4.close() + + def test_failed_startup_reaps_children_no_leak(self, monkeypatch): + """After a startup failure the forked children are killed and reaped.""" + l3 = _l3_child() + l3.init = _init_hangs # noqa: SLF001 + + captured: dict[str, list[int]] = {} + orig_abort = Worker._abort_hierarchical + + def spy_abort(self): + captured["pids"] = list(self._chip_pids) + list(self._sub_pids) + list(self._next_level_pids) + return orig_abort(self) + + monkeypatch.setattr(Worker, "_abort_hierarchical", spy_abort) + + w4 = Worker(level=4, num_sub_workers=0, startup_timeout_s=1.0) + w4.register(_trivial_orch) + w4.add_worker(l3) + w4.init() + try: + with _hard_timeout(_TEST_WALL_BUDGET_S): + with pytest.raises(RuntimeError): + w4.run(_trivial_orch) + + assert "pids" in captured and captured["pids"], "rollback did not run" + for pid in captured["pids"]: + with pytest.raises(ChildProcessError): + os.waitpid(pid, os.WNOHANG) + + # Rollback clears the process/mailbox bookkeeping so a later close() + # is a clean no-op. + assert w4._next_level_pids == [] + assert w4._next_level_shms == [] + assert w4._worker is None + finally: + w4.close() + + def test_ready_sibling_closed_gracefully_on_sibling_failure(self): + """A child that reached READY is closed gracefully (not SIGKILLed) when a sibling fails. + + The healthy L3 owns a nested sub-worker mailbox shm only it can unlink; + graceful SHUTDOWN lets it clean up. The failing L3 is delayed so the + healthy one is reliably READY first. + """ + good = _l3_child(num_sub_workers=1) + bad = _l3_child(num_sub_workers=1) + bad.init = _init_slow_raises # noqa: SLF001 + + w4 = Worker(level=4, num_sub_workers=0, startup_timeout_s=20.0) + w4.register(_trivial_orch) + w4.add_worker(good) + w4.add_worker(bad) + w4.init() + try: + with _hard_timeout(_TEST_WALL_BUDGET_S): + with pytest.raises(RuntimeError, match="slow inner init failure"): + w4.run(_trivial_orch) + + report = w4._last_rollback + assert report is not None + # The healthy sibling reached its serve loop and was closed + # gracefully (SHUTDOWN + reaped), not SIGKILLed. + assert len(report["graceful"]) >= 1 + assert set(report["graceful"]).isdisjoint(report["killed"]) + assert w4._next_level_pids == [] + finally: + w4.close() + + def test_second_child_failure_reaps_first(self): + """When one of several next-level children fails, all are torn down.""" + good = _l3_child() + bad = _l3_child() + bad.init = _init_raises # noqa: SLF001 + + w4 = Worker(level=4, num_sub_workers=0, startup_timeout_s=10.0) + w4.register(_trivial_orch) + w4.add_worker(good) + w4.add_worker(bad) + w4.init() + try: + with _hard_timeout(_TEST_WALL_BUDGET_S): + with pytest.raises(RuntimeError, match="injected inner init failure"): + w4.run(_trivial_orch) + assert w4._next_level_pids == [] + finally: + w4.close() + + +class TestSubStartupFailure: + def test_sub_child_exit_before_ready_raises(self, monkeypatch): + """A sub child that dies before entering its loop aborts startup. + + Injects a failure into the child's identity-table build (the sub's only + fallible pre-loop step) so the sub exits before publishing READY; the + parent's sub readiness barrier must catch it rather than silently + succeeding and hanging a later submit_sub. + """ + import simpler.worker as worker_mod # noqa: PLC0415 + + def _boom(*_a, **_k): + os._exit(7) + + monkeypatch.setattr(worker_mod, "_make_local_identity_tables", _boom) + + w3 = Worker(level=3, num_sub_workers=1, startup_timeout_s=10.0) + w3.register(lambda args: None) + w3.init() + try: + with _hard_timeout(_TEST_WALL_BUDGET_S): + with pytest.raises(RuntimeError, match="sub worker .* exited during init"): + w3.run(_trivial_orch) + finally: + w3.close() + + +class TestStartupConfigValidation: + def test_nonpositive_timeout_rejected(self): + with pytest.raises(ValueError, match="startup_timeout_s"): + Worker(level=4, num_sub_workers=0, startup_timeout_s=0) + + def test_nonfinite_timeout_rejected(self): + with pytest.raises(ValueError, match="finite"): + Worker(level=4, num_sub_workers=0, startup_timeout_s=float("inf")) + with pytest.raises(ValueError, match="finite"): + Worker(level=4, num_sub_workers=0, startup_timeout_s=float("nan")) + + +class TestReadyBarrierHappyPath: + """The barrier passes a healthy tree through and dispatch still works. + + These verify the next-level readiness barrier does not break a healthy + startup and that tasks dispatch afterwards. They do NOT assert full L4->L3->L2 + eager readiness: L3 descendants are still started lazily on first inner + run(), so INIT_READY here means the L3's own init succeeded. + """ + + def test_l4_l3_tree_comes_up_and_runs(self): + counter_shm, counter_buf = _make_shared_counter() + w4 = None + try: + l3 = Worker(level=3, num_sub_workers=1) + l3_sub = l3.register(lambda args: _increment_counter(counter_buf)) + + def l3_orch(orch, args, config): + orch.submit_sub(l3_sub) + + w4 = Worker(level=4, num_sub_workers=0, startup_timeout_s=30.0) + l3_handle = w4.register(l3_orch) + w4.add_worker(l3) + w4.init() + + def l4_orch(orch, args, config): + orch.submit_next_level(l3_handle, TaskArgs(), CallConfig()) + + w4.run(l4_orch) + assert _read_counter(counter_buf) == 1 + finally: + if w4 is not None: + w4.close() + counter_shm.close() + counter_shm.unlink() + + def test_multiple_l3_children_all_ready(self): + """Two next-level children both pass the barrier and dispatch. + + Each child increments its OWN counter (the counter is a non-atomic RMW + that would race if the two children shared one). + """ + a_shm, a_buf = _make_shared_counter() + b_shm, b_buf = _make_shared_counter() + w4 = None + try: + l3a = Worker(level=3, num_sub_workers=1) + a_sub = l3a.register(lambda args: _increment_counter(a_buf)) + + def l3a_orch(orch, args, config): + orch.submit_sub(a_sub) + + l3b = Worker(level=3, num_sub_workers=1) + b_sub = l3b.register(lambda args: _increment_counter(b_buf)) + + def l3b_orch(orch, args, config): + orch.submit_sub(b_sub) + + w4 = Worker(level=4, num_sub_workers=0, startup_timeout_s=30.0) + ha = w4.register(l3a_orch) + hb = w4.register(l3b_orch) + w4.add_worker(l3a) + w4.add_worker(l3b) + w4.init() + + def l4_orch(orch, args, config): + orch.submit_next_level(ha, TaskArgs(), CallConfig()) + orch.submit_next_level(hb, TaskArgs(), CallConfig()) + + w4.run(l4_orch) + assert _read_counter(a_buf) == 1 + assert _read_counter(b_buf) == 1 + finally: + if w4 is not None: + w4.close() + a_shm.close() + a_shm.unlink() + b_shm.close() + b_shm.unlink() + + +class _FakeChipRaises: + """Stand-in for ChipWorker whose init raises — no NPU touched.""" + + def init(self, *_a, **_k): + raise RuntimeError("injected chip init failure") + + def finalize(self): # pragma: no cover - never reached (init raises) + pass + + +class _FakeChipHangs: + def init(self, *_a, **_k): + time.sleep(3600) + + def finalize(self): # pragma: no cover + pass + + +def _sim_binaries_available() -> bool: + try: + from simpler_setup.runtime_builder import RuntimeBuilder # noqa: PLC0415 + + RuntimeBuilder("a2a3sim").get_binaries("tensormap_and_ringbuffer") + return True + except Exception: # noqa: BLE001 + return False + + +@pytest.mark.skipif(not _sim_binaries_available(), reason="a2a3sim runtime binaries not built") +class TestChipStartupFailure: + """Chip (L2) startup failure — device-free via a faked ChipWorker on a2a3sim. + + Constructing an L3 with ``device_ids`` only reads the prebuilt runtime + binaries; the forked chip child instantiates ``worker.ChipWorker``, which + the test replaces so no silicon is required. The failure trips before + ``dw.init()``, so the sim runtime is never actually driven. Exercises the + same parent-side readiness barrier as the next-level edge (the #1003 spin at + the former ``while ... != INIT_DONE`` was on this chip path). + """ + + def _make_l3(self, timeout_s): + return Worker( + level=3, + device_ids=[0], + platform="a2a3sim", + runtime="tensormap_and_ringbuffer", + num_sub_workers=0, + startup_timeout_s=timeout_s, + ) + + def test_chip_init_failure_raises_bounded(self, monkeypatch): + import simpler.worker as worker_mod # noqa: PLC0415 + + monkeypatch.setattr(worker_mod, "ChipWorker", _FakeChipRaises) + l3 = self._make_l3(timeout_s=10.0) + l3.register(lambda args: None) + l3.init() + try: + with _hard_timeout(_TEST_WALL_BUDGET_S): + with pytest.raises(RuntimeError, match="injected chip init failure"): + l3.run(lambda orch, args, config: None) + finally: + l3.close() + + def test_chip_init_hang_trips_deadline(self, monkeypatch): + import simpler.worker as worker_mod # noqa: PLC0415 + + monkeypatch.setattr(worker_mod, "ChipWorker", _FakeChipHangs) + l3 = self._make_l3(timeout_s=1.5) + l3.register(lambda args: None) + l3.init() + start = time.monotonic() + try: + with _hard_timeout(_TEST_WALL_BUDGET_S): + with pytest.raises(RuntimeError, match="deadline"): + l3.run(lambda orch, args, config: None) + assert 1.5 <= time.monotonic() - start < _TEST_WALL_BUDGET_S + finally: + l3.close()