Refactor: run native progress on the child loop - #1694
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:
📝 WalkthroughWalkthroughNative runs now use explicit ChangesNative run lifecycle
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant CAPI as C API
participant Runner as DeviceRunner
participant Progress as Child progress path
participant Context as NativeRunContext
CAPI->>Runner: enqueue_run
Runner->>Progress: start runner work
Progress->>Context: publish acceptance and completion
CAPI->>Runner: poll_run or drain_run
Runner-->>CAPI: completion result
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.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/common/worker/chip_worker.cpp (1)
704-717: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winReintroduce an in-flight launch marker before calling
launch_run_fn_.
launch_native_run()leaves the slot asPREPAREDbetween the validation lock and theLAUNCHEDwrite. In that window another thread can start the same token, andcleanup_native_runs_noexcept()treats it as cleanable and callsfinalize_run_fn_()while launch is still in flight. RestoreLAUNCHINGinside the first lock, mark it as skip/drain in cleanup if needed, and restore the pre-launch/reaped phase on the failure path.🤖 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/common/worker/chip_worker.cpp` around lines 704 - 717, Update launch_native_run() to set the selected NativeRunSlotState phase to LAUNCHING while holding the initial native_run_mu_ lock, before invoking launch_run_fn_. Ensure cleanup treats LAUNCHING slots as in-flight and skips or drains them instead of finalizing them. On launch failure, restore the slot to the appropriate PREPARED or REAPED phase based on polling, while retaining the existing wait result and error behavior.
🤖 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 `@src/common/worker/chip_worker.cpp`:
- Around line 704-717: Update launch_native_run() to set the selected
NativeRunSlotState phase to LAUNCHING while holding the initial native_run_mu_
lock, before invoking launch_run_fn_. Ensure cleanup treats LAUNCHING slots as
in-flight and skips or drains them instead of finalizing them. On launch
failure, restore the slot to the appropriate PREPARED or REAPED phase based on
polling, while retaining the existing wait result and error behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 1ea67b34-9dd9-4564-bc54-7bf26b5d796d
📒 Files selected for processing (23)
docs/chip-level-arch.mddocs/dfx/host-trace.mddocs/dynamic-linking.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/platform/onboard/host/device_runner.cppsrc/a5/platform/sim/host/device_runner.cppsrc/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/onboard/host/device_runner_base.cppsrc/common/platform/onboard/host/device_runner_base.hsrc/common/platform/sim/host/c_api_shared.cppsrc/common/platform/sim/host/device_runner_base.cppsrc/common/platform/sim/host/device_runner_base.hsrc/common/worker/chip_worker.cppsrc/common/worker/chip_worker.hsrc/common/worker/native_run_context.hsrc/common/worker/native_run_launch_signal.htests/ut/cpp/CMakeLists.txttests/ut/cpp/common/test_native_run_launch_signal.cpp
💤 Files with no reviewable changes (5)
- src/a2a3/platform/sim/host/device_runner.cpp
- tests/ut/cpp/CMakeLists.txt
- src/common/worker/native_run_launch_signal.h
- src/a5/platform/sim/host/device_runner.cpp
- tests/ut/cpp/common/test_native_run_launch_signal.cpp
8075ab2 to
4a997f5
Compare
2d8604f to
63dbd37
Compare
- enqueue native runs directly from the child-owned launch path and drain them from wait or finalize\n- remove the per-run executor handoff while preserving sticky acceptance and trace spans
63dbd37 to
8314b78
Compare
ChaoWao
left a comment
ChaoWao
left a comment
There was a problem hiding this comment.
Reviewed at 8314b780 (base 0e3851a4). This is the best-shaped version of B6c I've seen — the rebase onto #1714 both resolved the two red a5 checks and shrank the diff from 32 files / +1413 −1077 to 21 files / +156 −508. Deleting the TLS transfer mechanism outright rather than replacing it with a shim, and dropping Launching from both state machines, are the right calls. publish_acceptance keeping receipt.matches(identity()) as the only key preserves the plan's "acceptance only from LaunchReceipt" invariant, and the plan's explicit prohibition on temporary TLS park/adopt is not just respected but exceeded.
Two things I'd like addressed before merge, one of them a one-liner.
1. Must fix — drain_execution now runs without an attached device context
Both new drain call sites run before any attach_current_thread:
simpler_wait_run— onboardc_api_shared.cpp:822, sim:729— no attach anywhere in the function.simpler_finalize_run— onboard:848, sim:753— the function does attach, but at:859/:764, i.e. after the drain block.
Every other lifecycle entry point attaches first (prepare :714, launch :771, poll :807), and the code this replaces attached explicitly: the executor lambda's first act was attach_current_thread, with drain_execution inside the attach_rc == 0 branch.
The attach isn't decorative — attach_current_thread → rtSetDevice, and its own comment says "CANN device context is per-thread, so every caller must attach explicitly." drain_execution then does reap_run → aclrtSynchronizeStreamWithTimeout on both streams, retire_run_aicore_stream (stream destroy), print_handshake_results (device memory read), and via its RAII guard cleanup_execution → rtFree. All of those need the context.
Why CI is green anyway: the child progress loop is single-threaded, so the launch-time attach on the same thread covers both drains. That makes this latent, not currently-broken. But simpler_wait_run is public C API and its binding is _wait_native_run with nb::call_guard<nb::gil_scoped_release>() — explicitly designed to be called off the GIL — and nothing documents a same-thread requirement. native_run_lifecycle's ST calls it from the prepare thread, so it can't discriminate.
Fix is the same shape as the siblings, and cheap — rtSetDevice on an already-attached thread is idempotent, which is why poll calls it unconditionally every iteration:
int attach_rc = state->runner->attach_current_thread(state->runner->device_id());
if (attach_rc != 0) { /* record completion_rc, publish Complete, return */ }In finalize, moving the existing attach above the drain block would do it.
2. Should fix — the dep_gen host-graph's same-thread invariant is now implicit, and its docs describe the deleted mechanism
dep_gen_host_graph keeps its capture in static thread_local HostGraphState (dep_gen_host_graph.cpp:233). The graph is built during prepare (bind_callable_to_runtime_impl → run_host_orchestration → begin_capture(), runtime_maker.cpp:474) and emitted during drain (reap_run → dep_gen_host_graph_emit, device_runner.cpp:675). Pre-PR those were different threads — which is exactly why take_native_run_thread_state / adopt_native_run_thread_state existed.
Deleting that transfer is correct now that prepare and drain land on one thread. But the invariant that makes it correct is unstated and unenforced, and its failure mode is quiet: emit returns -3 and logs "no capture was adopted on this thread" while the run itself succeeds and no deps.json appears.
Two concrete bits:
dep_gen_host_graph.h:34,39-43(both arches) still document the deleted mechanism. The PR edits that paragraph but only swaps "executor's thread-local state" → "progress thread's local state", so it still claims prepare moves the graph into run-owned storage and launch adopts that snapshot — neither happens now. Same for the emit error string (dep_gen_host_graph.cpp:527), which still refers to "the run-owned snapshot". Perdoc-consistency.md§3 this needed a rewrite rather than a word swap.take_capture/adopt_capture/destroy_captureremain defined and exported in both a2a3 and a5 hbg runtimes with zero production callers — onlytests/ut/cpp/a2a3/test_dep_gen_host_graph.cppuses them. They read as live API.
Either resolution works for me: delete the now-dead capture API (preferred — removes the ambiguity entirely), or keep it and state the invariant as a present-tense fact at the emit site. What I'd avoid is shipping the current state, where the enforcement mechanism is gone but its documentation and exported surface both remain.
Worth noting the safety net is real: the hbg dep_gen ST is wired into a2a3 onboard CI with --enable-dep-gen (_st-npu-a2a3.yml:142) and asserts all 6 edges — so a genuine break would show up in st-onboard-a2a3.
Smaller items
3. Two deleted tests covered invariants that survive. test_native_run_launch_signal.cpp goes away wholesale. Three of its five cases tested the CV handoff and are correctly gone, but two didn't:
StaleReceiptCannotPublishAcceptance—publish_acceptancestill implements exactly this gate, and it's the load-bearing half of "acceptance only fromLaunchReceipt".- the first half of
PublishAcceptanceStoresOnceAndKeepsNotificationSticky— asserts the store lands with the right value.
Both port to publish_acceptance nearly verbatim, no threading needed. As it stands the acceptance-identity gate — the one mechanism stopping a stale run from marking a task accepted — loses its only direct test; test_native_run_execution.cpp covers the LaunchReceipt side but never the accepted_state store.
4. catch (...) {} around both drains swallows the reason. The rc still propagates (it stays at the -1 initializer), so the logic is right — but pre-PR that silent catch sat at a thread boundary where an escaping exception would terminate; here it's a plain call in a function that already returns rc. drain_execution is where op-timeouts and 507018s surface, so a LOG_ERROR in each handler is probably the highest-value line in the file for triage.
5. launched changed meaning in finalize, for the better — worth a line in the body. It went from phase != Prepared ("launch was attempted") to active_execution != nullptr ("launch produced an active execution"). These differ for a NotStarted launch failure, and the new one is more correct: such a run never touched a stream, so validate_runtime_impl should get -1 and set_gm_sm_ptr(nullptr) should run — which is what the new code does. A real behavior improvement hiding inside a refactor; a future bisect would otherwise attribute it to nothing.
CI
The body still says "Simulation validation was intentionally not run for this change" — that's now moot: all four sim jobs are green, as are ut-a5 and st-onboard-a5 (both red before the rebase), st-pod-onboard-a2a3, profiling-flags-smoke, and both ut hosts. Only st-onboard-a2a3 and one packaging/sim job are still in flight at the time of writing — and st-onboard-a2a3 is the one carrying the hbg dep_gen assertion above, so it's worth waiting on specifically.
Worth updating that line in the body once it settles, so the PR record doesn't understate what was actually verified.
|
@ChaoWao Addressed the B6c review:
|
2 participants
Summary
LaunchingphaseValidation
pip install --no-build-isolation -e .(current-source onboard build)pytest tests/ut/py/test_worker/test_host_worker.py -q(216 passed)task-submit: a2a3 native run lifecycle (1 passed)task-submit: a2a3 worker async FIFO (3 passed)LCW CANN editable build and
test_native_run_executionpassed at the current head.\n- The rebased CI passed all four simulation jobs, both host UT jobs,ut-a5,st-onboard-a5,st-pod-onboard-a2a3, and profiling smoke;st-onboard-a2a3includes the host-build-graph dep_gen assertion.