diff --git a/CMakeLists.txt b/CMakeLists.txt index 460c6a4a041..2c36fd20d83 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/batching) + list(APPEND _executorch_extensions extension_llm_batching) endif() if(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL) diff --git a/extension/llm/batching/CMakeLists.txt b/extension/llm/batching/CMakeLists.txt new file mode 100644 index 00000000000..3b36639973c --- /dev/null +++ b/extension/llm/batching/CMakeLists.txt @@ -0,0 +1,42 @@ +# 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. + +# 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_batching INTERFACE) +# std::optional and std::variant in the public headers. +target_compile_features(extension_llm_batching INTERFACE cxx_std_17) +target_include_directories( + extension_llm_batching INTERFACE ${_common_include_directories} +) +target_compile_options( + extension_llm_batching INTERFACE ${_common_compile_options} +) + +install( + TARGETS extension_llm_batching + EXPORT ExecuTorchTargets + DESTINATION ${CMAKE_INSTALL_LIBDIR} + INCLUDES + DESTINATION ${_common_include_directories} +) +install( + DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/ + DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/executorch/extension/llm/batching + FILES_MATCHING + PATTERN "*.h" +) + +if(BUILD_TESTING) + add_subdirectory(test) +endif() diff --git a/extension/llm/batching/scheduler.h b/extension/llm/batching/scheduler.h new file mode 100644 index 00000000000..b3c22acee6b --- /dev/null +++ b/extension/llm/batching/scheduler.h @@ -0,0 +1,367 @@ +/* + * 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 + +// 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 it is tracked and cancelled under +// +// 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. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace executorch { +namespace extension { +namespace llm { +namespace batching { + +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 so a queue entry and the index name the same step. +struct PendingRequest { + Request request; + // 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; + // False once dispatched. Only a waiting step is counted by queued_. + bool queued = true; +}; + +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); + } + + // 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); + return admit_(p); + } + + // 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; + } + + // 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); + } + + // 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(); + } + + // 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(); + } + + // 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: + // 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); + } + } 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; + } + queued_ += 1; + return true; + } + + // 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; + } + + // 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(); + } + } + + // 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; + } + // 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; + } + } + + // 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; + } + // 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; + 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_; + + // 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}; +}; + +} // namespace batching +} // namespace llm +} // namespace extension +} // namespace executorch 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/CMakeLists.txt b/extension/llm/batching/test/CMakeLists.txt new file mode 100644 index 00000000000..603fcf1e341 --- /dev/null +++ b/extension/llm/batching/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_batching_test SOURCES ${_test_srcs} EXTRA_LIBS + extension_llm_batching +) diff --git a/extension/llm/batching/test/scheduler_test.cpp b/extension/llm/batching/test/scheduler_test.cpp new file mode 100644 index 00000000000..c09ad450fce --- /dev/null +++ b/extension/llm/batching/test/scheduler_test.cpp @@ -0,0 +1,821 @@ +/* + * 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::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::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 { + +// 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; +} + +} // namespace + +// Guards every get() whose absence would block instead of failing. + +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)}; + 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})); + EXPECT_EQ(b.n_tokens(), 1 + 1 + 6 + 5); +} + +TEST(BatchTest, CarriesPayloadUninspected) { + Scheduler s{SchedulerParams(1, 8)}; + 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); + 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)}; + EXPECT_FALSE(s.submit(make_request(1, 10, 0, 0))); + EXPECT_EQ(s.queued(), 0u); +} + +TEST(SubmitTest, RejectsPrefillAboveChunkSize) { + Scheduler s{SchedulerParams(1, 4)}; + EXPECT_FALSE(s.submit(make_request(1, 10, 9, 0))); + EXPECT_EQ(s.queued(), 0u); +} + +TEST(SubmitTest, RejectsDuplicateRequestId) { + Scheduler s{SchedulerParams(2, 4)}; + EXPECT_TRUE(s.submit(make_request(7, 10, 1, 0))); + EXPECT_FALSE(s.submit(make_request(7, 20, 1, 0))); + + EXPECT_EQ(s.queued(), 1u) << "rejected step must not inflate the count"; + EXPECT_EQ(ids(s.get_work()), (std::vector{7})); +} + +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(); + + EXPECT_TRUE(s.submit(a)); + EXPECT_TRUE(s.submit(b)); + EXPECT_EQ(s.get_work().requests.size(), 2u); + + // 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)}; + 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()); + 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)}; + 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); + 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)}; + 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_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 + 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})); + EXPECT_EQ(b.n_tokens(), 3 + 4); +} + +// --- prefill scheduling ---------------------------------------------------- + +TEST(PrefillTest, ChunksOfOneSessionStayInOrder) { + Scheduler s{SchedulerParams(1, 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) { + EXPECT_TRUE(s.submit(make_request(1 + c, 10, 4, c * 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})); + 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 + 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 + 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); + 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 + 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})); +} + +// 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) { + EXPECT_TRUE(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); +} + +// --- cancelling a queued step ---------------------------------------------- + +TEST(CancelStepTest, DropsAQueuedStepBeforeItRuns) { + Scheduler s{SchedulerParams(2, 4)}; + EXPECT_TRUE(s.submit(make_request(1, 77, 1, 12))); + EXPECT_TRUE(s.pending(1)); + EXPECT_EQ(s.queued(), 1u); + + 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"; +} + +TEST(CancelStepTest, UnknownIdIsIgnored) { + Scheduler s{SchedulerParams(1, 4)}; + s.cancel(404); + EXPECT_EQ(s.queued(), 0u); + SUCCEED(); +} + +// 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)); + + ASSERT_EQ(ids(s.get_work()), (std::vector{9})); + EXPECT_FALSE(s.pending(9)) << "get_work released it"; + EXPECT_EQ(s.queued(), 0u); + + s.cancel(9); // no-op, and must not disturb the count + EXPECT_EQ(s.queued(), 0u); + EXPECT_FALSE(s.has_work()); +} + +// 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})); + + EXPECT_TRUE(s.submit(make_request(7, 10, 1, 1))); + EXPECT_TRUE(s.pending(7)); + EXPECT_EQ(ids(s.get_work()), (std::vector{7})); +} + +// --- cancellation +// --------------------------------------------------------------- + +TEST(CancelTest, CancelsOneStepAndLeavesOthersRunnable) { + Scheduler s{SchedulerParams(2, 4)}; + EXPECT_TRUE(s.submit(make_request(1, 10, 1, 0))); + EXPECT_TRUE(s.submit(make_request(2, 20, 1, 0))); + + s.cancel(1); + EXPECT_FALSE(s.pending(1)); + EXPECT_EQ(s.queued(), 1u); + EXPECT_EQ(ids(s.get_work()), (std::vector{2})); +} + +TEST(CancelTest, AllCancelledDecodesDrainToAnEmptyBatch) { + Scheduler s{SchedulerParams(2, 4)}; + 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_EQ(s.queued(), 0u); + EXPECT_FALSE(s.has_work()); + EXPECT_TRUE(s.get_work().empty()); +} + +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); + + EXPECT_EQ(s.queued(), 1u); + EXPECT_EQ(ids(s.get_work()), (std::vector{2})); +} + +TEST(CancelTest, ClearEmptiesEveryQueue) { + Scheduler s{SchedulerParams(1, 4)}; + EXPECT_TRUE(s.submit(make_request(1, 10, 1, 0))); + EXPECT_TRUE(s.submit(make_request(2, 20, 4, 0))); + + s.clear(); + EXPECT_EQ(s.queued(), 0u); + EXPECT_FALSE(s.has_work()); + EXPECT_TRUE(s.get_work().empty()); + EXPECT_FALSE(s.pending(1)); + EXPECT_FALSE(s.pending(2)); +} + +TEST(CancelTest, ClearSweepsQueuedAndInFlightTogether) { + Scheduler s{SchedulerParams(2, 4)}; + 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); + + 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.clear(); + EXPECT_EQ(s.queued(), 0u); + EXPECT_FALSE(s.has_work()); + for (RequestId id = 1; id <= 4; ++id) { + EXPECT_FALSE(s.pending(id)); + } + s.cancel(1); // cancelling after the sweep is ignored, not an error + SUCCEED(); +} + +TEST(CancelTest, UnknownIdIsIgnored) { + Scheduler s{SchedulerParams(1, 4)}; + s.cancel(404); + EXPECT_EQ(s.queued(), 0u); + SUCCEED(); +} + +// 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.cancel(1); + EXPECT_EQ(s.queued(), 0u); + EXPECT_FALSE(s.has_work()); +} + +// ...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)}; + 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.cancel(1); + EXPECT_EQ(s.queued(), 1u) << "the queued step is still queued"; + EXPECT_TRUE(s.has_work()); + EXPECT_EQ(ids(s.get_work()), (std::vector{2})); +} + +TEST(CancelTest, CancellingAQueuedStepStillDecrementsOnce) { + Scheduler s{SchedulerParams(2, 4)}; + 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.cancel(1); + EXPECT_EQ(s.queued(), 1u); + s.cancel(1); // already gone + EXPECT_EQ(s.queued(), 1u); +} + +TEST(CancelTest, CancellingEveryInFlightStepLeavesQueuedCountAtZero) { + Scheduler s{SchedulerParams(4, 4)}; + for (RequestId id = 1; id <= 4; ++id) { + 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.cancel(id); + } + EXPECT_EQ(s.queued(), 0u); + EXPECT_FALSE(s.has_work()); +} + +// 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(CancelTest, CancelledDecodeInTheMiddleOfTheQueue) { + Scheduler s{SchedulerParams(1, 4)}; + 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.cancel(2); + EXPECT_EQ(s.queued(), 2u); + + EXPECT_EQ(ids(s.get_work()), (std::vector{1})); + EXPECT_EQ(ids(s.get_work()), (std::vector{3})); + EXPECT_EQ(s.queued(), 0u); +} + +TEST(CancelTest, CancelledPrefillChunkInTheMiddleOfASession) { + Scheduler s{SchedulerParams(1, 4)}; + 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.cancel(2); + EXPECT_EQ(s.queued(), 2u); + + // Surviving chunks of the session keep their relative order. + EXPECT_EQ(ids(s.get_work()), (std::vector{1, 3})); + EXPECT_EQ(s.queued(), 0u); +} + +TEST(PrefillTest, SessionRejoinsTheRotationAfterDraining) { + Scheduler s{SchedulerParams(1, 4)}; + 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()); + + 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. + 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)}; + 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. + 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})); +} + +// --- has_work -------------------------------------------------------------- + +TEST(HasWorkTest, TracksQueuedStepsOnly) { + Scheduler s{SchedulerParams(1, 4)}; + EXPECT_FALSE(s.has_work()); + + 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.cancel(1); + EXPECT_FALSE(s.has_work()); +} + +// --- 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)}; + + // 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; + + auto expect_exact_count = [&](int op) { + EXPECT_EQ(s.queued(), waiting.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() % 3) { + case 0: { // submit + const RequestId id = s.next_request_id(); + ASSERT_TRUE(s.submit(make_request( + id, + static_cast(rng() % 4), + static_cast(1 + rng() % 4), + static_cast(rng() % 100)))); + waiting.insert(id); + submitted++; + break; + } + case 1: { // dispatch a batch + for (const Request& r : s.get_work().requests) { + 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 2: { // cancel something still waiting + if (waiting.empty()) { + break; + } + auto it = waiting.begin(); + std::advance(it, static_cast(rng() % waiting.size())); + s.cancel(*it); + waiting.erase(it); + cancelled++; + break; + } + default: + break; + } + } + + // Drain: every remaining step must reach a batch, and the count must land + // exactly on zero rather than wrapping past it. + while (!waiting.empty()) { + Batch b = s.get_work(); + if (b.empty()) { + break; + } + for (const Request& r : b.requests) { + waiting.erase(r.request_id); + dispatched++; + } + } + + EXPECT_TRUE(waiting.empty()); + EXPECT_EQ(s.queued(), 0u); + EXPECT_FALSE(s.has_work()); + EXPECT_EQ(dispatched + cancelled, submitted); +} + +// --- concurrency ----------------------------------------------------------- + +TEST(ConcurrencyTest, ProducersAndEngineMakeProgressWithoutLoss) { + Scheduler s{SchedulerParams(4, 8)}; + constexpr int kProducers = 4; + constexpr int kPer = 250; + std::atomic stop{false}; + std::atomic accepted{0}; + std::atomic dispatched{0}; + + std::thread engine([&] { + while (!stop.load()) { + Batch b = s.get_work(); + 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(); + 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) { + if (s.submit(make_request( + s.next_request_id(), + (t * 7 + i) % 5, + (i % 3 == 0) ? 1 : 4, + static_cast(i)))) { + accepted++; + } + } + }); + } + 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(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) { + 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) { + 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(); + + EXPECT_GT(observations.load(), 0) << "observer never ran"; + EXPECT_EQ(s.queued(), 0u); +}