Optimize: serialize provenance-guarded device ops per worker, not process-wide - #1702
Optimize: serialize provenance-guarded device ops per worker, not process-wide#1702lterrac wants to merge 1 commit 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:
🚥 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.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/simpler/orchestrator.py (1)
637-647: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick winRelease
_child_prov_lockbefore the native free.Lines [637-647] keep the process-wide
_child_prov_lockheld throughself._o.free(wid, p). A slow free on one worker therefore blocks frees and provenance operations for every other worker. Keep_child_prov_worker_lock(wid)around the native call, but scope_child_prov_lockto validation and revocation only.Proposed fix
- with self._worker._child_prov_worker_lock(wid), self._worker._child_prov_lock: + with self._worker._child_prov_worker_lock(wid): # Safety-first commit barrier: revoke provenance BEFORE the native # free. If the native free succeeds and an async unwind (e.g. a # KeyboardInterrupt delivered after the binding returns) fires before # a post-free clear could run, a freed address would stay live and a # later copy/dispatch would re-authorize it — a UAF. Revoking first # turns a native-free failure into a terminal leak (recoverable) but # never re-authorizes a maybe-freed address. - self._worker._child_prov_require_malloc_base(wid, p, api="free") - self._worker._child_prov_clear_malloc(wid, p) + with self._worker._child_prov_lock: + self._worker._child_prov_require_malloc_base(wid, p, api="free") + self._worker._child_prov_clear_malloc(wid, p) self._o.free(wid, p)🤖 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 637 - 647, Restructure the cleanup block around _child_prov_worker_lock so _child_prov_lock is held only while _child_prov_require_malloc_base and _child_prov_clear_malloc execute. Release _child_prov_lock before calling self._o.free(wid, p), while retaining _child_prov_worker_lock(wid) across the native free.
🤖 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.
Outside diff comments:
In `@python/simpler/orchestrator.py`:
- Around line 637-647: Restructure the cleanup block around
_child_prov_worker_lock so _child_prov_lock is held only while
_child_prov_require_malloc_base and _child_prov_clear_malloc execute. Release
_child_prov_lock before calling self._o.free(wid, p), while retaining
_child_prov_worker_lock(wid) across the native free.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 3c0cf1fb-c8f9-4f33-9847-a245f982e997
📒 Files selected for processing (2)
python/simpler/orchestrator.pypython/simpler/worker.py
|
@coderabbitai good catch — that was a real inconsistency in my own change, fixed in 19c89e5. The free path still combined both locks in a single Validated on 4 × 910B2 with the isolated upload harness (which frees every shard after each band, so the free path runs 16 times per run): no deadlock, threaded 17.6 / 30.2 GB/s vs serial 10.5 / 9.9 GB/s. |
|
Tip For best results, initiate chat on the files or code changes.
🐇✅ ✏️ Learnings added
If you found this review helpful, would you consider giving us a shout-out on X? Thank you for using CodeRabbit! You are interacting with an AI system. |
|
The lock split itself looks right, but I think the 57.1s → 10.1s number is no longer attributable to this diff on The measured baseline predates a second serializer on the same pathValidation was done on simpler So on the branch this PR actually sits on, Call pathFor an L3+
What that implies for the 8-thread uploadA thread pool spawned by the caller has no
Which matches the symptom described in the PR body — threads entering Suggested
Two smaller notes:
|
19c89e5 to
6dc98e4
Compare
|
@YunjiQin thank you — you were right, and the measurement backs every step of your analysis. I re-ran on this PR's base as you suggested and rebased the branch onto current Your hypothesis, measured
Serial is identical across the two bases, so this is not a setup difference — A draft for the second lock, kept separateRather than reply with a question I prototyped your suggestion and measured it, so the decision has a number attached. It is a separate commit (
What I am not sure about is whether "no run may be admitted while I run" is the whole invariant Two smaller points from your review:
The dependent pypto change (hw-native-sys/pypto#2292, concurrent shard upload) is a no-op without this, so there is no ordering constraint between them. |
alloc_stacked_tensor uploads shard i to worker i in a serial loop, so a rank-stacked resident weight moves at single-chip H2D bandwidth no matter how many chips the group spans. Each shard targets a different chip worker and nothing orders them, so drive them from a thread pool. Rolling back needs a little more care than the serial loop: a concurrent failure can land anywhere in the group, so the successes are no longer a prefix of ids. Collect them by index and free them against their own worker before re-raising, instead of zipping shards with ids positionally. Measured on 8 x 910B2 uploading DeepSeek V4 Flash W8A8's 346 GB of rank-stacked weights (per-shard 11.5 GB): upload 57.1s -> 10.4s (6.0 -> 33.2 GB/s) startup 82.0s -> 35.1s This needs the matching Simpler change (hw-native-sys/simpler#1702) to pay off: Simpler holds one process-wide lock across the native half of malloc / copy_to, which serializes the group regardless of how the caller issues it. Without that change this commit is a no-op, not a regression.
|
CI feedback addressed. Two failures, only one of them mine.
It now asserts the narrower exclusion the code actually provides:
The argument that the narrower form is sufficient: provenance is keyed by Verified on
Branch rebased onto current |
961e9c5 to
3c20734
Compare
Two locks are held across the *native* half of malloc / free / copy_to / copy_from, and either one alone is enough to serialize every device op of every next-level worker onto a single chip's pace. The bindings already release the GIL there, which makes the effect easy to miss: unrelated Python threads keep running, but a `copy_to` on chip 0 blocks a `malloc` on chip 1 for its whole duration. `_child_prov_lock` (provenance) stays the bookkeeping lock — it still makes each provenance mutation/read atomic, and the safety-first ordering is unchanged (record after a successful alloc, revoke before a native free) — and a per-worker lock is taken around the native call instead. Ops on the same worker stay mutually exclusive, so a copy can still never overlap that buffer's free; ops on different workers now overlap. The per-worker lock is always acquired before `_child_prov_lock` and never the reverse, so the pair cannot deadlock. `_submit_mu`, taken through `_control_reservation` (hw-native-sys#1541), is the other one: a control command that belongs to no run holds it across the native call, so with the provenance fix alone it becomes the serializer — all 8 threads enter `Orchestrator.copy_to` within 14 us of each other, then each child starts its copy within ~1 ms of the previous one finishing, at full link speed. What such a command needs is "no run may be admitted while I run", which is a property of the worker, and two commands on different chips can both have that at the same time. So `_submit_mu` becomes a shared/exclusive lock: run admission takes it exclusively, control takes it shared. Writer-preferring, so control traffic cannot starve a submit. The reservation's re-entrancy is untouched: it short-circuits on the thread-local set before reaching the lock. Measured on 8 x 910B2, one upload thread per shard, fresh shared-memory bands so every page is read cold exactly once: provenance locks only threaded 6.5-8.0 GB/s serial 8.6-10.3 + run-admission lock threaded 19.2-34.0 GB/s serial 9.4-10.5 This narrows an invariant an existing test pins down, so that test is updated rather than left passing by accident: `test_free_holds_lock_across_native_free` asserted that the *parent worker's* lock is held across the native free. It now asserts the narrower exclusion actually needed — that worker's own lock held across the native call, `_child_prov_lock` released, and the revoke committed first. Provenance is keyed by (worker_id, ptr) and revoked before the native free, so a concurrent dispatch reads the table under `_child_prov_lock` and finds the address already gone, or is about a different chip entirely. The shared/exclusive shape for hw-native-sys#1541's serializer is the part most open to discussion; it can be replaced with a different design without touching the provenance change.
3c20734 to
6259971
Compare
2 participants
Rebased onto current
main. Two commits: a fix, and a draft for a second lock that can be dropped independently.The problem
_child_prov_lockis held across the native half ofmalloc/free/copy_to/copy_from, so every device op of every next-level worker serializes on one lock belonging to the parent worker. The bindings already release the GIL there, which makes it easy to miss: unrelated Python threads keep running, but acopy_toon chip 0 blocks amallocon chip 1 for its whole duration.Measured with per-shard timestamps on both sides of an 8-chip upload: all 8 orchestrator threads enter
Orchestrator.copy_towithin 14 µs of each other, then each chip child starts its copy within ~1 ms of the previous one finishing, each at full link speed. InsideLocalMailboxEndpoint::control_copy_to,wait_lockis0.0000s— the per-worker mailbox mutexes are never contended, the threads simply arrive one at a time.Commit 1 — per-worker provenance locks
_child_prov_lockstays the bookkeeping lock (each provenance mutation/read still atomic, safety-first ordering unchanged: record after a successful alloc, revoke before a native free); a per-worker lock now wraps the native call. Same-worker ops stay mutually exclusive, different workers overlap. The per-worker lock is always taken before_child_prov_lock, never the reverse, so the two cannot deadlock.This modifies an existing test.
test_free_holds_lock_across_native_freepinned the wide behaviour; it now asserts the narrower exclusion the code provides — that worker's lock held across the native free,_child_prov_lockreleased, revoke committed first. Sufficient because provenance is keyed by(worker_id, ptr)and the revoke commits before the native free, so a concurrent dispatch reads the table under_child_prov_lockand finds the address already gone, or is about a different chip. Flagged explicitly: it pins a deliberate decision, so if the reasoning does not hold, the answer is to revert the free path rather than to keep the test green another way.Commit 2 — draft: shared/exclusive run-admission lock
This alone is not enough on this base, as @YunjiQin identified:
_control_reservation(#1541) takes_submit_muand holds it across the same native call, at the same per-worker granularity. My original numbers were measured on9922afdb, which predates #1541 — corrected below.A control command that belongs to no run needs "no run may be admitted while I run", a property of the worker; two commands on different chips can both have it at once. So
_submit_mubecomes shared/exclusive: run admission takes it exclusively, control takes it shared. Writer-preferring, so control traffic cannot starve a submit; the reservation's thread-local re-entrancy short-circuits before the lock and is untouched.Kept as a separate commit precisely so it can be dropped or replaced — it changes the serializer #1541 introduced, and the shape is the author's call.
Measured
8 × 910B2, one upload thread per shard, fresh shared-memory bands so every page is read cold exactly once (reusing a band measures the warm path and inflates everything):
Serial is identical across bases, so the difference is the lock and not the setup. For reference on the same node, 8 independent processes doing raw cold H2D from a shared mapping reach ~29–34 GB/s aggregate.
Tests
tests/ut/py/test_shared_exclusive_lock.py(new): shared holders overlap, each mode excludes the other, a waiting writer blocks new readers.tests/ut/py/test_worker/test_child_addr_guard.py: updated as described; the file passes in full (49 tests with the new one included).Open question
Whether "no run may be admitted while I run" is the whole invariant
_control_reservationcarries, or whether something also relies on control commands excluding each other. That decides commit 2 — happy to implement a different shape or hand it over.st-pod-onboard-a2a3fails onvector_add_mixed_l3/poll_native_run failed; PR #1722 fails identically on the same job while #1709 / #1714 / #1718 / #1723 pass, so it reads as a pod-runner flake rather than this branch.The dependent pypto change (hw-native-sys/pypto#2292) is a no-op without this, so there is no merge-order constraint between them.