Fix: close 4 of #1729's 5 deferred buffer-registry cleanup items - #1747
Merged
Merged
Conversation
…up items PR hw-native-sys#1729 (wire flip, merged aa1d7c7) deliberately deferred 6 cleanup items to its own PR comments rather than growing an already 126-file change. Item 6 (the simpler_setup.Tensor / simpler.task_interface.Tensor naming collision) was closed separately by hw-native-sys#1741. This closes 4 of the remaining 5; the 5th needs a C++ signature change and is left for a follow-up (see below). MappedArg.buffer ignored the descriptor's access mode and always returned a writable memoryview, including for FORK_COW backings, whose whole contract is that a write is invisible to the owner (copy-on-write splits the page privately). A callable that wrote through it lost data silently. buffer now returns a read-only view (memoryview.toreadonly()) when access is AccessMode.READ. New test: test_mapped_arg_buffer_is_read_only_for_a_read_access_descriptor (fails on the old code, passes after the fix). Note: torch.frombuffer does not itself honor a read-only memoryview -- it only warns and still allows the write -- so this closes the contract at the buffer-protocol layer; it does not stop a torch consumer from writing through its own tensor view. Fixing that would need actually protecting the COW pages (e.g. mprotect), out of scope here. ImportRegistry.materialize_blob and .materialize_args each rebuilt a snapshot of every identity the endpoint had ever materialized (self._by_identity in full) on every dispatch, via a now-deleted materialization_map() helper -- O(every buffer this chip child has ever seen) instead of O(this task's own tensor count), on the chip and L2-leaf dispatch path. Both callers only ever look up entries for tensors they independently re-parsed from the same blob/TaskArgs, so no entry outside the current call's own tensors was ever consulted. Both methods now build their returned dict directly from their own loop. New test: test_materialize_args_scopes_the_returned_map_to_this_calls_tensors, asserting a second call's returned map does not carry a first call's identity forward. ImportRegistry.unregister had zero callers and zero test references anywhere in the repo. The "import mapping released with handle lifecycle" invariant it was meant to serve doesn't have a lifecycle to attach to yet -- release_buffer() doesn't exist in this codebase. Deleted rather than left as an untested stand-in for a feature that isn't built. tests/st/{a2a3,a5}/tensormap_and_ringbuffer/test_l3_host_buffer_registration.py were still named after register_host_buffer, deleted long ago -- only the filename was a leftover; the class name (TestPostForkHostBufferZeroCopy) and docstrings already describe the current create_buffer + POSIX-shm mechanism accurately. Renamed both arch siblings in this commit to test_l3_post_fork_host_buffer.py. Also fixed a now-stale cross-reference in .docs/l3l4/memory-kinds.md and flagged (but did not chase down) an unrelated pre-existing gap it also pointed at: the ut test it names for kind3 registration no longer exists in the repo. Deferred to a follow-up PR: the chip task blob gets decoded twice on every dispatch -- once by read_args_from_blob (to drive ImportRegistry.materialize) and again by materialize_tensor_blob's own C++ read_blob call, on the same bytes. Closing that needs materialize_tensor_blob to accept the already-parsed view instead of re-reading raw bytes, which is a signature change on python/bindings/task_interface.cpp's hot dispatch path -- it deserves its own PR with dedicated dispatch-latency verification rather than riding along with these four independent one-line fixes. Verified: pytest tests/ut 1281 passed / 13 skipped / 0 failed; ruff check/format clean; a real a2a3 onboard run (test_l3_tensor_dispatch.py, 2 devices) exercising the changed chip-dispatch materialize path; test_l3_post_fork_host_buffer.py passing under its own a2a3sim platform restriction on both arch siblings.
ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (4)
📝 WalkthroughWalkthrough
ChangesBuffer materialization and post-fork host buffers
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Sequence Diagram(s)sequenceDiagram
participant SceneTestCase
participant TorchViews
participant _one_task_orch
participant Worker0
participant VectorKernel
SceneTestCase->>TorchViews: allocate and fill host buffers
SceneTestCase->>_one_task_orch: pass buffer handles
_one_task_orch->>Worker0: submit tensor arguments
Worker0->>VectorKernel: execute vector task
VectorKernel-->>Worker0: write output buffer
Worker0-->>SceneTestCase: return completed task
SceneTestCase->>TorchViews: compare output and release resources
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/simpler/buffer.py (1)
445-473: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winReplace
memoryview.toreadonly()with a read-only-safe contract for Torch args.
MappedArg.bufferis documented fortorch.frombuffer(arg.buffer, ...), buttorch.frombufferdoes not keep the tensor write-protected. ForAccessMode.READ, aFORK_SHMbacking can still be modified through the tensor and observed by the owner, and aFORK_COWbacking can mutate the child’s private copy instead of leaving the imported view unchanged. IfREADaccess must be enforced, do not expose a zero-copy Torch alias for read-only inputs; copy them before Torch conversion, or provide a read-only-safe tensor representation. Add a Torch regression test for the enforced contract.🤖 Prompt for 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. In `@python/simpler/buffer.py` around lines 445 - 473, Update MappedArg.buffer and the Torch argument conversion path to enforce AccessMode.READ without exposing a writable zero-copy alias; copy read-only inputs before torch.frombuffer or use an equivalent read-only-safe tensor representation, while preserving zero-copy behavior for writable access. Add a Torch regression test covering both FORK_SHM and FORK_COW read-only inputs and verifying that tensor writes cannot modify the imported view or owner-visible backing.
🤖 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 `@tests/st/a5/tensormap_and_ringbuffer/test_l3_post_fork_host_buffer.py`:
- Around line 125-130: Update the finally cleanup block in the A5 test to close
all three Buffer objects—ba, bb, and bout—after deleting the views a, b, and
out, matching the equivalent a2a3 test cleanup and ensuring shared-memory
allocations are released promptly.
---
Outside diff comments:
In `@python/simpler/buffer.py`:
- Around line 445-473: Update MappedArg.buffer and the Torch argument conversion
path to enforce AccessMode.READ without exposing a writable zero-copy alias;
copy read-only inputs before torch.frombuffer or use an equivalent
read-only-safe tensor representation, while preserving zero-copy behavior for
writable access. Add a Torch regression test covering both FORK_SHM and FORK_COW
read-only inputs and verifying that tensor writes cannot modify the imported
view or owner-visible backing.
🪄 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: 3f78db61-f308-490a-a6da-3e5e18fb2b61
📒 Files selected for processing (4)
python/simpler/buffer.pytests/st/a2a3/tensormap_and_ringbuffer/test_l3_post_fork_host_buffer.pytests/st/a5/tensormap_and_ringbuffer/test_l3_post_fork_host_buffer.pytests/ut/py/test_buffer.py
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/simpler/buffer.py (1)
445-473: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winReplace
memoryview.toreadonly()with a read-only-safe contract for Torch args.
MappedArg.bufferis documented fortorch.frombuffer(arg.buffer, ...), buttorch.frombufferdoes not keep the tensor write-protected. ForAccessMode.READ, aFORK_SHMbacking can still be modified through the tensor and observed by the owner, and aFORK_COWbacking can mutate the child’s private copy instead of leaving the imported view unchanged. IfREADaccess must be enforced, do not expose a zero-copy Torch alias for read-only inputs; copy them before Torch conversion, or provide a read-only-safe tensor representation. Add a Torch regression test for the enforced contract.🤖 Prompt for 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. In `@python/simpler/buffer.py` around lines 445 - 473, Update MappedArg.buffer and the Torch argument conversion path to enforce AccessMode.READ without exposing a writable zero-copy alias; copy read-only inputs before torch.frombuffer or use an equivalent read-only-safe tensor representation, while preserving zero-copy behavior for writable access. Add a Torch regression test covering both FORK_SHM and FORK_COW read-only inputs and verifying that tensor writes cannot modify the imported view or owner-visible backing.
🤖 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 `@tests/st/a5/tensormap_and_ringbuffer/test_l3_post_fork_host_buffer.py`:
- Around line 125-130: Update the finally cleanup block in the A5 test to close
all three Buffer objects—ba, bb, and bout—after deleting the views a, b, and
out, matching the equivalent a2a3 test cleanup and ensuring shared-memory
allocations are released promptly.
---
Outside diff comments:
In `@python/simpler/buffer.py`:
- Around line 445-473: Update MappedArg.buffer and the Torch argument conversion
path to enforce AccessMode.READ without exposing a writable zero-copy alias;
copy read-only inputs before torch.frombuffer or use an equivalent
read-only-safe tensor representation, while preserving zero-copy behavior for
writable access. Add a Torch regression test covering both FORK_SHM and FORK_COW
read-only inputs and verifying that tensor writes cannot modify the imported
view or owner-visible backing.
🪄 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: 3f78db61-f308-490a-a6da-3e5e18fb2b61
📒 Files selected for processing (4)
python/simpler/buffer.pytests/st/a2a3/tensormap_and_ringbuffer/test_l3_post_fork_host_buffer.pytests/st/a5/tensormap_and_ringbuffer/test_l3_post_fork_host_buffer.pytests/ut/py/test_buffer.py
🛑 Comments failed to post (1)
tests/st/a5/tensormap_and_ringbuffer/test_l3_post_fork_host_buffer.py (1)
125-130: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Close all three
Bufferobjects in the cleanup block.Lines 129 releases the PyTorch views but does not release
ba,bb, orbout. The equivalent a2a3 test closes each buffer after deleting its views. Repeated A5 simulator runs can retain shared-memory allocations until later worker teardown.Proposed fix
finally: # Drop the views before the worker closes: close() unlinks each backing, and a # live view over one keeps its shm alive past that. In finally so a failure # above still releases all three. del a, b, out + ba.close() + bb.close() + bout.close() assert result📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.finally: # Drop the views before the worker closes: close() unlinks each backing, and a # live view over one keeps its shm alive past that. In finally so a failure # above still releases all three. del a, b, out ba.close() bb.close() bout.close() assert result🤖 Prompt for 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. In `@tests/st/a5/tensormap_and_ringbuffer/test_l3_post_fork_host_buffer.py` around lines 125 - 130, Update the finally cleanup block in the A5 test to close all three Buffer objects—ba, bb, and bout—after deleting the views a, b, and out, matching the equivalent a2a3 test cleanup and ensuring shared-memory allocations are released promptly.
7 tasks
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.
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
PR #1729 (wire flip, merged
aa1d7c7d) deliberately deferred 6 small cleanup items to its own PR comments rather than growing an already 126-file change. Item 6 (aTensornaming collision) was closed separately by #1741. This closes 4 of the remaining 5, all confined topython/simpler/buffer.pyplus two test renames:MappedArg.bufferignored the descriptor's access mode and always returned a writablememoryview, including forFORK_COWbackings, whose contract is that a write is invisible to the owner (copy-on-write splits the page privately) — a callable that wrote through it lost data silently.buffernow returns a read-only view (memoryview.toreadonly()) when access isAccessMode.READ. New test written first against the old code to confirm it failed, then made to pass.ImportRegistry.materialize_blob/.materialize_argsrebuilt a snapshot of every identity the endpoint had ever materialized on every single dispatch — O(every buffer this chip child has ever seen) instead of O(this task's own tensor count), on the chip and L2-leaf dispatch path. Both now build their returned map directly from their own per-call loop; the now-unusedmaterialization_map()helper is deleted.ImportRegistry.unregisterhad zero callers and zero tests anywhere in the repo. The lifecycle it was meant to serve (release tied to handle lifetime) doesn't exist yet —release_buffer()isn't implemented in this codebase. Deleted rather than left as an untested stand-in for a feature that isn't built.test_l3_host_buffer_registration.pytest files (a2a3 + a5) were still named afterregister_host_buffer, deleted long ago — only the filename was stale. Renamed both arch siblings totest_l3_post_fork_host_buffer.pyin this commit.Deferred to a follow-up PR: the chip task blob gets decoded twice on every dispatch (once to drive
ImportRegistry.materialize, again insidematerialize_tensor_blob's own C++ parse of the same bytes). Closing that needs a signature change onpython/bindings/task_interface.cpp's hot dispatch path, and deserves its own PR with dedicated dispatch-latency verification rather than riding along with these four independent fixes.Test plan
pytest tests/ut— 1281 passed / 13 skipped / 0 failedruff check/ruff format --checkcleanMappedArg.bufferread-only fix, confirmed failing pre-fixmaterialize_args's returned map is scoped to the current call, not the endpoint's full historytest_l3_tensor_dispatch.py, 2 devices) exercising the changed chip-dispatch materialize patha2a3simplatform restriction; a5 sibling collects correctly🤖 Generated with Claude Code