Repository navigation
Track external wakeups in DebugScheduler - #123
Conversation
DebugScheduler had three defects when tasks are woken from outside it, for example when a CUDA host callback returns a pooled stream to a task waiting in acquire_stream: - The runnable queue had no lock, while the callback thread enqueued the rescheduled task concurrently with the driving thread. - help_while_waiting() panicked with DEADLOCK whenever the queue was empty, even though a wakeup was pending. - Tests that called run_all() and then assumed completion raced the callback. CudaTensorTest.StreamOrdered* aborted about once in 150 repetitions on main, destroying a still-suspended task. Add a balanced external-wait registration to the scheduler interface. An awaiter that hands its task to another party calls BasicTask::begin_external_wait() before publishing it and BasicTask::reschedule_after_external_wait() to wake it; the latter reschedules before ending the wait, so the scheduler never sees an empty queue with nothing pending while the task is in transit. The CUDA resource, stream and cuBLAS execution awaiters register their waits. Schedulers other than DebugScheduler treat registration as a no-op. DebugScheduler guards its queue and the wait count with a mutex and condition variable. Waits block, without a time limit, while the queue is empty and an external wakeup is pending, and report DEADLOCK only when nothing is runnable and nothing is pending. run_all() keeps its meaning as a debugging aid that runs until nothing is runnable; it does not wait for externally suspended tasks, so a caller that holds the resource a task needs can still drive the scheduler. The stream-ordered tensor tests now wait for their task instead of relying on run_all(). New tests cover a host-thread wakeup through run_all() and wait_for(), and a deterministic CUDA case in which the only stream is held by a host-function gate until the task is suspended on it. Also move test_runtime_floating_eq.cpp into the numerical-precision test module. It uses the numerics PrecisionContext registry, which includes the tensor layer, so uni20_common_tests failed to build in CUDA configurations for lack of the CUDA include path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Codex Review Summary: ❌ FailedThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for tracking and waiting on external wakeups (such as CUDA host callbacks) within the DebugScheduler and BasicTask async framework, updating the CUDA/cuBLAS awaiters and adding corresponding tests. The review feedback highlights critical race conditions where wakeup_.notify_all() is called outside of lock blocks, and where tasks can run to completion and destroy the scheduler before the external wait is fully released. To resolve these issues atomically, it is recommended to introduce a virtual reschedule_after_external_wait function to IScheduler and implement it in DebugScheduler to perform both rescheduling and wait-decrementing under a single lock acquisition.
| void end_external_wait() noexcept override | ||
| { | ||
| { | ||
| std::lock_guard lock(mutex_); | ||
| CHECK(external_waits_ > 0, "DebugScheduler external wait ended without a matching begin"); | ||
| --external_waits_; | ||
| } | ||
| wakeup_.notify_all(); | ||
| } |
There was a problem hiding this comment.
Calling wakeup_.notify_all() outside the lock block in end_external_wait creates a critical race condition. If the waiting thread is woken up, it can acquire the lock, see that external_waits_ == 0, return from wait_for_runnable_or_no_external_waits, and proceed to destroy the DebugScheduler object before notify_all() completes, leading to undefined behavior (use-after-free of the condition variable).
Additionally, calling BasicTask::reschedule before end_external_wait in reschedule_after_external_wait allows the task to run to completion and destroy the scheduler before the wait is released. To solve both issues atomically, we should override reschedule_after_external_wait in DebugScheduler to perform both operations under a single lock acquisition, and move notify_all() inside the lock block.
void end_external_wait() noexcept override
{
std::lock_guard lock(mutex_);
CHECK(external_waits_ > 0, "DebugScheduler external wait ended without a matching begin");
--external_waits_;
wakeup_.notify_all();
}
void reschedule_after_external_wait(BasicTask&& task) override
{
TRACE_MODULE(ASYNC, "Rescheduling a task after external wait", &task, task.coroutine_handle());
{
std::lock_guard lock(mutex_);
Handles_.push_back(std::move(task));
CHECK(external_waits_ > 0, "DebugScheduler external wait ended without a matching begin");
--external_waits_;
}
wakeup_.notify_all();
}There was a problem hiding this comment.
Fixed in 47d6191: DebugScheduler::complete_external_wait requeues, decrements and notifies in one critical section, so once the task can run the waking thread only releases the lock. Destroying a scheduler while an external party may still call into it, beyond that, is a lifetime-contract violation rather than a scheduler defect.
| void enqueue_task(BasicTask&& task) | ||
| { | ||
| { | ||
| std::lock_guard lock(mutex_); | ||
| Handles_.push_back(std::move(task)); | ||
| } | ||
| wakeup_.notify_all(); | ||
| } |
There was a problem hiding this comment.
Calling wakeup_.notify_all() outside the lock block in enqueue_task creates a similar race condition where the waiting thread can wake up, run the task to completion, and destroy the scheduler before notify_all() completes. Moving notify_all() inside the lock block ensures safe and robust notification.
void enqueue_task(BasicTask&& task)
{
std::lock_guard lock(mutex_);
Handles_.push_back(std::move(task));
wakeup_.notify_all();
}There was a problem hiding this comment.
notify_all() is now inside the lock here too (47d6191). Only external wakeups enqueue from other threads, and they now go through complete_external_wait, so this was for consistency rather than a reachable defect.
| inline void BasicTask::reschedule_after_external_wait(BasicTask task) | ||
| { | ||
| DEBUG_CHECK(task.handle_); | ||
| auto* scheduler = task.handle_.promise().scheduler(); | ||
| DEBUG_CHECK(scheduler, "unexpected: task scheduler is not set!"); | ||
| BasicTask::reschedule(std::move(task)); | ||
| scheduler->end_external_wait(); | ||
| } |
There was a problem hiding this comment.
Calling BasicTask::reschedule before scheduler->end_external_wait() introduces a critical race condition where the task can run to completion and destroy the scheduler before end_external_wait() is executed. To resolve this atomically, we should delegate to a single reschedule_after_external_wait call on the scheduler.
inline void BasicTask::reschedule_after_external_wait(BasicTask task)
{
DEBUG_CHECK(task.handle_);
auto* scheduler = task.handle_.promise().scheduler();
DEBUG_CHECK(scheduler, "unexpected: task scheduler is not set!");
task = BasicTask::make_sole_owner(std::move(task));
if (task)
{
scheduler->reschedule_after_external_wait(std::move(task));
}
else
{
scheduler->end_external_wait();
}
}There was a problem hiding this comment.
Partly agreed; fixed in 47d6191. The general premise that a scheduler may be destroyed while work is outstanding is outside the contract: schedulers do not track their tasks and must outlive them. But rescheduling before ending the wait did leave a real window. The woken task could finish, the program could correctly conclude that all its tasks were done and destroy the scheduler, and the waking thread would still call into it. end_external_wait() is replaced by IScheduler::complete_external_wait(BasicTask&&), which reschedules and releases the wait as one step, and BasicTask::reschedule_after_external_wait makes that single call.
| /// \brief Record that an external wakeup recorded by `begin_external_wait()` was delivered. | ||
| /// \details Called after the woken task, if any, has been rescheduled. | ||
| virtual void end_external_wait() noexcept {} | ||
| }; |
There was a problem hiding this comment.
Add a virtual reschedule_after_external_wait function to IScheduler to allow schedulers to atomically reschedule a task and end its external wait under a single lock acquisition, preventing race conditions during destruction.
/// \brief Record that an external wakeup recorded by begin_external_wait() was delivered.
/// \brief Called after the woken task, if any, has been rescheduled.
virtual void end_external_wait() noexcept {}
/// \brief Reschedule a task woken from outside and end its external wait atomically.
virtual void reschedule_after_external_wait(BasicTask&& task)
{
this->reschedule(std::move(task));
this->end_external_wait();
}
};There was a problem hiding this comment.
Adopted in 47d6191 as complete_external_wait(BasicTask&&), replacing end_external_wait() rather than adding a second entry point. The default just reschedules; DebugScheduler overrides it.
|
Review Bot: review of exact head Found 1 issue; details below. What the PR does. It adds a balanced external-wait registration to the scheduler interface: Finding. What I reviewed. The full diff. Verified (polaron, GCC 13.3, Debug,
Not checked.
Review tier: full (code or new pull request) · model: claude-opus-5-5 · effort: high |
| auto* scheduler = task.handle_.promise().scheduler(); | ||
| DEBUG_CHECK(scheduler, "unexpected: task scheduler is not set!"); | ||
| BasicTask::reschedule(std::move(task)); | ||
| scheduler->end_external_wait(); |
There was a problem hiding this comment.
Review Bot: review of exact head b08ab3b72de7d08d31e04b9f475e03645704c9cf
The woken task can finish, and its scheduler can be destroyed, before this end_external_wait() runs. Line 36 enqueues the task and releases DebugScheduler::mutex_. From then on the driving thread can run the task to completion, see its wait_for predicate become true, and destroy the scheduler. All of that can happen while the waking thread has still not reached line 37. Two things follow:
- the destructor
CHECK(external_waits_ == 0, ...)(debug_scheduler.hpp:83) fires even though the task has completed; or the waking thread decrements and notifies through a scheduler that is already destroyed; external_waits()can still read 1 after the awaited result is ready. The new CUDA test asserts exactly this with no synchronisation (test_cuda_tensor.cpp:570,EXPECT_EQ(scheduler.external_waits(), 0U)straight afterwait_for).
This confirms the concern raised in #123 (comment), here with a reproduction.
Reproduction. Host only, GCC 13.3 Debug, polaron. I added three probe tests to tests/async/test_debug_scheduler.cpp; they are not part of the PR. Each uses the PR's own ExternalWakeup pattern without the gate and the 20 ms sleep: a std::thread calls BasicTask::reschedule_after_external_wait at once. The main thread calls scheduler.run() and then scheduler.wait_for([&]{ return resumed; }). Then:
- stale count: it reads
external_waits()before joining the waker. Over 10 × 20 000 iterations it read 1 in 37 iterations after the task had resumed; - destroy: it lets the scheduler go out of scope and only then joins the waker. All 10/10 runs of 20 000 iterations aborted with
CHECK at debug_scheduler.hpp:83 ➡️ DebugScheduler destroyed while tasks await external wakeups, external_waits_ = 0. The message prints0because the waker decremented the counter between the check and the message, i.e. concurrently with destruction; - widened window: a
DebugSchedulersubclass sleeps 50 ms inend_external_wait()before calling the base, as preemption of the waker would. With it,wait_forreturned withresumed=1 external_waits=1.
Possible fix (tested locally). Do the enqueue and the decrement under a single lock acquisition, and notify while still holding the lock. In my probe I added a virtual IScheduler::reschedule_after_external_wait(BasicTask&&). Its default implementation keeps today's two calls; DebugScheduler overrides it to push the task, decrement and notify_all under mutex_. BasicTask::reschedule_after_external_wait applies make_sole_owner and then calls it. If no task is left it calls only end_external_wait(). I also moved the notify_all into the locked region in enqueue_task and end_external_wait. With those changes the same probes gave 0/10 aborts and 0 stale reads, and the 201 existing uni20_async_tests passed.
Not checked. I did not reproduce this on the CUDA path: the new CudaTensorTest.WaitBlocksForStreamReturnedByHostCallback passed 3000/3000 repetitions on a GV100. There the driving thread's work after the wakeup (stream lease, cudaLaunchHostFunc) makes the window small, and streams().synchronize() runs before the scheduler is destroyed. No TSan/ASan run, and no Clang build.
There was a problem hiding this comment.
Thanks, confirmed. This reviewed b08ab3b7; it's fixed in 47d6191 the same way you suggest. end_external_wait() became IScheduler::complete_external_wait(BasicTask&&) (the default just reschedules), and DebugScheduler pushes, decrements and notifies under one lock acquisition. BasicTask::reschedule_after_external_wait applies make_sole_owner and makes that single call. enqueue_task also notifies under the lock now.
Your probe is now a regression test in 6f0bebb: SchedulerMayBeDestroyedOnceTheExternallyWokenTaskHasRun. The waker wakes the task at once, and the scheduler is destroyed before the waker is joined, for 20000 iterations. With the previous scheduler files it fails in 3/3 runs, at iterations 749, 1528 and 988 (stale external_waits(), then the destructor CHECK). With the fix it passes 3/3, and uni20_async_tests passes 202/202.
There was a problem hiding this comment.
Review Bot: recheck of exact head 6f0bebba7cb4df61ebbb50ffdc4cb36cecef719e
Confirmed fixed. I ran the same test source against both scheduler versions (GCC 13.3 Debug, polaron). With the four scheduler headers from b08ab3b7, SchedulerMayBeDestroyedOnceTheExternallyWokenTaskHasRun aborts in 3/3 runs (iterations 2766, 1584, 4329). At this head it passes in 3/3 runs. My multi-task probes also pass at this head: 8 tasks woken twice each, fifo/reverse/random order, 5000 iterations, scheduler destroyed before the wakers are joined. They pass 12/12, against 9/12 aborted with the old headers. TSan reports no warnings on the DebugScheduler tests. Details are in the review comment for this head.
reschedule_after_external_wait rescheduled the task and then called end_external_wait() separately. In between, the driving thread could run the woken task to completion; a program that then correctly concluded that all its tasks had finished could destroy the scheduler while the waking thread was still about to touch it. It also let external_waits() and the destructor's balance check observe a pending wait for a task that had already run. Replace end_external_wait() with complete_external_wait(BasicTask&&), which reschedules the task when present and releases the wait as one operation. DebugScheduler requeues, decrements and notifies inside a single critical section, so once the task can run the waking thread only releases the lock. Other schedulers keep the default, which just reschedules. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The waking thread wakes the task at once, and the scheduler is destroyed after the task has run but before that thread is joined. Against the previous reschedule-then-end ordering this fails within about 750 to 1500 of the 20000 iterations, first with a stale external_waits() count and then with the destructor's balance check; with complete_external_wait it passes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
Review Bot: review of exact head No defect found. The earlier finding (#123 (comment), wake-before-release race in What changed since What I reviewed. The incremental diff ( Verified (polaron, GCC 13.3.0, Debug, the REVIEW-hint CMake options with
Not checked.
Review tier: full (code or new pull request) · model: claude-opus-5-5 · effort: high |
Independent of #121 and #122.
CudaTensorTest.StreamOrdered*aborted intermittently onmain(about once in 150 repetitions on polaron):submittedwas false, then a CHECK fired on destruction of a still-active task. The cause is inDebugScheduler, not the test alone:stream_count = 1. The initialization memset holds the only stream until a CUDA host callback returns it to the pool. If the task'sacquire_streamruns first it suspends, and the callback thread later reschedules it.DebugScheduler's queue had no lock, so that reschedule raced the driving thread.help_while_waiting()(behindget_wait) reported DEADLOCK whenever nothing was runnable, even with a wakeup pending.run_all()meaning completion. It returns when nothing is runnable.Change. A balanced external-wait registration on the scheduler interface:
BasicTask::begin_external_wait()before an awaiter hands its task to another party, andBasicTask::reschedule_after_external_wait()to wake it (reschedule first, then end the wait, so an empty queue is never observed with the task in transit). The CUDA resource, stream and cuBLAS execution awaiters register their waits; other schedulers treat registration as a no-op.DebugSchedulerguards its queue with a mutex and condition variable. Waits block, with no time limit, while a registered wakeup is pending, and DEADLOCK is reported only when it is provable: nothing runnable and nothing pending.run_all()keeps its meaning as a debugging aid that runs until nothing is runnable; it does not wait for externally suspended tasks. An earlier revision made it wait, andCublasExecutionTest.AsyncOutputRemainsPendingUntilCublasSubmissionCompletesthen hung: that test holds every cuBLAS handle on the driving thread and releases them only afterrun_all()returns. The stream-ordered tests now wait for their task withwait_for.Also moves
test_runtime_floating_eq.cppfromtests/commoninto the numerical-precision module. It uses the numericsPrecisionContextregistry, which includes the tensor layer (itsmatrix()returns aDenseMatrix), souni20_common_testsdid not build in CUDA configurations.Tests. New host tests drive a task woken from another thread through
run_all()(returns with the wakeup pending) andwait_for()(blocks for it). A new CUDA test holds the only stream with a host-function gate so the task is deterministically suspended on a callback wakeup before waiting.Validation (polaron, 2x Quadro GV100, CUDA 12.9, GCC 13 Debug; ryzen CPU-only GCC 13 Debug):
uni20_common_tests;ctest2760/2760.ctest2547/2547.CudaTensorTest.Stream*and the new CUDA wait test: 900 repetitions, 0 failures. Onmainthe same tests abort within about 150.scripts/check-docs.pypasses; clang-format clean on changed files. Not run: Clang, Release, CUDA 13.🤖 Generated with Claude Code