diff --git a/docs/comm-domain.md b/docs/comm-domain.md index e1f43b9048..ab661156a9 100644 --- a/docs/comm-domain.md +++ b/docs/comm-domain.md @@ -159,16 +159,15 @@ rewritten before decoding. Two sources are legal: | Source | How | Why it works | | ------ | --- | ------------ | -| **fork-inherited** | `tensor.share_memory_()` **before the local L3 children are forked** (i.e. before the first `Worker.run()`) | the child inherits the MAP_SHARED page at fork | +| **fork-inherited** | `tensor.share_memory_()` **before `Worker.init()`** (before the local L3 children are forked) | the child inherits the MAP_SHARED page at fork | | **worker-allocated post-fork** | `worker.create_host_buffer(nbytes)` after the children exist | born-shared memory attached into every local child, **zero-copy** | -The local L3 children are forked lazily on the **first** `run()`. A host tensor +The local L3 children are forked eagerly in `Worker.init()`. A host tensor created after that — the natural dynamic-shape serving pattern — is invisible to the children unless it lives in a `create_host_buffer` buffer: ```python -worker = Worker(level=3, ...); worker.register(chip); worker.init() -worker.run(orch0, ...) # forks the chips +worker = Worker(level=3, ...); worker.register(chip); worker.init() # forks the chips buf_h = worker.create_host_buffer(tokens * hidden_size * 4) # born-shared, post-fork buf_o = worker.create_host_buffer(batch * vocab * 4) @@ -215,7 +214,7 @@ the child — allocate it with `create_host_buffer` instead. - **`orch.copy_to` is the unmanaged low-level path.** `create_host_buffer` covers the `run` / `submit_next_level` host-tensor args. The explicit `orch.copy_to(src=tensor.data_ptr())` staging path (§5) is *not* validated — - its `src` must be fork-inherited (`.share_memory_()` before the first `run()`) + its `src` must be fork-inherited (`.share_memory_()` before `init()`) or a `create_host_buffer` buffer. - **Fork-inherited anonymous memory is copy-on-write, hence stale.** Even a tensor the child legitimately inherited is only useful as a *live* input if it diff --git a/docs/hierarchical_level_runtime.md b/docs/hierarchical_level_runtime.md index bc361ad970..6ebf6b1a5f 100644 --- a/docs/hierarchical_level_runtime.md +++ b/docs/hierarchical_level_runtime.md @@ -177,8 +177,9 @@ Communication channels: A higher-level `Worker` can register a lower-level `Worker` as a NEXT_LEVEL child through the same mailbox protocol L3 uses for chip children. The Python `Worker.add_worker(child)` stores an un-init'd child -Worker; on first `run()`, the parent forks a child process that inits the -inner Worker and enters a mailbox-polling loop (`_child_worker_loop`). +Worker; `init()` then eagerly forks a child process that recursively inits the +inner Worker (bringing up its whole subtree before publishing `INIT_READY`) +and enters a mailbox-polling loop (`_child_worker_loop`). ```python # L3 child: sub-only (or with chips via device_ids) @@ -214,8 +215,9 @@ Worker opens its own scope, executes the orch function with its own — recursion is symmetric. **Nested fork ordering**: L3's own children (sub/chip) are forked **inside** -the L4 child process, on L3's first `run()`. This keeps the process tree -clean: L4 parent → L3 child → L3's sub/chip grandchildren. +the L4 child process, in L3's eager `init()` (which the L4 child runs before it +publishes `INIT_READY`). This keeps the process tree clean: L4 parent → L3 child +→ L3's sub/chip grandchildren. **Mode per level is independent**: L3 can use PROCESS (chip children), while L4 also uses PROCESS (L3 Worker children). Each `Worker` picks its children's diff --git a/docs/task-flow.md b/docs/task-flow.md index 4c16a36350..da6fa9fb88 100644 --- a/docs/task-flow.md +++ b/docs/task-flow.md @@ -460,27 +460,30 @@ that maps to a Python orchestration function, not a `ChipCallable`. ### Fork sequence -L4's `init()` allocates the L4 Worker's HeapRing (before fork). -On first `run()`, the deferred `_start_hierarchical()`: +L4's `init()` allocates the L4 Worker's HeapRing (before fork), then eagerly +runs `_start_hierarchical()` — `init()` is the single startup point, and it +returns only once the whole tree is READY: 1. Forks one child process per L3 Worker child -2. **Inside the child**: `inner_worker.init()` creates the L3 Worker - (mmaps L3's own HeapRing), allocates L3's sub/chip mailboxes. L3's - own children are forked lazily on L3's first `run()`. -3. Child enters `_child_worker_loop(mailbox, registry, inner_worker)` -4. **Parent**: registers each mailbox with L4's Worker via - `add_next_level_worker_at(worker_id, mailbox_addr)` +2. **Inside the child**: `inner_worker.init()` is eager and recursive — it + creates the L3 Worker (mmaps L3's own HeapRing), forks L3's sub/chip + children, and blocks on their `INIT_READY` before it returns. +3. Child publishes `INIT_READY` (whole L3 subtree ready), then enters + `_child_worker_loop(mailbox, registry, inner_worker)` for dispatch +4. **Parent**: awaits each child's `INIT_READY`, then registers each mailbox + with L4's Worker via `add_next_level_worker_at(worker_id, mailbox_addr)` ```text L4 parent process ├─ Worker(4) + HeapRing (MAP_SHARED, inherited by L3 child) └─ fork ──────────────────► L3 child process - ├─ inner_worker.init() - │ └─ Worker(3) + L3's own HeapRing + ├─ inner_worker.init() (eager, recursive) + │ ├─ Worker(3) + L3's own HeapRing + │ └─ forks L3's sub/chip children, awaits + │ their INIT_READY + ├─ publish INIT_READY (subtree ready) └─ _child_worker_loop(mbox, registry, inner_worker) - └─ on first dispatch: - inner_worker.run(orch_fn, args, cfg) - └─ _start_hierarchical() forks L3's sub children + └─ on dispatch: inner_worker.run(orch_fn, args, cfg) ``` ### Dispatch walkthrough diff --git a/docs/worker-manager.md b/docs/worker-manager.md index 80bc16aa91..7be9c19d1c 100644 --- a/docs/worker-manager.md +++ b/docs/worker-manager.md @@ -258,23 +258,31 @@ the session runner reports `HELLO READY`. ### 4.1 Nested fork ordering (L4+ Worker children) -When an L4 Worker has L3 Worker children, the fork sequence nests: +When an L4 Worker has L3 Worker children, `init()` is eager and the fork +sequence nests recursively — the whole tree is READY when `L4.init()` returns: ```text L4 parent process ├─ _init_hierarchical(): Worker(4) + HeapRing mmap (before fork) - └─ _start_hierarchical() (on first run): + └─ _start_hierarchical() (inside init()): ├─ fork L3 child ────────► L3 child process: - │ inner_worker.init() ← Worker(3) + L3 HeapRing - │ _child_worker_loop() - │ └─ on first dispatch: inner_worker.run() - │ └─ _start_hierarchical() forks L3's sub/chip children + │ inner_worker.init() ← eager, recursive + │ ├─ Worker(3) + L3 HeapRing + │ └─ _start_hierarchical() forks L3's + │ sub/chip children and blocks on their + │ INIT_READY, THEN publishes INIT_READY + │ _child_worker_loop() (dispatch only) + ├─ await every L3 child's INIT_READY (whole subtree ready) └─ register mailbox with L4's Worker ``` Each inner Worker inits **inside its forked child process** so its own -children are forked from the correct parent. The L4 parent never sees L3's -sub/chip grandchildren — they're L3's responsibility. +children are forked from the correct parent. Because the inner `init()` is +eager and blocks on its descendants, a child publishes `INIT_READY` only after +its whole subtree is ready, so readiness propagates recursively up to +`L4.init()`. The L4 parent never sees L3's sub/chip grandchildren — they're +L3's responsibility; if startup fails, they are reclaimed via the child's +process-group cancellation domain, not the resource_tracker. **Key invariant**: `Worker(N)` and its HeapRing are created before any fork at level N. Children inherit the `MAP_SHARED` mmap at the same virtual diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index 8b250a244b..4be4b91aa4 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -596,6 +596,12 @@ NB_MODULE(_task_interface, m) { m.attr("MAX_TENSOR_DIMS") = MAX_TENSOR_DIMS; m.attr("MAX_REGISTERED_CALLABLE_IDS") = MAX_REGISTERED_CALLABLE_IDS; m.attr("RUNTIME_ENV_RING_COUNT") = RUNTIME_ENV_RING_COUNT; + // Byte size of a Tensor and the offset of its child_memory flag within it. + // A task-args blob stores Tensors as a raw memcpy array, so a Python-side + // blob walker locates tensor i's fields at i * TENSOR_STRIDE_BYTES without + // reimplementing the struct layout. + m.attr("TENSOR_STRIDE_BYTES") = static_cast(sizeof(Tensor)); + m.attr("TENSOR_CHILD_MEMORY_OFFSET") = static_cast(offsetof(Tensor, child_memory)); // --- Tensor --- // The unified strided tensor descriptor. Constructed contiguous via make() diff --git a/python/simpler/remote_l3_session.py b/python/simpler/remote_l3_session.py index 161366d891..66d42dbc2b 100644 --- a/python/simpler/remote_l3_session.py +++ b/python/simpler/remote_l3_session.py @@ -21,10 +21,12 @@ import importlib import json import os +import signal import socket import struct import sys import threading +import time import traceback from dataclasses import dataclass from multiprocessing import shared_memory @@ -910,9 +912,17 @@ def run_session(manifest: dict[str, Any], ready_fd: int) -> int: try: session_timeout_s = _session_timeout_s(manifest) manifest_dispatch_registry = _install_manifest_dispatcher_registry(manifest) - inner_worker.init() + # Register the inner L3 callables before init() so they are frozen into + # the eager startup snapshot and uploaded to the chip children before + # those children publish INIT_READY. init() is the single, eager startup + # point: it forks and readies the whole inner L3->L2 tree, so the runner + # reports ready (below) only once that subtree is up. manifest_inner_handles = _install_manifest_inner_registry(manifest, inner_worker) - inner_worker._start_hierarchical() # noqa: SLF001 -- session runner owns the fork-safe prestart. + # Bound the inner startup and mark it non-root so its chip/sub children + # inherit this runner's process group (see start_new_session in the + # daemon) rather than splitting into their own — the daemon reaps the + # whole L3->L2 subtree with one killpg on the runner. + inner_worker.init(_startup_deadline=time.monotonic() + session_timeout_s) listen_host = str(manifest.get("listen_host", "127.0.0.1")) command_sock = _bind_listener(listen_host) @@ -963,11 +973,18 @@ def run_session(manifest: dict[str, Any], ready_fd: int) -> int: pass +def _raise_keyboard_interrupt(_signum, _frame): + # A daemon cooperative-kill (SIGTERM) unwinds run_session so its finally + # closes the inner Worker — reaping the L3->L2 subtree — before exit. + raise KeyboardInterrupt + + def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser() parser.add_argument("--manifest", required=True) parser.add_argument("--ready-fd", required=True, type=int) ns = parser.parse_args(argv) + signal.signal(signal.SIGTERM, _raise_keyboard_interrupt) with open(ns.manifest, encoding="utf-8") as f: manifest = json.load(f) return run_session(manifest, ns.ready_fd) diff --git a/python/simpler/remote_l3_worker.py b/python/simpler/remote_l3_worker.py index 59dce0c886..a51e939704 100644 --- a/python/simpler/remote_l3_worker.py +++ b/python/simpler/remote_l3_worker.py @@ -11,9 +11,11 @@ from __future__ import annotations import argparse +import contextlib import json import os import select +import signal import socket import struct import subprocess @@ -97,14 +99,24 @@ def _wait_or_kill_runner(proc: subprocess.Popen[Any], *, timeout_s: float = 5.0) return except subprocess.TimeoutExpired: pass - try: - proc.kill() - except OSError: - pass + # Cooperative: SIGTERM lets the runner close its inner Worker and reap its + # own L3->L2 subtree (unlinking nested shms) before it exits. + with contextlib.suppress(OSError): + proc.terminate() try: proc.wait(timeout=timeout_s) + return except subprocess.TimeoutExpired: pass + # Hard backstop: SIGKILL the whole runner process group. The runner is a + # session leader (start_new_session) and its inner chip/sub children inherit + # its group, so this reaps the entire subtree rather than orphaning it. + with contextlib.suppress(OSError, ProcessLookupError): + os.killpg(proc.pid, signal.SIGKILL) + with contextlib.suppress(OSError): + proc.kill() + with contextlib.suppress(subprocess.TimeoutExpired): + proc.wait(timeout=timeout_s) def _reap_session_runner(proc: subprocess.Popen[Any]) -> None: @@ -138,17 +150,15 @@ def _start_session(manifest: dict[str, Any]) -> dict[str, Any]: ], pass_fds=(ready_w,), close_fds=True, + # Own session/group so the daemon can killpg the whole runner + # subtree (runner + eagerly-forked inner L3->L2 children). + start_new_session=True, ) os.close(ready_w) ready_w = -1 try: ready = _read_runner_ready(ready_r, timeout_s) except BaseException: - if proc.poll() is None: - try: - proc.kill() - except OSError: - pass _wait_or_kill_runner(proc) raise ready["pid"] = int(proc.pid) diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 74dd5f1348..9eae4dd38d 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -60,6 +60,7 @@ def my_l4_orch(orch, args, config): from __future__ import annotations import bisect +import contextlib import ctypes import importlib import json @@ -81,6 +82,7 @@ def my_l4_orch(orch, args, config): from _task_interface import ( # pyright: ignore[reportMissingImports] MAX_REGISTERED_CALLABLE_IDS, RUNTIME_ENV_RING_COUNT, + TENSOR_CHILD_MEMORY_OFFSET, WorkerType, _l3_child_onboard_region_close, _l3_child_onboard_region_create, @@ -450,11 +452,14 @@ def _rewrite_blob_host_addrs(buf: memoryview, blob_off: int, ranges: list[tuple[ """Redirect registered host pointers in a task-args blob to child mappings. ``ranges`` is ``(parent_lo, parent_hi, child_base)`` for each host buffer the - child has mapped via _CTRL_MAP_HOST. For every tensor whose ``buffer.addr`` - (a parent VA) lands in a registered range, rewrite it in place to - ``child_base + (addr - parent_lo)`` so the runtime dereferences the child's - own mapping. Tensors outside every range (fork-inherited or child-allocated) - are left untouched. See _BLOB_TENSOR_STRIDE for the wire layout. + child has mapped via _CTRL_MAP_HOST. For every host tensor whose + ``buffer.addr`` (a parent VA) lands in a registered range, rewrite it in + place to ``child_base + (addr - parent_lo)`` so the runtime dereferences the + child's own mapping. Tensors outside every range (fork-inherited or + child-allocated) are left untouched. A ``child_memory`` tensor carries a + child-owned device pointer, never a host VA, so it is skipped even when its + address numerically falls inside a registered host range — rewriting it would + corrupt the device pointer. See _BLOB_TENSOR_STRIDE for the wire layout. """ tensor_count = struct.unpack_from(" None: """Unified TASK_READY / CONTROL_REQUEST / SHUTDOWN state machine. @@ -1358,13 +1366,13 @@ def _run_chip_main_loop( # noqa: PLR0912, PLR0913, PLR0915 -- unified TASK_READ Returning a non-zero code overrides the kernel's success. TASK_READY carries a callable digest. The child resolves it to a - target-local slot and runs it. The slot must already be prepared via - ``_CTRL_PREPARE`` (the explicit registration path the parent pushes after - init() to stage the H2D upload + device-orch load); a TASK_READY for an - unprepared slot is a control-flow error and fails the task rather than - lazily preparing it. + target-local slot and runs it. The slot must already be prepared: initial + startup-snapshot ChipCallables are prepared before INIT_READY (carried in via + ``prepared``), and callables registered dynamically after startup arrive via + ``_CTRL_PREPARE``. A TASK_READY for an unprepared slot is a control-flow + error and fails the task rather than lazily preparing it. """ - prepared: set[int] = set() + prepared = prepared if prepared is not None else set() l3_l2_region_store = _L2HostL3L2RegionStore() # Post-fork host buffers mapped into this child. `host_buf_table` # owns the mmap per token (for unmap + teardown); `host_buf_ranges` is the @@ -1576,6 +1584,24 @@ def _chip_process_loop( # noqa: PLR0913 -- fork-child entry: all context (bins, _mailbox_store_i32(_buffer_field_addr(buf, _OFF_STATE), _INIT_FAILED) return + # Prepare every ChipCallable in the startup snapshot before publishing + # INIT_READY, so the H2D upload + device-orch load is charged inside the + # readiness barrier and the first task dispatch pays no upload. The set of + # prepared cids carries into the main loop, which requires a cid be prepared + # before it dispatches. The parent therefore issues no post-READY + # control_prepare for the initial snapshot. + prepared: set[int] = set() + try: + for cid, target in registry.items(): + if isinstance(target, ChipCallable): + _ensure_prepared(cw, registry, prepared, int(cid), device_id=device_id) + except Exception as e: + _tb.print_exc() + _write_error(buf, 1, _format_exc(f"chip_process dev={device_id} prepare", e)) + _mailbox_store_i32(_buffer_field_addr(buf, _OFF_STATE), _INIT_FAILED) + cw.finalize() + return + mailbox_addr = ctypes.addressof(ctypes.c_char.from_buffer(buf)) state_addr = mailbox_addr + _OFF_STATE # Signal init success. The parent's readiness barrier waits for every chip @@ -1598,6 +1624,7 @@ def _chip_process_loop( # noqa: PLR0913 -- fork-child entry: all context (bins, identity_refs, chip_platform=platform, chip_runtime=runtime, + prepared=prepared, ) finally: cw.finalize() @@ -1736,7 +1763,13 @@ def _child_worker_loop( break -def _forked_child_main(buf: memoryview, label: str, setup, serve) -> None: +class _StartupCancelled(BaseException): + """Raised inside a forked child when the parent cooperatively cancels its + startup (SIGTERM). Unwinds the child's own ``setup`` — recursively rolling + back any grandchildren it already forked — before it exits.""" + + +def _forked_child_main(buf: memoryview, label: str, setup, serve, make_group_leader: bool = False) -> None: """Run a forked child to completion, always terminating via ``os._exit``. ``setup()`` runs the child's fallible init and returns an opaque context; @@ -1744,6 +1777,13 @@ def _forked_child_main(buf: memoryview, label: str, setup, serve) -> None: success the child publishes INIT_READY (the parent's readiness barrier unblocks) and ``serve(ctx)`` runs the mailbox loop. + ``make_group_leader`` puts the child in its own process group so the startup + root can reap the whole subtree (this child plus every descendant it forks) + with one ``killpg``; deeper descendants inherit the group and do not set + their own. During ``setup`` a SIGTERM is a cooperative cancel: it raises + ``_StartupCancelled``, which unwinds ``setup`` (recursively tearing down any + grandchildren and their nested shms) before the child exits. + 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 @@ -1752,14 +1792,32 @@ def _forked_child_main(buf: memoryview, label: str, setup, serve) -> None: """ import traceback as _tb # noqa: PLC0415 + if make_group_leader: + with contextlib.suppress(OSError): + os.setpgid(0, 0) + state_addr = _buffer_field_addr(buf, _OFF_STATE) + + def _on_cancel(_signum, _frame): + raise _StartupCancelled() + + prev_term = signal.signal(signal.SIGTERM, _on_cancel) try: ctx = setup() + except _StartupCancelled: + # Parent cancelled us mid-init; setup() already unwound its own subtree. + _tb.print_exc() + os._exit(1) 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) + # Serving is torn down via the SHUTDOWN mailbox state, not the cancel signal. + # signal.signal returns None when the prior handler was not installed from + # Python (e.g. a C library / host default); restore SIG_DFL in that case so + # the round-trip does not raise TypeError. + signal.signal(signal.SIGTERM, prev_term if prev_term is not None else signal.SIG_DFL) _mailbox_store_i32(state_addr, _INIT_READY) try: serve(ctx) @@ -1828,6 +1886,15 @@ def __init__( self._hierarchical_start_state = "not_started" self._hierarchical_start_mu = threading.Lock() self._hierarchical_start_cv = threading.Condition(self._hierarchical_start_mu) + # Absolute time.monotonic() deadline for the current startup epoch, set + # once at init() and shared by every child group and recursive descendant + # so the whole tree comes up within a single startup_timeout_s budget. + self._startup_deadline: float = 0.0 + # True on the worker whose init() the user called (the startup root). + # The root's direct children are process-group leaders, so the root can + # killpg a whole subtree; nested (recursive) workers inherit their + # parent's group and rely on the root's killpg as the hard backstop. + self._is_startup_root: bool = True # Optional CallConfig whose ring sizing is pre-warmed at init() so the # first run() with the same sizing skips the (~800ms) cold prebuilt @@ -1912,16 +1979,20 @@ def _allocate_next_level_worker_id(self) -> int: return worker_id def add_remote_worker(self, spec: RemoteWorkerSpec) -> int: - if self._initialized: - raise RuntimeError("Worker.add_remote_worker after init") - if self.level < 4: - raise TypeError("Worker.add_remote_worker: remote L3 workers require a level >= 4 parent") - if not isinstance(spec, RemoteWorkerSpec): - raise TypeError("Worker.add_remote_worker expects a RemoteWorkerSpec") - worker_id = self._allocate_next_level_worker_id() - self._remote_worker_specs.append(spec) - self._remote_worker_ids.append(worker_id) - return worker_id + # Hold the lifecycle lock across the state check and the topology + # mutation so a concurrent init() cannot freeze the topology snapshot + # between them. + with self._hierarchical_start_cv: + if self._initialized or self._hierarchical_start_state != "not_started": + raise RuntimeError("Worker.add_remote_worker after init") + if self.level < 4: + raise TypeError("Worker.add_remote_worker: remote L3 workers require a level >= 4 parent") + if not isinstance(spec, RemoteWorkerSpec): + raise TypeError("Worker.add_remote_worker expects a RemoteWorkerSpec") + worker_id = self._allocate_next_level_worker_id() + self._remote_worker_specs.append(spec) + self._remote_worker_ids.append(worker_id) + return worker_id @staticmethod def _parse_remote_endpoint(endpoint: str) -> tuple[str, int]: @@ -2058,7 +2129,6 @@ def _require_remote_worker_started(self, worker_id: int) -> None: raise RuntimeError("remote memory APIs require Worker.init() before allocation or copy") if int(worker_id) not in set(self._remote_worker_ids): raise ValueError("remote memory APIs require a remote worker id returned by add_remote_worker") - self._start_hierarchical() if self._worker is None: raise RuntimeError("remote memory APIs require a started hierarchical Worker") @@ -2523,6 +2593,28 @@ def _resolve_handle( with self._registry_lock: return self._resolve_handle_locked(handle, expected_namespace=expected_namespace) + def _register_into_snapshot_or_wait(self, reg: _CallableRegistration) -> CallableHandle | None: + """Linearize a level>=3 register against the startup epoch. + + Waits out an in-progress init() (INITIALIZING); a FAILED epoch raises. A + pre-start registration is installed into the startup snapshot and its + handle returned; once the hierarchy is READY, returns None so the caller + takes its post-start control-broadcast path. + """ + with self._hierarchical_start_cv: + while self._hierarchical_start_state == "starting": + self._hierarchical_start_cv.wait() + if self._hierarchical_start_state == "failed": + raise RuntimeError("Worker hierarchical startup failed; close this Worker and create a new one") + pre_start = self._hierarchical_start_state != "started" and not getattr( + self, "_hierarchical_started", False + ) + if pre_start: + with self._registry_lock: + handle, _is_new = self._install_registration_locked(reg) + return handle + return None + def register(self, target, *, workers: list[int] | None = None) -> CallableHandle: """Register a callable for dispatch and return an opaque handle. @@ -2545,24 +2637,18 @@ def register(self, target, *, workers: list[int] | None = None) -> CallableHandl "Worker.register(RemoteCallable): workers must name remote worker ids returned by " "add_remote_worker" ) - if not self._initialized: - with self._registry_lock: - handle, _is_new = self._install_registration_locked(reg) + # Linearize against the startup epoch exactly like the local path: a + # register that races an in-progress init() waits for it, then a + # pre-start registration lands in the snapshot while a post-READY one + # goes through the remote prepare/commit control path. + handle = self._register_into_snapshot_or_wait(reg) + if handle is not None: return handle return self._post_start_register_remote(reg) if self.level >= 3: - with self._hierarchical_start_cv: - while self._hierarchical_start_state == "starting": - self._hierarchical_start_cv.wait() - if self._hierarchical_start_state == "failed": - raise RuntimeError("Worker hierarchical startup failed; close this Worker and create a new one") - pre_start = self._hierarchical_start_state != "started" and not getattr( - self, "_hierarchical_started", False - ) - if pre_start: - with self._registry_lock: - handle, _is_new = self._install_registration_locked(reg) - return handle + handle = self._register_into_snapshot_or_wait(reg) + if handle is not None: + return handle if not isinstance(target, ChipCallable): return self._post_start_register_python(reg) @@ -2640,7 +2726,6 @@ def _post_start_register_remote( # noqa: PLR0912 -- two-phase remote register/c handle, _is_new = self._install_registration_locked(reg) return handle - self._start_hierarchical() if self._worker is None: raise RuntimeError("Worker.register(RemoteCallable): hierarchical worker is not started") @@ -2970,10 +3055,10 @@ def _post_init_register(self, target: ChipCallable, digest: bytes, *, is_new: bo ``mailbox_mu_`` so the broadcast serializes against any in-flight dispatch on each child mailbox. No Python lock required. """ - # Chip children are forked lazily on the first Worker.run() via - # _start_hierarchical; before that point the chip mailboxes have no - # reader and a CTRL_REGISTER broadcast would deadlock. In that pre-fork - # window, the startup snapshot carries the digest and target bytes. + # Until init() has started the hierarchy the chip mailboxes have no + # reader, so a CTRL_REGISTER broadcast would deadlock; a registration in + # that window is instead carried by the startup snapshot and + # COW-inherited by the children forked in init(). if not getattr(self, "_hierarchical_started", False): return assert self._worker is not None @@ -3169,7 +3254,6 @@ def _unregister_remote_handle(self, handle: CallableHandle) -> None: errors: list[str] = [] try: if self._initialized: - self._start_hierarchical() assert self._worker is not None for worker_id in worker_ids: try: @@ -3298,54 +3382,104 @@ def add_worker(self, worker: Worker) -> int: raise RuntimeError("Worker.add_worker() requires level >= 4") if self._config.get("device_ids", []): raise RuntimeError("Worker.add_worker() cannot be combined with device_ids on the same Worker") - if self._initialized: - raise RuntimeError("Worker.add_worker() must be called before init()") if worker._initialized: raise RuntimeError("Child worker must not be initialized before add_worker()") - worker_id = self._allocate_next_level_worker_id() - self._next_level_workers.append(worker) - self._next_level_worker_ids.append(worker_id) - return worker_id + # Hold the lifecycle lock across the state check and the topology + # mutation so a concurrent init() cannot freeze the topology snapshot + # between them. + with self._hierarchical_start_cv: + if self._initialized or self._hierarchical_start_state != "not_started": + raise RuntimeError("Worker.add_worker() must be called before init()") + worker_id = self._allocate_next_level_worker_id() + self._next_level_workers.append(worker) + self._next_level_worker_ids.append(worker_id) + return worker_id # ------------------------------------------------------------------ # init — auto-discovery # ------------------------------------------------------------------ - def init(self, prewarm_config=None) -> None: - """Initialize the worker. + def init(self, prewarm_config=None, *, _startup_deadline: float | None = None) -> None: + """Initialize the worker and bring its whole subtree to READY. + + For an L3+ worker ``init`` is the single startup submission point: it + forks every local child (sub / chip / next-level), waits for the whole + subtree — recursively, for L4+ — to publish INIT_READY, activates any + remote L3 sessions, starts the C++ scheduler, and only then publishes + READY in one atomic commit. It returns with the tree ready to run, or + raises after a bounded rollback that leaves no child, mailbox, or session + behind. ``run`` / ``create_host_buffer`` / the remote register/memory + APIs never trigger startup. Args: prewarm_config: Optional CallConfig. When given, its ring sizing (``runtime_env.ring_task_window`` / ``ring_heap`` / ``ring_dep_pool``) is built + cached so the first ``run`` with the same sizing skips the (~800ms) cold prebuilt runtime-arena build. - Timing is level-dependent: an L2 worker prewarms here in - ``init``; an L3+ worker has no chip children until the hierarchy - starts on the first ``run``, so it prewarms there — during - hierarchy startup, alongside the callable SO upload and before the - first task dispatches. A no-op for runtimes without a prebuilt - arena (host_build_graph). ``None`` (default) disables prewarm. + An L2 worker prewarms here; an L3+ worker prewarms each chip child + during hierarchy startup, before it publishes INIT_READY. A no-op + for runtimes without a prebuilt arena (host_build_graph). ``None`` + (default) disables prewarm. + _startup_deadline: Internal. Absolute ``time.monotonic()`` deadline + inherited from a parent's startup epoch so a recursive descendant + consumes the parent's remaining budget instead of restarting the + timeout. ``None`` starts a fresh epoch. """ - if self._initialized: - raise RuntimeError("Worker already initialized") - # Validate up front so a bad config fails here regardless of level, even - # though an L3+ worker does not build the arena until the first run. if prewarm_config is not None: prewarm_config.validate() - self._prewarm_config = prewarm_config + # Claim the startup epoch atomically: NEW -> INITIALIZING under the + # lifecycle lock so a concurrent init / register / close observes one + # linear transition and never a half-built Worker. The INITIALIZING / + # READY / FAILED state machine is a hierarchical (level >= 3) concept; an + # L2 worker inits synchronously in-process and never enters "starting". + with self._hierarchical_start_cv: + if self._initialized: + raise RuntimeError("Worker already initialized") + if self._hierarchical_start_state == "starting": + raise RuntimeError("Worker.init() is already in progress") + if self._hierarchical_start_state == "failed": + raise RuntimeError("Worker startup failed; close this Worker and create a new one") + self._prewarm_config = prewarm_config + if self.level >= 3: + self._hierarchical_start_state = "starting" + self._is_startup_root = _startup_deadline is None + own_deadline = time.monotonic() + self._startup_timeout_s + # A recursive descendant caps its own timeout at the parent's + # remaining budget so the whole tree fits one startup_timeout_s. + self._startup_deadline = ( + own_deadline if _startup_deadline is None else min(_startup_deadline, own_deadline) + ) + self._hierarchical_start_cv.notify_all() try: if self.level == 2: self._init_level2() elif self.level >= 3: self._init_hierarchical() + self._start_hierarchical() else: raise ValueError(f"Worker: level {self.level} not supported") except BaseException: - self._cleanup_partial_init() + # Commit FAILED even if rollback itself raises, so a waiter blocked on + # the lifecycle condition (register / close) is always released. + try: + self._cleanup_partial_init() + finally: + if self.level >= 3: + with self._hierarchical_start_cv: + self._hierarchical_start_state = "failed" + self._hierarchical_start_cv.notify_all() raise - self._initialized = True + # Atomic READY commit: publish _initialized together with the started + # state so no thread ever observes a started hierarchy while + # _initialized is still False. + with self._hierarchical_start_cv: + self._initialized = True + if self.level >= 3: + self._hierarchical_started = True + self._hierarchical_start_state = "started" + self._hierarchical_start_cv.notify_all() def _init_level2(self) -> None: from simpler_setup.runtime_builder import RuntimeBuilder # noqa: PLC0415 @@ -3423,6 +3557,12 @@ def _init_hierarchical(self) -> None: else: self._worker = _Worker(self.level, int(heap_ring_size)) + # Open remote L3 sessions (the remote runner brings up its own L3->L2 + # tree and answers the daemon handshake). Opening only exchanges a + # transient socket message — it starts no local thread. Registering the + # socket into the C++ Worker (which spawns the RemoteL3Endpoint health + # thread) is deferred to _start_hierarchical, after the last local fork, + # to preserve the fork-before-thread invariant. opened_remote_sessions: list[_RemoteSession] = [] try: for worker_id, spec in zip(self._remote_worker_ids, self._remote_worker_specs, strict=True): @@ -3434,17 +3574,6 @@ def _init_hierarchical(self) -> None: spec=spec, worker_id=worker_id, session_id=session_id, timeout_s=timeout_s ) opened_remote_sessions.append(session) - assert self._worker is not None - self._worker.add_remote_l3_socket( - worker_id, - session_id, - spec.transport, - session.command_host, - session.command_port, - session.health_host, - session.health_port, - timeout_s, - ) self._remote_sessions.append(session) opened_remote_sessions.pop() except BaseException: @@ -3454,220 +3583,218 @@ def _init_hierarchical(self) -> None: self._hierarchical_started = False def _start_hierarchical(self) -> None: # noqa: PLR0912 -- three parallel fork loops (sub/chip/next) + bootstrap wait + scheduler register/init; branches track the fork order documented in the body - """Fork child processes and start C++ scheduler. Called on first run().""" + """Fork every local child, await the subtree, register endpoints, start the scheduler. + + Called only by init(), which owns the lifecycle state. Any failure here + propagates to init(), whose single rollback entry (_cleanup_partial_init) + closes the C++ Worker (if the scheduler started) and tears the whole + epoch down. The readiness barriers raise on any child failure/exit/hang. + """ device_ids = self._config.get("device_ids", []) n_sub = self._config.get("num_sub_workers", 0) + deadline = self._startup_deadline - # 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 - # cannot return through the pre-start path after this point. - with self._hierarchical_start_cv: - while self._hierarchical_start_state == "starting": - self._hierarchical_start_cv.wait() - if self._hierarchical_start_state == "started": - return - if self._hierarchical_start_state == "failed": - raise RuntimeError("Worker hierarchical startup failed; close this Worker and create a new one") - self._hierarchical_start_state = "starting" - with self._registry_lock: - identity_snapshot = [ - (digest, state.target, state.ref_count, state.kind, state.target_namespace) - for digest, state in self._identity_registry.items() - ] - self._hierarchical_start_cv.notify_all() + # Freeze the startup registry snapshot. init() already holds the epoch in + # the "starting" state, so a concurrent register/unregister is blocked on + # the lifecycle condition and cannot slip a mutation in after this point. + with self._registry_lock: + identity_snapshot = [ + (digest, state.target, state.ref_count, state.kind, state.target_namespace) + for digest, state in self._identity_registry.items() + ] - self._startup_reaped_pids = set() - self._startup_ready_pids = set() + 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 + # 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 - def _setup(): - return _make_local_identity_tables( - identity_snapshot, - callable_kind=("PYTHON_SERIALIZED", "PYTHON_IMPORT"), - target_namespace="LOCAL_PYTHON", - ) + 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. - chip_log_level, chip_log_info_v = _simpler_log.get_current_config() - if device_ids: - for idx, dev_id in enumerate(device_ids): - pid = os.fork() - if pid == 0: - buf = self._chip_shms[idx].buf - assert buf is not None - # _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) - - # Cross-chip init barrier. ChipWorker.init can have a long - # right tail (e.g. PTO2_RING_HEAP=4 GiB pushes per-rank - # device_malloc beyond the host stream sync budget); without - # this barrier a fast-init chip starts its aclrtSyncStream - # window N seconds before a slow peer reaches the same - # 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. 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 - # HeapRing and allocates its own child mailboxes), then enter - # _child_worker_loop. The inner Worker's own children are forked - # lazily on first run() inside _child_worker_loop, so the process - # tree nests correctly: L4 → L3 child → L3's chip/sub children. - for idx, inner_worker in enumerate(self._next_level_workers): + _forked_child_main( + buf, + f"sub worker {i}", + _setup, + lambda t, b=buf: _sub_worker_loop(b, *t), + make_group_leader=self._is_startup_root, + ) + 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", deadline) + + # 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. + chip_log_level, chip_log_info_v = _simpler_log.get_current_config() + if device_ids: + for idx, dev_id in enumerate(device_ids): pid = os.fork() if pid == 0: - buf = self._next_level_shms[idx].buf + buf = self._chip_shms[idx].buf assert buf is not None + if self._is_startup_root: + with contextlib.suppress(OSError): + os.setpgid(0, 0) + # _chip_process_loop publishes INIT_READY/INIT_FAILED itself + # (around cw.init + ChipCallable prepare). 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) + + # Cross-chip init barrier. ChipWorker.init can have a long right tail + # (e.g. PTO2_RING_HEAP=4 GiB pushes per-rank device_malloc beyond the + # host stream sync budget); without this barrier a fast-init chip + # starts its aclrtSyncStream window N seconds before a slow peer + # reaches the same 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. A chip that fails or dies during init + # raises here rather than spinning forever. + self._await_children_ready(self._chip_shms, self._chip_pids, "chip", deadline) + + # Fork next-level Worker children (L4+ with Worker children). + # Each child process eagerly inits the inner Worker, which forks its own + # chip/sub (and, for L5+, deeper next-level) children and blocks on their + # readiness before returning — so the process tree nests correctly (L4 → + # L3 child → L3's chip/sub grandchildren) and INIT_READY propagates up + # only after the whole subtree is ready. + for idx, inner_worker in enumerate(self._next_level_workers): + pid = os.fork() + if pid == 0: + buf = self._next_level_shms[idx].buf + assert buf is not None - 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) + def _setup(inner=inner_worker): + # Propagate the fork-constant prewarm sizing and the shared + # startup deadline so the inner subtree comes up within the + # parent's remaining budget. INIT_READY is published only + # after BOTH the inner init (its whole subtree) and the + # identity-table build succeed, so the parent never observes + # READY for a child that then dies in fallible post-init + # setup. A failure after inner.init() succeeded tears the + # inner subtree back down before propagating, so a fallible + # post-init step leaves no orphaned grandchildren / shms. + inner.init(prewarm_config=self._prewarm_config, _startup_deadline=deadline) + try: return _make_local_identity_tables( identity_snapshot, callable_kind=("PYTHON_SERIALIZED", "PYTHON_IMPORT"), target_namespace="LOCAL_PYTHON", ) + except BaseException: + with contextlib.suppress(BaseException): + inner.close() + raise - _forked_child_main( - buf, - f"next_level worker {idx}", - _setup, - lambda tables, b=buf, inner=inner_worker: _child_worker_loop(b, *tables, inner), - ) - 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. - dw = self._worker - assert dw is not None + _forked_child_main( + buf, + f"next_level worker {idx}", + _setup, + lambda tables, b=buf, inner=inner_worker: _child_worker_loop(b, *tables, inner), + make_group_leader=self._is_startup_root, + ) + else: + self._next_level_pids.append(pid) + + # The recursive readiness edge: a next-level child's own init blocks on + # its descendants, so its INIT_READY means the whole subtree is ready. A + # failure, exit, or hang aborts startup here. + self._await_children_ready(self._next_level_shms, self._next_level_pids, "next_level", deadline) + + # Last local fork is done. Now — and only now — register remote L3 + # endpoints, which spawns the RemoteL3Endpoint health thread; doing it + # here keeps every local fork ahead of every C++/Python thread. Each + # remote attach consumes the shared startup budget. + for session, spec in zip(self._remote_sessions, self._remote_worker_specs, strict=True): + remaining = deadline - time.monotonic() + if remaining <= 0: + raise RuntimeError("remote L3 endpoint attach: startup deadline exceeded") + assert self._worker is not None + self._worker.add_remote_l3_socket( + session.worker_id, + session.session_id, + spec.transport, + session.command_host, + session.command_port, + session.health_host, + session.health_port, + min(self._remote_session_timeout_s(), remaining), + ) - # Register chip workers as NEXT_LEVEL (L3) - if device_ids: - for shm in self._chip_shms: - dw.add_next_level_worker(_mailbox_addr(shm)) + # _Worker was constructed in _init_hierarchical (pre-fork) so children + # inherit the HeapRing MAP_SHARED mmap. Register PROCESS-mode workers via + # the unified mailbox. + dw = self._worker + assert dw is not None - # Register Worker children as NEXT_LEVEL (L4+) - if self._next_level_shms and not hasattr(dw, "add_next_level_worker_at"): - raise RuntimeError("explicit NEXT_LEVEL worker ids require a rebuilt _task_interface module") - for idx, shm in enumerate(self._next_level_shms): - worker_id = self._next_level_worker_ids[idx] - dw.add_next_level_worker_at(worker_id, _mailbox_addr(shm)) + # Register chip workers as NEXT_LEVEL (L3) + if device_ids: + for shm in self._chip_shms: + dw.add_next_level_worker(_mailbox_addr(shm)) - for shm in self._sub_shms: - dw.add_sub_worker(_mailbox_addr(shm)) - - # Start Scheduler + WorkerThreads (C++ threads start here, after fork) - dw.init() - scheduler_started = True - - self._orch = Orchestrator(dw.get_orchestrator(), self) - - # Pre-warm every chip child: for each registered ChipCallable digest, - # send `_CTRL_PREPARE` to all chip children so the first - # `submit_next_level` does not pay the H2D upload cost. Sub fns / - # orch fns do not need pre-warming — the registry is already - # COW-inherited. - if device_ids: - for digest, target, _ref_count, kind, namespace in identity_snapshot: - if kind == "CHIP_CALLABLE" and namespace == "LOCAL_CHIP" and isinstance(target, ChipCallable): - for worker_id in range(len(self._chip_shms)): - dw.control_prepare(worker_id, digest) - - # The prebuilt runtime-arena is prewarmed inside each chip child's - # cw.init (the fork-constant sizing is COW-inherited via - # _chip_process_loop's prewarm_config arg) — no control command here. - - self._hierarchical_started = True - with self._hierarchical_start_cv: - self._hierarchical_start_state = "started" - self._hierarchical_start_cv.notify_all() - 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 + # Register Worker children as NEXT_LEVEL (L4+) + if self._next_level_shms and not hasattr(dw, "add_next_level_worker_at"): + raise RuntimeError("explicit NEXT_LEVEL worker ids require a rebuilt _task_interface module") + for idx, shm in enumerate(self._next_level_shms): + worker_id = self._next_level_worker_ids[idx] + dw.add_next_level_worker_at(worker_id, _mailbox_addr(shm)) + + for shm in self._sub_shms: + dw.add_sub_worker(_mailbox_addr(shm)) + + # Start Scheduler + WorkerThreads (C++ threads start here, after fork) + dw.init() - def _await_children_ready(self, shms, pids, kind: str) -> None: + self._orch = Orchestrator(dw.get_orchestrator(), self) + + # Every ChipCallable in the startup snapshot was already uploaded by its + # chip child before that child published INIT_READY (see + # _chip_process_loop), and the runtime arena was prewarmed there too — so + # there is no post-scheduler control_prepare on the startup path. + + def _await_children_ready(self, shms, pids, kind: str, deadline: float) -> 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 @@ -3675,11 +3802,12 @@ def _await_children_ready(self, shms, pids, kind: str) -> None: 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. + rollback never re-SIGKILLs a possibly-reused PID). ``deadline`` is the + single startup-epoch deadline shared by every child group and every + recursive descendant, so a deep tree cannot multiply the timeout; a + child that hangs past it aborts the epoch. 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 = [] @@ -3721,44 +3849,73 @@ def _await_children_ready(self, shms, pids, kind: str) -> None: # Hierarchical abort # ------------------------------------------------------------------ - 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. - - 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). + def _abort_hierarchical(self) -> None: # noqa: PLR0912 -- graceful/cooperative-cancel-then-killpg rollback across sub/chip/next-level, bounded-wait, reap, free shms + """Tear down the whole forked subtree + shms after a bootstrap failure. + + Called from the init() failure path — the single rollback entry + (_cleanup_partial_init) — so `dw.init()` may or may not have run. + + Teardown proceeds in two bounded phases within one cleanup budget: + + 1. Cooperative. A child that reached its serve loop (READY) is asked to + close gracefully via SHUTDOWN so it finalizes its device / unlinks its + own nested shms. A next-level child still inside ``inner.init()`` + (mid-init, and possibly already the parent of grandchildren) is sent + SIGTERM, which unwinds its ``inner.init()`` and recursively reclaims + its grandchildren and their nested shms. + 2. Hard backstop. Any child still alive past the cleanup deadline is + reaped. As the startup root, ``killpg`` takes the whole subtree — the + child and every descendant that inherited its process group — so a + mid-init grandchild is reaped here rather than left to the + multiprocessing resource_tracker; a nested (non-root) worker SIGKILLs + the direct pid and relies on the root's killpg. PIDs the barrier + already reaped are excluded so a reused PID is never signalled. """ reaped = set(self._startup_reaped_pids) graceful: list[int] = [] + cancelled: 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. + # A next-level child may have published INIT_READY in the window between + # the barrier last polling it and aborting on a failing sibling — its + # mailbox still reads INIT_READY (only the barrier resets it to IDLE). + # Promote it to READY so it is torn down gracefully (SHUTDOWN unlinks its + # own nested shms) rather than cooperatively cancelled after it has + # already restored the default SIGTERM disposition. for idx, pid in enumerate(self._next_level_pids): - if pid in reaped or pid not in self._startup_ready_pids: + if pid in reaped or pid 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: + if buf is not None and _mailbox_load_i32(_buffer_field_addr(buf, _OFF_STATE)) == _INIT_READY: + self._startup_ready_pids.add(pid) + + # Phase 1a: READY children (sub / chip / next-level) close gracefully. + for pids_list, shms_list in ( + (self._next_level_pids, self._next_level_shms), + (self._chip_pids, self._chip_shms), + (self._sub_pids, self._sub_shms), + ): + for idx, pid in enumerate(pids_list): + if pid in reaped or pid not in self._startup_ready_pids: + continue + 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) + graceful.append(pid) + + # Phase 1b: mid-init next-level children get a cooperative cancel so they + # unwind inner.init() and recursively reclaim their own subtree. + for pid in self._next_level_pids: + if pid in reaped or pid in self._startup_ready_pids: continue - _mailbox_store_i32(_buffer_field_addr(buf, _OFF_STATE), _SHUTDOWN) - graceful.append(pid) - if graceful: + with contextlib.suppress(ProcessLookupError, OSError): + os.kill(pid, signal.SIGTERM) + cancelled.append(pid) + + waiting = set(graceful) | set(cancelled) + if waiting: gdeadline = time.monotonic() + _ROLLBACK_GRACEFUL_TIMEOUT_S - waiting = set(graceful) while waiting and time.monotonic() <= gdeadline: for pid in list(waiting): try: @@ -3772,27 +3929,25 @@ def _abort_hierarchical(self) -> None: # noqa: PLR0912 -- graceful-close-then-k reaped.add(pid) if waiting: time.sleep(_STARTUP_POLL_INTERVAL_S) - # Any graceful child still alive overran the deadline; fall through - # to SIGKILL below. + # Phase 2: hard backstop for any survivor. A not-yet-reaped pid still + # holds its slot (no reuse), so killpg on the root reaps the survivor's + # whole group (it + inherited-group descendants) safely. pids = list(self._chip_pids) + list(self._sub_pids) + list(self._next_level_pids) for pid in pids: if pid in reaped: continue - try: + if self._is_startup_root: + with contextlib.suppress(ProcessLookupError, OSError): + os.killpg(pid, signal.SIGKILL) + with contextlib.suppress(ProcessLookupError, OSError): os.kill(pid, signal.SIGKILL) - killed.append(pid) - except ProcessLookupError: - pass - except OSError: - pass + killed.append(pid) for pid in pids: if pid in reaped and pid not in killed: continue - try: + with contextlib.suppress(ChildProcessError): os.waitpid(pid, 0) - except ChildProcessError: - pass self._last_rollback = { "graceful": [p for p in graceful if p not in killed], @@ -3850,10 +4005,6 @@ def _cleanup_partial_init(self) -> None: self._hierarchical_started = False self._comm_base_ready = False self._initialized = False - with self._hierarchical_start_cv: - if self._hierarchical_start_state != "started": - self._hierarchical_start_state = "not_started" - self._hierarchical_start_cv.notify_all() @property def live_domains(self) -> dict[str, CommDomainHandle]: @@ -4451,7 +4602,7 @@ def create_host_buffer(self, nbytes: int) -> HostBuffer: """Allocate a born-shared host buffer, attached into every local L3 child, that a later ``run()`` reads/writes with **no per-run copy**. - Local L3 children are forked lazily on the first ``run()``; memory created + Local L3 children are forked during ``init()``; host memory allocated afterwards is not in their address space. This hands you memory that is *born* in a shm already attached into every child, so there is nothing to copy: the child reads and writes the same physical pages the parent sees. @@ -4476,9 +4627,13 @@ def create_host_buffer(self, nbytes: int) -> HostBuffer: raise TypeError("create_host_buffer requires a level >= 3 Worker") if not self._initialized: raise RuntimeError("create_host_buffer requires Worker.init() before allocation") - self._start_hierarchical() - if not self._chip_shms: - raise RuntimeError("create_host_buffer requires forked chip children (none are configured)") + # A born-shared buffer is mapped into every direct process child (chip + # and sub alike, via _broadcast_host_control). Only a truly childless L3 + # has nowhere to attach it. + if not self._chip_shms and not self._sub_shms: + raise RuntimeError( + "create_host_buffer requires at least one forked chip or sub child (this Worker has none)" + ) assert self._worker is not None nbytes = int(nbytes) @@ -4760,7 +4915,6 @@ def run(self, callable, args=None, config=None) -> None: self._chip_worker._run_slot(state.slot_id, args, cfg) return None - self._start_hierarchical() assert self._orch is not None assert self._worker is not None # Drop any error stashed by a previous run() so this call starts @@ -4837,6 +4991,13 @@ def host_dlopen_count(self) -> int: # ------------------------------------------------------------------ def close(self) -> None: # noqa: PLR0912 -- parallel teardown for _worker + sub/chip/next/bootstrap shms with ordering constraints documented inline + # A close() that races an in-progress init() must not no-op on the still + # -false _initialized flag and leave the epoch's tree running: wait for + # the epoch to reach READY (then tear it down) or FAILED (already rolled + # back by init()). + with self._hierarchical_start_cv: + while self._hierarchical_start_state == "starting": + self._hierarchical_start_cv.wait() if not self._initialized: return diff --git a/simpler_setup/scene_test.py b/simpler_setup/scene_test.py index a99e02a1d4..04ed0fa2d9 100644 --- a/simpler_setup/scene_test.py +++ b/simpler_setup/scene_test.py @@ -167,6 +167,121 @@ def tensor_names(self) -> list[str]: return [s.name for s in self._specs if isinstance(s, Tensor)] +class _RehostedTaskArgs: + """Move a builder's host tensors into born-shared child buffers. + + ``Worker.init()`` is eager: the L3 chip/sub children are forked in ``init()``, + before ``generate_args()`` runs, so a plain post-init host tensor's raw VA is + not in any child's address space. Each host tensor is rehosted into its own + ``create_host_buffer`` (born-shared, mapped into every direct child) with + dtype / shape / value preserved, and the builder is rebound to a view over + that buffer so multi-round reset, dispatch, and golden compare all read and + write the same physical pages the children see. + + Each tensor gets an independent buffer (no aliasing of one registered range). + A non-contiguous / non-faithfully-representable layout is rejected rather + than silently copied to contiguous. ``release()`` frees every buffer in LIFO + order; a partial-construction failure rolls the builder back and frees what + was already allocated. + """ + + def __init__(self, worker, test_args: TaskArgsBuilder): + import torch # noqa: PLC0415 + + self._worker = worker + self._torch = torch + self._buffers: list = [] # (HostBuffer, view) in creation order + self._originals: dict[str, Any] = {} # name -> pre-rehost tensor + self._test_args = test_args + self._reject_aliased_tensors(test_args) + try: + new_specs = [] + for spec in test_args._specs: + # Only non-empty host tensors carry bytes across the process edge; + # an empty tensor is never dereferenced by the child, so it is + # left untouched rather than allocating a zero-length buffer. + if isinstance(spec, Tensor) and isinstance(spec.value, torch.Tensor) and spec.value.numel() > 0: + view = self._rehost_one(spec.value) + self._originals[spec.name] = test_args._data[spec.name] + test_args._data[spec.name] = view + new_specs.append(Tensor(spec.name, view)) + else: + new_specs.append(spec) + test_args._specs = new_specs + except BaseException: + self.release() + raise + + def _reject_aliased_tensors(self, test_args: TaskArgsBuilder) -> None: + # Two args whose storage byte-ranges overlap encode an OverlapMap + # dependency that independent born-shared buffers cannot preserve, so + # reject rather than silently split them into separate storage. + torch = self._torch + ranges: list = [] # (name, lo, hi) + for spec in test_args._specs: + if not (isinstance(spec, Tensor) and isinstance(spec.value, torch.Tensor)): + continue + t = spec.value + if t.numel() == 0: + continue + lo = t.data_ptr() + hi = lo + t.numel() * t.element_size() + for oname, olo, ohi in ranges: + if lo < ohi and olo < hi: + raise ValueError( + f"SceneTest rehost: tensors {spec.name!r} and {oname!r} alias overlapping storage; " + "an aliased layout is not faithfully representable across the process edge — build " + "them as independent tensors in generate_args()" + ) + ranges.append((spec.name, lo, hi)) + + def _rehost_one(self, t): + torch = self._torch + if t.device.type != "cpu": + raise ValueError( + f"SceneTest rehost: a host tensor crossing a process edge must be a CPU tensor, got " + f"device {t.device}; a device tensor must be declared child_memory, not rehosted to host" + ) + if not t.is_contiguous(): + raise ValueError( + "SceneTest rehost: a host tensor crossing a process edge must be contiguous to move " + "into a born-shared child buffer; a non-contiguous / aliased layout is not faithfully " + "representable — build it contiguous in generate_args()" + ) + nbytes = t.numel() * t.element_size() + buf = self._worker.create_host_buffer(nbytes) + try: + view = torch.frombuffer(buf.buffer, dtype=t.dtype, count=t.numel()).view(t.shape) + view.copy_(t) + except BaseException: + self._worker.free_host_buffer(buf) + raise + self._buffers.append((buf, view)) + return view + + def release(self) -> None: + # Restore the builder's original entries (dropping the born-shared view + # refs), then free each buffer in LIFO order after releasing its view. + for name, orig in self._originals.items(): + self._test_args._data[name] = orig + self._test_args._specs = [ + Tensor(s.name, self._originals[s.name]) if isinstance(s, Tensor) and s.name in self._originals else s + for s in self._test_args._specs + ] + self._originals.clear() + while self._buffers: + buf, view = self._buffers.pop() + del view + try: + buf.buffer.release() + except (ValueError, BufferError): + pass + try: + self._worker.free_host_buffer(buf) + except Exception as exc: # noqa: BLE001 -- best-effort cleanup; a leak here must not mask the test result, but process-control exceptions still propagate + logger.warning("SceneTest rehost cleanup: free_host_buffer failed: %s", exc) + + # --------------------------------------------------------------------------- # CallableNamespace — dot-access container for L3 callables # --------------------------------------------------------------------------- @@ -1109,51 +1224,59 @@ def _run_and_validate_l3( # noqa: PLR0913 -- threads CLI diagnostic flags + L3 golden_args = test_args.clone() self.compute_golden(golden_args, params) - # Save initial tensor values for reset between rounds - all_tensor_names = test_args.tensor_names() - initial_tensors = {} - if rounds > 1: - for name in all_tensor_names: - initial_tensors[name] = getattr(test_args, name).clone() - - # Build CallableNamespace: compiled ChipCallables + sub callable IDs - ns = CallableNamespace({**compiled_callables, **sub_handles}) - - # Get orch function (plain function from CALLABLE) - orch_fn = self.CALLABLE["orchestration"] - - # Execute rounds. As for L2 (see _run_and_validate_l2), per-round timing - # is obtained offline from the `[STRACE]` stderr markers via - # `strace_timing --rounds-table`; the L3 chip children emit their own - # markers (grouped by (pid, inv)), so multi-round works without any - # inline fd capture here. - for round_idx in range(rounds): - if round_idx > 0: - for name, initial in initial_tensors.items(): - getattr(test_args, name).copy_(initial) - - # See _run_and_validate_l2: the per-round masking is dead code - # under the existing upstream gate. Keep parity by passing through. - config = self._build_config( - config_dict, - enable_l2_swimlane=enable_l2_swimlane, - enable_dump_args=enable_dump_args, - enable_pmu=enable_pmu, - enable_dep_gen=enable_dep_gen, - enable_scope_stats=enable_scope_stats, - output_prefix=output_prefix, - ) + # Eager Worker.init() forked the chip/sub children before generate_args + # ran, so move test_args' host tensors into born-shared buffers the + # children can see. Golden was cloned above from the original host data; + # reset, dispatch, and compare below all operate on the rehosted views. + rehosted = _RehostedTaskArgs(worker, test_args) + try: + # Save initial tensor values for reset between rounds + all_tensor_names = test_args.tensor_names() + initial_tensors = {} + if rounds > 1: + for name in all_tensor_names: + initial_tensors[name] = getattr(test_args, name).clone() + + # Build CallableNamespace: compiled ChipCallables + sub callable IDs + ns = CallableNamespace({**compiled_callables, **sub_handles}) + + # Get orch function (plain function from CALLABLE) + orch_fn = self.CALLABLE["orchestration"] + + # Execute rounds. As for L2 (see _run_and_validate_l2), per-round + # timing is obtained offline from the `[STRACE]` stderr markers via + # `strace_timing --rounds-table`; the L3 chip children emit their own + # markers (grouped by (pid, inv)), so multi-round works without any + # inline fd capture here. + for round_idx in range(rounds): + if round_idx > 0: + for name, initial in initial_tensors.items(): + getattr(test_args, name).copy_(initial) + + # See _run_and_validate_l2: the per-round masking is dead code + # under the existing upstream gate. Keep parity by passing through. + config = self._build_config( + config_dict, + enable_l2_swimlane=enable_l2_swimlane, + enable_dump_args=enable_dump_args, + enable_pmu=enable_pmu, + enable_dep_gen=enable_dep_gen, + enable_scope_stats=enable_scope_stats, + output_prefix=output_prefix, + ) - # Orch fn signature: (orch, args, cfg) — inner fn forwards to - # the user's scene orch which takes (orch, callables, task_args, config). - def task_orch(orch, _args, _cfg, _ns=ns, _test_args=test_args, _config=config): - orch_fn(orch, _ns, _test_args, _config) + # Orch fn signature: (orch, args, cfg) — inner fn forwards to + # the user's scene orch which takes (orch, callables, task_args, config). + def task_orch(orch, _args, _cfg, _ns=ns, _test_args=test_args, _config=config): + orch_fn(orch, _ns, _test_args, _config) - with _temporary_env(self._resolve_env()): - worker.run(task_orch) + with _temporary_env(self._resolve_env()): + worker.run(task_orch) - if not skip_golden: - _compare_outputs(test_args, golden_args, all_tensor_names, self.RTOL, self.ATOL) + if not skip_golden: + _compare_outputs(test_args, golden_args, all_tensor_names, self.RTOL, self.ATOL) + finally: + rehosted.release() # ------------------------------------------------------------------ # pytest auto test method diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py b/tests/st/a2a3/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py index d591b67663..72e5edca0b 100644 --- a/tests/st/a2a3/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py +++ b/tests/st/a2a3/tensormap_and_ringbuffer/dynamic_register/test_dynamic_register.py @@ -158,10 +158,10 @@ def test_prepare_new_identity_after_start_then_run(st_platform, st_device_ids): config.block_dim = 3 config.aicpu_thread_num = 4 - # 1. Run pre_handle once to force _start_hierarchical (forks chip - # children, runs the CTRL_PREPARE prewarm loop). This puts the - # chip child into _run_chip_main_loop, the only state in which - # a CTRL_REGISTER broadcast can be ACKed. + # 1. Run pre_handle once against the snapshot-prepared chip callable. + # init() already forked the chip children and prepared the startup + # snapshot, so they are in _run_chip_main_loop — the only state in + # which the CTRL_REGISTER broadcast in step 2 can be ACKed. def orch_pre(o, _args, _cfg): o.submit_next_level(pre_handle, chip_args_pre, config) diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_host_buffer_registration.py b/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_host_buffer_registration.py index 9610eecd77..4b8fbca969 100644 --- a/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_host_buffer_registration.py +++ b/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_host_buffer_registration.py @@ -97,21 +97,17 @@ class TestPostForkHostBufferZeroCopy(SceneTestCase): {"name": "post_fork_zero_copy", "platforms": ["a2a3sim"]}, ] - def _force_fork(self, worker, chip_handle): - a = torch.full((SIZE,), 2.0, dtype=DTYPE).share_memory_() - b = torch.full((SIZE,), 3.0, dtype=DTYPE).share_memory_() - out = torch.zeros(SIZE, dtype=DTYPE).share_memory_() - worker.run(_one_task_orch(chip_handle, a, b, out), args=None, config=CallConfig()) - assert torch.allclose(out, _golden(a, b), rtol=self.RTOL, atol=self.ATOL) - def test_run(self, st_worker): """Zero-copy: buffers allocated AFTER the fork via ``create_host_buffer``, - filled in place, run, and read back — all without a per-run copy.""" + filled in place, run, and read back — all without a per-run copy. + + ``Worker.init()`` is eager, so the chip child is already forked when the + buffers are created; a born-shared ``create_host_buffer`` is the mapped + path a post-init host tensor takes to reach that child. + """ worker = st_worker chip_handle = type(self)._st_chip_handles["vector"] - self._force_fork(worker, chip_handle) - nbytes = SIZE * DTYPE.itemsize # element count × dtype size, not a magic 4 ba = worker.create_host_buffer(nbytes) bb = worker.create_host_buffer(nbytes) diff --git a/tests/ut/py/test_callable_identity.py b/tests/ut/py/test_callable_identity.py index 5ed5847607..2f17d96c6d 100644 --- a/tests/ut/py/test_callable_identity.py +++ b/tests/ut/py/test_callable_identity.py @@ -478,6 +478,23 @@ def __init__(self, *args): def add_remote_l3_socket(self, worker_id, *args): self.remote_worker_ids.append(worker_id) + # Eager init() forks the local L3 child and starts the C++ scheduler, so + # the mock must satisfy the register/init/orchestrator surface. + def add_sub_worker(self, *args): + pass + + def add_next_level_worker(self, *args): + pass + + def add_next_level_worker_at(self, *args): + pass + + def init(self): + pass + + def get_orchestrator(self): + return None + def close(self): self.closed = True @@ -1014,7 +1031,6 @@ def test_remote_sim_prepare_callable_control_roundtrip(): workers=[worker_id], ) worker.init() - worker._start_hierarchical() # noqa: SLF001 -- exercise PREPARE_CALLABLE before TASK dispatch. assert worker._worker is not None worker._worker.control_prepare(worker_id, handle.digest) @@ -1200,7 +1216,6 @@ def test_remote_sim_inner_python_import_register_runs_sub_task(): workers=[worker_id], ) worker.init() - worker._start_hierarchical() assert worker._worker is not None inner_digest = hashid_to_digest(_INNER_SUB_HASHID) target = b"tests.ut.py.test_callable_identity:_remote_inner_sub_noop" diff --git a/tests/ut/py/test_remote_l3_lifecycle.py b/tests/ut/py/test_remote_l3_lifecycle.py index 1e9b454c9d..12cd4f10ca 100644 --- a/tests/ut/py/test_remote_l3_lifecycle.py +++ b/tests/ut/py/test_remote_l3_lifecycle.py @@ -47,18 +47,26 @@ class FakePopen: pid = 12345 def __init__(self, *args, **kwargs): + self.terminated = False self.killed = False self.wait_calls = 0 def poll(self): return -9 if self.killed else None + def terminate(self): + self.terminated = True + def kill(self): self.killed = True def wait(self, timeout=None): + # Model a runner that does not exit on its own or on the cooperative + # SIGTERM, so cleanup must escalate to the killpg/kill backstop. self.wait_calls += 1 - return -9 if self.killed else 0 + raise remote_l3_worker.subprocess.TimeoutExpired( + cmd="runner", timeout=timeout if timeout is not None else 0.0 + ) fake_proc = FakePopen() @@ -74,6 +82,8 @@ def fake_read_ready(fd, timeout_s): with pytest.raises(TimeoutError): remote_l3_worker._start_session(_manifest()) + # Cooperative SIGTERM first, then the hard SIGKILL backstop. + assert fake_proc.terminated assert fake_proc.killed assert fake_proc.wait_calls >= 1 @@ -116,10 +126,7 @@ class FakeWorker: def __init__(self, *args, **kwargs): self.closed = False - def init(self): - pass - - def _start_hierarchical(self): + def init(self, *args, **kwargs): pass def close(self): diff --git a/tests/ut/py/test_scene_test_rehost.py b/tests/ut/py/test_scene_test_rehost.py new file mode 100644 index 0000000000..df2440ac7e --- /dev/null +++ b/tests/ut/py/test_scene_test_rehost.py @@ -0,0 +1,132 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""Device-free tests for the SceneTest host-tensor rehost adapter. + +Eager ``Worker.init()`` forks the L3 chip/sub children before ``generate_args`` +runs, so the harness rehosts each host tensor into a born-shared child buffer. +These tests fake the buffer allocator to check the adapter's transactional +contract without any worker/hardware: value/shape fidelity, LIFO release, +partial-construction rollback, and non-contiguous rejection. +""" + +from __future__ import annotations + +from multiprocessing.shared_memory import SharedMemory + +import pytest +import torch + +from simpler_setup.scene_test import TaskArgsBuilder, Tensor, _RehostedTaskArgs + + +class _FakeHostBuffer: + def __init__(self, nbytes: int): + self.shm = SharedMemory(create=True, size=nbytes) + self.buffer = self.shm.buf + + +class _FakeWorker: + """Stands in for a started L3 Worker's host-buffer allocator.""" + + def __init__(self, fail_on_create: int | None = None): + self.created: list[_FakeHostBuffer] = [] + self.freed: list[_FakeHostBuffer] = [] + self._fail_on_create = fail_on_create + + def create_host_buffer(self, nbytes: int) -> _FakeHostBuffer: + if self._fail_on_create is not None and len(self.created) >= self._fail_on_create: + raise RuntimeError("injected create_host_buffer failure") + buf = _FakeHostBuffer(nbytes) + self.created.append(buf) + return buf + + def free_host_buffer(self, buf: _FakeHostBuffer) -> None: + self.freed.append(buf) + buf.shm.close() + buf.shm.unlink() + + +def test_rehost_preserves_values_and_frees_lifo(): + ta = TaskArgsBuilder( + Tensor("a", torch.arange(4, dtype=torch.float32)), + Tensor("b", torch.zeros(4, dtype=torch.float32)), + ) + w = _FakeWorker() + rehosted = _RehostedTaskArgs(w, ta) + try: + # Values preserved and the builder now points at born-shared views. + assert torch.equal(ta.a, torch.arange(4, dtype=torch.float32)) + assert len(w.created) == 2 + # A write lands in the born-shared shm (what a child would read). + ta.b.copy_(torch.full((4,), 5.0)) + shared = w.created[1].buffer + assert shared is not None + # SharedMemory may round its mapping up to a page, so read only the + # tensor's own elements. + assert list(memoryview(shared).cast("f"))[:4] == [5.0] * 4 + finally: + rehosted.release() + # Both buffers freed, and the builder is restored to plain host tensors. + assert len(w.freed) == 2 + assert torch.equal(ta.a, torch.arange(4, dtype=torch.float32)) + + +def test_rehost_partial_failure_rolls_back(): + orig_a = torch.zeros(4, dtype=torch.float32) + orig_b = torch.ones(4, dtype=torch.float32) + ta = TaskArgsBuilder( + Tensor("a", orig_a), + Tensor("b", orig_b), + Tensor("c", torch.zeros(4, dtype=torch.float32)), + ) + w = _FakeWorker(fail_on_create=2) # third allocation fails + with pytest.raises(RuntimeError, match="injected create_host_buffer failure"): + _RehostedTaskArgs(w, ta) + # The two successfully-created buffers are freed, and the builder is + # restored to its original tensors (no half-rehosted state). + assert len(w.freed) == 2 + assert ta.a is orig_a + assert ta.b is orig_b + assert [s.value for s in ta.specs if isinstance(s, Tensor)][:2] == [orig_a, orig_b] + + +def test_rehost_rejects_aliased_tensors(): + base = torch.zeros(8, dtype=torch.float32) + ta = TaskArgsBuilder(Tensor("a", base[:4]), Tensor("b", base[2:6])) + w = _FakeWorker() + with pytest.raises(ValueError, match="alias overlapping storage"): + _RehostedTaskArgs(w, ta) + # Rejected before any allocation. + assert w.created == [] + assert w.freed == [] + + +def test_rehost_skips_empty_tensor(): + empty = torch.zeros(0, dtype=torch.float32) + ta = TaskArgsBuilder(Tensor("a", torch.arange(4, dtype=torch.float32)), Tensor("e", empty)) + w = _FakeWorker() + rehosted = _RehostedTaskArgs(w, ta) + try: + # Only the non-empty tensor is rehosted; the empty one is left untouched. + assert len(w.created) == 1 + assert ta.e is empty + finally: + rehosted.release() + + +def test_rehost_rejects_noncontiguous(): + noncontig = torch.zeros(4, 4, dtype=torch.float32)[:, ::2] + assert not noncontig.is_contiguous() + ta = TaskArgsBuilder(Tensor("a", noncontig)) + w = _FakeWorker() + with pytest.raises(ValueError, match="contiguous"): + _RehostedTaskArgs(w, ta) + # Rejected before any allocation — nothing to leak. + assert w.created == [] + assert w.freed == [] diff --git a/tests/ut/py/test_worker/test_host_addr_rewrite.py b/tests/ut/py/test_worker/test_host_addr_rewrite.py new file mode 100644 index 0000000000..fde974cc0b --- /dev/null +++ b/tests/ut/py/test_worker/test_host_addr_rewrite.py @@ -0,0 +1,73 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""Host-pointer blob rewrite must not touch child_memory (device) tensors. + +``_rewrite_blob_host_addrs`` redirects registered host pointers in a task-args +blob into a forked child's own mapping. A tensor whose ``buffer.addr`` is a +child-owned device pointer (``child_memory == 1``) carries a device VA, which +may numerically fall inside a registered host range yet must never be rewritten +— doing so corrupts the device pointer. The rewrite therefore keys on the +per-tensor ``child_memory`` flag, not on the numeric address alone. +""" + +import struct + +from _task_interface import TENSOR_CHILD_MEMORY_OFFSET, TENSOR_STRIDE_BYTES +from simpler.worker import _BLOB_HEADER_BYTES, _rewrite_blob_host_addrs + +_PARENT_LO = 0x7F00_0000_0000 +_PARENT_HI = 0x7F00_0010_0000 +_CHILD_BASE = 0x7E00_0000_0000 + + +def _make_blob(tensors: list[tuple[int, int]]) -> bytearray: + """Build a task-args blob: [int32 T][int32 S][Tensor*T]. + + ``tensors`` is a list of ``(addr, child_memory)``. Only the two fields the + rewrite reads (buffer.addr at offset 0, child_memory at its struct offset) + are populated; the rest of each 128-byte tensor stays zero. + """ + n = len(tensors) + buf = bytearray(_BLOB_HEADER_BYTES + n * TENSOR_STRIDE_BYTES) + struct.pack_into(" int: + off = _BLOB_HEADER_BYTES + i * TENSOR_STRIDE_BYTES + return struct.unpack_from("= 64 + finally: + buf.buffer.release() + w.free_host_buffer(buf) + finally: + w.close() + + def test_childless_l3_rejects_create_host_buffer(self): + w = Worker(level=3, num_sub_workers=0) + w.register(lambda args: None) + w.init() + try: + with pytest.raises(RuntimeError, match="at least one forked chip or sub child"): + w.create_host_buffer(64) + finally: + w.close() diff --git a/tests/ut/py/test_worker/test_host_worker.py b/tests/ut/py/test_worker/test_host_worker.py index 68a387d984..30d829a560 100644 --- a/tests/ut/py/test_worker/test_host_worker.py +++ b/tests/ut/py/test_worker/test_host_worker.py @@ -133,7 +133,7 @@ def init(self, device_id, bins, *, log_level, log_info_v, prewarm_config=None): def finalize(self) -> None: events.append(("finalize",)) - def fake_run_chip_main_loop(cw, *_args, chip_platform, chip_runtime): + def fake_run_chip_main_loop(cw, *_args, chip_platform, chip_runtime, prepared=None): events.append(("main_loop", cw, chip_platform, chip_runtime)) monkeypatch.setattr(worker_mod, "ChipWorker", FakeChipWorker) @@ -225,32 +225,33 @@ def test_close_releases_registered_callables(self): assert hw._identity_registry == {} assert hw._live_handles == {} - def test_prepare_python_fn_after_init_before_start_succeeds(self): - # init() allocates mailboxes but does not fork children. Python - # callables prepared in this window still land in the startup - # snapshot consumed by the first run(). - hw = Worker(level=3, num_sub_workers=0) + def test_register_python_fn_before_init_lands_in_snapshot(self): + # init() is eager, so there is no init-before-start window. A Python + # callable registered before init() is frozen into the startup snapshot + # and COW-inherited by the sub child during init — no run() needed. + hw = Worker(level=3, num_sub_workers=1) + handle = hw.register(lambda args: None) hw.init() try: - handle = hw.register(lambda args: None) assert _slot_for(hw, handle) in hw._callable_registry finally: hw.close() - def test_prepare_python_fn_after_init_before_start_does_not_broadcast(self): - class BroadcastTrap: - def broadcast_control_all(self, *args, **kwargs): - raise AssertionError("pre-start Python prepare must not broadcast") - + def test_pre_init_register_does_not_broadcast(self): + # A registration issued before init() takes the pre-start snapshot path + # and must not attempt a control broadcast — there is no started + # hierarchy to broadcast to yet. hw = Worker(level=3, num_sub_workers=1) + + def _trap(*_a, **_k): + raise AssertionError("pre-init Python register must not broadcast") + + hw._broadcast_py_control = _trap + handle = hw.register(lambda args: None) hw.init() - real_worker = hw._worker try: - hw._worker = BroadcastTrap() - handle = hw.register(lambda args: None) assert _slot_for(hw, handle) in hw._callable_registry finally: - hw._worker = real_worker hw.close() def test_prepare_python_fn_after_start_no_python_children_raises(self): @@ -373,73 +374,51 @@ def do_unregister(): hw._hierarchical_start_cv.wait = original_wait hw.close() - def test_prepare_blocks_startup_snapshot_from_not_started_window(self): - hw = Worker(level=3, num_sub_workers=0) + def test_register_during_initializing_waits_then_takes_post_start_path(self): + # A register that races an in-progress startup epoch (INITIALIZING) must + # block on the lifecycle condition rather than mutate the registry, then + # resolve via the post-start broadcast path once the epoch commits READY. + hw = Worker(level=3, num_sub_workers=1) hw.init() + try: + with hw._hierarchical_start_cv: + hw._hierarchical_start_state = "starting" - real_registry_lock = hw._registry_lock - register_waiting = threading.Event() - release_register = threading.Event() - startup_snapshot_attempted = threading.Event() - result: list[CallableHandle] = [] - errors: list[BaseException] = [] - - class BlockingRegistryLock: - def __enter__(self): - thread_name = threading.current_thread().name - if thread_name == "register-thread": - register_waiting.set() - if not release_register.wait(timeout=2.0): - raise TimeoutError("test timed out waiting to release register") - elif thread_name == "startup-thread": - startup_snapshot_attempted.set() - return real_registry_lock.__enter__() - - def __exit__(self, exc_type, exc, tb): - return real_registry_lock.__exit__(exc_type, exc, tb) - - def locked(self): - return real_registry_lock.locked() - - hw._registry_lock = BlockingRegistryLock() + errors: list[BaseException] = [] + result: list[object] = [] + wait_entered = threading.Event() + original_wait = hw._hierarchical_start_cv.wait - def do_register(): - try: - result.append(hw.register(lambda args: None)) - except BaseException as exc: # noqa: BLE001 - errors.append(exc) + def wait_with_signal(timeout=None): + wait_entered.set() + return original_wait(timeout) - def do_startup(): - try: - hw._start_hierarchical() - except BaseException as exc: # noqa: BLE001 - errors.append(exc) + hw._hierarchical_start_cv.wait = wait_with_signal - register_thread = threading.Thread(target=do_register, name="register-thread") - startup_thread = threading.Thread(target=do_startup, name="startup-thread") - try: - register_thread.start() - assert register_waiting.wait(timeout=2.0) + def do_register(): + try: + result.append(hw.register(lambda args: None)) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) - startup_thread.start() - assert not startup_snapshot_attempted.wait(timeout=0.2) + t = threading.Thread(target=do_register) + t.start() + assert wait_entered.wait(timeout=2.0) + # Still INITIALIZING: the register must be parked, not installed. + assert len(hw._identity_registry) == 0 - release_register.set() - register_thread.join(timeout=2.0) - startup_thread.join(timeout=2.0) + with hw._hierarchical_start_cv: + hw._hierarchical_started = True + hw._hierarchical_start_state = "started" + hw._hierarchical_start_cv.notify_all() + t.join(timeout=2.0) - assert not register_thread.is_alive() - assert not startup_thread.is_alive() + assert not t.is_alive() assert errors == [] assert len(result) == 1 - assert _slot_for(hw, result[0]) == 0 - assert startup_snapshot_attempted.is_set() - assert hw._hierarchical_start_state == "started" finally: - release_register.set() - register_thread.join(timeout=2.0) - startup_thread.join(timeout=2.0) - hw._registry_lock = real_registry_lock + if "original_wait" in locals(): + hw._hierarchical_start_cv.wait = original_wait hw.close() def test_prepare_chip_callable_after_init_no_chips_succeeds(self): diff --git a/tests/ut/py/test_worker/test_startup_readiness.py b/tests/ut/py/test_worker/test_startup_readiness.py index a9c51608f8..d8713ea7a3 100644 --- a/tests/ut/py/test_worker/test_startup_readiness.py +++ b/tests/ut/py/test_worker/test_startup_readiness.py @@ -32,6 +32,7 @@ import os import signal import struct +import threading import time from contextlib import contextmanager from multiprocessing.shared_memory import SharedMemory @@ -67,6 +68,15 @@ def _handler(_signum, _frame): signal.signal(signal.SIGALRM, old) +def _run_catch(fn): + """Run ``fn`` in a thread body, returning None on success or the exception.""" + try: + fn() + return None + except BaseException as e: # noqa: BLE001 + return e + + def _make_shared_counter(): shm = SharedMemory(create=True, size=4) buf = shm.buf @@ -127,12 +137,11 @@ def test_inner_init_failure_raises_bounded_error(self): 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) + w4.init() assert time.monotonic() - start < _TEST_WALL_BUDGET_S finally: w4.close() @@ -145,11 +154,10 @@ def test_inner_exit_before_ready_raises(self): 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) + w4.init() finally: w4.close() @@ -161,12 +169,11 @@ def test_startup_deadline_fires_on_hung_child(self): 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) + w4.init() elapsed = time.monotonic() - start assert 1.5 <= elapsed < _TEST_WALL_BUDGET_S finally: @@ -189,11 +196,10 @@ def spy_abort(self): 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) + w4.init() assert "pids" in captured and captured["pids"], "rollback did not run" for pid in captured["pids"]: @@ -223,11 +229,10 @@ def test_ready_sibling_closed_gracefully_on_sibling_failure(self): 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) + w4.init() report = w4._last_rollback assert report is not None @@ -249,11 +254,10 @@ def test_second_child_failure_reaps_first(self): 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) + w4.init() assert w4._next_level_pids == [] finally: w4.close() @@ -277,11 +281,10 @@ def _boom(*_a, **_k): 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) + w3.init() finally: w3.close() @@ -302,9 +305,9 @@ 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. + startup and that tasks dispatch afterwards. init() is eager and recursive, so + a next-level child's INIT_READY means its whole subtree (its own sub / chip + children) came up during the parent's init(), not on first run. """ def test_l4_l3_tree_comes_up_and_runs(self): @@ -378,6 +381,346 @@ def l4_orch(orch, args, config): b_shm.unlink() +class TestEagerInitContract: + """init() is the single, eager startup point for level >= 3. + + Children come up during init(); run() / the other former lazy triggers no + longer start the hierarchy, and init() without any run() closes cleanly. + """ + + def test_l3_sub_children_ready_after_init_no_run(self): + hw = Worker(level=3, num_sub_workers=2) + hw.register(lambda args: None) + hw.init() + try: + assert hw._hierarchical_started is True + assert len(hw._sub_pids) == 2 + # Every sub child is forked and alive (READY) before any run(). + for pid in hw._sub_pids: + assert os.waitpid(pid, os.WNOHANG) == (0, 0) + finally: + hw.close() + + def test_init_then_close_without_run(self): + hw = Worker(level=3, num_sub_workers=1) + hw.register(lambda args: None) + hw.init() + hw.close() + assert hw._sub_pids == [] + assert hw._worker is None + + def test_l4_l3_sub_subtree_ready_after_init_no_run(self): + l3 = _l3_child(num_sub_workers=1) + w4 = Worker(level=4, num_sub_workers=0, startup_timeout_s=30.0) + w4.register(_trivial_orch) + w4.add_worker(l3) + w4.init() + try: + # The L4's direct L3 child is READY, which — because init() is + # recursive — means the L3 already forked and readied its own sub + # grandchild before publishing INIT_READY. + assert w4._hierarchical_started is True + assert len(w4._next_level_pids) == 1 + assert os.waitpid(w4._next_level_pids[0], os.WNOHANG) == (0, 0) + finally: + w4.close() + + def test_run_does_not_start_hierarchy(self, monkeypatch): + counter_shm, counter_buf = _make_shared_counter() + hw = Worker(level=3, num_sub_workers=1) + sub = hw.register(lambda args: _increment_counter(counter_buf)) + hw.init() + + orig = Worker._start_hierarchical + calls = {"n": 0} + + def spy(self): + calls["n"] += 1 + return orig(self) + + monkeypatch.setattr(Worker, "_start_hierarchical", spy) + try: + + def orch(o, a, c): + o.submit_sub(sub) + + hw.run(orch) + assert calls["n"] == 0 + assert _read_counter(counter_buf) == 1 + finally: + hw.close() + counter_shm.close() + counter_shm.unlink() + + +class TestSubtreeCancellation: + """§4.6 cancellation domain: a mid-init subtree is deterministically reaped. + + A next-level child is a process-group leader whose forked descendants + inherit the group, so the startup root's rollback reaps the whole subtree + (the child plus its grandchildren) with killpg rather than leaking orphans + to the multiprocessing resource_tracker. + """ + + def test_stuck_midinit_subtree_killpg_reaps_grandchild(self, monkeypatch): + import simpler.worker as worker_mod # noqa: PLC0415 + + # Shrink the cooperative-cleanup grace so the stuck survivor hits the + # killpg backstop quickly. + monkeypatch.setattr(worker_mod, "_ROLLBACK_GRACEFUL_TIMEOUT_S", 1.0) + + gshm = SharedMemory(create=True, size=8) + gbuf = gshm.buf + assert gbuf is not None + struct.pack_into("q", gbuf, 0, 0) + + def _fork_grandchild_then_ignore_cancel(*_a, _gbuf=gbuf, **_k): + pid = os.fork() + if pid == 0: + # Grandchild: inherits the L3 child's process group (no setpgid). + while True: + try: + time.sleep(3600) + except BaseException: # noqa: BLE001, PERF203 + pass + struct.pack_into("q", _gbuf, 0, pid) + # Swallow the cooperative cancel so this subtree becomes a stuck + # survivor only killpg can reap. + while True: + try: + time.sleep(3600) + except BaseException: # noqa: BLE001, PERF203 + pass + + l3 = _l3_child() + l3.init = _fork_grandchild_then_ignore_cancel # noqa: SLF001 + + w4 = Worker(level=4, num_sub_workers=0, startup_timeout_s=1.0) + w4.register(_trivial_orch) + w4.add_worker(l3) + try: + with _hard_timeout(_TEST_WALL_BUDGET_S): + with pytest.raises(RuntimeError, match="deadline"): + w4.init() + + gpid = struct.unpack_from("q", gbuf, 0)[0] + assert gpid > 0, "grandchild was never forked" + # killpg reaped the whole process group: the grandchild is gone. + deadline = time.monotonic() + 5.0 + alive = True + while time.monotonic() < deadline: + try: + os.kill(gpid, 0) + except ProcessLookupError: + alive = False + break + time.sleep(0.05) + assert not alive, "grandchild survived — killpg backstop did not reap the subtree" + finally: + w4.close() + gshm.close() + gshm.unlink() + + +class TestApiLinearizationDuringInit: + """§4.4/§5.2: init() is one atomic epoch; concurrent API calls linearize. + + Each test pauses init() inside _start_hierarchical (state == "starting") and + drives a second API call to observe the INITIALIZING behavior. + """ + + @staticmethod + def _pause_start(entered, release): + orig = Worker._start_hierarchical + + def slow(self): + entered.set() + if not release.wait(timeout=10.0): + raise TimeoutError("start not released") + return orig(self) + + return slow + + def test_register_blocks_during_initializing_then_completes(self, monkeypatch): + entered = threading.Event() + release = threading.Event() + monkeypatch.setattr(Worker, "_start_hierarchical", self._pause_start(entered, release)) + + w = Worker(level=3, num_sub_workers=1, startup_timeout_s=30.0) + w.register(lambda args: None) + init_err: list = [] + it = threading.Thread(target=lambda: init_err.append(_run_catch(w.init))) + it.start() + try: + with _hard_timeout(_TEST_WALL_BUDGET_S): + assert entered.wait(3.0) + reg_out: list[object] = [] + reg_started = threading.Event() + + def do_reg(): + reg_started.set() + reg_out.append(w.register(lambda args: None)) + + rt = threading.Thread(target=do_reg) + rt.start() + assert reg_started.wait(3.0) + time.sleep(0.3) + assert reg_out == [], "register must block while INITIALIZING" + release.set() + it.join(10.0) + rt.join(10.0) + assert init_err == [None] + assert len(reg_out) == 1 + finally: + release.set() + it.join(5.0) + w.close() + + def test_init_failure_wakes_register_waiter_with_startup_error(self, monkeypatch): + entered = threading.Event() + release = threading.Event() + + def boom(_self): + entered.set() + release.wait(timeout=10.0) + raise RuntimeError("injected start failure") + + monkeypatch.setattr(Worker, "_start_hierarchical", boom) + + w = Worker(level=3, num_sub_workers=1, startup_timeout_s=30.0) + w.register(lambda args: None) + init_err: list = [] + it = threading.Thread(target=lambda: init_err.append(_run_catch(w.init))) + it.start() + try: + with _hard_timeout(_TEST_WALL_BUDGET_S): + assert entered.wait(3.0) + reg_err: list = [] + reg_started = threading.Event() + + def do_reg(): + reg_started.set() + reg_err.append(_run_catch(lambda: w.register(lambda args: None))) + + rt = threading.Thread(target=do_reg) + rt.start() + assert reg_started.wait(3.0) + time.sleep(0.2) + release.set() + it.join(10.0) + rt.join(10.0) + assert any("injected start failure" in str(e) for e in init_err) + assert any(e is not None and "startup failed" in str(e) for e in reg_err) + finally: + release.set() + it.join(5.0) + w.close() + + def test_add_worker_rejected_during_initializing(self, monkeypatch): + entered = threading.Event() + release = threading.Event() + monkeypatch.setattr(Worker, "_start_hierarchical", self._pause_start(entered, release)) + + w4 = Worker(level=4, num_sub_workers=0, startup_timeout_s=30.0) + w4.register(_trivial_orch) + w4.add_worker(_l3_child()) + it = threading.Thread(target=w4.init) + it.start() + try: + with _hard_timeout(_TEST_WALL_BUDGET_S): + assert entered.wait(3.0) + with pytest.raises(RuntimeError, match="before init"): + w4.add_worker(_l3_child()) + release.set() + it.join(10.0) + finally: + release.set() + it.join(5.0) + w4.close() + + def test_close_waits_for_in_progress_init(self, monkeypatch): + entered = threading.Event() + release = threading.Event() + monkeypatch.setattr(Worker, "_start_hierarchical", self._pause_start(entered, release)) + + w = Worker(level=3, num_sub_workers=1, startup_timeout_s=30.0) + w.register(lambda args: None) + it = threading.Thread(target=w.init) + it.start() + try: + with _hard_timeout(_TEST_WALL_BUDGET_S): + assert entered.wait(3.0) + closed = threading.Event() + ct = threading.Thread(target=lambda: (w.close(), closed.set())) + ct.start() + # close() must not no-op while init() is still in progress. + assert not closed.wait(0.3) + release.set() + ct.join(10.0) + it.join(10.0) + assert closed.is_set() + assert w._sub_pids == [] + finally: + release.set() + it.join(5.0) + + +class _FakeChipOk: + """Stand-in for a ChipWorker whose init succeeds — no NPU touched.""" + + def init(self, *_a, **_k): + pass + + def _register_callable_at_slot(self, *_a, **_k): # pragma: no cover + pass + + def finalize(self): + pass + + +class TestLevel2Lifecycle: + """An L2 worker's init()/close() must not deadlock on the level>=3-only + epoch state machine. + + Regression: init() left the L2 worker's lifecycle state at "starting" (only + level>=3 committed "started"), so close()'s wait-out-"starting" hung forever + — which timed out the first L2 test of every sim / onboard suite. + """ + + def _make_l2(self, monkeypatch): + import simpler.worker as worker_mod # noqa: PLC0415 + + import simpler_setup.runtime_builder as rb_mod # noqa: PLC0415 + + class _FakeBuilder: + def __init__(self, *_a, **_k): + pass + + def get_binaries(self, *_a, **_k): + return object() + + # Mock the device/runtime layer so this stays device-free (no sim + # binaries required) — the regression is purely the lifecycle state. + monkeypatch.setattr(worker_mod, "ChipWorker", _FakeChipOk) + monkeypatch.setattr(rb_mod, "RuntimeBuilder", _FakeBuilder) + return Worker(level=2, device_id=0, platform="a2a3sim", runtime="tensormap_and_ringbuffer") + + def test_l2_init_then_close_does_not_hang(self, monkeypatch): + w = self._make_l2(monkeypatch) + with _hard_timeout(_TEST_WALL_BUDGET_S): + w.init() + assert w._initialized is True + w.close() + assert w._initialized is False + # A second close is a clean no-op (does not re-block on the epoch cv). + w.close() + + def test_l2_close_without_init_is_noop(self, monkeypatch): + w = self._make_l2(monkeypatch) + with _hard_timeout(_TEST_WALL_BUDGET_S): + w.close() + + class _FakeChipRaises: """Stand-in for ChipWorker whose init raises — no NPU touched.""" @@ -434,11 +777,10 @@ def test_chip_init_failure_raises_bounded(self, monkeypatch): 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) + l3.init() finally: l3.close() @@ -448,12 +790,11 @@ def test_chip_init_hang_trips_deadline(self, monkeypatch): 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) + l3.init() assert 1.5 <= time.monotonic() - start < _TEST_WALL_BUDGET_S finally: l3.close()