[LLM serving 2/x] Add batching-backed serving session lifecycle

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