Fix ort global variable destruction order

src/generators.cpp

@@ -11,6 +11,23 @@
 
 namespace Generators {
 
+static bool _ = (Ort::InitApi(), false);
+
+OrtGlobals::OrtGlobals() : env_{OrtEnv::Create()} {}
+
+std::unique_ptr<OrtGlobals>& GetOrtGlobals() {
+  static auto globals = std::make_unique<OrtGlobals>();
+  return globals;
+}
+
+void Shutdown() {
+  GetOrtGlobals().reset();
+}
+
+OrtEnv& GetOrtEnv() {
+  return *GetOrtGlobals()->env_;
+}
+
 // IEEE 752-2008 binary16 format, 1 sign bit, 5 bit exponent, 10 bit fraction
 float Float16ToFloat32(uint16_t v) {
   // Extract sign, exponent, and fraction from numpy.float16

src/generators.h

@@ -114,6 +114,23 @@ struct Generator {
   bool computed_logits_{};  // Set to true in ComputeLogits() and false after appending a token to ensure a 1 to 1 call ratio
 };
 
+struct OrtGlobals {
+  OrtGlobals();
+
+  std::unique_ptr<OrtEnv> env_;
+#if USE_CUDA
+  std::unique_ptr<OrtMemoryInfo> memory_info_cuda_;
+  std::unique_ptr<Ort::Allocator> allocator_cuda_;
+#endif
+ private:
+  OrtGlobals(const OrtGlobals&) = delete;
+  void operator=(const OrtGlobals&) = delete;
+};
+
+std::unique_ptr<OrtGlobals>& GetOrtGlobals();
+void Shutdown();  // Do this once at exit, Ort code will fail after this call
+OrtEnv& GetOrtEnv();
+
 std::shared_ptr<Model> CreateModel(OrtEnv& ort_env, const char* config_path);
 std::shared_ptr<GeneratorParams> CreateGeneratorParams(const Model& model);
 std::shared_ptr<GeneratorParams> CreateGeneratorParams();  // For benchmarking purposes only

src/models/model.cpp

@@ -187,14 +187,12 @@ std::vector<std::string> Tokenizer::DecodeBatch(std::span<const int32_t> sequenc
 // has been destroyed. Without this, we will crash in the Onnxruntime BFCArena code when deleting tensors due to the
 // arena already being destroyed.
 Ort::Allocator* GetCudaAllocator(OrtSession& session) {
-  static std::unique_ptr<OrtMemoryInfo> memory_info_cuda_;
-  static std::unique_ptr<Ort::Allocator> allocator_cuda_;
-
-  if (!allocator_cuda_) {
-    memory_info_cuda_ = OrtMemoryInfo::Create("Cuda", OrtAllocatorType::OrtDeviceAllocator, 0, OrtMemType::OrtMemTypeDefault);
-    allocator_cuda_ = Ort::Allocator::Create(session, *memory_info_cuda_);
+  auto& globals = *GetOrtGlobals();
+  if (!globals.allocator_cuda_) {
+    globals.memory_info_cuda_ = OrtMemoryInfo::Create("Cuda", OrtAllocatorType::OrtDeviceAllocator, 0, OrtMemType::OrtMemTypeDefault);
+    globals.allocator_cuda_ = Ort::Allocator::Create(session, *globals.memory_info_cuda_);
   }
-  return allocator_cuda_.get();
+  return globals.allocator_cuda_.get();
 }
 #endif
 

src/ort_genai_c.cpp

@@ -13,16 +13,6 @@
 
 namespace Generators {
 
-std::unique_ptr<OrtEnv> g_ort_env;
-
-OrtEnv& GetOrtEnv() {
-  if (!g_ort_env) {
-    Ort::InitApi();
-    g_ort_env = OrtEnv::Create();
-  }
-  return *g_ort_env;
-}
-
 struct Result {
   explicit Result(const char* what) : what_{what} {}
   std::string what_;
@@ -39,6 +29,13 @@ extern "C" {
     return reinterpret_cast<OgaResult*>(std::make_unique<Generators::Result>(e.what()).release()); \
   }
 
+OgaResult* OGA_API_CALL OgaShutdown() {
+  OGA_TRY
+  Generators::Shutdown();
+  return nullptr;
+  OGA_CATCH
+}
+
 const char* OGA_API_CALL OgaResultGetError(const OgaResult* result) {
   return reinterpret_cast<const Generators::Result*>(result)->what_.c_str();
 }

src/ort_genai_c.h

@@ -40,6 +40,12 @@ typedef struct OgaSequences OgaSequences;
 typedef struct OgaTokenizer OgaTokenizer;
 typedef struct OgaTokenizerStream OgaTokenizerStream;
 
+/* \brief Call this on process exit to cleanly shutdown the genai library & its onnxruntime usage
+ * \return Error message contained in the OgaResult. The const char* is owned by the OgaResult
+ *         and can will be freed when the OgaResult is destroyed.
+ */
+OGA_EXPORT OgaResult* OGA_API_CALL OgaShutdown();
+
 /*
  * \param[in] result OgaResult that contains the error message.
  * \return Error message contained in the OgaResult. The const char* is owned by the OgaResult

src/python/python.cpp

@@ -24,15 +24,6 @@ pybind11::array_t<T> ToPython(std::span<T> v) {
 
 namespace Generators {
 
-std::unique_ptr<OrtEnv> g_ort_env;
-
-OrtEnv& GetOrtEnv() {
-  if (!g_ort_env) {
-    g_ort_env = OrtEnv::Create();
-  }
-  return *g_ort_env;
-}
-
 // A roaming array is one that can be in CPU or GPU memory, and will copy the memory as needed to be used from anywhere
 template <typename T>
 struct PyRoamingArray : RoamingArray<T> {
@@ -186,6 +177,14 @@ PYBIND11_MODULE(onnxruntime_genai, m) {
 
     )pbdoc";
 
+  // Add a cleanup call to happen before global variables are destroyed
+  static int unused{};  // The capsule needs something to reference
+  pybind11::capsule cleanup(
+      &unused, "cleanup", [](PyObject*) {
+        Generators::Shutdown();
+      });
+  m.add_object("_cleanup", cleanup);
+
   // So that python users can catch OrtExceptions specifically
   pybind11::register_exception<Ort::Exception>(m, "OrtException");
 
@@ -203,9 +202,6 @@ PYBIND11_MODULE(onnxruntime_genai, m) {
       .def("set_search_options", &PyGeneratorParams::SetSearchOptions)  // See config.h 'struct Search' for the options
       .def("try_use_cuda_graph_with_max_batch_size", &PyGeneratorParams::TryUseCudaGraphWithMaxBatchSize);
 
-  // We need to init the OrtApi before we can use it
-  Ort::InitApi();
-
   pybind11::class_<TokenizerStream>(m, "TokenizerStream")
       .def("decode", [](TokenizerStream& t, int32_t token) { return t.Decode(token); });
 

src/smartptrs.h

@@ -115,7 +115,7 @@ struct cuda_stream_holder {
 #else
 struct cuda_stream_holder {
   void Create() {
-    assert(false);
+    throw std::runtime_error("Trying to create a cuda stream in a non cuda build");
   }
 
   operator cudaStream_t() const { return v_; }

test/main.cpp

@@ -5,20 +5,16 @@
 #include <generators.h>
 #include <iostream>
 
-extern std::unique_ptr<OrtEnv> g_ort_env;
-
 int main(int argc, char** argv) {
   std::cout << "Generators Utility Library" << std::endl;
   std::cout << "Initializing OnnxRuntime... ";
   std::cout.flush();
   try {
-    Ort::InitApi();
-    g_ort_env = OrtEnv::Create();
     std::cout << "done" << std::endl;
     ::testing::InitGoogleTest(&argc, argv);
     int result = RUN_ALL_TESTS();
     std::cout << "Shutting down OnnxRuntime... ";
-    g_ort_env.reset();
+    Generators::Shutdown();
     std::cout << "done" << std::endl;
     return result;
   } catch (const std::exception& e) {

test/model_tests.cpp

@@ -10,7 +10,6 @@
 #ifndef MODEL_PATH
 #define MODEL_PATH "../../test/test_models/"
 #endif
-std::unique_ptr<OrtEnv> g_ort_env;
 
 // To generate this file:
 // python convert_generation.py --model_type gpt2 -m hf-internal-testing/tiny-random-gpt2 --output tiny_gpt2_greedysearch_fp16.onnx --use_gpu --max_length 20
@@ -33,7 +32,7 @@ TEST(ModelTests, GreedySearchGptFp32) {
   // To generate this file:
   // python convert_generation.py --model_type gpt2 -m hf-internal-testing/tiny-random-gpt2 --output tiny_gpt2_greedysearch_fp16.onnx --use_gpu --max_length 20
   // And copy the resulting gpt2_init_past_fp32.onnx file into these two files (as it's the same for gpt2)
-  auto model = Generators::CreateModel(*g_ort_env,
+  auto model = Generators::CreateModel(Generators::GetOrtEnv(),
                                        MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
 
   auto params = Generators::CreateGeneratorParams(*model);
@@ -74,7 +73,7 @@ TEST(ModelTests, BeamSearchGptFp32) {
   //        --output tiny_gpt2_beamsearch_fp16.onnx --use_gpu --max_length 20
   // (with separate_gpt2_decoder_for_init_run set to False as it is now set to True by default)
 
-  auto model = Generators::CreateModel(*g_ort_env, MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
+  auto model = Generators::CreateModel(Generators::GetOrtEnv(), MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
 
   auto params = Generators::CreateGeneratorParams(*model);
   params->batch_size = static_cast<int>(input_ids_shape[0]);
@@ -119,7 +118,7 @@ void Test_GreedySearch_Gpt_Cuda(const char* model_path, const char* model_label)
       0, 0, 0, 52, 204, 204, 204, 204, 204, 204,
       0, 0, 195, 731, 731, 114, 114, 114, 114, 114};
 
-  auto model = Generators::CreateModel(*g_ort_env, model_path);
+  auto model = Generators::CreateModel(Generators::GetOrtEnv(), model_path);
 
   auto params = Generators::CreateGeneratorParams(*model);
   params->batch_size = static_cast<int>(input_ids_shape[0]);
@@ -164,7 +163,7 @@ void Test_BeamSearch_Gpt_Cuda(const char* model_path, const char* model_label) {
   // python convert_generation.py --model_type gpt2 -m hf-internal-testing/tiny-random-gpt2
   //        --output tiny_gpt2_beamsearch_fp16.onnx --use_gpu --max_length 20
   // (with separate_gpt2_decoder_for_init_run set to False as it is now set to True by default)
-  auto model = Generators::CreateModel(*g_ort_env, model_path);
+  auto model = Generators::CreateModel(Generators::GetOrtEnv(), model_path);
 
   auto params = Generators::CreateGeneratorParams(*model);
   params->batch_size = static_cast<int>(input_ids_shape[0]);
@@ -215,7 +214,7 @@ Print all primes between 1 and n
 
   std::cout << "With prompt:" << prompt << "\r\n";
 
-  auto model = Generators::CreateModel(*g_ort_env, MODEL_PATH "phi-2");
+  auto model = Generators::CreateModel(Generators::GetOrtEnv(), MODEL_PATH "phi-2");
   auto tokenizer = model->CreateTokenizer();
   auto tokens = tokenizer->Encode(prompt);
 
@@ -253,7 +252,7 @@ Print all primes between 1 and n
 
   std::cout << "With prompt:" << prompt << "\r\n";
 
-  auto model = Generators::CreateModel(*g_ort_env, MODEL_PATH "phi-2");
+  auto model = Generators::CreateModel(Generators::GetOrtEnv(), MODEL_PATH "phi-2");
   auto tokenizer = model->CreateTokenizer();
   auto tokens = tokenizer->Encode(prompt);
 

test/sampling_benchmark.cpp

@@ -14,13 +14,11 @@
 #define MODEL_PATH "../../test/test_models/"
 #endif
 
-extern std::unique_ptr<OrtEnv> g_ort_env;
-
 // Defined in sampling_tests.cpp
 void CreateRandomLogits(float* logits, int num_large, int vocab_size, int batch_size, std::mt19937& engine);
 
 TEST(Benchmarks, BenchmarkRandomizedSamplingTopPCpu) {
-  auto model = Generators::CreateModel(*g_ort_env, MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
+  auto model = Generators::CreateModel(Generators::GetOrtEnv(), MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
   int vocab_size = 32000;  // vocab size of llama
   int batch_size = 1;
   std::vector<int32_t> input_ids{0, 1, 2, 3, 4};
@@ -54,7 +52,7 @@ TEST(Benchmarks, BenchmarkRandomizedSamplingTopPCpu) {
 }
 
 TEST(Benchmarks, BenchmarkRandomizedSamplingTopKCpu) {
-  auto model = Generators::CreateModel(*g_ort_env, MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
+  auto model = Generators::CreateModel(Generators::GetOrtEnv(), MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
   int vocab_size = 32000;  // vocab size of llama
   int batch_size = 1;
   int k = 5;
@@ -91,7 +89,7 @@ TEST(Benchmarks, BenchmarkRandomizedSamplingTopKCpu) {
 }
 
 TEST(Benchmarks, BenchmarkRandomizedSamplingTopPAndKCpu) {
-  auto model = Generators::CreateModel(*g_ort_env, MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
+  auto model = Generators::CreateModel(Generators::GetOrtEnv(), MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
   int vocab_size = 32000;  // vocab size of llama
   int batch_size = 1;
   float p = 0.95f;
@@ -132,7 +130,7 @@ TEST(Benchmarks, BenchmarkRandomizedSamplingTopPAndKCpu) {
 #include "tests_helper.cuh"
 
 TEST(Benchmarks, BenchmarkRandomizedSamplingTopPCuda) {
-  auto model = Generators::CreateModel(*g_ort_env, MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
+  auto model = Generators::CreateModel(Generators::GetOrtEnv(), MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
   int vocab_size = 32000;  // vocab size of llama
   int batch_size = 1;
   std::vector<int32_t> input_ids{0, 1, 2, 3, 4};
@@ -175,10 +173,7 @@ TEST(Benchmarks, BenchmarkRandomizedSamplingTopPCuda) {
 }
 
 TEST(Benchmarks, BenchmarkRandomizedSamplingTopKCuda) {
-  std::unique_ptr<OrtEnv> g_ort_env;
-  Ort::InitApi();
-  g_ort_env = OrtEnv::Create();
-  auto model = Generators::CreateModel(*g_ort_env, MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
+  auto model = Generators::CreateModel(Generators::GetOrtEnv(), MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
   int vocab_size = 32000;  // vocab size of llama
   int batch_size = 1;
   int k = 5;
@@ -218,10 +213,7 @@ TEST(Benchmarks, BenchmarkRandomizedSamplingTopKCuda) {
 }
 
 TEST(Benchmarks, BenchmarkRandomizedSamplingTopPAndKCuda) {
-  std::unique_ptr<OrtEnv> g_ort_env;
-  Ort::InitApi();
-  g_ort_env = OrtEnv::Create();
-  auto model = Generators::CreateModel(*g_ort_env, MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
+  auto model = Generators::CreateModel(Generators::GetOrtEnv(), MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
   int vocab_size = 32000;  // vocab size of llama
   int batch_size = 1;
   float p = 0.95f;
@@ -266,10 +258,7 @@ TEST(Benchmarks, BenchmarkRandomizedSamplingTopPAndKCuda) {
 }
 
 TEST(Benchmarks, BenchmarkRandomizedSelectTopCuda) {
-  std::unique_ptr<OrtEnv> g_ort_env;
-  Ort::InitApi();
-  g_ort_env = OrtEnv::Create();
-  auto model = Generators::CreateModel(*g_ort_env, MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
+  auto model = Generators::CreateModel(Generators::GetOrtEnv(), MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32");
   int vocab_size = 32000;  // vocab size of llama
   int batch_size = 12;
   std::vector<int32_t> input_ids{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}; // Needs to match batch_size