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..1c4857fa3f6 --- /dev/null +++ b/extension/llm/serving/serving_runtime.cpp @@ -0,0 +1,342 @@ +/* + * 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 + +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; + // Queue transfers must not destroy user captures under mutex_. + std::unique_ptr 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(); }); + } + } + + 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) { + 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_); + 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) { + inbox_.push_back(std::move(command)); + ++outstanding_; + } + } + if (error) { + complete(std::move(command.completion), std::move(error)); + } else { + cv_.notify_one(); + } + } + + 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_; + } + complete(std::move(command.completion), 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(); +} + +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)); +} + +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..672cafac7a0 --- /dev/null +++ b/extension/llm/serving/serving_runtime.h @@ -0,0 +1,110 @@ +/* + * 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 +#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, 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: + // 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. 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 + // 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. + 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, + // 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, 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: + 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..8cf55e202f8 --- /dev/null +++ b/extension/llm/serving/test/serving_runtime_test.cpp @@ -0,0 +1,983 @@ +/* + * 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 +#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::LifecycleCallback; +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::chrono::seconds timeout = kTimeout) { + std::unique_lock lock(mutex_); + return cv_.wait_for(lock, timeout, [&] { 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 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 { + 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, 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); + 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