Add: the BufferHandle/Tensor wire ABI, its codec, and create_buffer - #1599
Add: the BufferHandle/Tensor wire ABI, its codec, and create_buffer#1599YunjiQin wants to merge 2 commits into
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:
📝 WalkthroughWalkthroughIntroduces a typed BufferHandle and tensor-view ABI with canonical identities, backend descriptors, validation, blob serialization, Python materialization, worker-owned shared-memory allocation, bindings, documentation, and C++/Python tests. ChangesBufferHandle ABI
Estimated code review effort: 4 (Complex) | ~60 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
tests/ut/cpp/types/test_buffer_handle.cpp (1)
71-82: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
EnumValuesAreFrozendoesn't pinBackendKind::FORK_COW.Every other
BackendKindvalue (0-4) is checked here butFORK_COW = 5is missing, even though the test's stated purpose is to "pin the sizes, enum values, and the blob codec from the outside". Trivial to add for parity with the rest of this test.✅ Proposed addition
EXPECT_EQ(static_cast<uint8_t>(BackendKind::DEVICE_MALLOC), 4); + EXPECT_EQ(static_cast<uint8_t>(BackendKind::FORK_COW), 5); }🤖 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/ut/cpp/types/test_buffer_handle.cpp` around lines 71 - 82, Add an assertion in BufferHandleAbi.EnumValuesAreFrozen that verifies BackendKind::FORK_COW converts to uint8_t value 5, preserving the existing checks for BackendKind values 0–4.src/common/task_interface/buffer_handle.h (1)
313-328: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low value
BufferRefBlobView::ref(i)has no bounds check oni.
ref(i)memcpy'ssizeof(BufferRef)bytes atref_bytes + i * sizeof(BufferRef)with no check thati < ref_count(ori >= 0). Current callers (bufferref_blob_descriptors/refs/scalarsintask_interface.cpp) all loopi < view.ref_count, so this isn't exploitable today, but the function is documented as the single shared extraction point for three receive boundaries, so a future caller passing an unchecked index would read out of the validated region.Optional: add an
i < ref_countassert/throw insideref()itself for defense-in-depth.🤖 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 `@src/common/task_interface/buffer_handle.h` around lines 313 - 328, Add bounds validation to BufferRefBlobView::ref before computing the byte offset or calling memcpy, rejecting indices below zero or greater than or equal to ref_count with the established assertion or exception mechanism. Preserve the existing BufferRef validation and extraction behavior for valid indices.
🤖 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 `@docs/buffer-handle-abi.md`:
- Around line 174-186: Add a VMM_WINDOW row to the Backends table, documenting
that it materializes as the device VA carved by allocate_domain and is used for
communication-domain window buffers. Keep it alongside DEVICE_MALLOC and
preserve the table’s wire-ABI completeness statement.
In `@python/simpler/buffer_handle.py`:
- Around line 239-242: Update the layout comment immediately above
_BUFFER_REF_TAIL to use the actual wire sizes: BufferRef 144 B and
BufferHandleDescriptor 88 B, while preserving the existing field description and
size assertion.
In `@src/common/task_interface/buffer_handle.h`:
- Around line 80-96: Update validate_buffer_ref to reject any BufferRef whose
backend_kind is FORK_COW unless access is exactly READ, preserving the existing
address_space/backend_kind validation for all other combinations. Add or extend
C++ validation tests to cover FORK_COW with WRITE and READWRITE access as
rejected, while retaining acceptance for FORK_COW with READ.
---
Nitpick comments:
In `@src/common/task_interface/buffer_handle.h`:
- Around line 313-328: Add bounds validation to BufferRefBlobView::ref before
computing the byte offset or calling memcpy, rejecting indices below zero or
greater than or equal to ref_count with the established assertion or exception
mechanism. Preserve the existing BufferRef validation and extraction behavior
for valid indices.
In `@tests/ut/cpp/types/test_buffer_handle.cpp`:
- Around line 71-82: Add an assertion in BufferHandleAbi.EnumValuesAreFrozen
that verifies BackendKind::FORK_COW converts to uint8_t value 5, preserving the
existing checks for BackendKind values 0–4.
🪄 Autofix (Beta)
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: d442b3d1-f5cd-43eb-b613-298f4efd8dd7
📒 Files selected for processing (12)
docs/README.mddocs/buffer-handle-abi.mdmkdocs.ymlpython/bindings/task_interface.cpppython/simpler/buffer_handle.pypython/simpler/worker.pysrc/common/task_interface/buffer_handle.hsrc/common/task_interface/data_type.hsrc/common/task_interface/tensor.htests/ut/cpp/CMakeLists.txttests/ut/cpp/types/test_buffer_handle.cpptests/ut/py/test_buffer_handle.py
💤 Files with no reviewable changes (1)
- src/common/task_interface/tensor.h
| ## Backends | ||
|
|
||
| | Backend | Materializes to | Used for | | ||
| | ------- | --------------- | -------- | | ||
| | `POSIX_SHM` | a named shm mapped into the consumer | `create_buffer` | | ||
| | `FORK_SHM` | the same VA, no map | a pre-fork `MAP_SHARED` host buffer (e.g. `share_memory_()`), writable from the child | | ||
| | `FORK_COW` | the same VA, no map | a pre-fork plain host buffer: copy-on-write, so **READ only** | | ||
| | `DEVICE_MALLOC` | the device pointer, no map (chip-local) | device memory on a next-level worker | | ||
| | `REMOTE_SIDECAR` | resolved via the remote transport | an arg to a remote L3; the descriptor rides in the sidecar | | ||
|
|
||
| The enum is complete because it is wire ABI: a value must mean the same thing in | ||
| every process of a run, so backends are declared here even where the code that | ||
| allocates or materializes them arrives later. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Backends table omits VMM_WINDOW.
The table lists POSIX_SHM, FORK_SHM, FORK_COW, DEVICE_MALLOC, REMOTE_SIDECAR but not VMM_WINDOW, even though wrap_vmm_window() (device VA wrap, address_space=DEVICE) is already implemented in python/simpler/buffer_handle.py as part of this same PR. Worth adding a row (e.g. "the device VA carved by allocate_domain | comm-domain window buffers") for completeness alongside DEVICE_MALLOC.
🤖 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 `@docs/buffer-handle-abi.md` around lines 174 - 186, Add a VMM_WINDOW row to
the Backends table, documenting that it materializes as the device VA carved by
allocate_domain and is used for communication-domain window buffers. Keep it
alongside DEVICE_MALLOC and preserve the table’s wire-ABI completeness
statement.
| # BufferRef (272 B): BufferHandleDescriptor(216) + byte_offset u64, ndims u32, shapes[MAX] u32, | ||
| # strides[MAX] u32, dtype u8, _pad[3]. | ||
| _BUFFER_REF_TAIL = struct.Struct(f"<QI{MAX_TENSOR_DIMS}I{MAX_TENSOR_DIMS}IB3x") | ||
| assert BUFFER_HANDLE_DESCRIPTOR_BYTES + _BUFFER_REF_TAIL.size == BUFFER_REF_BYTES, "BufferRef layout drift" |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Stale/incorrect size comment.
The comment says "BufferRef (272 B): BufferHandleDescriptor(216)", but the actual wire sizes (per buffer_handle.h's static_asserts and the assert two lines below) are BufferRef=144B and BufferHandleDescriptor=88B — matching every other comment/doc in this PR. Looks like a leftover from an earlier draft of the layout.
✏️ Proposed fix
-# BufferRef (272 B): BufferHandleDescriptor(216) + byte_offset u64, ndims u32, shapes[MAX] u32,
-# strides[MAX] u32, dtype u8, _pad[3].
+# BufferRef (144 B): BufferHandleDescriptor(88) + byte_offset u64, ndims u32, shapes[MAX] u32,
+# strides[MAX] u32, dtype u8, _pad[3].📝 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.
| # BufferRef (272 B): BufferHandleDescriptor(216) + byte_offset u64, ndims u32, shapes[MAX] u32, | |
| # strides[MAX] u32, dtype u8, _pad[3]. | |
| _BUFFER_REF_TAIL = struct.Struct(f"<QI{MAX_TENSOR_DIMS}I{MAX_TENSOR_DIMS}IB3x") | |
| assert BUFFER_HANDLE_DESCRIPTOR_BYTES + _BUFFER_REF_TAIL.size == BUFFER_REF_BYTES, "BufferRef layout drift" | |
| # BufferRef (144 B): BufferHandleDescriptor(88) + byte_offset u64, ndims u32, shapes[MAX] u32, | |
| # strides[MAX] u32, dtype u8, _pad[3]. | |
| _BUFFER_REF_TAIL = struct.Struct(f"<QI{MAX_TENSOR_DIMS}I{MAX_TENSOR_DIMS}IB3x") | |
| assert BUFFER_HANDLE_DESCRIPTOR_BYTES + _BUFFER_REF_TAIL.size == BUFFER_REF_BYTES, "BufferRef layout drift" |
🤖 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_handle.py` around lines 239 - 242, Update the layout
comment immediately above _BUFFER_REF_TAIL to use the actual wire sizes:
BufferRef 144 B and BufferHandleDescriptor 88 B, while preserving the existing
field description and size assertion.
| // Materialization backend of a handle. The consumer resolves a BufferRef to a local address via the | ||
| // import registry keyed by canonical identity; this tag selects how. REMOTE_SIDECAR is reserved for | ||
| // P2 and rejected on decode in P1. Values are frozen; 6.. reserved (unknown tag => reject). | ||
| // | ||
| // FORK_SHM and FORK_COW materialize identically — the body is a base VA the child already has, | ||
| // inherited across the fork — but the kernel's write semantics are opposite, so they are separate | ||
| // tags rather than one tag plus a hint. A child's write to a MAP_SHARED page lands in the physical | ||
| // page the parent reads; a write to a copy-on-write page splits it into a private copy the parent | ||
| // never sees, silently. FORK_COW therefore grants READ only, and that is enforced on decode. | ||
| enum class BackendKind : uint8_t { | ||
| FORK_SHM = 0, | ||
| POSIX_SHM = 1, | ||
| VMM_WINDOW = 2, | ||
| REMOTE_SIDECAR = 3, | ||
| DEVICE_MALLOC = 4, | ||
| FORK_COW = 5, | ||
| }; |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
validate_buffer_ref doesn't enforce the FORK_COW read-only invariant it documents.
The BackendKind doc comment states "FORK_COW therefore grants READ only, and that is enforced on decode." But validate_buffer_ref — described as "the single implementation behind all three trust boundaries" (builder, blob decode, materialization) — never checks access against backend_kind; it only gates address_space × backend_kind. A wire BufferRef with backend_kind=FORK_COW and access=WRITE/READWRITE currently passes this validator.
Today this is latent: python/simpler/buffer_handle.py's BufferHandleDescriptor.__post_init__ independently re-checks this on every Python-side unpack, so the two consumption paths shown in this PR (mapped_args_from_blob → Tensor.unpack → BufferHandleDescriptor.unpack) are still protected. But a pure-C++ consumer (the builder TaskArgs.add_tensor, or any future device-side/dispatch consumer explicitly called out in this validator's own docstring) gets no such protection, and the header's own claim about validate_buffer_ref is factually wrong for this field. The C++ unit test suite also has no case for it (EnumValuesAreFrozen/test_decode_rejects_corrupted_field equivalents never exercise FORK_COW + non-READ access).
Since the PR explicitly frames this validator as needing "isolated review before consumer migration," this is worth closing now rather than after a future consumer starts relying on it.
🔒 Proposed fix
// address_space x backend_kind capability gate. REMOTE_SIDECAR is legal in either space.
const auto backend = static_cast<BackendKind>(h.backend_kind);
const bool device_space = h.address_space == static_cast<uint8_t>(AddressSpace::DEVICE);
if (backend != BackendKind::REMOTE_SIDECAR) {
const bool device_backend = backend == BackendKind::VMM_WINDOW || backend == BackendKind::DEVICE_MALLOC;
if (device_backend != device_space) reject("invalid Tensor: unsupported address_space x backend_kind");
}
+
+ // FORK_COW is copy-on-write: a write splits the page into a private copy the owner never sees,
+ // silently. Grant READ only, mirroring BufferHandleDescriptor.__post_init__ in
+ // simpler.buffer_handle so the two decode paths cannot drift apart.
+ if (backend == BackendKind::FORK_COW && h.access != static_cast<uint8_t>(AccessMode::READ)) {
+ reject("invalid Tensor: FORK_COW backend grants READ only");
+ }Also applies to: 236-273
🤖 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 `@src/common/task_interface/buffer_handle.h` around lines 80 - 96, Update
validate_buffer_ref to reject any BufferRef whose backend_kind is FORK_COW
unless access is exactly READ, preserving the existing
address_space/backend_kind validation for all other combinations. Add or extend
C++ validation tests to cover FORK_COW with WRITE and READWRITE access as
rejected, while retaining acceptance for FORK_COW with READ.
b2e2bc9 to
57a89fd
Compare
At L3 and above a task argument is currently a raw pointer plus a `child_memory` bool, which cannot say which process the address is valid in, which backing it belongs to, or whether the holder may write it. This adds the typed ABI that replaces it, with nothing consuming it yet. Three wire types in one header, each pinned by static_assert on size and every field offset, so a layout drift fails the build rather than the run: - `CanonicalIdentity` (32 B) — owner_instance_id + buffer_id + generation. The key an owner registry and every consumer import cache share. Fixed-length with no length field, so hashing and comparison cannot read past it, and padding is excluded from equality so two decodes of one backing key alike. - `BufferHandleDescriptor` (88 B) — the self-describing backing: address space, access, backend kind, extent, and a length-delimited backend body. - `BufferRef` (144 B) — the blob element: the descriptor embedded whole plus a strided view. No materialized address; a consumer resolves it on receipt. `validate_buffer_ref` is the single validator every trust boundary runs. It bounds `body_len` and `ndims`, checks the address_space × backend matrix, requires strides > 0 and a known dtype, and rejects a view that extends past its backing. The blob codec lives beside the types it encodes, and its reader validates each element as it is extracted, so a decoder cannot hand out an element nobody checked. `Worker.create_buffer` allocates a POSIX-shm backing under a fresh identity; the Worker owns every handle it created and unlinks the backings on close(). `ImportRegistry` is the consumer half — map-once, keyed by identity. Also bounds-checks `get_element_size`, which indexed its table with a raw u8 that now reaches it from the wire, and moves `MAX_TENSOR_DIMS` to `data_type.h` so the view types can share it with `Tensor`. The ABI is frozen once it ships, which is why it lands on its own: the byte layout, the enum values, and the validator are what this change is asking to have reviewed. Naming a `Tensor` and submitting it, the per-consumer materialization split, and the submit-time guards follow separately.
YunjiQin
changed the title
|
Went through the CodeRabbit findings — all five held up against the current code, none was a false positive. Fixed in 9fe7bb9. The one that mattered:
On the docstring-coverage warning (32% vs 80%): added docstrings where they state a contract that is not evident from the signature — the two wire enums, Verified after the fixes: pyut 913 passed, cpput 67/67, a2a3sim scene tests 43 passed. |
`validate_buffer_ref` is the one gate every receive boundary runs, but the FORK_COW rule lived only in the Python mirror's descriptor constructor. A descriptor arriving over the wire as FORK_COW with a write grant therefore passed the C++ decode: the consumer's first write would split the page into a private copy and the owner would observe nothing, with no error anywhere. Both mirrors of the ABI now reject it, and the test covers the WRITE and READWRITE grants as well as the FORK_SHM case that stays legal. `BufferRefBlobView::ref` also bounds-checks its index. Every caller loops to `ref_count` today, so nothing reaches past the validated region, but the method is the documented extraction point for three separate consumers and a future one has no way to see that the bound is the caller's job. Also: pin `FORK_COW = 5` alongside the other enumerators, which the frozen-value test had skipped; give the `BufferRef` layout comment the post-slimming sizes it describes (144 / 88, not 272 / 216 — the assertion below it was already right); add the missing `VMM_WINDOW` row to the backend table, which listed five of the six wire values; and document the two wire enums plus the handle and registry lifecycle calls, where the contract is not evident from the signature.
1 participant
PR-1 of the P1-B split: the wire ABI on its own
Split out of #1453 at reviewer request — that PR is ~3× the repo's oversize threshold for core
churn, and this is the piece that most needs to be read in its own context: a frozen wire ABI.
Field sets, enum values, and byte offsets are hard to change once they ship, and in #1453 they sat
underneath a large consumer migration.
One commit,
+2218/−3across 12 files — almost pure addition and zero behavior change. Nothingin the tree dispatches a
Tensoryet.What lands here
Three wire types in
src/common/task_interface/buffer_handle.h,each pinned by
static_asserton size and every field offset so a layout drift fails the build:CanonicalIdentityowner_instance_id+buffer_id+generation. The key the owner registry and every consumer import cache share. Fixed-length, no length field, padding excluded from equality/hash.BufferHandleDescriptorBufferRefOne validator.
validate_buffer_refboundsbody_lenandndims, checks theaddress_space × backend_kindmatrix, requires strides > 0 and a known dtype, and rejects a viewextending past its backing. The blob codec lives beside the types it encodes and validates each
element on extraction, so no decoder can hand out an unchecked element.
Owner and consumer halves.
Worker.create_bufferallocates a POSIX-shm backing under a freshidentity and unlinks it on
close();ImportRegistryis the consumer side (map-once, by identity).Also: bounds-checks
get_element_size(it indexed its table with a rawu8that now arrivesfrom the wire — a genuine OOB read), and moves
MAX_TENSOR_DIMSintodata_type.hso the viewtypes can share it.
Reviewing this
The three questions worth spending time on:
owner_worker_path_idis deliberatelydiagnostic-only and takes part in no routing, visibility, or identity decision.
validate_buffer_refcomplete? It is the only thing standing between a corrupted orhostile blob and a pointer dereference.
test_decode_rejects_corrupted_fieldwalks afield-by-field corruption matrix through the decoder;
test_decode_rejects_arbitrary_bytesthrows 256 random buffers plus all-zeros and all-ones at it — none may be accepted, none may
crash.
owner_instance_idis a per-incarnation random draw, so abuffer_idreused by a later process cannot collide with a live identity.generationstartsat 1 and 0 is rejected on decode.
The design rationale for keeping
TensorandChipTensoras two types — with the two rejectedunifications and why — is in
docs/buffer-handle-abi.md.What is deliberately NOT here
create_bufferships an unwired capability: you can allocate a handle, name a view over it, andmaterialize a blob, but
TaskArgsstill carriesChipTensor, so nothing dispatches aTensoryet.That is the intended cut — submitting one, the per-consumer split (chip leaf / Python sub-worker /
L4→L3 re-export), the device-memory allocators (
alloc_shared_tensor,alloc_child_tensor), andthe submit-time guards land with the dispatch wire.
docs/buffer-handle-abi.mddescribes the model end to end, including the parts this PR does notship — the submit path, the three-way consumer split, the submit-time checks. Its "Scope / status"
section says plainly which of those the tree has and which the wire flip adds, so the page is a
complete reference without implying the whole thing is live. Two other places name the device POD
ChipTensor, which is the model's term for it; its Python binding is still exported asTensorandis renamed with the wire flip — the page says so where it matters.
The Python binding surface for the wire type (
materialize_bufferref_blob) and theTensor→ChipTensorbinding rename were prepared and then deliberately held back from this PR,to keep it to the ABI and its own tests.
Verification
test_buffer_handle(layout, validation, codec).Follow-ups from the same split: capability matrix +
access ⊆ grantedguards; device-memoryownership Orchestrator→Worker; deleting the
create_host_buffer/MAP_HOSTsubsystem; and the wireflip itself (#1453, to be rebased onto this).