[LLM serving 2/x] Add batching-backed serving session lifecycle (#23277)

Introduce an experimental ServingRuntime that manages named sessions on
top of the existing batching runtime. It owns one Runner and scheduler,
while physical session state and cleanup remain with batching::Session.

* Add asynchronous open, close, and reset operations, with bounded
session capacity and pending operations processed in admission order.
* Reset destroys and reopens the physical session under the same key.
Reopening failure leaves the key explicitly unavailable until reset
succeeds or close releases it.
* Use RAII for physical session closure, and provide synchronized
readiness information and idempotent shutdown that settles accepted
operations.
* Add thread-safe Runner::initialization_state() reporting Pending,
Ready, or Failed. Initialization failures and exceptions refuse work and
settle queued opens without an executor wrapper.

This PR covers lifecycle management only. Generation, tokenization,
streaming, and prefix reuse are added separately; no new scheduler or
decode loop is introduced.

Author: metascroy

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;
   }

extension/llm/batching/runner.cpp

@@ -290,6 +290,10 @@ class RunnerImpl : public std::enable_shared_from_this<RunnerImpl> {
       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<RunnerImpl> {
   std::condition_variable stopped_cv_;
   std::vector<Command> inbox_;
   std::atomic<Lifecycle> lifecycle_{Lifecycle::Running};
+  std::atomic<InitializationState> 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.

extension/llm/batching/runner.h

@@ -236,6 +236,8 @@ class ET_EXPERIMENTAL Session {
   std::unique_ptr<SessionState> 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.

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<std::mutex> 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<std::mutex> lock(mutex_);
+    return cv_.wait_for(lock, kTimeout, [this] { return entered_; });
+  }
+
+  void release_initialize() {
+    {
+      std::lock_guard<std::mutex> 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<InitializeOutcome> {
+};
+
+} // 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<std::future<std::optional<Session>>> 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<std::future<std::optional<Session>>> 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<InitializeOutcome>& 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";
+    });

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()