Fix: reserve NEXT_LEVEL workers for ready groups - #1565
Conversation
📝 WalkthroughWalkthroughNEXT_LEVEL group dispatch now reserves every target worker while the group FIFO head is blocked, preventing conflicting singles from dispatching until all targets are idle. Documentation and a staged-idle regression test describe and verify the updated behavior. ChangesNEXT_LEVEL group reservation
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Scheduler
participant GroupFIFO
participant WorkerThreads
Scheduler->>GroupFIFO: inspect group FIFO head
Scheduler->>WorkerThreads: reserve all target workers
WorkerThreads-->>Scheduler: report worker availability
Scheduler->>GroupFIFO: pop group when every target is idle
Scheduler->>WorkerThreads: dispatch group and clear reservations
Scheduler->>WorkerThreads: dispatch singles to remaining workers
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 |
Review — approve with nitsReviewed What I verified rather than assumed
Should fix1. "Unrelated worker queues progress" has zero test coverage. That claim is in the PR body and now in three docs, but 2.
No test covers that path either — zero tests in the file submit more than one group. Optionally, keeping the duplicate check on the local if (std::find(workers.begin(), workers.end(), worker) != workers.end()) {
throw std::runtime_error("Scheduler::dispatch_next_level_group: duplicate target worker");
}
reserved_next_level_worker_ids_.insert(worker_id);(The check is unreachable either way — Consider3. The 4. 5. Pre-existing, out of scope, flagging because this PR is why people will look at that loop: The one thing that needs a decisionThis PR adds a new invariant to
Documenting it is the right call — it is the precondition the fix now depends on. The problem is that the repo currently violates it in 10 places. Why the invariant is load-bearing
// Phase 2: device barrier — each rank notifies every peer that its
// stage-in is visible, then waits until every peer has notified us.
for (int peer = 0; peer < nranks; ++peer) { ... TNOTIFY(sig, 1, AtomicAdd); }
for (int peer = 0; peer < nranks; ++peer) { ... TWAIT(sig, 1, GE); }A collective rank does not return until every peer rank has been dispatched. So "all N ranks co-resident" is a hard precondition — which is the definition of the group API's all-or-nothing dispatch. Yet every collective in the tree is submitted as N independent How that deadlocks after this PRTwo workers, a 2-rank collective
Before this PR the same sequence made progress: at It is not reachable todayThe shape needs a group and a
So this PR cannot deadlock anything in-tree. I'm not asking to hold the merge on it. But three fairly natural next steps introduce it:
Why it would be nasty to debugThe scene at hang time: reserved workers sitting idle (resources free, nothing running), one worker at 100% AIV spinning in Proposed follow-upsA. Move collectives onto the group APICriterion: kernel contains
The edit is mechanical — accumulate in the loop, submit once after: args_list = []
for i in range(nranks):
chip_args = TaskArgs()
...
args_list.append(chip_args)
orch.submit_next_level_group(chip, args_list, config, workers=list(range(nranks)))Two things worth noting, both in favour:
This should be its own PR — 10 sites, touches the ST suite, needs onboard validation. Not for #1565. B. A stall watchdog on the reservationTime alone can't separate deadlock from a legitimately long op. The discriminator is structural: a reserved worker that is Sketch (host-side, so // Well below the 45s op-execute timeout, so the report always precedes the
// 507018 that would otherwise be the only symptom.
constexpr auto kReservationStallWarnAfter = std::chrono::seconds(5);
void Scheduler::check_reservation_stall(TaskSlot group_slot) {
const auto now = std::chrono::steady_clock::now();
if (reserved_next_level_worker_ids_ != reported_reservation_) { // set changed → progress
reported_reservation_ = reserved_next_level_worker_ids_;
reservation_since_ = now;
reservation_stall_reported_ = false;
return;
}
if (reservation_stall_reported_) return;
if (now - reservation_since_ < kReservationStallWarnAfter) return;
bool withholding = false; // an idle reserved target with queued singles
for (int32_t id : cfg_.ready_next_level_queues->worker_ids()) { ... }
if (!withholding) return;
reservation_stall_reported_ = true;
std::fprintf(stderr, "[simpler][scheduler] NEXT_LEVEL group slot=%d has held its target "
"reservation for >5s without launching. targets:%s\n"
" If a task on a BUSY target is waiting for a peer task queued on a reserved target, "
"this will never resolve. The usual cause is a collective submitted as N separate "
"submit_next_level() calls instead of one submit_next_level_group().\n", ...);
}Hook it on the blocked return in if (!all_workers_idle) {
check_reservation_stall(slot);
return;
}Needs a small Testable in the existing fixture by simply never calling I'd argue against making it self-healing (releasing the reservation on timeout). That contradicts #1551's "waits and retries rather than yielding its share back", brings the starvation bug back in intermittent form, and would permanently mask the submission bug that caused it. Better to hang, but hang legibly. Suggested split
Nice work on this one — the core change is 13 lines, the docs landed in the same commit across all three affected files, and documenting the new precondition instead of leaving it latent is what let the rest of this analysis happen at all. Cross-checked with |
09d77a0 to
92ff363
Compare
A blocked NEXT_LEVEL group reserves only the head group's targets from later singles until the complete set becomes idle. The reservation flows directly from the group pass to the singles pass, preserving all-or-nothing dispatch while unrelated queues continue. Scheduler regressions cover staged target completion, unrelated workers, and consecutive groups. Peer-waiting collectives submit each complete rank cohort through the group API so they satisfy the same-level placement contract. Fixes hw-native-sys#1551
|
@ChaoWao Thanks for the detailed review. I folded the changes that belong to
I kept the stall watchdog and the pre-existing scheduler condition-variable |
|
Approving — the scheduler change is sound, the collectives migration landed with Two notes on the update itself before the follow-ups:
Recording the follow-ups here so they don't get lost. None of them block this PR. P1 —
|
2 participants
A blocked NEXT_LEVEL group now reserves every target from later singles until
the complete set becomes idle. The reservation is round-local and passed
directly from the group pass to the singles pass, preserving all-or-nothing
group dispatch while unrelated worker queues progress.
Scheduler regressions cover targets draining one at a time, unrelated-worker
progress, and consecutive groups with overlapping targets.
In-tree collectives whose ranks wait for same-level peers now submit the
complete cohort through
submit_next_level_group(), matching the reservationcontract and preventing a running rank from waiting on a peer withheld behind
a ready group.
Fixes #1551