Add: broadcast import-cache release when Worker.release_buffer() runs - #1769
Merged
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:
📝 WalkthroughWalkthrough
ChangesBuffer release propagation
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant OwnerWorker
participant ChildWorker
participant ImportRegistry
OwnerWorker->>OwnerWorker: Close owner buffer
OwnerWorker->>ChildWorker: Broadcast canonical identity
ChildWorker->>ImportRegistry: Unregister imported identity
ChildWorker->>ChildWorker: Forward release to descendants
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: 1
🤖 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/worker.py`:
- Line 9835: Make _release_import_recursive atomic with task admission:
serialize release through cache invalidation or mark the identity revoked before
releasing the admission lock, and reject subsequent submissions. Update
_submit_l2_locked to publish _chip_run_touched_identities before materializing
the mapping, ensuring release cannot miss an admitted run. Add a concurrent
release-versus-submit regression test covering L2 and L3+ behavior.
🪄 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: c6d3a598-81b9-48ed-993c-110342846eec
📒 Files selected for processing (6)
python/simpler/buffer.pypython/simpler/worker.pytests/ut/py/test_buffer.pytests/ut/py/test_remote_l3_lifecycle.pytests/ut/py/test_worker/test_create_buffer.pytests/ut/py/test_worker/test_release_buffer.py
ChaoWao
force-pushed
the
import-release-broadcast
branch
from
August 11, 2026 03:22
edc30cf to
927c1f1
Compare
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.
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
Worker.release_buffer()now broadcasts an import-cache release to every consumer of the identity — the deferred item from Fix: reject Buffer release while an in-flight run still references it #1751's design (§8 of the P1-B corrected design: "import mapping is released along with the handle's lifecycle, not dragged toWorker.close()"). Without this, a long-running chip/SUB child or nested Worker keeps a released Buffer's shm mapping resident for its entire process lifetime — a slow/dev/shmleak for any workload that creates/releases many buffers.ImportRegistry.unregister(identity)(buffer.py) drops one cached mapping; idempotent when the endpoint never materialized it.broadcast_control_all/Worker._broadcast_py_controlchannel that_CTRL_PY_REGISTERet al. already use — it already reaches bothNEXT_LEVEL(chip children and nested Workers, registered identically viaadd_next_level_worker) andSUB. New sub_cmd_CTRL_IMPORT_RELEASErides that channel.CanonicalIdentity's three meaningful fields packed field-by-field (_pack_identity_wire/_unpack_identity_wire), not the identity's own bytes —CanonicalIdentity's binding deliberately exposes nopack(), so this reconstructs the wire form the same wayremote_l3_protocol.pyalready does for the cross-machine wire._child_worker_loop(nestedNEXT_LEVELWorker) forwards one more hop down via a newWorker._release_import_recursive(), which also drops the same-processself._chip_import_registryentry an L2 direct-chip Worker may hold._broadcast_unregister: a slow or dead child must not block or failrelease_buffer().Testing
ImportRegistry.unregister()(present / absent / re-materialize-after-drop)fake_chip_l3harness) proving the wire round-trip works against a live process, not just mockstest_l3_tensor_dispatch.py, 2 chips) confirms no regression in the shared mailbox control loop the new branches were added to