fix(scheduler): tag sync_start drain ack/election state by generation

[Leaf-Salix]

Summary

sync_start drain coordinates scheduler threads through a set of reusable
shared fields (sync_start_pending, drain_worker_elected,
drain_ack_mask, drain_stage_go, drain_stage_done_mask). None of these
fields carried a generation/attempt tag. When the elected owner found
global resources insufficient, it reset the ack/election fields back to
their sentinel values so all threads could retry — but a scheduler thread
that had already cleared the old attempt's ack barrier, and was merely
delayed before performing its own election CAS, could resume after the
reset and win the CAS for the new attempt using stale state. That stale
thread would then publish stage_go and wait on stage_done from peers
that were still accumulating acks for the new attempt, splitting the
barrier. Nothing in the protocol could detect or recover from this split,
so the run hung until the AICPU op-execute timeout fired
(RuntimeError: run failed with code 507018), matching the report in
#1455.

This is a textbook ABA problem: every individual field value is valid at
any instant, but the fields can be observed as belonging to two different
drain attempts by two different threads at the same time.

Fix

• Add a drain_attempt counter (atomic<uint64_t>) to SyncStartDrainState,
  bumped every time the ack/election state is reset — both when a brand
  new sync_start task enters drain and when the elected owner resets due
  to insufficient resources.
• Replace the single shared drain_ack_mask bitmask with a per-thread
  drain_ack_attempts_[] array: each thread stamps its ack with the
  attempt it belongs to. The "all acked" check
  (sync_start_drain_ack_mask_for_attempt) only counts acks whose stamped
  attempt matches the attempt currently in progress, so a late ack from an
  old attempt can never satisfy a newer attempt's barrier.
• Wrap the raw election CAS in sync_start_drain_try_elect: after winning
  the CAS on drain_worker_elected, it re-checks drain_attempt. If the
  attempt has moved on since the caller captured it, the function
  self-reverts the CAS (releases drain_worker_elected back to 0) and
  reports "not elected", so a stale winner never gets to act as owner for
  the wrong generation.
• Every ack-barrier / stage-go / reopen spin-wait in handle_drain_mode
  now also bails out as soon as it observes drain_attempt has changed,
  instead of only relying on is_completed() / drain_worker_elected == 0.

Applied identically to all three drain-protocol implementations that share
this pattern: a2a3/tensormap_and_ringbuffer, a2a3/host_build_graph, and
a5/tensormap_and_ringbuffer.

Deterministic reproduction

Ported the reproduction described in #1455 into an env-gated test hook
(SIMPLER_DRAIN_ABA_TEST=1, a2a3/tensormap_and_ringbuffer only): once per
drain run, scheduler thread 1 parks right after its own ack barrier but
before its election attempt, resumes only once a sibling's
insufficient-resource reset has advanced drain_attempt, and waits for the
other threads to have re-acked under the new attempt before falling
through — recreating the exact stale-owner interleaving from the issue.
Against the pre-fix protocol this reliably reproduces the split-barrier
hang; against the fix, sync_start_drain_try_elect rejects the stale
election attempt and the run completes normally.

New, self-contained scene test:
tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_drain_aba_hook
(one holder MIX task occupying a cluster, one sync_start MIX task needing
all of them — forces at least one insufficient-resource retry on the first
attempt, matching the issue's repro steps).

The hook also guards against silently proving nothing: if a full spin
budget passes without a sibling's reset ever advancing drain_attempt,
the race window was never exercised. Rather than falling through and
letting the scene test pass on the unraced happy path, the hook latches a
dedicated PTO2_ERROR_DRAIN_ABA_TEST_NOT_EXERCISED code and fails the run
— so a future change to the test's resource math that breaks the repro
fails loudly instead of quietly stopping to prove anything.

Documentation

Updated docs/dynamic-linking.md's description of
SchedulerContext::deinit()'s reset responsibilities, which still named
the now-removed drain_ack_mask field.

Tests

Unit tests (host-side, no hardware)

New SyncStartDrainAttemptTest suite (tests/ut/cpp/{a2a3,a5}/test_scheduler_state.cpp)
directly exercises the two new helpers:

• LateAckCannotSatisfyNextAttemptBarrier — an ack stamped with an old
  attempt is excluded from the current attempt's mask even after the
  current attempt's other threads have acked.
• StaleParticipantCannotOwnNextAttempt — a thread whose captured attempt
  no longer matches the live drain_attempt fails to elect, and does not
  leave drain_worker_elected in a won state.

Verified with a from-scratch clone of the pushed branch (independent of the
working checkout that produced it) and a fresh CMake configure/build:
60/60 no-hardware C++ unit tests pass (ctest -LE requires_hardware),
including the full pre-existing test_scheduler_state / test_a5_scheduler_state
/ test_trb_runtime_temp_buffer suites — no regressions.

• Toolchain: AppleClang 16.0.0 (clang-1600.0.26.6), CMake 4.4.0.
• Test binaries linked against GoogleTest via this repo's
  tests/ut/cpp/CMakeLists.txt.

Hardware A/B test (real a2a3 silicon, Ascend 910-series chip)

Ran the new spmd_sync_start_drain_aba_hook scene test with
SIMPLER_DRAIN_ABA_TEST=1 against two builds, each installed from a clean
editable install with the onboard toolchain on PATH, with the onboard
build cache cleared before each install and verified by strings-checking
the compiled libaicpu_kernel.so for the expected symbols before running
(to rule out a stale/cached binary from an unrelated install):

Build	Result
Pre-fix drain protocol (upstream/main at d0bc661a, unmodified) + the same ABA-forcing hook	FAILED — RuntimeError: run failed with code 507018 after ~62s (most recent run: 61.90s), matching #1455
This PR's fixed drain protocol + the same hook	PASSED in ~7s, repeated over 5 consecutive runs (including after the ABA-hook hardening and test-directory cleanup commits) with no failures

Re-confirmed both sides freshly (clean onboard rebuild + rerun) after the
final round of changes to this branch, with upstream/main re-checked via
git ls-remote immediately beforehand to confirm the baseline commit
(d0bc661a) had not moved since this branch was created — the failing
and passing results above are current, not carried over from an earlier,
possibly stale build.

Environment: --platform=a2a3 on real device (not the a2a3sim simulator),
single-device execution, CANN toolkit cann-9.0.0, Python 3.10.9,
pytest 9.0.3 (plugins: pytest-forked 1.6.0, anyio 4.12.1, pytest-xdist
3.8.0), run with -s -v --forked.

Base commit for this branch: upstream/main at d0bc661a (hw-native-sys/simpler).

Related Issue

Fixes hw-native-sys/simpler#1455

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 42cf8c26-94e5-4103-8891-eccefbc6874e

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

The sync-start drain protocol now uses attempt-scoped acknowledgements and elections to reject stale retry state. Runtime wiring exposes a deterministic ABA test mode, with lifecycle resets, unit coverage, and an SPMD reproducer.

Changes

Sync-start drain protocol

Layer / File(s)	Summary
Attempt-scoped drain state src/*/scheduler/scheduler_types.h	drain_attempt and per-thread acknowledgement tracking replace the shared acknowledgement mask; helper functions validate barriers and elections by attempt.
Attempt-aware scheduler flow src/*/scheduler/scheduler_completion.cpp	Drain retries, waits, stale-election handling, resource fallback, and finalization now terminate or advance when the attempt changes.
Lifecycle and ABA test wiring src/*/scheduler/scheduler_cold_path.cpp, src/a2a3/runtime/tensormap_and_ringbuffer/{host,runtime}/*, docs/dynamic-linking.md	Scheduler state is reset for attempts and acknowledgements, while SIMPLER_DRAIN_ABA_TEST is propagated into the device launch descriptor.
Unit and SPMD validation tests/ut/cpp/*/test_scheduler_state.cpp, tests/st/a2a3/.../spmd_sync_start_drain_aba_hook/*	Tests cover late acknowledgements and stale elections; the new orchestration and pressure kernels reproduce the interleaving.

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

Possibly related PRs

• hw-native-sys/simpler#1423: Both changes revise a5 sync-start drain coordination and election/reset behavior.

Sequence Diagram(s)

sequenceDiagram
  participant SchedulerThreads
  participant DrainState
  participant Runtime
  SchedulerThreads->>DrainState: enter drain and advance drain_attempt
  SchedulerThreads->>DrainState: publish per-thread acknowledgement
  SchedulerThreads->>DrainState: elect matching-attempt worker
  DrainState-->>SchedulerThreads: accept attempt or invalidate cohort
  SchedulerThreads->>Runtime: finalize or retry sync-start drain

Loading

Poem

A bunny hops through attempts anew,
Old ACKs fade; fresh rounds come through.
Elections guard the drain-door tight,
ABA shadows vanish from sight.
Kernels spin and tests declare:
The scheduler’s safe in its carrot lair!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 33.33% 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
Title check	✅ Passed	The title concisely describes the main change: generation-tagging sync_start drain ack/election state.
Description check	✅ Passed	The description matches the diff and explains the ABA fix, test hook, documentation, and validation.
Linked Issues check	✅ Passed	The changes implement the #1455 requirements: attempt tagging, stale-thread rejection, and deterministic ABA validation.
Out of Scope Changes check	✅ Passed	No clearly unrelated code changes appear; the added tests, docs, and helper/runtime support all serve the ABA fix.

---

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

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 10436722dc

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[coderabbitai]

Actionable comments posted: 4

🤖 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/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp`:
- Around line 472-487: In SchedulerContext::enter_drain_mode, advance and
publish drain_attempt before clearing drain_worker_elected, matching the
ordering used by handle_drain_mode and the shared parent implementation. Keep
pending_task publication and the remaining coordination resets intact, but
ensure stale-round election checks cannot observe the new task before the
attempt changes.

In
`@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp`:
- Around line 503-521: In SchedulerContext::enter_drain_mode, advance and
publish drain_attempt before clearing drain_worker_elected. Preserve the
existing reset sequence and pending_task publication, but ensure the new
generation is visible before election state is reset, matching the ordering
already used by handle_drain_mode.

In
`@src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp`:
- Around line 416-430: Reorder the initialization in
SchedulerContext::enter_drain_mode so drain_attempt is incremented and published
before clearing drain_worker_elected, matching the ordering used by
handle_drain_mode’s insufficient-resource path and the other scheduler variants.
Preserve the existing drain-state initialization and ownership logic.

In
`@tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_drain_aba_hook/test_spmd_sync_start_drain_aba_hook.py`:
- Line 19: Update the test setup around SIMPLER_DRAIN_ABA_TEST to assign "1"
unconditionally, ensuring the ABA hook is exercised regardless of the inherited
environment. Scope the override to this test and restore the prior environment
value during teardown, including removing the variable when it was previously
unset.

🪄 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: 88fb1af7-bff2-4d96-a640-1199859fa07a

📥 Commits

Reviewing files that changed from the base of the PR and between d0bc661 and 1043672.

📒 Files selected for processing (23)

• docs/dynamic-linking.md
• src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp
• src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp
• src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h
• src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_types.h
• src/a2a3/runtime/tensormap_and_ringbuffer/common/pto_runtime_status.h
• src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp
• src/a2a3/runtime/tensormap_and_ringbuffer/runtime/runtime.h
• src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp
• src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp
• src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h
• src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h
• src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/runtime.cpp
• src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp
• src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp
• src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h
• src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h
• tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_drain_aba_hook/kernels/aic/kernel_spmd_mix_pressure.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_drain_aba_hook/kernels/aiv/kernel_spmd_mix_pressure.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_drain_aba_hook/kernels/orchestration/spmd_sync_start_drain_aba_hook_orch.cpp
• tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_drain_aba_hook/test_spmd_sync_start_drain_aba_hook.py
• tests/ut/cpp/a2a3/test_scheduler_state.cpp
• tests/ut/cpp/a5/test_scheduler_state.cpp