fix(scheduler): make sync_start drain retries generation-safe - #1461
Conversation
📝 WalkthroughWalkthroughThe 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. ChangesSync-start drain protocol
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
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
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 |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 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[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
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
📒 Files selected for processing (23)
docs/dynamic-linking.mdsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cppsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cppsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.hsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_types.hsrc/a2a3/runtime/tensormap_and_ringbuffer/common/pto_runtime_status.hsrc/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/runtime.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/runtime.cppsrc/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cppsrc/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cppsrc/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.hsrc/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.htests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_drain_aba_hook/kernels/aic/kernel_spmd_mix_pressure.cpptests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_drain_aba_hook/kernels/aiv/kernel_spmd_mix_pressure.cpptests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_drain_aba_hook/kernels/orchestration/spmd_sync_start_drain_aba_hook_orch.cpptests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_drain_aba_hook/test_spmd_sync_start_drain_aba_hook.pytests/ut/cpp/a2a3/test_scheduler_state.cpptests/ut/cpp/a5/test_scheduler_state.cpp
51dd422 to
e7d4362
Compare
ChaoZheng109
left a comment
ChaoZheng109
left a comment
There was a problem hiding this comment.
Reviewed the generation-tagging approach for the #1455 ABA — the fix looks correct and the tri-part rollout (tmr / host_build_graph / a5) is consistent. A few redundancy / simplification notes below; none affect correctness. Separately, I'm running a hardware benchmark covering both plain cases and sync_start cases to confirm no serious regression, and will follow up with the numbers.
| // insufficient resources). | ||
| while (drain_state_.drain_stage_go.load(std::memory_order_acquire) == 0) { | ||
| if (is_completed()) return; | ||
| if (drain_state_.sync_start_pending.load(std::memory_order_acquire) == 0) return; |
There was a problem hiding this comment.
Likely-dead defensive guard. A non-elected thread only spins here while drain_stage_go == 0, i.e. before the owner has released staging. The owner sets sync_start_pending = 0 only in finalize, which is gated on drain_stage_done_mask being complete — but a thread parked in this loop hasn't set its own stage_done bit yet, so the owner cannot have finalized while this thread is here. The only other writer that zeroes sync_start_pending is teardown, already covered by is_completed(). So this condition looks unreachable inside this loop. Suggest removing it or documenting the intended case. Same in a5 :582 and host_build_graph :737.
There was a problem hiding this comment.
Agreed. I removed the sync_start_pending == 0 escape condition from all three scheduler variants.
A valid owner cannot finalize the task before every participating thread has completed the staging barrier. Therefore a follower waiting for stage_go cannot observe the pending task being cleared for the same attempt.
Task teardown is handled separately by the existing completion path, so this guard was both unreachable for a valid attempt and unnecessary for teardown.
There was a problem hiding this comment.
Follow-up for the current version,
34a250e:
The unreachable sync_start_pending == 0 escape remains removed in all three
scheduler variants. The follower's stage_go wait now exits only for runtime
completion or an attempt-generation change.
I also removed the follower's later pending_task reload/null guard. The fixed
coordinator cannot clear the pending task until every participant has staged
and published its stage_done bit, so neither guard is reachable for a valid
attempt; teardown remains covered by is_completed().
| }; | ||
| static_assert(sizeof(SyncStartDrainState) == 64); | ||
|
|
||
| inline uint32_t sync_start_drain_ack_mask_for_attempt( |
There was a problem hiding this comment.
Perf note (not redundant). The old drain_ack_mask gave the ack barrier an O(1) atomic load per spin iteration; sync_start_drain_ack_mask_for_attempt now performs thread_count acquire-loads every iteration. N is small (<=4) so the absolute cost is tiny, but this is on the drain spin path, so it's worth confirming the sync_start benchmark cases don't regress. Optional micro-opt: only rescan the per-thread array when the cheap single drain_attempt load still matches the current round.
There was a problem hiding this comment.
Thanks for pointing this out. I replaced the production O(N) acknowledgement scan rather than adding another full rescan after the attempt load.
The production synchronization now uses a generation-tagged binary-tree reduction:
- each worker waits for at most two child tokens;
- each worker publishes one subtree-ready token;
- the root publishes the generation-ready token;
- non-root workers wait for that root token.
This gives O(1) fan-in per worker and an O(log N) critical path, without shared read-modify-write contention.
The only remaining O(N) scan is an explicitly test-only helper used by the ABA test hook to inspect independently published arrival tokens. It is not reachable from the production scheduler path.
I am posting the updated current-base/current-head benchmark results in the PR conversation. Those results are mixed, so I am keeping the performance concern open.
There was a problem hiding this comment.
Follow-up for the current version,
34a250e:
The production O(N) acknowledgement scan remains removed. The final protocol
uses a generation-tagged binary-tree reduction:
- each scheduler waits for at most two child tokens;
- each scheduler publishes one subtree token;
- followers wait for the root token before entering the staging decision;
- thread 0 is the fixed coordinator, so there is no election CAS.
This gives O(1) fan-in per scheduler and an O(log N) critical path without a
shared RMW hotspot. The only remaining O(N) scan is explicitly test-only and
is used by the deterministic ABA hook.
I also tested alternatives before retaining the tree. A packed shared O(1)
atomic regressed normal/stress Device time by about 16%/19%, removing the root
broadcast regressed it by about 10%/8%, and a rotating coordinator showed no
stable benefit.
The latest stress-only follow-up used five independent
A100-B100-B100-A100 cycles, for 1,000 rounds per version. Its pooled PR delta
was about +4%, while the paired cycle deltas ranged from -3.25% to +18.90% and
the cycle-level interval crossed zero because the case is strongly multimodal.
The full result and corrected interpretation are recorded
here.
|
Re: my note above about benchmarking — this is better run on the author's side so the numbers come from a clean, controlled setup. @Leaf-Salix, could you run a hardware benchmark (a2a3, This PR changes the sync_start drain hot path — the ack barrier goes from a single O(1) Please make sure it covers both:
One build gotcha worth flagging: after switching to this branch, rebuilding with only |
03852b0 to
1ee0152
Compare
|
Follow-up to @ChaoZheng109's benchmark request. After rebasing, I discarded the earlier comparison against the old merge-base and reran the benchmark against the exact current commits:
Setup:
Positive deltas below mean the PR is slower:
The alternating-matmul and paged-attention Effective/Sched results are essentially flat. The stress sync-start case improves substantially. However, The production O(N) acknowledgement scan has been removed, but the current implementation still needs further optimization for the normal sync-start path. I am keeping this performance concern open while investigating that regression. |
456f793 to
34a250e
Compare
|
Follow-up to the performance review after updating the implementation. The branch is now rebased and squashed to one commit:
The hardware runs below used upstream What changedI removed the production O(N) acknowledgement scan and then tested several
The fixed-coordinator protocol was audited across a2a3 TMR, a2a3 The fixed coordinator does make drain finalization sensitive to temporary Validation
Final performance against current upstreamBoth trees were fully rebuilt in separate checkout-local venvs, including their Plain/non-drain cases (
Sync-start cases (
The earlier normal sync_start and batch-paged-attention regressions are no |
|
Follow-up to the performance comment above: I expanded Versions:
The new run used five independent Task: Pooled 1,000-round result
Paired results by cycle
Treating each cycle, rather than each round, as the correlated sampling unit, The round-level distribution is strongly multimodal:
This is not a small number of outliers, so 10% trimming does not remove the I therefore no longer interpret the earlier |
| const char *serial_env = std::getenv("SIMPLER_TMR_SERIAL_ORCH_SCHED_ENABLE"); | ||
| runtime->dev.serial_orch_sched = | ||
| serial_env && (serial_env[0] == '1' || serial_env[0] == 't' || serial_env[0] == 'T'); | ||
| const char *drain_aba_env = std::getenv("SIMPLER_DRAIN_ABA_TEST"); |
There was a problem hiding this comment.
Remove the SIMPLER_DRAIN_ABA_TEST runtime hook entirely — the fix is already covered without it, and this hook spreads test-only code across production files.
The generation-tagging fix (drain_attempt + tree ack barrier + tri-state election) is self-contained and stands on its own. Its regression protection already exists without any runtime hook:
SyncStartDrainAttemptTest(tests/ut/cpp/{a2a3,a5}/test_scheduler_state.cpp) directly exercises the two failure modes —LateAckCannotSatisfyNextAttemptBarrierandStaleParticipantCannotOwnNextAttempt— against the helpers, no hardware/hook needed.- The normal + stress scene tests (
spmd_sync_start,spmd_sync_start_stress) cover the real concurrent drain path; stress run repeatedly is a probabilistic regression net.
Against that, the hook is a deterministic one-time proof whose permanent cost is out of proportion to its benefit:
- ~70 lines of test-only code smeared across 9 production files, including cross-cutting ones unrelated to the scheduler fix —
common/runtime_status/error_names.h,runtime/runtime.h,common/pto_runtime_status.h,runtime/shared/runtime.cpp,host/runtime_maker.cpp. - A new env-var behavior gate (
SIMPLER_DRAIN_ABA_TEST) — this is exactly the permanent configuration surface that.claude/rules/env-macro-gating.mdsays needs explicit maintainer sign-off and should be avoided; a stray=1in any environment deliberately stalls scheduler thread 1 on every drain. - A production error code
PTO2_ERROR_DRAIN_ABA_TEST_NOT_EXERCISED = 104occupying the shared status enum purely to serve the test. - Test-injection logic (thread-1 park loop + O(N) test scan + NOT_EXERCISED latch) inside the drain hot path
handle_drain_mode. - Asymmetric: the hook is a2a3 TMR only, while the fix lands in all three variants — so it doesn't even cover the code it guards uniformly.
Note that #ifdef-gating it would only keep it out of the shipped binary and off the hot path — the source lines and the cross-cutting file churn still remain in the tree and the diff. Since correctness is already validated, the cleanest outcome is to delete the hook outright rather than gate it. The one-time "pre-fix hangs / post-fix passes" evidence is valuable as a record — keep it in the PR description or a docs/investigations/ entry, not permanently in the code.
Concretely, please remove:
host/runtime_maker.cpp— theSIMPLER_DRAIN_ABA_TESTgetenv block +LOG_INFO_V0(this anchor).runtime/runtime.h—bool drain_aba_test_mode;runtime/shared/runtime.cpp—dev.drain_aba_test_mode = false;common/pto_runtime_status.h—PTO2_ERROR_DRAIN_ABA_TEST_NOT_EXERCISED.common/runtime_status/error_names.h— the threeDRAIN_ABA_TEST_NOT_EXERCISEDcases.runtime/scheduler/scheduler_context.h—drain_aba_test_mode_,drain_test_victim_armed_.runtime/scheduler/scheduler_cold_path.cpp— thedrain_aba_test_mode_assignment/reset.runtime/scheduler/scheduler_completion.cpp— thedrain_aba_test_mode_branches, the thread-1 park block, and the arrival-token publication used only by it.runtime/scheduler/scheduler_types.h— the test-onlysync_start_drain_test_ack_mask_for_attemptscan (keep the production token helpers).tests/st/a2a3/.../spmd_sync_start_drain_aba_hook/and thePTO2_ERROR_DRAIN_ABA_TEST_NOT_EXERCISEDentry intests/ut/cpp/common/test_error_code_names.cpp.
Keep the unit tests in test_scheduler_state.cpp — those are the right, hook-free regression barrier.
There was a problem hiding this comment.
Thanks, agreed. I removed the runtime hook entirely in 9429177.
The deletion covers all of the test-only production surface listed here:
- the
SIMPLER_DRAIN_ABA_TESTgetenv/log and launch-descriptor plumbing; drain_aba_test_mode_,drain_test_victim_armed_, the arrival publication, thread-1 park loop, andNOT_EXERCISEDlatch in the scheduler path;PTO2_ERROR_DRAIN_ABA_TEST_NOT_EXERCISEDand its name/description/hint mappings;- the test-only O(N) scan/arrival-token helpers from all three scheduler variants;
- the dedicated
spmd_sync_start_drain_aba_hookscene and the common error-code test entry.
Only the production attempt/subtree-token helpers remain. LateAckCannotSatisfyNextAttemptBarrier is retained for both A2/A3 and A5 and now exercises only those production helpers.
One detail changed after the review snapshot: StaleParticipantCannotOwnNextAttempt exercised the former tri-state election. The current implementation has removed drain_worker_elected, sync_start_drain_try_elect(), SyncStartDrainElectResult, and the election CAS from all three variants; scheduler thread 0 is now the fixed coordinator. That election-specific test therefore no longer has a production mechanism to exercise, so I did not reintroduce a test-only helper or a dummy test for it.
I preserved the one-time deterministic evidence in docs/investigations/2026-07-sync-start-drain-retry-aba.md, including the hook interleaving, pre-fix failure, protocol alternatives, and performance results.
Validation after removing the hook:
test_scheduler_state,test_a5_scheduler_state, andtest_error_code_namespassed;- A5sim normal/stress passed;
- A3
spmd_sync_start,spmd_sync_start_stress, and host-build-graph sync-start coverage passed throughtask-submit; - modified-file pre-commit, including clang-tidy and markdownlint, passed.
- Tag per-thread tree acknowledgements with the drain attempt - Use scheduler thread 0 as the fixed staging coordinator - Remove the deterministic runtime hook and retain hook-free regression coverage - Preserve the 64-byte shared-state layout and record the investigation
34a250e to
9429177
Compare
Leaf-Salix
changed the title
3 participants
Summary
Fixes #1455.
The
sync_startdrain protocol reused acknowledgement, election, and staging state across insufficient-resource retries without identifying which retry attempt produced each value. A delayed scheduler could therefore cross the old acknowledgement barrier, resume after another scheduler reset the state, and combine locally valid values from two attempts. The resulting split barrier could not make progress and eventually surfaced as AICPU timeout507018.Fix
drain_attemptand advance it whenever a new drain acknowledgement round starts.The protocol is applied consistently to:
tensormap_and_ringbuffer;host_build_graph;tensormap_and_ringbuffer.Regression coverage
The final tree contains no
SIMPLER_DRAIN_ABA_TESTruntime hook, test-only production error code, or test-only scan in the scheduler hot path.SyncStartDrainAttemptTest.LateAckCannotSatisfyNextAttemptBarriercovers generation isolation directly against the production token helpers for A2/A3 and A5.spmd_sync_startscene tests exercise the real concurrent drain path.docs/investigations/2026-07-sync-start-drain-retry-aba.md.Validation
Final pushed head:
9429177, based on upstream1c475a70.test_scheduler_state,test_a5_scheduler_state, andtest_error_code_namespassed.sync_startpassed.task-submitpassed:task_20260727_053540_328311513473;task_20260727_053615_329667622840;task_20260727_053659_33145449300.Performance
A locked-device 100-round ABBA comparison of the retained fixed-coordinator tree against upstream
a53cc168measured normalsync_starteffectively flat (Device -0.91%, Effective -0.83%) and stress about 3% faster (Device -2.93%, Effective -3.02%).A later five-cycle, 1,000-sample stress comparison against
1c475a70measured the PR about 4% slower in the pooled mean, but paired cycle deltas ranged from -3.25% to +18.90% and the cycle-level confidence interval crossed zero. The case was strongly multimodal, so neither the earlier improvement nor the later regression was established as a stable effect. Plain non-drain cases showed no consistent regression.Alternatives were also measured and rejected: a packed shared atomic regressed normal/stress Device time by about 16%/19%, removing the root broadcast regressed it by about 10%/8%, and rotating the coordinator produced no stable benefit. The generation-tagged tree with fixed coordinator therefore remains the best measured design.