Add session config to return an error if model needs to be compiled (microsoft#24416)

adrianlizarraga · web-flow · commit fcb4866d15d8 · 2025-04-18T09:02:03.000-07:00
### Description
Adds session config option (`"session.disable_model_compile"`) that
disables model compilation during session initialization.

If this option is set to "1", inference session creation will fail with
error code ORT_MODEL_REQUIRES_COMPILATION if compilation is required to
run the model on any Execution Provider added to the session. Only the
following kinds of models are valid when this option is set to "1":
- Pre-compiled models that have EPContext nodes for the compiling
Execution Providers in the session.
- Non-compiled models that run only on non-compiling Execution
Providers, like CPU EP.

### Example usage
The following example (taken from a unit test) tries to load a model
that requires compilation with a session that disables compilation. The
session creation fails with error code `ORT_MODEL_REQUIRES_COMPILATION`.
Then, the example compiles the model and loads the compiled model
successfully.

```C++
  // Taken from a unit test ...

  // Initialize session options with QNN EP
  Ort::SessionOptions session_options;
  ProviderOptions provider_options;
  provider_options["backend_type"] = "htp";
  provider_options["offload_graph_io_quantization"] = "0";

  session_options.AppendExecutionProvider("QNN", provider_options);
  session_options.AddConfigEntry(kOrtSessionOptionsDisableEpCompile, "1");  // Disable model compilation!

  // Create an inference session that fails with error ORT_MODEL_REQUIRES_COMPILATION
  try {
    Ort::Session session(*ort_env, input_model_file, session_options);
    FAIL() &lt;&lt; "Expected Session creation to fail but it succeeded";  // Should not get here!
  } catch (const Ort::Exception&amp; excpt) {
    OrtErrorCode error_code = excpt.GetOrtErrorCode();
    std::string_view error_msg = excpt.what();
    ASSERT_EQ(error_code, ORT_MODEL_REQUIRES_COMPILATION);
    ASSERT_THAT(error_msg, testing::HasSubstr(kQnnExecutionProvider));
  }

  // Session creation failed because the model was not pre-compiled.
  // Try to compile it now.

  // Create model compilation options from the session options.
  Ort::ModelCompilationOptions compile_options(*ort_env, session_options);
  compile_options.SetInputModelPath(input_model_file);
  compile_options.SetOutputModelPath(output_model_file);

  // Compile the model.
  Ort::Status status = Ort::CompileModel(*ort_env, compile_options);
  ASSERT_TRUE(status.IsOK()) &lt;&lt; status.GetErrorMessage();

  // Should be able to create a session with the compiled model and the original session options.
  Ort::Session session(*ort_env, output_model_file, session_options);
```

### Motivation and Context
Compiling models can take a very long time. Want to have a session
option that requires input models that do not need to be compiled.
diff --git a/include/onnxruntime/core/common/status.h b/include/onnxruntime/core/common/status.h
@@ -45,6 +45,7 @@ enum StatusCode {
   INVALID_GRAPH = 10,
   EP_FAIL = 11,
   MODEL_LOAD_CANCELED = 12,
+  MODEL_REQUIRES_COMPILATION = 13,
 };
 
 constexpr const char* StatusCodeToString(StatusCode status) noexcept {
@@ -75,6 +76,8 @@ constexpr const char* StatusCodeToString(StatusCode status) noexcept {
       return "EP_FAIL";
     case StatusCode::MODEL_LOAD_CANCELED:
       return "MODEL_LOAD_CANCELED";
+    case StatusCode::MODEL_REQUIRES_COMPILATION:
+      return "MODEL_REQUIRES_COMPILATION";
     default:
       return "GENERAL ERROR";
   }
@@ -109,6 +112,8 @@ constexpr HRESULT StatusCodeToHRESULT(StatusCode status) noexcept {
       return HRESULT_FROM_WIN32(ERROR_INTERNAL_ERROR);
     case StatusCode::MODEL_LOAD_CANCELED:
       return HRESULT_FROM_WIN32(ERROR_CANCELLED);
+    case StatusCode::MODEL_REQUIRES_COMPILATION:
+      return HRESULT_FROM_WIN32(ERROR_NOT_SUPPORTED);
     default:
       return E_FAIL;
   }
diff --git a/include/onnxruntime/core/session/onnxruntime_c_api.h b/include/onnxruntime/core/session/onnxruntime_c_api.h
@@ -256,6 +256,7 @@ typedef enum OrtErrorCode {
   ORT_INVALID_GRAPH,
   ORT_EP_FAIL,
   ORT_MODEL_LOAD_CANCELED,
+  ORT_MODEL_REQUIRES_COMPILATION,
 } OrtErrorCode;
 
 typedef enum OrtOpAttrType {
diff --git a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h
@@ -348,3 +348,19 @@ static const char* const kOrtSessionOptionsQDQMatMulNBitsAccuracyLevel = "sessio
 // “Default”: OS determines the scheduling priority and processor performance to service this workload. [Default]
 // “Efficient”: OS treats this workload is efficiency oriented with low scheduling priority and efficient processor performance.
 static const char* const kOrtEpDynamicOptionsWorkloadType = "ep.dynamic.workload_type";
+
+// Disables model compilation during session initialization.
+//
+// If this option is set to "1", inference session creation will fail with error code ORT_MODEL_REQUIRES_COMPILATION
+// if compilation is required to run the model on any Execution Provider added to the session.
+// Only the following kinds of models are valid when this option is set to "1":
+//   - Pre-compiled models that have EPContext nodes for the compiling Execution Providers in the session.
+//   - Non-compiled models that run only on non-compiling Execution Providers, like CPU EP.
+//
+// See \href https://onnxruntime.ai/docs/execution-providers/EP-Context-Design.html for details about
+// compiled models with EPContext nodes.
+//
+// Option values:
+// - "0": EP compile is not disabled. [DEFAULT]
+// - "1": EP compile is disabled.
+static const char* const kOrtSessionOptionsDisableModelCompile = "session.disable_model_compile";
diff --git a/onnxruntime/core/framework/graph_partitioner.cc b/onnxruntime/core/framework/graph_partitioner.cc
@@ -423,7 +423,8 @@ static Status PartitionOnnxFormatModelImpl(Graph& graph, FuncManager& func_mgr,
                                            const layout_transformation::DebugGraphFn& debug_graph_fn,
                                            const CheckLoadCancellationFn& check_load_cancellation_fn,
                                            const logging::Logger& logger, IResourceAccountant* resource_accountant,
-                                           const GraphOptimizerRegistry& graph_optimizer_registry) {
+                                           const GraphOptimizerRegistry& graph_optimizer_registry,
+                                           bool disable_model_compile) {
   // handle testing edge case where optimizers or constant lifting results in graph with no nodes.
   // doing it here saves all providers checking for this in GetCapability
   if (graph.NumberOfNodes() == 0) {
@@ -440,7 +441,7 @@ static Status PartitionOnnxFormatModelImpl(Graph& graph, FuncManager& func_mgr,
                                                        transform_layout_fn, debug_graph_fn,
                                                        check_load_cancellation_fn,
                                                        logger, resource_accountant,
-                                                       graph_optimizer_registry));
+                                                       graph_optimizer_registry, disable_model_compile));
     }
   }
 
@@ -529,6 +530,16 @@ static Status PartitionOnnxFormatModelImpl(Graph& graph, FuncManager& func_mgr,
     }
   }
 
+  // Helper function that returns true if any of the nodes assigned to a compiling EP is not already compiled.
+  auto graph_viewer_has_non_compiled_node = [](const GraphViewer& graph_viewer) -> bool {
+    for (const auto& node : graph_viewer.Nodes()) {
+      if (node.OpType() != "EPContext") {
+        return true;
+      }
+    }
+    return false;
+  };
+
   // NOTE: if mode_ is kAssignOnly, nodes_to_compile will be empty at this point due to logic in PlaceNode
   // even with single node, EP might still want to compile it.
   // for example, it want to JIT an optimized kernel for LSTM with a given shape.
@@ -548,7 +559,14 @@ static Status PartitionOnnxFormatModelImpl(Graph& graph, FuncManager& func_mgr,
       for (size_t j = 0, end = nodes_to_compile.size(); j < end; j++) {
         auto* node = nodes_to_compile[j];
         const auto& cur_capability = *capabilities_to_compile[j];
-        viewers.push_back(std::make_unique<GraphViewer>(graph, *cur_capability.sub_graph));
+        auto graph_viewer = std::make_unique<GraphViewer>(graph, *cur_capability.sub_graph);
+
+        if (disable_model_compile && graph_viewer_has_non_compiled_node(*graph_viewer)) {
+          return ORT_MAKE_STATUS(ONNXRUNTIME, MODEL_REQUIRES_COMPILATION, "User disabled EP compilation but EP '",
+                                 type, "' needs to compile one or more nodes.");
+        }
+
+        viewers.push_back(std::move(graph_viewer));
         nodes_and_viewers.push_back(IExecutionProvider::FusedNodeAndGraph{*node, *viewers.back()});
       }
 
@@ -896,7 +914,7 @@ static Status PartitionOnnxFormatModel(const PartitionParams& partition_params,
                                        KernelRegistryManager& kernel_registry_manager,
                                        const std::optional<ResourceAccountantMap>& acc_map,
                                        const GraphOptimizerRegistry& graph_optimizer_registry,
-                                       const logging::Logger& logger) {
+                                       const logging::Logger& logger, bool disable_model_compile) {
   bool modified_graph = false;
 
   auto& graph = partition_params.graph.get();
@@ -921,7 +939,8 @@ static Status PartitionOnnxFormatModel(const PartitionParams& partition_params,
                                                        transform_layout_function,
                                                        partition_params.debug_graph_fn,
                                                        check_load_cancellation_fn,
-                                                       logger, resource_accountant, graph_optimizer_registry));
+                                                       logger, resource_accountant, graph_optimizer_registry,
+                                                       disable_model_compile));
     }
 
     // expand any nodes that have an ONNX function definition but no matching ORT kernel.
@@ -1195,8 +1214,10 @@ Status GraphPartitioner::Partition(Graph& graph, FuncManager& func_mgr,
     std::optional<ResourceAccountantMap> ep_acc_map;
     ORT_RETURN_IF_ERROR(NodeStatsRecorder::CreateAccountants(config_options, graph.ModelPath(), ep_acc_map));
 
+    bool disable_model_compile = config_options.GetConfigOrDefault(kOrtSessionOptionsDisableModelCompile, "0") == "1";
     ORT_RETURN_IF_ERROR(PartitionOnnxFormatModel(partition_params, mode, providers_, kernel_registry_mgr_,
-                                                 ep_acc_map, *graph_optimizer_registry_, logger));
+                                                 ep_acc_map, *graph_optimizer_registry_, logger,
+                                                 disable_model_compile));
 
     if (ep_context_gen_options.enable) {
       ORT_RETURN_IF_ERROR(CreateEpContextModel(providers_, graph, ep_context_gen_options, logger));
diff --git a/onnxruntime/core/session/model_compilation_options.cc b/onnxruntime/core/session/model_compilation_options.cc
@@ -23,6 +23,7 @@ ModelCompilationOptions::ModelCompilationOptions(const OrtEnv& env, const OrtSes
 
   // Shouldn't fail because the key/value strings are below the maximum string length limits in ConfigOptions.
   ORT_ENFORCE(session_options_.value.config_options.AddConfigEntry(kOrtSessionOptionEpContextEnable, "1").IsOK());
+  ORT_ENFORCE(session_options_.value.config_options.AddConfigEntry(kOrtSessionOptionsDisableModelCompile, "0").IsOK());
 }
 
 void ModelCompilationOptions::SetInputModelPath(const std::string& input_model_path) {
@@ -195,6 +196,7 @@ Status ModelCompilationOptions::CheckOutputModelSettings() const {
 
 Status ModelCompilationOptions::Check() const {
   ORT_ENFORCE(session_options_.value.ep_context_gen_options.enable);
+  ORT_ENFORCE(session_options_.value.config_options.GetConfigOrDefault(kOrtSessionOptionsDisableModelCompile, "0") == "0");
   ORT_RETURN_IF_ERROR(CheckInputModelSettings());
   ORT_RETURN_IF_ERROR(CheckOutputModelSettings());
   return Status::OK();
diff --git a/onnxruntime/test/providers/qnn/qnn_ep_context_test.cc b/onnxruntime/test/providers/qnn/qnn_ep_context_test.cc
@@ -272,6 +272,62 @@ static void CheckEpContextNodeCounts(void* model_buffer, size_t model_buffer_siz
   std::filesystem::remove(output_model_path);
 }
 
+// Test workflow that:
+//   - Creates session that disables EP compilation.
+//   - Session creation fails because input model is not pre-compiled.
+//   - Uses OrtCompileApi to compile the model.
+//   - Recreates session with the compiled model.
+TEST_F(QnnHTPBackendTests, CompileApi_DisableEpCompile_ThenCompileExplicitly) {
+  const ORTCHAR_T* input_model_file = ORT_TSTR("./compileapi_disable_compile_input.onnx");
+  const ORTCHAR_T* output_model_file = ORT_TSTR("./compileapi_disable_compile_output.onnx");
+  std::filesystem::remove(input_model_file);
+  std::filesystem::remove(output_model_file);
+
+  // Create a test model and save it to a file.
+  TestModel test_model;
+  CreateTestModel(BuildGraphWithQAndNonQ(false), 21, logging::Severity::kERROR, test_model);
+  ASSERT_STATUS_OK(test_model.Save(input_model_file));
+
+  // Initialize session options with QNN EP
+  Ort::SessionOptions so;
+  ProviderOptions provider_options;
+  provider_options["backend_type"] = "htp";
+  provider_options["offload_graph_io_quantization"] = "0";
+
+  so.AppendExecutionProvider("QNN", provider_options);
+  so.AddConfigEntry(kOrtSessionOptionsDisableModelCompile, "1");  // Disable model compilation!
+
+  // Create an inference session that fails with error ORT_MODEL_REQUIRES_COMPILATION
+  try {
+    Ort::Session session(*ort_env, input_model_file, so);
+    FAIL() << "Expected Session creation to fail but it succeeded";  // Should not get here!
+  } catch (const Ort::Exception& excpt) {
+    OrtErrorCode error_code = excpt.GetOrtErrorCode();
+    std::string_view error_msg = excpt.what();
+    ASSERT_EQ(error_code, ORT_MODEL_REQUIRES_COMPILATION);
+    ASSERT_THAT(error_msg, testing::HasSubstr(kQnnExecutionProvider));
+  }
+
+  // Session creation failed because the model was not pre-compiled.
+  // Try to compile it now.
+
+  // Create model compilation options from the session options.
+  Ort::ModelCompilationOptions compile_options(*ort_env, so);
+  compile_options.SetInputModelPath(input_model_file);
+  compile_options.SetOutputModelPath(output_model_file);
+
+  // Compile the model.
+  Ort::Status status = Ort::CompileModel(*ort_env, compile_options);
+  ASSERT_TRUE(status.IsOK()) << status.GetErrorMessage();
+
+  // Make sure the compiled model was generated and has the expected number of EPContext nodes.
+  ASSERT_TRUE(std::filesystem::exists(output_model_file));
+  CheckEpContextNodeCounts(output_model_file, 2, 2);
+
+  // Should be able to create a session with the compiled model and the original session options.
+  EXPECT_NO_THROW((Ort::Session(*ort_env, output_model_file, so)));
+}
+
 // Test using the CompileModel() API with settings:
 //   - input model file
 //   - output model file
