Add: attach registered device Buffer sources at L2 submission

[ChaoWao]

Summary

Caller-owned device storage can enter the direct L2 Buffer/Tensor/TaskArgs path without a Worker-specific buffer constructor:

source = Buffer.wrap(address=ptr, nbytes=capacity, location=worker.device_location)
args = TaskArgs()
args.add_tensor(Tensor(source, shapes=(count,), dtype=DataType.FLOAT32))
worker.submit(callable_handle, args).result()
source.close()

Wrapping creates a source registration. Submission installs a private consumer snapshot and reserves the source before materialization. The Buffer identity and descriptor stay unchanged; existing geometry, grant, transfer, exact-worker and import checks remain active. Close and acquisition share the source lock, so a cached attachment cannot revive a released source. Rewrapping an idle address creates a new identity. Live source overlap is refused; overlap with an owned allocation is refused on attachment.

The current context provider is Worker.device_location on an initialized direct L2 Program Worker. The token identifies that process and Worker incarnation, not merely a device ordinal. Foreign Workers, forked processes and closed contexts are refused. This deliberately does not attach an externally owned RTS context or add a cross-Worker/Remote grant. Those consumers need their own context provider and authorization path; this PR does not mark U1b or kernel eager complete.

Lifetime and limits

• No physical allocation/free or Python allocation-owner retention. The caller guarantees the actual device, capacity, stable address and physical lifetime, including recovery.
• Worker.free rejects borrowed sources. Worker copy operations do not consume source registrations in this PR; the external allocator supplies IO.
• Materialization failure releases the source use. An exception after native entry or a reported run error retains it conservatively: the Python path has no failed-call completion proof. These registrations cannot be closed/reused through this API, and a reset is not silently treated as proof. The context token is invalidated after such errors.
• Unawaited successful runs release their uses after a successful lane close. Program Worker teardown can reset its device.

Follow-up to #2480, related to #2460. Replaces the proposed worker.borrow_device_buffer surface in draft #2475; it does not include the independent native fix already merged as #2478. No native ABI changes or planning documents are included.

Validation

• Initial source-path tests fail on main; the two context-invalidation cases also fail before their guard is added. Disabling source-use retention makes a close during materialization succeed and fails the regression test.
• 123 focused Buffer/L2 tests passed, including concurrent close during materialization, descriptor/context refusal, address reuse, registration rollback, native-error retention unawaited-run teardown, close between cache lookup and reservation, and a real fork with the source mutex held.
• HBG readiness: six cases passed on each of a2a3sim and a5sim. Simulation uses external ctypes storage absent from the native allocation table; onboard uses low-level ChipWorker allocations present in that table. Producers are awaited before consumers, so neither result demonstrates joined launch.
• A2/A3 onboard: six cases passed via task-submit (task_20260929_022352_171443715530, device 8). A5 onboard was not run locally. The later context-error guard is covered by unit tests; no new hardware success-path behavior was added after this run.
• Pre-commit hooks passed.

The cached-close test fails before the reservation validates already-attached source identities under its lock. The fork test times out before the process check is moved ahead of mutex acquisition.

• Final commit a938898b: complete Python unit suite 2833 passed, 65 skipped after rebuilding the editable package at that commit.

[coderabbitai]

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Advanced

Run ID: 570810a8-8d12-4424-aba6-c39f9151e4e8

📥 Commits

Reviewing files that changed from the base of the PR and between 35e9331 and a938898.

📒 Files selected for processing (6)

• docs/buffer-abi.md
• docs/user/reference/python-api.md
• python/simpler/buffer.py
• python/simpler/worker.py
• tests/st/host_build_graph_readiness/test_host_readiness.py
• tests/ut/py/test_worker/test_l2_argument_binding.py

Included review availability: This review used your included allowance. Your plan provides up to 1 included review per hour; 0 remain after this review.

---

📝 Walkthrough

Walkthrough

The change adds Buffer.wrap and Worker.device_location for registering externally owned device memory with an initialized direct-L2 Worker. L2 submission validates and reserves those buffers, with reservation handling tied to run completion and Worker teardown.

Changes

Borrowed Device Sources

Layer / File(s)	Summary
Register device source ranges python/simpler/buffer.py, docs/buffer-abi.md, docs/user/reference/python-api.md	Adds BufferLocation and Buffer.wrap. Registration validates the source range, access grant, context, and descriptor snapshot. Documentation describes supported contexts and caller responsibility for external allocation lifetime.
Attach sources to L2 runs python/simpler/worker.py, tests/ut/py/test_worker/test_l2_argument_binding.py, tests/st/host_build_graph_readiness/test_host_readiness.py, docs/buffer-abi.md	L2 submission attaches and reserves registered sources. Successful completion releases reservations; unproven completion retains them. Tests cover validation, attachment, failure handling, teardown, and a host-readiness source case.

Priority: ➖ Normal

Estimated code review effort: 4 (Complex) | ~60 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant PythonCaller
  participant Worker
  participant Buffer
  participant DeviceBufferSource
  participant L2RunLane
  PythonCaller->>Worker: Request device_location
  Worker-->>PythonCaller: Return BufferLocation
  PythonCaller->>Buffer: Wrap address with Buffer.wrap
  Buffer->>DeviceBufferSource: Register source range
  PythonCaller->>Worker: Submit Tensor containing Buffer
  Worker->>DeviceBufferSource: Validate and reserve source
  Worker->>L2RunLane: Submit native run
  L2RunLane-->>Worker: Report successful completion
  Worker->>DeviceBufferSource: Release source use at finalization

Loading

Merge Risk: ⚪ Minimal · up to a9388

This change adds an opt-in API for registering externally owned device memory for direct L2 submissions. The supplied context shows no concrete failure introduced by the change. Existing behavior is unchanged for callers that do not use the new API.

Security Architecture Review

Security architecture risk: 🟡 Moderate · up to a9388

Borrowed device memory is useful, but it moves an important safety guarantee to the caller: the library cannot establish that a supplied address belongs to the caller or remains allocated. Submission lifecycle controls reduce the risk of premature release through this API.

Retained concerns

• Medium · security · inferred: A caller with a direct L2 Worker can register a device address without proving ownership of the physical allocation. The resulting identity passes Worker provenance checks; if mutually untrusted allocations are addressable on that device, this could permit access outside the caller’s allocation. Such shared exposure is not established by the available evidence.

Security review details

Security Blast Radius

• inferred — The immediate new authority belongs to code able to obtain an initialized direct L2 Worker. Whether an arbitrary registered address can reach another security principal’s memory depends on device address isolation and allocator trust, neither of which is established here.

Security Findings and Attack Paths

• inferred — A caller could supply an address it does not own, receive a fresh valid registration, and submit that pointer. This is a conditional attack path, not a verified cross-tenant read or write: the caller must already have direct L2 access and the device must make the targeted address reachable.

Trust Boundaries and Controls

• observed — Process and live-context checks, descriptor equality, same-Worker dispatch checks, and use-counted close prevent several ways to reuse or redirect a registration. They do not independently attest the physical origin or lifetime of its address.

Resilience and Maintainability Implications

• observed — Unproven native completion keeps the registration reserved rather than allowing close or reuse. The external allocator must independently keep the physical memory valid during that period.

Hardening Proposals

• proposed — If direct L2 callers or allocators can be mutually untrusted, require allocator-issued ownership and lifetime proof for registration, or explicitly restrict borrowed-address authority to trusted allocators and isolate their device address spaces.

🚥 Pre-merge checks | ✅ 4

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: attaching registered device Buffer sources during L2 submission.
Description check	✅ Passed	The description directly explains the new caller-owned device storage path, submission behavior, lifetime rules, scope limits, and validation results.
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.

Warning

Some tools did not complete. Review the errors below.

🔧 Ruff (0.16.6)

python/simpler/worker.py

�[1;31mruff failed�[0m
�[1mCause:�[0m Required version ==0.14.8 does not match the running version 0.16.6

python/simpler/buffer.py

�[1;31mruff failed�[0m
�[1mCause:�[0m Required version ==0.14.8 does not match the running version 0.16.6

tests/ut/py/test_worker/test_l2_argument_binding.py

�[1;31mruff failed�[0m
�[1mCause:�[0m Required version ==0.14.8 does not match the running version 0.16.6

• 1 others

---

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

_{I’m a rabbit with a buffer to wrap,
A device-bound address tucked into my map.
The run takes its turn, then releases the hold,
While I nibble a carrot beside the run lane, bold.
The memory stays yours, as the docs clearly say,
And I hop through the tests at the end of the day.}

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