Fix: reject Buffer release while an in-flight run still references it - #1751
Merged
ChaoWao merged 1 commit intoAug 10, 2026
Merged
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThe change tracks canonical tensor buffer identities during hierarchical dispatch. ChangesHierarchical buffer lifecycle
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant RemoteBufferEntry
participant Worker
participant RunResources
RemoteBufferEntry->>Worker: release_buffer(data)
Worker->>RunResources: check referenced identities
RunResources-->>Worker: report unsettled reference
Worker-->>RemoteBufferEntry: reject or complete release
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/simpler/remote_l3_session.py`:
- Line 193: Update _RemoteBufferEntry.close so self.owner remains set while
calling owner.release_buffer(self.data), and clear it only after the release
succeeds; preserve the owner for retries when release_buffer rejects or raises.
Add a test covering a retry of _RemoteBufferEntry.close after an in-flight
rejection, including verification that the owner registry entry remains until
successful release.
In `@python/simpler/worker.py`:
- Around line 9644-9658: Serialize release_buffer with run graph construction
using the existing graph-construction synchronization boundary, so
_record_touched_identities and release exclusion are atomic. Ensure
_submit_l3_locked keeps the accepted handle protected until orchestration
records the Tensor identity, preventing release_buffer from unlinking the Buffer
during the race. Add a concurrent regression test that blocks the orchestration
callback, races release_buffer before submit_next_level() records the identity,
and verifies the Buffer remains protected.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 3862f5b5-0b66-4694-a6e7-15b8ff2bee82
📒 Files selected for processing (6)
python/simpler/orchestrator.pypython/simpler/remote_l3_session.pypython/simpler/worker.pytests/ut/py/test_remote_l3_lifecycle.pytests/ut/py/test_worker/test_create_buffer.pytests/ut/py/test_worker/test_release_buffer.py
The design doc's P1-B lifecycle table wants release_buffer() to refuse a release while an unfinished consumer still holds a reference. That check didn't exist: Worker._release_buffer() (private, worker.py:9644) already closed a Buffer and dropped its registry entry -- used by _release_all_buffers() at Worker.close() and by remote_l3_session.py's RemoteBufferHandle.close() for a session-scoped buffer -- but nothing tracked whether a currently-dispatched task still named that Buffer's identity. _RunResources (worker.py:3486) gains touched_identities: every identity a NEXT_LEVEL Tensor arg carries during a run's orchestration. L2 never needs this -- a run completes synchronously inside submit() (RunHandle._completed, never added to _accepted_run_handles), so there is never an in-flight window there. The two async L3+ dispatch entry points, Orchestrator.submit_next_level and .submit_next_level_group, already walk every tensor arg once per call (via Worker._child_ptrs_in_args, for the existing device-pointer provenance guard); a new Worker._record_touched_identities walks the same args once more into the current run's touched set, a no-op outside a run's context. _release_buffer is renamed to public release_buffer (the design doc's literal ask) and now checks, before closing: is this identity in the touched set of any accepted-but-not-yet-cleaned-up run? If so it raises immediately rather than releasing -- reject, not block, matching the codebase's existing fail-fast admission-fence style. The check takes _submit_mu first: _submit_l3_locked adds a run's handle to _accepted_run_handles before its orchestration callback runs (the callback is what populates touched_identities via submit_next_level), so without _submit_mu, release_buffer could observe that handle mid-callback with an empty touched_identities set and release a Buffer the callback is about to dispatch a Tensor over. _submit_mu already serializes graph construction for exactly this reason (worker.py:9741, "Graph callbacks stay serialized"), so taking it here means the check only ever runs between callbacks, never mid-callback with an incomplete touched set. Updated release_buffer's two call sites (remote_l3_session.py, dropping the now-unneeded SLF001 suppression) and the existing test_create_buffer.py coverage of the renamed method (registry-drop-on-release, entry-survives-a-failed-close). remote_l3_session.py's RemoteBufferHandle.close() also had to stop clearing self.owner before calling release_buffer(): release_buffer can now legitimately raise (the in-flight rejection), and a retried close() must still go through the owning Worker rather than falling through to a direct self.data.close() that bypasses the guard and leaves the registry entry behind. self.owner now clears only after release_buffer() actually succeeds. Two other pieces of the same design-doc epic were explicitly scoped out after checking the code: buffer_id reuse + a generation bump on reuse (buffer_id is a uint64_t that _next_buffer_id() only ever increments -- no resource-exhaustion pressure justifies building reuse machinery that nothing needs), and broadcasting release to consumer processes so they drop a cached import proactively (needs a cross-process mechanism similar to Worker.unregister's callable-cleanup broadcast; a consumer that later re-resolves a released identity fails loudly instead of corrupting silently, so leaving this for a separate change is not unsafe). New tests (tests/ut/py/test_worker/test_release_buffer.py): submit_next_level and _group correctly record touched identities device-free (reusing test_child_addr_guard.py's fake-C++-orchestrator harness) and are a no-op with no run context open; release_buffer rejects while a fake in-flight handle references the identity (and the registry entry survives the rejection), succeeds once that handle is marked done or absent, stays idempotent, and correctly blocks behind a concurrently-racing orchestration callback rather than releasing mid-callback (confirmed this last test fails without the _submit_mu fix). test_remote_l3_lifecycle.py gains a matching case for the owner-retention fix: close() rejected while in flight, owner still set, registry entry still present; close() retried after the run completes succeeds and clears both. Verified: pytest tests/ut 1297 passed / 13 skipped / 0 failed; ruff check/format clean; a real a2a3 onboard run (test_l3_tensor_dispatch.py, 2 devices) exercising the touched-identity walk on the real dispatch path, plus test_l3_create_buffer.py under its own a2a3sim platform restriction.
ChaoWao
force-pushed
the
release-buffer-in-flight-retain
branch
from
August 10, 2026 01:46
82f6341 to
739d126
Compare
5 tasks
ChaoWao
added a commit
to ChaoWao/simpler-fork
that referenced
this pull request
Aug 10, 2026
My own review comment on hw-native-sys#1751 flagged a real gap: release_buffer()'s in-flight check only ever looks at self._accepted_run_handles, and _submit_l2_locked (worker.py, the direct-chip dispatch path submit() uses at L2) never adds its RunHandle there. create_buffer() only requires level >= 2, so an L2 Worker can hold a real, registered Buffer and dispatch chip runs against it with zero protection from release_buffer(). Checked how exploitable this is today: zero production callers do Worker(level=2).create_buffer() then race release_buffer() against a concurrent submit() -- this is a latent trap for future usage, not a firing bug. It doesn't explode today only because ImportRegistry.materialize() maps its own separate mmap for the identity, and POSIX unlink() only removes the shm's name -- an already-open mapping keeps working. That's an undocumented, untested coincidence, not a guarantee. Also corrected the assumption hw-native-sys#1751 shipped -- "L2 never needs this, a run completes synchronously inside submit()" -- which is false: the direct-chip lane permits one active plus one prepared compatible run, so up to two L2 runs can be in flight at once (per the W1b/W1c async-pipeline work, hw-native-sys#1748/hw-native-sys#1750). The real reason release_buffer()'s existing check can't see L2 runs is narrower: L2 uses a separate run-id namespace (self._chip_run_seq, tracked in self._chip_runs) and never touches _accepted_run_handles/_submit_mu -- _submit_locked returns from the L2 branch before reaching the with self._submit_mu: block. Rather than fold L2 into _accepted_run_handles (whose other readers -- _cleanup_bearing_predecessor, the live-handle scan for direct-control ordering, whole-run FIFO teardown draining -- carry L3-specific assumptions built around _orch-issued run ids and orchestration callbacks), this mirrors _chip_runs' own lifecycle with a parallel dict, _chip_run_touched_identities, added/removed at the same two points (_submit_l2_locked, _finalize_run_handle's L2 branch, and Worker.close()'s teardown) under the existing _registry_lock. release_buffer() now runs a second, independent check against it after the existing L3+ one. Extracted _identities_in_args as a shared static helper so _record_touched_identities (L3+) and the new L2 code walk tensor args the same way instead of duplicating the loop. New tests (tests/ut/py/test_worker/test_release_buffer.py): _submit_l2_locked records the touched identity (and an empty set for args=None); _finalize_run_handle clears it; release_buffer rejects while an L2 run's identity is present (registry entry survives the rejection) and succeeds once it's gone. Fixed two existing bare-Worker test helpers (test_create_buffer.py, test_remote_l3_lifecycle.py) that construct a Worker via __new__ and manually set internals -- they now also set the new dict. Verified: pytest tests/ut 1302 passed / 13 skipped / 0 failed; ruff check/format clean; a real a2a3 onboard run (pipeline_slots/test_pipeline_slots.py, an L2 direct-chip scene test) exercising the touched-identity walk on the real L2 dispatch path.
ChaoWao
added a commit
to ChaoWao/simpler-fork
that referenced
this pull request
Aug 10, 2026
My own review comment on hw-native-sys#1751 flagged a real gap: release_buffer()'s in-flight check only ever looks at self._accepted_run_handles, and _submit_l2_locked (worker.py, the direct-chip dispatch path submit() uses at L2) never adds its RunHandle there. create_buffer() only requires level >= 2, so an L2 Worker can hold a real, registered Buffer and dispatch chip runs against it with zero protection from release_buffer(). Checked how exploitable this is today: zero production callers do Worker(level=2).create_buffer() then race release_buffer() against a concurrent submit() -- this is a latent trap for future usage, not a firing bug. It doesn't explode today only because ImportRegistry.materialize() maps its own separate mmap for the identity, and POSIX unlink() only removes the shm's name -- an already-open mapping keeps working. That's an undocumented, untested coincidence, not a guarantee. Also corrected the assumption hw-native-sys#1751 shipped -- "L2 never needs this, a run completes synchronously inside submit()" -- which is false: the direct-chip lane permits one active plus one prepared compatible run, so up to two L2 runs can be in flight at once (per the W1b/W1c async-pipeline work, hw-native-sys#1748/hw-native-sys#1750). The real reason release_buffer()'s existing check can't see L2 runs is narrower: L2 uses a separate run-id namespace (self._chip_run_seq, tracked in self._chip_runs) and never touches _accepted_run_handles/_submit_mu -- _submit_locked returns from the L2 branch before reaching the with self._submit_mu: block. Rather than fold L2 into _accepted_run_handles (whose other readers -- _cleanup_bearing_predecessor, the live-handle scan for direct-control ordering, whole-run FIFO teardown draining -- carry L3-specific assumptions built around _orch-issued run ids and orchestration callbacks), this mirrors _chip_runs' own lifecycle with a parallel dict, _chip_run_touched_identities, added/removed at the same two points (_submit_l2_locked, _finalize_run_handle's L2 branch, and Worker.close()'s teardown) under the existing _registry_lock. release_buffer() now runs a second, independent check against it after the existing L3+ one. Extracted _identities_in_args as a shared static helper so _record_touched_identities (L3+) and the new L2 code walk tensor args the same way instead of duplicating the loop. _submit_l2_locked publishes the touched-identities entry BEFORE calling _submit_chip_run_direct, not after: writing it post-dispatch left a window where a concurrent release_buffer() could see no entry at all for a run already running on the chip, since the entry that would have blocked it didn't exist yet. On a dispatch failure the entry is popped back out. _finalize_run_handle's L2 branch and Worker.close()'s teardown now take _registry_lock around every read/clear of _chip_runs and _chip_run_touched_identities (previously the membership check in _finalize_run_handle and the two clears in close() ran unlocked), using pop(..., None) so a concurrent close() can never make either raise. New tests (tests/ut/py/test_worker/test_release_buffer.py): _submit_l2_locked records the touched identity (and an empty set for args=None); _finalize_run_handle clears it; release_buffer rejects while an L2 run's identity is present (registry entry survives the rejection) and succeeds once it's gone; touched identities are visible to release_buffer() while dispatch is still blocked mid-call, proving the publish-before-dispatch ordering. Fixed two existing bare-Worker test helpers (test_create_buffer.py, test_remote_l3_lifecycle.py) that construct a Worker via __new__ and manually set internals -- they now also set the new dict. Verified: pytest tests/ut full suite passed; ruff check/format clean; pyright clean; a real a2a3 onboard run (pipeline_slots/test_pipeline_slots.py, an L2 direct-chip scene test) exercising the touched-identity walk on the real L2 dispatch path.
ChaoWao
added a commit
that referenced
this pull request
Aug 10, 2026
…#1757) My own review comment on #1751 flagged a real gap: release_buffer()'s in-flight check only ever looks at self._accepted_run_handles, and _submit_l2_locked (worker.py, the direct-chip dispatch path submit() uses at L2) never adds its RunHandle there. create_buffer() only requires level >= 2, so an L2 Worker can hold a real, registered Buffer and dispatch chip runs against it with zero protection from release_buffer(). Checked how exploitable this is today: zero production callers do Worker(level=2).create_buffer() then race release_buffer() against a concurrent submit() -- this is a latent trap for future usage, not a firing bug. It doesn't explode today only because ImportRegistry.materialize() maps its own separate mmap for the identity, and POSIX unlink() only removes the shm's name -- an already-open mapping keeps working. That's an undocumented, untested coincidence, not a guarantee. Also corrected the assumption #1751 shipped -- "L2 never needs this, a run completes synchronously inside submit()" -- which is false: the direct-chip lane permits one active plus one prepared compatible run, so up to two L2 runs can be in flight at once (per the W1b/W1c async-pipeline work, #1748/#1750). The real reason release_buffer()'s existing check can't see L2 runs is narrower: L2 uses a separate run-id namespace (self._chip_run_seq, tracked in self._chip_runs) and never touches _accepted_run_handles/_submit_mu -- _submit_locked returns from the L2 branch before reaching the with self._submit_mu: block. Rather than fold L2 into _accepted_run_handles (whose other readers -- _cleanup_bearing_predecessor, the live-handle scan for direct-control ordering, whole-run FIFO teardown draining -- carry L3-specific assumptions built around _orch-issued run ids and orchestration callbacks), this mirrors _chip_runs' own lifecycle with a parallel dict, _chip_run_touched_identities, added/removed at the same two points (_submit_l2_locked, _finalize_run_handle's L2 branch, and Worker.close()'s teardown) under the existing _registry_lock. release_buffer() now runs a second, independent check against it after the existing L3+ one. Extracted _identities_in_args as a shared static helper so _record_touched_identities (L3+) and the new L2 code walk tensor args the same way instead of duplicating the loop. _submit_l2_locked publishes the touched-identities entry BEFORE calling _submit_chip_run_direct, not after: writing it post-dispatch left a window where a concurrent release_buffer() could see no entry at all for a run already running on the chip, since the entry that would have blocked it didn't exist yet. On a dispatch failure the entry is popped back out. _finalize_run_handle's L2 branch and Worker.close()'s teardown now take _registry_lock around every read/clear of _chip_runs and _chip_run_touched_identities (previously the membership check in _finalize_run_handle and the two clears in close() ran unlocked), using pop(..., None) so a concurrent close() can never make either raise. New tests (tests/ut/py/test_worker/test_release_buffer.py): _submit_l2_locked records the touched identity (and an empty set for args=None); _finalize_run_handle clears it; release_buffer rejects while an L2 run's identity is present (registry entry survives the rejection) and succeeds once it's gone; touched identities are visible to release_buffer() while dispatch is still blocked mid-call, proving the publish-before-dispatch ordering. Fixed two existing bare-Worker test helpers (test_create_buffer.py, test_remote_l3_lifecycle.py) that construct a Worker via __new__ and manually set internals -- they now also set the new dict. Verified: pytest tests/ut full suite passed; ruff check/format clean; pyright clean; a real a2a3 onboard run (pipeline_slots/test_pipeline_slots.py, an L2 direct-chip scene test) exercising the touched-identity walk on the real L2 dispatch path.
lterrac
added a commit
to lterrac/simpler
that referenced
this pull request
Aug 10, 2026
Three things serialize the device-memory ops, in series, so removing any one of them alone measures as noise — which is why this took a while to pin down. A `copy_to` on chip 0 blocks a `malloc` on chip 1 for its whole duration, and the eight chips of an 8-way upload run strictly back to back. `_child_prov_lock` stays the bookkeeping lock — it still makes each provenance mutation/read atomic, and the safety-first ordering is unchanged (record after a successful alloc, revoke before a native free) — and a per-worker lock is taken around the native call instead. Ops on the same worker stay mutually exclusive, so a copy can still never overlap that buffer's free; ops on different workers now overlap. The per-worker lock is always acquired before `_child_prov_lock` and never the reverse, so the pair cannot deadlock. `_submit_mu`, taken through `_control_reservation`, is the other one: a control command that belongs to no run holds it across the native call, so with the provenance fix alone it becomes the serializer. What such a command needs is "no run may be admitted while I run", which is a property of the worker, and two commands on different chips can both have that at the same time. So `_submit_mu` becomes a shared/exclusive lock: run admission takes it exclusively, control takes it shared. Writer-preferring, so control traffic cannot starve a submit. The reservation's re-entrancy is untouched: it short-circuits on the thread-local set before reaching the lock. The third is the **GIL**. `Worker.malloc / free / copy_to / copy_from` are bound as plain lambdas in `worker_bind.h` with no call guard, so the interpreter lock is held for the whole native call while 31 other methods in that same file already release it. With the two Python locks split but the GIL still held, eight threads still cannot overlap: measured `overlap_factor` (sum of per-copy wall times over the wall time of the batch) stayed at 1.1-1.6 out of a possible 8. These four get `nb::call_guard<nb::gil_scoped_release>()`; none of them re-enters Python, and the descriptors are converted before the call. Re-applied on the post-hw-native-sys#1650/hw-native-sys#1729 structure rather than rebased textually. The provenance and native-call logic moved out of `Orchestrator` into `Worker`, so the split now lives in `Worker.alloc_child_tensor / free / copy_to / copy_from` and `orchestrator.py` is left exactly as main has it. `release_buffer` (hw-native-sys#1751) took `_submit_mu` bare, which a textual merge would have compiled and then crashed on, since `_SharedExclusiveLock` has no `__enter__`; it now takes it exclusively, keeping the ordering it had as a plain lock. Both traps were called out by @ChaoZheng109 in review. Three tests stand a plain `threading.Lock` in for `_submit_mu` or take it as a context manager, so they are updated to the real type and to `.exclusive()` (`test_create_buffer.py`, `test_remote_l3_lifecycle.py`, `test_release_buffer.py::test_serializes_with_a_racing_orchestration_callback`). They pin the serializer's *identity*, not its granularity, and the exclusive form is what graph construction now takes, so the property each one asserts is unchanged. This also narrows an invariant an existing test pins down, so that test is updated rather than left passing by accident: `test_free_holds_lock_across_native_free` asserted that the *parent worker's* lock is held across the native free. It now asserts the narrower exclusion actually needed — that worker's own lock held across the native call, `_child_prov_lock` released, and the revoke committed first. Provenance is keyed by (worker_id, ptr) and revoked before the native free, so a concurrent dispatch reads the table under `_child_prov_lock` and finds the address already gone, or is about a different chip entirely. Measured on this code, on 8 x 910B2, with a pure-simpler harness (no pypto, no kernels): an L3 Worker over 8 chips, one born-shared 2.15 GB host buffer per chip, uploaded through `Worker.copy_to` outside a run — the same control path the resident-weight upload uses. The same work is done once with a thread per chip and once strictly sequentially; `overlap_factor` is the sum of the per-copy wall times over the wall time of the batch, so 1 means "back to back" and 8 means "fully overlapped". | | threaded | serial | overlap_factor | |---|---:|---:|---:| | main, unpatched | 19.99-22.11 GB/s | 20.69-22.82 GB/s | 1.4-1.8 | | + the two lock splits | 20.86 GB/s | 21.80 GB/s | 1.1-1.6 | | + the GIL guards (this PR) | **86.85-90.03 GB/s** | 21.47-21.81 GB/s | **6.4-6.5** | **4.14x** over sequential on the same build, where before there was none: the eight copies now all start together instead of queueing. Two things keep that honest. The serial column is the control and stays in 20.7-22.8 GB/s across all three builds, so the gain is concurrency and not a faster machine. And absolute throughput drifts about 10% between sessions — which is why the claim rests on the threaded/serial ratio measured *within* a build, and on the overlap factor, rather than on any single absolute number. The table also shows why this took three attempts to see. The two lock splits move neither throughput nor overlap; an earlier attempt at the GIL guards alone measured as noise too. With three serializers in series, removing any one of them changes nothing measurable, and only the last one removed appears to "cause" the win.
lterrac
added a commit
to lterrac/simpler
that referenced
this pull request
Aug 10, 2026
Three things serialize the device-memory ops, in series, so removing any one of them alone measures as noise — which is why this took a while to pin down. A `copy_to` on chip 0 blocks a `malloc` on chip 1 for its whole duration, and the eight chips of an 8-way upload run strictly back to back. `_child_prov_lock` stays the bookkeeping lock — it still makes each provenance mutation/read atomic, and the safety-first ordering is unchanged (record after a successful alloc, revoke before a native free) — and a per-worker lock is taken around the native call instead. Ops on the same worker stay mutually exclusive, so a copy can still never overlap that buffer's free; ops on different workers now overlap. The per-worker lock is always acquired before `_child_prov_lock` and never the reverse, so the pair cannot deadlock. `_submit_mu`, taken through `_control_reservation`, is the other one: a control command that belongs to no run holds it across the native call, so with the provenance fix alone it becomes the serializer. What such a command needs is "no run may be admitted while I run", which is a property of the worker, and two commands on different chips can both have that at the same time. So `_submit_mu` becomes a shared/exclusive lock: run admission takes it exclusively, control takes it shared. Writer-preferring, so control traffic cannot starve a submit. The reservation's re-entrancy is untouched: it short-circuits on the thread-local set before reaching the lock. The third is the **GIL**. `Worker.malloc / free / copy_to / copy_from` are bound as plain lambdas in `worker_bind.h` with no call guard, so the interpreter lock is held for the whole native call while 31 other methods in that same file already release it. With the two Python locks split but the GIL still held, eight threads still cannot overlap: measured `overlap_factor` (sum of per-copy wall times over the wall time of the batch) stayed at 1.1-1.6 out of a possible 8. These four get `nb::call_guard<nb::gil_scoped_release>()`; none of them re-enters Python, and the descriptors are converted before the call. Re-applied on the post-hw-native-sys#1650/hw-native-sys#1729 structure rather than rebased textually. The provenance and native-call logic moved out of `Orchestrator` into `Worker`, so the split now lives in `Worker.alloc_child_tensor / free / copy_to / copy_from` and `orchestrator.py` is left exactly as main has it. `release_buffer` (hw-native-sys#1751) took `_submit_mu` bare, which a textual merge would have compiled and then crashed on, since `_SharedExclusiveLock` has no `__enter__`; it now takes it exclusively, keeping the ordering it had as a plain lock. Both traps were called out by @ChaoZheng109 in review. Three tests stand a plain `threading.Lock` in for `_submit_mu` or take it as a context manager, so they are updated to the real type and to `.exclusive()` (`test_create_buffer.py`, `test_remote_l3_lifecycle.py`, `test_release_buffer.py::test_serializes_with_a_racing_orchestration_callback`). They pin the serializer's *identity*, not its granularity, and the exclusive form is what graph construction now takes, so the property each one asserts is unchanged. This also narrows an invariant an existing test pins down, so that test is updated rather than left passing by accident: `test_free_holds_lock_across_native_free` asserted that the *parent worker's* lock is held across the native free. It now asserts the narrower exclusion actually needed — that worker's own lock held across the native call, `_child_prov_lock` released, and the revoke committed first. Provenance is keyed by (worker_id, ptr) and revoked before the native free, so a concurrent dispatch reads the table under `_child_prov_lock` and finds the address already gone, or is about a different chip entirely. Measured on this code, on 8 x 910B2, with a pure-simpler harness (no pypto, no kernels): an L3 Worker over 8 chips, one born-shared 2.15 GB host buffer per chip, uploaded through `Worker.copy_to` outside a run — the same control path the resident-weight upload uses. The same work is done once with a thread per chip and once strictly sequentially; `overlap_factor` is the sum of the per-copy wall times over the wall time of the batch, so 1 means "back to back" and 8 means "fully overlapped". | | threaded | serial | overlap_factor | |---|---:|---:|---:| | main, unpatched | 19.99-22.11 GB/s | 20.69-22.82 GB/s | 1.4-1.8 | | + the two lock splits | 20.86 GB/s | 21.80 GB/s | 1.1-1.6 | | + the GIL guards (this PR) | **86.85-90.03 GB/s** | 21.47-21.81 GB/s | **6.4-6.5** | **4.14x** over sequential on the same build, where before there was none: the eight copies now all start together instead of queueing. Two things keep that honest. The serial column is the control and stays in 20.7-22.8 GB/s across all three builds, so the gain is concurrency and not a faster machine. And absolute throughput drifts about 10% between sessions — which is why the claim rests on the threaded/serial ratio measured *within* a build, and on the overlap factor, rather than on any single absolute number. The table also shows why this took three attempts to see. The two lock splits move neither throughput nor overlap; an earlier attempt at the GIL guards alone measured as noise too. With three serializers in series, removing any one of them changes nothing measurable, and only the last one removed appears to "cause" the win.
ChaoWao
added a commit
to ChaoWao/simpler-fork
that referenced
this pull request
Aug 11, 2026
hw-native-sys#1751 deferred this explicitly: releasing a Buffer only unlinks its shm on the owner side, but a consumer that once materialized it (a forked chip or SUB child, or a nested NEXT_LEVEL Worker) keeps that mapping resident for its entire process lifetime -- release never told it to drop the cache. The authoritative design (.docs/worker-memory-model/p1b-corrected-design.md §8) states the requirement directly: import mapping is supposed to be released along with the handle's lifecycle, not dragged to Worker.close(). For a long-running worker that creates/releases many buffers over its life, every consumer's resident mapping is a slow leak of /dev/shm capacity that never gets reclaimed until the consumer process itself exits. ImportRegistry (buffer.py) gains unregister(identity): pop the cached mapping if this endpoint made one, close its shm, no-op otherwise -- hw-native-sys#1747 deleted the previous unregister() as dead code with zero production callers; this reintroduces one with a real caller. The broadcast itself needs no new C++: the codebase already has a generic cross-process control channel (WorkerManager::broadcast_control_all, driven from Python via Worker._broadcast_py_control) that _CTRL_PY_REGISTER / _CTRL_PY_UNREGISTER / _CTRL_PY_IMPORT_REGISTER already use, and it reaches both WorkerType.NEXT_LEVEL (which covers chip children and nested Workers uniformly -- both are registered through the same add_next_level_worker call) and WorkerType.SUB. A new sub_cmd, _CTRL_IMPORT_RELEASE, rides that existing channel; the digest-sized control slot carries a CanonicalIdentity's three meaningful fields (owner_instance_id, buffer_id, generation) packed by a new _pack_identity_wire/_unpack_identity_wire pair, not the identity's own bytes -- CanonicalIdentity's binding deliberately exposes no pack() (a raw byte dump once let a registry key on wire padding and split one backing in two), so the wire form is reconstructed field-by-field here, the same way remote_l3_protocol.py already encodes one for the cross-machine wire. Receiving-side branches: _run_chip_main_loop and _sub_worker_loop each call import_registry.unregister(identity) directly; _child_worker_loop (a nested NEXT_LEVEL Worker) forwards one more hop down via the new Worker._release_import_recursive(), which also drops the same-process self._chip_import_registry entry an L2 direct-chip Worker may hold for its own buffers. release_buffer() calls it once buffer.close() has actually succeeded, so a failed close never tells a descendant to drop a mapping the owner still considers live. The broadcast is best-effort throughout, mirroring _broadcast_unregister: a child that never materialized the identity has nothing to drop, and a slow or dead child must not block or fail release_buffer() -- the Buffer is already closed on the owner side by the time it runs. _submit_l2_locked now publishes _chip_run_touched_identities BEFORE calling _materialize_l2_args, not just before native dispatch: _materialize_l2_args is what populates self._chip_import_registry, the very cache this PR's broadcast now pops on release. Publishing only around dispatch (as hw-native-sys#1757 left it) still left a window where release_buffer() could see no in-flight run while a submit already in progress had cached the mapping, pass its check, and pop that mapping out from under a dispatch that had not reached native execution yet -- self._chip_import_registry never existed as a release_buffer() target before this PR, so this window is newly reachable, not a pre-existing gap. New tests: ImportRegistry.unregister() present/absent/re-materialize-after- drop (test_buffer.py); a real-forked-chip-child integration test via the device-free fake_chip_l3 harness proving the wire round-trip (sub_cmd numbering, CanonicalIdentity packing) actually works against a live process, not just mocks; a regression test blocking _materialize_l2_args mid-call and confirming release_buffer() already rejects at that point, not only after materialize returns (test_release_buffer.py) -- confirmed against the pre-fix ordering first: release_buffer() did not raise, and the blocked submit thread then hit FileNotFoundError reopening the shm release had already unlinked out from under it. Two bare-Worker test helpers (test_create_buffer.py, test_remote_l3_lifecycle.py) construct a Worker via __new__ and manually set internals -- they now also set _chip_import_registry and _worker so release_buffer() (which now touches both) keeps working against them. Verified: pyut 1315 passed / 13 skipped / 0 failed; ruff check/format and pyright clean on every touched file; a real a2a3 onboard run (test_l3_tensor_dispatch.py, 2 chips) confirms no regression in the shared mailbox control loop the new sub_cmd branches were added to.
ChaoWao
added a commit
that referenced
this pull request
Aug 11, 2026
…#1769) #1751 deferred this explicitly: releasing a Buffer only unlinks its shm on the owner side, but a consumer that once materialized it (a forked chip or SUB child, or a nested NEXT_LEVEL Worker) keeps that mapping resident for its entire process lifetime -- release never told it to drop the cache. The authoritative design (.docs/worker-memory-model/p1b-corrected-design.md §8) states the requirement directly: import mapping is supposed to be released along with the handle's lifecycle, not dragged to Worker.close(). For a long-running worker that creates/releases many buffers over its life, every consumer's resident mapping is a slow leak of /dev/shm capacity that never gets reclaimed until the consumer process itself exits. ImportRegistry (buffer.py) gains unregister(identity): pop the cached mapping if this endpoint made one, close its shm, no-op otherwise -- #1747 deleted the previous unregister() as dead code with zero production callers; this reintroduces one with a real caller. The broadcast itself needs no new C++: the codebase already has a generic cross-process control channel (WorkerManager::broadcast_control_all, driven from Python via Worker._broadcast_py_control) that _CTRL_PY_REGISTER / _CTRL_PY_UNREGISTER / _CTRL_PY_IMPORT_REGISTER already use, and it reaches both WorkerType.NEXT_LEVEL (which covers chip children and nested Workers uniformly -- both are registered through the same add_next_level_worker call) and WorkerType.SUB. A new sub_cmd, _CTRL_IMPORT_RELEASE, rides that existing channel; the digest-sized control slot carries a CanonicalIdentity's three meaningful fields (owner_instance_id, buffer_id, generation) packed by a new _pack_identity_wire/_unpack_identity_wire pair, not the identity's own bytes -- CanonicalIdentity's binding deliberately exposes no pack() (a raw byte dump once let a registry key on wire padding and split one backing in two), so the wire form is reconstructed field-by-field here, the same way remote_l3_protocol.py already encodes one for the cross-machine wire. Receiving-side branches: _run_chip_main_loop and _sub_worker_loop each call import_registry.unregister(identity) directly; _child_worker_loop (a nested NEXT_LEVEL Worker) forwards one more hop down via the new Worker._release_import_recursive(), which also drops the same-process self._chip_import_registry entry an L2 direct-chip Worker may hold for its own buffers. release_buffer() calls it once buffer.close() has actually succeeded, so a failed close never tells a descendant to drop a mapping the owner still considers live. The broadcast is best-effort throughout, mirroring _broadcast_unregister: a child that never materialized the identity has nothing to drop, and a slow or dead child must not block or fail release_buffer() -- the Buffer is already closed on the owner side by the time it runs. _submit_l2_locked now publishes _chip_run_touched_identities BEFORE calling _materialize_l2_args, not just before native dispatch: _materialize_l2_args is what populates self._chip_import_registry, the very cache this PR's broadcast now pops on release. Publishing only around dispatch (as #1757 left it) still left a window where release_buffer() could see no in-flight run while a submit already in progress had cached the mapping, pass its check, and pop that mapping out from under a dispatch that had not reached native execution yet -- self._chip_import_registry never existed as a release_buffer() target before this PR, so this window is newly reachable, not a pre-existing gap. New tests: ImportRegistry.unregister() present/absent/re-materialize-after- drop (test_buffer.py); a real-forked-chip-child integration test via the device-free fake_chip_l3 harness proving the wire round-trip (sub_cmd numbering, CanonicalIdentity packing) actually works against a live process, not just mocks; a regression test blocking _materialize_l2_args mid-call and confirming release_buffer() already rejects at that point, not only after materialize returns (test_release_buffer.py) -- confirmed against the pre-fix ordering first: release_buffer() did not raise, and the blocked submit thread then hit FileNotFoundError reopening the shm release had already unlinked out from under it. Two bare-Worker test helpers (test_create_buffer.py, test_remote_l3_lifecycle.py) construct a Worker via __new__ and manually set internals -- they now also set _chip_import_registry and _worker so release_buffer() (which now touches both) keeps working against them. Verified: pyut 1315 passed / 13 skipped / 0 failed; ruff check/format and pyright clean on every touched file; a real a2a3 onboard run (test_l3_tensor_dispatch.py, 2 chips) confirms no regression in the shared mailbox control loop the new sub_cmd branches were added to.
lterrac
added a commit
to lterrac/simpler
that referenced
this pull request
Aug 11, 2026
Three things serialize the device-memory ops, in series, so removing any one of them alone measures as noise — which is why this took a while to pin down. A `copy_to` on chip 0 blocks a `malloc` on chip 1 for its whole duration, and the eight chips of an 8-way upload run strictly back to back. `_child_prov_lock` stays the bookkeeping lock — it still makes each provenance mutation/read atomic, and the safety-first ordering is unchanged (record after a successful alloc, revoke before a native free) — and a per-worker lock is taken around the native call instead. Ops on the same worker stay mutually exclusive, so a copy can still never overlap that buffer's free; ops on different workers now overlap. The per-worker lock is always acquired before `_child_prov_lock` and never the reverse, so the pair cannot deadlock. `_submit_mu`, taken through `_control_reservation`, is the other one: a control command that belongs to no run holds it across the native call, so with the provenance fix alone it becomes the serializer. What such a command needs is "no run may be admitted while I run", which is a property of the worker, and two commands on different chips can both have that at the same time. So `_submit_mu` becomes a shared/exclusive lock: run admission takes it exclusively, control takes it shared. Writer-preferring, so control traffic cannot starve a submit. The reservation's re-entrancy is untouched: it short-circuits on the thread-local set before reaching the lock. The third is the **GIL**. `Worker.malloc / free / copy_to / copy_from` are bound as plain lambdas in `worker_bind.h` with no call guard, so the interpreter lock is held for the whole native call while 31 other methods in that same file already release it. With the two Python locks split but the GIL still held, eight threads still cannot overlap: measured `overlap_factor` (sum of per-copy wall times over the wall time of the batch) stayed at 1.1-1.6 out of a possible 8. These four get `nb::call_guard<nb::gil_scoped_release>()`; none of them re-enters Python, and the descriptors are converted before the call. Re-applied on the post-hw-native-sys#1650/hw-native-sys#1729 structure rather than rebased textually. The provenance and native-call logic moved out of `Orchestrator` into `Worker`, so the split now lives in `Worker.alloc_child_tensor / free / copy_to / copy_from` and `orchestrator.py` is left exactly as main has it. `release_buffer` (hw-native-sys#1751) took `_submit_mu` bare, which a textual merge would have compiled and then crashed on, since `_SharedExclusiveLock` has no `__enter__`; it now takes it exclusively, keeping the ordering it had as a plain lock. Both traps were called out by @ChaoZheng109 in review. Three tests stand a plain `threading.Lock` in for `_submit_mu` or take it as a context manager, so they are updated to the real type and to `.exclusive()` (`test_create_buffer.py`, `test_remote_l3_lifecycle.py`, `test_release_buffer.py::test_serializes_with_a_racing_orchestration_callback`). They pin the serializer's *identity*, not its granularity, and the exclusive form is what graph construction now takes, so the property each one asserts is unchanged. This also narrows an invariant an existing test pins down, so that test is updated rather than left passing by accident: `test_free_holds_lock_across_native_free` asserted that the *parent worker's* lock is held across the native free. It now asserts the narrower exclusion actually needed — that worker's own lock held across the native call, `_child_prov_lock` released, and the revoke committed first. Provenance is keyed by (worker_id, ptr) and revoked before the native free, so a concurrent dispatch reads the table under `_child_prov_lock` and finds the address already gone, or is about a different chip entirely. Measured on this code, on 8 x 910B2, with a pure-simpler harness (no pypto, no kernels): an L3 Worker over 8 chips, one born-shared 2.15 GB host buffer per chip, uploaded through `Worker.copy_to` outside a run — the same control path the resident-weight upload uses. The same work is done once with a thread per chip and once strictly sequentially; `overlap_factor` is the sum of the per-copy wall times over the wall time of the batch, so 1 means "back to back" and 8 means "fully overlapped". | | threaded | serial | overlap_factor | |---|---:|---:|---:| | main, unpatched | 19.99-22.11 GB/s | 20.69-22.82 GB/s | 1.4-1.8 | | + the two lock splits | 20.86 GB/s | 21.80 GB/s | 1.1-1.6 | | + the GIL guards (this PR) | **86.85-90.03 GB/s** | 21.47-21.81 GB/s | **6.4-6.5** | **4.14x** over sequential on the same build, where before there was none: the eight copies now all start together instead of queueing. Two things keep that honest. The serial column is the control and stays in 20.7-22.8 GB/s across all three builds, so the gain is concurrency and not a faster machine. And absolute throughput drifts about 10% between sessions — which is why the claim rests on the threaded/serial ratio measured *within* a build, and on the overlap factor, rather than on any single absolute number. The table also shows why this took three attempts to see. The two lock splits move neither throughput nor overlap; an earlier attempt at the GIL guards alone measured as noise too. With three serializers in series, removing any one of them changes nothing measurable, and only the last one removed appears to "cause" the win.
ChaoZheng109
pushed a commit
that referenced
this pull request
Aug 11, 2026
Three things serialize the device-memory ops, in series, so removing any one of them alone measures as noise — which is why this took a while to pin down. A `copy_to` on chip 0 blocks a `malloc` on chip 1 for its whole duration, and the eight chips of an 8-way upload run strictly back to back. `_child_prov_lock` stays the bookkeeping lock — it still makes each provenance mutation/read atomic, and the safety-first ordering is unchanged (record after a successful alloc, revoke before a native free) — and a per-worker lock is taken around the native call instead. Ops on the same worker stay mutually exclusive, so a copy can still never overlap that buffer's free; ops on different workers now overlap. The per-worker lock is always acquired before `_child_prov_lock` and never the reverse, so the pair cannot deadlock. `_submit_mu`, taken through `_control_reservation`, is the other one: a control command that belongs to no run holds it across the native call, so with the provenance fix alone it becomes the serializer. What such a command needs is "no run may be admitted while I run", which is a property of the worker, and two commands on different chips can both have that at the same time. So `_submit_mu` becomes a shared/exclusive lock: run admission takes it exclusively, control takes it shared. Writer-preferring, so control traffic cannot starve a submit. The reservation's re-entrancy is untouched: it short-circuits on the thread-local set before reaching the lock. The third is the **GIL**. `Worker.malloc / free / copy_to / copy_from` are bound as plain lambdas in `worker_bind.h` with no call guard, so the interpreter lock is held for the whole native call while 31 other methods in that same file already release it. With the two Python locks split but the GIL still held, eight threads still cannot overlap: measured `overlap_factor` (sum of per-copy wall times over the wall time of the batch) stayed at 1.1-1.6 out of a possible 8. These four get `nb::call_guard<nb::gil_scoped_release>()`; none of them re-enters Python, and the descriptors are converted before the call. Re-applied on the post-#1650/#1729 structure rather than rebased textually. The provenance and native-call logic moved out of `Orchestrator` into `Worker`, so the split now lives in `Worker.alloc_child_tensor / free / copy_to / copy_from` and `orchestrator.py` is left exactly as main has it. `release_buffer` (#1751) took `_submit_mu` bare, which a textual merge would have compiled and then crashed on, since `_SharedExclusiveLock` has no `__enter__`; it now takes it exclusively, keeping the ordering it had as a plain lock. Both traps were called out by @ChaoZheng109 in review. Three tests stand a plain `threading.Lock` in for `_submit_mu` or take it as a context manager, so they are updated to the real type and to `.exclusive()` (`test_create_buffer.py`, `test_remote_l3_lifecycle.py`, `test_release_buffer.py::test_serializes_with_a_racing_orchestration_callback`). They pin the serializer's *identity*, not its granularity, and the exclusive form is what graph construction now takes, so the property each one asserts is unchanged. This also narrows an invariant an existing test pins down, so that test is updated rather than left passing by accident: `test_free_holds_lock_across_native_free` asserted that the *parent worker's* lock is held across the native free. It now asserts the narrower exclusion actually needed — that worker's own lock held across the native call, `_child_prov_lock` released, and the revoke committed first. Provenance is keyed by (worker_id, ptr) and revoked before the native free, so a concurrent dispatch reads the table under `_child_prov_lock` and finds the address already gone, or is about a different chip entirely. Measured on this code, on 8 x 910B2, with a pure-simpler harness (no pypto, no kernels): an L3 Worker over 8 chips, one born-shared 2.15 GB host buffer per chip, uploaded through `Worker.copy_to` outside a run — the same control path the resident-weight upload uses. The same work is done once with a thread per chip and once strictly sequentially; `overlap_factor` is the sum of the per-copy wall times over the wall time of the batch, so 1 means "back to back" and 8 means "fully overlapped". | | threaded | serial | overlap_factor | |---|---:|---:|---:| | main, unpatched | 19.99-22.11 GB/s | 20.69-22.82 GB/s | 1.4-1.8 | | + the two lock splits | 20.86 GB/s | 21.80 GB/s | 1.1-1.6 | | + the GIL guards (this PR) | **86.85-90.03 GB/s** | 21.47-21.81 GB/s | **6.4-6.5** | **4.14x** over sequential on the same build, where before there was none: the eight copies now all start together instead of queueing. Two things keep that honest. The serial column is the control and stays in 20.7-22.8 GB/s across all three builds, so the gain is concurrency and not a faster machine. And absolute throughput drifts about 10% between sessions — which is why the claim rests on the threaded/serial ratio measured *within* a build, and on the overlap factor, rather than on any single absolute number. The table also shows why this took three attempts to see. The two lock splits move neither throughput nor overlap; an earlier attempt at the GIL guards alone measured as noise too. With three serializers in series, removing any one of them changes nothing measurable, and only the last one removed appears to "cause" the win.
4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The memory-model design doc's P1-B lifecycle table wants a released Buffer's
release_buffer()to refuse the release while an unfinished consumer still holds a reference. That check didn't exist. Mid-implementation I foundWorker._release_buffer()(private) already closes a Buffer and drops its registry entry — used by_release_all_buffers()atWorker.close()and byremote_l3_session.py'sRemoteBufferHandle.close()for a session-scoped buffer — so this builds on that instead of adding a second, parallel public method:_RunResourcesgainstouched_identities: every identity a NEXT_LEVELTensorarg carries during a run's orchestration. L2 never needs this — a run completes synchronously insidesubmit()(RunHandle._completed, never added to_accepted_run_handles), so there's never an in-flight window there.Orchestrator.submit_next_level/.submit_next_level_groupalready walk every tensor arg once per call (for the existing device-pointer provenance guard); a newWorker._record_touched_identitieswalks the same args once more into the current run's touched set._release_bufferis renamed to publicrelease_buffer(the design doc's literal ask) and now checks, before closing: is this identity in the touched set of any accepted-but-not-yet-cleaned-up run? If so it raises immediately — reject, not block, matching the codebase's existing fail-fast admission-fence style.Two other pieces of the same design-doc epic were explicitly scoped out after checking the code:
buffer_idreuse + a generation bump on reuse —buffer_idis auint64_tthat_next_buffer_id()only ever increments; no resource-exhaustion pressure justifies building reuse machinery nothing needs.Worker.unregister's callable-cleanup broadcast. A consumer that later re-resolves a released identity fails loudly instead of corrupting silently, so leaving this for a separate change is not unsafe.Test plan
tests/ut/py/test_worker/test_release_buffer.py):submit_next_level/_groupcorrectly record touched identities device-free (reusingtest_child_addr_guard.py's fake-C++-orchestrator harness) and are a no-op with no run context open;release_bufferrejects while a fake in-flight handle references the identity (registry entry survives the rejection), succeeds once that handle is done/absent, stays idempotentpytest tests/ut— 1295 passed / 13 skipped / 0 failedruff check/ruff format --checkcleantest_l3_tensor_dispatch.py, 2 devices) exercising the touched-identity walk on the real dispatch pathtest_l3_create_buffer.pyunder its owna2a3simplatform restriction🤖 Generated with Claude Code