Refactor: isolate hierarchical worker run state

docs/comm-domain.md

@@ -58,26 +58,26 @@ The handle is a context manager. Its lifecycle has **two distinct states**:
   exits). Further indexing (`handle[i]`) raises. This is the *user-visible*
   state: "do not hand this domain to any new `submit_*`."
 - **`freed`** — the backend `comm_release_domain_windows` has actually run and
-  the device memory is gone. This happens **after** `Worker.run` drains the
-  DAG, never inside the `with` block.
+  the device memory is gone. This happens **after** the owning run's completion
+  fence, never inside the `with` block.
 
 This split exists because `submit_next_level()` only *enqueues* DAG work;
-`Worker.run()` does not drain until the orch function returns. If `release()`
-freed memory immediately on `with`-exit, a still-queued task that captured the
-domain's `device_ctx` / `buffer_ptrs` would read freed memory. So **release is
-deferred**: `release()` flips `released` and queues the backend free; the real
-free runs after drain, when every task that could reference the window has
-completed.
+`Worker.run()` does not wait for completion until the orch function returns.
+If `release()` freed memory immediately on `with`-exit, a still-queued task that
+captured the domain's `device_ctx` / `buffer_ptrs` would read freed memory. So
+**release is deferred**: `release()` flips `released` and queues the backend
+free; the real free runs after the run fence, when every task that could
+reference the window has completed.
 
 Mental model: like `with open(f) as fh: ...` — the user-visible close is
 lexical (end of block), the physical teardown is managed for you. Use
 `handle.released` to guard against accidental reuse; use `handle.freed` only if
 you must assert physical teardown.
 
-Cleanup is **drain-safe**: even if a chip task fails and `drain()` re-raises,
-`Worker.run` still executes the pending releases and sweeps any live domains the
-orch fn forgot to release (LIFO), so a failed run cannot strand backend
-allocations into the next run.
+Cleanup is **failure-safe**: even if a chip task fails and the run wait
+re-raises, `Worker.run` still executes pending releases and sweeps any live
+domains the orch fn forgot to release (LIFO), so a failed run cannot strand
+backend allocations into the next run.
 
 ---
 
@@ -94,7 +94,7 @@ called many times:
 
 - the **base communicator is created once** and reused — it is *not* rebuilt
   per `run` or per domain;
-- only the **per-domain windows** are allocated (and freed after drain) on each
+- only the **per-domain windows** are allocated (and freed after the run fence) on each
   `allocate_domain` / `run`. Each allocation gets a fresh `allocation_id` so
   concurrent or sequential domains never collide on IPC handshake / barrier
   names.

docs/hierarchical_level_runtime.md

@@ -157,7 +157,7 @@ what flows through `ChipWorker::run`.
                  │                                 │   success → COMPLETED
                  │                                 │   failure → FAILED + poison downstream
                  │                                 │   try_consume → ring release
-                 │ drain() ◄── notify when all done │
+                 │ wait_run(id) ◄── run becomes terminal
 ```
 
 Communication channels:
@@ -212,9 +212,9 @@ When L4's `WorkerThread` writes a task frame to the L3 child's mailbox, the
 frame carries the callable hash digest plus `config` and `args_blob`. The child
 loop reads the digest, resolves it through its local identity table to a private
 orch-function slot, and calls `inner_worker.run(orch_fn, args, cfg)`. The inner
-Worker opens its own scope, executes the orch function with its own
-`Orchestrator`, and drains. Each level's orch fn receives its own Orchestrator
-— recursion is symmetric.
+Worker opens its own run and scope, executes the orch function with its own
+`Orchestrator`, closes submission, and waits for that run. Each level's orch fn
+receives its own Orchestrator — recursion is symmetric.
 
 **Nested fork ordering**: L3's own children (sub/chip) are forked **inside**
 the L4 child process, in L3's eager `init()` (which the L4 child runs before it

docs/orchestrator.md

@@ -7,9 +7,9 @@ internals, not public submit arguments. See
 [callable-identity-registration.md](callable-identity-registration.md).
 
 The Orchestrator is the **DAG builder**. It runs single-threaded on the user's
-thread (inside `Worker::run` between `scope_begin` and `drain`) and owns the
-three data structures that turn a sequence of `submit_*` calls into a scheduled
-DAG: `Ring`, `TensorMap`, and `Scope`.
+thread while a run is open for submission and owns the three data structures
+that turn a sequence of `submit_*` calls into a scheduled DAG: `Ring`,
+`TensorMap`, and `Scope`.
 
 For the high-level role of the Orchestrator among the three engine components,
 see [hierarchical_level_runtime.md](hierarchical_level_runtime.md). For what
@@ -51,14 +51,18 @@ public:
     // --- Intermediate-buffer allocation (runtime-owned lifetime) ---
     Tensor alloc(const std::vector<uint32_t> &shape, DataType dtype);
 
-    // --- Internal lifecycle (invoked by Worker::run only, bound as _scope_begin
-    //     / _scope_end / _drain in the Python facade) ---
+    // --- Internal lifecycle (invoked by Worker::run only) ---
+    RunId begin_run();
+    void close_run_submission(RunId run_id);
+    void fail_run_submission(RunId run_id, std::exception_ptr error);
+    void wait_run(RunId run_id);
+    bool run_done(RunId run_id) const;
+    void release_run(RunId run_id);
     void scope_begin();
     void scope_end();
-    void drain();
 
 private:
-    // ... components: Ring, TensorMap, Scope, slot pool, active_tasks_ counter
+    // ... components: Ring, TensorMap, Scope, and the RunState registry
 };
 
 struct SubmitResult { TaskSlot task_slot; };  // internal only; not bound to Python
@@ -67,9 +71,10 @@ struct SubmitResult { TaskSlot task_slot; };  // internal only; not bound to Pyt
 **Status**: `submit_sub` takes only `(CallableIdentity, args)` — no
 `config`, since SUB has no per-call config.
 
-`scope_begin` / `scope_end` / `drain` are invoked from Python `Worker.run` via
-`_scope_begin` / `_scope_end` / `_drain` bindings. They are not part of the
-user-facing orch-fn API.
+The run lifecycle and outer `scope_begin` / `scope_end` are invoked from Python
+`Worker.run` through private bindings. They are not part of the user-facing
+orch-fn API. `Worker.run` remains blocking: it closes submission, waits for the
+matching run fence, performs run-owned cleanup, and returns `None`.
 
 Remote L3 submit adds two hidden pieces of metadata: final eligible worker-id
 sets and optional `RemoteTaskArgsSidecar` entries aligned by tensor index.
@@ -161,9 +166,9 @@ small POD copied by value. `callable` is a `uint64_t` opaque handle (see
 **Step 3 — tag walk**: The only place tags are consumed. After this step tags
 are never inspected again; they are not carried into the slot's stored
 `task_args` value during dispatch (see [task-flow.md](task-flow.md) §3).
-Local tensors key TensorMap by `(LOCAL_HOST, ptr)` or
-`(LOCAL_CHILD, worker, ptr)`. Remote tensors with sidecars key by
-`(address_kind, owner_worker_id, buffer_id, generation, offset)`.
+Every TensorMap key starts with the current `RunId`. Local tensor identity is
+then `(LOCAL_HOST, ptr)` or `(LOCAL_CHILD, worker, ptr)`. Remote tensors with
+sidecars use `(address_kind, owner_worker_id, buffer_id, generation, offset)`.
 
 | Tag | `tensormap.lookup` | `tensormap.insert` |
 | --- | ------------------ | ------------------ |
@@ -276,8 +281,7 @@ SubmitResult Orchestrator::submit_sub(const CallableIdentity &callable, TaskArgs
    state lives in parent-process heap (never crossed into child workers),
    so the ring-index addressing scheme L2 needs for shmem descriptors
    buys us nothing here. A monotonic `int32_t` gives ~2 billion ids per
-   `reset_to_empty()` interval, reset to 0 at the end of every
-   `Worker.run()`.
+   globally quiescent compaction interval.
 2. **`MAX_RING_DEPTH = 4` independent shared-memory heap slabs**
    (Strict-1; matches L2's `PTO2_MAX_RING_DEPTH`). Each slab has its own
    `mmap(MAP_SHARED | MAP_ANONYMOUS)` region, bump cursor, FIFO
@@ -299,7 +303,7 @@ SubmitResult Orchestrator::submit_sub(const CallableIdentity &callable, TaskArgs
    records its `ring_idx` and `ring_slot_idx` (position within that
    ring's FIFO order). `std::deque::push_back` never invalidates pointers
    to existing elements, so the pointer returned by `slot_state(id)`
-   stays valid until `reset_to_empty()` drops the whole deque.
+   stays valid until globally quiescent `reset_to_empty()` drops the deque.
 
 ```cpp
 struct AllocResult {
@@ -351,12 +355,13 @@ next-oldest in-ring slot is released, walking the ring's `heap_tail`
 forward. Rings never touch each other — inner-scope tasks reclaim
 without waiting for an outer-scope task to finish.
 
-**End-of-run reset**: `Orchestrator::drain()` waits for
-`active_tasks_` to hit 0, then calls `ring.reset_to_empty()` which
-drops the whole slot-state deque *and* rewinds every ring's cursors /
-`released[]` / `slot_heap_end[]` back to 0. Memory per `Worker.run()`
-is bounded by that run's peak alive task count; nothing accumulates
-across runs.
+**Run completion and compaction**: every committed slot increments its owning
+`RunState.active_tasks`; `on_consumed` decrements that same run exactly once. A
+run becomes terminal only after submission is closed and its count reaches
+zero. Slot and heap reclamation still happens incrementally through
+`ring.release(slot)`. `release_run()` may call `ring.reset_to_empty()` only
+when no registered runs or live slots remain, so one run never resets storage
+owned by another.
 
 **Locking**: each ring has its own `mu` / `cv`; the shared
 `next_task_id_` and slot deque are guarded by a separate `slots_mu_`.
@@ -457,29 +462,28 @@ threshold transitions to CONSUMED inline; others stay COMPLETED or PENDING
 until the scheduler and consumers finish their own releases. This mirrors
 L2's `pto2_scope_end`.
 
-Users who need a synchronous wait for *all* in-flight tasks must call
-`drain()` (or let `Worker::run` finish — its outer `scope_end` is followed
-by `drain()` before the call returns). There is deliberately no
-per-scope drain primitive: the extra machinery (per-scope active counter
-and cv) would only pay for itself in patterns we do not have yet.
+The internal run fence, not `scope_end`, provides synchronous completion.
+`Worker.run` closes its outer scope, closes submission, and waits for that
+run's active count to reach zero before returning. There is deliberately no
+per-scope wait primitive.
 
 ---
 
 ## 7. TensorMap
 
-The TensorMap maps `tensor_base_ptr → current_producer_slot`. It drives
-automatic dependency inference.
+The TensorMap maps `(RunId, TensorKey) → current_producer_slot`. It drives
+automatic dependency inference without resolving a producer from another run.
 
 ```cpp
 class TensorMap {
 public:
-    TaskSlot lookup(uint64_t base_ptr) const;         // returns INVALID if absent
-    void     insert(uint64_t base_ptr, TaskSlot sid); // overwrites; previous
-                                                      // producer remains wire-referenced
-    void     erase(uint64_t base_ptr);                // called when producer
-                                                      // reaches CONSUMED
+    TaskSlot lookup(RunId run_id, TensorKey key) const;
+    void insert(RunId run_id, TensorKey key, TaskSlot sid);
+    void erase_task_outputs(RunId run_id,
+                            const std::vector<TensorKey> &keys);
 private:
-    std::unordered_map<uint64_t, TaskSlot> map_;
+    std::mutex mu_;
+    std::unordered_map<RunTensorKey, TaskSlot, RunTensorKeyHash> map_;
 };
 ```
 
@@ -503,11 +507,9 @@ private:
 
 ### Thread safety
 
-TensorMap is written only by the Orch thread (in `submit_*`) and modified by
-the Scheduler thread via `erase` (on CONSUMED). Since `submit_*` and `erase`
-for different entries are non-overlapping in practice, a single mutex guards
-the map in the current implementation. If contention becomes a concern, a
-concurrent hash map can replace it.
+TensorMap is written by the Orch thread in `submit_*` and modified by the
+Scheduler thread when a slot becomes CONSUMED. A mutex serializes lookup,
+insert, erase, and size operations across those threads.
 
 ---
 
@@ -577,7 +579,8 @@ Tensor Orchestrator::alloc(const std::vector<uint32_t> &shape, DataType dtype) {
     // 2. Register as this slot's output so downstream tensors with the same
     //    data pointer look up this slot as producer.
     uint64_t key = reinterpret_cast<uint64_t>(ar.heap_ptr);
-    tensormap_.insert(key, ar.slot);
+    s.run_id = current_run_id;
+    tensormap_.insert(current_run_id, key, ar.slot);
     s.output_keys.push_back(key);
     // 3. No fanin — alloc has no work to wait on.
     s.fanin_count = 0;
@@ -591,7 +594,7 @@ Tensor Orchestrator::alloc(const std::vector<uint32_t> &shape, DataType dtype) {
     s.fanout_released = 1;
     // 6. Straight to COMPLETED — no dispatch needed.
     s.state = TaskState::COMPLETED;
-    active_tasks_++;
+    current_run.active_tasks++;
     return Tensor{key, shape, dtype};
 }
 ```
@@ -695,9 +698,9 @@ instead of stalling forever. Default timeout: 10 s.
 
 ## 9. Invariants
 
-1. **Orch is single-threaded**: only one thread ever calls `submit_*` or holds
-   the `Orchestrator`. No locking is needed on TensorMap, Scope, or Ring-head
-   for self-writes.
+1. **Orch is single-threaded**: only one thread builds a run and calls
+   `submit_*` at a time. TensorMap still uses a mutex because Scheduler-driven
+   consumption erases entries concurrently.
 2. **Tags are consumed at submit**: `task_args.tag(i)` is read only inside
    `submit_*`. Phases after submit (slot storage, dispatch, execution) do not
    see tags.

docs/remote-l3-worker-design.md

@@ -143,7 +143,7 @@ Relevant code paths:
   - Tags are stripped after submit.
 - `docs/comm-domain.md`
   - Dynamic communication domains already model deferred release after
-    `drain()`.
+    their owning run's completion fence.
 
 The Scheduler should not inspect transport details, but it does need enough
 metadata to avoid dispatching a task to an endpoint that cannot run it. Remote
@@ -178,11 +178,10 @@ the endpoint, reports endpoint errors, and notifies the Scheduler with an
 explicit success/failure outcome.
 
 Directed ready queues, exact group dispatch, fanin/fanout, and ring release
-remain in the common runtime. The first-error-wins policy remains only as the error
-reporting policy for choosing which root error `drain()` raises. The important
-change is that completion is no longer implicitly success; every endpoint,
-including `LocalMailboxEndpoint`, must report an explicit success/failure
-outcome.
+remain in the common runtime. First-error-wins is scoped to each `RunState` and
+chooses which root error that run's wait raises. Completion is not implicitly
+success; every endpoint, including `LocalMailboxEndpoint`, reports an explicit
+success/failure outcome.
 
 ## Fork-Safe Remote Process Model
 
@@ -517,17 +516,17 @@ Required parent-side behavior:
   `task_failure` instead of reporting a successful completion.
 - Non-zero task or endpoint errors become candidates for the worker's first
   reported error.
-- The worker still notifies the Scheduler so `drain()` cannot hang.
+- The worker still notifies the Scheduler so the originating run can finish.
 - The notification carries an outcome: success, task failure, or endpoint
   failure.
 - Failed slots transition to a failed/poisoned state rather than successful
   `COMPLETED`.
 - Downstream consumers of a failed producer are marked failed/skipped and are
   not dispatched.
-- `drain()` waits for bookkeeping and cleanup, then raises the first root
-  error with remote host, worker id, hashid, and sequence in the message.
+- The run fence waits for bookkeeping and cleanup, then raises that run's first
+  root error with remote host, worker id, hashid, and sequence in the message.
 
-Local mailbox dispatch keeps first-error-wins only for final error reporting.
+Local mailbox dispatch keeps per-run first-error-wins only for final reporting.
 It must not mark a failed child dispatch as successful `COMPLETED`. The remote
 buffer path and the local adapter both use the same poisoned dependency
 propagation before dependent tasks are exposed to failed producer outputs.

docs/scheduler.md

@@ -66,9 +66,10 @@ while (true) {
 }
 ```
 
-`loop_mutex` covers completion and dispatch slot access. `Orchestrator::drain`
-uses the same mutex while releasing/resetting slots, preventing slot reuse
-while the Scheduler still holds a reference.
+`loop_mutex` covers completion and dispatch slot access.
+`Orchestrator::release_run` uses the same mutex during optional globally
+quiescent compaction, preventing slot removal while the Scheduler holds a
+reference.
 
 ## 3. Directed NEXT_LEVEL dispatch
 

docs/task-flow.md

@@ -369,7 +369,7 @@ The mailbox layout, fork ordering, and child loop are in
 
 | Region | Lives in | Used by | Lifetime |
 | ------ | -------- | ------- | -------- |
-| `Ring` slot-state pool (`std::deque<unique_ptr<TaskSlotState>>`) | parent heap | Orchestrator, Scheduler, WorkerThread parent side | monotonic task-id; reset at `Worker.run` drain |
+| `Ring` slot-state pool (`std::deque<unique_ptr<TaskSlotState>>`) | parent heap | Orchestrator, Scheduler, WorkerThread parent side | monotonic task-id; compacted only when globally quiescent |
 | `slot.task_args` (single) or `task_args_list[N]` (group, vector-backed) | parent heap | same | until slot reaches CONSUMED |
 | per-WT mailbox | shm MAP_SHARED | parent WorkerThread writes, child reads | lifetime of WorkerThread |
 | **HeapRing[0..3]** (user OUTPUT auto-alloc + `orch.alloc`) | **4 separate shm MAP_SHARED mmaps**, one per scope-layer ring | output to user code; inherited by forked children | per-ring FIFO via `rings_[r].last_alive`; scope depth picks the ring |
@@ -378,9 +378,9 @@ The mailbox layout, fork ordering, and child loop are in
 
 Slot state lives inside `Ring` as `std::deque<std::unique_ptr<…>>` so
 `push_back` never invalidates pointers to live slots.
-`ring.slot_state(id)` hands out a stable pointer for every live slot;
-`drain()` calls `ring.reset_to_empty()` to drop all slot state at the
-end of each `Worker.run`, bounding per-run memory.
+`ring.slot_state(id)` hands out a stable pointer for every live slot. Each slot
+is reclaimed individually when it reaches CONSUMED. The deque is reset only
+when the worker has no registered runs or live slots.
 
 The HeapRing is **partitioned into `MAX_RING_DEPTH = 4` independent
 rings** (Strict-1; matches L2's `PTO2_MAX_RING_DEPTH`). Each ring is its
@@ -502,10 +502,10 @@ L4 parent process
 | 6 | L3 child | `inner_worker.run(my_l3_orch, args, cfg)` → `scope_begin` → `my_l3_orch(orch3, ...)` |
 | 7 | L3 `Orchestrator.submit_sub` | `l3_sub_handle` digest dispatched to L3's own sub worker child |
 | 8 | L3 sub child | child resolves digest to its local Python callable and executes `verify_result()` |
-| 9 | L3 drain | all L3 tasks complete; `scope_end` + `drain` return |
+| 9 | L3 run fence | all L3 tasks complete; `scope_end` + `wait_run` return |
 | 10 | L3 child | `inner_worker.run()` returns; `_child_worker_loop` writes `TASK_DONE` |
 | 11 | L4 LocalMailboxEndpoint | sees `TASK_DONE`; returns success completion |
-| 12 | L4 drain | L4 scope_end + drain; `w4.run()` returns |
+| 12 | L4 run fence | L4 `scope_end` + `wait_run`; `w4.run()` returns |
 
 Each level's orch fn receives **its own** `Orchestrator` — the recursion is
 symmetric. `Worker` code does not branch on `level`; the level is only a