Skip to content

Refactor: isolate resources per worker run - #1466

Merged
ChaoWao merged 1 commit into
hw-native-sys:mainfrom
Crane-Liu:codex/worker-async-pr6-run-resources
Jul 27, 2026
Merged

Refactor: isolate resources per worker run#1466
ChaoWao merged 1 commit into
hw-native-sys:mainfrom
Crane-Liu:codex/worker-async-pr6-run-resources

Conversation

@Crane-Liu

@Crane-Liu Crane-Liu commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

Depends on #1464 (merged).

Summary

  • bind the Python resources a run creates — remote slot references, dynamic CommDomains, L3-L2 regions, orchestration host-buffer registrations — to the RunHandle that owns them
  • release only the completed run's resources at its fence; a run finishing no longer frees what another run still owns
  • reject allocate_domain outside graph construction, and move the deferred-release queue onto the owning run with it
  • keep Worker-level registries for uniqueness checks and the close() sweep
  • keep K=1 submission: a later submit() still drains the prior run, so this adds no acceptance flight or multi-run admission

Ownership

_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.

Two things make that single pointer safe, and both are load-bearing:

  • _submit_mu makes it single-writer across concurrent submit() 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 concurrent handle.wait() on a prior run has already completed by then.

Domains end at their run's fence

A CommDomainHandle's release routes to the _RunResources of the run that allocated it, captured at allocation. That leaves no owner outside graph construction, so allocate_domain now rejects that case rather than allocating something nothing will free.

The Worker-level _pending_release_domains list goes away with it. Once the fence-owned path became the only drain, _execute_pending_domain_releases was always called with a run's resources and close() 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_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 same lock, so a
release 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 _released flag first leaves the sweep to skip
setting that flag and go straight to a CTRL_RELEASE_DOMAIN the release is
already making — a duplicate release of live backend memory, not a no-op.

_release_domain_now performs the release under its own lock and records the
result. 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, drop
it from _live_domains, and on the close() path continue into
_worker.close() while the release is still driving the mailboxes. That is
precisely what close() orders its domain sweep before the scheduler teardown
to avoid, and it would also break the public CommDomainHandle.freed contract
("backend release has executed").

A failure is replayed to later callers rather than retried, so the sweeps keep
the handle in _live_domains as a detectable residual instead of reporting a
release that never happened.

The reverse order was already safe: a sweep sets _released before 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.

  • changed-file pre-commit hooks: passed
  • Python unit tests: 836 passed, 2 skipped
  • C++ unit tests: 61/61 passed
  • a2a3 onboard, the exact CI command (examples tests/st minus the two SDMA demos, 4 devices under task-submit): 117 passed, 1 skipped, 0 failed — unchanged before and after the domain-ownership change, so the new allocate_domain guard interrupts no real L3 path
  • all five new tests were mutation-checked. Reverting the per-run cleanup fails test_run_finalization_releases_only_its_resources; replacing the allocate_domain guard with a silent default _RunResources() fails test_allocate_domain_outside_graph_construction_is_rejected; both halves of the retirement path — dropping the branch, and never setting the flag — fail test_domain_released_after_its_run_retired_is_freed_inline; and removing either domain_lock or the straggler drain fails test_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 gives allocation 7 reached the backend 2 times, claiming but freeing outside the lock gives the second caller returned while the owner was still releasing, and recording success for a failed release gives a failed release must not be reported as freed

The two red a5 jobs on the previous head were the a5 runner, not this change: st-onboard-a5 predates #1481 (which skips that job while the runner is under repair) and ut-a5 failed at HcclCommInitRootInfo failed: 1 in test_two_rank_comm_lifecycle / test_two_rank_allocate_release_round_trip. This PR touches no C++ or binding file, and ut-a5 is green on current main.

@coderabbitai

coderabbitai Bot commented Jul 24, 2026

Copy link
Copy Markdown

Review Change Stack

📝 Walkthrough

Walkthrough

The PR replaces drain-based orchestration with run-scoped lifecycle tracking, Python RunHandle submission, run-aware tensor bookkeeping, failure callbacks, and resource cleanup. It also adds pipeline-contract metadata and per-slot device stream management, with corresponding documentation and tests.

Changes

Run-fenced execution

Layer / File(s) Summary
Run state and tensor bookkeeping
src/common/hierarchical/*, tests/ut/cpp/hierarchical/test_orchestrator.cpp, tests/ut/cpp/hierarchical/test_tensormap.cpp
Orchestrator runs now track lifecycle, task counts, errors, and run-scoped TensorMap entries.
Failure propagation
src/common/hierarchical/scheduler.*, src/common/hierarchical/worker*, tests/ut/cpp/hierarchical/test_scheduler.cpp
Task and endpoint failures are routed through scheduler callbacks into the originating run.
Python run handles and cleanup
python/simpler/worker.py, python/bindings/worker_bind.h
Worker.submit() returns RunHandle; run-owned regions, remote references, domains, and native runs are finalized independently.
Lifecycle documentation and tests
docs/*, python/simpler/orchestrator.py, python/simpler/task_interface.py, tests/ut/py/test_worker/*
Documentation and tests describe run fences, explicit outcomes, timeout behavior, close handling, and deferred cleanup.

Pipeline execution

Layer / File(s) Summary
Pipeline contracts
src/common/worker/pto_runtime_c_api.h, src/a2a3/runtime/*/runtime_maker.cpp, src/common/worker/chip_worker.*
Runtime pipeline metadata is exported, optionally loaded, validated, and exposed through ChipWorker.
Slot-based device streams
src/a2a3/platform/onboard/host/device_runner.*, src/common/platform/onboard/host/device_runner_base.*
Device execution creates paired per-slot streams, launches kernels on them, synchronizes both streams, and destroys them through RAII-aware reap paths.

Estimated code review effort: 5 (Critical) | ~120 minutes

Possibly related PRs

Poem

A rabbit submits a run with a hop,
A fence waits softly for work to stop.
Tensor keys wear run IDs bright,
Streams pair up and launch just right.
Cleanup follows, no leaks in sight.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 18.22% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
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.
Title check ✅ Passed The title is concise and accurately summarizes the main change: isolating resources per worker run.
Description check ✅ Passed The description is detailed but clearly aligned with the resource-isolation changes in the pull request.

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

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

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (1)
src/a2a3/platform/onboard/host/device_runner.cpp (1)

226-234: 🚀 Performance & Scalability | 🔵 Trivial

Per-run stream create/destroy replaces the previously-reused persistent streams.

Each run() invocation now issues rtStreamCreate/rtStreamDestroy for a fresh AICPU/AICore pair instead of reusing DeviceRunnerBase's persistent streams across runs. This is clearly deliberate (K=2 scaffolding per the device_runner.h comments), but it does add two extra CANN stream lifecycle round-trips to every run() call on the hot path. Worth keeping an eye on measured per-run latency if run() 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

📥 Commits

Reviewing files that changed from the base of the PR and between 6284040 and 7719922.

📒 Files selected for processing (39)
  • docs/comm-domain.md
  • docs/hierarchical_level_runtime.md
  • docs/orchestrator.md
  • docs/remote-l3-worker-design.md
  • docs/scheduler.md
  • docs/task-flow.md
  • python/bindings/worker_bind.h
  • python/simpler/orchestrator.py
  • python/simpler/task_interface.py
  • python/simpler/worker.py
  • src/a2a3/platform/onboard/host/device_runner.cpp
  • src/a2a3/platform/onboard/host/device_runner.h
  • src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp
  • src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp
  • src/common/hierarchical/orchestrator.cpp
  • src/common/hierarchical/orchestrator.h
  • src/common/hierarchical/ring.cpp
  • src/common/hierarchical/ring.h
  • src/common/hierarchical/scheduler.cpp
  • src/common/hierarchical/scheduler.h
  • src/common/hierarchical/tensormap.cpp
  • src/common/hierarchical/tensormap.h
  • src/common/hierarchical/types.cpp
  • src/common/hierarchical/types.h
  • src/common/hierarchical/worker.cpp
  • src/common/hierarchical/worker_manager.cpp
  • src/common/hierarchical/worker_manager.h
  • src/common/platform/onboard/host/device_runner_base.cpp
  • src/common/platform/onboard/host/device_runner_base.h
  • src/common/worker/chip_worker.cpp
  • src/common/worker/chip_worker.h
  • src/common/worker/pto_runtime_c_api.h
  • tests/ut/cpp/hierarchical/test_orchestrator.cpp
  • tests/ut/cpp/hierarchical/test_scheduler.cpp
  • tests/ut/cpp/hierarchical/test_tensormap.cpp
  • tests/ut/py/test_worker/test_error_propagation.py
  • tests/ut/py/test_worker/test_host_worker.py
  • tests/ut/py/test_worker/test_l3_l2_orch_comm.py
  • tests/ut/py/test_worker/test_startup_readiness.py

Comment thread python/simpler/worker.py
@Crane-Liu
Crane-Liu force-pushed the codex/worker-async-pr6-run-resources branch from 7719922 to c22f0ab Compare July 24, 2026 23:59
@ChaoWao
ChaoWao force-pushed the codex/worker-async-pr6-run-resources branch from c22f0ab to 3d23abc Compare July 26, 2026 15:28
@ChaoWao ChaoWao changed the title (6/8) Refactor: isolate resources per worker run Refactor: isolate resources per worker run Jul 26, 2026
@ChaoWao

ChaoWao commented Jul 26, 2026

Copy link
Copy Markdown
Collaborator

Revision: rebased onto merged main, plus per-run domain ownership

Pushed 3d23abcf. Rebased by cherry-picking the tip onto merged main — the five superseded ancestor commits drop out, leaving 2 files / +220/-46. One conflict in test_host_worker.py, where main (#1482) and this branch each appended a test at the same spot; both kept.

The one thing I changed

Worker._pending_release_domains was left without a drain. _execute_pending_domain_releases has exactly one call site (_finalize_run_handle), and it now always passes a run's _RunResourcesRunHandle._resources is default-constructed, so it is never None. close() never touched that list either. So the resources is None branch could only ever queue a handle nothing would free.

It is unreachable today: _allocate_domain is only reachable through Orchestrator, which only exists during graph construction, where the pointer is set. But an unreachable defensive branch that silently leaks is worse than no branch — it implies a supported path that isn't.

So the branch and the field are gone, and allocate_domain now states the contract instead of pretending to cover it:

# A domain's lifetime ends at the fence of the run that allocated it,
# so there is no owner to charge it to outside graph construction.
resources = self._building_run_resources
if resources is None:
    raise RuntimeError("allocate_domain is only valid while a run's graph is being built")

The check sits after the existing level >= 3 / _worker is not None validations so the fail-fast ordering the surrounding comment calls out is preserved. _release_domain_handle and _execute_pending_domain_releases now take _RunResources as a required parameter.

_release_all_live_domains, _cleanup_l3_l2_regions and _release_active_remote_slot_refs keep their optional parameter — those fallbacks are reachable, from close() and _cleanup_partial_init(), and they do get drained.

New test: test_allocate_domain_outside_graph_construction_is_rejected. Mutation-checked — replacing the raise with a silent _RunResources() makes it fail.

Two notes for the record, no change made

  • A domain must not outlive its run. _release_fn closes over the allocating run's _RunResources, so a release() after that run's fence appends to a drained-and-abandoned list — and _release_domain_handle has already popped the handle from _live_domains, so close()'s sweep cannot catch it either. That is the intended contract of this PR, but it is currently only derivable from the closure. Worth one line on CommDomainHandle.
  • The field comment on _building_run_resources names only _submit_mu. That covers submit-vs-submit. The nesting case is actually covered by _submit_mu being non-reentrant, and the submit-vs-concurrent-wait() case by submit() draining prior handles before setting the pointer. Both are elsewhere.

Validation

826 passed / 2 skipped Python, 61/61 C++, and the exact onboard CI command on 4 devices: 117 passed, 1 skipped, 0 failed — identical before and after the domain change, so the new guard interrupts no real L3 path.

The two red a5 jobs on the previous head were the runner

Earned, not assumed: this PR touches no C++ or binding file; ut-a5 is green on current main; and the failure is HcclCommInitRootInfo failed: 1 in test_two_rank_comm_lifecycle / test_two_rank_allocate_release_round_trip — HCCL bring-up on the same a5 card whose st-onboard-a5 job #1481 now skips while it is under repair. That run also predates #1481.

@ChaoWao

ChaoWao commented Jul 26, 2026

Copy link
Copy Markdown
Collaborator

CI status on 3d23abcf: everything green except ut, which is main's red, not this PR#1504 (already open) fixes it.

test_orphan_child_reaping::test_children_exit_when_parent_is_killed scrapes the forked children's pids out of the parent subprocess' combined stdout+stderr with if tok.isdigit(). The multiprocessing.resource_tracker warning that follows the parent's deliberate SIGKILL — "There appear to be 3 leaked shared_memory objects" — contributes a bare 3, so the parse reports four pids:

E  AssertionError: expected exactly 3 forked sub-worker pids, parent reported [6510, 6511, 6512, 3]

Earned rather than assumed:

  • the test arrived in Fix: exit a forked worker child once its parent is gone #1495 and does not exist on this branch — CI runs the merge commit, so it comes entirely from main;
  • it reproduces on plain main locally (1 failure in 20 runs at 3dc617aa), with the identical [..., ..., ..., 3] signature;
  • the stray token is 3 both on main and here, so nothing in this PR changed the leak count;
  • this PR touches only worker.py's per-run resource bookkeeping and one test file, neither of which is on that path.

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: st-onboard-a2a3, ut-a2a3, ut-a5 (the two a5 jobs that were red on the previous head are now green — that was the a5 runner, as diagnosed above), st-sim-a2a3, st-sim-a5, packaging-matrix, profiling-flags-smoke, pre-commit. st-onboard-a5 skips per #1481.

@ChaoWao

ChaoWao commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Revision: a domain released after its run retired now frees inline

Pushed 1ac34f98, rebased onto merged main (22f2d533).

The fix

codex review flagged worker.py:5234-5235 and I confirmed it with a probe before changing anything:

after finalize:                 freed = ['d'] | handle.freed = True
after late release:
  in worker._live_domains?      False       <- popped, so close()'s live sweep cannot see it
  queued on retired run?        ['late']    <- on a queue that is never drained again
after a LATER run's fence:       h2.freed = False
after close-style live sweep:   h2.freed = False

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.

_RunResources.retired is set once the owning fence has drained the queues for the last time. _release_domain_handle sees it and frees inline, with the same log-and-continue error handling the deferred path uses:

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 freed = ['d', 'late'], h2.freed = True, empty queue.

test_domain_released_after_its_run_retired_is_freed_inline pins it, mutation-checked from both sides — dropping the retired branch fails it, and never setting the flag fails it.

Reachability, stated plainly

This needs user code holding a CommDomainHandle past its run. Nothing in-repo does, and close() waits for _active_ops to drain before touching accepted handles, so it cannot interleave with graph construction either. codex described it as a race; I traced the call paths and could not find a concurrent window through the supported API, so I would call it out-of-contract usage rather than a live race — but the silent leak is real either way, and it is a failure mode this PR introduced.

Validation

  • Python unit: 832 passed, 2 skipped
  • C++ unit: 61/61
  • a2a3 onboard, the exact CI command on 4 devices: 117 passed, 1 skipped, 0 failed — unchanged across all three revisions of this PR
  • changed-file pre-commit: passed

One flaky test, tracked separately

Two of my full local suite runs failed on test_callable_identity.py::test_remote_sim_* with ConnectionRefusedError, in the window where an onboard job and two review CLIs were also running on this box. I could not reproduce it afterwards: 9 branch runs passed vs 2 failed, including three run in parallel with each other, and main is 5/5 clean. The tests spawn a remote_l3_worker subprocess, time.sleep(0.3), then connect — nothing this PR touches can influence when that daemon binds its port. Filed as #1508 with the five call sites and a suggested readiness poll. The same family produced the earlier ut (macos-latest) red on this PR, which cleared on re-run.

@ChaoWao
ChaoWao force-pushed the codex/worker-async-pr6-run-resources branch from 1ac34f9 to f1476c2 Compare July 27, 2026 04:15
@ChaoWao

ChaoWao commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Revision: the retirement window is closed with a lock, not narrowed

Pushed f1476c2a. The external review was right, and my previous revision only made the hole smaller.

Proof it was still open

I forced the exact interleaving with a dict subclass that blocks the releasing thread between its retired read and its append:

retired            : True
freed              : ['other']     <- only the other domain
h.freed            : False         <- the raced handle, never freed
stranded on queue  : ['d']         <- on a queue nothing reads again
reachable via close: False         <- the release itself popped it from _live_domains

The retired flag alone moved the window from "any time after the fence" to "the few bytecodes between the check and the append". A narrowed race is worse than a wide one — harder to reproduce, harder to attribute.

One correction to the report's step order, for anyone reproducing it: the per-run live_domains.pop happens before the retired read, and the worker-wide pop after it. The outcome is the same, but a repro has to follow the real order to hit the window.

Fix

_RunResources.domain_lock makes the release transition and retirement mutually exclusive. The release does its whole transition under it and frees outside it; retirement flips the flag and claims any straggler under the same lock:

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 CTRL_RELEASE_DOMAIN for a domain already gone. _free_domain_after_fence returns early on handle.freed.

The concurrency test the report asked for

test_domain_released_while_its_run_retires_is_still_freed parks the fence in its lock-free live-domain sweep while a second thread runs a release to completion, then releases the fence. Parking it inside the lock instead just deadlocks — my first repro script only completed because of a 5 s wait timeout, which proves nothing about the ordering, so the test gates on a point the lock does not cover.

Mutation-checked both ways: removing domain_lock and removing the straggler drain each fail it with AssertionError: a release racing retirement was stranded on a drained queue. 40 consecutive runs, 0 failures — the test itself is not flaky.

Validation

  • Python unit: 833 passed, 2 skipped
  • C++ unit: 61/61
  • a2a3 onboard, the exact CI command on 4 devices: 117 passed, 1 skipped, 0 failed

One earlier onboard run on this commit hit TestL3Group::test_run with a child signal 11 and a preceding mailbox has an unresolved timed-out control command. Not reproduced — 4/4 in isolation and a clean full rerun on the same commit, with three clean full runs on earlier revisions of this PR. One occurrence in five; I have not been able to tie it to this change, and I am flagging it rather than filing it as this PR's.

@ChaoWao
ChaoWao force-pushed the codex/worker-async-pr6-run-resources branch from f1476c2 to afaefe2 Compare July 27, 2026 05:07
@ChaoWao

ChaoWao commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Revision: one allocation, one backend release

Pushed afaefe2c. The double-free finding is right; the interleaving in the report is not the one that reaches it, and finding the real one is what gave the test teeth.

The stated order is already safe

_release_all_live_domains sets handle._released = True before the backend call, and CommDomainHandle.release() no-ops when _released is already set. So sweep-snapshots → retire → release cannot double free: the release returns at the first line. I wrote the test that way first and it passed against a build with the claim removed — no teeth.

(The report's first two steps also cannot interleave: _release_all_live_domains and _retire_run_domains are sequential _step calls on the fence thread, so the sweep has returned before retirement runs.)

The order that does reach it

Release first, sweep second. The release wins _released and takes the post-fence path; the sweep, still holding that handle in a snapshot it took earlier, sees handle.released already True, skips setting it, and goes straight to _release_domain_now. Both issue CTRL_RELEASE_DOMAIN for live backend memory.

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:

E  AssertionError: allocation 7 was released 2 times

Fix

_release_domain_now claims the allocation under _domain_free_mu and performs the release outside it:

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 freed a moment before the owning path's backend call returns; that is stated on the method.

Validation

  • Python unit: 834 passed, 2 skipped; the new test 30/30 stable
  • C++ unit: 61/61
  • a2a3 onboard, the exact CI command on 4 devices: 117 passed, 1 skipped, 0 failed

ut-a5 is red on the a5 runner and is being disregarded for now per the maintainer.

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.
@ChaoWao
ChaoWao force-pushed the codex/worker-async-pr6-run-resources branch from afaefe2 to 18cc5a0 Compare July 27, 2026 06:50
@ChaoWao

ChaoWao commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Revision: the second caller waits for the outcome, it does not just skip

Pushed 18cc5a04. This finding is right, and it caught a mistake I made for a bad reason: in the previous revision I moved the backend release outside the claim lock because holding it made my test take 5 s. That traded a correctness invariant for test runtime, and I even wrote a comment rationalising it.

What claim-and-return broke

close() orders its domain sweep before _worker.close() for a stated reason — "once dw.close() runs the chip mailboxes are unusable and we can no longer drive CTRL_RELEASE_DOMAIN". That ordering assumes the sweep returns when the releases have completed, not when they have been claimed. Claim-and-return let close() mark the handle freed, drop it from _live_domains, and proceed into teardown with a release still driving those mailboxes. It also contradicted the public CommDomainHandle.freed contract, and erased the residual on failure — the sweep keeps an un-freed handle in _live_domains precisely so close() reports it.

Fix

_release_domain_now performs the release under _domain_free_mu and records the result in _domain_free_results. A second caller blocks, then replays the outcome:

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] = None

Concurrent 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 close()'s ordering guarantee is worth more than their concurrency.

The three properties, each with its own test and its own mutation

property test mutation that catches it
one backend call per allocation ..._free_a_domain_once no claim → allocation 7 reached the backend 2 times
the second caller waits, freed stays false while in flight test_second_domain_release_waits_for_the_first claim then free outside the lock → the second caller returned while the owner was still releasing
a failure is replayed, residual stays visible test_failed_domain_release_is_replayed_to_a_second_caller record success regardless → a failed release must not be reported as freed

Two notes on how these are driven, since my first attempt at each was wrong:

  • The waiting and failure tests drive _release_domain_now directly. Going through release() does not work: the post-fence path pops the handle from _live_domains before the backend call, so a concurrent sweep snapshots nothing, blocks on nothing, and leaves no residual to assert on.
  • ..._free_a_domain_once still goes through release(), because the double free specifically needs the release to win the _released flag ahead of the sweep.

Validation

  • Python unit: 836 passed, 2 skipped; the three domain-release tests 40/40 stable
  • C++ unit: 61/61
  • a2a3 onboard, the exact CI command on 4 devices: 117 passed, 1 skipped, 0 failed

@ChaoWao
ChaoWao merged commit da75d35 into hw-native-sys:main Jul 27, 2026
15 of 16 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants