Fix: release prepared collectors and keep arming failures safe

[ChaoWao]

Summary

Follow-up to #1700, fixing two defects in the launch seam it introduced.

Prepared collectors were never released unless the run launched. A prepared
execution initializes its diagnostics collectors during native prepare, but
cleanup_execution() gated finalize_collectors() on a launched flag that
abandon_prepared_execution() and the prepare rollback both pass as false. Three
reachable paths leaked as a result:

• prepare fails after one or more init_* succeeded;
• the run is prepared but never launched (finalize without launch, or the
  !entered_run abandon after an attach_current_thread failure);
• launch returns NotStarted after the AICore closure already started the
  collectors — mark_submitted failure, or launch_aicore_kernel failure, i.e.
  the 207001 path this file carries a recovery routine for. Here the mgmt/poll
  threads are already running and nothing stops them.

Each leaked the collector's shared memory and its host registration, and left
is_initialized() set — and every collector refuses a second initialize() while
its shm is still mapped (return -1), so a leaked collector makes the next
prepare on that runner fail outright.

The gate was guarding against a state that cannot occur: an overlapping
predecessor can own no collector, because try_reserve_native_run admits a
prepared successor only when the successor declares no diagnostics and the
predecessor's own permits_prepared_successor says the same. With no overlap
possible, every finalize_collectors() branch is already is_initialized()-guarded
and is a no-op in exactly the case the gate was protecting. Released on every path
now. a5 keeps the flag for the sticky terminal poll state, which only a launched
run may publish.

An arming failure poisoned the device. ExactLaunchTransaction classified any
exception escaping the AICore step as an uncertain partial launch. But that step
arms the run before it submits — latching launch shape, resetting the device-wall
buffer, starting collectors, publishing stream ownership, and on sim the whole
dlsym'd setter block — so a std::bad_alloc from the core_types vector or a
std::system_error from create_thread marked the device unusable and forced a
reset at finalize, with no stream ever touched. contracts.md: "Failure before the
first stream submission is FAILED_SAFE."

Each submit callback now catches its own arming failures and returns non-zero, so
the transaction's exception path means what the contract says: the submission
itself was attempted and the outcome is uncertain. The callback contract is stated
on ExactLaunchTransaction.

Testing

Regression barrier for the collector leak — test_native_run_lifecycle.py gains
a diagnostics-enabled prepare → finalize-without-launch → prepare cycle. The leak is
onboard-only (sim's abandon_prepared_execution already released collectives
unconditionally via cleanup_active_run()), so it was verified on a2a3 silicon both
ways:

• with the launched gate restored byte-for-byte, on devices 3,5 —
  fails, with exactly the error the leak causes:

  ERROR init: [scope_stats_collector.cpp:62] ScopeStatsCollector already initialized
  ERROR prepare_execution: [device_runner.cpp:425] init_scope_stats failed: -1
  RuntimeError: prepare_native_run failed with code -1 (... run_epoch=4)
  

• with the fix — passes.

The test's STRACE invocation counts move 3→5 (sim) and 7→9 (onboard); the two extra
invocations are the abandoned prepares, which never reach simpler_run.runner_run,
so launched_count is unchanged.

The arming-classification fix ships without a regression test. The
ArmingFailureReportedAsRcStaysSafe case added here documents the callback contract,
but it passes against the pre-fix code too — the if (result.rc != 0) return result;
early return predates this change — so it is not a barrier. Guarding it would mean
injecting an allocation or thread-spawn failure into the production closures, which
has no hook today. Flagging rather than implying coverage.

Other evidence, all on the current head:

• C++ unit tests — 89/89 pass
• Simulation — a2a3sim and a5sim scene suites exit 0
• Hardware — a2a3 onboard suite via task-submit on free devices, exit 0

An earlier onboard run on --device auto reported three 507018s; all were at
LoadAicpuOp::BootstrapDispatcher inside simpler_init, before any run is prepared
or launched, and all on device 1, which task-submit --list showed held by another
user's live task. Re-running pinned to free devices is the green result above.

[coderabbitai]

📝 Walkthrough

Walkthrough

The change makes launch arming exception-safe across onboard and simulated runners. Pre-submission failures return -1 without starting work. Cleanup now finalizes owned collectors during rollback and drain paths.

Changes

Execution lifecycle

Layer / File(s)	Summary
Launch transaction contract and validation src/common/worker/native_run_execution.h, tests/ut/cpp/common/test_native_run_execution.cpp	The launch contract documents submission outcomes. A unit test verifies that an arming std::bad_alloc returns -1, preserves NotStarted, and avoids AICPU submission.
Exception-safe launch arming src/a2a3/platform/{onboard,sim}/host/device_runner.cpp, src/a5/platform/{onboard,sim}/host/device_runner.cpp	Setup, collector startup, initialization, stream publication, and handshake preparation are guarded before submission or simulated-core creation. Exceptions are logged and converted to -1.
Rollback and collector cleanup ownership src/a2a3/platform/onboard/host/device_runner.{cpp,h}, src/a5/platform/onboard/host/device_runner.{cpp,h}	cleanup_execution no longer accepts launched. Owned collectors are finalized during rollback and drain cleanup, while terminal poll publication remains limited to launched executions.

Estimated code review effort: 4 (Complex) | ~45 minutes

Possibly related PRs

• hw-native-sys/simpler#1447: Modifies collector cleanup in the same a2a3 and a5 device runners.
• hw-native-sys/simpler#1462: Overlaps the a2a3 launch and cleanup lifecycle.
• hw-native-sys/simpler#1700: Builds on the same launch transaction and PreparedExecution cleanup paths.

Sequence Diagram(s)

sequenceDiagram
  participant DeviceRunner
  participant Collectors
  participant ExecutionStream
  participant Worker
  DeviceRunner->>Collectors: Start per-run collectors
  DeviceRunner->>ExecutionStream: Publish prepared execution
  DeviceRunner->>Worker: Submit execution
  DeviceRunner-->>DeviceRunner: Return -1 if arming throws
  DeviceRunner->>Collectors: Finalize owned collectors during cleanup

Loading

Poem

A rabbit checks the launcher's gate,
No thread begins when setup waits.
Collectors rest when errors call,
Streams stay clean through drain and fall.
Hop, hop—the rollback path is straight!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% 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 clearly summarizes the two main changes: collector cleanup and safe handling of arming failures.
Description check	✅ Passed	The description directly explains both defects, the fixes, affected paths, and test evidence.

---

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

• X
• Mastodon
• Reddit
• LinkedIn

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

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

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

585-632: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Arming guard is correct; the failure log text is narrower than the guarded region.

The try block covers mark_submitted(slot) at Line 604. mark_submitted publishes the stream set. If a later statement in the block throws, the catch at Line 629 logs "arming failed before any stream submission", which is inaccurate for that window.

The window is small. The statements after mark_submitted only read runtime.get_workers() and write aicore_done. The rollback path still retires the AICore stream through abandon_prepared_execution, so no resource leaks. Consider moving mark_submitted and the aicore_done reset out of the try block, or adjusting the log text, so the message matches the guarded region.

🤖 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 585 - 632,
Adjust the guarded launch sequence around mark_submitted and the aicore_done
reset so the catch message accurately describes all operations it covers. Either
move mark_submitted and the subsequent worker-state reset outside the try block,
or broaden the catch log in the launch_run path to avoid claiming that no stream
submission occurred.

🤖 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 `@src/a5/platform/onboard/host/device_runner.cpp`:
- Around line 341-388: Move the run_poll_slot_ and run_poll_state_ stores out of
the rollbackable arming try block in exact_launch_transaction. Perform them only
after arming completes successfully and immediately before the actual submission
path, ensuring exceptions from dep_gen_collector_.start or core-type allocation
cannot publish a stale slot/state pair.

---

Nitpick comments:
In `@src/a2a3/platform/onboard/host/device_runner.cpp`:
- Around line 585-632: Adjust the guarded launch sequence around mark_submitted
and the aicore_done reset so the catch message accurately describes all
operations it covers. Either move mark_submitted and the subsequent worker-state
reset outside the try block, or broaden the catch log in the launch_run path to
avoid claiming that no stream submission occurred.

🪄 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: 4ddd43a5-6940-48d9-9d25-60e68a89878c

📥 Commits

Reviewing files that changed from the base of the PR and between 64fe416 and 92da7d0.

📒 Files selected for processing (8)

• src/a2a3/platform/onboard/host/device_runner.cpp
• src/a2a3/platform/onboard/host/device_runner.h
• src/a2a3/platform/sim/host/device_runner.cpp
• src/a5/platform/onboard/host/device_runner.cpp
• src/a5/platform/onboard/host/device_runner.h
• src/a5/platform/sim/host/device_runner.cpp
• src/common/worker/native_run_execution.h
• tests/ut/cpp/common/test_native_run_execution.cpp