Add: reject kind4 device pointers dispatched or freed on the wrong worker - #1430
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
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 Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThe change adds exact child-pointer provenance tracking, validates worker ownership for memory operations and task dispatch, records and revokes domain allocations, and adds comprehensive unit coverage for L2, L3, grouped submission, and stale-pointer behavior. ChangesChild pointer provenance
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant Orchestrator
participant Worker
participant CppOrchestrator
Caller->>Orchestrator: submit child-managed TaskArgs
Orchestrator->>Worker: resolve candidates and check provenance
Worker-->>Orchestrator: approve or reject target
Orchestrator->>CppOrchestrator: submit validated task
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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.
🧹 Nitpick comments (1)
python/simpler/orchestrator.py (1)
196-199: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional: avoid rebinding the
workerparameter. Line 197 reassigns theintparameterworkertoself._worker. It's correct today becausecpp_worker_id = int(worker)captures the int first, but the name now carries two meanings and the same pattern recurs insubmit_next_level_group; a later edit that readsworkerexpecting the pin id below this point would break silently. Consider a distinct local (e.g.worker_obj).♻️ Suggested rename
cpp_worker_id = int(worker) - worker = self._worker - captured_refs = worker._capture_remote_sidecar_refs(remote_sidecar) if worker is not None else [] - child_ptrs = worker._child_ptrs_in_args(c_args) if worker is not None else [] + worker_obj = self._worker + captured_refs = worker_obj._capture_remote_sidecar_refs(remote_sidecar) if worker_obj is not None else [] + child_ptrs = worker_obj._child_ptrs_in_args(c_args) if worker_obj is not None else []🤖 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/orchestrator.py` around lines 196 - 199, In the affected orchestration flow, avoid rebinding the worker parameter after computing cpp_worker_id. Store self._worker in a distinct local such as worker_obj and use that local for _capture_remote_sidecar_refs and _child_ptrs_in_args; apply the same naming correction in submit_next_level_group while preserving the existing None fallback behavior.
🤖 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.
Nitpick comments:
In `@python/simpler/orchestrator.py`:
- Around line 196-199: In the affected orchestration flow, avoid rebinding the
worker parameter after computing cpp_worker_id. Store self._worker in a distinct
local such as worker_obj and use that local for _capture_remote_sidecar_refs and
_child_ptrs_in_args; apply the same naming correction in submit_next_level_group
while preserving the existing None fallback behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 694eaeb1-a272-4a98-817f-e458089d5622
📒 Files selected for processing (3)
python/simpler/orchestrator.pypython/simpler/worker.pytests/ut/py/test_worker/test_child_addr_guard.py
e254c29 to
c5d61e2
Compare
|
Thanks — addressed all four findings (force-pushed). 1 [high] 2 [high] CommDomain release had no provenance commit barrier. Fixed. The revoke moved to the start of 3 [medium] transaction-failure tests. Added: native 4 [low] doc scope. PR body now says "every kind4 consumer (L2 Verification: 36 device-free UT (all 4 fix-specific tests verified failing against the pre-fix source), 494/2-skip broad sweep, ruff + pyright clean, no wire/C++/ABI change. |
c5d61e2 to
d8e7f6a
Compare
|
Round 2 addressed (force-pushed). 1 [high] free() not transactional → UAF on async unwind. Fixed with a safety-first commit barrier: 2 [medium] group swallowed illegal Should-fix (both done):
Verification: 39 device-free UT — I reverted both round-2 hunks and confirmed the 4 fix-specific tests fail with the right assertions, then restored. 497/2-skip broad sweep, ruff + pyright clean, no wire/C++/ABI change. |
d8e7f6a to
4936744
Compare
|
Round 3 addressed (force-pushed 1 [blocker] group could pin two members to the same worker. Fixed. After the per-member affinity write-back, a group with two members resolving to the same owner (a duplicate non-negative affinity — e.g. two child args on the same chip) is now rejected ("a group must dispatch to distinct workers") rather than emitting 2 [blocker] free contract vs docs mismatch. Fixed everywhere to the implemented safety-first semantics (validate → revoke → native free; native-free failure = terminal leak, no retry):
Non-blocking: Verification: 40 device-free UT — the group-duplicate fix verified failing (DID NOT RAISE) against reverted source; 498/2-skip broad sweep; ruff + pyright clean. CI: agreed on your read — last round's |
4936744 to
8e95959
Compare
|
Round 4 addressed (force-pushed 1 [high] empty entry could re-authorize a freed pointer. Fixed, fail-closed:
2 [medium] provenance analysis outside the remote-ref rollback window. Fixed: the (fallible) kind4 analysis now runs before Should-fix / track:
Verification: 45 device-free UT — both blocker fixes verified failing against reverted source (empty-entry DID-NOT-RAISE; capture called-1-time); 503/2-skip broad sweep; ruff + pyright clean; no wire/C++/ABI change. |
…rker A child (kind4) device pointer is a bare int today — no owner, no generation (tensor.h child_memory is one bit). Once such a pointer entered a TaskArgs it could be freed, copied, or run on a worker that never allocated it: a raw device VA is not globally unique, so worker B silently touched worker A's (or illegal) device memory. Guard ① (hw-native-sys#1397) only stopped host-addr rewrite from corrupting a device VA; nothing checked the pointer's owning worker. Track live child pointers by their exact (worker_id, ptr) provenance and validate every kind4 consumer (device ops and dispatch; L2 Worker.run is not a consumer): - Provenance table on the Worker: a typed (worker_id, ptr) -> entry map (malloc_owned + domain_allocation_ids), so a malloc base and an aliasing CommDomain buffer are distinct roles. Each op is atomic under one private lock; ordering is safety-first (record after alloc, revoke before free). Async interruption is fail-closed: a role is fully set before its entry is inserted, the last role is deleted directly (no role-less entry ever exists), and every live check tests the roles, not key presence, so an entry momentarily left empty never re-authorizes a freed pointer. Cleared on close(). - malloc: recorded after the backend malloc succeeds. free: a safety-first commit barrier — validate an exact live malloc base, then revoke provenance BEFORE the native free, so an async unwind (e.g. a KeyboardInterrupt after the binding returns) can never leave a freed address live; a native-free failure becomes a terminal leak, never a re-authorized maybe-freed address. copy_* require the device-side pointer live on that worker. Covers both Worker.malloc (L2 + L3) and the direct orch.malloc/copy/free paths (single Orchestrator choke), L2 and L3 unified. - submit_next_level / submit_next_level_group: a child_memory argument must resolve to exactly one eligible target worker (judged on the resolved eligibility, not a raw worker=-1) and be a live allocation there; 0 or >= 2 candidates is ambiguous and rejected. The resolved owner is passed to C++ as the effective affinity so the child TensorKey is keyed by its owner, not the raw -1 (otherwise the same buffer submitted once as -1 and once as W gets two keys and its dep is missed). The group path materialises a full per-member affinity only when workers is empty/None; a non-empty workers list is length- checked against args (never silently padded, which would bypass the C++ length check). A group must dispatch to distinct workers, so two members that resolve to the same owner (a duplicate non-negative affinity) are rejected rather than silently serialized on one WorkerThread. - CommDomain window / buffer pointers enter provenance on allocate_domain and are revoked in _release_domain_now BEFORE the backend free (a commit barrier): the deferred window stays dispatchable, but once physical release begins the pointers are undispatchable, so a concurrent dispatch cannot validate a being-freed pointer and a partial backend-release failure leaves them dropped rather than "live forever". Boundary: catches stale-before-reuse, NOT strict ABA (a re-malloc of the same VA becomes live again) — that needs P1 generation handles. Python-only, no wire / C++ / ABI change; the raw-C++ deep-bypass (worker._orch._o.malloc) is out of scope. The remote-slot-ref capture runs after (not before) the kind4 provenance analysis, so an analysis failure cannot strand captured refs outside the rollback try and defer a remote free forever. Adds 45 device-free UT (provenance table incl. empty-entry fail-closed, target resolution, orchestrator memory ops incl. free commit-barrier + native-error + lock-held-across-native, single/group dispatch + resolved affinity + group length/duplicate-worker validation, capture-after-analysis ordering, CommDomain dispatch and release-before-free commit barrier, L2 path).
1 participant
Scope
Device-address guard ② (
worker-memory-model.md§9). A child (kind4) devicepointer is a bare
inttoday — no owner, no generation (tensor.hchild_memoryis one bit). Once it enters aTaskArgsit can be freed,copied, or run on a worker that never allocated it; a raw device VA is not
globally unique, so worker B could silently touch worker A's (or illegal)
device memory. Guard ① (#1397) only stopped host-addr rewrite — nothing checked
the owning worker.
This tracks live child pointers by their exact
(worker_id, ptr)provenanceand validates every kind4 consumer — the device memory ops and dispatch
(L2
Worker.runis not a consumer). Python-only; no wire / C++ / ABI change.Delivered
Worker:(worker_id, ptr) -> {malloc_owned, domain_allocation_ids}. Each op is atomic under one private lock. Ordering issafety-first: malloc records after the native alloc;
freeand domainrelease revoke before the native free.
freevalidates an exact live mallocbase, revokes provenance, then does the native free (a commit barrier, L2
and L3 identical). A native-free failure is a terminal leak — provenance is
not restored and an explicit retry is rejected; an async unwind after a
successful free can never leave a freed address live.
copy_*require thedevice-side pointer live on that worker. Covers both
Worker.mallocand thedirect
orch.malloc/copy/freepath (singleOrchestratorchoke).child_memoryargument mustresolve to exactly one eligible target worker (resolved eligibility, not raw
-1) and be live there;0or>= 2candidates is rejected. The resolvedowner is passed to C++ as the effective affinity so the child
TensorKeyis keyed by its owner, not
-1. The group path materialises a full per-memberaffinity only when
workersis empty/None; a non-empty list is length-checked(never padded), and a group that resolves two members to the same worker is
rejected (a group must dispatch to distinct workers).
allocate_domainand are revoked in
_release_domain_nowbefore the backend free (commitbarrier): the deferred window stays dispatchable, but once release begins the
pointers are undispatchable and a partial backend-release failure leaves them
dropped rather than "live forever".
Boundary (explicit)
generation handles.
worker._orch._o.malloc) out of scope (needs C++).RemoteBufferHandle) already owner+generation guarded — exempt.Testing
test_child_addr_guard.py): provenance table,target resolution, orchestrator memory ops incl. free commit-barrier +
native-error terminal-leak, single/group dispatch + resolved affinity +
group length/duplicate-worker validation, CommDomain dispatch and
release-before-free commit barrier, L2 path. Fix-specific tests verified
failing against pre-change source.
test_worker/+ task/callable): 498 passed / 2 skipped— no regression. ruff + pyright clean.
worker=).st-sim-*,st-onboard-*).