From ce0e06ec4da0da27432d7368c9361a42b35ff63e Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Tue, 29 Sep 2026 11:07:19 -0700 Subject: [PATCH 1/2] [LLM] Add batching-backed serving session lifecycle Give serving clients stable logical session keys while keeping physical execution and cleanup in the existing batching runtime. Reset replaces state under the same key and reports reopening failure explicitly; bounded admission and ordered lifecycle operations make later concurrent generation safe to add separately. Runner owns thread-safe Pending/Ready/Failed initialization state and contains initialization exceptions; serving uses the original executor directly. Test Plan: Serving and batching CMake targets passed. Standalone validation passed 20 lifecycle tests and 164 batching tests, 100 normal lifecycle repetitions, 25 ThreadSanitizer repetitions, and 19 lifecycle tests with exceptions disabled. Initialization exceptions settle as not-ready errors rather than terminating the engine thread. Authored with AI assistance from Devmate. --- extension/llm/batching/executor.h | 3 +- extension/llm/batching/runner.cpp | 25 +- extension/llm/batching/runner.h | 7 + extension/llm/batching/test/runner_test.cpp | 183 +++++- extension/llm/serving/CMakeLists.txt | 24 +- extension/llm/serving/serving_runtime.cpp | 289 +++++++++ extension/llm/serving/serving_runtime.h | 96 +++ extension/llm/serving/test/CMakeLists.txt | 16 + .../llm/serving/test/serving_runtime_test.cpp | 590 ++++++++++++++++++ 9 files changed, 1218 insertions(+), 15 deletions(-) create mode 100644 extension/llm/serving/serving_runtime.cpp create mode 100644 extension/llm/serving/serving_runtime.h create mode 100644 extension/llm/serving/test/CMakeLists.txt create mode 100644 extension/llm/serving/test/serving_runtime_test.cpp diff --git a/extension/llm/batching/executor.h b/extension/llm/batching/executor.h index c8228932ea8..8710a0dc030 100644 --- a/extension/llm/batching/executor.h +++ b/extension/llm/batching/executor.h @@ -72,7 +72,8 @@ class ET_EXPERIMENTAL Executor { // Optional one-time setup, called on the engine thread before any other // method. // - // false = the runner admits no work, so open_session_async() reports nullopt. + // false or an exception = the runner admits no work, so + // open_session_async() reports nullopt. virtual bool initialize() { return true; } diff --git a/extension/llm/batching/runner.cpp b/extension/llm/batching/runner.cpp index 45c04812cfd..dd9738e4c4c 100644 --- a/extension/llm/batching/runner.cpp +++ b/extension/llm/batching/runner.cpp @@ -290,6 +290,10 @@ class RunnerImpl : public std::enable_shared_from_this { GenConfig config, GenerationCallback on_update); + InitializationState initialization_state() const noexcept { + return initialization_state_.load(std::memory_order_acquire); + } + // Engine-thread data, so only stable once that thread is joined. EngineMetrics metrics() const { return metrics_; @@ -510,6 +514,8 @@ class RunnerImpl : public std::enable_shared_from_this { std::condition_variable stopped_cv_; std::vector inbox_; std::atomic lifecycle_{Lifecycle::Running}; + std::atomic initialization_state_{ + InitializationState::Pending}; bool join_in_progress_ = false; std::thread::id engine_thread_id_; @@ -636,6 +642,10 @@ void Runner::shutdown() { impl_->shutdown(); } +InitializationState Runner::initialization_state() const noexcept { + return impl_->initialization_state(); +} + EngineMetrics Runner::metrics() const { return impl_->metrics(); } @@ -711,8 +721,21 @@ void RunnerImpl::run_() { // handed a session for an executor that did not come up, and so one-time // setup is not charged to whichever generation happened to go first. const MetricsTime init_start = MetricsClock::now(); - const bool ready = executor_.initialize(); + bool ready = false; +#if ET_HAS_EXCEPTIONS + try { +#endif + ready = executor_.initialize(); +#if ET_HAS_EXCEPTIONS + } catch (...) { + // Use the same shutdown/drain path as an explicit initialization failure. + } +#endif metrics_.init_us = us_between(init_start, MetricsClock::now()); + // Publish the reason before a stopped Runner can refuse an open. + initialization_state_.store( + ready ? InitializationState::Ready : InitializationState::Failed, + std::memory_order_release); if (!ready) { // Stop without running work. The drain below still answers whatever was // queued while this was starting, so no caller is left waiting. diff --git a/extension/llm/batching/runner.h b/extension/llm/batching/runner.h index 51f6a5dc2d2..d09cc3469bb 100644 --- a/extension/llm/batching/runner.h +++ b/extension/llm/batching/runner.h @@ -236,6 +236,8 @@ class ET_EXPERIMENTAL Session { std::unique_ptr state_; }; +enum class ET_EXPERIMENTAL InitializationState { Pending, Ready, Failed }; + class ET_EXPERIMENTAL Runner { public: // Takes the scheduler, one per runner, which also supplies the prefill chunk @@ -252,6 +254,11 @@ class ET_EXPERIMENTAL Runner { Runner(const Runner&) = delete; Runner& operator=(const Runner&) = delete; + // Any thread. Pending until executor initialization returns; Failed if it + // returns false or throws. The result is retained after shutdown, so Ready + // describes successful initialization, not whether work is still accepted. + InitializationState initialization_state() const noexcept; + // Any thread; queued to the engine thread and acked. // // nullopt = the executor is at capacity, or the runner is shutting down. diff --git a/extension/llm/batching/test/runner_test.cpp b/extension/llm/batching/test/runner_test.cpp index 070f8085e55..3045b4d8163 100644 --- a/extension/llm/batching/test/runner_test.cpp +++ b/extension/llm/batching/test/runner_test.cpp @@ -36,6 +36,7 @@ using executorch::extension::llm::batching::FinishReason; using executorch::extension::llm::batching::GenConfig; using executorch::extension::llm::batching::GenerationHandle; using executorch::extension::llm::batching::GenerationUpdate; +using executorch::extension::llm::batching::InitializationState; using executorch::extension::llm::batching::Position; using executorch::extension::llm::batching::Runner; using executorch::extension::llm::batching::SamplingParams; @@ -2668,19 +2669,189 @@ TEST(InitializeTest, FailureStopsTheRunnerAndRefusesSessions) { << "a session must not be opened on an executor that failed to start"; } -// Opens queued while the engine was starting still have to be answered, or a -// caller blocked on the future would hang. -TEST(InitializeTest, FailureDoesNotStrandAQueuedOpen) { - FakeExecutor executor; - executor.fail_initialize = true; +namespace { + +enum class InitializeOutcome { + Success, + Failure, +#if ET_HAS_EXCEPTIONS + Exception, + NonStandardException, +#endif +}; + +class GatedInitializeExecutor : public FakeExecutor { + public: + explicit GatedInitializeExecutor(InitializeOutcome outcome) + : outcome_(outcome) { + fail_initialize = outcome_ == InitializeOutcome::Failure; + } + + bool initialize() override { + const bool ready = FakeExecutor::initialize(); + { + std::unique_lock lock(mutex_); + entered_ = true; + cv_.notify_all(); + cv_.wait(lock, [this] { return released_; }); + } +#if ET_HAS_EXCEPTIONS + if (outcome_ == InitializeOutcome::Exception) { + throw std::runtime_error("initialize failed"); + } + if (outcome_ == InitializeOutcome::NonStandardException) { + throw 7; + } +#endif + return ready; + } + + bool wait_for_initialize() { + std::unique_lock lock(mutex_); + return cv_.wait_for(lock, kTimeout, [this] { return entered_; }); + } + + void release_initialize() { + { + std::lock_guard lock(mutex_); + released_ = true; + } + cv_.notify_all(); + } + + private: + const InitializeOutcome outcome_; + std::mutex mutex_; + std::condition_variable cv_; + bool entered_ = false; + bool released_ = false; +}; + +class InitializeStateTest : public ::testing::TestWithParam { +}; + +} // namespace + +TEST_P(InitializeStateTest, PendingSettlesQueuedOpensAndRetainsResult) { + GatedInitializeExecutor executor(GetParam()); Fixture fixture(executor); + // Keep assertions nonfatal until the gate is released so teardown can join. + EXPECT_TRUE(executor.wait_for_initialize()); + EXPECT_EQ( + fixture.runner.initialization_state(), InitializationState::Pending); std::vector>> opens; for (int i = 0; i < 4; ++i) { opens.push_back(fixture.runner.open_session_async()); + EXPECT_EQ( + opens.back().wait_for(std::chrono::seconds(0)), + std::future_status::timeout); } + EXPECT_TRUE(executor.opened().empty()); + EXPECT_EQ( + fixture.runner.initialization_state(), InitializationState::Pending); + executor.release_initialize(); + + const bool succeeds = GetParam() == InitializeOutcome::Success; + const auto expected = + succeeds ? InitializationState::Ready : InitializationState::Failed; for (auto& opened : opens) { ASSERT_EQ(opened.wait_for(kTimeout), std::future_status::ready); - EXPECT_FALSE(opened.get().has_value()); + auto session = opened.get(); + EXPECT_EQ(session.has_value(), succeeds); + EXPECT_EQ(fixture.runner.initialization_state(), expected); + if (session) { + EXPECT_TRUE(session->valid()); + } + } + if (!succeeds) { + auto refused = fixture.runner.open_session_async(); + ASSERT_EQ( + refused.wait_for(std::chrono::seconds(0)), std::future_status::ready); + EXPECT_FALSE(refused.get()); + } + + fixture.runner.shutdown(); + EXPECT_EQ(fixture.runner.initialization_state(), expected); + fixture.runner.shutdown(); + EXPECT_EQ(fixture.runner.initialization_state(), expected); + auto refused = fixture.runner.open_session_async(); + ASSERT_EQ( + refused.wait_for(std::chrono::seconds(0)), std::future_status::ready); + EXPECT_FALSE(refused.get()); + EXPECT_EQ(executor.initialize_calls(), 1); + EXPECT_FALSE(executor.called_before_initialize()); + EXPECT_EQ(executor.opened().size(), succeeds ? 4u : 0u); + EXPECT_EQ(executor.closed().size(), executor.opened().size()); + EXPECT_EQ(executor.open_count(), 0); + EXPECT_EQ(fixture.runner.metrics().steps, 0u); +} + +TEST_P( + InitializeStateTest, + ShutdownWhilePendingStillRetainsInitializationResult) { + GatedInitializeExecutor executor(GetParam()); + Fixture fixture(executor); + EXPECT_TRUE(executor.wait_for_initialize()); + std::vector>> opens; + opens.push_back(fixture.runner.open_session_async()); + EXPECT_EQ( + opens.back().wait_for(std::chrono::seconds(0)), + std::future_status::timeout); + + std::thread stopping([&] { fixture.runner.shutdown(); }); + bool stopped_admission = false; + const auto deadline = std::chrono::steady_clock::now() + kTimeout; + while (std::chrono::steady_clock::now() < deadline) { + auto probe = fixture.runner.open_session_async(); + if (probe.wait_for(std::chrono::milliseconds(1)) == + std::future_status::ready) { + stopped_admission = !probe.get(); + break; + } + opens.push_back(std::move(probe)); } + EXPECT_TRUE(stopped_admission); + EXPECT_EQ( + fixture.runner.initialization_state(), InitializationState::Pending); + executor.release_initialize(); + stopping.join(); + + for (auto& opened : opens) { + ASSERT_EQ(opened.wait_for(kTimeout), std::future_status::ready); + EXPECT_FALSE(opened.get()); + } + EXPECT_EQ( + fixture.runner.initialization_state(), + GetParam() == InitializeOutcome::Success ? InitializationState::Ready + : InitializationState::Failed); + EXPECT_EQ(executor.initialize_calls(), 1); + EXPECT_TRUE(executor.opened().empty()); + EXPECT_EQ(fixture.runner.metrics().steps, 0u); } + +INSTANTIATE_TEST_SUITE_P( + InitializeTest, + InitializeStateTest, + ::testing::Values( + InitializeOutcome::Success, +#if ET_HAS_EXCEPTIONS + InitializeOutcome::Exception, + InitializeOutcome::NonStandardException, +#endif + InitializeOutcome::Failure), + [](const ::testing::TestParamInfo& info) { + switch (info.param) { + case InitializeOutcome::Success: + return "Success"; + case InitializeOutcome::Failure: + return "Failure"; +#if ET_HAS_EXCEPTIONS + case InitializeOutcome::Exception: + return "Exception"; + case InitializeOutcome::NonStandardException: + return "NonStandardException"; +#endif + } + return "Unknown"; + }); diff --git a/extension/llm/serving/CMakeLists.txt b/extension/llm/serving/CMakeLists.txt index b54b56a7588..7a60c9ff932 100644 --- a/extension/llm/serving/CMakeLists.txt +++ b/extension/llm/serving/CMakeLists.txt @@ -4,16 +4,22 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -add_library(extension_llm_serving INTERFACE) -target_compile_features(extension_llm_serving INTERFACE cxx_std_17) +add_library(extension_llm_serving STATIC serving_runtime.cpp) +target_compile_features(extension_llm_serving PUBLIC cxx_std_17) target_include_directories( - extension_llm_serving INTERFACE ${_common_include_directories} + extension_llm_serving PUBLIC ${_common_include_directories} +) +target_compile_options(extension_llm_serving PUBLIC ${_common_compile_options}) +target_link_libraries( + extension_llm_serving PUBLIC extension_llm_batching executorch_core ) -target_link_libraries(extension_llm_serving INTERFACE executorch_core) - -install(TARGETS extension_llm_serving EXPORT ExecuTorchTargets) -install(FILES types.h prompt_history.h +install( + TARGETS extension_llm_serving + EXPORT ExecuTorchTargets + DESTINATION ${CMAKE_INSTALL_LIBDIR} +) +install(FILES types.h prompt_history.h serving_runtime.h DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/executorch/extension/llm/serving ) @@ -26,3 +32,7 @@ if(NOT EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER) DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/executorch/extension/llm/runner ) endif() + +if(BUILD_TESTING) + add_subdirectory(test) +endif() diff --git a/extension/llm/serving/serving_runtime.cpp b/extension/llm/serving/serving_runtime.cpp new file mode 100644 index 00000000000..4d00eabf992 --- /dev/null +++ b/extension/llm/serving/serving_runtime.cpp @@ -0,0 +1,289 @@ +/* + * 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 + +namespace executorch { +namespace extension { +namespace llm { +namespace serving { + +struct ServingRuntime::Impl { + enum class Operation { Open, Close, Reset }; + enum class Lifecycle { Running, Stopping, Stopped }; + enum class SessionPhase { Opening, Open, Reopening, Unavailable }; + + struct Slot { + SessionPhase phase = SessionPhase::Opening; + std::optional session; + }; + + struct Command { + Operation operation; + std::string key; + std::promise completion; + }; + + Impl( + batching::Executor& executor, + std::unique_ptr scheduler, + ServingRuntimeConfig config) + : config_(config), + valid_config_( + scheduler && config.max_sessions > 0 && + config.max_pending_operations > 0) { + if (valid_config_) { + runner_ = + std::make_unique(executor, std::move(scheduler)); + control_ = std::thread([this] { run(); }); + } + } + + std::future submit(Operation operation, std::string key) { + Command command{operation, std::move(key), {}}; + auto future = command.completion.get_future(); + { + std::lock_guard lock(mutex_); + LifecycleResult error; + if (!valid_config_) { + error = + ServingError{ErrorCode::InvalidArgument, "invalid runtime config"}; + } else if ( + lifecycle_ != Lifecycle::Running || + runner_->initialization_state() == + batching::InitializationState::Failed) { + error = ServingError{ErrorCode::NotReady, "runtime is not ready"}; + } else if (command.key.empty()) { + error = + ServingError{ErrorCode::InvalidArgument, "session key is empty"}; + } else if (outstanding_ == config_.max_pending_operations) { + error = ServingError{ + ErrorCode::CapacityExceeded, "lifecycle operation limit reached"}; + } + if (error) { + command.completion.set_value(std::move(error)); + return future; + } + inbox_.push_back(std::move(command)); + ++outstanding_; + } + cv_.notify_one(); + return future; + } + + LifecycleResult process(Command& command) { + std::optional retired; + std::unique_lock lock(mutex_); + if (lifecycle_ != Lifecycle::Running) { + return ServingError{ErrorCode::NotReady, "runtime is stopping"}; + } + auto it = sessions_.find(command.key); + if (command.operation == Operation::Close) { + if (it != sessions_.end()) { + retired = std::move(it->second.session); + sessions_.erase(it); + } + lock.unlock(); + retired.reset(); + return std::nullopt; + } + const bool opening = command.operation == Operation::Open; + if (opening && it != sessions_.end()) { + if (it->second.phase == SessionPhase::Open) { + return std::nullopt; + } + return ServingError{ + ErrorCode::NotReady, + "session is unavailable; reset to retry or close to release its slot"}; + } + if (opening) { + if (sessions_.size() == config_.max_sessions) { + return ServingError{ + ErrorCode::CapacityExceeded, "session limit reached"}; + } + it = sessions_.try_emplace(command.key).first; + } else { + if (it == sessions_.end()) { + return ServingError{ + ErrorCode::SessionNotFound, "session key not found"}; + } + it->second.phase = SessionPhase::Reopening; + retired = std::move(it->second.session); + it->second.session.reset(); + } + + // Only this thread changes registry entries. Keep the slot reserved while + // waiting, but never hold the admission mutex across Runner calls or waits. + lock.unlock(); + retired.reset(); + std::optional opened; + LifecycleResult error; +#if ET_HAS_EXCEPTIONS + try { +#endif + opened = runner_->open_session_async().get(); +#if ET_HAS_EXCEPTIONS + } catch (...) { + error = ServingError{ErrorCode::Internal, "session open failed"}; + } +#endif + lock.lock(); + if (lifecycle_ != Lifecycle::Running) { + lock.unlock(); + opened.reset(); + return ServingError{ErrorCode::NotReady, "runtime is stopping"}; + } + if (opened) { + it->second.session = std::move(opened); + it->second.phase = SessionPhase::Open; + return std::nullopt; + } + if (opening) { + sessions_.erase(it); + } else { + it->second.phase = SessionPhase::Unavailable; + } + const auto code = error ? error->code + : runner_->initialization_state() == + batching::InitializationState::Failed + ? ErrorCode::NotReady + : ErrorCode::CapacityExceeded; + return ServingError{ + code, + opening + ? "session open failed" + : "session reopen failed; old state is gone and key is unavailable"}; + } + + void run() { + for (;;) { + Command command; + { + std::unique_lock lock(mutex_); + cv_.wait(lock, [this] { + return lifecycle_ != Lifecycle::Running || !inbox_.empty(); + }); + if (inbox_.empty()) { + break; + } + command = std::move(inbox_.front()); + inbox_.pop_front(); + } + auto result = process(command); + { + std::lock_guard lock(mutex_); + if (lifecycle_ != Lifecycle::Running) { + result = ServingError{ErrorCode::NotReady, "runtime is stopping"}; + } + --outstanding_; + command.completion.set_value(std::move(result)); + } + } + decltype(sessions_) retired; + { + std::lock_guard lock(mutex_); + retired.swap(sessions_); + } + // Session destruction may enqueue Runner commands; keep it outside mutex_. + retired.clear(); + } + + ServingInfo info() const { + std::lock_guard lock(mutex_); + return { + valid_config_ && lifecycle_ == Lifecycle::Running && + runner_->initialization_state() == + batching::InitializationState::Ready, + config_.max_context_length, + config_.max_sessions, + sessions_.size()}; + } + + void shutdown() { + { + std::unique_lock lock(mutex_); + if (lifecycle_ != Lifecycle::Running) { + stopped_cv_.wait( + lock, [this] { return lifecycle_ == Lifecycle::Stopped; }); + return; + } + lifecycle_ = Lifecycle::Stopping; + } + cv_.notify_one(); + // Stop Runner before joining control: control may be awaiting an open. + // Runner settles that future and closes even an unpublished late session. + if (runner_) { + runner_->shutdown(); + control_.join(); + } + { + std::lock_guard lock(mutex_); + lifecycle_ = Lifecycle::Stopped; + } + stopped_cv_.notify_all(); + } + + const ServingRuntimeConfig config_; + const bool valid_config_; + std::unique_ptr runner_; + mutable std::mutex mutex_; + std::condition_variable cv_; + std::condition_variable stopped_cv_; + Lifecycle lifecycle_ = Lifecycle::Running; + std::size_t outstanding_ = 0; + std::deque inbox_; + std::unordered_map sessions_; + std::thread control_; +}; + +ServingRuntime::ServingRuntime( + batching::Executor& executor, + std::unique_ptr scheduler, + ServingRuntimeConfig config) + : impl_(std::make_unique(executor, std::move(scheduler), config)) {} + +ServingRuntime::~ServingRuntime() { + shutdown(); +} + +std::future ServingRuntime::open_session_async( + std::string key) { + return impl_->submit(Impl::Operation::Open, std::move(key)); +} + +std::future ServingRuntime::close_session_async( + std::string key) { + return impl_->submit(Impl::Operation::Close, std::move(key)); +} + +std::future ServingRuntime::reset_session_async( + std::string key) { + return impl_->submit(Impl::Operation::Reset, std::move(key)); +} + +ServingInfo ServingRuntime::info() const { + return impl_->info(); +} + +void ServingRuntime::shutdown() { + impl_->shutdown(); +} + +} // namespace serving +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/serving/serving_runtime.h b/extension/llm/serving/serving_runtime.h new file mode 100644 index 00000000000..1397ea5d413 --- /dev/null +++ b/extension/llm/serving/serving_runtime.h @@ -0,0 +1,96 @@ +/* + * 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 + +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace executorch { +namespace extension { +namespace llm { +namespace serving { + +struct ET_EXPERIMENTAL ServingRuntimeConfig { + std::size_t max_sessions = 1; + // Reported metadata; 0 means unknown. + std::size_t max_context_length = 0; + // Queued plus executing lifecycle operations. Must be non-zero. + std::size_t max_pending_operations = 64; +}; + +// nullopt acknowledges success. Dropping a future does not cancel its +// operation. +using LifecycleResult ET_EXPERIMENTAL = std::optional; + +class ET_EXPERIMENTAL ServingRuntime { + public: + // Owns one Runner and its scheduler. The borrowed executor must outlive + // shutdown/destruction and must not be driven by another Runner concurrently. + // Initialization happens on the engine thread. A null scheduler or zero + // session/operation limit leaves an inert runtime: info().ready is false and + // lifecycle operations return InvalidArgument, without starting threads. + ServingRuntime( + batching::Executor& executor, + std::unique_ptr scheduler, + ServingRuntimeConfig config); + ~ServingRuntime(); + + ServingRuntime(const ServingRuntime&) = delete; + ServingRuntime& operator=(const ServingRuntime&) = delete; + ServingRuntime(ServingRuntime&&) = delete; + ServingRuntime& operator=(ServingRuntime&&) = delete; + + // Any thread. Keys must be non-empty. Accepted operations run in admission + // order, including operations submitted before a previous future completes. + // Queue saturation returns CapacityExceeded without changing session state. + // Opens may queue during initialization; success means an empty session is + // owned, or the key was already open. An initial open failure releases its + // slot. Opening an unavailable key fails with NotReady; use reset to retry. + std::future open_session_async(std::string key); + + // Idempotent, including absent keys. Success releases the logical slot and + // requests Session closure through RAII, not synchronous physical cleanup. + std::future close_session_async(std::string key); + + // Cold replacement, retaining the key and logical slot throughout. Success + // means reopening completed. Failure destroys the old state and leaves the + // key unavailable, still reserving its slot until reset succeeds or close + // releases it. A missing key returns SessionNotFound. Never uses prefix + // reuse. + std::future reset_session_async(std::string key); + + // A synchronized snapshot. Ready means initialized and accepting operations, + // not that capacity is available. Active slots include opening, reopening, + // and unavailable keys, but not commands still waiting to reserve a slot. + ServingInfo info() const; + + // Idempotent, including concurrent callers. Stops admission, settles all + // accepted futures, and joins the control and Runner threads. No new work is + // accepted afterward. Call only from external threads, never engine + // callbacks. Destruction calls shutdown; as usual, object lifetime must be + // externally synchronized against callers still accessing the runtime. + void shutdown(); + + private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace serving +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/serving/test/CMakeLists.txt b/extension/llm/serving/test/CMakeLists.txt new file mode 100644 index 00000000000..0e240ef4579 --- /dev/null +++ b/extension/llm/serving/test/CMakeLists.txt @@ -0,0 +1,16 @@ +# 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) + +et_cxx_test( + extension_llm_serving_test SOURCES serving_runtime_test.cpp EXTRA_LIBS + extension_llm_serving +) diff --git a/extension/llm/serving/test/serving_runtime_test.cpp b/extension/llm/serving/test/serving_runtime_test.cpp new file mode 100644 index 00000000000..29c927a010b --- /dev/null +++ b/extension/llm/serving/test/serving_runtime_test.cpp @@ -0,0 +1,590 @@ +/* + * 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 + +#include +#include + +#include + +using executorch::extension::llm::batching::DecodeFirstScheduler; +using executorch::extension::llm::batching::Position; +using executorch::extension::llm::batching::SessionId; +using executorch::extension::llm::batching::testing::FakeExecutor; +using executorch::extension::llm::serving::ErrorCode; +using executorch::extension::llm::serving::LifecycleResult; +using executorch::extension::llm::serving::ServingError; +using executorch::extension::llm::serving::ServingRuntime; +using executorch::extension::llm::serving::ServingRuntimeConfig; + +namespace { + +constexpr std::chrono::seconds kTimeout{5}; + +static_assert(!std::is_copy_constructible::value); +static_assert(!std::is_move_constructible::value); + +class Gate { + public: + void hold() { + std::lock_guard lock(mutex_); + held_ = true; + } + + void arrive() { + std::unique_lock lock(mutex_); + ++arrivals_; + cv_.notify_all(); + cv_.wait(lock, [this] { return !held_; }); + } + + bool wait_for(std::size_t arrivals) { + std::unique_lock lock(mutex_); + return cv_.wait_for(lock, kTimeout, [&] { return arrivals_ >= arrivals; }); + } + + void release() { + { + std::lock_guard lock(mutex_); + held_ = false; + } + cv_.notify_all(); + } + + private: + std::mutex mutex_; + std::condition_variable cv_; + std::size_t arrivals_ = 0; + bool held_ = false; +}; + +class LifecycleExecutor : public FakeExecutor { + public: + bool initialize() override { + note_thread(); + initializing.arrive(); +#if ET_HAS_EXCEPTIONS + if (throw_initialize) { + throw std::runtime_error("initialization failed"); + } +#endif + return FakeExecutor::initialize(); + } + + std::optional open_session() override { + note_thread(); + opening.arrive(); + ++open_attempts; + if (refuse_next_open.exchange(false)) { + return std::nullopt; + } + auto session = FakeExecutor::open_session(); + opened_session.arrive(); + return session; + } + + void close_session(SessionId session) override { + note_thread(); + closing.arrive(); + FakeExecutor::close_session(session); + } + + std::optional clone(SessionId, Position) override { + ++clone_calls; + return std::nullopt; + } + + void release_all() { + initializing.release(); + opening.release(); + opened_session.release(); + closing.release(); + } + + std::vector threads() const { + std::lock_guard lock(mutex_); + return threads_; + } + + Gate initializing; + Gate opening; + Gate opened_session; + Gate closing; + bool throw_initialize = false; + std::atomic refuse_next_open{false}; + std::atomic open_attempts{0}; + std::atomic clone_calls{0}; + + private: + void note_thread() { + std::lock_guard lock(mutex_); + threads_.push_back(std::this_thread::get_id()); + } + + mutable std::mutex mutex_; + std::vector threads_; +}; + +class ServingRuntimeTest : public ::testing::Test { + protected: + void start(ServingRuntimeConfig config = {2, 128, 16}) { + runtime = std::make_unique( + executor, DecodeFirstScheduler::create(), config); + } + + void TearDown() override { + // Release gates even after a fatal assertion, before the runtime joins. + executor.release_all(); + runtime.reset(); + EXPECT_EQ(executor.clone_calls.load(), 0); + EXPECT_TRUE(executor.seen().empty()); + } + + LifecycleExecutor executor; + std::unique_ptr runtime; +}; + +struct JoiningThreads { + explicit JoiningThreads(LifecycleExecutor& executor) : executor(executor) {} + + ~JoiningThreads() { + executor.release_all(); + for (auto& thread : threads) { + if (thread.joinable()) { + thread.join(); + } + } + } + + LifecycleExecutor& executor; + std::vector threads; +}; + +LifecycleResult result(std::future future) { + if (future.wait_for(kTimeout) != std::future_status::ready) { + ADD_FAILURE() << "lifecycle operation did not settle"; + return ServingError{ErrorCode::Internal, "test timeout"}; + } + return future.get(); +} + +void expect_error(std::future future, ErrorCode code) { + auto error = result(std::move(future)); + ASSERT_TRUE(error); + EXPECT_EQ(error->code, code); + EXPECT_FALSE(error->message.empty()); +} + +template +bool wait_until(Predicate predicate) { + const auto deadline = std::chrono::steady_clock::now() + kTimeout; + while (!predicate()) { + if (std::chrono::steady_clock::now() >= deadline) { + return false; + } + std::this_thread::yield(); + } + return true; +} + +TEST_F(ServingRuntimeTest, OpenQueuesDuringInitializationAndReportsReadiness) { + executor.initializing.hold(); + executor.opening.hold(); + start({1, 4096, 4}); + ASSERT_TRUE(executor.initializing.wait_for(1)); + const auto info = runtime->info(); + EXPECT_FALSE(info.ready); + EXPECT_EQ(info.max_sessions, 1u); + EXPECT_EQ(info.max_context_length, 4096u); + EXPECT_EQ(info.active_sessions, 0u); + + auto opened = runtime->open_session_async("session"); + EXPECT_EQ( + opened.wait_for(std::chrono::seconds(0)), std::future_status::timeout); + executor.initializing.release(); + ASSERT_TRUE(executor.opening.wait_for(1)); + EXPECT_TRUE(runtime->info().ready); + EXPECT_EQ(runtime->info().active_sessions, 1u); + EXPECT_EQ( + opened.wait_for(std::chrono::seconds(0)), std::future_status::timeout); + executor.opening.release(); + EXPECT_FALSE(result(std::move(opened))); + EXPECT_TRUE(runtime->info().ready); // Full capacity is still initialized. + EXPECT_EQ(executor.initialize_calls(), 1); + EXPECT_FALSE(executor.called_before_initialize()); +} + +TEST_F(ServingRuntimeTest, DuplicateOpenAndCloseAreIdempotent) { + executor.opening.hold(); + start(); + auto first = runtime->open_session_async("session"); + ASSERT_TRUE(executor.opening.wait_for(1)); + auto duplicate = runtime->open_session_async("session"); + executor.opening.release(); + EXPECT_FALSE(result(std::move(first))); + EXPECT_FALSE(result(std::move(duplicate))); + EXPECT_FALSE(result(runtime->open_session_async("session"))); + EXPECT_EQ(executor.open_attempts.load(), 1); + EXPECT_EQ(runtime->info().active_sessions, 1u); + EXPECT_FALSE(result(runtime->close_session_async("session"))); + EXPECT_FALSE(result(runtime->close_session_async("session"))); + runtime->shutdown(); + EXPECT_EQ(executor.closed(), executor.opened()); +} + +TEST_F(ServingRuntimeTest, SameKeyOperationsPreserveAdmissionOrder) { + executor.opening.hold(); + executor.capacity = 1; + start({1, 128, 8}); + auto opened = runtime->open_session_async("session"); + ASSERT_TRUE(executor.opening.wait_for(1)); + auto reset = runtime->reset_session_async("session"); + auto closed = runtime->close_session_async("session"); + auto reopened = runtime->open_session_async("session"); + executor.opening.release(); + EXPECT_FALSE(result(std::move(opened))); + EXPECT_FALSE(result(std::move(reset))); + EXPECT_FALSE(result(std::move(closed))); + EXPECT_FALSE(result(std::move(reopened))); + EXPECT_EQ(executor.opened().size(), 3u); + EXPECT_EQ(executor.closed().size(), 2u); + EXPECT_EQ(runtime->info().active_sessions, 1u); +} + +TEST_F(ServingRuntimeTest, OpenCannotBypassPendingResetAndClose) { + start(); + EXPECT_FALSE(result(runtime->open_session_async("session"))); + executor.opening.hold(); + auto reset = runtime->reset_session_async("session"); + ASSERT_TRUE(executor.opening.wait_for(2)); + auto closed = runtime->close_session_async("session"); + auto opened = runtime->open_session_async("session"); + EXPECT_EQ( + opened.wait_for(std::chrono::seconds(0)), std::future_status::timeout); + executor.opening.release(); + EXPECT_FALSE(result(std::move(reset))); + EXPECT_FALSE(result(std::move(closed))); + EXPECT_FALSE(result(std::move(opened))); + EXPECT_EQ(executor.opened().size(), 3u); +} + +TEST_F(ServingRuntimeTest, CloseAcknowledgesLogicalNotPhysicalCleanup) { + start(); + EXPECT_FALSE(result(runtime->open_session_async("session"))); + executor.closing.hold(); + EXPECT_FALSE(result(runtime->close_session_async("session"))); + ASSERT_TRUE(executor.closing.wait_for(1)); + EXPECT_EQ(runtime->info().active_sessions, 0u); + EXPECT_EQ(executor.open_count(), 1); + EXPECT_FALSE(result(runtime->close_session_async("session"))); + executor.closing.release(); + runtime->shutdown(); + EXPECT_EQ(executor.closed(), executor.opened()); +} + +TEST_F(ServingRuntimeTest, LogicalCapacityRejectsWithoutCallingExecutor) { + executor.opening.hold(); + start({1, 128, 4}); + auto first = runtime->open_session_async("first"); + ASSERT_TRUE(executor.opening.wait_for(1)); + auto second = runtime->open_session_async("second"); + executor.opening.release(); + EXPECT_FALSE(result(std::move(first))); + expect_error(std::move(second), ErrorCode::CapacityExceeded); + EXPECT_EQ(executor.open_attempts.load(), 1); + EXPECT_FALSE(result(runtime->close_session_async("first"))); + EXPECT_FALSE(result(runtime->open_session_async("second"))); + EXPECT_EQ(runtime->info().active_sessions, 1u); +} + +TEST_F(ServingRuntimeTest, ExecutorRefusalReleasesInitialReservation) { + executor.capacity = 1; + start({2, 0, 4}); + EXPECT_FALSE(result(runtime->open_session_async("first"))); + expect_error( + runtime->open_session_async("second"), ErrorCode::CapacityExceeded); + EXPECT_EQ(runtime->info().active_sessions, 1u); + EXPECT_FALSE(result(runtime->close_session_async("first"))); + EXPECT_FALSE(result(runtime->open_session_async("second"))); +} + +TEST_F(ServingRuntimeTest, ResetRetainsSlotUntilColdReplacementCompletes) { + executor.capacity = 1; + start({1, 128, 4}); + EXPECT_FALSE(result(runtime->open_session_async("session"))); + const auto original = executor.opened().front(); + executor.closing.hold(); + executor.opening.hold(); + auto reset = runtime->reset_session_async("session"); + ASSERT_TRUE(executor.closing.wait_for(1)); + EXPECT_EQ(runtime->info().active_sessions, 1u); + EXPECT_EQ( + reset.wait_for(std::chrono::seconds(0)), std::future_status::timeout); + executor.closing.release(); + ASSERT_TRUE(executor.opening.wait_for(2)); + EXPECT_EQ(executor.open_count(), 0); + EXPECT_EQ(executor.closed(), (std::vector{original})); + EXPECT_EQ(runtime->info().active_sessions, 1u); + auto other = runtime->open_session_async("other"); + EXPECT_EQ( + reset.wait_for(std::chrono::seconds(0)), std::future_status::timeout); + executor.opening.release(); + EXPECT_FALSE(result(std::move(reset))); + expect_error(std::move(other), ErrorCode::CapacityExceeded); + ASSERT_EQ(executor.opened().size(), 2u); + EXPECT_NE(executor.opened().back(), original); + EXPECT_EQ(executor.open_count(), 1); +} + +TEST_F(ServingRuntimeTest, FailedResetRetainsUnavailableSlotAndCanRetry) { + start({1, 128, 4}); + EXPECT_FALSE(result(runtime->open_session_async("session"))); + executor.refuse_next_open = true; + auto failed = result(runtime->reset_session_async("session")); + ASSERT_TRUE(failed); + EXPECT_EQ(failed->code, ErrorCode::CapacityExceeded); + EXPECT_NE(failed->message.find("old state is gone"), std::string::npos); + EXPECT_NE(failed->message.find("unavailable"), std::string::npos); + EXPECT_EQ(runtime->info().active_sessions, 1u); + EXPECT_TRUE(runtime->info().ready); + EXPECT_EQ(executor.open_count(), 0); + EXPECT_EQ(executor.closed().size(), 1u); + expect_error(runtime->open_session_async("session"), ErrorCode::NotReady); + expect_error( + runtime->open_session_async("other"), ErrorCode::CapacityExceeded); + EXPECT_EQ(executor.open_attempts.load(), 2); + EXPECT_FALSE(result(runtime->reset_session_async("session"))); + EXPECT_EQ(executor.open_attempts.load(), 3); + EXPECT_EQ(runtime->info().active_sessions, 1u); + EXPECT_FALSE(result(runtime->open_session_async("session"))); + EXPECT_EQ(executor.open_attempts.load(), 3); +} + +TEST_F(ServingRuntimeTest, CloseReleasesUnavailableSlot) { + start({1, 128, 4}); + EXPECT_FALSE(result(runtime->open_session_async("session"))); + executor.refuse_next_open = true; + expect_error( + runtime->reset_session_async("session"), ErrorCode::CapacityExceeded); + EXPECT_FALSE(result(runtime->close_session_async("session"))); + EXPECT_EQ(runtime->info().active_sessions, 0u); + EXPECT_FALSE(result(runtime->open_session_async("other"))); + runtime->shutdown(); + EXPECT_EQ(executor.closed(), executor.opened()); +} + +TEST_F(ServingRuntimeTest, QueueBoundIncludesExecutingOperation) { + executor.opening.hold(); + start({2, 128, 2}); + auto first = runtime->open_session_async("first"); + ASSERT_TRUE(executor.opening.wait_for(1)); + auto second = runtime->open_session_async("second"); + auto rejected = runtime->close_session_async("first"); + EXPECT_EQ( + rejected.wait_for(std::chrono::seconds(0)), std::future_status::ready); + expect_error(std::move(rejected), ErrorCode::CapacityExceeded); + EXPECT_EQ(runtime->info().active_sessions, 1u); + executor.opening.release(); + EXPECT_FALSE(result(std::move(first))); + EXPECT_FALSE(result(std::move(second))); + EXPECT_EQ(runtime->info().active_sessions, 2u); + EXPECT_FALSE(result(runtime->close_session_async("first"))); +} + +TEST_F(ServingRuntimeTest, InitializationFailureSettlesQueuedAndNewOperations) { + executor.fail_initialize = true; + executor.initializing.hold(); + start(); + ASSERT_TRUE(executor.initializing.wait_for(1)); + auto first = runtime->open_session_async("first"); + auto second = runtime->open_session_async("second"); + executor.initializing.release(); + expect_error(std::move(first), ErrorCode::NotReady); + expect_error(std::move(second), ErrorCode::NotReady); + expect_error(runtime->open_session_async("third"), ErrorCode::NotReady); + EXPECT_FALSE(runtime->info().ready); + EXPECT_EQ(runtime->info().active_sessions, 0u); + EXPECT_EQ(executor.open_attempts.load(), 0); + runtime->shutdown(); +} + +#if ET_HAS_EXCEPTIONS +TEST_F(ServingRuntimeTest, InitializationExceptionBecomesNotReady) { + executor.throw_initialize = true; + executor.initializing.hold(); + start(); + ASSERT_TRUE(executor.initializing.wait_for(1)); + auto opened = runtime->open_session_async("session"); + executor.initializing.release(); + expect_error(std::move(opened), ErrorCode::NotReady); + expect_error(runtime->open_session_async("another"), ErrorCode::NotReady); + EXPECT_FALSE(runtime->info().ready); + EXPECT_EQ(runtime->info().active_sessions, 0u); + EXPECT_EQ(executor.open_attempts.load(), 0); + runtime->shutdown(); +} +#endif + +TEST_F(ServingRuntimeTest, ShutdownDiscardsLateOpenAndSettlesFullQueue) { + executor.opened_session.hold(); + start({2, 128, 3}); + auto opened = runtime->open_session_async("session"); + ASSERT_TRUE(executor.opened_session.wait_for(1)); + auto reset = runtime->reset_session_async("session"); + auto closed = runtime->close_session_async("session"); + EXPECT_EQ(executor.open_count(), 1); + EXPECT_TRUE(runtime->info().ready); + std::atomic completed{0}; + JoiningThreads callers(executor); + for (int i = 0; i < 4; ++i) { + callers.threads.emplace_back([&] { + runtime->shutdown(); + ++completed; + }); + } + ASSERT_TRUE(wait_until([&] { return !runtime->info().ready; })); + EXPECT_EQ(completed.load(), 0); + expect_error(runtime->open_session_async("new"), ErrorCode::NotReady); + executor.opened_session.release(); + ASSERT_TRUE(wait_until([&] { return completed.load() == 4; })); + expect_error(std::move(opened), ErrorCode::NotReady); + expect_error(std::move(reset), ErrorCode::NotReady); + expect_error(std::move(closed), ErrorCode::NotReady); + EXPECT_EQ(runtime->info().active_sessions, 0u); + EXPECT_EQ(executor.closed(), executor.opened()); + runtime->shutdown(); + expect_error(runtime->close_session_async("session"), ErrorCode::NotReady); + expect_error(runtime->reset_session_async("session"), ErrorCode::NotReady); +} + +TEST_F(ServingRuntimeTest, ShutdownDuringResetClosesOldAndLateReplacementOnce) { + start(); + EXPECT_FALSE(result(runtime->open_session_async("session"))); + executor.opened_session.hold(); + auto reset = runtime->reset_session_async("session"); + ASSERT_TRUE(executor.opened_session.wait_for(2)); + EXPECT_EQ(executor.opened().size(), 2u); + EXPECT_EQ(executor.closed().size(), 1u); + JoiningThreads callers(executor); + callers.threads.emplace_back([&] { runtime->shutdown(); }); + ASSERT_TRUE(wait_until([&] { return !runtime->info().ready; })); + executor.opened_session.release(); + callers.threads.front().join(); + expect_error(std::move(reset), ErrorCode::NotReady); + EXPECT_EQ(executor.closed(), executor.opened()); + EXPECT_EQ(executor.open_count(), 0); + EXPECT_EQ(runtime->info().active_sessions, 0u); +} + +TEST_F(ServingRuntimeTest, ConcurrentAdmissionOpensOneSession) { + constexpr int kCallers = 8; + executor.opening.hold(); + start({1, 128, kCallers}); + std::vector> futures(kCallers); + std::atomic admitted{0}; + JoiningThreads callers(executor); + for (int i = 0; i < kCallers; ++i) { + callers.threads.emplace_back([&, i] { + futures[i] = runtime->open_session_async("session"); + ++admitted; + }); + } + ASSERT_TRUE(wait_until([&] { return admitted.load() == kCallers; })); + ASSERT_TRUE(executor.opening.wait_for(1)); + executor.opening.release(); + for (auto& future : futures) { + EXPECT_FALSE(result(std::move(future))); + } + EXPECT_EQ(executor.open_attempts.load(), 1); + EXPECT_EQ(runtime->info().active_sessions, 1u); +} + +TEST_F(ServingRuntimeTest, DestructionJoinsAndClosesAllOwnedSessionsOnce) { + start(); + EXPECT_FALSE(result(runtime->open_session_async("first"))); + EXPECT_FALSE(result(runtime->open_session_async("second"))); + EXPECT_FALSE(result(runtime->reset_session_async("first"))); + runtime.reset(); + EXPECT_EQ(executor.open_count(), 0); + auto opened = executor.opened(); + auto closed = executor.closed(); + std::sort(opened.begin(), opened.end()); + std::sort(closed.begin(), closed.end()); + EXPECT_EQ(opened, closed); + const auto threads = executor.threads(); + ASSERT_FALSE(threads.empty()); + EXPECT_NE(threads.front(), std::this_thread::get_id()); + EXPECT_TRUE(std::all_of(threads.begin(), threads.end(), [&](auto id) { + return id == threads.front(); + })); +} + +TEST_F(ServingRuntimeTest, AbandonedFutureDoesNotAbandonSessionOwnership) { + executor.opening.hold(); + start(); + { + auto abandoned = runtime->open_session_async("session"); + ASSERT_TRUE(executor.opening.wait_for(1)); + } + auto duplicate = runtime->open_session_async("session"); + executor.opening.release(); + EXPECT_FALSE(result(std::move(duplicate))); + EXPECT_EQ(executor.open_attempts.load(), 1); + EXPECT_EQ(runtime->info().active_sessions, 1u); + runtime.reset(); + EXPECT_EQ(executor.closed(), executor.opened()); +} + +TEST_F(ServingRuntimeTest, InvalidKeysAndMissingResetDoNotAllocateSlots) { + start(); + expect_error(runtime->open_session_async(""), ErrorCode::InvalidArgument); + expect_error(runtime->close_session_async(""), ErrorCode::InvalidArgument); + expect_error(runtime->reset_session_async(""), ErrorCode::InvalidArgument); + expect_error( + runtime->reset_session_async("missing"), ErrorCode::SessionNotFound); + EXPECT_FALSE(result(runtime->close_session_async("missing"))); + EXPECT_EQ(runtime->info().active_sessions, 0u); + EXPECT_EQ(executor.open_attempts.load(), 0); +} + +TEST_F(ServingRuntimeTest, InvalidConfigurationIsInertWithoutExceptions) { + for (const auto config : + {ServingRuntimeConfig{0, 128, 4}, ServingRuntimeConfig{1, 128, 0}}) { + start(config); + EXPECT_FALSE(runtime->info().ready); + expect_error( + runtime->open_session_async("session"), ErrorCode::InvalidArgument); + expect_error( + runtime->reset_session_async("session"), ErrorCode::InvalidArgument); + expect_error( + runtime->close_session_async("session"), ErrorCode::InvalidArgument); + runtime->shutdown(); + runtime.reset(); + } + runtime = std::make_unique( + executor, nullptr, ServingRuntimeConfig{}); + EXPECT_FALSE(runtime->info().ready); + expect_error( + runtime->open_session_async("session"), ErrorCode::InvalidArgument); + EXPECT_EQ(executor.initialize_calls(), 0); +} + +} // namespace From b6eb36bffc0b733522a1ef3887ce033556699630 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Wed, 30 Sep 2026 15:56:12 -0700 Subject: [PATCH 2/2] up --- extension/llm/serving/serving_runtime.cpp | 77 +++- extension/llm/serving/serving_runtime.h | 30 +- .../llm/serving/test/serving_runtime_test.cpp | 397 +++++++++++++++++- 3 files changed, 482 insertions(+), 22 deletions(-) diff --git a/extension/llm/serving/serving_runtime.cpp b/extension/llm/serving/serving_runtime.cpp index 4d00eabf992..1c4857fa3f6 100644 --- a/extension/llm/serving/serving_runtime.cpp +++ b/extension/llm/serving/serving_runtime.cpp @@ -16,6 +16,7 @@ #include #include +#include namespace executorch { namespace extension { @@ -35,7 +36,8 @@ struct ServingRuntime::Impl { struct Command { Operation operation; std::string key; - std::promise completion; + // Queue transfers must not destroy user captures under mutex_. + std::unique_ptr completion; }; Impl( @@ -53,12 +55,44 @@ struct ServingRuntime::Impl { } } + static void complete( + std::unique_ptr completion, + LifecycleResult result) { + if (!completion || !*completion) { + return; + } +#if ET_HAS_EXCEPTIONS + try { +#endif + (*completion)(std::move(result)); +#if ET_HAS_EXCEPTIONS + } catch (...) { + ET_LOG(Error, "Lifecycle completion callback threw"); + } +#endif + } + std::future submit(Operation operation, std::string key) { - Command command{operation, std::move(key), {}}; - auto future = command.completion.get_future(); + auto completion = std::make_shared>(); + auto future = completion->get_future(); + submit( + operation, + std::move(key), + [completion = std::move(completion)](LifecycleResult result) { + completion->set_value(std::move(result)); + }); + return future; + } + + void + submit(Operation operation, std::string key, LifecycleCallback on_complete) { + Command command{ + operation, + std::move(key), + std::make_unique(std::move(on_complete))}; + LifecycleResult error; { std::lock_guard lock(mutex_); - LifecycleResult error; if (!valid_config_) { error = ServingError{ErrorCode::InvalidArgument, "invalid runtime config"}; @@ -74,15 +108,16 @@ struct ServingRuntime::Impl { error = ServingError{ ErrorCode::CapacityExceeded, "lifecycle operation limit reached"}; } - if (error) { - command.completion.set_value(std::move(error)); - return future; + if (!error) { + inbox_.push_back(std::move(command)); + ++outstanding_; } - inbox_.push_back(std::move(command)); - ++outstanding_; } - cv_.notify_one(); - return future; + if (error) { + complete(std::move(command.completion), std::move(error)); + } else { + cv_.notify_one(); + } } LifecycleResult process(Command& command) { @@ -190,8 +225,8 @@ struct ServingRuntime::Impl { result = ServingError{ErrorCode::NotReady, "runtime is stopping"}; } --outstanding_; - command.completion.set_value(std::move(result)); } + complete(std::move(command.completion), std::move(result)); } decltype(sessions_) retired; { @@ -260,16 +295,34 @@ ServingRuntime::~ServingRuntime() { shutdown(); } +void ServingRuntime::open_session_async( + std::string key, + LifecycleCallback on_complete) { + impl_->submit(Impl::Operation::Open, std::move(key), std::move(on_complete)); +} + std::future ServingRuntime::open_session_async( std::string key) { return impl_->submit(Impl::Operation::Open, std::move(key)); } +void ServingRuntime::close_session_async( + std::string key, + LifecycleCallback on_complete) { + impl_->submit(Impl::Operation::Close, std::move(key), std::move(on_complete)); +} + std::future ServingRuntime::close_session_async( std::string key) { return impl_->submit(Impl::Operation::Close, std::move(key)); } +void ServingRuntime::reset_session_async( + std::string key, + LifecycleCallback on_complete) { + impl_->submit(Impl::Operation::Reset, std::move(key), std::move(on_complete)); +} + std::future ServingRuntime::reset_session_async( std::string key) { return impl_->submit(Impl::Operation::Reset, std::move(key)); diff --git a/extension/llm/serving/serving_runtime.h b/extension/llm/serving/serving_runtime.h index 1397ea5d413..672cafac7a0 100644 --- a/extension/llm/serving/serving_runtime.h +++ b/extension/llm/serving/serving_runtime.h @@ -9,6 +9,7 @@ #pragma once #include +#include #include #include #include @@ -28,13 +29,15 @@ struct ET_EXPERIMENTAL ServingRuntimeConfig { std::size_t max_sessions = 1; // Reported metadata; 0 means unknown. std::size_t max_context_length = 0; - // Queued plus executing lifecycle operations. Must be non-zero. + // Queued plus executing lifecycle operations, excluding completion + // callbacks. Must be non-zero. std::size_t max_pending_operations = 64; }; // nullopt acknowledges success. Dropping a future does not cancel its // operation. using LifecycleResult ET_EXPERIMENTAL = std::optional; +using LifecycleCallback ET_EXPERIMENTAL = std::function; class ET_EXPERIMENTAL ServingRuntime { public: @@ -55,15 +58,24 @@ class ET_EXPERIMENTAL ServingRuntime { ServingRuntime& operator=(ServingRuntime&&) = delete; // Any thread. Keys must be non-empty. Accepted operations run in admission - // order, including operations submitted before a previous future completes. - // Queue saturation returns CapacityExceeded without changing session state. + // order. Queue saturation reports CapacityExceeded without changing state. + // For all lifecycle operations, callbacks run once on control after + // processing and admission release, or inline on the caller for immediate + // rejection. They run without runtime locks and may race with submission's + // return. Callbacks must do short, nonblocking work: no waits for runtime + // work or runtime shutdown/destruction. Exceptions are logged, not retried. + // An empty callback discards the result. Futures adapt this same completion + // path. + // // Opens may queue during initialization; success means an empty session is // owned, or the key was already open. An initial open failure releases its // slot. Opening an unavailable key fails with NotReady; use reset to retry. + void open_session_async(std::string key, LifecycleCallback on_complete); std::future open_session_async(std::string key); // Idempotent, including absent keys. Success releases the logical slot and // requests Session closure through RAII, not synchronous physical cleanup. + void close_session_async(std::string key, LifecycleCallback on_complete); std::future close_session_async(std::string key); // Cold replacement, retaining the key and logical slot throughout. Success @@ -71,6 +83,7 @@ class ET_EXPERIMENTAL ServingRuntime { // key unavailable, still reserving its slot until reset succeeds or close // releases it. A missing key returns SessionNotFound. Never uses prefix // reuse. + void reset_session_async(std::string key, LifecycleCallback on_complete); std::future reset_session_async(std::string key); // A synchronized snapshot. Ready means initialized and accepting operations, @@ -78,11 +91,12 @@ class ET_EXPERIMENTAL ServingRuntime { // and unavailable keys, but not commands still waiting to reserve a slot. ServingInfo info() const; - // Idempotent, including concurrent callers. Stops admission, settles all - // accepted futures, and joins the control and Runner threads. No new work is - // accepted afterward. Call only from external threads, never engine - // callbacks. Destruction calls shutdown; as usual, object lifetime must be - // externally synchronized against callers still accessing the runtime. + // Idempotent, including concurrent callers. Stops admission, completes all + // accepted operations, and joins the control and Runner threads. Accepted + // callbacks have returned and their runtime-owned captures are released. + // No new work is accepted afterward. Call only from external threads, never + // engine or lifecycle callbacks. Destruction calls shutdown; as usual, object + // lifetime must be synchronized against callers still accessing the runtime. void shutdown(); private: diff --git a/extension/llm/serving/test/serving_runtime_test.cpp b/extension/llm/serving/test/serving_runtime_test.cpp index 29c927a010b..8cf55e202f8 100644 --- a/extension/llm/serving/test/serving_runtime_test.cpp +++ b/extension/llm/serving/test/serving_runtime_test.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -22,6 +23,7 @@ #include #include +#include #include @@ -30,6 +32,7 @@ using executorch::extension::llm::batching::Position; using executorch::extension::llm::batching::SessionId; using executorch::extension::llm::batching::testing::FakeExecutor; using executorch::extension::llm::serving::ErrorCode; +using executorch::extension::llm::serving::LifecycleCallback; using executorch::extension::llm::serving::LifecycleResult; using executorch::extension::llm::serving::ServingError; using executorch::extension::llm::serving::ServingRuntime; @@ -56,9 +59,9 @@ class Gate { cv_.wait(lock, [this] { return !held_; }); } - bool wait_for(std::size_t arrivals) { + bool wait_for(std::size_t arrivals, std::chrono::seconds timeout = kTimeout) { std::unique_lock lock(mutex_); - return cv_.wait_for(lock, kTimeout, [&] { return arrivals_ >= arrivals; }); + return cv_.wait_for(lock, timeout, [&] { return arrivals_ >= arrivals; }); } void release() { @@ -145,21 +148,59 @@ class LifecycleExecutor : public FakeExecutor { class ServingRuntimeTest : public ::testing::Test { protected: + void SetUp() override { + executorch::runtime::runtime_init(); + } + void start(ServingRuntimeConfig config = {2, 128, 16}) { runtime = std::make_unique( executor, DecodeFirstScheduler::create(), config); } + void watch_callbacks() { + // Reentry regressions can deadlock inside info()/submit() or teardown. + callback_watchdog = std::thread([this] { + if (!test_finished.wait_for(1, kTimeout * 6)) { + ADD_FAILURE() << "callback test or teardown deadlocked"; + std::abort(); + } + }); + } + void TearDown() override { // Release gates even after a fatal assertion, before the runtime joins. + callback_gate.release(); + cleanup_gate.release(); executor.release_all(); + for (auto& caller : callback_callers) { + if (caller.joinable()) { + caller.join(); + } + } + if (runtime) { + runtime->shutdown(); + } runtime.reset(); EXPECT_EQ(executor.clone_calls.load(), 0); EXPECT_TRUE(executor.seen().empty()); + test_finished.arrive(); + if (callback_watchdog.joinable()) { + callback_watchdog.join(); + } } LifecycleExecutor executor; std::unique_ptr runtime; + Gate callback_gate; + Gate cleanup_gate; + std::vector callback_callers; + Gate test_finished; + std::thread callback_watchdog; + struct SmallCallbackState { + std::atomic destructions{0}; + std::atomic calls{0}; + std::thread::id thread; + } small_callback_state; }; struct JoiningThreads { @@ -553,6 +594,358 @@ TEST_F(ServingRuntimeTest, AbandonedFutureDoesNotAbandonSessionOwnership) { EXPECT_EQ(executor.closed(), executor.opened()); } +TEST_F(ServingRuntimeTest, CallbackAndFutureOperationsPreserveAdmissionOrder) { + watch_callbacks(); + executor.opening.hold(); + callback_gate.hold(); + start({1, 128, 4}); + const auto caller = std::this_thread::get_id(); + auto calls = std::make_shared>(0); + runtime->open_session_async("session", [this, calls, caller](auto error) { + EXPECT_FALSE(error); + EXPECT_EQ(calls->fetch_add(1), 0); + const auto control = std::this_thread::get_id(); + EXPECT_NE(control, caller); + for (auto engine : executor.threads()) { + EXPECT_NE(control, engine); + } + EXPECT_EQ(runtime->info().active_sessions, 1u); + callback_gate.arrive(); + }); + ASSERT_TRUE(executor.opening.wait_for(1)); + auto reset = runtime->reset_session_async("session"); + runtime->close_session_async("session", [this, calls, caller](auto error) { + EXPECT_FALSE(error); + EXPECT_EQ(calls->fetch_add(1), 1); + EXPECT_NE(std::this_thread::get_id(), caller); + for (auto engine : executor.threads()) { + EXPECT_NE(std::this_thread::get_id(), engine); + } + EXPECT_EQ(runtime->info().active_sessions, 0u); + }); + auto reopened = runtime->open_session_async("session"); + executor.opening.release(); + ASSERT_TRUE(callback_gate.wait_for(1)); + EXPECT_EQ( + reset.wait_for(std::chrono::seconds(0)), std::future_status::timeout); + EXPECT_EQ(executor.open_attempts.load(), 1); + callback_gate.release(); + EXPECT_FALSE(result(std::move(reset))); + EXPECT_FALSE(result(std::move(reopened))); + EXPECT_EQ(runtime->info().active_sessions, 1u); + EXPECT_EQ(executor.opened().size(), 3u); + runtime->shutdown(); + EXPECT_EQ(calls->load(), 2); + EXPECT_EQ(executor.closed(), executor.opened()); +} + +TEST_F(ServingRuntimeTest, CallbackCanSubmitFollowupAfterReleasingOnlyPermit) { + watch_callbacks(); + start({1, 128, 1}); + const auto caller = std::this_thread::get_id(); + auto calls = std::make_shared>(0); + runtime->open_session_async("session", [this, calls, caller](auto error) { + EXPECT_FALSE(error); + EXPECT_EQ(calls->fetch_add(1), 0); + const auto control = std::this_thread::get_id(); + EXPECT_NE(control, caller); + for (auto engine : executor.threads()) { + EXPECT_NE(control, engine); + } + EXPECT_EQ(runtime->info().active_sessions, 1u); + runtime->reset_session_async("session", [this, calls, control](auto reset) { + EXPECT_FALSE(reset); + EXPECT_EQ(calls->fetch_add(1), 1); + EXPECT_EQ(std::this_thread::get_id(), control); + EXPECT_EQ(runtime->info().active_sessions, 1u); + runtime->close_session_async( + "session", [this, calls, control](auto closed) { + EXPECT_FALSE(closed); + EXPECT_EQ(calls->fetch_add(1), 2); + EXPECT_EQ(std::this_thread::get_id(), control); + EXPECT_EQ(runtime->info().active_sessions, 0u); + callback_gate.arrive(); + }); + }); + }); + ASSERT_TRUE(callback_gate.wait_for(1)); + runtime->shutdown(); + EXPECT_EQ(calls->load(), 3); + EXPECT_EQ(executor.open_attempts.load(), 2); + EXPECT_EQ(executor.closed(), executor.opened()); +} + +TEST_F(ServingRuntimeTest, CallbackRejectionsAreInlineAndAllowReentry) { + watch_callbacks(); + using Operation = void (ServingRuntime::*)(std::string, LifecycleCallback); + const Operation operations[] = { + &ServingRuntime::open_session_async, + &ServingRuntime::reset_session_async, + &ServingRuntime::close_session_async}; + enum class Rejection { Validation, Inert, Stopped, Saturated }; + for (auto rejection : + {Rejection::Validation, + Rejection::Inert, + Rejection::Stopped, + Rejection::Saturated}) { + SCOPED_TRACE(static_cast(rejection)); + executor.opening.hold(); + start({rejection == Rejection::Inert ? 0u : 1u, 128, 1}); + if (rejection == Rejection::Stopped) { + runtime->shutdown(); + } else if (rejection == Rejection::Saturated) { + runtime->open_session_async("occupied", {}); + ASSERT_TRUE(executor.opening.wait_for(1)); + } + const std::string key = rejection == Rejection::Validation ? "" : "session"; + const auto expected = rejection == Rejection::Stopped ? ErrorCode::NotReady + : rejection == Rejection::Saturated ? ErrorCode::CapacityExceeded + : ErrorCode::InvalidArgument; + const auto caller = std::this_thread::get_id(); + auto calls = std::make_shared>(0); + for (auto operation : operations) { + const auto before = calls->load(); + (runtime.get()->*operation)( + key, [this, calls, caller, expected, key](auto error) { + EXPECT_EQ(std::this_thread::get_id(), caller); + ASSERT_TRUE(error); + EXPECT_EQ(error->code, expected); + EXPECT_FALSE(error->message.empty()); + (void)runtime->info(); + ++*calls; + runtime->close_session_async( + key, [calls, caller, expected](auto nested) { + EXPECT_EQ(std::this_thread::get_id(), caller); + ASSERT_TRUE(nested); + EXPECT_EQ(nested->code, expected); + ++*calls; + }); + }); + EXPECT_EQ(calls->load(), before + 2); + } + executor.opening.release(); + runtime->shutdown(); + EXPECT_EQ(calls->load(), 6); + runtime.reset(); + } +} + +TEST_F(ServingRuntimeTest, QueuedCallbacksSettleOnceOnInitFailureAndShutdown) { + watch_callbacks(); + for (bool shutdown : {false, true}) { + SCOPED_TRACE(shutdown); + executor.fail_initialize = !shutdown; + executor.initializing.hold(); + executor.opened_session.hold(); + const auto previous_threads = executor.threads().size(); + start({3, 128, 3}); + struct State { + std::atomic calls[3]{}; + std::mutex mutex; + std::thread::id control; + }; + auto state = std::make_shared(); + const auto caller = std::this_thread::get_id(); + auto completion = [this, state, caller, previous_threads]( + int index, LifecycleResult error) { + EXPECT_EQ(state->calls[index].fetch_add(1), 0); + EXPECT_NE(std::this_thread::get_id(), caller); + const auto engines = executor.threads(); + for (auto i = previous_threads; i < engines.size(); ++i) { + EXPECT_NE(std::this_thread::get_id(), engines[i]); + } + { + std::lock_guard lock(state->mutex); + if (state->control == std::thread::id{}) { + state->control = std::this_thread::get_id(); + } + EXPECT_EQ(std::this_thread::get_id(), state->control); + } + EXPECT_FALSE(runtime->info().ready); + ASSERT_TRUE(error); + EXPECT_EQ(error->code, ErrorCode::NotReady); + }; + runtime->open_session_async( + "first", [completion](auto error) { completion(0, error); }); + if (shutdown) { + executor.initializing.release(); + ASSERT_TRUE(executor.opened_session.wait_for(1)); + runtime->reset_session_async( + "first", [completion](auto error) { completion(1, error); }); + runtime->close_session_async( + "first", [completion](auto error) { completion(2, error); }); + } else { + runtime->open_session_async( + "second", [completion](auto error) { completion(1, error); }); + runtime->open_session_async( + "third", [completion](auto error) { completion(2, error); }); + } + expect_error( + runtime->open_session_async("overflow"), ErrorCode::CapacityExceeded); + if (shutdown) { + callback_callers.emplace_back([this] { runtime->shutdown(); }); + ASSERT_TRUE(wait_until([this] { return !runtime->info().ready; })); + } + executor.release_all(); + ASSERT_TRUE(wait_until([state] { + return state->calls[0].load() == 1 && state->calls[1].load() == 1 && + state->calls[2].load() == 1; + })); + runtime->shutdown(); + for (auto& thread : callback_callers) { + thread.join(); + } + callback_callers.clear(); + for (auto& count : state->calls) { + EXPECT_EQ(count.load(), 1); + } + EXPECT_EQ(runtime->info().active_sessions, 0u); + EXPECT_EQ(executor.closed(), executor.opened()); + runtime.reset(); + } +} + +#if ET_HAS_EXCEPTIONS +TEST_F(ServingRuntimeTest, ThrowingCallbacksDoNotRollbackOrStopControl) { + watch_callbacks(); + start({1, 128, 1}); + auto calls = std::make_shared>(0); + runtime->open_session_async("session", [this, calls](auto error) { + EXPECT_FALSE(error); + EXPECT_EQ(runtime->info().active_sessions, 1u); + ++*calls; + callback_gate.arrive(); + throw std::runtime_error("admitted callback failed"); + }); + ASSERT_TRUE(callback_gate.wait_for(1)); + EXPECT_FALSE(result(runtime->open_session_async("session"))); + EXPECT_EQ(executor.open_attempts.load(), 1); + EXPECT_EQ(runtime->info().active_sessions, 1u); + const auto caller = std::this_thread::get_id(); + EXPECT_NO_THROW( + runtime->reset_session_async("", [this, calls, caller](auto error) { + ASSERT_TRUE(error); + EXPECT_EQ(error->code, ErrorCode::InvalidArgument); + EXPECT_EQ(std::this_thread::get_id(), caller); + EXPECT_EQ(runtime->info().active_sessions, 1u); + ++*calls; + throw std::runtime_error("rejection callback failed"); + })); + EXPECT_EQ(calls->load(), 2); + EXPECT_FALSE(result(runtime->reset_session_async("session"))); + EXPECT_EQ(executor.open_attempts.load(), 2); + EXPECT_FALSE(result(runtime->close_session_async("session"))); + runtime->shutdown(); + EXPECT_EQ(calls->load(), 2); + EXPECT_EQ(executor.closed(), executor.opened()); +} +#endif + +TEST_F(ServingRuntimeTest, EmptyCallbacksDiscardNotificationNotOperation) { + watch_callbacks(); + executor.opening.hold(); + start({1, 128, 4}); + runtime->open_session_async("session", {}); + ASSERT_TRUE(executor.opening.wait_for(1)); + runtime->reset_session_async("session", {}); + runtime->close_session_async("session", {}); + auto reopened = runtime->open_session_async("session"); + executor.opening.release(); + EXPECT_FALSE(result(std::move(reopened))); + EXPECT_EQ(executor.open_attempts.load(), 3); + EXPECT_EQ(runtime->info().active_sessions, 1u); + runtime->open_session_async("", {}); + runtime->shutdown(); + EXPECT_EQ(executor.closed(), executor.opened()); +} + +TEST_F(ServingRuntimeTest, SmallCallbackDestructionCanReenterInfo) { + watch_callbacks(); + executor.opening.hold(); + start(); + // Keep this nothrow-copyable and small enough for libc++'s inline storage. + struct SmallCallback { + ServingRuntime* runtime; + SmallCallbackState* state; + + ~SmallCallback() { + (void)runtime->info(); + ++state->destructions; + } + + void operator()(LifecycleResult error) const { + EXPECT_FALSE(error); + state->thread = std::this_thread::get_id(); + ++state->calls; + } + }; + static_assert(sizeof(SmallCallback) == 2 * sizeof(void*)); + static_assert(std::is_nothrow_copy_constructible::value); + runtime->open_session_async( + "session", SmallCallback{runtime.get(), &small_callback_state}); + ASSERT_TRUE(executor.opening.wait_for(1)); + const auto before_completion = small_callback_state.destructions.load(); + auto closed = runtime->close_session_async("session"); + executor.opening.release(); + EXPECT_FALSE(result(std::move(closed))); + runtime->shutdown(); + EXPECT_GT(small_callback_state.destructions.load(), before_completion); + EXPECT_EQ(small_callback_state.calls.load(), 1); + EXPECT_NE(small_callback_state.thread, std::this_thread::get_id()); + for (auto engine : executor.threads()) { + EXPECT_NE(small_callback_state.thread, engine); + } +} + +TEST_F( + ServingRuntimeTest, + ShutdownWaitsForCallbackAndCaptureAfterPermitRelease) { + watch_callbacks(); + start({1, 128, 1}); + callback_gate.hold(); + auto cleaned = std::make_shared>(false); + auto capture = std::shared_ptr(new int(0), [this, cleaned](int* value) { + cleanup_gate.arrive(); + *cleaned = true; + delete value; + }); + auto calls = std::make_shared>(0); + runtime->open_session_async("session", [this, capture, calls](auto error) { + (void)capture; + EXPECT_FALSE(error); + EXPECT_EQ(runtime->info().active_sessions, 1u); + ++*calls; + callback_gate.arrive(); + }); + ASSERT_TRUE(callback_gate.wait_for(1)); + cleanup_gate.hold(); + capture.reset(); + // Admission succeeds while the previous callback still owns its captures. + auto closed = runtime->close_session_async("session"); + EXPECT_EQ( + closed.wait_for(std::chrono::seconds(0)), std::future_status::timeout); + auto stopped = std::make_shared>(0); + for (int i = 0; i < 2; ++i) { + callback_callers.emplace_back([this, stopped] { + runtime->shutdown(); + ++*stopped; + }); + } + ASSERT_TRUE(wait_until([this] { return !runtime->info().ready; })); + EXPECT_EQ(stopped->load(), 0); + EXPECT_FALSE(cleaned->load()); + callback_gate.release(); + ASSERT_TRUE(cleanup_gate.wait_for(1)); + EXPECT_EQ(stopped->load(), 0); + EXPECT_FALSE(cleaned->load()); + cleanup_gate.release(); + ASSERT_TRUE(wait_until([stopped] { return stopped->load() == 2; })); + expect_error(std::move(closed), ErrorCode::NotReady); + EXPECT_TRUE(cleaned->load()); + EXPECT_EQ(calls->load(), 1); + EXPECT_EQ(executor.closed(), executor.opened()); +} + TEST_F(ServingRuntimeTest, InvalidKeysAndMissingResetDoNotAllocateSlots) { start(); expect_error(runtime->open_session_async(""), ErrorCode::InvalidArgument);