Refactor: isolate resources per worker run - #1466
Conversation
📝 WalkthroughWalkthroughThe PR replaces drain-based orchestration with run-scoped lifecycle tracking, Python ChangesRun-fenced execution
Pipeline execution
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
src/a2a3/platform/onboard/host/device_runner.cpp (1)
226-234: 🚀 Performance & Scalability | 🔵 TrivialPer-run stream create/destroy replaces the previously-reused persistent streams.
Each
run()invocation now issuesrtStreamCreate/rtStreamDestroyfor a fresh AICPU/AICore pair instead of reusingDeviceRunnerBase's persistent streams across runs. This is clearly deliberate (K=2 scaffolding per thedevice_runner.hcomments), but it does add two extra CANN stream lifecycle round-trips to everyrun()call on the hot path. Worth keeping an eye on measured per-run latency ifrun()is called at high frequency; if it turns out material, a stream-pool-per-slot (create once, reuse until slot count changes) would avoid the per-call churn while keeping the per-slot ownership model.Also applies to: 454-481
🤖 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 `@src/a2a3/platform/onboard/host/device_runner.cpp` around lines 226 - 234, Measure the per-run overhead introduced by create_run_stream_set and the run_stream_cleanup guard in run(). If stream lifecycle churn is material, replace per-invocation creation/destruction with a stream pool owned per pipeline slot, reusing streams until the slot count changes while preserving per-slot ownership and cleanup.
🤖 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/worker.py`:
- Around line 6038-6044: Update Worker.close() so the accepted run-handle
draining loop bounds each handle.wait() call using the remaining
_ROLLBACK_GRACEFUL_TIMEOUT_S budget already used for _active_ops. When that
budget expires before all handles drain, stop waiting and take the retryable
teardown path by preserving drain_complete = False and setting
_teardown_attempted = False.
---
Nitpick comments:
In `@src/a2a3/platform/onboard/host/device_runner.cpp`:
- Around line 226-234: Measure the per-run overhead introduced by
create_run_stream_set and the run_stream_cleanup guard in run(). If stream
lifecycle churn is material, replace per-invocation creation/destruction with a
stream pool owned per pipeline slot, reusing streams until the slot count
changes while preserving per-slot ownership and cleanup.
🪄 Autofix (Beta)
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: 2fe8b959-5750-4048-a5cd-0fad445a9bd3
📒 Files selected for processing (39)
docs/comm-domain.mddocs/hierarchical_level_runtime.mddocs/orchestrator.mddocs/remote-l3-worker-design.mddocs/scheduler.mddocs/task-flow.mdpython/bindings/worker_bind.hpython/simpler/orchestrator.pypython/simpler/task_interface.pypython/simpler/worker.pysrc/a2a3/platform/onboard/host/device_runner.cppsrc/a2a3/platform/onboard/host/device_runner.hsrc/a2a3/runtime/host_build_graph/host/runtime_maker.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cppsrc/common/hierarchical/orchestrator.cppsrc/common/hierarchical/orchestrator.hsrc/common/hierarchical/ring.cppsrc/common/hierarchical/ring.hsrc/common/hierarchical/scheduler.cppsrc/common/hierarchical/scheduler.hsrc/common/hierarchical/tensormap.cppsrc/common/hierarchical/tensormap.hsrc/common/hierarchical/types.cppsrc/common/hierarchical/types.hsrc/common/hierarchical/worker.cppsrc/common/hierarchical/worker_manager.cppsrc/common/hierarchical/worker_manager.hsrc/common/platform/onboard/host/device_runner_base.cppsrc/common/platform/onboard/host/device_runner_base.hsrc/common/worker/chip_worker.cppsrc/common/worker/chip_worker.hsrc/common/worker/pto_runtime_c_api.htests/ut/cpp/hierarchical/test_orchestrator.cpptests/ut/cpp/hierarchical/test_scheduler.cpptests/ut/cpp/hierarchical/test_tensormap.cpptests/ut/py/test_worker/test_error_propagation.pytests/ut/py/test_worker/test_host_worker.pytests/ut/py/test_worker/test_l3_l2_orch_comm.pytests/ut/py/test_worker/test_startup_readiness.py
7719922 to
c22f0ab
Compare
c22f0ab to
3d23abc
Compare
Revision: rebased onto merged
|
|
CI status on
Earned rather than assumed:
It is deterministic on the GitHub runners and rare on our aarch64 box, which is why a re-run did not clear it. #1504 replaces the log scrape with a dedicated pid file. Green on this branch: |
3d23abc to
c5c50ca
Compare
c5c50ca to
1ac34f9
Compare
Revision: a domain released after its run retired now frees inlinePushed The fix
Both recovery paths are closed to it at once. Before this PR the handle landed on the Worker-level queue and the next run's fence picked it up — but that cross-run pickup is exactly the coupling this PR exists to remove, so restoring it is not an option.
if resources.retired:
# Deferral exists so tasks that captured this domain still see live
# memory; the owning run's fence has already passed, so there is
# nothing left to defer behind — and its queue is never read again.After the fix the same probe reports
Reachability, stated plainlyThis needs user code holding a Validation
One flaky test, tracked separatelyTwo of my full local suite runs failed on |
1ac34f9 to
f1476c2
Compare
Revision: the retirement window is closed with a lock, not narrowedPushed Proof it was still openI forced the exact interleaving with a The One correction to the report's step order, for anyone reproducing it: the per-run Fix
with resources.domain_lock:
resources.retired = True
stragglers = list(resources.pending_release_domains)
resources.pending_release_domains.clear()
for handle in stragglers:
self._free_domain_after_fence(handle)A release either lands on the queue before retirement or observes it and frees itself. No third outcome. Also fixed while here: the previous revision's inline free was unguarded. The end-of-run sweep works from a snapshot, so a release taken after that snapshot could issue a second The concurrency test the report asked for
Mutation-checked both ways: removing Validation
One earlier onboard run on this commit hit |
f1476c2 to
afaefe2
Compare
Revision: one allocation, one backend releasePushed The stated order is already safe
(The report's first two steps also cannot interleave: The order that does reach itRelease first, sweep second. The release wins The test parks the sweep on a different handle, so it is stopped after its snapshot and before it reaches the contested one — which is the shape the report asked for, reached from the other side. Against a build without the claim: Fix
with self._domain_free_mu:
if handle.allocation_id in self._domains_freed:
return
self._domains_freed.add(handle.allocation_id)
# Claimed, not held: the backend release runs outside the lock so
# concurrent releases of *different* allocations do not serialize.
self._release_domain_claimed(handle)Holding the lock across the backend call also works, but it serializes every domain free and made the test pass only by hitting a 5 s timeout — the releasing thread was blocked on the lock rather than exercising the claim. Claim-then-free is 0.28 s and deterministic. The claim is permanent even if the release raises, which matches the sweeps' existing no-retry contract. A path that finds the allocation claimed returns at once, so it may mark the handle Validation
|
Bind the Python resources a run creates — remote slot references, dynamic CommDomains, L3-L2 regions, and orchestration host-buffer registrations — to the RunHandle that owns them, so completing one run releases only that run's resources. Previously the cleanup at a run's fence swept Worker-level registries and could free a resource another run still owned. _RunResources holds the four groups; Worker._building_run_resources points at the set being filled while a run's graph is built, and the creation helpers charge new objects to it. _submit_mu makes that pointer single-writer, and submit drains every prior handle before setting it, so no fence-owned cleanup runs against a set still being filled. A nested submit on the same Worker would block on _submit_mu before reaching the pointer. A domain's lifetime ends at the fence of the run that allocated it, so allocate_domain outside graph construction has no owner to charge and is rejected. The deferred-release queue moves onto the run with it: the Worker-level list had no drain once the fence-owned path replaced it, so a handle queued there would never have been freed. Deferral exists so tasks that captured a domain still see live memory through execution. Once the owning run's fence has claimed its domains there is nothing left to defer behind, so a release arriving after that frees inline. domain_lock makes the two sides mutually exclusive: without it a release can read the retirement flag as False, be preempted, and append to a queue the fence has already drained for the last time — reachable from neither the fence nor close(), which that same release blinded by dropping the handle from _live_domains. Retirement claims any straggler under the lock. One allocation gets one backend release, and every caller waits for its outcome. Both a sweep and a post-fence release can reach the same handle: a sweep works from a snapshot, and a release that wins the handle's released flag first leaves the sweep to skip that flag and go straight to a CTRL_RELEASE_DOMAIN the release is already making. _release_domain_now performs the release under its own lock and records the result, so a second caller neither issues a duplicate release nor returns before the first finished — returning early would let it mark the handle freed, drop it from _live_domains, and on the close() path tear down the mailboxes the release is still using, which is why close() orders its domain sweep before _worker.close(). A failure is replayed rather than retried, so the sweeps keep the handle as a detectable residual. Worker-level registries stay for uniqueness checks and for the close() sweep, which still reaches resources created outside a run. K=1 submission is unchanged: a later submit() still drains the prior run, so this adds no acceptance flight or multi-run admission.
afaefe2 to
18cc5a0
Compare
Revision: the second caller waits for the outcome, it does not just skipPushed What claim-and-return broke
Fix
with self._domain_free_mu:
if handle.allocation_id in self._domain_free_results:
failure = self._domain_free_results[handle.allocation_id]
if failure is not None:
raise failure
return
try:
self._release_domain_claimed(handle)
except BaseException as exc:
self._domain_free_results[handle.allocation_id] = exc
raise
self._domain_free_results[handle.allocation_id] = NoneConcurrent releases of different allocations now serialize. That is the right trade here — domain frees are rare and already effectively serial on the fence thread, and The three properties, each with its own test and its own mutation
Two notes on how these are driven, since my first attempt at each was wrong:
Validation
|
Depends on #1464 (merged).
Summary
CommDomains, L3-L2 regions, orchestration host-buffer registrations — to theRunHandlethat owns themallocate_domainoutside graph construction, and move the deferred-release queue onto the owning run with itclose()sweepsubmit()still drains the prior run, so this adds no acceptance flight or multi-run admissionOwnership
_RunResourcesholds the four groups.Worker._building_run_resourcespoints at the set being filled while a run's graph is built, and the creation helpers charge new objects to it.Two things make that single pointer safe, and both are load-bearing:
_submit_mumakes it single-writer across concurrentsubmit()calls, and is non-reentrant, so a nested submit on the same Worker blocks before it could reach the pointer.submit()drains every prior handle before setting it, so no fence-owned cleanup can run against a set still being filled — a concurrenthandle.wait()on a prior run has already completed by then.Domains end at their run's fence
A
CommDomainHandle's release routes to the_RunResourcesof the run that allocated it, captured at allocation. That leaves no owner outside graph construction, soallocate_domainnow rejects that case rather than allocating something nothing will free.The Worker-level
_pending_release_domainslist goes away with it. Once the fence-owned path became the only drain,_execute_pending_domain_releaseswas always called with a run's resources andclose()never touched the Worker list — anything queued there would have leaked silently. Deleting the branch is what makes that unrepresentable.A release that arrives after its run retired
Deferral exists so tasks that captured a domain still see live memory through
execution. Once the owning run's fence has claimed its domains there is nothing
left to defer behind, so a release arriving after that frees inline.
domain_lockmakes the two sides mutually exclusive. Without it a release canread the retirement flag as False, be preempted, and append to a queue the
fence has already drained for the last time — reachable from neither the fence
nor
close(), which that same release blinded by dropping the handle from_live_domains. Retirement claims any straggler under the same lock, so arelease either lands on the queue before retirement or observes it and frees
itself; there is no third outcome.
The post-fence free tolerates a handle the end-of-run sweep already took: that
sweep works from a snapshot, so a release taken after it reaches the free path
for a domain that is already gone.
One allocation, one backend release — and every caller waits for it
Two paths can reach the same handle. A sweep works from a snapshot, and a
release that wins the handle's
_releasedflag first leaves the sweep to skipsetting that flag and go straight to a
CTRL_RELEASE_DOMAINthe release isalready making — a duplicate release of live backend memory, not a no-op.
_release_domain_nowperforms the release under its own lock and records theresult. A second caller for the same allocation blocks until the first
finishes, then sees its outcome. Returning early — claiming and freeing outside
the lock — is not enough: the second caller would mark the handle
freed, dropit from
_live_domains, and on theclose()path continue into_worker.close()while the release is still driving the mailboxes. That isprecisely what
close()orders its domain sweep before the scheduler teardownto avoid, and it would also break the public
CommDomainHandle.freedcontract("backend release has executed").
A failure is replayed to later callers rather than retried, so the sweeps keep
the handle in
_live_domainsas a detectable residual instead of reporting arelease that never happened.
The reverse order was already safe: a sweep sets
_releasedbefore freeing,which makes a later
release()a no-op.Scope
TASK_ACCEPTED, asynchronous flight, Gate A/B, K=2 and arena banking remain later changes. Nothing here changes when a run is admitted.Validation
Rebased onto merged
main; the five superseded ancestor commits are dropped, so this is a single commit of 2 files, +568/-51.a2a3onboard, the exact CI command (examples tests/stminus the two SDMA demos, 4 devices undertask-submit): 117 passed, 1 skipped, 0 failed — unchanged before and after the domain-ownership change, so the newallocate_domainguard interrupts no real L3 pathtest_run_finalization_releases_only_its_resources; replacing theallocate_domainguard with a silent default_RunResources()failstest_allocate_domain_outside_graph_construction_is_rejected; both halves of the retirement path — dropping the branch, and never setting the flag — failtest_domain_released_after_its_run_retired_is_freed_inline; and removing eitherdomain_lockor the straggler drain failstest_domain_released_while_its_run_retires_is_still_freed, which parks the fence in its lock-free live-domain sweep while a second thread releases; and the three domain-release properties are each caught by their own mutation — no claim at all givesallocation 7 reached the backend 2 times, claiming but freeing outside the lock givesthe second caller returned while the owner was still releasing, and recording success for a failed release givesa failed release must not be reported as freedThe two red a5 jobs on the previous head were the a5 runner, not this change:
st-onboard-a5predates #1481 (which skips that job while the runner is under repair) andut-a5failed atHcclCommInitRootInfo failed: 1intest_two_rank_comm_lifecycle/test_two_rank_allocate_release_round_trip. This PR touches no C++ or binding file, andut-a5is green on currentmain.