Fix: block scheduler while group heads wait

docs/scheduler.md

@@ -40,11 +40,34 @@ NEXT_LEVEL queue after waiting in PENDING.
 Each `ReadyQueue` is a mutex-protected non-blocking FIFO, partitioned by run:
 a task is enqueued under its own `run_id`, and the Scheduler pops only from the
 partition of the run that currently holds the FIFO head. That is what keeps two
-admitted runs from interleaving their device work while both are live. Root
-submission, worker completion, and stop requests notify the Scheduler condition
-variable; its wait predicate checks the completion FIFO, the dispatchable run's
-partitions, and the stop flag. Ready queues have no blocking pop or shutdown
-state.
+admitted runs from interleaving their device work while both are live.
+
+Root submission, run activation, stop requests, a worker publishing itself
+idle, and group-member completions that do not enqueue a terminal task
+completion all advance a wake generation under the Scheduler
+condition-variable mutex. Terminal task completions push the completion FIFO
+under the same mutex. The wait predicate checks for a completion or an
+unconsumed wake generation.
+Blocked queue heads therefore remain queued without keeping the predicate
+permanently true. Ready queues have no blocking pop or shutdown state.
+
+The predicate is edge-triggered, and that is a contract on everything that
+feeds it: **any state change that turns already-queued work into placeable
+work must push a completion or advance the generation.** The predicate no
+longer re-reads queue occupancy or worker state, so a change that only mutates
+those is invisible to a parked Scheduler.
+
+Two producers are easy to miss because the change happens outside the
+Scheduler thread and outside `Orchestrator`. `dispatch_ready` re-queues work it
+could not place — through `enqueue_ready_cb`, which by design does not notify —
+so the retry depends entirely on a later edge. And a `WorkerThread` publishes
+its completion *before* it publishes its lane state: it calls `on_complete` and
+only then stores `active_inflight_`/`inflight_`. That order is deliberate, so a
+stopping Scheduler cannot read a worker as no longer busy while its final
+completion is still unqueued — which means the completion wake alone can land
+on a worker that still reads as occupied, and the dispatch pass it triggers
+finds nothing placeable. The `on_idle` callback fires after that publication
+and supplies the edge that makes the retry happen.
 
 Popping a slot is not the same as owning it. A run whose graph callback throws
 fails and consumes its own unstarted slots, so the Scheduler claims each slot
@@ -57,7 +80,7 @@ The Scheduler drains completions before dispatching new work:
 
 ```cpp
 while (true) {
-    wait_until_completion_ready_or_stop();
+    wait_until_completion_or_new_wake_generation();
 
     while (completion_queue has an item) {
         on_task_complete(item);

src/common/hierarchical/scheduler.cpp

@@ -103,13 +103,24 @@ void Scheduler::start(const Config &cfg) {
         throw std::invalid_argument("Scheduler::start: null config fields");
     cfg_ = cfg;
 
+    {
+        // run()'s observed generation restarts at zero, so any advance here
+        // arms the first round.
+        std::lock_guard<std::mutex> lk(completion_mu_);
+        ++wake_generation_;
+    }
+    dispatch_round_count_.store(0, std::memory_order_relaxed);
     stop_requested_.store(false, std::memory_order_relaxed);
     running_.store(true, std::memory_order_release);
     sched_thread_ = std::thread(&Scheduler::run, this);
 }
 
 void Scheduler::request_stop() {
     stop_requested_.store(true, std::memory_order_release);
+    {
+        std::lock_guard<std::mutex> lk(completion_mu_);
+        ++wake_generation_;
+    }
     completion_cv_.notify_all();
 }
 
@@ -136,7 +147,7 @@ void Scheduler::worker_done(WorkerCompletion completion) {
     if (s.is_group()) {
         WorkerCompletion terminal = completion;
         {
-            std::lock_guard<std::mutex> lk(s.group_mu);
+            std::unique_lock<std::mutex> lk(s.group_mu);
             const int32_t group_size = s.group_size();
             PreparedGroupVectors prepared =
                 prepare_group_vectors_locked(s, group_size, GroupMemberState::NOT_DISPATCHED);
@@ -152,7 +163,11 @@ void Scheduler::worker_done(WorkerCompletion completion) {
             int32_t index = invalid_group_index ? -1 : terminal.group_index;
             if (index >= 0 && index < group_size) {
                 GroupMemberState &member_state = s.group_member_states[static_cast<size_t>(index)];
-                if (is_terminal_group_state(member_state)) return;
+                if (is_terminal_group_state(member_state)) {
+                    lk.unlock();
+                    notify_ready();
+                    return;
+                }
 
                 if (terminal.outcome == EndpointOutcome::SUCCESS) {
                     member_state = GroupMemberState::SUCCESS;
@@ -185,7 +200,11 @@ void Scheduler::worker_done(WorkerCompletion completion) {
                 }
             }
 
-            if (s.group_terminal_count.load(std::memory_order_acquire) < group_size) return;
+            if (s.group_terminal_count.load(std::memory_order_acquire) < group_size) {
+                lk.unlock();
+                notify_ready();
+                return;
+            }
 
             if (s.group_failed) {
                 int32_t failure_index = s.group_first_failure_index;
@@ -216,50 +235,26 @@ void Scheduler::worker_done(WorkerCompletion completion) {
 }
 
 void Scheduler::notify_ready() {
-    std::lock_guard<std::mutex> lk(completion_mu_);
-    completion_cv_.notify_one();
-}
-
-bool stageable_successor_ready(const NextLevelReadyQueues &ready_queues, const WorkerManager &manager, RunId run_id) {
-    // B3b stages singles only. If the successor has a group head, retain the
-    // established all-or-nothing group priority instead of letting a staged
-    // single occupy one of its reserved workers after FIFO promotion.
-    if (!ready_queues.groups_empty(run_id)) return false;
-    for (int32_t worker_id : ready_queues.worker_ids()) {
-        WorkerThread *worker = manager.get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id);
-        if (worker != nullptr && worker->can_stage() && !ready_queues.single_empty(worker_id, run_id)) {
-            return true;
-        }
+    {
+        std::lock_guard<std::mutex> lk(completion_mu_);
+        ++wake_generation_;
     }
-    return false;
+    completion_cv_.notify_one();
 }
 
 // =============================================================================
 // Scheduler loop
 // =============================================================================
 
 void Scheduler::run() {
+    uint64_t observed_wake_generation = 0;
     while (true) {
-        // Wait until there's something to process
         {
             std::unique_lock<std::mutex> lk(completion_mu_);
-            completion_cv_.wait(lk, [this] {
-                bool ready = false;
-                if (cfg_.active_run_cb) {
-                    RunId active = cfg_.active_run_cb();
-                    ready = active != INVALID_RUN_ID &&
-                            (!cfg_.ready_next_level_queues->empty(active) || !cfg_.ready_sub_queue->empty(active) ||
-                             cfg_.manager->needs_activation(active));
-                    if (!ready && cfg_.preparable_run_cb) {
-                        RunId preparable = cfg_.preparable_run_cb();
-                        ready = preparable != INVALID_RUN_ID && !cfg_.manager->has_staged_run(preparable) &&
-                                stageable_successor_ready(*cfg_.ready_next_level_queues, *cfg_.manager, preparable);
-                    }
-                } else {
-                    ready = !cfg_.ready_next_level_queues->empty() || !cfg_.ready_sub_queue->empty();
-                }
-                return !completion_queue_.empty() || ready || stop_requested_.load(std::memory_order_acquire);
+            completion_cv_.wait(lk, [this, &observed_wake_generation] {
+                return !completion_queue_.empty() || wake_generation_ != observed_wake_generation;
             });
+            observed_wake_generation = wake_generation_;
         }
 
         // Hold loop_mu_ across the entire slot-touching body so quiescent
@@ -343,7 +338,6 @@ void Scheduler::on_task_complete(const WorkerCompletion &completion) {
         // lost: submit's publication compares the pair under the same lock.
         if (try_mark_ready(cs)) {
             cfg_.enqueue_ready_cb(consumer);
-            completion_cv_.notify_one();
         }
     }
 
@@ -417,6 +411,7 @@ void Scheduler::try_consume(TaskSlot slot) {
 // sched_thread_ with no surrounding handler, any throw is fatal to the whole
 // worker tree (std::terminate), not a per-task failure.
 void Scheduler::dispatch_ready() {
+    dispatch_round_count_.fetch_add(1, std::memory_order_relaxed);
     std::optional<RunId> run_snapshot;
     if (cfg_.active_run_cb) {
         RunId active_run = cfg_.active_run_cb();

src/common/hierarchical/scheduler.h

@@ -22,13 +22,21 @@
  *   Orch: submit() → directed NEXT_LEVEL queue or shared SUB queue + notify
  *
  *   Scheduler thread:
- *     wait on cv (ready queue OR completion queue OR stop requested)
+ *     wait on cv (completion queue OR unconsumed wake generation)
  *     drain completion_queue → on_task_complete → fanout release → ready_queue
  *     launch directed NEXT_LEVEL tasks, then freely scheduled SUB tasks
  *
  *   WorkerThread (managed by WorkerManager):
  *     loop: task_queue.pop() → endpoint.run(dispatch) →
  *           completion callback → Scheduler.worker_done(completion)
+ *           → lane state published → idle callback → Scheduler.notify_ready()
+ *
+ * The wait is edge-triggered, so it carries an obligation on everything that
+ * feeds it: any state change that turns already-queued work into placeable
+ * work must push a completion or advance the wake generation. Queue occupancy
+ * is no longer re-read by the predicate, so a change that only mutates it —
+ * a requeue from dispatch, a worker publishing itself idle, a run reaching the
+ * FIFO head — is invisible until someone posts the matching edge.
  */
 
 #pragma once
@@ -69,7 +77,6 @@ struct WorkerDispatch;
  * window this closes is exactly the code between the queue pop and the launch.
  */
 bool claim_for_dispatch(TaskSlotState &s);
-bool stageable_successor_ready(const NextLevelReadyQueues &ready_queues, const WorkerManager &manager, RunId run_id);
 
 class Scheduler {
 public:
@@ -103,6 +110,9 @@ class Scheduler {
     void stop();
 
     bool running() const { return running_.load(std::memory_order_acquire); }
+    // Diagnostic only — counts dispatch passes so an observer can tell a parked
+    // scheduler from one spinning on unplaceable work. Orders nothing.
+    uint64_t dispatch_round_count() const { return dispatch_round_count_.load(std::memory_order_relaxed); }
 
     // Called by WorkerManager (from WorkerThread) after endpoint run() reaches
     // a terminal outcome.
@@ -126,10 +136,12 @@ class Scheduler {
     std::queue<WorkerCompletion> completion_queue_;
     std::mutex completion_mu_;
     std::condition_variable completion_cv_;
+    uint64_t wake_generation_{0};
 
     std::thread sched_thread_;
     std::atomic<bool> stop_requested_{false};
     std::atomic<bool> running_{false};
+    std::atomic<uint64_t> dispatch_round_count_{0};
 
     void run();
     void on_task_complete(const WorkerCompletion &completion);

src/common/hierarchical/types.cpp

@@ -273,22 +273,6 @@ bool NextLevelReadyQueues::try_front_group(RunId run_id, TaskSlot &out) { return
 bool NextLevelReadyQueues::try_pop_group(TaskSlot &out) { return group_queue_.try_pop(out); }
 bool NextLevelReadyQueues::try_pop_group(RunId run_id, TaskSlot &out) { return group_queue_.try_pop(run_id, out); }
 
-bool NextLevelReadyQueues::empty() const {
-    if (!group_queue_.empty()) return false;
-    for (const auto &queue : queues_) {
-        if (!queue->empty()) return false;
-    }
-    return true;
-}
-
-bool NextLevelReadyQueues::empty(RunId run_id) const {
-    if (!group_queue_.empty(run_id)) return false;
-    for (const auto &queue : queues_) {
-        if (!queue->empty(run_id)) return false;
-    }
-    return true;
-}
-
 bool NextLevelReadyQueues::groups_empty(RunId run_id) const { return group_queue_.empty(run_id); }
 
 bool NextLevelReadyQueues::single_empty(int32_t worker_id, RunId run_id) const {

src/common/hierarchical/types.h

@@ -543,8 +543,6 @@ class NextLevelReadyQueues {
     bool try_front_group(RunId run_id, TaskSlot &out);
     bool try_pop_group(TaskSlot &out);
     bool try_pop_group(RunId run_id, TaskSlot &out);
-    bool empty() const;
-    bool empty(RunId run_id) const;
     bool groups_empty(RunId run_id) const;
     bool single_empty(int32_t worker_id, RunId run_id) const;
     bool singles_empty(RunId run_id) const;

src/common/hierarchical/worker.cpp

@@ -106,6 +106,9 @@ void Worker::init() {
         },
         [this](WorkerDispatch dispatch) {
             orchestrator_.mark_task_accepted(dispatch.task_slot);
+        },
+        [this] {
+            scheduler_.notify_ready();
         }
     );
     ready_next_level_queues_.reset(manager_.next_level_worker_ids());

src/common/hierarchical/worker_manager.cpp

@@ -269,12 +269,14 @@ char *LocalMailboxEndpoint::task_frame(size_t index) const {
 
 void WorkerThread::start(
     Ring *ring, const std::function<void(WorkerCompletion)> &on_complete,
-    const std::function<void(WorkerDispatch)> &on_accept, std::unique_ptr<WorkerEndpoint> endpoint
+    const std::function<void(WorkerDispatch)> &on_accept, const std::function<void()> &on_idle,
+    std::unique_ptr<WorkerEndpoint> endpoint
 ) {
     if (!endpoint) throw std::invalid_argument("WorkerThread::start: null endpoint");
     ring_ = ring;
     on_complete_ = on_complete;
     on_accept_ = on_accept;
+    on_idle_ = on_idle;
     endpoint_ = std::move(endpoint);
     shutdown_ = false;
     if (endpoint_->caps().max_inflight_tasks == 0) {
@@ -393,12 +395,6 @@ bool WorkerThread::has_staged_run(RunId run_id) const {
     return staged_run_id_.load(std::memory_order_relaxed) == run_id;
 }
 
-bool WorkerThread::needs_activation(RunId run_id) const {
-    std::lock_guard<std::mutex> lane_lk(lane_mu_);
-    return staged_run_id_.load(std::memory_order_relaxed) == run_id &&
-           activated_run_id_.load(std::memory_order_relaxed) != run_id;
-}
-
 bool WorkerThread::can_stage() const {
     std::lock_guard<std::mutex> lane_lk(lane_mu_);
     return caps().supports_frame_staging && staged_run_id_.load(std::memory_order_relaxed) == INVALID_RUN_ID;
@@ -578,10 +574,16 @@ void WorkerThread::loop() {
             }
         }
 
+        // on_complete_ runs before the lane state is published so a stopping
+        // scheduler cannot read this worker as no longer busy while its final
+        // completion is still unqueued. That leaves the reverse window — the
+        // scheduler placing work sees a stale non-idle lane — which is what
+        // on_idle_ closes.
         on_complete_(std::move(completion));
         active_inflight_.store(false, std::memory_order_release);
         inflight_.fetch_sub(1, std::memory_order_acq_rel);
         cv_.notify_one();
+        if (on_idle_) on_idle_();
     }
 }
 
@@ -639,6 +641,7 @@ void WorkerThread::finish_progress_dispatch(const WorkerEndpointProgress &progre
     }
     inflight_.fetch_sub(1, std::memory_order_acq_rel);
     cv_.notify_one();
+    if (on_idle_) on_idle_();
 }
 
 void WorkerThread::fail_progress_driver(const std::string &reason) noexcept {
@@ -1095,7 +1098,9 @@ void WorkerManager::add_next_level_endpoint(std::unique_ptr<WorkerEndpoint> endp
 
 void WorkerManager::add_sub(void *mailbox, int child_pid) { sub_entries_.push_back(LocalSubEntry{mailbox, child_pid}); }
 
-void WorkerManager::start(Ring *ring, const OnCompleteFn &on_complete, const OnAcceptFn &on_accept) {
+void WorkerManager::start(
+    Ring *ring, const OnCompleteFn &on_complete, const OnAcceptFn &on_accept, const OnIdleFn &on_idle
+) {
     if (ring == nullptr) throw std::invalid_argument("WorkerManager::start: null ring");
 
     std::vector<int32_t> next_level_worker_ids;
@@ -1125,7 +1130,7 @@ void WorkerManager::start(Ring *ring, const OnCompleteFn &on_complete, const OnA
             auto endpoint = std::make_unique<LocalMailboxEndpoint>(
                 entry.worker_id, entry.mailbox, entry.child_pid, entry.task_frame_count
             );
-            wt->start(ring, on_complete, on_accept, std::move(endpoint));
+            wt->start(ring, on_complete, on_accept, on_idle, std::move(endpoint));
             next_level_threads_.push_back(std::move(wt));
         }
     };
@@ -1136,14 +1141,14 @@ void WorkerManager::start(Ring *ring, const OnCompleteFn &on_complete, const OnA
             auto endpoint = std::make_unique<LocalMailboxEndpoint>(
                 static_cast<int32_t>(i), entries[i].mailbox, entries[i].child_pid
             );
-            wt->start(ring, on_complete, on_accept, std::move(endpoint));
+            wt->start(ring, on_complete, on_accept, on_idle, std::move(endpoint));
             threads.push_back(std::move(wt));
         }
     };
     make_next_level_threads();
     for (auto &endpoint : next_level_endpoint_entries_) {
         auto wt = std::make_unique<WorkerThread>();
-        wt->start(ring, on_complete, on_accept, std::move(endpoint));
+        wt->start(ring, on_complete, on_accept, on_idle, std::move(endpoint));
         next_level_threads_.push_back(std::move(wt));
     }
     next_level_endpoint_entries_.clear();
@@ -1636,12 +1641,6 @@ bool WorkerManager::has_staged_run(RunId run_id) const {
     return false;
 }
 
-bool WorkerManager::needs_activation(RunId run_id) const {
-    for (const auto &worker : next_level_threads_)
-        if (worker->needs_activation(run_id)) return true;
-    return false;
-}
-
 bool WorkerManager::activate_prepared_run(RunId run_id) {
     bool activated = false;
     for (const auto &worker : next_level_threads_)