Fix: harden the Buffer/Tensor wire gate and the create_buffer child check - #1703
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 changes harden tensor extent and descriptor handling, add explicit fork backend selection, clarify buffer materialization rules, and allow valid next-level worker topologies to create shared buffers. Tests cover overflow, backend constraints, worker topology, identity, and cleanup. ChangesBuffer ABI and worker integration
Estimated code review effort: 3 (Moderate) | ~20 minutes 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 |
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 `@tests/ut/cpp/types/test_buffer.cpp`:
- Around line 287-293: Update the test setup before the `b = a` assignment so
`a.byte_offset` is 1 instead of 0, ensuring the end-offset addition overflows
and the existing `tensor_extent_bytes(a)` and `tensors_overlap(a, b)` assertions
cover that branch.
🪄 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: 65703234-c971-493d-ad83-d5c835207ad5
📒 Files selected for processing (7)
docs/buffer-abi.mdpython/simpler/buffer.pypython/simpler/worker.pysrc/common/task_interface/buffer.htests/ut/cpp/types/test_buffer.cpptests/ut/py/test_buffer.pytests/ut/py/test_worker/test_create_buffer.py
`tensor_extent_bytes` summed `(shapes[i]-1)*strides[i]` into a uint64. Both fields are u32, so a single product already reaches ~2^64 and the multiply by the element size overflows on top of it: the extent wrapped to a small value and `validate_tensor` accepted the view as in-bounds. One dimension is enough — shapes=(2147483649,), strides=(2147483648,), FLOAT32 validated against a 4-byte backing while addressing ~16 EiB — and the path is reachable from Python through `buffer.tensor(...)`. The arithmetic now saturates, and an extent that lands on the saturation sentinel is refused by name rather than compared against `nbytes`, so a descriptor claiming an absurd `nbytes` cannot buy the view back. The frozen field layout is untouched. `tensors_overlap` inherited the same wrap, where the consequence is worse than a rejected argument: two fully overlapping views compare as disjoint, which at the dependency layer is a missed edge rather than an error. Its end offsets saturate too. `BufferDescriptor::operator==` bounded a memcmp by an unvalidated `body_len`. It is clamped to DESC_MAX_BYTES, so the one length field in the header cannot bound a read past the array it indexes. The arbitrary-bytes pass asserted magic / generation / body_len / ndims on its survivors but not the footprint invariant, which is why 4096 random blobs never caught this; it asserts it now. `validate_tensor`'s comment claimed the validator stands behind materialization as well. It does not — `ImportRegistry.materialize` takes an already-decoded descriptor and adds no endpoint check — and a header that freezes an ABI should not overstate its own gate. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ccess `wrap_fork_inherited` derived `backend_kind` from `access`, tagging FORK_SHM whenever the grant was not READ. FORK_SHM and FORK_COW are opposite kernel write semantics, not two spellings of one grant, and inferring either from the other makes a read-only MAP_SHARED backing inexpressible: it is tagged FORK_COW, and FORK_COW's READ-only rule then locks it there. The caller holds the mmap and is the only party that knows which it is, so it states the tag; the default pair stays the safe one, FORK_COW with READ. The function had no caller, so nothing depends on the old signature. Also corrects text describing machinery that does not exist. `ImportRegistry.materialize` pointed a caller holding raw bytes at `BufferDescriptor.unpack`, and the test module said it pins a `pack`/`unpack` round trip; neither exists, and withholding the encoding is exactly what keeps construction the only way in. `docs/buffer-abi.md` named the wire type `simpler.task_interface.Tensor` while the same page's status note, the module and a test all say it stays in `simpler.buffer` until the cutover. The `h.shm.buf` sample now records that it is transitional: byte access belongs on the view, since a device backing has no `shm` and code written against one forks by backend. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…rkers The L3+ child check counted chip and sub children only, so an L4 whose children are local L3 Workers was refused. A next-level child is a forked process that maps a POSIX_SHM backing by name exactly as a chip or sub child does, so it can consume the buffer; `_next_level_shms` counts now, and the error message names all three shapes. `create_buffer` had no test. The new ones cover the gate in each child shape, the childless L3+ refusal, the L2 leaf that needs no child at all, buffer id uniqueness and the single nonce within one incarnation, and that a failing close is reported while its registry entry stays for the cleanup journal to retry. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
42d6ea7 to
29b3527
Compare
`TaskArgs` carried the GM-address-bearing `ChipTensor`, so a sender had to know its receiver's address space: `_rewrite_blob_host_addrs` patched addresses by numeric range and mis-rewrote device pointers that happened to fall inside a registered host range, which a `child_memory` skip then papered over. This flips the wire to the address-free `Tensor` hw-native-sys#1599 froze. A consumer resolves a backing by canonical identity — map-once, exact — and no address crosses a process boundary at all. The wire is the mailbox blob that was already there. `TaskArgs`, `TaskArgsView`, `write_blob` and `read_blob` change what they carry, not how many of them exist; `ChipTensor` survives only inside `ChipStorageTaskArgs`, which the L2 leaf materializes into and hands to runtime.so. `TaskArgsView::tensors` now runs `validate_tensor` on every element it decodes: the element gained wire invariants a decoder can check, and since the bound types expose their fields but not their bytes, this is the only place a mailbox blob becomes a Tensor. `ChipWorker` takes only `ChipStorageTaskArgs`. The blob round trip it sat behind — C++ serialising bytes Python handed straight back to C++ — is gone, with the `TaskArgsView` overloads that existed to unpack it. A control-plane copy names both of its ends. `copy_to`/`copy_from` take a `Buffer` on each side and write the two descriptors into the control frame; the child resolves them through the same `ImportRegistry` its task arguments go through. That replaces both the raw host address, meaningless across a fork once the MAP_HOST pre-attach is gone, and the shm-name path that briefly stood in for it — a second resolution rule for the one thing this ABI exists to resolve. Every receive path keeps its container. A nested next-level child re-exports each backing to a local handle and still hands its orchestration function a `TaskArgs`, scalars included; a remote L3 runner resolves its sidecar descriptors into one too, so an orchestration function can forward the args it was given whichever way it was reached. Remote buffers are allocated through `create_buffer`, so what a runner hands back carries an identity and a descriptor rather than being a bare `SharedMemory` belonging to neither mechanism. `simpler_setup.torch_interop.make_tensor_arg` becomes `make_chip_tensor_arg`: it returns the chip POD where `Worker.make_tensor_arg` returns the wire `Tensor`, and one name for both was the second public meaning rule 13 forbids. `wrap_fork_inherited` callers name their backend explicitly now that hw-native-sys#1703 stopped inferring it from `access`. Each already knew the answer in a comment — the HeapRing backings are MAP_SHARED, and `make_tensor_arg` follows the `shared` it computes: at L2 the consumer is this process, so a write reaches the owner trivially and FORK_COW's contract is the one that would be false there.
`TaskArgs` carried the GM-address-bearing `ChipTensor`, so a sender had to know its receiver's address space: `_rewrite_blob_host_addrs` patched addresses by numeric range and mis-rewrote device pointers that happened to fall inside a registered host range, which a `child_memory` skip then papered over. This flips the wire to the address-free `Tensor` hw-native-sys#1599 froze. A consumer resolves a backing by canonical identity — map-once, exact — and no address crosses a process boundary at all. The wire is the mailbox blob that was already there. `TaskArgs`, `TaskArgsView`, `write_blob` and `read_blob` change what they carry, not how many of them exist; `ChipTensor` survives only inside `ChipStorageTaskArgs`, which the L2 leaf materializes into and hands to runtime.so. `TaskArgsView::tensors` now runs `validate_tensor` on every element it decodes: the element gained wire invariants a decoder can check, and since the bound types expose their fields but not their bytes, this is the only place a mailbox blob becomes a Tensor. `ChipWorker` takes only `ChipStorageTaskArgs`. The blob round trip it sat behind — C++ serialising bytes Python handed straight back to C++ — is gone, with the `TaskArgsView` overloads that existed to unpack it. A control-plane copy names both of its ends. `copy_to`/`copy_from` take a `Buffer` on each side and write the two descriptors into the control frame; the child resolves them through the same `ImportRegistry` its task arguments go through. That replaces both the raw host address, meaningless across a fork once the MAP_HOST pre-attach is gone, and the shm-name path that briefly stood in for it — a second resolution rule for the one thing this ABI exists to resolve. Every receive path keeps its container. A nested next-level child re-exports each backing to a local handle and still hands its orchestration function a `TaskArgs`, scalars included; a remote L3 runner resolves its sidecar descriptors into one too, so an orchestration function can forward the args it was given whichever way it was reached. Remote buffers are allocated through `create_buffer`, so what a runner hands back carries an identity and a descriptor rather than being a bare `SharedMemory` belonging to neither mechanism. `simpler_setup.torch_interop.make_tensor_arg` becomes `make_chip_tensor_arg`: it returns the chip POD where `Worker.make_tensor_arg` returns the wire `Tensor`, and one name for both was the second public meaning rule 13 forbids. `wrap_fork_inherited` callers name their backend explicitly now that hw-native-sys#1703 stopped inferring it from `access`. Each already knew the answer in a comment — the HeapRing backings are MAP_SHARED, and `make_tensor_arg` follows the `shared` it computes: at L2 the consumer is this process, so a write reaches the owner trivially and FORK_COW's contract is the one that would be false there. `docs/buffer-abi.md` is the published page describing this wire, so it moves with it. Its status note said `add_tensor` still takes a `ChipTensor` and that `Tensor` is deliberately absent from `simpler.task_interface`; its scope section said the dispatch wire is not connected and the submit-time checks are unreachable. All three described the tree before this change. What remains absent is now stated positively: the two other allocators, and the endpoint x `address_space` check inside `materialize` — a device backing resolved there yields a pointer meaningful only on its owner chip, and that is enforced today where a task is submitted rather than where it is materialized. Three names outlived what this change removes. A `reserve_slot` comment still offered `create_host_buffer` as an example of an allocator the caller owns, after the last call to it disappeared. `MappedArg.buffer` labelled `FORK_SHM` as copy-on-write; it is the MAP_SHARED one, and that distinction is the whole reason the two fork backends are separate — the corrected wording already sits in `materialize` in the same file. The comments describing `ChipTensor::child_memory` name a field now called `address_space`. The Python keyword stays `child_memory` and now says why: it is the name of a u8 on the remote-L3 tensor wire, which renaming the keyword would not change.
) `TaskArgs` carried the GM-address-bearing `ChipTensor`, so a sender had to know its receiver's address space: `_rewrite_blob_host_addrs` patched addresses by numeric range and mis-rewrote device pointers that happened to fall inside a registered host range, which a `child_memory` skip then papered over. This flips the wire to the address-free `Tensor` #1599 froze. A consumer resolves a backing by canonical identity — map-once, exact — and no address crosses a process boundary at all. The wire is the mailbox blob that was already there. `TaskArgs`, `TaskArgsView`, `write_blob` and `read_blob` change what they carry, not how many of them exist; `ChipTensor` survives only inside `ChipStorageTaskArgs`, which the L2 leaf materializes into and hands to runtime.so. `TaskArgsView::tensors` now runs `validate_tensor` on every element it decodes: the element gained wire invariants a decoder can check, and since the bound types expose their fields but not their bytes, this is the only place a mailbox blob becomes a Tensor. `ChipWorker` takes only `ChipStorageTaskArgs`. The blob round trip it sat behind — C++ serialising bytes Python handed straight back to C++ — is gone, with the `TaskArgsView` overloads that existed to unpack it. A control-plane copy names both of its ends. `copy_to`/`copy_from` take a `Buffer` on each side and write the two descriptors into the control frame; the child resolves them through the same `ImportRegistry` its task arguments go through. That replaces both the raw host address, meaningless across a fork once the MAP_HOST pre-attach is gone, and the shm-name path that briefly stood in for it — a second resolution rule for the one thing this ABI exists to resolve. Every receive path keeps its container. A nested next-level child re-exports each backing to a local handle and still hands its orchestration function a `TaskArgs`, scalars included; a remote L3 runner resolves its sidecar descriptors into one too, so an orchestration function can forward the args it was given whichever way it was reached. Remote buffers are allocated through `create_buffer`, so what a runner hands back carries an identity and a descriptor rather than being a bare `SharedMemory` belonging to neither mechanism. `simpler_setup.torch_interop.make_tensor_arg` becomes `make_chip_tensor_arg`: it returns the chip POD where `Worker.make_tensor_arg` returns the wire `Tensor`, and one name for both was the second public meaning rule 13 forbids. `wrap_fork_inherited` callers name their backend explicitly now that #1703 stopped inferring it from `access`. Each already knew the answer in a comment — the HeapRing backings are MAP_SHARED, and `make_tensor_arg` follows the `shared` it computes: at L2 the consumer is this process, so a write reaches the owner trivially and FORK_COW's contract is the one that would be false there. `docs/buffer-abi.md` is the published page describing this wire, so it moves with it. Its status note said `add_tensor` still takes a `ChipTensor` and that `Tensor` is deliberately absent from `simpler.task_interface`; its scope section said the dispatch wire is not connected and the submit-time checks are unreachable. All three described the tree before this change. What remains absent is now stated positively: the two other allocators, and the endpoint x `address_space` check inside `materialize` — a device backing resolved there yields a pointer meaningful only on its owner chip, and that is enforced today where a task is submitted rather than where it is materialized. Three names outlived what this change removes. A `reserve_slot` comment still offered `create_host_buffer` as an example of an allocator the caller owns, after the last call to it disappeared. `MappedArg.buffer` labelled `FORK_SHM` as copy-on-write; it is the MAP_SHARED one, and that distinction is the whole reason the two fork backends are separate — the corrected wording already sits in `materialize` in the same file. The comments describing `ChipTensor::child_memory` name a field now called `address_space`. The Python keyword stays `child_memory` and now says why: it is the name of a u8 on the remote-L3 tensor wire, which renaming the keyword would not change.
Follow-up to #1599: the gate the ABI rests on
#1599 froze the byte layout and shipped
validate_tensoras the receive-side gate — every laterboundary is meant to run it, and after the wire flip it runs on every argument that arrives. Three
things in it do not hold up, and a frozen ABI is the wrong place to leave them. None of this
changes a field, an offset, or an enum value: every
static_assertinbuffer.his untouched.①
validate_tensoraccepts a view that spans exabytes over a 4-byte backingtensor_extent_bytessummed(shapes[i]-1)*strides[i]into auint64. Both fields areu32, so asingle product already reaches ~2^64, and the multiply by the element size overflows on top of that.
The extent wrapped to a small value and the bound check passed. One dimension is enough, and it
is reachable straight from Python:
The arithmetic saturates, and an extent that lands on the sentinel is refused by name rather
than compared against
nbytes— otherwise a descriptor claiming an absurdnbytesbuys the viewback. This breaks design invariant 4 (
view footprint ⊆ nbytes), which the design doc lists asexecutable.
tensors_overlapinherited the same wrap, and there the consequence is worse than a rejectedargument: two fully overlapping views compare as disjoint, i.e. a missed dependency edge — the
old-address-key failure mode the identity key exists to remove. Its end offsets saturate too.
The 4096-iteration arbitrary-bytes pass did not catch this because its survivor assertions covered
magic/generation/body_len/ndimsbut not the footprint invariant — the one most atrisk. It asserts it now.
②
create_bufferrefuses an L4 whose children are local L3 WorkersThe L3+ child check counted
_chip_shmsand_sub_shmsonly. A next-level Worker child is a forkedprocess that maps a
POSIX_SHMbacking by name exactly as a chip or sub child does, so it canconsume the buffer.
_next_level_shmscounts now.create_bufferhad no test at all;tests/ut/py/test_worker/test_create_buffer.pycovers the gatein each child shape, the childless refusal, the L2 leaf, id uniqueness within one incarnation, and
the per-buffer best-effort release (a failing close is reported, its entry kept for the journal to
retry).
③
wrap_fork_inheritedinferred the backend fromaccessFORK_SHM if access != READ else FORK_COW. The two are opposite kernel write semantics, not twospellings of one grant — putting the physical semantics in
backend_kindrather than leaving it tobe inferred from an orthogonal field is the reason they are separate tags. Inferring it makes a
read-only
MAP_SHAREDbacking inexpressible: it is taggedFORK_COW, whose READ-only rule thenlocks it there. The caller holds the mmap and is the only party that knows which it is, so it states
the tag; the default pair stays the safe one. The function had no caller, so nothing depends on the
old signature.
④ Text that describes machinery which does not exist
ImportRegistry.materializepointed a caller holding raw bytes atBufferDescriptor.unpack, andtests/ut/py/test_buffer.pysaid it pins apack/unpackround trip. Neither exists —withholding the encoding is exactly what makes construction the only way in.
validate_tensor's comment claimed it stands behind materialization. It does not:ImportRegistry.materializetakes an already-decoded descriptor and adds no endpoint check. Thecomment now says so and points at the separate change.
docs/buffer-abi.mdnamed the wire typesimpler.task_interface.Tensor, contradicting the samepage's status note, the module, and
test_wire_tensor_stays_off_the_public_submit_surface.h.shm.bufsample now records that it is transitional and why: byte access belongs on theview, since a device backing has no
shmand code written against one forks by backend.BufferDescriptor::operator==bounded amemcmpby an unvalidatedbody_len; clamped. Notreachable today (construction validates), but it is the only length field in the header that
bounds a read, and the cutover is what makes raw decoded descriptors exist.
Verification
cpputtest_buffer— 15/15, including the two new overflow / overlap cases and thestrengthened arbitrary-bytes pass
pyutfull suite — 1142 passed, 13 skippedtests/ut/py/test_buffer.py+ the newtest_create_buffer.py— 31 passedDeliberately not here
Each is a separate concern, and the first two are the ones I would take next:
address_spacematrix.materializestill returns a device VA to a hostendpoint. That is a behavioural gate with its own test surface, not a comment fix.
wrap_vmm_window,remote_sidecar_tensor,host_ptr_nbytesandtensors_overlapstill have no caller. Whether they survive the freeze is the author's call, notsomething to settle in a fix PR.
BufferDescriptor/Tensorbind value equality without value hashing, and their__eq__is notnb::is_operator(), so a comparison against a foreign type raisesTypeErrorinstead ofreturning
NotImplemented.remote_sidecar_tensorbuildsowner_instance_idfromowner_worker_id, against the full-widthrandom draw its own minting function requires.
🤖 Generated with Claude Code