Set runtime options for Terminating session (#1011)

Add API option for setting any runtime options. Particularly, add
Terminate session option for the user, the session remains in the
terminated state until the user disables terminate_session.

To terminate session:
SetRuntimeOptionsConfig("terminate_session", "1")

To unset:
SetRuntimeOptionsConfig("terminate_session", "0")

---------

Co-authored-by: Pranav Sharma <[email protected]>

Author: 2 people

documents/Runtime_option.md

@@ -0,0 +1,15 @@
+# Runtime Options
+
+This file will provide details on the usage of SetRuntimeOption API. It will list all the current key value pairs which can be used as an input for this API.
+
+## Set Terminate
+
+Set Terminate is a runtime option to terminate the current session or continue/restart an already terminated session. There are two valid ways to call Set Terminate.
+
+To enable terminate, the valid pair is: ("set_terminate", "1")
+
+To disable terminate, the valid pair is: ("set_terminate", "0")
+
+Key: "set_terminate"
+
+Accepted values: ("0", "1")

src/generators.cpp

@@ -30,6 +30,11 @@ std::string CurrentModulePath() {
 }
 #endif
 
+void ThrowErrorIfSessionTerminated(bool is_session_terminated) {
+  if (is_session_terminated)
+    throw std::runtime_error("Session in Terminated state, exiting!");
+}
+
 namespace Generators {
 
 #if USE_CUDA
@@ -284,6 +289,7 @@ Generator::Generator(const Model& model, const GeneratorParams& params) : model_
 }
 
 void Generator::ComputeLogits() {
+  ThrowErrorIfSessionTerminated(state_->session_terminated_);
   if (computed_logits_)
     throw std::runtime_error("ComputeLogits called again without calling GenerateNextToken first");
 
@@ -301,7 +307,25 @@ void Generator::ComputeLogits() {
   search_->ApplyRepetitionPenalty(search.repetition_penalty);
 }
 
+void Generator::SetRuntimeOption(const char* key, const char* value) {
+  // TODO: Need a better way to handle different keys
+  // We can create a config manager to host all configurations and do comparison at that point
+  if (strcmp(key, "terminate_session") == 0) {
+    if (strcmp(value, "0") == 0) {
+      state_->UnsetTerminate();
+    } else if (strcmp(value, "1") == 0) {
+      state_->SetTerminate();
+    } else {
+      // Value not expected
+      throw std::runtime_error(std::string("terminate_session key value unexpected: ") + value);
+    }
+  } else {
+    throw std::runtime_error(std::string("SetRuntimeOption key is not expected: ") + key);
+  }
+}
+
 bool Generator::IsDone() const {
+  ThrowErrorIfSessionTerminated(state_->session_terminated_);
   if (computed_logits_)
     throw std::runtime_error("IsDone() can't be called in the middle of processing logits");
 
@@ -313,7 +337,12 @@ bool Generator::IsDone() const {
   return is_done;
 }
 
+bool Generator::IsSessionTerminated() const {
+  return state_->session_terminated_;
+}
+
 void Generator::GenerateNextToken() {
+  ThrowErrorIfSessionTerminated(state_->session_terminated_);
   if (!computed_logits_)
     throw std::runtime_error("Must call ComputeLogits before GenerateNextToken");
   computed_logits_ = false;

src/generators.h

@@ -40,6 +40,8 @@ using cudaStream_t = void*;
 #include "runtime_settings.h"
 #include "tensor.h"
 
+void ThrowErrorIfSessionTerminated(bool is_session_terminated);
+
 namespace Generators {
 struct Model;
 struct State;
@@ -108,7 +110,9 @@ struct Generator : LeakChecked<Generator> {
   Generator(const Model& model, const GeneratorParams& params);
 
   bool IsDone() const;
+  void SetRuntimeOption(const char* key, const char* value);
   void ComputeLogits();
+  bool IsSessionTerminated() const;
   void GenerateNextToken();
 
   DeviceMemorySpan<int32_t> GetSequence(size_t index) const;

src/models/model.cpp

@@ -65,7 +65,18 @@ void State::Run(OrtSession& session, int new_batch_size) {
   }
 }
 
+void State::SetTerminate() {
+  session_terminated_ = true;
+  run_options_->SetTerminate();
+}
+
+void State::UnsetTerminate() {
+  session_terminated_ = false;
+  run_options_->UnsetTerminate();
+}
+
 OrtValue* State::GetInput(const char* name) {
+  ThrowErrorIfSessionTerminated(session_terminated_);
   for (size_t i = 0; i < input_names_.size(); i++) {
     if (std::strcmp(input_names_[i], name) == 0) {
       return inputs_[i];
@@ -75,6 +86,7 @@ OrtValue* State::GetInput(const char* name) {
 }
 
 OrtValue* State::GetOutput(const char* name) {
+  ThrowErrorIfSessionTerminated(session_terminated_);
   for (size_t i = 0; i < output_names_.size(); i++) {
     if (std::strcmp(output_names_[i], name) == 0) {
       return outputs_[i];

src/models/model.h

@@ -34,6 +34,9 @@ struct State {
   virtual const CapturedGraphInfo* GetCapturedGraphInfo() const { return nullptr; }
   virtual void Finalize() {}
 
+  void SetTerminate();
+  void UnsetTerminate();
+  mutable bool session_terminated_{};
   OrtValue* GetInput(const char* name);
 
   virtual OrtValue* GetOutput(const char* name);

src/ort_genai.h

@@ -258,6 +258,10 @@ struct OgaGenerator : OgaAbstract {
     return OgaGenerator_IsDone(this);
   }
 
+  bool IsSessionTerminated() const {
+    return OgaGenerator_IsSessionTerminated(this);
+  }
+
   void ComputeLogits() {
     OgaCheckResult(OgaGenerator_ComputeLogits(this));
   }
@@ -266,6 +270,10 @@ struct OgaGenerator : OgaAbstract {
     OgaCheckResult(OgaGenerator_GenerateNextToken(this));
   }
 
+  void SetRuntimeOption(const char* key, const char* value) {
+    OgaCheckResult(OgaGenerator_SetRuntimeOption(this, key, value));
+  }
+
   size_t GetSequenceCount(size_t index) const {
     return OgaGenerator_GetSequenceCount(this, index);
   }

src/ort_genai_c.cpp

@@ -271,6 +271,10 @@ bool OGA_API_CALL OgaGenerator_IsDone(const OgaGenerator* generator) {
   return reinterpret_cast<const Generators::Generator*>(generator)->IsDone();
 }
 
+bool OGA_API_CALL OgaGenerator_IsSessionTerminated(const OgaGenerator* generator) {
+  return reinterpret_cast<const Generators::Generator*>(generator)->IsSessionTerminated();
+}
+
 OgaResult* OGA_API_CALL OgaGenerator_ComputeLogits(OgaGenerator* generator) {
   OGA_TRY
   reinterpret_cast<Generators::Generator*>(generator)->ComputeLogits();
@@ -285,6 +289,13 @@ OgaResult* OGA_API_CALL OgaGenerator_GenerateNextToken(OgaGenerator* generator)
   OGA_CATCH
 }
 
+OgaResult* OGA_API_CALL OgaGenerator_SetRuntimeOption(OgaGenerator* generator, const char* key, const char* value) {
+  OGA_TRY
+  reinterpret_cast<Generators::Generator*>(generator)->SetRuntimeOption(key, value);
+  return nullptr;
+  OGA_CATCH
+}
+
 OgaResult* OGA_API_CALL OgaGenerator_GetOutput(const OgaGenerator* oga_generator, const char* name, OgaTensor** out) {
   OGA_TRY
   auto& generator = *reinterpret_cast<const Generators::Generator*>(oga_generator);

src/ort_genai_c.h

@@ -278,6 +278,7 @@ OGA_EXPORT void OGA_API_CALL OgaDestroyGenerator(OgaGenerator* generator);
  * \return True if the generator has finished generating all the sequences, false otherwise.
  */
 OGA_EXPORT bool OGA_API_CALL OgaGenerator_IsDone(const OgaGenerator* generator);
+OGA_EXPORT bool OGA_API_CALL OgaGenerator_IsSessionTerminated(const OgaGenerator* generator);
 
 /*
  * \brief Computes the logits from the model based on the input ids and the past state. The computed logits are stored in the generator.
@@ -287,6 +288,8 @@ OGA_EXPORT bool OGA_API_CALL OgaGenerator_IsDone(const OgaGenerator* generator);
 OGA_EXPORT OgaResult* OGA_API_CALL OgaGenerator_ComputeLogits(OgaGenerator* generator);
 OGA_EXPORT OgaResult* OGA_API_CALL OgaGenerator_GenerateNextToken(OgaGenerator* generator);
 
+OGA_EXPORT OgaResult* OGA_API_CALL OgaGenerator_SetRuntimeOption(OgaGenerator* generator, const char* key, const char* value);
+
 /*
  * \brief Returns a copy of the model output identified by the given name as an OgaTensor on CPU. The buffer is owned by returned OgaTensor
  *       and will be released when the OgaTensor is destroyed

test/c_api_tests.cpp

@@ -5,6 +5,8 @@
 #include <iostream>
 #include <ort_genai.h>
 #include "../src/span.h"
+#include <thread>
+#include <vector>
 
 #ifndef MODEL_PATH
 #define MODEL_PATH "../../test/test_models/"
@@ -283,6 +285,55 @@ TEST(CAPITests, GetOutputCAPI) {
   generator->GenerateNextToken();
 }
 
+TEST(CAPITests, SetTerminate) {
+#if TEST_PHI2
+
+  auto GeneratorSetTerminateCall = [](OgaGenerator* generator) {
+    // Set Terminate
+    generator->SetRuntimeOption("terminate_session", "1");
+  };
+
+  auto GenerateOutput = [](OgaGenerator* generator, std::unique_ptr<OgaTokenizerStream> tokenizer_stream) {
+    try {
+      while (!generator->IsDone()) {
+        generator->ComputeLogits();
+        generator->GenerateNextToken();
+      }
+    }
+    catch (const std::exception& e) {
+      EXPECT_EQ(generator->IsSessionTerminated(), true);
+      std::cout << "Session Terminated: " << e.what() << std::endl;
+    }
+  };
+
+  auto model = OgaModel::Create(PHI2_PATH);
+  auto tokenizer = OgaTokenizer::Create(*model);
+  auto tokenizer_stream = OgaTokenizerStream::Create(*tokenizer);
+
+  const char* input_string = "She sells sea shells by the sea shore.";
+  auto input_sequences = OgaSequences::Create();
+  tokenizer->Encode(input_string, *input_sequences);
+  auto params = OgaGeneratorParams::Create(*model);
+  params->SetInputSequences(*input_sequences);
+  params->SetSearchOption("max_length", 40);
+
+  auto generator = OgaGenerator::Create(*model, *params);
+  EXPECT_EQ(generator->IsSessionTerminated(), false);
+  std::vector<std::thread> threads;
+  threads.push_back(std::thread(GenerateOutput, generator.get(), std::move(tokenizer_stream)));
+  threads.push_back(std::thread(GeneratorSetTerminateCall, generator.get()));
+
+  for (auto& th : threads) {
+    std::cout << "Waiting for threads completion" << std::endl;
+    th.join();  // Wait for each thread to finish
+  }
+  EXPECT_EQ(generator->IsSessionTerminated(), true);
+  // Unset terminate
+  generator->SetRuntimeOption("terminate_session", "0");
+  EXPECT_EQ(generator->IsSessionTerminated(), false);
+#endif
+}
+
 #if TEST_PHI2
 
 struct Phi2Test {