Attach the device before drain and retire the dead dep_gen capture API - #1725
Merged
Conversation
simpler_wait_run and simpler_finalize_run reached drain_execution without
an attached device context. drain_execution synchronizes and destroys
streams, reads device memory via print_handshake_results, and frees device
allocations through its cleanup guard — all per-thread CANN state. Only the
single-threaded child progress loop kept this from surfacing: the
launch-time attach happened to cover both drains. simpler_wait_run is
public C API, bound as _wait_native_run under gil_scoped_release, so no
caller is obliged to stay on that thread.
Attach at both entry points, matching prepare, launch and poll. rtSetDevice
is idempotent on an already-attached thread, which is why poll already
calls it unconditionally. In finalize the existing attach moves above the
drain so one call covers both the drain and validate_runtime_impl. Both
attaches sit inside a try: attach_current_thread is not noexcept and these
are C API boundaries.
The silent `catch (...) {}` around each drain now logs. Drain is where
op-execute timeouts and 507018 cascades surface, and the rc alone does not
say which.
Deleting the executor handoff left the dep_gen host graph correct but
undocumented: the capture is thread_local, built by the orchestration
during prepare and read by emit during drain, which is sound only because
both now land on one thread. That invariant was unstated while the headers
still described the take/adopt transfer that used to enforce it, and
take_capture / adopt_capture / destroy_capture stayed exported in both
arches with no production callers. State the invariant where the graph
lives, and remove the dead API rather than leave it reading as live.
test_native_run_launch_signal.cpp was deleted with the CV handoff it
covered, but its acceptance-identity case tested publish_acceptance, which
survives — and that gate is the only thing stopping a stale run from
marking a task accepted. test_native_run_acceptance.cpp restores direct
coverage; mutating the receipt.matches check fails it. The dep_gen UT is
rewritten from the handoff mechanism to the same-thread invariant,
including the cross-thread case that must emit nothing.
Verified: 73/73 cpp UT; 218 host-worker py UT; a2a3 onboard
host_build_graph 24 passed / 2 skipped, tensormap 42 passed; a2a3sim
host_build_graph 20 passed / 4 skipped. The dep_gen ST asserting all 6
edges passes onboard, confirming deps.json is still produced.
📝 WalkthroughWalkthroughThe PR changes host dependency-graph capture from cross-thread snapshot transfer to same-thread, thread-local emission. Native and simulated run finalization now attach the calling thread before draining and validation. Tests cover graph emission and native-run acceptance. ChangesHost graph lifecycle
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant RunCaller
participant simpler_finalize_run
participant DeviceRunner
participant HostBuildGraph
RunCaller->>simpler_finalize_run: finalize run
simpler_finalize_run->>DeviceRunner: attach calling thread
simpler_finalize_run->>DeviceRunner: drain execution
DeviceRunner->>HostBuildGraph: emit same-thread graph
simpler_finalize_run->>DeviceRunner: validate runtime state
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@docs/dfx/dep-gen.md`:
- Line 398: Update the dep_gen support matrix to consistently include a5
host-direct orchestration: revise the statements around the a2a3-only host
support and the claim that a5 lacks dep_gen, while retaining the existing a5
device-runner hookup details. State the supported a5 behavior once and ensure
all three related entries agree.
🪄 Autofix
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: 13356f8c-0ae1-42d2-9071-03bc42121165
📒 Files selected for processing (12)
docs/dfx/dep-gen.mdsrc/a2a3/platform/onboard/host/device_runner.cppsrc/a2a3/platform/sim/host/device_runner.cppsrc/a2a3/runtime/host_build_graph/host/dep_gen_host_graph.cppsrc/a2a3/runtime/host_build_graph/runtime/dep_gen_host_graph.hsrc/a5/runtime/host_build_graph/host/dep_gen_host_graph.cppsrc/a5/runtime/host_build_graph/runtime/dep_gen_host_graph.hsrc/common/platform/onboard/host/c_api_shared.cppsrc/common/platform/sim/host/c_api_shared.cpptests/ut/cpp/CMakeLists.txttests/ut/cpp/a2a3/test_dep_gen_host_graph.cpptests/ut/cpp/common/test_native_run_acceptance.cpp
This was referenced Aug 8, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Follow-ups to the review on #1694, which merged as
4f417732.1.
drain_executionran without an attached device contextsimpler_wait_runandsimpler_finalize_runboth reacheddrain_executionbefore any
attach_current_thread. Drain synchronizes and destroys streams,reads device memory via
print_handshake_results, and frees deviceallocations through its cleanup guard — all per-thread CANN state, and
attach_current_thread's own comment says "CANN device context isper-thread, so every caller must attach explicitly."
This was latent rather than broken: the child progress loop is
single-threaded, so the launch-time attach happened to cover both drains.
But
simpler_wait_runis public C API, bound as_wait_native_runundernb::call_guard<nb::gil_scoped_release>(), and nothing documents asame-thread requirement.
Both entry points now attach first, matching prepare / launch / poll.
rtSetDeviceis idempotent on an already-attached thread, which is why pollalready calls it unconditionally every iteration. In finalize the existing
attach simply moves above the drain, so one call covers the drain and
validate_runtime_implboth.Each attach sits inside a
try:attach_current_threadis notnoexcept(onboard calls
rtSetDevice; sim'sacquire_deviceallocates), and theseare C API boundaries where an escaping exception would
terminate.2.
catch (...) {}around the drains now logsThe rc still propagates unchanged. Drain is where op-execute timeouts and
507018 cascades surface, and the rc alone doesn't say which — so the reason
is worth a line.
3. dep_gen host graph: invariant stated, dead API removed
Deleting the executor handoff left the capture correct but undocumented.
The graph is
thread_local, built by the orchestration during prepare andread by
emitduring drain — sound only because both now land on the childprogress loop's single thread. That invariant was unstated while the headers
still described the
take/adopttransfer that used to enforce it, andtake_capture/adopt_capture/destroy_captureremained exported inboth arches with zero production callers.
Taking the reviewer's preferred resolution: state the invariant where the
graph lives, and delete the dead API rather than leave it reading as live.
Header prose, the emit error string, the device-runner comments and
docs/dfx/dep-gen.mdare rewritten rather than word-swapped.4. Restored test coverage for a surviving invariant
test_native_run_launch_signal.cppwas deleted along with the CV handoff itcovered. Three of its five cases went with the mechanism, correctly — but the
acceptance-identity case tested
publish_acceptance, which survives, andthat gate is the only thing stopping a stale run from marking a task
accepted.
test_native_run_acceptance.cpprestores direct coverage(matching receipt stores, stale receipt does not, null
accepted_stateisaccepted); mutating the
receipt.matchescheck makes it fail, so it is notvacuous.
The dep_gen UT is rewritten from the deleted handoff mechanism to the
surviving same-thread invariant, including the cross-thread case that must
emit nothing rather than a partial
deps.json.Not addressed
CodeRabbit asked to reintroduce a
LAUNCHINGin-flight marker inchip_worker.cpp.chip_worker.h:385-389states all ChipWorker callers runon the thread that called
init(), sofinalize()cannot race an in-flightlaunch_native_run()and the phase gap it describes is not reachable. Thisalso matches the review, which called dropping
Launchingfrom both statemachines the right call.
Validation
73/73cpp UT (ctest -L no_hardware)218passed —tests/ut/py/test_worker/test_host_worker.pyhost_build_graph24 passed / 2 skipped;tensormap_and_ringbuffer42 passedhost_build_graph20 passed / 4 skipped--enable-dep-gen, confirmingdeps.jsonis still produced on thesingle-thread path.