Add support for ArenaCfg configuration options

Author: pskiran1

src/onnxruntime.cc

@@ -302,6 +302,30 @@ ModelState::ModelState(TRITONBACKEND_Model* triton_model)
     }
   }
 
+  // Set use_device_allocator_for_initializers
+  {
+    triton::common::TritonJson::Value params;
+    if (ModelConfig().Find("parameters", &params)) {
+      triton::common::TritonJson::Value json_value;
+      const char* use_device_allocator_for_initializers_key =
+          "session.use_device_allocator_for_initializers";
+      if (params.Find(use_device_allocator_for_initializers_key, &json_value)) {
+        std::string string_value;
+        THROW_IF_BACKEND_MODEL_ERROR(
+            json_value.MemberAsString("string_value", &string_value));
+
+        LOG_MESSAGE(
+            TRITONSERVER_LOG_VERBOSE,
+            (std::string("Configuring ") +
+             use_device_allocator_for_initializers_key + " to " + string_value)
+                .c_str());
+        THROW_IF_BACKEND_MODEL_ORT_ERROR(ort_api->AddSessionConfigEntry(
+            soptions, use_device_allocator_for_initializers_key,
+            string_value.c_str()));
+      }
+    }
+  }
+
   // memory configs
   // enable/disable mem arena
   {
@@ -762,8 +786,90 @@ ModelState::LoadModel(
         rel_cuda_options(cuda_options, ort_api->ReleaseCUDAProviderOptions);
     cuda_options_map["device_id"] = std::to_string(instance_group_device_id);
     cuda_options_map["has_user_compute_stream"] = stream != nullptr ? "1" : "0";
+
+    // Memory arena config
+    OrtArenaCfg* arena_cfg = nullptr;
+    {
+      triton::common::TritonJson::Value params;
+      if (model_config_.Find("parameters", &params)) {
+        triton::common::TritonJson::Value json_value;
+        std::vector<const char*> keys;
+        std::vector<size_t> values;
+        if (params.Find("max_mem", &json_value)) {
+          std::string string_value;
+          THROW_IF_BACKEND_MODEL_ERROR(
+              json_value.MemberAsString("string_value", &string_value));
+          keys.push_back("max_mem");
+          size_t value;
+          RETURN_IF_ERROR(ParseUnsignedLongLongValue(string_value, &value));
+          values.push_back(value);
+        }
+        if (params.Find("arena_extend_strategy", &json_value)) {
+          std::string string_value;
+          THROW_IF_BACKEND_MODEL_ERROR(
+              json_value.MemberAsString("string_value", &string_value));
+          keys.push_back("arena_extend_strategy");
+          size_t value;
+          RETURN_IF_ERROR(ParseUnsignedLongLongValue(string_value, &value));
+          values.push_back(value);
+        }
+        if (params.Find("initial_chunk_size_bytes", &json_value)) {
+          std::string string_value;
+          THROW_IF_BACKEND_MODEL_ERROR(
+              json_value.MemberAsString("string_value", &string_value));
+          keys.push_back("initial_chunk_size_bytes");
+          size_t value;
+          RETURN_IF_ERROR(ParseUnsignedLongLongValue(string_value, &value));
+          values.push_back(value);
+        }
+        if (params.Find("initial_growth_chunk_size_bytes", &json_value)) {
+          std::string string_value;
+          THROW_IF_BACKEND_MODEL_ERROR(
+              json_value.MemberAsString("string_value", &string_value));
+          keys.push_back("initial_growth_chunk_size_bytes");
+          size_t value;
+          RETURN_IF_ERROR(ParseUnsignedLongLongValue(string_value, &value));
+          values.push_back(value);
+        }
+        if (params.Find("max_dead_bytes_per_chunk", &json_value)) {
+          std::string string_value;
+          THROW_IF_BACKEND_MODEL_ERROR(
+              json_value.MemberAsString("string_value", &string_value));
+          keys.push_back("max_dead_bytes_per_chunk");
+          size_t value;
+          RETURN_IF_ERROR(ParseUnsignedLongLongValue(string_value, &value));
+          values.push_back(value);
+        }
+        if (params.Find("max_power_of_two_extend_bytes", &json_value)) {
+          std::string string_value;
+          THROW_IF_BACKEND_MODEL_ERROR(
+              json_value.MemberAsString("string_value", &string_value));
+          keys.push_back("max_power_of_two_extend_bytes");
+          size_t value;
+          RETURN_IF_ERROR(ParseUnsignedLongLongValue(string_value, &value));
+          values.push_back(value);
+        }
+        if (!keys.empty()) {
+          RETURN_IF_ORT_ERROR(ort_api->CreateArenaCfgV2(
+              keys.data(), values.data(), keys.size(), &arena_cfg));
+
+          std::ostringstream oss;
+          for (size_t i = 0; i < keys.size(); ++i) {
+            oss << keys[i] << "=" << values[i] << ", ";
+          }
+          LOG_MESSAGE(
+              TRITONSERVER_LOG_VERBOSE,
+              (std::string("Updated arena config options: ") + oss.str())
+                  .c_str());
+        }
+      }
+    }
+    std::unique_ptr<OrtArenaCfg, decltype(ort_api->ReleaseArenaCfg)>
+        rel_arena_cfg(arena_cfg, ort_api->ReleaseArenaCfg);
     RETURN_IF_ORT_ERROR(ort_api->UpdateCUDAProviderOptionsWithValue(
-        rel_cuda_options.get(), "default_memory_arena_cfg", nullptr));
+        rel_cuda_options.get(), "default_memory_arena_cfg",
+        rel_arena_cfg.get()));
+
     {
       // Parse CUDA EP configurations directly from the parameters field.
       // This is deprecated with adding support for CUDA EP in the