Fix: mint owner_instance_id after fork; check device backing owner

[ChaoWao]

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 is 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).

Known bound (flagged by review, documented rather than closed here): 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 needs
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 now states this bound explicitly.

Test plan

• pytest tests/ut — 1279 passed, 13 skipped, 0 failed
• ruff check / ruff format --check, pyright — clean
• Two new tests exercise materialize() directly (bypassing submit),
  mirroring Update: cut task args over to the self-describing Tensor wire ABI #1729's existing sub-worker test: a host endpoint rejects a
  DEVICE descriptor, and a device endpoint rejects one minted by a
  different chip's owner
• Real a2a3 onboard run (2 devices) across host_build_graph and
  tensormap_and_ringbuffer examples — exercises the touched chip-fork
  path, no failures

[coderabbitai]

📝 Walkthrough

Walkthrough

ImportContext adds endpoint and chip-instance ownership checks for device-backed buffer imports. Workers now propagate owner instance IDs into chip execution paths, and tests cover host rejection, foreign-chip rejection, and updated loop setup.

Changes

Device import ownership

Layer / File(s)	Summary
Import context and materialization validation python/simpler/buffer.py	ImportContext records host status and the owning chip instance ID. ImportRegistry.materialize() rejects device descriptors on host endpoints and descriptors owned by another chip.
Worker identity and endpoint wiring python/simpler/worker.py	Workers configure host and chip import contexts. Initialization regenerates the owner instance ID and passes it through chip-process and L2 execution paths.
Materialization and worker-loop coverage tests/ut/py/test_buffer.py, tests/ut/py/test_worker/test_endpoint_capability.py, tests/ut/py/test_worker/test_host_worker.py	Tests use explicit chip contexts and cover host rejection, foreign-chip rejection, and owner-ID propagation.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant HostWorker
  participant ChipProcess
  participant ImportRegistry
  participant DeviceDescriptor
  HostWorker->>HostWorker: regenerate owner_instance_id
  HostWorker->>ChipProcess: pass owner_instance_id
  ChipProcess->>ImportRegistry: configure device endpoint context
  ImportRegistry->>DeviceDescriptor: validate backing owner
  ImportRegistry-->>ChipProcess: materialize or reject descriptor

Loading

Possibly related PRs

• hw-native-sys/simpler#1599: Modifies ImportRegistry and buffer-handle ownership identity semantics.
• hw-native-sys/simpler#1703: Modifies device-backing validation in ImportRegistry.
• hw-native-sys/simpler#1704: Extends buffer import validation in python/simpler/buffer.py.

Poem

A rabbit checks each buffer’s claim,
Host and chip must name their game.
A fresh ID guards every run,
Foreign tensors pass through none.
Imports now follow ownership’s track—
No stray device backing sneaks back.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description check	✅ Passed	The description clearly explains the owner identity fix, endpoint ownership checks, affected code paths, tests, and validation results.
Title check	✅ Passed	The title clearly summarizes both main changes: reminting owner_instance_id after fork and validating device backing ownership.

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

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/buffer.py`:
- Around line 564-573: Replace DEVICE backing validation in buffer.py lines
564-573 with a serialized chip-specific ownership identifier, update the
ownership contract at buffer.py lines 546-548, and remove the chip-topology
claim at lines 584-585. In worker.py lines 2470-2475 configure each chip
endpoint with its own identity, and at line 7107 pass the identity for idx
rather than the shared Worker identity. Add a test in
tests/ut/py/test_worker/test_endpoint_capability.py lines 110-118 rejecting
materialization between sibling chip contexts that share a Worker identity but
have different chip identities.

🪄 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: 1b6b06af-c16d-432e-b1f2-cf08de20d014

📥 Commits

Reviewing files that changed from the base of the PR and between 6dc5bc7 and 128c7ff.

📒 Files selected for processing (5)

• python/simpler/buffer.py
• python/simpler/worker.py
• tests/ut/py/test_buffer.py
• tests/ut/py/test_worker/test_endpoint_capability.py
• tests/ut/py/test_worker/test_host_worker.py

[ChaoWao]

Review — #1744 (G4 + ImportContext, the last two steps of the capability-judgment chain)

This is G4 (mint owner_instance_id after fork) and ImportContext (materialize-time endpoint × address_space check) delivered together. Both were "新 PR · 未创建" on the board until now.

CI triage (2 failures, both confirmed infra — not this diff)

Pulled the actual logs, not guessed:

• ut-a5: ensure_pto_isa_root → OSError: PTO-ISA not available — the pto-isa clone over HTTPS timed out on a5-npu-1. Nothing in this diff touches pto_isa.pin or any pto-isa include path.
• st-onboard-a5: failed before checkout even ran — the runner itself couldn't download the actions/cache GitHub Action from codeload.github.com (3 retries, 100s timeout each). Pure runner-network issue.

Both on the same a5-npu-1 self-hosted runner — matches the flaky-infra pattern already seen elsewhere on this PR's run history. st-onboard-a2a3 was still pending when I checked; worth confirming it lands green before merge. Recommend rerunning the two a5 jobs.

Code verified against the PR body's claims

• Mint-point relocation is real, not just claimed: Worker.init() now re-mints _owner_instance_id right after NEW → INITIALIZING; __init__'s original mint is kept only as a fallback (correctly commented as existing for test doubles that set _lifecycle directly without calling init()). _run_chip_main_loop/_chip_process_loop gained an owner_instance_id parameter threaded from the parent at fork time — confirmed via the call site update in Worker (self._owner_instance_id passed in).
• ImportContext is wired into exactly the three construction sites that exist — _sub_worker_loop (host), _run_chip_main_loop (device, given the parent's nonce at fork), and the self._chip_import_registry lazy-init path (in-process chip worker, uses its own nonce). Checked with git grep -n "ImportRegistry(" across the whole tree: no fourth call site was missed, and none still constructs with zero args.
• Fail-closed default: materialize() treats context is None the same as is_host_endpoint=True — rejects any DEVICE backing rather than silently permitting one. Since every remaining call site now passes an explicit context, this is dead-but-safe, not a live gap.
• The CodeRabbit finding is the one we already discussed and it's closed the way we agreed: ImportContext's docstring states the Worker-grained-not-chip-grained bound explicitly, matching option 1 (document, don't block) — verified the thread itself: resolved=true, and the docstring text on disk matches what the PR body claims almost verbatim.
• The two new tests exercise materialize() directly, not just plumbing: test_host_endpoint_materialize_refuses_a_device_tensor_directly and test_chip_materialization_refuses_a_foreign_chips_device_tensor both bypass mapped_args_from_blob/submit-time dispatch entirely, mirroring the shape of Update: cut task args over to the self-describing Tensor wire ABI #1729's existing sub-worker test. This is the right level — CodeRabbit's finding was specifically about a path that reaches materialize() directly.

Minor note

Branch is one commit behind current main (merge-base 6dc5bc75, main is now at 6e2731dc post-#1588). No file overlap with #1588's diff (chip_run_lane/chip_worker/runtime_maker vs. buffer.py/worker.py here) so this is cosmetic, not a merge risk — gh already reports mergeable: MERGEABLE.

Verdict

No Must-fix, no Should-fix beyond what's already in the PR body. Once the two a5 infra flakes are rerun clean and st-onboard-a2a3 lands, this is ready. (Not filing this as an approval — GitHub blocks self-approval on your own PR, same as #1743.)