Optimize: serialize provenance-guarded device ops per worker, not process-wide - #1702
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 |
3c20734 to
6259971
Compare
|
This branch is 19 commits behind
tl;dr: rebase onto |
6259971 to
d220247
Compare
|
@ChaoZheng109 thank you — both traps you named were real, and one of them was real twice. Force-pushed. Please read this as a rewrite, not a rebase. The base moved, the edited functions moved, there is a new C++ change, and every number is re-measured with a different harness. If you reviewed the previous version, none of your reading carries over. Your two traps1. The functions moved. Re-applied in 2. The bare The same trap also existed in three tests, which your review could not have seen and I only found by running the suite: A third serializer, and it is why the old numbers did not carry overRe-measuring on this base gave no gain at all from the two lock splits, which sent me looking. So there are three serializers in series, and removing any one of them measures as noise. That is the whole reason this looked like a dead end twice: the two lock splits alone did nothing, and an earlier attempt at the GIL guards alone did nothing either. This adds C++ to the PR ( Measured on this codePure simpler, no pypto, no kernels: an L3 Worker over 8 chips, one born-shared 2.15 GB host buffer per chip, uploaded through
4.14x over sequential within the same build. Two caveats I would rather state than have you find: the serial column is the control and stays in 20.7-22.8 GB/s across all three builds, and absolute throughput drifts ~10% between sessions — so the claim rests on the in-build threaded/serial ratio and on the overlap factor, not on any single absolute figure. TestsFull One piece of context for scope: pypto currently pins simpler |
d220247 to
d33bd45
Compare
|
Rebased onto Of the five commits that landed in the meantime, exactly one touches this PR's files: #1757, "give L2 the same in-flight protection release_buffer() gives L3+" — i.e. the very method this PR changes. The textual conflicts were two twins in test fixtures, where The part that needed judgement rather than merging is Re-verified on the rebased base, since a test run on the old base would not carry over a rewrite of this method: full |
Three things serialize the device-memory ops, in series, so removing any one of them alone measures as noise — which is why this took a while to pin down. A `copy_to` on chip 0 blocks a `malloc` on chip 1 for its whole duration, and the eight chips of an 8-way upload run strictly back to back. `_child_prov_lock` 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`, 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. 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. The third is the **GIL**. `Worker.malloc / free / copy_to / copy_from` are bound as plain lambdas in `worker_bind.h` with no call guard, so the interpreter lock is held for the whole native call while 31 other methods in that same file already release it. With the two Python locks split but the GIL still held, eight threads still cannot overlap: measured `overlap_factor` (sum of per-copy wall times over the wall time of the batch) stayed at 1.1-1.6 out of a possible 8. These four get `nb::call_guard<nb::gil_scoped_release>()`; none of them re-enters Python, and the descriptors are converted before the call. Re-applied on the post-hw-native-sys#1650/hw-native-sys#1729 structure rather than rebased textually. The provenance and native-call logic moved out of `Orchestrator` into `Worker`, so the split now lives in `Worker.alloc_child_tensor / free / copy_to / copy_from` and `orchestrator.py` is left exactly as main has it. `release_buffer` (hw-native-sys#1751) took `_submit_mu` bare, which a textual merge would have compiled and then crashed on, since `_SharedExclusiveLock` has no `__enter__`; it now takes it exclusively, keeping the ordering it had as a plain lock. Both traps were called out by @ChaoZheng109 in review. Three tests stand a plain `threading.Lock` in for `_submit_mu` or take it as a context manager, so they are updated to the real type and to `.exclusive()` (`test_create_buffer.py`, `test_remote_l3_lifecycle.py`, `test_release_buffer.py::test_serializes_with_a_racing_orchestration_callback`). They pin the serializer's *identity*, not its granularity, and the exclusive form is what graph construction now takes, so the property each one asserts is unchanged. This also 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. Measured on this code, on 8 x 910B2, with a pure-simpler harness (no pypto, no kernels): an L3 Worker over 8 chips, one born-shared 2.15 GB host buffer per chip, uploaded through `Worker.copy_to` outside a run — the same control path the resident-weight upload uses. The same work is done once with a thread per chip and once strictly sequentially; `overlap_factor` is the sum of the per-copy wall times over the wall time of the batch, so 1 means "back to back" and 8 means "fully overlapped". | | threaded | serial | overlap_factor | |---|---:|---:|---:| | main, unpatched | 19.99-22.11 GB/s | 20.69-22.82 GB/s | 1.4-1.8 | | + the two lock splits | 20.86 GB/s | 21.80 GB/s | 1.1-1.6 | | + the GIL guards (this PR) | **86.85-90.03 GB/s** | 21.47-21.81 GB/s | **6.4-6.5** | **4.14x** over sequential on the same build, where before there was none: the eight copies now all start together instead of queueing. Two things keep that honest. The serial column is the control and stays in 20.7-22.8 GB/s across all three builds, so the gain is concurrency and not a faster machine. And absolute throughput drifts about 10% between sessions — which is why the claim rests on the threaded/serial ratio measured *within* a build, and on the overlap factor, rather than on any single absolute number. The table also shows why this took three attempts to see. The two lock splits move neither throughput nor overlap; an earlier attempt at the GIL guards alone measured as noise too. With three serializers in series, removing any one of them changes nothing measurable, and only the last one removed appears to "cause" the win.
d33bd45 to
ab7ed7b
Compare
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. What this is worth, and when ---------------------------- On the Simpler commit this repo currently pins (3165cc89) this change is a **no-op, not a regression**: Simpler serializes the native half of malloc / copy_to three times over — a process-wide provenance lock, the run-admission lock taken by control commands, and the GIL, which its bindings did not release for those four entry points. Any one of the three is enough to flatten the group, so issuing the shards concurrently changes nothing there. Measured at 0%. hw-native-sys/simpler#1702 removes all three. In its own 8-chip harness, driving the same control path this function drives, the per-copy overlap factor goes from 1.4-1.8 to 6.4-6.5 and threaded throughput from ~21 GB/s to 87-90 GB/s, 4.14x over sequential on the same build. That is the gain this commit exists to expose. It is not reachable from here yet. This repo pins Simpler 82 commits behind that merge, and the pin cannot simply be advanced: Simpler moved device memory onto Buffer handles (hw-native-sys#1650 / hw-native-sys#1729), while this file still calls `self._w.malloc(nbytes, worker_id)` and `self._w.copy_to(dst_ptr, src_ptr, nbytes, worker_id)` with raw pointers. Advancing the pin is a port, and until it lands there is no stack on which the end-to-end effect of this commit can be measured. An earlier revision of this message quoted 57.1s -> 10.4s for the upload and 82.0s -> 35.1s for startup on DeepSeek V4 Flash W8A8. **Those figures are withdrawn.** They were taken on a stack that combined the pre-refactor Simpler with an out-of-tree lock patch and a serving-side loader change, and none of the three is what this commit would ship against. They will be re-measured end to end once the pin advances.
3 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.