Add Set terminate Option for user

src/generators.cpp

@@ -154,6 +154,9 @@ Generator::Generator(const Model& model, const GeneratorParams& params) : model_
 }
 
 void Generator::ComputeLogits() {
+  if (state_->params_->session_terminated) {
+    throw std::runtime_error("Session in Terminated state, exiting!");
+  }
   if (computed_logits_)
     throw std::runtime_error("ComputeLogits called again without calling GenerateNextToken first");
 
@@ -172,6 +175,9 @@ void Generator::ComputeLogits() {
 }
 
 bool Generator::IsDone() const {
+  if (state_->params_->session_terminated) {
+    throw std::runtime_error("Session in Terminated state, exiting!");
+  }
   if (computed_logits_)
     throw std::runtime_error("IsDone() can't be called in the middle of processing logits");
 
@@ -184,6 +190,9 @@ bool Generator::IsDone() const {
 }
 
 void Generator::GenerateNextToken() {
+  if (state_->params_->session_terminated) {
+    throw std::runtime_error("Session in Terminated state, exiting!");
+  }
   if (!computed_logits_)
     throw std::runtime_error("Must call ComputeLogits before GenerateNextToken");
   computed_logits_ = false;

src/generators.h

@@ -66,6 +66,7 @@ struct GeneratorParams : std::enable_shared_from_this<GeneratorParams>, LeakChec
   int batch_size{1};
   int max_batch_size{0};
   bool use_cuda_graph{};
+  mutable bool session_terminated{false};
   int sequence_length{};
   int BatchBeamSize() const { return search.num_beams * batch_size; }
 

src/models/model.cpp

@@ -37,9 +37,9 @@ static std::string CurrentModulePath() {
 
 namespace Generators {
 
-State::State(const GeneratorParams& params, const Model& model)
+State::State(const GeneratorParams& params_, const Model& model)
     : model_{model},
-      params_{params.shared_from_this()} {}
+      params_{params_.shared_from_this()} {}
 
 void State::Run(OrtSession& session, OrtRunOptions& run_options, int new_batch_size) {
   auto captured_graph_info = GetCapturedGraphInfo();
@@ -76,7 +76,20 @@ void State::Run(OrtSession& session, OrtRunOptions& run_options, int new_batch_s
   }
 }
 
+void State::SetTerminate() {
+  params_->session_terminated = true;
+  model_.run_options_->SetTerminate();
+}
+
+void State::UnsetTerminate() {
+  params_->session_terminated = false;
+  model_.run_options_->UnsetTerminate();
+}
+
 OrtValue* State::GetInput(const char* name) {
+  if (params_->session_terminated) {
+    throw std::runtime_error("Session in Terminated state, exiting!");
+  }
   for (size_t i = 0; i < input_names_.size(); i++) {
     if (std::strcmp(input_names_[i], name) == 0) {
       return inputs_[i];
@@ -86,6 +99,9 @@ OrtValue* State::GetInput(const char* name) {
 }
 
 OrtValue* State::GetOutput(const char* name) {
+  if (params_->session_terminated) {
+    throw std::runtime_error("Session in Terminated state, exiting!");
+  }
   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

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

src/ort_genai.h

@@ -244,6 +244,14 @@ struct OgaGenerator : OgaAbstract {
     OgaCheckResult(OgaGenerator_GenerateNextToken(this));
   }
 
+  void SetTerminate() {
+    OgaCheckResult(OgaGenerator_SetTerminate(this));
+  }
+
+  void UnsetTerminate() {
+    OgaCheckResult(OgaGenerator_UnsetTerminate(this));
+  }
+
   size_t GetSequenceCount(size_t index) const {
     return OgaGenerator_GetSequenceCount(this, index);
   }

src/ort_genai_c.cpp

@@ -265,6 +265,22 @@ OgaResult* OGA_API_CALL OgaGenerator_GenerateNextToken(OgaGenerator* generator)
   OGA_CATCH
 }
 
+OgaResult* OGA_API_CALL OgaGenerator_SetTerminate(OgaGenerator* oga_generator) {
+  OGA_TRY
+  auto& generator = *reinterpret_cast<const Generators::Generator*>(oga_generator);
+  generator.state_->SetTerminate();
+  return nullptr;
+  OGA_CATCH
+}
+
+OgaResult* OGA_API_CALL OgaGenerator_UnsetTerminate(OgaGenerator* oga_generator) {
+  OGA_TRY
+  auto& generator = *reinterpret_cast<const Generators::Generator*>(oga_generator);
+  generator.state_->UnsetTerminate();
+  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

@@ -253,6 +253,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_SetTerminate(OgaGenerator* generator);
+OGA_EXPORT OgaResult* OGA_API_CALL OgaGenerator_UnsetTerminate(OgaGenerator* generator);
 
 /*
  * \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