Add: eager, transactional, recursive Worker.init() for L3+

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

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

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

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

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<int>(sizeof(Tensor));
+    m.attr("TENSOR_CHILD_MEMORY_OFFSET") = static_cast<int>(offsetof(Tensor, child_memory));
 
     // --- Tensor ---
     // The unified strided tensor descriptor. Constructed contiguous via make()

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)

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)