From fbc09b4a6b870605ef6276b5696b168d3d8edb56 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Fri, 21 Aug 2026 13:58:07 -0700 Subject: [PATCH 1/3] up --- CMakeLists.txt | 2 + extension/llm/scheduler/CMakeLists.txt | 41 + extension/llm/scheduler/scheduler.h | 498 ++++++++ extension/llm/scheduler/test/CMakeLists.txt | 18 + .../llm/scheduler/test/scheduler_test.cpp | 1003 +++++++++++++++++ 5 files changed, 1562 insertions(+) create mode 100644 extension/llm/scheduler/CMakeLists.txt create mode 100644 extension/llm/scheduler/scheduler.h create mode 100644 extension/llm/scheduler/test/CMakeLists.txt create mode 100644 extension/llm/scheduler/test/scheduler_test.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 460c6a4a041..642351637bd 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1001,6 +1001,8 @@ if(EXECUTORCH_BUILD_EXTENSION_LLM) list(APPEND _executorch_extensions tokenizers) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/cache) list(APPEND _executorch_extensions extension_llm_cache) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/scheduler) + list(APPEND _executorch_extensions extension_llm_scheduler) endif() if(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL) diff --git a/extension/llm/scheduler/CMakeLists.txt b/extension/llm/scheduler/CMakeLists.txt new file mode 100644 index 00000000000..975a9173cea --- /dev/null +++ b/extension/llm/scheduler/CMakeLists.txt @@ -0,0 +1,41 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Step scheduler for batched LLM serving. Header-only and free of ExecuTorch +# runtime types, so it is an INTERFACE target: consumers get the include path +# and nothing to link. + +if(NOT EXECUTORCH_ROOT) + set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..) +endif() + +add_library(extension_llm_scheduler INTERFACE) +# std::optional and std::variant in the public headers. +target_compile_features(extension_llm_scheduler INTERFACE cxx_std_17) +target_include_directories( + extension_llm_scheduler INTERFACE ${_common_include_directories} +) +target_compile_options( + extension_llm_scheduler INTERFACE ${_common_compile_options} +) + +install( + TARGETS extension_llm_scheduler + EXPORT ExecuTorchTargets + DESTINATION ${CMAKE_INSTALL_LIBDIR} + INCLUDES + DESTINATION ${_common_include_directories} +) +install( + DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/ + DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/executorch/extension/llm/scheduler + FILES_MATCHING + PATTERN "*.h" +) + +if(BUILD_TESTING) + add_subdirectory(test) +endif() diff --git a/extension/llm/scheduler/scheduler.h b/extension/llm/scheduler/scheduler.h new file mode 100644 index 00000000000..72d4b74e6dd --- /dev/null +++ b/extension/llm/scheduler/scheduler.h @@ -0,0 +1,498 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +// Step scheduler for batched LLM serving. A Request is one step for one +// sequence -- a single decode token, or one prefill chunk -- not a whole +// generation. Callers submit a step, await its sampled token, and submit the +// next one. +// +// get_work() takes decodes first, up to max_decode_sequences, then spends the +// rest of max_batch_size on prefill, so prefill never delays a queued decode. +// Decode is one arrival-order queue. Prefill is a FIFO per session plus a +// rotation over sessions, and a pass takes at most one chunk from each, so a +// long prompt advances a chunk at a time instead of monopolising the batch. +// +// No tensors, cache, or model here: where a step's KV lives belongs to the +// planner. Callers split prefill into chunks of at most max_prefill_chunk_size, +// which needs the session's position sequence and so is caller state. +// +// The scheduler reads exactly three things from a Request: +// +// tokens.size() budget arithmetic, and decode vs prefill routing +// session_id which prefill FIFO, and a slot in the rotation +// request_id the key its result is settled through +// +// Everything else lives in RequestParams / ResponsePayload and is carried from +// submit() +// to Batch and back without being inspected, so payload fields can be added +// without touching any scheduling logic. +// +// Every public method is guarded by mutex_. The Scheduler must outlive every +// thread using it; params() hands out a reference into it. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace executorch { +namespace extension { +namespace llm { +namespace scheduler { + +using Token = std::int64_t; +using RequestId = std::int64_t; +using SessionId = std::int64_t; +using Position = std::int32_t; + +// --- Payload: carried through the scheduler untouched ---------------------- + +// Carried per step, so a caller may change sampling between turns. +struct SamplingParams { + float temperature = 0.0f; + float top_p = 1.0f; + std::int32_t top_k = 0; + std::uint64_t seed = 0; +}; + +// Which positions of a step the model must produce output for. A verify step +// needs every row -- it inspects the prediction at each drafted position -- +// while a decode or prefill chunk only needs the last. A sampled step yields +// one token per row; an unsampled one yields one distribution per row. +enum class OutputRows { + Last, + All, +}; + +// Raw model output for a step whose Request asked not to sample +struct LogitsBlock { + std::vector data; + std::int32_t n_rows = 0; + std::int32_t vocab = 0; +}; +using LogitsPtr = std::shared_ptr; + +// How to execute a step. Promote to a variant if a second step kind ever needs +// different fields rather than different values. +struct RequestParams { + Position position = 0; + // Absent means: do not sample, return the raw distribution. Rejection + // sampling and constrained decoding need the distribution; greedy + // speculative verification does not, since argmax at every row is just + // sampling at temperature zero. + std::optional sampling; + OutputRows output_rows = OutputRows::Last; +}; + +// Which alternative is held follows from whether the Request asked to sample. +// A sampled step yields one token per output row, so a greedy verify round +// returns the prediction at every drafted position; deciding how many to +// accept, and where the session therefore continues, is the caller's job. +using ResponsePayload = std::variant, LogitsPtr>; + +// --- Scheduling ------------------------------------------------------------ + +// One step for one sequence. request_id names the submission and is always +// unique, so the same work submitted twice is two requests that both run and +// both get answered. Only request_id, session_id, and tokens.size() are read +// by the scheduler; `params` is carried. +struct Request { + RequestId request_id = 0; + SessionId session_id = 0; + std::vector tokens; + RequestParams params; + + std::int32_t n_tokens() const { + return static_cast(tokens.size()); + } + + // A one-token prefill is indistinguishable from a decode and is treated as + // one: same single row of output, one decode slot. + bool is_decode() const { + return tokens.size() == 1; + } +}; + +struct Response { + RequestId request_id = 0; + SessionId session_id = 0; + ResponsePayload payload; +}; + +// One batch of steps for a single forward. Decodes first, then prefills, in +// admission order. +struct Batch { + std::vector requests; + + bool empty() const { + return requests.empty(); + } + std::int32_t n_tokens() const { + std::int32_t n = 0; + for (const Request& r : requests) { + n += r.n_tokens(); + } + return n; + } +}; + +class SchedulerParams { + public: + explicit SchedulerParams( + std::int32_t max_decode_sequences = 32, + std::int32_t max_prefill_chunk_size = 256) + : max_decode_sequences_(max_decode_sequences), + max_prefill_chunk_size_(max_prefill_chunk_size) { + if (max_decode_sequences_ <= 0 || max_prefill_chunk_size_ <= 0) { + throw std::invalid_argument("SchedulerParams: limits must be positive"); + } + // max_batch_size() is derived, so the combination has to be representable. + // Silently wrapping would hand get_work() a negative budget, which admits + // nothing and leaves every request unresolved. + const std::int64_t budget = + 2 * static_cast(max_prefill_chunk_size_) + + max_decode_sequences_; + if (budget > std::numeric_limits::max()) { + throw std::invalid_argument( + "SchedulerParams: 2 * max_prefill_chunk_size + max_decode_sequences " + "overflows int32"); + } + } + + // Decodes admitted per batch; the rest wait. Bounds the working set the + // planner must hold at once. + std::int32_t max_decode_sequences() const { + return max_decode_sequences_; + } + // Largest prefill chunk accepted. A larger submit is rejected, not split. + std::int32_t max_prefill_chunk_size() const { + return max_prefill_chunk_size_; + } + // Room for two full prefill chunks beside a saturated decode batch. Exceeding + // max_decode_sequences is load-bearing: it means a full-size chunk still fits + // once decodes are taken. + std::int32_t max_batch_size() const { + return 2 * max_prefill_chunk_size_ + max_decode_sequences_; + } + + private: + std::int32_t max_decode_sequences_; + std::int32_t max_prefill_chunk_size_; +}; + +// Held by shared_ptr because std::promise is move-only. +struct PendingRequest { + Request request; + // A deque cannot cheaply erase from the middle, so fail() marks and the + // queues drop on the way past. + bool cancelled = false; + // Still waiting in a queue, as opposed to handed to a Batch. Only a queued + // step is counted by queued_, so this is what fail() checks before + // decrementing it. + bool queued = true; + std::promise promise; +}; + +using PendingPtr = std::shared_ptr; +using PendingQueue = std::deque; + +class Scheduler { + public: + explicit Scheduler(SchedulerParams params) : params_(params) {} + + // Unique within this Scheduler, which is the scope request_id must be unique + // over. Callable from any thread. + RequestId next_request_id() { + return next_request_id_.fetch_add(1, std::memory_order_relaxed); + } + + // The future carries the sampled token, or an exception if the step is + // rejected: no tokens, a prefill above max_prefill_chunk_size, or a + // request_id already outstanding. + std::future submit(Request request) { + auto p = std::make_shared(); + p->request = std::move(request); + std::future fut = p->promise.get_future(); + + const std::int32_t n = p->request.n_tokens(); + if (n == 0) { + set_error_(*p, "step carries no tokens"); + return fut; + } + if (!p->request.is_decode() && n > params_.max_prefill_chunk_size()) { + set_error_(*p, "prefill chunk exceeds max_prefill_chunk_size"); + return fut; + } + + std::lock_guard g(mutex_); + // emplace drops a duplicate silently, which would leave this promise never + // settled and the caller's future blocked forever. + auto [slot, inserted] = pending_requests_.emplace(p->request.request_id, p); + if (!inserted) { + set_error_(*p, "request_id already outstanding"); + return fut; + } + + try { + if (p->request.is_decode()) { + decode_queue_.push_back(p); + } else { + const SessionId sid = p->request.session_id; + // Rotation before map: a rotation entry with no session is harmless, + // since get_work() drops it, whereas a session missing from the + // rotation is never visited and its chunks never run. + if (prefill_by_session_.find(sid) == prefill_by_session_.end()) { + prefill_rotation_.push_back(sid); + } + prefill_by_session_[sid].push_back(p); + } + } catch (...) { + // Registered but unqueued would be unschedulable, so its future would + // never settle. Undo and reject instead. + pending_requests_.erase(slot); + set_error_(*p, "failed to queue step"); + return fut; + } + queued_ += 1; + return fut; + } + + // Whether get_work() would return a non-empty batch. + bool has_work() const { + std::lock_guard g(mutex_); + return queued_ > 0; + } + + // Decodes first, up to max_decode_sequences, then prefill round-robin across + // sessions until max_batch_size is spent. Empty when nothing is queued. + Batch get_work() { + Batch batch; + std::lock_guard g(mutex_); + std::int32_t budget = params_.max_batch_size(); + + take_decodes_(batch, budget); + while (budget > 0 && take_prefill_pass_(batch, budget)) { + } + return batch; + } + + // Settle a sampled step: one token per output row, so a greedy verify round + // reports the target's prediction at every drafted position. An unknown id is + // ignored, so a completion racing fail_all() is not an error. + void complete(RequestId request_id, std::vector tokens) { + Response r; + r.payload = std::move(tokens); + settle_(request_id, std::move(r)); + } + + // Settle an unsampled step with its raw distribution. + void complete(RequestId request_id, LogitsPtr logits) { + Response r; + r.payload = std::move(logits); + settle_(request_id, std::move(r)); + } + + // Fails one step, queued or in flight. Cancelling a whole generation is the + // caller's loop over its outstanding ids. + void fail(RequestId request_id, const std::string& what) { + PendingPtr p; + { + std::lock_guard g(mutex_); + auto it = pending_requests_.find(request_id); + if (it == pending_requests_.end()) { + return; + } + p = it->second; + p->cancelled = true; + // An in-flight step was already uncounted by get_work(); decrementing + // again would underflow queued_ or mask other queued work. + if (p->queued) { + p->queued = false; + queued_ -= 1; + } + pending_requests_.erase(it); + } + set_error_(*p, what); + } + + // For shutdown. + void fail_all(const std::string& what) { + std::vector doomed; + { + std::lock_guard g(mutex_); + for (auto& entry : pending_requests_) { + entry.second->cancelled = true; + doomed.push_back(entry.second); + } + pending_requests_.clear(); + queued_ = 0; + decode_queue_.clear(); + prefill_by_session_.clear(); + prefill_rotation_.clear(); + } + for (PendingPtr& p : doomed) { + set_error_(*p, what); + } + } + + // Whether `request_id` is still queued or in flight. + bool pending(RequestId request_id) const { + std::lock_guard g(mutex_); + return pending_requests_.find(request_id) != pending_requests_.end(); + } + + // Steps waiting in a queue, excluding those in flight. + std::size_t queued() const { + std::lock_guard g(mutex_); + return queued_; + } + + const SchedulerParams& params() const { + return params_; + } + + private: + // Fills in identity and hands `r` to the waiter. Erasing under the lock + // before settling is what stops complete() and fail() both resolving the + // same promise. + void settle_(RequestId request_id, Response&& r) { + PendingPtr p; + { + std::lock_guard g(mutex_); + auto it = pending_requests_.find(request_id); + if (it == pending_requests_.end()) { + return; + } + p = it->second; + pending_requests_.erase(it); + r.request_id = request_id; + r.session_id = p->request.session_id; + } + p->promise.set_value(std::move(r)); // outside the lock: wakes the waiter + } + + static void set_error_(PendingRequest& p, const std::string& what) { + p.promise.set_exception( + std::make_exception_ptr(std::runtime_error("scheduler: " + what))); + } + + // fail() already settled these and decremented queued_; they are only still + // here because a deque cannot erase from the middle. + static void drop_cancelled_(PendingQueue& q) { + while (!q.empty() && q.front()->cancelled) { + q.pop_front(); + } + } + + // Decodes in arrival order. A decode is one token and max_batch_size exceeds + // max_decode_sequences, so the budget cannot run out here. + // Caller holds mutex_. + void take_decodes_(Batch& batch, std::int32_t& budget) { + for (std::int32_t n = 0; n < params_.max_decode_sequences(); ++n) { + drop_cancelled_(decode_queue_); + if (decode_queue_.empty()) { + return; + } + // Copy into the batch first: it is the only step here that can throw, and + // popping first would lose the request entirely. + batch.requests.push_back(decode_queue_.front()->request); + decode_queue_.front()->queued = false; + decode_queue_.pop_front(); + budget -= 1; + queued_ -= 1; + } + } + + // At most one chunk per session. Returns whether anything was taken. + // Caller holds mutex_. + bool take_prefill_pass_(Batch& batch, std::int32_t& budget) { + bool progress = false; + std::vector deferred; // head did not fit + std::vector served; + + // Nothing is put back into the rotation here, so it shrinks by one per + // iteration and each session is visited at most once. + while (budget > 0 && !prefill_rotation_.empty()) { + const SessionId sid = prefill_rotation_.front(); + prefill_rotation_.pop_front(); + auto it = prefill_by_session_.find(sid); + if (it == prefill_by_session_.end()) { + continue; + } + PendingQueue& dq = it->second; + drop_cancelled_(dq); + if (dq.empty()) { + prefill_by_session_.erase(it); + continue; + } + const std::int32_t n = dq.front()->request.n_tokens(); + if (n > budget) { + // Pass the turn on rather than stop, so a smaller chunk behind can + // still use the budget. + deferred.push_back(sid); + continue; + } + // Copy into the batch before mutating anything else; see take_decodes_. + batch.requests.push_back(dq.front()->request); + dq.front()->queued = false; + dq.pop_front(); + budget -= n; + queued_ -= 1; + progress = true; + if (dq.empty()) { + prefill_by_session_.erase(it); + } else { + served.push_back(sid); + } + } + + // Deferred sessions were popped from the front, so they belong ahead of + // whatever the pass never reached; served sessions go to the back. Everyone + // who got nothing keeps their order and outranks everyone who did. + for (auto it = deferred.rbegin(); it != deferred.rend(); ++it) { + prefill_rotation_.push_front(*it); + } + for (SessionId sid : served) { + prefill_rotation_.push_back(sid); + } + return progress; + } + + SchedulerParams params_; + mutable std::mutex mutex_; + + PendingQueue decode_queue_; + // A session is in prefill_rotation_ exactly when it has an entry here, and + // the entry is erased as soon as its deque drains, so the rotation cannot + // accumulate duplicates. + std::unordered_map prefill_by_session_; + std::deque prefill_rotation_; + + // Index from submit() until complete()/fail(); the queues hold the same + // shared_ptr. + std::unordered_map pending_requests_; + std::size_t queued_ = 0; + std::atomic next_request_id_{1}; +}; + +} // namespace scheduler +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/scheduler/test/CMakeLists.txt b/extension/llm/scheduler/test/CMakeLists.txt new file mode 100644 index 00000000000..5a13866da35 --- /dev/null +++ b/extension/llm/scheduler/test/CMakeLists.txt @@ -0,0 +1,18 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +cmake_minimum_required(VERSION 3.19) + +set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../..) + +include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake) + +set(_test_srcs scheduler_test.cpp) + +et_cxx_test( + extension_llm_scheduler_test SOURCES ${_test_srcs} EXTRA_LIBS + extension_llm_scheduler +) diff --git a/extension/llm/scheduler/test/scheduler_test.cpp b/extension/llm/scheduler/test/scheduler_test.cpp new file mode 100644 index 00000000000..cafea7d7b2d --- /dev/null +++ b/extension/llm/scheduler/test/scheduler_test.cpp @@ -0,0 +1,1003 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +using executorch::extension::llm::scheduler::Batch; +using executorch::extension::llm::scheduler::LogitsBlock; +using executorch::extension::llm::scheduler::LogitsPtr; +using executorch::extension::llm::scheduler::OutputRows; +using executorch::extension::llm::scheduler::Position; +using executorch::extension::llm::scheduler::Request; +using executorch::extension::llm::scheduler::RequestId; +using executorch::extension::llm::scheduler::RequestParams; +using executorch::extension::llm::scheduler::Response; +using executorch::extension::llm::scheduler::SamplingParams; +using executorch::extension::llm::scheduler::Scheduler; +using executorch::extension::llm::scheduler::SchedulerParams; +using executorch::extension::llm::scheduler::SessionId; +using executorch::extension::llm::scheduler::Token; + +namespace { + +// A step of `n_tokens` for `session` at `position`, sampled by default. Token +// values are irrelevant to scheduling, so they are all the same. +Request make_request( + RequestId request_id, + SessionId session, + int n_tokens, + Position position, + OutputRows rows = OutputRows::Last, + bool sample = true) { + Request r; + r.request_id = request_id; + r.session_id = session; + r.tokens.assign(static_cast(n_tokens), 7); + r.params.position = position; + r.params.output_rows = rows; + if (sample) { + r.params.sampling = SamplingParams{}; + } + return r; +} + +std::vector ids(const Batch& b) { + std::vector out; + out.reserve(b.requests.size()); + for (const Request& r : b.requests) { + out.push_back(r.request_id); + } + return out; +} + +bool settled(std::future& f) { + return f.wait_for(std::chrono::seconds(0)) == std::future_status::ready; +} + +bool holds_tokens(const Response& r) { + return std::holds_alternative>(r.payload); +} + +const std::vector& tokens_of(const Response& r) { + return std::get>(r.payload); +} + +LogitsPtr logits_of(const Response& r) { + return std::get(r.payload); +} + +// Bounded wait. A regression that leaves a future unresolved must fail +// diagnostically rather than hang the suite -- several tests here exist +// precisely to catch non-settlement. +constexpr std::chrono::seconds kSettleTimeout{5}; + +testing::AssertionResult Settles(std::future& f) { + if (f.wait_for(kSettleTimeout) != std::future_status::ready) { + return testing::AssertionFailure() + << "future did not settle within " << kSettleTimeout.count() << "s"; + } + return testing::AssertionSuccess(); +} + +} // namespace + +// Guards every get() whose absence would block instead of failing. +#define ASSERT_SETTLES(f) ASSERT_TRUE(Settles(f)) + +namespace {} // namespace + +// --- SchedulerParams ------------------------------------------------------- + +TEST(SchedulerParamsTest, DefaultsAndDerivedBatchSize) { + SchedulerParams p; + EXPECT_EQ(p.max_decode_sequences(), 32); + EXPECT_EQ(p.max_prefill_chunk_size(), 256); + EXPECT_EQ(p.max_batch_size(), 2 * 256 + 32); +} + +TEST(SchedulerParamsTest, BatchSizeIsDerivedNotStored) { + SchedulerParams p(2, 4); + EXPECT_EQ(p.max_batch_size(), 2 * 4 + 2); +} + +// The decode loop omits a budget check because this always holds. +TEST(SchedulerParamsTest, BatchSizeAlwaysExceedsDecodeCap) { + for (std::int32_t d = 1; d <= 8; ++d) { + for (std::int32_t c = 1; c <= 8; ++c) { + SchedulerParams p(d, c); + EXPECT_GT(p.max_batch_size(), p.max_decode_sequences()); + EXPECT_GE(p.max_batch_size() - p.max_decode_sequences(), c); + } + } +} + +TEST(SchedulerParamsTest, RejectsNonPositiveLimits) { + EXPECT_THROW(SchedulerParams(0, 4), std::invalid_argument); + EXPECT_THROW(SchedulerParams(2, 0), std::invalid_argument); + EXPECT_THROW(SchedulerParams(-1, 4), std::invalid_argument); + EXPECT_THROW(SchedulerParams(2, -1), std::invalid_argument); +} + +// A wrapped max_batch_size() would be a negative budget: get_work() would admit +// nothing and every request would sit unresolved forever. +TEST(SchedulerParamsTest, RejectsCombinationsThatOverflowBatchSize) { + constexpr std::int32_t kMax = std::numeric_limits::max(); + EXPECT_THROW(SchedulerParams(1, kMax), std::invalid_argument); + EXPECT_THROW(SchedulerParams(kMax, kMax), std::invalid_argument); + EXPECT_THROW(SchedulerParams(kMax, kMax / 2), std::invalid_argument); +} + +TEST(SchedulerParamsTest, AcceptsTheLargestRepresentableCombination) { + constexpr std::int32_t kMax = std::numeric_limits::max(); + // 2 * chunk + decodes == kMax exactly. + const std::int32_t chunk = (kMax - 1) / 2; + SchedulerParams p(1, chunk); + EXPECT_GT(p.max_batch_size(), 0); + EXPECT_EQ(p.max_batch_size(), 2 * chunk + 1); + EXPECT_THROW(SchedulerParams(2, chunk), std::invalid_argument); +} + +// --- Request / Batch shape ------------------------------------------------- + +TEST(RequestTest, OneTokenStepIsADecode) { + EXPECT_TRUE(make_request(1, 10, 1, 0).is_decode()); + EXPECT_FALSE(make_request(1, 10, 2, 0).is_decode()); +} + +TEST(RequestTest, OutputRowsIsCarriedNotInterpreted) { + EXPECT_EQ( + make_request(1, 10, 8, 0, OutputRows::Last).params.output_rows, + OutputRows::Last); + EXPECT_EQ( + make_request(1, 10, 5, 0, OutputRows::All).params.output_rows, + OutputRows::All); +} + +TEST(RequestTest, PayloadSupportsAllFourQuadrants) { + EXPECT_TRUE( + make_request(1, 10, 1, 0, OutputRows::Last, true).params.sampling); + EXPECT_TRUE(make_request(2, 10, 5, 0, OutputRows::All, true).params.sampling); + EXPECT_FALSE( + make_request(3, 10, 1, 0, OutputRows::Last, false).params.sampling); + EXPECT_FALSE( + make_request(4, 10, 5, 0, OutputRows::All, false).params.sampling); +} + +TEST(BatchTest, EmptyBatchReportsZeros) { + Scheduler s{SchedulerParams(1, 4)}; + Batch b = s.get_work(); + EXPECT_TRUE(b.empty()); + EXPECT_EQ(b.n_tokens(), 0); +} + +TEST(BatchTest, SumsTokensAcrossSteps) { + Scheduler s{SchedulerParams(2, 8)}; + s.submit(make_request(1, 10, 1, 0)); + s.submit(make_request(2, 20, 1, 0)); + s.submit(make_request(3, 30, 6, 0)); + s.submit(make_request(4, 40, 5, 100, OutputRows::All)); + + Batch b = s.get_work(); + EXPECT_EQ(ids(b), (std::vector{1, 2, 3, 4})); + EXPECT_EQ(b.n_tokens(), 1 + 1 + 6 + 5); +} + +TEST(BatchTest, CarriesPayloadUninspected) { + Scheduler s{SchedulerParams(1, 8)}; + s.submit(make_request(1, 40, 5, 100, OutputRows::All, /*sample=*/false)); + + Batch b = s.get_work(); + ASSERT_EQ(b.requests.size(), 1u); + EXPECT_EQ(b.requests[0].params.position, 100); + EXPECT_EQ(b.requests[0].params.output_rows, OutputRows::All); + EXPECT_FALSE(b.requests[0].params.sampling.has_value()); +} + +// --- submit ---------------------------------------------------------------- + +TEST(SubmitTest, RejectsEmptyStep) { + Scheduler s{SchedulerParams(1, 4)}; + auto f = s.submit(make_request(1, 10, 0, 0)); + ASSERT_SETTLES(f); + EXPECT_THROW((void)f.get(), std::runtime_error); + EXPECT_EQ(s.queued(), 0u); +} + +TEST(SubmitTest, RejectsPrefillAboveChunkSize) { + Scheduler s{SchedulerParams(1, 4)}; + auto f = s.submit(make_request(1, 10, 9, 0)); + ASSERT_SETTLES(f); + EXPECT_THROW((void)f.get(), std::runtime_error); + EXPECT_EQ(s.queued(), 0u); +} + +TEST(SubmitTest, RejectsDuplicateRequestIdRatherThanHanging) { + Scheduler s{SchedulerParams(2, 4)}; + auto first = s.submit(make_request(7, 10, 1, 0)); + auto dup = s.submit(make_request(7, 20, 1, 0)); + + ASSERT_SETTLES(dup); + EXPECT_THROW((void)dup.get(), std::runtime_error); + EXPECT_EQ(s.queued(), 1u) << "rejected step must not inflate the count"; + EXPECT_EQ(ids(s.get_work()), (std::vector{7})); + s.complete(7, std::vector{42}); + ASSERT_SETTLES(first); + EXPECT_EQ(tokens_of(first.get()), (std::vector{42})); +} + +TEST(SubmitTest, RequestIdIsReusableOnceSettled) { + Scheduler s{SchedulerParams(2, 4)}; + auto first = s.submit(make_request(7, 10, 1, 0)); + s.get_work(); + s.complete(7, std::vector{1}); + ASSERT_SETTLES(first); + EXPECT_EQ(tokens_of(first.get())[0], 1); + EXPECT_FALSE(s.pending(7)); + + auto second = s.submit(make_request(7, 10, 1, 1)); + s.get_work(); + s.complete(7, std::vector{2}); + ASSERT_SETTLES(second); + EXPECT_EQ(tokens_of(second.get())[0], 2); +} + +TEST(SubmitTest, DuplicateWorkWithDistinctIdsBothRun) { + Scheduler s{SchedulerParams(2, 4)}; + Request a = make_request(0, 10, 1, 5); + Request b = make_request(0, 10, 1, 5); + a.request_id = s.next_request_id(); + b.request_id = s.next_request_id(); + + auto fa = s.submit(a); + auto fb = s.submit(b); + EXPECT_EQ(s.get_work().requests.size(), 2u); + + s.complete(a.request_id, std::vector{111}); + s.complete(b.request_id, std::vector{222}); + ASSERT_SETTLES(fa); + EXPECT_EQ(tokens_of(fa.get())[0], 111); + ASSERT_SETTLES(fb); + EXPECT_EQ(tokens_of(fb.get())[0], 222); +} + +TEST(SubmitTest, RejectedStepEntersNoQueue) { + Scheduler s{SchedulerParams(1, 4)}; + auto empty = s.submit(make_request(1, 10, 0, 0)); + auto toobig = s.submit(make_request(2, 10, 9, 0)); + ASSERT_SETTLES(empty); + EXPECT_THROW((void)empty.get(), std::runtime_error); + ASSERT_SETTLES(toobig); + EXPECT_THROW((void)toobig.get(), std::runtime_error); + + EXPECT_EQ(s.queued(), 0u); + EXPECT_FALSE(s.has_work()); + EXPECT_FALSE(s.pending(1)); + EXPECT_FALSE(s.pending(2)); + EXPECT_TRUE(s.get_work().empty()); +} + +TEST(SubmitTest, NextRequestIdIsUnique) { + Scheduler s{SchedulerParams(2, 4)}; + std::set seen; + for (int i = 0; i < 1000; ++i) { + EXPECT_TRUE(seen.insert(s.next_request_id()).second); + } +} + +TEST(SubmitTest, NextRequestIdIsUniqueAcrossThreads) { + Scheduler s{SchedulerParams(2, 4)}; + constexpr int kThreads = 8; + constexpr int kPer = 500; + std::vector> out(kThreads); + std::vector threads; + for (int t = 0; t < kThreads; ++t) { + threads.emplace_back([&, t] { + for (int i = 0; i < kPer; ++i) { + out[t].push_back(s.next_request_id()); + } + }); + } + for (std::thread& t : threads) { + t.join(); + } + std::set seen; + for (const auto& v : out) { + for (RequestId id : v) { + EXPECT_TRUE(seen.insert(id).second); + } + } + EXPECT_EQ(seen.size(), static_cast(kThreads * kPer)); +} + +// --- decode scheduling ----------------------------------------------------- + +TEST(DecodeTest, ServedInArrivalOrderUpToTheCap) { + Scheduler s{SchedulerParams(2, 4)}; + s.submit(make_request(1, 10, 1, 0)); + s.submit(make_request(2, 20, 1, 0)); + s.submit(make_request(3, 30, 1, 0)); + + EXPECT_EQ(ids(s.get_work()), (std::vector{1, 2})); + EXPECT_EQ(s.queued(), 1u); + EXPECT_EQ(ids(s.get_work()), (std::vector{3})); +} + +// A shorter queue must not give a later arrival a head start. +TEST(DecodeTest, StaysFifoAcrossDrainAndRefill) { + Scheduler s{SchedulerParams(2, 4)}; + s.submit(make_request(1, 10, 1, 0)); + s.submit(make_request(2, 20, 1, 0)); + s.submit(make_request(3, 30, 1, 0)); + EXPECT_EQ(ids(s.get_work()), (std::vector{1, 2})); + + s.submit(make_request(4, 40, 1, 0)); + EXPECT_EQ(ids(s.get_work()), (std::vector{3, 4})); +} + +TEST(DecodeTest, BeatsPrefillAndStillLeavesRoomForAFullChunk) { + Scheduler s{SchedulerParams(3, 4)}; // batch = 11 + s.submit(make_request(1, 10, 1, 0)); + s.submit(make_request(2, 20, 1, 0)); + s.submit(make_request(3, 30, 1, 0)); + s.submit(make_request(50, 90, 4, 0)); + + Batch b = s.get_work(); + EXPECT_EQ(ids(b), (std::vector{1, 2, 3, 50})); + EXPECT_EQ(b.n_tokens(), 3 + 4); +} + +// --- prefill scheduling ---------------------------------------------------- + +TEST(PrefillTest, ChunksOfOneSessionStayInOrder) { + Scheduler s{SchedulerParams(1, 2)}; + s.submit(make_request(1, 7, 2, 0)); + s.submit(make_request(2, 7, 2, 2)); + EXPECT_EQ(ids(s.get_work()), (std::vector{1, 2})); +} + +TEST(PrefillTest, LongPromptCannotHogTheBatch) { + Scheduler s{SchedulerParams(2, 4)}; + for (int c = 0; c < 4; ++c) { + s.submit(make_request(1 + c, 10, 4, c * 4)); + } + s.submit(make_request(5, 20, 4, 0)); + s.submit(make_request(6, 20, 4, 4)); + + EXPECT_EQ(ids(s.get_work()), (std::vector{1, 5})); + EXPECT_EQ(ids(s.get_work()), (std::vector{2, 6})); + EXPECT_EQ(ids(s.get_work()), (std::vector{3, 4})) + << "session 20 drained, so session 10 may take two chunks"; + EXPECT_EQ(s.queued(), 0u); +} + +TEST(PrefillTest, LoneSessionFillsTheBatchAcrossPasses) { + Scheduler s{SchedulerParams(1, 4)}; // batch = 9 + s.submit(make_request(1, 10, 4, 0)); + s.submit(make_request(2, 10, 4, 4)); + EXPECT_EQ(ids(s.get_work()), (std::vector{1, 2})); +} + +TEST(PrefillTest, StopsWhenTheNextChunkDoesNotFit) { + Scheduler s{SchedulerParams(1, 4)}; // batch = 9 + s.submit(make_request(1, 10, 4, 0)); + s.submit(make_request(2, 20, 4, 0)); + s.submit(make_request(3, 30, 4, 0)); + + Batch b = s.get_work(); + EXPECT_EQ(b.requests.size(), 2u); + EXPECT_EQ(s.queued(), 1u); +} + +// A session skipped on size was reached; one the pass never got to was not. +// Both got nothing, so they must keep their original relative order. +TEST(PrefillTest, DeferredSessionOutranksOneNeverReached) { + Scheduler s{SchedulerParams(1, 4)}; // batch = 9 + s.submit(make_request(1, 10, 4, 0)); // served, 9 -> 5 + s.submit(make_request(2, 20, 3, 0)); // served, 5 -> 2 + s.submit(make_request(3, 30, 4, 0)); // 4 > 2, deferred + s.submit(make_request(4, 40, 2, 0)); // served, 2 -> 0, pass exits + s.submit(make_request(5, 50, 4, 0)); // never reached + + EXPECT_EQ(ids(s.get_work()), (std::vector{1, 2, 4})); + EXPECT_EQ(ids(s.get_work()), (std::vector{3, 5})); +} + +TEST(PrefillTest, DeferredSessionsKeepTheirOrder) { + Scheduler s{SchedulerParams(1, 4)}; // batch = 9 + s.submit(make_request(1, 10, 4, 0)); + s.submit(make_request(2, 20, 4, 0)); + s.submit(make_request(3, 30, 4, 0)); + s.submit(make_request(4, 40, 3, 0)); + + EXPECT_EQ(ids(s.get_work()), (std::vector{1, 2})); + EXPECT_EQ(ids(s.get_work()), (std::vector{3, 4})); +} + +// Two long prompts must not starve a third: served sessions rotate to the back. +TEST(PrefillTest, RotationIsFairAcrossCalls) { + Scheduler s{SchedulerParams(1, 4)}; + RequestId id = 1; + for (SessionId session : {10, 20, 30}) { + for (int c = 0; c < 8; ++c) { + s.submit(make_request(id++, session, 4, c * 4)); + } + } + + std::map served; + for (int call = 0; call < 9; ++call) { + for (const Request& r : s.get_work().requests) { + served[r.session_id]++; + } + } + EXPECT_EQ(served[10], 6); + EXPECT_EQ(served[20], 6); + EXPECT_EQ(served[30], 6); +} + +// --- completion ------------------------------------------------------------ + +TEST(CompleteTest, SampledStepReturnsTokensAndNoLogits) { + Scheduler s{SchedulerParams(2, 4)}; + auto f = s.submit(make_request(1, 77, 1, 12)); + s.get_work(); + s.complete(1, std::vector{999}); + + ASSERT_SETTLES(f); + Response r = f.get(); + EXPECT_EQ(r.request_id, 1); + EXPECT_EQ(r.session_id, 77); + EXPECT_EQ(tokens_of(r), (std::vector{999})); + EXPECT_TRUE(holds_tokens(r)); +} + +// Greedy verification: one token per drafted position. How many are accepted, +// and where the session continues, is the caller's business. +TEST(CompleteTest, GreedyVerifyReturnsOneTokenPerOutputRow) { + Scheduler s{SchedulerParams(1, 8)}; + auto f = s.submit(make_request(2, 77, 5, 100, OutputRows::All)); + + Batch b = s.get_work(); + ASSERT_EQ(b.requests.size(), 1u); + EXPECT_EQ(b.requests[0].params.output_rows, OutputRows::All); + s.complete(2, std::vector{11, 22, 33, 44, 55}); + + ASSERT_SETTLES(f); + Response r = f.get(); + EXPECT_EQ(tokens_of(r), (std::vector{11, 22, 33, 44, 55})) + << "a verify round must report the prediction at each drafted position, " + "in order"; + EXPECT_TRUE(holds_tokens(r)); +} + +TEST(CompleteTest, UnsampledStepReturnsLogitsAndNoTokens) { + Scheduler s{SchedulerParams(1, 8)}; + auto f = + s.submit(make_request(3, 77, 4, 200, OutputRows::All, /*sample=*/false)); + s.get_work(); + + auto block = std::make_shared(); + block->n_rows = 4; + block->vocab = 3; + block->data = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}; + s.complete(3, LogitsPtr(block)); + + ASSERT_SETTLES(f); + Response r = f.get(); + EXPECT_FALSE(holds_tokens(r)); + LogitsPtr got = logits_of(r); + ASSERT_NE(got, nullptr); + EXPECT_EQ(got->n_rows, 4); + EXPECT_EQ(got->vocab, 3); + ASSERT_EQ(got->data.size(), 12u); + EXPECT_FLOAT_EQ(got->data[11], 11.0f); +} + +TEST(CompleteTest, PendingIsTrueUntilSettled) { + Scheduler s{SchedulerParams(1, 4)}; + auto f = s.submit(make_request(9, 77, 4, 12)); + EXPECT_TRUE(s.pending(9)); + s.get_work(); + EXPECT_TRUE(s.pending(9)) << "in flight still counts as pending"; + s.complete(9, std::vector{1}); + EXPECT_FALSE(s.pending(9)); + EXPECT_FALSE(s.pending(12345)); +} + +TEST(CompleteTest, UnknownIdIsIgnoredByBothOverloads) { + Scheduler s{SchedulerParams(1, 4)}; + s.complete(999, std::vector{1}); + s.complete(999, LogitsPtr{}); + SUCCEED(); +} + +// --- failure --------------------------------------------------------------- + +TEST(FailTest, FailsOneStepAndLeavesOthersRunnable) { + Scheduler s{SchedulerParams(2, 4)}; + auto doomed = s.submit(make_request(1, 10, 1, 0)); + auto ok = s.submit(make_request(2, 10, 1, 1)); + + s.fail(1, "cancelled"); + EXPECT_EQ(s.queued(), 1u); + EXPECT_FALSE(s.pending(1)); + EXPECT_TRUE(s.pending(2)); + ASSERT_SETTLES(doomed); + EXPECT_THROW((void)doomed.get(), std::runtime_error); + + EXPECT_EQ(ids(s.get_work()), (std::vector{2})); + s.complete(2, std::vector{5}); + ASSERT_SETTLES(ok); + EXPECT_EQ(tokens_of(ok.get())[0], 5); +} + +TEST(FailTest, AllCancelledDecodesDrainToAnEmptyBatch) { + Scheduler s{SchedulerParams(2, 4)}; + auto a = s.submit(make_request(1, 10, 1, 0)); + auto b = s.submit(make_request(2, 20, 1, 0)); + s.fail(1, "x"); + s.fail(2, "x"); + + EXPECT_TRUE(s.get_work().empty()); + EXPECT_EQ(s.queued(), 0u); + ASSERT_SETTLES(a); + EXPECT_THROW((void)a.get(), std::runtime_error); + ASSERT_SETTLES(b); + EXPECT_THROW((void)b.get(), std::runtime_error); +} + +TEST(FailTest, CancelledPrefillLeavesTheRotation) { + Scheduler s{SchedulerParams(1, 2)}; + auto doomed = s.submit(make_request(1, 7, 2, 0)); + s.submit(make_request(2, 8, 2, 0)); + + s.fail(1, "cancelled"); + EXPECT_EQ(s.queued(), 1u); + EXPECT_EQ(ids(s.get_work()), (std::vector{2})); + ASSERT_SETTLES(doomed); + EXPECT_THROW((void)doomed.get(), std::runtime_error); +} + +TEST(FailTest, FailAllSettlesEverythingAndEmptiesTheQueues) { + Scheduler s{SchedulerParams(1, 4)}; + auto a = s.submit(make_request(1, 10, 1, 0)); + auto b = s.submit(make_request(2, 20, 4, 0)); + + s.fail_all("shutdown"); + EXPECT_EQ(s.queued(), 0u); + EXPECT_FALSE(s.has_work()); + EXPECT_TRUE(s.get_work().empty()); + ASSERT_SETTLES(a); + EXPECT_THROW((void)a.get(), std::runtime_error); + ASSERT_SETTLES(b); + EXPECT_THROW((void)b.get(), std::runtime_error); +} + +// fail_all() assigns queued_ = 0 rather than decrementing, so the mix cannot +// underflow -- but it must still settle in-flight steps, not just queued ones. +TEST(FailTest, FailAllSettlesQueuedAndInFlightTogether) { + Scheduler s{SchedulerParams(2, 4)}; + auto inflight_a = s.submit(make_request(1, 10, 1, 0)); + auto inflight_b = s.submit(make_request(2, 20, 1, 0)); + ASSERT_EQ(s.get_work().requests.size(), 2u); + ASSERT_EQ(s.queued(), 0u); + + auto queued_a = s.submit(make_request(3, 30, 1, 0)); + auto queued_b = s.submit(make_request(4, 40, 4, 0)); + ASSERT_EQ(s.queued(), 2u); + + s.fail_all("shutdown"); + + EXPECT_EQ(s.queued(), 0u); + EXPECT_FALSE(s.has_work()); + EXPECT_TRUE(s.get_work().empty()); + for (auto* f : {&inflight_a, &inflight_b, &queued_a, &queued_b}) { + ASSERT_SETTLES(*f); + EXPECT_THROW((void)f->get(), std::runtime_error); + } + EXPECT_FALSE(s.pending(1)); + EXPECT_FALSE(s.pending(3)); + + // A completion arriving after the sweep is ignored, not a double-settle. + s.complete(1, std::vector{1}); + SUCCEED(); +} + +TEST(FailTest, UnknownIdIsIgnored) { + Scheduler s{SchedulerParams(1, 4)}; + s.fail(999, "nobody"); + SUCCEED(); +} + +// get_work() already uncounted the step, so fail() must not decrement again. +TEST(FailTest, FailingInFlightStepDoesNotUnderflowQueued) { + Scheduler s{SchedulerParams(2, 4)}; + auto a = s.submit(make_request(1, 10, 1, 0)); + auto b = s.submit(make_request(2, 20, 1, 0)); + ASSERT_EQ(s.queued(), 2u); + s.get_work(); + ASSERT_EQ(s.queued(), 0u); + + s.fail(1, "in-flight fault"); + EXPECT_EQ(s.queued(), 0u); + EXPECT_FALSE(s.has_work()); + ASSERT_SETTLES(a); + EXPECT_THROW((void)a.get(), std::runtime_error); + + s.complete(2, std::vector{5}); + ASSERT_SETTLES(b); + EXPECT_EQ(tokens_of(b.get())[0], 5); +} + +// The same double-decrement would otherwise hide genuinely queued work and +// stall an engine that waits on has_work(). +TEST(FailTest, FailingInFlightStepDoesNotHideQueuedWork) { + Scheduler s{SchedulerParams(1, 4)}; + auto inflight = s.submit(make_request(1, 10, 1, 0)); + s.get_work(); + auto waiting = s.submit(make_request(2, 20, 1, 0)); + ASSERT_EQ(s.queued(), 1u); + + s.fail(1, "in-flight fault"); + EXPECT_EQ(s.queued(), 1u); + EXPECT_TRUE(s.has_work()); + ASSERT_SETTLES(inflight); + EXPECT_THROW((void)inflight.get(), std::runtime_error); + + EXPECT_EQ(ids(s.get_work()), (std::vector{2})); + s.complete(2, std::vector{7}); + ASSERT_SETTLES(waiting); + EXPECT_EQ(tokens_of(waiting.get())[0], 7); +} + +TEST(FailTest, FailingAQueuedStepStillDecrementsOnce) { + Scheduler s{SchedulerParams(2, 4)}; + auto a = s.submit(make_request(1, 10, 1, 0)); + auto b = s.submit(make_request(2, 20, 1, 0)); + ASSERT_EQ(s.queued(), 2u); + + s.fail(1, "cancelled"); + EXPECT_EQ(s.queued(), 1u); + ASSERT_SETTLES(a); + EXPECT_THROW((void)a.get(), std::runtime_error); + EXPECT_EQ(ids(s.get_work()), (std::vector{2})); + EXPECT_EQ(s.queued(), 0u); + s.complete(2, std::vector{1}); + ASSERT_SETTLES(b); + EXPECT_EQ(tokens_of(b.get())[0], 1); +} + +TEST(FailTest, FailingEveryInFlightStepLeavesQueuedCountAtZero) { + Scheduler s{SchedulerParams(4, 4)}; + std::vector> futures; + for (RequestId id = 1; id <= 4; ++id) { + futures.push_back(s.submit(make_request(id, id, 1, 0))); + } + ASSERT_EQ(s.get_work().requests.size(), 4u); + ASSERT_EQ(s.queued(), 0u); + + for (RequestId id = 1; id <= 4; ++id) { + s.fail(id, "fault"); + } + EXPECT_EQ(s.queued(), 0u); + EXPECT_FALSE(s.has_work()); + for (auto& f : futures) { + ASSERT_SETTLES(f); + EXPECT_THROW((void)f.get(), std::runtime_error); + } +} + +// A cancelled entry is only dropped once it reaches the head of its queue, so +// one buried behind live steps has to survive until then without disturbing +// them or the count. +TEST(FailTest, CancelledDecodeInTheMiddleOfTheQueue) { + Scheduler s{SchedulerParams(1, 4)}; + auto a = s.submit(make_request(1, 10, 1, 0)); + auto b = s.submit(make_request(2, 20, 1, 0)); + auto c = s.submit(make_request(3, 30, 1, 0)); + + s.fail(2, "cancelled"); + EXPECT_EQ(s.queued(), 2u); + ASSERT_SETTLES(b); + EXPECT_THROW((void)b.get(), std::runtime_error); + + EXPECT_EQ(ids(s.get_work()), (std::vector{1})); + EXPECT_EQ(ids(s.get_work()), (std::vector{3})); + EXPECT_EQ(s.queued(), 0u); + + s.complete(1, std::vector{1}); + s.complete(3, std::vector{3}); + ASSERT_SETTLES(a); + EXPECT_EQ(tokens_of(a.get())[0], 1); + ASSERT_SETTLES(c); + EXPECT_EQ(tokens_of(c.get())[0], 3); +} + +TEST(FailTest, CancelledPrefillChunkInTheMiddleOfASession) { + Scheduler s{SchedulerParams(1, 4)}; + auto a = s.submit(make_request(1, 7, 4, 0)); + auto b = s.submit(make_request(2, 7, 4, 4)); + auto c = s.submit(make_request(3, 7, 4, 8)); + + s.fail(2, "cancelled"); + EXPECT_EQ(s.queued(), 2u); + ASSERT_SETTLES(b); + EXPECT_THROW((void)b.get(), std::runtime_error); + + // Surviving chunks of the session keep their relative order. + EXPECT_EQ(ids(s.get_work()), (std::vector{1, 3})); + EXPECT_EQ(s.queued(), 0u); +} + +// A session is erased from the rotation when its last chunk is taken; a later +// submit has to put it back. +TEST(PrefillTest, SessionRejoinsTheRotationAfterDraining) { + Scheduler s{SchedulerParams(1, 4)}; + s.submit(make_request(1, 10, 4, 0)); + EXPECT_EQ(ids(s.get_work()), (std::vector{1})); + EXPECT_TRUE(s.get_work().empty()); + + s.submit(make_request(2, 10, 4, 4)); + EXPECT_EQ(ids(s.get_work()), (std::vector{2})); + + // And it interleaves normally with a second session afterwards. + s.submit(make_request(3, 10, 4, 8)); + s.submit(make_request(4, 20, 4, 0)); + EXPECT_EQ(ids(s.get_work()), (std::vector{3, 4})); +} + +TEST(PrefillTest, SessionWhoseOnlyChunkIsCancelledLeavesNoStaleRotationEntry) { + Scheduler s{SchedulerParams(1, 4)}; + auto doomed = s.submit(make_request(1, 10, 4, 0)); + s.fail(1, "cancelled"); + ASSERT_SETTLES(doomed); + EXPECT_THROW((void)doomed.get(), std::runtime_error); + EXPECT_EQ(s.queued(), 0u); + EXPECT_TRUE(s.get_work().empty()); + + // The stale rotation entry, if any, must not swallow the session's next turn. + s.submit(make_request(2, 10, 4, 0)); + s.submit(make_request(3, 20, 4, 0)); + EXPECT_EQ(ids(s.get_work()), (std::vector{2, 3})); +} + +// --- has_work -------------------------------------------------------------- + +TEST(HasWorkTest, TracksQueuedStepsOnly) { + Scheduler s{SchedulerParams(1, 4)}; + EXPECT_FALSE(s.has_work()); + + auto f = s.submit(make_request(1, 10, 1, 0)); + EXPECT_TRUE(s.has_work()); + + s.get_work(); + EXPECT_FALSE(s.has_work()) << "in flight is not queued"; + s.complete(1, std::vector{1}); + EXPECT_FALSE(s.has_work()); + EXPECT_TRUE(settled(f)); +} + +// --- randomized invariants ------------------------------------------------- + +// queued() must equal exactly the steps submitted but not yet taken into a +// batch. An upper bound alone would miss under-counting, which is the failure +// mode that silently hides work from an engine waiting on has_work(). +TEST(InvariantTest, QueuedCountStaysConsistentUnderRandomOps) { + std::mt19937 rng(1234); + Scheduler s{SchedulerParams(3, 4)}; + + std::map> outstanding; + std::vector in_flight; + int submitted = 0, completed = 0, failed = 0; + + // Every outstanding step is either still queued or in flight, so the + // scheduler's count must be exactly the difference. + auto expect_exact_count = [&](int op) { + ASSERT_GE(outstanding.size(), in_flight.size()) << "model broken at " << op; + EXPECT_EQ(s.queued(), outstanding.size() - in_flight.size()) + << "queued() diverged from the model at op " << op; + }; + + for (int i = 0; i < 4000; ++i) { + expect_exact_count(i); + if (::testing::Test::HasFailure()) { + return; + } + + switch (rng() % 4) { + case 0: { // submit + RequestId id = s.next_request_id(); + SessionId session = static_cast(rng() % 4); + int n = static_cast(1 + rng() % 4); + outstanding.emplace( + id, + s.submit(make_request( + id, session, n, static_cast(rng() % 100)))); + submitted++; + break; + } + case 1: { // drain a batch + for (const Request& r : s.get_work().requests) { + in_flight.push_back(r.request_id); + } + break; + } + case 2: { // complete something in flight + if (in_flight.empty()) { + break; + } + std::size_t k = rng() % in_flight.size(); + RequestId id = in_flight[k]; + in_flight.erase(in_flight.begin() + static_cast(k)); + s.complete(id, std::vector{1}); + auto it = outstanding.find(id); + if (it != outstanding.end()) { + ASSERT_SETTLES(it->second); + (void)it->second.get(); + outstanding.erase(it); + completed++; + } + break; + } + case 3: { // fail something, queued or in flight + if (outstanding.empty()) { + break; + } + auto it = outstanding.begin(); + std::advance( + it, static_cast(rng() % outstanding.size())); + RequestId id = it->first; + s.fail(id, "random"); + ASSERT_SETTLES(it->second); + EXPECT_THROW((void)it->second.get(), std::runtime_error); + outstanding.erase(it); + in_flight.erase( + std::remove(in_flight.begin(), in_flight.end(), id), + in_flight.end()); + failed++; + break; + } + default: + break; + } + } + + // Drain: everything still outstanding must settle, and the count must land + // exactly on zero rather than wrapping past it. + while (!outstanding.empty()) { + Batch b = s.get_work(); + if (b.empty()) { + break; + } + for (const Request& r : b.requests) { + s.complete(r.request_id, std::vector{1}); + auto it = outstanding.find(r.request_id); + if (it != outstanding.end()) { + ASSERT_SETTLES(it->second); + (void)it->second.get(); + outstanding.erase(it); + completed++; + } + } + } + + EXPECT_TRUE(outstanding.empty()); + EXPECT_EQ(s.queued(), 0u); + EXPECT_FALSE(s.has_work()); + EXPECT_EQ(completed + failed, submitted); +} + +// --- concurrency ----------------------------------------------------------- + +// Producers submit and await; one engine thread drains and completes. +TEST(ConcurrencyTest, ProducersAndEngineMakeProgressWithoutLoss) { + Scheduler s{SchedulerParams(4, 8)}; + constexpr int kProducers = 4; + constexpr int kPer = 250; + std::atomic stop{false}; + std::atomic completed{0}; + std::atomic lost{0}; + + std::thread engine([&] { + while (!stop.load()) { + Batch b = s.get_work(); + for (const Request& r : b.requests) { + s.complete(r.request_id, std::vector{42}); + } + if (b.empty()) { + std::this_thread::yield(); + } + } + Batch b = s.get_work(); + for (const Request& r : b.requests) { + s.complete(r.request_id, std::vector{42}); + } + }); + + std::vector producers; + for (int t = 0; t < kProducers; ++t) { + producers.emplace_back([&, t] { + for (int i = 0; i < kPer; ++i) { + Request r = make_request( + s.next_request_id(), + (t * 7 + i) % 5, + (i % 3 == 0) ? 1 : 4, + static_cast(i)); + auto f = s.submit(std::move(r)); + // Bounded: a lost request must end this producer rather than block it, + // or the join below would never return and the engine never stop. + if (f.wait_for(kSettleTimeout) != std::future_status::ready) { + lost++; + return; + } + if (tokens_of(f.get()).size() == 1u) { + completed++; + } + } + }); + } + for (std::thread& t : producers) { + t.join(); + } + stop.store(true); + engine.join(); + + EXPECT_EQ(lost.load(), 0) << "a request was never settled"; + EXPECT_EQ(completed.load(), kProducers * kPer); + EXPECT_EQ(s.queued(), 0u); +} + +TEST(ConcurrencyTest, ObserversAreSafeDuringScheduling) { + Scheduler s{SchedulerParams(2, 4)}; + std::atomic stop{false}; + std::atomic observing{false}; + std::atomic observations{0}; + + std::thread observer([&] { + observing.store(true); + while (!stop.load()) { + (void)s.has_work(); + (void)s.queued(); + (void)s.pending(1); + (void)s.params().max_batch_size(); + observations++; + } + }); + // Without this the main loop can finish before the observer starts, leaving + // the test passing vacuously. + while (!observing.load()) { + std::this_thread::yield(); + } + + for (int i = 0; i < 500; ++i) { + auto f = s.submit(make_request(s.next_request_id(), 10, 1, 0)); + Batch b = s.get_work(); + for (const Request& r : b.requests) { + s.complete(r.request_id, std::vector{1}); + } + ASSERT_SETTLES(f); + (void)f.get(); + } + stop.store(true); + observer.join(); + + EXPECT_GT(observations.load(), 0) << "observer never ran"; + EXPECT_EQ(s.queued(), 0u); +} From 2eb059ea19f6cefaf65e2fd8c117b88493422bab Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Fri, 21 Aug 2026 15:13:05 -0700 Subject: [PATCH 2/3] up --- CMakeLists.txt | 4 +-- .../{scheduler => batching}/CMakeLists.txt | 19 ++++++------ .../llm/{scheduler => batching}/scheduler.h | 4 +-- .../test/CMakeLists.txt | 4 +-- .../test/scheduler_test.cpp | 30 +++++++++---------- 5 files changed, 31 insertions(+), 30 deletions(-) rename extension/llm/{scheduler => batching}/CMakeLists.txt (58%) rename extension/llm/{scheduler => batching}/scheduler.h (99%) rename extension/llm/{scheduler => batching}/test/CMakeLists.txt (82%) rename extension/llm/{scheduler => batching}/test/scheduler_test.cpp (97%) diff --git a/CMakeLists.txt b/CMakeLists.txt index 642351637bd..2c36fd20d83 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1001,8 +1001,8 @@ if(EXECUTORCH_BUILD_EXTENSION_LLM) list(APPEND _executorch_extensions tokenizers) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/cache) list(APPEND _executorch_extensions extension_llm_cache) - add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/scheduler) - list(APPEND _executorch_extensions extension_llm_scheduler) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/batching) + list(APPEND _executorch_extensions extension_llm_batching) endif() if(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL) diff --git a/extension/llm/scheduler/CMakeLists.txt b/extension/llm/batching/CMakeLists.txt similarity index 58% rename from extension/llm/scheduler/CMakeLists.txt rename to extension/llm/batching/CMakeLists.txt index 975a9173cea..3b36639973c 100644 --- a/extension/llm/scheduler/CMakeLists.txt +++ b/extension/llm/batching/CMakeLists.txt @@ -4,26 +4,27 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -# Step scheduler for batched LLM serving. Header-only and free of ExecuTorch -# runtime types, so it is an INTERFACE target: consumers get the include path -# and nothing to link. +# Batched LLM serving: the step scheduler, and the runner that drives batches +# through a model. Currently header-only and free of ExecuTorch runtime types, +# so it is an INTERFACE target; this becomes a real library once the runner +# lands with a .cpp. if(NOT EXECUTORCH_ROOT) set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..) endif() -add_library(extension_llm_scheduler INTERFACE) +add_library(extension_llm_batching INTERFACE) # std::optional and std::variant in the public headers. -target_compile_features(extension_llm_scheduler INTERFACE cxx_std_17) +target_compile_features(extension_llm_batching INTERFACE cxx_std_17) target_include_directories( - extension_llm_scheduler INTERFACE ${_common_include_directories} + extension_llm_batching INTERFACE ${_common_include_directories} ) target_compile_options( - extension_llm_scheduler INTERFACE ${_common_compile_options} + extension_llm_batching INTERFACE ${_common_compile_options} ) install( - TARGETS extension_llm_scheduler + TARGETS extension_llm_batching EXPORT ExecuTorchTargets DESTINATION ${CMAKE_INSTALL_LIBDIR} INCLUDES @@ -31,7 +32,7 @@ install( ) install( DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/ - DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/executorch/extension/llm/scheduler + DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/executorch/extension/llm/batching FILES_MATCHING PATTERN "*.h" ) diff --git a/extension/llm/scheduler/scheduler.h b/extension/llm/batching/scheduler.h similarity index 99% rename from extension/llm/scheduler/scheduler.h rename to extension/llm/batching/scheduler.h index 72d4b74e6dd..0fcd60944ec 100644 --- a/extension/llm/scheduler/scheduler.h +++ b/extension/llm/batching/scheduler.h @@ -55,7 +55,7 @@ namespace executorch { namespace extension { namespace llm { -namespace scheduler { +namespace batching { using Token = std::int64_t; using RequestId = std::int64_t; @@ -492,7 +492,7 @@ class Scheduler { std::atomic next_request_id_{1}; }; -} // namespace scheduler +} // namespace batching } // namespace llm } // namespace extension } // namespace executorch diff --git a/extension/llm/scheduler/test/CMakeLists.txt b/extension/llm/batching/test/CMakeLists.txt similarity index 82% rename from extension/llm/scheduler/test/CMakeLists.txt rename to extension/llm/batching/test/CMakeLists.txt index 5a13866da35..603fcf1e341 100644 --- a/extension/llm/scheduler/test/CMakeLists.txt +++ b/extension/llm/batching/test/CMakeLists.txt @@ -13,6 +13,6 @@ include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake) set(_test_srcs scheduler_test.cpp) et_cxx_test( - extension_llm_scheduler_test SOURCES ${_test_srcs} EXTRA_LIBS - extension_llm_scheduler + extension_llm_batching_test SOURCES ${_test_srcs} EXTRA_LIBS + extension_llm_batching ) diff --git a/extension/llm/scheduler/test/scheduler_test.cpp b/extension/llm/batching/test/scheduler_test.cpp similarity index 97% rename from extension/llm/scheduler/test/scheduler_test.cpp rename to extension/llm/batching/test/scheduler_test.cpp index cafea7d7b2d..05e8ec26789 100644 --- a/extension/llm/scheduler/test/scheduler_test.cpp +++ b/extension/llm/batching/test/scheduler_test.cpp @@ -6,7 +6,7 @@ * LICENSE file in the root directory of this source tree. */ -#include +#include #include #include @@ -21,20 +21,20 @@ #include -using executorch::extension::llm::scheduler::Batch; -using executorch::extension::llm::scheduler::LogitsBlock; -using executorch::extension::llm::scheduler::LogitsPtr; -using executorch::extension::llm::scheduler::OutputRows; -using executorch::extension::llm::scheduler::Position; -using executorch::extension::llm::scheduler::Request; -using executorch::extension::llm::scheduler::RequestId; -using executorch::extension::llm::scheduler::RequestParams; -using executorch::extension::llm::scheduler::Response; -using executorch::extension::llm::scheduler::SamplingParams; -using executorch::extension::llm::scheduler::Scheduler; -using executorch::extension::llm::scheduler::SchedulerParams; -using executorch::extension::llm::scheduler::SessionId; -using executorch::extension::llm::scheduler::Token; +using executorch::extension::llm::batching::Batch; +using executorch::extension::llm::batching::LogitsBlock; +using executorch::extension::llm::batching::LogitsPtr; +using executorch::extension::llm::batching::OutputRows; +using executorch::extension::llm::batching::Position; +using executorch::extension::llm::batching::Request; +using executorch::extension::llm::batching::RequestId; +using executorch::extension::llm::batching::RequestParams; +using executorch::extension::llm::batching::Response; +using executorch::extension::llm::batching::SamplingParams; +using executorch::extension::llm::batching::Scheduler; +using executorch::extension::llm::batching::SchedulerParams; +using executorch::extension::llm::batching::SessionId; +using executorch::extension::llm::batching::Token; namespace { From b7131d70f08fba1ffdadd9165215a8c8138f7a16 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Fri, 21 Aug 2026 21:37:46 -0700 Subject: [PATCH 3/3] up --- extension/llm/batching/scheduler.h | 385 ++++------- extension/llm/batching/step.h | 136 ++++ .../llm/batching/test/scheduler_test.cpp | 618 ++++++------------ 3 files changed, 481 insertions(+), 658 deletions(-) create mode 100644 extension/llm/batching/step.h diff --git a/extension/llm/batching/scheduler.h b/extension/llm/batching/scheduler.h index 0fcd60944ec..b3c22acee6b 100644 --- a/extension/llm/batching/scheduler.h +++ b/extension/llm/batching/scheduler.h @@ -8,31 +8,26 @@ #pragma once -// Step scheduler for batched LLM serving. A Request is one step for one -// sequence -- a single decode token, or one prefill chunk -- not a whole -// generation. Callers submit a step, await its sampled token, and submit the -// next one. -// -// get_work() takes decodes first, up to max_decode_sequences, then spends the -// rest of max_batch_size on prefill, so prefill never delays a queued decode. -// Decode is one arrival-order queue. Prefill is a FIFO per session plus a -// rotation over sessions, and a pass takes at most one chunk from each, so a -// long prompt advances a chunk at a time instead of monopolising the batch. -// -// No tensors, cache, or model here: where a step's KV lives belongs to the -// planner. Callers split prefill into chunks of at most max_prefill_chunk_size, -// which needs the session's position sequence and so is caller state. +// Orders steps into batches. Decodes first, up to max_decode_sequences, then +// the rest of max_batch_size on prefill, so prefill never delays a queued +// decode. Decode is one arrival-order queue. Prefill is a FIFO per session +// plus a rotation over sessions, and a pass takes at most one chunk from each, +// so a long prompt advances a chunk at a time instead of monopolising the +// batch. // // The scheduler reads exactly three things from a Request: // // tokens.size() budget arithmetic, and decode vs prefill routing // session_id which prefill FIFO, and a slot in the rotation -// request_id the key its result is settled through +// request_id the key it is tracked and cancelled under // -// Everything else lives in RequestParams / ResponsePayload and is carried from -// submit() -// to Batch and back without being inspected, so payload fields can be added -// without touching any scheduling logic. +// Everything else is carried from submit() to Batch, one way, without being +// inspected -- the executor reads it, not the scheduler -- so payload fields +// can be added without touching any scheduling logic. +// +// Results do not come back through here, and a step stops being tracked the +// moment it is handed to a Batch: whoever drains get_work() owns it from then +// on. cancel() therefore only reaches steps still waiting for a batch. // // Every public method is guarded by mutex_. The Scheduler must outlive every // thread using it; params() hands out a reference into it. @@ -40,119 +35,22 @@ #include #include #include -#include #include #include #include -#include #include #include #include #include -#include #include +#include + namespace executorch { namespace extension { namespace llm { namespace batching { -using Token = std::int64_t; -using RequestId = std::int64_t; -using SessionId = std::int64_t; -using Position = std::int32_t; - -// --- Payload: carried through the scheduler untouched ---------------------- - -// Carried per step, so a caller may change sampling between turns. -struct SamplingParams { - float temperature = 0.0f; - float top_p = 1.0f; - std::int32_t top_k = 0; - std::uint64_t seed = 0; -}; - -// Which positions of a step the model must produce output for. A verify step -// needs every row -- it inspects the prediction at each drafted position -- -// while a decode or prefill chunk only needs the last. A sampled step yields -// one token per row; an unsampled one yields one distribution per row. -enum class OutputRows { - Last, - All, -}; - -// Raw model output for a step whose Request asked not to sample -struct LogitsBlock { - std::vector data; - std::int32_t n_rows = 0; - std::int32_t vocab = 0; -}; -using LogitsPtr = std::shared_ptr; - -// How to execute a step. Promote to a variant if a second step kind ever needs -// different fields rather than different values. -struct RequestParams { - Position position = 0; - // Absent means: do not sample, return the raw distribution. Rejection - // sampling and constrained decoding need the distribution; greedy - // speculative verification does not, since argmax at every row is just - // sampling at temperature zero. - std::optional sampling; - OutputRows output_rows = OutputRows::Last; -}; - -// Which alternative is held follows from whether the Request asked to sample. -// A sampled step yields one token per output row, so a greedy verify round -// returns the prediction at every drafted position; deciding how many to -// accept, and where the session therefore continues, is the caller's job. -using ResponsePayload = std::variant, LogitsPtr>; - -// --- Scheduling ------------------------------------------------------------ - -// One step for one sequence. request_id names the submission and is always -// unique, so the same work submitted twice is two requests that both run and -// both get answered. Only request_id, session_id, and tokens.size() are read -// by the scheduler; `params` is carried. -struct Request { - RequestId request_id = 0; - SessionId session_id = 0; - std::vector tokens; - RequestParams params; - - std::int32_t n_tokens() const { - return static_cast(tokens.size()); - } - - // A one-token prefill is indistinguishable from a decode and is treated as - // one: same single row of output, one decode slot. - bool is_decode() const { - return tokens.size() == 1; - } -}; - -struct Response { - RequestId request_id = 0; - SessionId session_id = 0; - ResponsePayload payload; -}; - -// One batch of steps for a single forward. Decodes first, then prefills, in -// admission order. -struct Batch { - std::vector requests; - - bool empty() const { - return requests.empty(); - } - std::int32_t n_tokens() const { - std::int32_t n = 0; - for (const Request& r : requests) { - n += r.n_tokens(); - } - return n; - } -}; - class SchedulerParams { public: explicit SchedulerParams( @@ -197,17 +95,14 @@ class SchedulerParams { std::int32_t max_prefill_chunk_size_; }; -// Held by shared_ptr because std::promise is move-only. +// Held by shared_ptr so a queue entry and the index name the same step. struct PendingRequest { Request request; - // A deque cannot cheaply erase from the middle, so fail() marks and the - // queues drop on the way past. + // A deque cannot cheaply erase from the middle, so dropping a step marks it + // and the queue skips it on the way past. bool cancelled = false; - // Still waiting in a queue, as opposed to handed to a Batch. Only a queued - // step is counted by queued_, so this is what fail() checks before - // decrementing it. + // False once dispatched. Only a waiting step is counted by queued_. bool queued = true; - std::promise promise; }; using PendingPtr = std::shared_ptr; @@ -223,55 +118,15 @@ class Scheduler { return next_request_id_.fetch_add(1, std::memory_order_relaxed); } - // The future carries the sampled token, or an exception if the step is - // rejected: no tokens, a prefill above max_prefill_chunk_size, or a - // request_id already outstanding. - std::future submit(Request request) { + // Queue a step. False = rejected and nothing changed: no tokens, a prefill + // above max_prefill_chunk_size, or a request_id already waiting. + // + // Results are not delivered here: whoever drains get_work() runs the batch + // and already holds them, so a step carries no return channel of its own. + bool submit(Request request) { auto p = std::make_shared(); p->request = std::move(request); - std::future fut = p->promise.get_future(); - - const std::int32_t n = p->request.n_tokens(); - if (n == 0) { - set_error_(*p, "step carries no tokens"); - return fut; - } - if (!p->request.is_decode() && n > params_.max_prefill_chunk_size()) { - set_error_(*p, "prefill chunk exceeds max_prefill_chunk_size"); - return fut; - } - - std::lock_guard g(mutex_); - // emplace drops a duplicate silently, which would leave this promise never - // settled and the caller's future blocked forever. - auto [slot, inserted] = pending_requests_.emplace(p->request.request_id, p); - if (!inserted) { - set_error_(*p, "request_id already outstanding"); - return fut; - } - - try { - if (p->request.is_decode()) { - decode_queue_.push_back(p); - } else { - const SessionId sid = p->request.session_id; - // Rotation before map: a rotation entry with no session is harmless, - // since get_work() drops it, whereas a session missing from the - // rotation is never visited and its chunks never run. - if (prefill_by_session_.find(sid) == prefill_by_session_.end()) { - prefill_rotation_.push_back(sid); - } - prefill_by_session_[sid].push_back(p); - } - } catch (...) { - // Registered but unqueued would be unschedulable, so its future would - // never settle. Undo and reject instead. - pending_requests_.erase(slot); - set_error_(*p, "failed to queue step"); - return fut; - } - queued_ += 1; - return fut; + return admit_(p); } // Whether get_work() would return a non-empty batch. @@ -293,66 +148,31 @@ class Scheduler { return batch; } - // Settle a sampled step: one token per output row, so a greedy verify round - // reports the target's prediction at every drafted position. An unknown id is - // ignored, so a completion racing fail_all() is not an error. - void complete(RequestId request_id, std::vector tokens) { - Response r; - r.payload = std::move(tokens); - settle_(request_id, std::move(r)); - } - - // Settle an unsampled step with its raw distribution. - void complete(RequestId request_id, LogitsPtr logits) { - Response r; - r.payload = std::move(logits); - settle_(request_id, std::move(r)); - } - - // Fails one step, queued or in flight. Cancelling a whole generation is the - // caller's loop over its outstanding ids. - void fail(RequestId request_id, const std::string& what) { - PendingPtr p; - { - std::lock_guard g(mutex_); - auto it = pending_requests_.find(request_id); - if (it == pending_requests_.end()) { - return; - } - p = it->second; - p->cancelled = true; - // An in-flight step was already uncounted by get_work(); decrementing - // again would underflow queued_ or mask other queued work. - if (p->queued) { - p->queued = false; - queued_ -= 1; - } - pending_requests_.erase(it); - } - set_error_(*p, what); + // Drop a step that is still queued, so it is never handed to a Batch. A step + // already dispatched is not here to drop -- get_work() released it when it + // put it in the batch -- so this is a no-op for one in flight, and for an id + // that never existed. Callers that abandon a step can therefore call it + // without first knowing which state it is in. + void cancel(RequestId request_id) { + std::lock_guard g(mutex_); + (void)take_(request_id); } - // For shutdown. - void fail_all(const std::string& what) { - std::vector doomed; - { - std::lock_guard g(mutex_); - for (auto& entry : pending_requests_) { - entry.second->cancelled = true; - doomed.push_back(entry.second); - } - pending_requests_.clear(); - queued_ = 0; - decode_queue_.clear(); - prefill_by_session_.clear(); - prefill_rotation_.clear(); - } - for (PendingPtr& p : doomed) { - set_error_(*p, what); + // Drop every queued step. For shutdown. + void clear() { + std::lock_guard g(mutex_); + for (auto& entry : pending_requests_) { + entry.second->cancelled = true; } + pending_requests_.clear(); + queued_ = 0; + decode_queue_.clear(); + prefill_by_session_.clear(); + prefill_rotation_.clear(); } - // Whether `request_id` is still queued or in flight. + // True while the step is waiting for a batch. A dispatched step is no longer + // tracked, so this does not mean "submitted and unfinished". bool pending(RequestId request_id) const { std::lock_guard g(mutex_); return pending_requests_.find(request_id) != pending_requests_.end(); @@ -369,32 +189,76 @@ class Scheduler { } private: - // Fills in identity and hands `r` to the waiter. Erasing under the lock - // before settling is what stops complete() and fail() both resolving the - // same promise. - void settle_(RequestId request_id, Response&& r) { - PendingPtr p; - { - std::lock_guard g(mutex_); - auto it = pending_requests_.find(request_id); - if (it == pending_requests_.end()) { - return; + // Validate, register and queue. False = rejected and nothing changed. + bool admit_(const PendingPtr& p) { + const std::int32_t n = p->request.n_tokens(); + if (n == 0) { + return false; + } + if (!p->request.is_decode() && n > params_.max_prefill_chunk_size()) { + return false; + } + + std::lock_guard g(mutex_); + // emplace drops a duplicate silently, which would leave the step + // registered under an id its owner already uses. + auto [slot, inserted] = pending_requests_.emplace(p->request.request_id, p); + if (!inserted) { + return false; + } + + try { + if (p->request.is_decode()) { + decode_queue_.push_back(p); + } else { + const SessionId sid = p->request.session_id; + // Rotation before map: a rotation entry with no session is harmless, + // since get_work() drops it, whereas a session missing from the + // rotation is never visited and its chunks never run. + if (prefill_by_session_.find(sid) == prefill_by_session_.end()) { + prefill_rotation_.push_back(sid); + } + prefill_by_session_[sid].push_back(p); } - p = it->second; - pending_requests_.erase(it); - r.request_id = request_id; - r.session_id = p->request.session_id; + } catch (...) { + // Registered but unqueued would be unschedulable: it would never appear + // in a batch, yet block its id. Undo and reject instead. + pending_requests_.erase(slot); + return false; } - p->promise.set_value(std::move(r)); // outside the lock: wakes the waiter + queued_ += 1; + return true; } - static void set_error_(PendingRequest& p, const std::string& what) { - p.promise.set_exception( - std::make_exception_ptr(std::runtime_error("scheduler: " + what))); + // Hand a step to a Batch: it leaves the queue and pending_requests_ at once, + // since that index only exists to find a step still waiting. Takes the id + // separately because the caller has already moved the request into the + // batch. Caller holds mutex_. + void dispatch_(const PendingPtr& p, RequestId request_id) { + p->queued = false; + queued_ -= 1; + pending_requests_.erase(request_id); + } + + // Remove a waiting step and return it, or null if it is not waiting -- which + // includes one already dispatched. Caller holds mutex_. + PendingPtr take_(RequestId request_id) { + auto it = pending_requests_.find(request_id); + if (it == pending_requests_.end()) { + return nullptr; + } + PendingPtr p = it->second; + // Marked so the queue drops it on the way past, since a deque cannot erase + // from the middle. + p->cancelled = true; + p->queued = false; + queued_ -= 1; + pending_requests_.erase(it); + return p; } - // fail() already settled these and decremented queued_; they are only still - // here because a deque cannot erase from the middle. + // cancel() already dropped these and decremented queued_; they are only + // still here because a deque cannot erase from the middle. static void drop_cancelled_(PendingQueue& q) { while (!q.empty() && q.front()->cancelled) { q.pop_front(); @@ -410,13 +274,17 @@ class Scheduler { if (decode_queue_.empty()) { return; } - // Copy into the batch first: it is the only step here that can throw, and - // popping first would lose the request entirely. - batch.requests.push_back(decode_queue_.front()->request); - decode_queue_.front()->queued = false; + // Move rather than copy: dispatch_ releases the step on the next line, + // so the source is discarded either way, and a prefill chunk carries up + // to max_prefill_chunk_size tokens. Request's move is noexcept, so + // push_back still gives the strong guarantee -- nothing is lost if it + // throws. The id is read first, since the request is gone after. + const PendingPtr p = decode_queue_.front(); + const RequestId id = p->request.request_id; + batch.requests.push_back(std::move(p->request)); + dispatch_(p, id); decode_queue_.pop_front(); budget -= 1; - queued_ -= 1; } } @@ -449,12 +317,13 @@ class Scheduler { deferred.push_back(sid); continue; } - // Copy into the batch before mutating anything else; see take_decodes_. - batch.requests.push_back(dq.front()->request); - dq.front()->queued = false; + // Moved, not copied; see take_decodes_. + const PendingPtr p = dq.front(); + const RequestId id = p->request.request_id; + batch.requests.push_back(std::move(p->request)); + dispatch_(p, id); dq.pop_front(); budget -= n; - queued_ -= 1; progress = true; if (dq.empty()) { prefill_by_session_.erase(it); @@ -485,8 +354,8 @@ class Scheduler { std::unordered_map prefill_by_session_; std::deque prefill_rotation_; - // Index from submit() until complete()/fail(); the queues hold the same - // shared_ptr. + // Every step waiting for a batch, from submit() until it is dispatched or + // cancelled; the queues hold the same shared_ptr. std::unordered_map pending_requests_; std::size_t queued_ = 0; std::atomic next_request_id_{1}; diff --git a/extension/llm/batching/step.h b/extension/llm/batching/step.h new file mode 100644 index 00000000000..9cb171bffd9 --- /dev/null +++ b/extension/llm/batching/step.h @@ -0,0 +1,136 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +// The vocabulary of batched LLM serving: what a step is, how to execute one, +// and what running one produces. Shared by the scheduler that orders steps, +// the executor that runs them, and the runner that drives both. +// +// A Request is one step for one sequence -- a single decode token, or one +// prefill chunk -- not a whole generation. +// +// No tensors, no cache, no model, no locks. Where a step's KV lives is the +// cache's business, and how steps are ordered is the scheduler's. + +#include +#include +#include +#include +#include + +namespace executorch { +namespace extension { +namespace llm { +namespace batching { + +using Token = std::int64_t; +using RequestId = std::int64_t; +using SessionId = std::int64_t; +using Position = std::int32_t; + +// --- How to execute a step ------------------------------------------------- + +// Carried per step, so a caller may change sampling between turns. +// +// Every field is a literal value; there is no "implementation default" +// sentinel. This differs from llm::SamplingConfig in +// extension/llm/runner/llm_session.h, where temperature == -1 means "let the +// implementation choose" and is validated as a distinct legal value. Anything +// bridging the two must resolve that sentinel first: passing -1 through here +// would be read as a negative temperature rather than a request for a default. +struct SamplingParams { + // 0 = greedy. Higher is more random. + float temperature = 0.0f; + float top_p = 1.0f; + std::int32_t top_k = 0; // 0 = disabled + std::uint64_t seed = 0; // 0 = unset +}; + +// Which positions of a step the model must produce output for. A verify step +// needs every row -- it inspects the prediction at each drafted position -- +// while a decode or prefill chunk only needs the last. A sampled step yields +// one token per row; an unsampled one yields one distribution per row. +enum class OutputRows { + Last, + All, +}; + +// How to execute a step. Promote to a variant if a second step kind ever needs +// different fields rather than different values. +struct RequestParams { + Position position = 0; + // Absent means: do not sample, return the raw distribution. Rejection + // sampling and constrained decoding need the distribution; greedy + // speculative verification does not, since argmax at every row is just + // sampling at temperature zero. + std::optional sampling; + OutputRows output_rows = OutputRows::Last; +}; + +// --- What running a step produces ------------------------------------------ + +// Raw model output for a step whose Request asked not to sample. +struct LogitsBlock { + std::vector data; + std::int32_t n_rows = 0; + std::int32_t vocab = 0; +}; +using LogitsPtr = std::shared_ptr; + +// Which alternative the executor produces follows from whether the Request +// asked to sample. A sampled step yields one token per output row, so a greedy +// verify round returns the prediction at every drafted position; deciding how +// many to accept, and where the session therefore continues, belongs to +// whoever ran the batch. This never passes through the scheduler. +using ResponsePayload = std::variant, LogitsPtr>; + +// --- The step itself ------------------------------------------------------- + +// One step for one sequence. request_id names the submission and is always +// unique, so the same work submitted twice is two requests that both run. +// Only request_id, session_id, and tokens.size() are read by the scheduler; +// `params` travels through it to the executor, one way, untouched. +struct Request { + RequestId request_id = 0; + SessionId session_id = 0; + std::vector tokens; + RequestParams params; + + std::int32_t n_tokens() const { + return static_cast(tokens.size()); + } + + // A one-token prefill is indistinguishable from a decode and is treated as + // one: same single row of output, one decode slot. + bool is_decode() const { + return tokens.size() == 1; + } +}; + +// One batch of steps for a single forward. Decodes first, then prefills, in +// admission order. +struct Batch { + std::vector requests; + + bool empty() const { + return requests.empty(); + } + std::int32_t n_tokens() const { + std::int32_t n = 0; + for (const Request& r : requests) { + n += r.n_tokens(); + } + return n; + } +}; + +} // namespace batching +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/batching/test/scheduler_test.cpp b/extension/llm/batching/test/scheduler_test.cpp index 05e8ec26789..c09ad450fce 100644 --- a/extension/llm/batching/test/scheduler_test.cpp +++ b/extension/llm/batching/test/scheduler_test.cpp @@ -29,7 +29,6 @@ using executorch::extension::llm::batching::Position; using executorch::extension::llm::batching::Request; using executorch::extension::llm::batching::RequestId; using executorch::extension::llm::batching::RequestParams; -using executorch::extension::llm::batching::Response; using executorch::extension::llm::batching::SamplingParams; using executorch::extension::llm::batching::Scheduler; using executorch::extension::llm::batching::SchedulerParams; @@ -68,39 +67,9 @@ std::vector ids(const Batch& b) { return out; } -bool settled(std::future& f) { - return f.wait_for(std::chrono::seconds(0)) == std::future_status::ready; -} - -bool holds_tokens(const Response& r) { - return std::holds_alternative>(r.payload); -} - -const std::vector& tokens_of(const Response& r) { - return std::get>(r.payload); -} - -LogitsPtr logits_of(const Response& r) { - return std::get(r.payload); -} - -// Bounded wait. A regression that leaves a future unresolved must fail -// diagnostically rather than hang the suite -- several tests here exist -// precisely to catch non-settlement. -constexpr std::chrono::seconds kSettleTimeout{5}; - -testing::AssertionResult Settles(std::future& f) { - if (f.wait_for(kSettleTimeout) != std::future_status::ready) { - return testing::AssertionFailure() - << "future did not settle within " << kSettleTimeout.count() << "s"; - } - return testing::AssertionSuccess(); -} - } // namespace // Guards every get() whose absence would block instead of failing. -#define ASSERT_SETTLES(f) ASSERT_TRUE(Settles(f)) namespace {} // namespace @@ -190,10 +159,10 @@ TEST(BatchTest, EmptyBatchReportsZeros) { TEST(BatchTest, SumsTokensAcrossSteps) { Scheduler s{SchedulerParams(2, 8)}; - s.submit(make_request(1, 10, 1, 0)); - s.submit(make_request(2, 20, 1, 0)); - s.submit(make_request(3, 30, 6, 0)); - s.submit(make_request(4, 40, 5, 100, OutputRows::All)); + EXPECT_TRUE(s.submit(make_request(1, 10, 1, 0))); + EXPECT_TRUE(s.submit(make_request(2, 20, 1, 0))); + EXPECT_TRUE(s.submit(make_request(3, 30, 6, 0))); + EXPECT_TRUE(s.submit(make_request(4, 40, 5, 100, OutputRows::All))); Batch b = s.get_work(); EXPECT_EQ(ids(b), (std::vector{1, 2, 3, 4})); @@ -202,7 +171,8 @@ TEST(BatchTest, SumsTokensAcrossSteps) { TEST(BatchTest, CarriesPayloadUninspected) { Scheduler s{SchedulerParams(1, 8)}; - s.submit(make_request(1, 40, 5, 100, OutputRows::All, /*sample=*/false)); + EXPECT_TRUE( + s.submit(make_request(1, 40, 5, 100, OutputRows::All, /*sample=*/false))); Batch b = s.get_work(); ASSERT_EQ(b.requests.size(), 1u); @@ -215,48 +185,23 @@ TEST(BatchTest, CarriesPayloadUninspected) { TEST(SubmitTest, RejectsEmptyStep) { Scheduler s{SchedulerParams(1, 4)}; - auto f = s.submit(make_request(1, 10, 0, 0)); - ASSERT_SETTLES(f); - EXPECT_THROW((void)f.get(), std::runtime_error); + EXPECT_FALSE(s.submit(make_request(1, 10, 0, 0))); EXPECT_EQ(s.queued(), 0u); } TEST(SubmitTest, RejectsPrefillAboveChunkSize) { Scheduler s{SchedulerParams(1, 4)}; - auto f = s.submit(make_request(1, 10, 9, 0)); - ASSERT_SETTLES(f); - EXPECT_THROW((void)f.get(), std::runtime_error); + EXPECT_FALSE(s.submit(make_request(1, 10, 9, 0))); EXPECT_EQ(s.queued(), 0u); } -TEST(SubmitTest, RejectsDuplicateRequestIdRatherThanHanging) { +TEST(SubmitTest, RejectsDuplicateRequestId) { Scheduler s{SchedulerParams(2, 4)}; - auto first = s.submit(make_request(7, 10, 1, 0)); - auto dup = s.submit(make_request(7, 20, 1, 0)); + EXPECT_TRUE(s.submit(make_request(7, 10, 1, 0))); + EXPECT_FALSE(s.submit(make_request(7, 20, 1, 0))); - ASSERT_SETTLES(dup); - EXPECT_THROW((void)dup.get(), std::runtime_error); EXPECT_EQ(s.queued(), 1u) << "rejected step must not inflate the count"; EXPECT_EQ(ids(s.get_work()), (std::vector{7})); - s.complete(7, std::vector{42}); - ASSERT_SETTLES(first); - EXPECT_EQ(tokens_of(first.get()), (std::vector{42})); -} - -TEST(SubmitTest, RequestIdIsReusableOnceSettled) { - Scheduler s{SchedulerParams(2, 4)}; - auto first = s.submit(make_request(7, 10, 1, 0)); - s.get_work(); - s.complete(7, std::vector{1}); - ASSERT_SETTLES(first); - EXPECT_EQ(tokens_of(first.get())[0], 1); - EXPECT_FALSE(s.pending(7)); - - auto second = s.submit(make_request(7, 10, 1, 1)); - s.get_work(); - s.complete(7, std::vector{2}); - ASSERT_SETTLES(second); - EXPECT_EQ(tokens_of(second.get())[0], 2); } TEST(SubmitTest, DuplicateWorkWithDistinctIdsBothRun) { @@ -266,26 +211,19 @@ TEST(SubmitTest, DuplicateWorkWithDistinctIdsBothRun) { a.request_id = s.next_request_id(); b.request_id = s.next_request_id(); - auto fa = s.submit(a); - auto fb = s.submit(b); + EXPECT_TRUE(s.submit(a)); + EXPECT_TRUE(s.submit(b)); EXPECT_EQ(s.get_work().requests.size(), 2u); - s.complete(a.request_id, std::vector{111}); - s.complete(b.request_id, std::vector{222}); - ASSERT_SETTLES(fa); - EXPECT_EQ(tokens_of(fa.get())[0], 111); - ASSERT_SETTLES(fb); - EXPECT_EQ(tokens_of(fb.get())[0], 222); + // Dispatch released them, so neither is tracked any more. + EXPECT_FALSE(s.pending(a.request_id)); + EXPECT_FALSE(s.pending(b.request_id)); } TEST(SubmitTest, RejectedStepEntersNoQueue) { Scheduler s{SchedulerParams(1, 4)}; - auto empty = s.submit(make_request(1, 10, 0, 0)); - auto toobig = s.submit(make_request(2, 10, 9, 0)); - ASSERT_SETTLES(empty); - EXPECT_THROW((void)empty.get(), std::runtime_error); - ASSERT_SETTLES(toobig); - EXPECT_THROW((void)toobig.get(), std::runtime_error); + EXPECT_FALSE(s.submit(make_request(1, 10, 0, 0))); + EXPECT_FALSE(s.submit(make_request(2, 10, 9, 0))); EXPECT_EQ(s.queued(), 0u); EXPECT_FALSE(s.has_work()); @@ -331,9 +269,9 @@ TEST(SubmitTest, NextRequestIdIsUniqueAcrossThreads) { TEST(DecodeTest, ServedInArrivalOrderUpToTheCap) { Scheduler s{SchedulerParams(2, 4)}; - s.submit(make_request(1, 10, 1, 0)); - s.submit(make_request(2, 20, 1, 0)); - s.submit(make_request(3, 30, 1, 0)); + EXPECT_TRUE(s.submit(make_request(1, 10, 1, 0))); + EXPECT_TRUE(s.submit(make_request(2, 20, 1, 0))); + EXPECT_TRUE(s.submit(make_request(3, 30, 1, 0))); EXPECT_EQ(ids(s.get_work()), (std::vector{1, 2})); EXPECT_EQ(s.queued(), 1u); @@ -343,21 +281,21 @@ TEST(DecodeTest, ServedInArrivalOrderUpToTheCap) { // A shorter queue must not give a later arrival a head start. TEST(DecodeTest, StaysFifoAcrossDrainAndRefill) { Scheduler s{SchedulerParams(2, 4)}; - s.submit(make_request(1, 10, 1, 0)); - s.submit(make_request(2, 20, 1, 0)); - s.submit(make_request(3, 30, 1, 0)); + EXPECT_TRUE(s.submit(make_request(1, 10, 1, 0))); + EXPECT_TRUE(s.submit(make_request(2, 20, 1, 0))); + EXPECT_TRUE(s.submit(make_request(3, 30, 1, 0))); EXPECT_EQ(ids(s.get_work()), (std::vector{1, 2})); - s.submit(make_request(4, 40, 1, 0)); + EXPECT_TRUE(s.submit(make_request(4, 40, 1, 0))); EXPECT_EQ(ids(s.get_work()), (std::vector{3, 4})); } TEST(DecodeTest, BeatsPrefillAndStillLeavesRoomForAFullChunk) { Scheduler s{SchedulerParams(3, 4)}; // batch = 11 - s.submit(make_request(1, 10, 1, 0)); - s.submit(make_request(2, 20, 1, 0)); - s.submit(make_request(3, 30, 1, 0)); - s.submit(make_request(50, 90, 4, 0)); + EXPECT_TRUE(s.submit(make_request(1, 10, 1, 0))); + EXPECT_TRUE(s.submit(make_request(2, 20, 1, 0))); + EXPECT_TRUE(s.submit(make_request(3, 30, 1, 0))); + EXPECT_TRUE(s.submit(make_request(50, 90, 4, 0))); Batch b = s.get_work(); EXPECT_EQ(ids(b), (std::vector{1, 2, 3, 50})); @@ -368,18 +306,18 @@ TEST(DecodeTest, BeatsPrefillAndStillLeavesRoomForAFullChunk) { TEST(PrefillTest, ChunksOfOneSessionStayInOrder) { Scheduler s{SchedulerParams(1, 2)}; - s.submit(make_request(1, 7, 2, 0)); - s.submit(make_request(2, 7, 2, 2)); + EXPECT_TRUE(s.submit(make_request(1, 7, 2, 0))); + EXPECT_TRUE(s.submit(make_request(2, 7, 2, 2))); EXPECT_EQ(ids(s.get_work()), (std::vector{1, 2})); } TEST(PrefillTest, LongPromptCannotHogTheBatch) { Scheduler s{SchedulerParams(2, 4)}; for (int c = 0; c < 4; ++c) { - s.submit(make_request(1 + c, 10, 4, c * 4)); + EXPECT_TRUE(s.submit(make_request(1 + c, 10, 4, c * 4))); } - s.submit(make_request(5, 20, 4, 0)); - s.submit(make_request(6, 20, 4, 4)); + EXPECT_TRUE(s.submit(make_request(5, 20, 4, 0))); + EXPECT_TRUE(s.submit(make_request(6, 20, 4, 4))); EXPECT_EQ(ids(s.get_work()), (std::vector{1, 5})); EXPECT_EQ(ids(s.get_work()), (std::vector{2, 6})); @@ -390,16 +328,16 @@ TEST(PrefillTest, LongPromptCannotHogTheBatch) { TEST(PrefillTest, LoneSessionFillsTheBatchAcrossPasses) { Scheduler s{SchedulerParams(1, 4)}; // batch = 9 - s.submit(make_request(1, 10, 4, 0)); - s.submit(make_request(2, 10, 4, 4)); + EXPECT_TRUE(s.submit(make_request(1, 10, 4, 0))); + EXPECT_TRUE(s.submit(make_request(2, 10, 4, 4))); EXPECT_EQ(ids(s.get_work()), (std::vector{1, 2})); } TEST(PrefillTest, StopsWhenTheNextChunkDoesNotFit) { Scheduler s{SchedulerParams(1, 4)}; // batch = 9 - s.submit(make_request(1, 10, 4, 0)); - s.submit(make_request(2, 20, 4, 0)); - s.submit(make_request(3, 30, 4, 0)); + EXPECT_TRUE(s.submit(make_request(1, 10, 4, 0))); + EXPECT_TRUE(s.submit(make_request(2, 20, 4, 0))); + EXPECT_TRUE(s.submit(make_request(3, 30, 4, 0))); Batch b = s.get_work(); EXPECT_EQ(b.requests.size(), 2u); @@ -422,10 +360,10 @@ TEST(PrefillTest, DeferredSessionOutranksOneNeverReached) { TEST(PrefillTest, DeferredSessionsKeepTheirOrder) { Scheduler s{SchedulerParams(1, 4)}; // batch = 9 - s.submit(make_request(1, 10, 4, 0)); - s.submit(make_request(2, 20, 4, 0)); - s.submit(make_request(3, 30, 4, 0)); - s.submit(make_request(4, 40, 3, 0)); + EXPECT_TRUE(s.submit(make_request(1, 10, 4, 0))); + EXPECT_TRUE(s.submit(make_request(2, 20, 4, 0))); + EXPECT_TRUE(s.submit(make_request(3, 30, 4, 0))); + EXPECT_TRUE(s.submit(make_request(4, 40, 3, 0))); EXPECT_EQ(ids(s.get_work()), (std::vector{1, 2})); EXPECT_EQ(ids(s.get_work()), (std::vector{3, 4})); @@ -437,7 +375,7 @@ TEST(PrefillTest, RotationIsFairAcrossCalls) { RequestId id = 1; for (SessionId session : {10, 20, 30}) { for (int c = 0; c < 8; ++c) { - s.submit(make_request(id++, session, 4, c * 4)); + EXPECT_TRUE(s.submit(make_request(id++, session, 4, c * 4))); } } @@ -452,329 +390,245 @@ TEST(PrefillTest, RotationIsFairAcrossCalls) { EXPECT_EQ(served[30], 6); } -// --- completion ------------------------------------------------------------ +// --- cancelling a queued step ---------------------------------------------- -TEST(CompleteTest, SampledStepReturnsTokensAndNoLogits) { +TEST(CancelStepTest, DropsAQueuedStepBeforeItRuns) { Scheduler s{SchedulerParams(2, 4)}; - auto f = s.submit(make_request(1, 77, 1, 12)); - s.get_work(); - s.complete(1, std::vector{999}); + EXPECT_TRUE(s.submit(make_request(1, 77, 1, 12))); + EXPECT_TRUE(s.pending(1)); + EXPECT_EQ(s.queued(), 1u); - ASSERT_SETTLES(f); - Response r = f.get(); - EXPECT_EQ(r.request_id, 1); - EXPECT_EQ(r.session_id, 77); - EXPECT_EQ(tokens_of(r), (std::vector{999})); - EXPECT_TRUE(holds_tokens(r)); + s.cancel(1); + EXPECT_FALSE(s.pending(1)); + EXPECT_EQ(s.queued(), 0u); + EXPECT_FALSE(s.has_work()); + EXPECT_TRUE(s.get_work().empty()) << "the queued entry must be dropped"; } -// Greedy verification: one token per drafted position. How many are accepted, -// and where the session continues, is the caller's business. -TEST(CompleteTest, GreedyVerifyReturnsOneTokenPerOutputRow) { - Scheduler s{SchedulerParams(1, 8)}; - auto f = s.submit(make_request(2, 77, 5, 100, OutputRows::All)); - - Batch b = s.get_work(); - ASSERT_EQ(b.requests.size(), 1u); - EXPECT_EQ(b.requests[0].params.output_rows, OutputRows::All); - s.complete(2, std::vector{11, 22, 33, 44, 55}); - - ASSERT_SETTLES(f); - Response r = f.get(); - EXPECT_EQ(tokens_of(r), (std::vector{11, 22, 33, 44, 55})) - << "a verify round must report the prediction at each drafted position, " - "in order"; - EXPECT_TRUE(holds_tokens(r)); +TEST(CancelStepTest, UnknownIdIsIgnored) { + Scheduler s{SchedulerParams(1, 4)}; + s.cancel(404); + EXPECT_EQ(s.queued(), 0u); + SUCCEED(); } -TEST(CompleteTest, UnsampledStepReturnsLogitsAndNoTokens) { - Scheduler s{SchedulerParams(1, 8)}; - auto f = - s.submit(make_request(3, 77, 4, 200, OutputRows::All, /*sample=*/false)); - s.get_work(); +// Dispatch releases a step, so cancelling one already in a batch has nothing +// to do. Callers abandoning a step can call this without knowing which state +// it is in. +TEST(CancelStepTest, DispatchedStepIsNoLongerTracked) { + Scheduler s{SchedulerParams(1, 4)}; + EXPECT_TRUE(s.submit(make_request(9, 77, 4, 12))); + EXPECT_TRUE(s.pending(9)); - auto block = std::make_shared(); - block->n_rows = 4; - block->vocab = 3; - block->data = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}; - s.complete(3, LogitsPtr(block)); + ASSERT_EQ(ids(s.get_work()), (std::vector{9})); + EXPECT_FALSE(s.pending(9)) << "get_work released it"; + EXPECT_EQ(s.queued(), 0u); - ASSERT_SETTLES(f); - Response r = f.get(); - EXPECT_FALSE(holds_tokens(r)); - LogitsPtr got = logits_of(r); - ASSERT_NE(got, nullptr); - EXPECT_EQ(got->n_rows, 4); - EXPECT_EQ(got->vocab, 3); - ASSERT_EQ(got->data.size(), 12u); - EXPECT_FLOAT_EQ(got->data[11], 11.0f); + s.cancel(9); // no-op, and must not disturb the count + EXPECT_EQ(s.queued(), 0u); + EXPECT_FALSE(s.has_work()); } -TEST(CompleteTest, PendingIsTrueUntilSettled) { - Scheduler s{SchedulerParams(1, 4)}; - auto f = s.submit(make_request(9, 77, 4, 12)); - EXPECT_TRUE(s.pending(9)); - s.get_work(); - EXPECT_TRUE(s.pending(9)) << "in flight still counts as pending"; - s.complete(9, std::vector{1}); - EXPECT_FALSE(s.pending(9)); - EXPECT_FALSE(s.pending(12345)); -} +// The id is free for reuse the moment the step is dispatched, since nothing +// tracks it any more. +TEST(CancelStepTest, DispatchFreesTheIdForReuse) { + Scheduler s{SchedulerParams(2, 4)}; + EXPECT_TRUE(s.submit(make_request(7, 10, 1, 0))); + ASSERT_EQ(ids(s.get_work()), (std::vector{7})); -TEST(CompleteTest, UnknownIdIsIgnoredByBothOverloads) { - Scheduler s{SchedulerParams(1, 4)}; - s.complete(999, std::vector{1}); - s.complete(999, LogitsPtr{}); - SUCCEED(); + EXPECT_TRUE(s.submit(make_request(7, 10, 1, 1))); + EXPECT_TRUE(s.pending(7)); + EXPECT_EQ(ids(s.get_work()), (std::vector{7})); } -// --- failure --------------------------------------------------------------- +// --- cancellation +// --------------------------------------------------------------- -TEST(FailTest, FailsOneStepAndLeavesOthersRunnable) { +TEST(CancelTest, CancelsOneStepAndLeavesOthersRunnable) { Scheduler s{SchedulerParams(2, 4)}; - auto doomed = s.submit(make_request(1, 10, 1, 0)); - auto ok = s.submit(make_request(2, 10, 1, 1)); + EXPECT_TRUE(s.submit(make_request(1, 10, 1, 0))); + EXPECT_TRUE(s.submit(make_request(2, 20, 1, 0))); - s.fail(1, "cancelled"); - EXPECT_EQ(s.queued(), 1u); + s.cancel(1); EXPECT_FALSE(s.pending(1)); - EXPECT_TRUE(s.pending(2)); - ASSERT_SETTLES(doomed); - EXPECT_THROW((void)doomed.get(), std::runtime_error); - + EXPECT_EQ(s.queued(), 1u); EXPECT_EQ(ids(s.get_work()), (std::vector{2})); - s.complete(2, std::vector{5}); - ASSERT_SETTLES(ok); - EXPECT_EQ(tokens_of(ok.get())[0], 5); } -TEST(FailTest, AllCancelledDecodesDrainToAnEmptyBatch) { +TEST(CancelTest, AllCancelledDecodesDrainToAnEmptyBatch) { Scheduler s{SchedulerParams(2, 4)}; - auto a = s.submit(make_request(1, 10, 1, 0)); - auto b = s.submit(make_request(2, 20, 1, 0)); - s.fail(1, "x"); - s.fail(2, "x"); + EXPECT_TRUE(s.submit(make_request(1, 10, 1, 0))); + EXPECT_TRUE(s.submit(make_request(2, 20, 1, 0))); + s.cancel(1); + s.cancel(2); - EXPECT_TRUE(s.get_work().empty()); EXPECT_EQ(s.queued(), 0u); - ASSERT_SETTLES(a); - EXPECT_THROW((void)a.get(), std::runtime_error); - ASSERT_SETTLES(b); - EXPECT_THROW((void)b.get(), std::runtime_error); + EXPECT_FALSE(s.has_work()); + EXPECT_TRUE(s.get_work().empty()); } -TEST(FailTest, CancelledPrefillLeavesTheRotation) { - Scheduler s{SchedulerParams(1, 2)}; - auto doomed = s.submit(make_request(1, 7, 2, 0)); - s.submit(make_request(2, 8, 2, 0)); +TEST(CancelTest, CancelledPrefillLeavesTheRotation) { + Scheduler s{SchedulerParams(1, 4)}; + EXPECT_TRUE(s.submit(make_request(1, 10, 4, 0))); + EXPECT_TRUE(s.submit(make_request(2, 20, 4, 0))); + s.cancel(1); - s.fail(1, "cancelled"); EXPECT_EQ(s.queued(), 1u); EXPECT_EQ(ids(s.get_work()), (std::vector{2})); - ASSERT_SETTLES(doomed); - EXPECT_THROW((void)doomed.get(), std::runtime_error); } -TEST(FailTest, FailAllSettlesEverythingAndEmptiesTheQueues) { +TEST(CancelTest, ClearEmptiesEveryQueue) { Scheduler s{SchedulerParams(1, 4)}; - auto a = s.submit(make_request(1, 10, 1, 0)); - auto b = s.submit(make_request(2, 20, 4, 0)); + EXPECT_TRUE(s.submit(make_request(1, 10, 1, 0))); + EXPECT_TRUE(s.submit(make_request(2, 20, 4, 0))); - s.fail_all("shutdown"); + s.clear(); EXPECT_EQ(s.queued(), 0u); EXPECT_FALSE(s.has_work()); EXPECT_TRUE(s.get_work().empty()); - ASSERT_SETTLES(a); - EXPECT_THROW((void)a.get(), std::runtime_error); - ASSERT_SETTLES(b); - EXPECT_THROW((void)b.get(), std::runtime_error); + EXPECT_FALSE(s.pending(1)); + EXPECT_FALSE(s.pending(2)); } -// fail_all() assigns queued_ = 0 rather than decrementing, so the mix cannot -// underflow -- but it must still settle in-flight steps, not just queued ones. -TEST(FailTest, FailAllSettlesQueuedAndInFlightTogether) { +TEST(CancelTest, ClearSweepsQueuedAndInFlightTogether) { Scheduler s{SchedulerParams(2, 4)}; - auto inflight_a = s.submit(make_request(1, 10, 1, 0)); - auto inflight_b = s.submit(make_request(2, 20, 1, 0)); + EXPECT_TRUE(s.submit(make_request(1, 10, 1, 0))); + EXPECT_TRUE(s.submit(make_request(2, 20, 1, 0))); ASSERT_EQ(s.get_work().requests.size(), 2u); ASSERT_EQ(s.queued(), 0u); - auto queued_a = s.submit(make_request(3, 30, 1, 0)); - auto queued_b = s.submit(make_request(4, 40, 4, 0)); + EXPECT_TRUE(s.submit(make_request(3, 30, 1, 0))); + EXPECT_TRUE(s.submit(make_request(4, 40, 4, 0))); ASSERT_EQ(s.queued(), 2u); - s.fail_all("shutdown"); - + s.clear(); EXPECT_EQ(s.queued(), 0u); EXPECT_FALSE(s.has_work()); - EXPECT_TRUE(s.get_work().empty()); - for (auto* f : {&inflight_a, &inflight_b, &queued_a, &queued_b}) { - ASSERT_SETTLES(*f); - EXPECT_THROW((void)f->get(), std::runtime_error); + for (RequestId id = 1; id <= 4; ++id) { + EXPECT_FALSE(s.pending(id)); } - EXPECT_FALSE(s.pending(1)); - EXPECT_FALSE(s.pending(3)); - - // A completion arriving after the sweep is ignored, not a double-settle. - s.complete(1, std::vector{1}); + s.cancel(1); // cancelling after the sweep is ignored, not an error SUCCEED(); } -TEST(FailTest, UnknownIdIsIgnored) { +TEST(CancelTest, UnknownIdIsIgnored) { Scheduler s{SchedulerParams(1, 4)}; - s.fail(999, "nobody"); + s.cancel(404); + EXPECT_EQ(s.queued(), 0u); SUCCEED(); } -// get_work() already uncounted the step, so fail() must not decrement again. -TEST(FailTest, FailingInFlightStepDoesNotUnderflowQueued) { - Scheduler s{SchedulerParams(2, 4)}; - auto a = s.submit(make_request(1, 10, 1, 0)); - auto b = s.submit(make_request(2, 20, 1, 0)); - ASSERT_EQ(s.queued(), 2u); - s.get_work(); +// An in-flight step was already uncounted by get_work(). Decrementing again +// would wrap queued_ and leave has_work() permanently true. +TEST(CancelTest, CancellingInFlightStepDoesNotUnderflowQueued) { + Scheduler s{SchedulerParams(1, 4)}; + EXPECT_TRUE(s.submit(make_request(1, 10, 1, 0))); + ASSERT_EQ(s.get_work().requests.size(), 1u); ASSERT_EQ(s.queued(), 0u); - s.fail(1, "in-flight fault"); + s.cancel(1); EXPECT_EQ(s.queued(), 0u); EXPECT_FALSE(s.has_work()); - ASSERT_SETTLES(a); - EXPECT_THROW((void)a.get(), std::runtime_error); - - s.complete(2, std::vector{5}); - ASSERT_SETTLES(b); - EXPECT_EQ(tokens_of(b.get())[0], 5); } -// The same double-decrement would otherwise hide genuinely queued work and -// stall an engine that waits on has_work(). -TEST(FailTest, FailingInFlightStepDoesNotHideQueuedWork) { +// ...and the other direction: it must not drop a count that belongs to work +// still waiting, or the engine would sleep with a full queue. +TEST(CancelTest, CancellingInFlightStepDoesNotHideQueuedWork) { Scheduler s{SchedulerParams(1, 4)}; - auto inflight = s.submit(make_request(1, 10, 1, 0)); - s.get_work(); - auto waiting = s.submit(make_request(2, 20, 1, 0)); + EXPECT_TRUE(s.submit(make_request(1, 10, 1, 0))); + ASSERT_EQ(s.get_work().requests.size(), 1u); + EXPECT_TRUE(s.submit(make_request(2, 20, 1, 0))); ASSERT_EQ(s.queued(), 1u); - s.fail(1, "in-flight fault"); - EXPECT_EQ(s.queued(), 1u); + s.cancel(1); + EXPECT_EQ(s.queued(), 1u) << "the queued step is still queued"; EXPECT_TRUE(s.has_work()); - ASSERT_SETTLES(inflight); - EXPECT_THROW((void)inflight.get(), std::runtime_error); - EXPECT_EQ(ids(s.get_work()), (std::vector{2})); - s.complete(2, std::vector{7}); - ASSERT_SETTLES(waiting); - EXPECT_EQ(tokens_of(waiting.get())[0], 7); } -TEST(FailTest, FailingAQueuedStepStillDecrementsOnce) { +TEST(CancelTest, CancellingAQueuedStepStillDecrementsOnce) { Scheduler s{SchedulerParams(2, 4)}; - auto a = s.submit(make_request(1, 10, 1, 0)); - auto b = s.submit(make_request(2, 20, 1, 0)); + EXPECT_TRUE(s.submit(make_request(1, 10, 1, 0))); + EXPECT_TRUE(s.submit(make_request(2, 20, 1, 0))); ASSERT_EQ(s.queued(), 2u); - s.fail(1, "cancelled"); + s.cancel(1); + EXPECT_EQ(s.queued(), 1u); + s.cancel(1); // already gone EXPECT_EQ(s.queued(), 1u); - ASSERT_SETTLES(a); - EXPECT_THROW((void)a.get(), std::runtime_error); - EXPECT_EQ(ids(s.get_work()), (std::vector{2})); - EXPECT_EQ(s.queued(), 0u); - s.complete(2, std::vector{1}); - ASSERT_SETTLES(b); - EXPECT_EQ(tokens_of(b.get())[0], 1); } -TEST(FailTest, FailingEveryInFlightStepLeavesQueuedCountAtZero) { +TEST(CancelTest, CancellingEveryInFlightStepLeavesQueuedCountAtZero) { Scheduler s{SchedulerParams(4, 4)}; - std::vector> futures; for (RequestId id = 1; id <= 4; ++id) { - futures.push_back(s.submit(make_request(id, id, 1, 0))); + EXPECT_TRUE(s.submit(make_request(id, id, 1, 0))); } ASSERT_EQ(s.get_work().requests.size(), 4u); ASSERT_EQ(s.queued(), 0u); for (RequestId id = 1; id <= 4; ++id) { - s.fail(id, "fault"); + s.cancel(id); } EXPECT_EQ(s.queued(), 0u); EXPECT_FALSE(s.has_work()); - for (auto& f : futures) { - ASSERT_SETTLES(f); - EXPECT_THROW((void)f.get(), std::runtime_error); - } } // A cancelled entry is only dropped once it reaches the head of its queue, so // one buried behind live steps has to survive until then without disturbing // them or the count. -TEST(FailTest, CancelledDecodeInTheMiddleOfTheQueue) { +TEST(CancelTest, CancelledDecodeInTheMiddleOfTheQueue) { Scheduler s{SchedulerParams(1, 4)}; - auto a = s.submit(make_request(1, 10, 1, 0)); - auto b = s.submit(make_request(2, 20, 1, 0)); - auto c = s.submit(make_request(3, 30, 1, 0)); + EXPECT_TRUE(s.submit(make_request(1, 10, 1, 0))); + EXPECT_TRUE(s.submit(make_request(2, 20, 1, 0))); + EXPECT_TRUE(s.submit(make_request(3, 30, 1, 0))); - s.fail(2, "cancelled"); + s.cancel(2); EXPECT_EQ(s.queued(), 2u); - ASSERT_SETTLES(b); - EXPECT_THROW((void)b.get(), std::runtime_error); EXPECT_EQ(ids(s.get_work()), (std::vector{1})); EXPECT_EQ(ids(s.get_work()), (std::vector{3})); EXPECT_EQ(s.queued(), 0u); - - s.complete(1, std::vector{1}); - s.complete(3, std::vector{3}); - ASSERT_SETTLES(a); - EXPECT_EQ(tokens_of(a.get())[0], 1); - ASSERT_SETTLES(c); - EXPECT_EQ(tokens_of(c.get())[0], 3); } -TEST(FailTest, CancelledPrefillChunkInTheMiddleOfASession) { +TEST(CancelTest, CancelledPrefillChunkInTheMiddleOfASession) { Scheduler s{SchedulerParams(1, 4)}; - auto a = s.submit(make_request(1, 7, 4, 0)); - auto b = s.submit(make_request(2, 7, 4, 4)); - auto c = s.submit(make_request(3, 7, 4, 8)); + EXPECT_TRUE(s.submit(make_request(1, 7, 4, 0))); + EXPECT_TRUE(s.submit(make_request(2, 7, 4, 4))); + EXPECT_TRUE(s.submit(make_request(3, 7, 4, 8))); - s.fail(2, "cancelled"); + s.cancel(2); EXPECT_EQ(s.queued(), 2u); - ASSERT_SETTLES(b); - EXPECT_THROW((void)b.get(), std::runtime_error); // Surviving chunks of the session keep their relative order. EXPECT_EQ(ids(s.get_work()), (std::vector{1, 3})); EXPECT_EQ(s.queued(), 0u); } -// A session is erased from the rotation when its last chunk is taken; a later -// submit has to put it back. TEST(PrefillTest, SessionRejoinsTheRotationAfterDraining) { Scheduler s{SchedulerParams(1, 4)}; - s.submit(make_request(1, 10, 4, 0)); + EXPECT_TRUE(s.submit(make_request(1, 10, 4, 0))); EXPECT_EQ(ids(s.get_work()), (std::vector{1})); EXPECT_TRUE(s.get_work().empty()); - s.submit(make_request(2, 10, 4, 4)); + EXPECT_TRUE(s.submit(make_request(2, 10, 4, 4))); EXPECT_EQ(ids(s.get_work()), (std::vector{2})); // And it interleaves normally with a second session afterwards. - s.submit(make_request(3, 10, 4, 8)); - s.submit(make_request(4, 20, 4, 0)); + EXPECT_TRUE(s.submit(make_request(3, 10, 4, 8))); + EXPECT_TRUE(s.submit(make_request(4, 20, 4, 0))); EXPECT_EQ(ids(s.get_work()), (std::vector{3, 4})); } TEST(PrefillTest, SessionWhoseOnlyChunkIsCancelledLeavesNoStaleRotationEntry) { Scheduler s{SchedulerParams(1, 4)}; - auto doomed = s.submit(make_request(1, 10, 4, 0)); - s.fail(1, "cancelled"); - ASSERT_SETTLES(doomed); - EXPECT_THROW((void)doomed.get(), std::runtime_error); + EXPECT_TRUE(s.submit(make_request(1, 10, 4, 0))); + s.cancel(1); EXPECT_EQ(s.queued(), 0u); EXPECT_TRUE(s.get_work().empty()); // The stale rotation entry, if any, must not swallow the session's next turn. - s.submit(make_request(2, 10, 4, 0)); - s.submit(make_request(3, 20, 4, 0)); + EXPECT_TRUE(s.submit(make_request(2, 10, 4, 0))); + EXPECT_TRUE(s.submit(make_request(3, 20, 4, 0))); EXPECT_EQ(ids(s.get_work()), (std::vector{2, 3})); } @@ -784,14 +638,13 @@ TEST(HasWorkTest, TracksQueuedStepsOnly) { Scheduler s{SchedulerParams(1, 4)}; EXPECT_FALSE(s.has_work()); - auto f = s.submit(make_request(1, 10, 1, 0)); + EXPECT_TRUE(s.submit(make_request(1, 10, 1, 0))); EXPECT_TRUE(s.has_work()); s.get_work(); EXPECT_FALSE(s.has_work()) << "in flight is not queued"; - s.complete(1, std::vector{1}); + s.cancel(1); EXPECT_FALSE(s.has_work()); - EXPECT_TRUE(settled(f)); } // --- randomized invariants ------------------------------------------------- @@ -803,15 +656,13 @@ TEST(InvariantTest, QueuedCountStaysConsistentUnderRandomOps) { std::mt19937 rng(1234); Scheduler s{SchedulerParams(3, 4)}; - std::map> outstanding; - std::vector in_flight; - int submitted = 0, completed = 0, failed = 0; + // The scheduler tracks exactly the steps waiting for a batch: dispatch and + // cancel both release one. + std::set waiting; + int submitted = 0, dispatched = 0, cancelled = 0; - // Every outstanding step is either still queued or in flight, so the - // scheduler's count must be exactly the difference. auto expect_exact_count = [&](int op) { - ASSERT_GE(outstanding.size(), in_flight.size()) << "model broken at " << op; - EXPECT_EQ(s.queued(), outstanding.size() - in_flight.size()) + EXPECT_EQ(s.queued(), waiting.size()) << "queued() diverged from the model at op " << op; }; @@ -821,57 +672,36 @@ TEST(InvariantTest, QueuedCountStaysConsistentUnderRandomOps) { return; } - switch (rng() % 4) { + switch (rng() % 3) { case 0: { // submit - RequestId id = s.next_request_id(); - SessionId session = static_cast(rng() % 4); - int n = static_cast(1 + rng() % 4); - outstanding.emplace( + const RequestId id = s.next_request_id(); + ASSERT_TRUE(s.submit(make_request( id, - s.submit(make_request( - id, session, n, static_cast(rng() % 100)))); + static_cast(rng() % 4), + static_cast(1 + rng() % 4), + static_cast(rng() % 100)))); + waiting.insert(id); submitted++; break; } - case 1: { // drain a batch + case 1: { // dispatch a batch for (const Request& r : s.get_work().requests) { - in_flight.push_back(r.request_id); - } - break; - } - case 2: { // complete something in flight - if (in_flight.empty()) { - break; - } - std::size_t k = rng() % in_flight.size(); - RequestId id = in_flight[k]; - in_flight.erase(in_flight.begin() + static_cast(k)); - s.complete(id, std::vector{1}); - auto it = outstanding.find(id); - if (it != outstanding.end()) { - ASSERT_SETTLES(it->second); - (void)it->second.get(); - outstanding.erase(it); - completed++; + EXPECT_EQ(waiting.erase(r.request_id), 1u) + << "a batch contained a step the model did not have queued"; + EXPECT_FALSE(s.pending(r.request_id)) << "dispatch must release it"; + dispatched++; } break; } - case 3: { // fail something, queued or in flight - if (outstanding.empty()) { + case 2: { // cancel something still waiting + if (waiting.empty()) { break; } - auto it = outstanding.begin(); - std::advance( - it, static_cast(rng() % outstanding.size())); - RequestId id = it->first; - s.fail(id, "random"); - ASSERT_SETTLES(it->second); - EXPECT_THROW((void)it->second.get(), std::runtime_error); - outstanding.erase(it); - in_flight.erase( - std::remove(in_flight.begin(), in_flight.end(), id), - in_flight.end()); - failed++; + auto it = waiting.begin(); + std::advance(it, static_cast(rng() % waiting.size())); + s.cancel(*it); + waiting.erase(it); + cancelled++; break; } default: @@ -879,76 +709,59 @@ TEST(InvariantTest, QueuedCountStaysConsistentUnderRandomOps) { } } - // Drain: everything still outstanding must settle, and the count must land + // Drain: every remaining step must reach a batch, and the count must land // exactly on zero rather than wrapping past it. - while (!outstanding.empty()) { + while (!waiting.empty()) { Batch b = s.get_work(); if (b.empty()) { break; } for (const Request& r : b.requests) { - s.complete(r.request_id, std::vector{1}); - auto it = outstanding.find(r.request_id); - if (it != outstanding.end()) { - ASSERT_SETTLES(it->second); - (void)it->second.get(); - outstanding.erase(it); - completed++; - } + waiting.erase(r.request_id); + dispatched++; } } - EXPECT_TRUE(outstanding.empty()); + EXPECT_TRUE(waiting.empty()); EXPECT_EQ(s.queued(), 0u); EXPECT_FALSE(s.has_work()); - EXPECT_EQ(completed + failed, submitted); + EXPECT_EQ(dispatched + cancelled, submitted); } // --- concurrency ----------------------------------------------------------- -// Producers submit and await; one engine thread drains and completes. TEST(ConcurrencyTest, ProducersAndEngineMakeProgressWithoutLoss) { Scheduler s{SchedulerParams(4, 8)}; constexpr int kProducers = 4; constexpr int kPer = 250; std::atomic stop{false}; - std::atomic completed{0}; - std::atomic lost{0}; + std::atomic accepted{0}; + std::atomic dispatched{0}; std::thread engine([&] { while (!stop.load()) { Batch b = s.get_work(); - for (const Request& r : b.requests) { - s.complete(r.request_id, std::vector{42}); - } + dispatched += static_cast(b.requests.size()); if (b.empty()) { std::this_thread::yield(); } } + // Drain after stop: work submitted between the last check and the store + // would otherwise never be dispatched. Batch b = s.get_work(); - for (const Request& r : b.requests) { - s.complete(r.request_id, std::vector{42}); - } + dispatched += static_cast(b.requests.size()); }); std::vector producers; for (int t = 0; t < kProducers; ++t) { producers.emplace_back([&, t] { for (int i = 0; i < kPer; ++i) { - Request r = make_request( - s.next_request_id(), - (t * 7 + i) % 5, - (i % 3 == 0) ? 1 : 4, - static_cast(i)); - auto f = s.submit(std::move(r)); - // Bounded: a lost request must end this producer rather than block it, - // or the join below would never return and the engine never stop. - if (f.wait_for(kSettleTimeout) != std::future_status::ready) { - lost++; - return; - } - if (tokens_of(f.get()).size() == 1u) { - completed++; + if (s.submit(make_request( + s.next_request_id(), + (t * 7 + i) % 5, + (i % 3 == 0) ? 1 : 4, + static_cast(i)))) { + accepted++; } } }); @@ -956,12 +769,21 @@ TEST(ConcurrencyTest, ProducersAndEngineMakeProgressWithoutLoss) { for (std::thread& t : producers) { t.join(); } + // Let the engine catch up before stopping it, so the count is meaningful. + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (dispatched.load() < accepted.load() && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::yield(); + } stop.store(true); engine.join(); - EXPECT_EQ(lost.load(), 0) << "a request was never settled"; - EXPECT_EQ(completed.load(), kProducers * kPer); + EXPECT_EQ(accepted.load(), kProducers * kPer) << "a submit was rejected"; + EXPECT_EQ(dispatched.load(), accepted.load()) + << "a step was never dispatched"; EXPECT_EQ(s.queued(), 0u); + EXPECT_FALSE(s.has_work()); } TEST(ConcurrencyTest, ObserversAreSafeDuringScheduling) { @@ -987,13 +809,9 @@ TEST(ConcurrencyTest, ObserversAreSafeDuringScheduling) { } for (int i = 0; i < 500; ++i) { - auto f = s.submit(make_request(s.next_request_id(), 10, 1, 0)); - Batch b = s.get_work(); - for (const Request& r : b.requests) { - s.complete(r.request_id, std::vector{1}); - } - ASSERT_SETTLES(f); - (void)f.get(); + const RequestId id = s.next_request_id(); + ASSERT_TRUE(s.submit(make_request(id, 10, 1, 0))); + (void)s.get_work(); // dispatch releases each step } stop.store(true); observer.join();