Update: cut task args over to the self-describing Tensor wire ABI - #1729
Merged
Conversation
|
Important Review skippedToo many files! This PR contains 126 files, which is 26 over the limit of 100. To get a review, reduce the PR to 100 files or fewer by splitting it into smaller PRs or changing its base branch. Upgrade to a paid plan to raise the limit. This review couldn't start because sufficient usage credits or metered capacity aren't available. Add credits or update usage-based reviews in the billing tab, then retry. ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (126)
You can disable this status message by setting the 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 |
YunjiQin
force-pushed
the
p1-b-wire-cutover
branch
from
August 7, 2026 03:54
69938e2 to
c9f5db5
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. `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.
ChaoWao
force-pushed
the
p1-b-wire-cutover
branch
from
August 7, 2026 14:40
c9f5db5 to
2822c4d
Compare
Collaborator
|
Rebased onto The rebase was not mechanical — one silent bug it would have introduced
In The rest:
Four fixes
Verified locally on the rebased commit
Not done here (deliberate)
One process note: CodeRabbit skipped this PR entirely ("125 files exceed the limit of 100"), so it has had no automated review. |
This was referenced Aug 7, 2026
ChaoWao
added a commit
that referenced
this pull request
Aug 8, 2026
simpler_setup.Tensor (a NamedTuple test-arg spec used by TaskArgsBuilder) and simpler.task_interface.Tensor (the wire-ABI struct PR #1729 wired up) now name two unrelated public types the same thing in two importable namespaces, violating codestyle.md rule 13 (public C++/Python types share one canonical name and semantics). #1729 landed with zero Python call sites on the ABI Tensor, so renaming the long-established, 119-call-site test-tooling type is the low-risk direction: it is a mechanical rename with no behavior change, versus fighting the design docs' canonical name for the wire element. Renamed the class in simpler_setup/scene_test.py, its export in simpler_setup/__init__.py, and every constructor/isinstance call site across examples/, tests/, and docs/testing.md, docs/user/README.md, .claude/rules/project-layout.md. docs/user/README.md's "one name to watch" caveat about the collision is removed since it no longer applies. Verified: ruff check/format clean, pyut 1251 passed / 13 skipped / 0 failed, and a live a2a3sim run through TaskArgsBuilder/compute_golden confirms the renamed spec still round-trips through the harness.
4 tasks
ChaoWao
added a commit
to ChaoWao/simpler-fork
that referenced
this pull request
Aug 8, 2026
Two gaps in the buffer identity / materialize path, both left open by hw-native-sys#1729. owner_instance_id was minted before the Worker it names existed: Worker.__init__ minted the nonce at Python object construction time -- before add_worker(), before any os.fork(), before init(). For a next-level child Worker that nonce was fixed in the parent process and only later copied (via fork COW) into the child that actually becomes its owner, so nothing tied the identity to a real fork having happened. Worker.init() now re-mints owner_instance_id right after its NEW -> INITIALIZING transition. A next-level child's init() runs only inside the process that was forked to host it (_start_hierarchical calls inner.init(...) only in the os.fork() child branch), so the nonce a Worker's buffers and endpoint identity actually use is never older than its own fork. __init__ keeps its mint as a fallback: several test files build a "ready" Worker by setting _lifecycle directly without ever calling init(), and still need a well-formed nonce. materialize() trusted a raw device pointer with no endpoint check: ImportRegistry.materialize()'s DEVICE_MALLOC/VMM_WINDOW branch decoded the descriptor body straight into a pointer and handed it back -- its own docstring said as much: "a DEVICE backing resolved here yields a device pointer, which is only meaningful on its owner chip. The endpoint x address_space matrix is a separate change; until it lands, that invariant rests on the caller." Submit-time checks already reject a device tensor before dispatch, but nothing stopped a caller that reaches materialize() directly. ImportContext(is_host_endpoint, owning_chip_instance_id) closes that: materialize() now refuses any DEVICE backing on a host endpoint outright, and on a device endpoint refuses one whose owner_instance_id doesn't match the nonce that endpoint was set up to serve. Wired into the three ImportRegistry() construction sites -- _sub_worker_loop (host), _run_chip_main_loop (device, given the parent's owner_instance_id at fork), and _run_l2_materialized (a device_ids-bearing L2 leaf, checked against its own nonce). Added test_host_endpoint_materialize_refuses_a_device_tensor_directly and test_chip_materialization_refuses_a_foreign_chips_device_tensor, mirroring hw-native-sys#1729's existing sub-worker analog but exercising materialize() itself (bypassing submit) for both rows of the endpoint matrix. Updated three existing test_buffer.py cases and two test_host_worker.py harness call sites to pass ImportContext / the new owner_instance_id parameter. Verified: pytest tests/ut 1279 passed / 13 skipped / 0 failed; ruff check/format and pyright clean; a real a2a3 onboard run (2 devices) across host_build_graph and tensormap_and_ringbuffer examples exercised the touched chip-fork path with no failures.
ChaoWao
added a commit
to ChaoWao/simpler-fork
that referenced
this pull request
Aug 8, 2026
Two gaps in the buffer identity / materialize path, both left open by hw-native-sys#1729. owner_instance_id was minted before the Worker it names existed: Worker.__init__ minted the nonce at Python object construction time -- before add_worker(), before any os.fork(), before init(). For a next-level child Worker that nonce was fixed in the parent process and only later copied (via fork COW) into the child that actually becomes its owner, so nothing tied the identity to a real fork having happened. Worker.init() now re-mints owner_instance_id right after its NEW -> INITIALIZING transition. A next-level child's init() runs only inside the process that was forked to host it (_start_hierarchical calls inner.init(...) only in the os.fork() child branch), so the nonce a Worker's buffers and endpoint identity actually use is never older than its own fork. __init__ keeps its mint as a fallback: several test files build a "ready" Worker by setting _lifecycle directly without ever calling init(), and still need a well-formed nonce. materialize() trusted a raw device pointer with no endpoint check: ImportRegistry.materialize()'s DEVICE_MALLOC/VMM_WINDOW branch decoded the descriptor body straight into a pointer and handed it back -- its own docstring said as much: "a DEVICE backing resolved here yields a device pointer, which is only meaningful on its owner chip. The endpoint x address_space matrix is a separate change; until it lands, that invariant rests on the caller." Submit-time checks already reject a device tensor before dispatch, but nothing stopped a caller that reaches materialize() directly. ImportContext(is_host_endpoint, owning_chip_instance_id) closes that: materialize() now refuses any DEVICE backing on a host endpoint outright, and on a device endpoint refuses one whose owner_instance_id doesn't match the nonce that endpoint was set up to serve. Wired into the three ImportRegistry() construction sites -- _sub_worker_loop (host), _run_chip_main_loop (device, given the parent's owner_instance_id at fork), and _run_l2_materialized (a device_ids-bearing L2 leaf, checked against its own nonce). This check is Worker-grained, not chip-grained: owner_instance_id is minted once per Worker incarnation, not once per chip, so a Worker with more than one entry in device_ids gives every one of its chip children the same nonce -- the wire BufferDescriptor has no field that distinguishes sibling chips (owner_worker_id is host-side-only free/copy provenance, never serialized). A DEVICE backing minted for chip 0 of such a Worker therefore also passes this check on sibling chip 1. It still rejects a different Worker's device buffer, and any host endpoint outright; the exact-chip half of the endpoint x address_space matrix for a multi-device Worker stays with the existing submit-time (target_worker_id, ptr) check in orchestrator.py, which this backstop complements rather than replaces. Closing that gap fully would need either a wire ABI field the frozen P1-A identity doesn't have, or a live per-chip pointer allowlist -- both bigger than a materialize-time backstop should take on; ImportContext's docstring states this bound explicitly. Added test_host_endpoint_materialize_refuses_a_device_tensor_directly and test_chip_materialization_refuses_a_foreign_chips_device_tensor, mirroring hw-native-sys#1729's existing sub-worker analog but exercising materialize() itself (bypassing submit) for both rows of the endpoint matrix. Updated three existing test_buffer.py cases and two test_host_worker.py harness call sites to pass ImportContext / the new owner_instance_id parameter. Verified: pytest tests/ut 1279 passed / 13 skipped / 0 failed; ruff check/format and pyright clean; a real a2a3 onboard run (2 devices) across host_build_graph and tensormap_and_ringbuffer examples exercised the touched chip-fork path with no failures.
ChaoWao
added a commit
that referenced
this pull request
Aug 8, 2026
…1744) Two gaps in the buffer identity / materialize path, both left open by #1729. owner_instance_id was minted before the Worker it names existed: Worker.__init__ minted the nonce at Python object construction time -- before add_worker(), before any os.fork(), before init(). For a next-level child Worker that nonce was fixed in the parent process and only later copied (via fork COW) into the child that actually becomes its owner, so nothing tied the identity to a real fork having happened. Worker.init() now re-mints owner_instance_id right after its NEW -> INITIALIZING transition. A next-level child's init() runs only inside the process that was forked to host it (_start_hierarchical calls inner.init(...) only in the os.fork() child branch), so the nonce a Worker's buffers and endpoint identity actually use is never older than its own fork. __init__ keeps its mint as a fallback: several test files build a "ready" Worker by setting _lifecycle directly without ever calling init(), and still need a well-formed nonce. materialize() trusted a raw device pointer with no endpoint check: ImportRegistry.materialize()'s DEVICE_MALLOC/VMM_WINDOW branch decoded the descriptor body straight into a pointer and handed it back -- its own docstring said as much: "a DEVICE backing resolved here yields a device pointer, which is only meaningful on its owner chip. The endpoint x address_space matrix is a separate change; until it lands, that invariant rests on the caller." Submit-time checks already reject a device tensor before dispatch, but nothing stopped a caller that reaches materialize() directly. ImportContext(is_host_endpoint, owning_chip_instance_id) closes that: materialize() now refuses any DEVICE backing on a host endpoint outright, and on a device endpoint refuses one whose owner_instance_id doesn't match the nonce that endpoint was set up to serve. Wired into the three ImportRegistry() construction sites -- _sub_worker_loop (host), _run_chip_main_loop (device, given the parent's owner_instance_id at fork), and _run_l2_materialized (a device_ids-bearing L2 leaf, checked against its own nonce). This check is Worker-grained, not chip-grained: owner_instance_id is minted once per Worker incarnation, not once per chip, so a Worker with more than one entry in device_ids gives every one of its chip children the same nonce -- the wire BufferDescriptor has no field that distinguishes sibling chips (owner_worker_id is host-side-only free/copy provenance, never serialized). A DEVICE backing minted for chip 0 of such a Worker therefore also passes this check on sibling chip 1. It still rejects a different Worker's device buffer, and any host endpoint outright; the exact-chip half of the endpoint x address_space matrix for a multi-device Worker stays with the existing submit-time (target_worker_id, ptr) check in orchestrator.py, which this backstop complements rather than replaces. Closing that gap fully would need either a wire ABI field the frozen P1-A identity doesn't have, or a live per-chip pointer allowlist -- both bigger than a materialize-time backstop should take on; ImportContext's docstring states this bound explicitly. Added test_host_endpoint_materialize_refuses_a_device_tensor_directly and test_chip_materialization_refuses_a_foreign_chips_device_tensor, mirroring #1729's existing sub-worker analog but exercising materialize() itself (bypassing submit) for both rows of the endpoint matrix. Updated three existing test_buffer.py cases and two test_host_worker.py harness call sites to pass ImportContext / the new owner_instance_id parameter. Verified: pytest tests/ut 1279 passed / 13 skipped / 0 failed; ruff check/format and pyright clean; a real a2a3 onboard run (2 devices) across host_build_graph and tensormap_and_ringbuffer examples exercised the touched chip-fork path with no failures.
6 tasks
ChaoWao
added a commit
that referenced
this pull request
Aug 8, 2026
PR #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 #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.
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
2 participants
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.
The wire flip
#1599 froze the address-free
Tensor; nothing dispatched one. This makesTaskArgscarry it.
Before,
TaskArgscarried the GM-address-bearingChipTensor, so a sender had to knowits receiver's address space.
_rewrite_blob_host_addrspatched addresses by numericrange and mis-rewrote device pointers that happened to fall inside a registered host
range — patched by adding a
child_memoryskip whose own comment recorded the hazard.A consumer now resolves a backing by canonical identity: map-once, exact, and no address
crosses a process boundary at all.
User guide:
docs/buffer-abi.md.One wire, not two
The wire is the mailbox blob that was already in
task_args.h.TaskArgs,TaskArgsView,write_blobandread_blobchange what they carry, not how many ofthem exist;
ChipTensorsurvives only insideChipStorageTaskArgs, which the L2 leafmaterializes into and hands to
runtime.so.TaskArgsView::tensorsnow runsvalidate_tensoron every element it decodes. That isload-bearing: 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. There is no second way in to forget it on.ChipWorkertakes onlyChipStorageTaskArgs. The blob round trip it used to sit behind— C++ serialising bytes that Python handed straight back to C++ — is gone, along with the
TaskArgsViewoverloads that existed to unpack it.A copy names both of its ends
copy_to/copy_fromtake aBufferon each side and write the two descriptors intothe control frame; the child resolves them through the same
ImportRegistryits taskarguments 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, which was a second resolution rule for the one thing this ABI exists to resolve.
ControlCopyRequestis 184 B against 64 KiB of control room,static_assert-ed, and bothdescriptors are validated before the child acts on either.
Every receive path keeps its container
This is where the port was initially wrong, twice, and both are fixed here:
its orchestration function a
TaskArgs, scalars included. Re-export changes what isinside the container, not the container. An intermediate shape of
list[Tensor]silentlydropped scalars — the compute-leaf path had kept them, so the two receive paths disagreed.
TaskArgstoo, so anorchestration function can forward the args it was given whichever way it was reached.
It could not before:
_materialize_task_argsproduced the chip POD while_split_next_level_argsaccepts only the wire type, so a remote L3 could not submit toits own chip children at all.
Remote buffers are allocated through
create_buffer, so what a runner hands back carriesan identity and a descriptor instead of being a bare
SharedMemorybelonging to neithermechanism.
Naming
simpler_setup.torch_interop.make_tensor_arg→make_chip_tensor_arg. It returns the chipPOD where
Worker.make_tensor_argreturns the wireTensor; one name for both was thesecond public meaning rule 13 forbids.
wrap_fork_inheritedcallers name their backend explicitly now that #1703 stopped inferringit from
access. Each already knew the answer in a comment: the HeapRing backings areMAP_SHARED, and
make_tensor_argfollows thesharedit already computes — at L2 theconsumer is this process, so a write reaches the owner trivially and FORK_COW's contract is
the one that would be false there.
Verification
pyut— 1230 passed, 6 skippedcpput— 86/86. Note this was 85/86 with 30 cases never executing: a new testhelper tagged a descriptor
POSIX_SHMwithout a body, the gate above threw from aworker thread, and
std::terminatetook the binary down mid-suite. Bisected to thefirst commit of this work, not pre-existing as twice claimed.
a2a3sim/a5simfull scene suites — rc=0, 0 failures on bothexamples/workers/l4/vector_add_mixed_l3on a2a3 silicon, both halves:local max_diff=0.000e+00,remote max_diff=0.000e+00. This is the_st-pod.ymlpayload; a loopback daemon stood in for the second machine.
--platform a2a3/a5overexamples tests/st) — not runKnown, and deliberately not in this PR
IMPORT_BUFFERmaps the owner's whole shm soentry.addris the whole backing's base,while
entry.offsetis set as if it were the exported sub-range's base. A non-zeroworker.remote_export(handle, offset=N)therefore resolves N bytes short. Pre-existing— this PR reproduces the arithmetic deliberately — and untested either way: every test
exports at offset 0.
own registry. Carrying
BufferDescriptors on the sidecar wire is P2.REMOTE_SIDECARremains rejected at materialization, as Add: the Buffer/Tensor wire ABI and owner-side create_buffer #1599 froze it.Reviewing this
The single commit is large because the flip is not separable — the wire type, its decoders,
and every call site have to move together or nothing builds. The order that makes it
readable:
src/common/task_interface/task_args.h(what the wire carries and where it isvalidated), then
python/simpler/buffer.pyandworker.py(the receive paths), thenremote_l3_session.py, then the call sites.