Fixes / improvements based on testing and code analysis

Author: Aditya Rastogi

onnxruntime/core/session/model_compilation_options.cc

@@ -12,6 +12,8 @@
 #include "core/common/path_string.h"
 #include "core/framework/allocator.h"
 #include "core/framework/ep_context_options.h"
+#include "core/platform/env.h"
+#include "core/session/inference_session_utils.h"
 #include "core/session/onnxruntime_session_options_config_keys.h"
 #include "core/session/environment.h"
 
@@ -289,11 +291,30 @@ Status ModelCompilationOptions::Check() const {
     // Model::LoadFromModelEditorApiModel -> Graph::Resolve()
   }
 
+  // ORT_LOAD_CONFIG_FROM_MODEL is not supported for OrtModel input.
+  // Check early so we fail before session construction.
+  if (input_from_model) {
+    const Env& os_env = Env::Default();
+    if (os_env.GetEnvironmentVar(inference_session_utils::kOrtLoadConfigFromModelEnvVar) == "1") {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                             "The environment variable ORT_LOAD_CONFIG_FROM_MODEL=1 is set, but loading "
+                             "config from model is not supported for in-memory OrtModel input. "
+                             "OrtModel is programmatically constructed and has no embedded ORT config. "
+                             "Unset ORT_LOAD_CONFIG_FROM_MODEL or use file/buffer input instead.");
+    }
+  }
+
   // Check output model settings.
   const epctx::ModelGenOptions& ep_context_gen_options = session_options_.value.ep_context_gen_options;
   bool has_no_output_model_location = std::holds_alternative<std::monostate>(
       ep_context_gen_options.output_model_location);
 
+  // Also treat an empty output file path as "no location" since it's not usable.
+  const auto* output_path = ep_context_gen_options.TryGetOutputModelPath();
+  if (!has_no_output_model_location && output_path != nullptr && output_path->empty()) {
+    has_no_output_model_location = true;
+  }
+
   // Determine if we can derive an output path from the input
   bool can_derive_output_path = input_from_file;
   bool model_has_path = false;
@@ -315,9 +336,8 @@ Status ModelCompilationOptions::Check() const {
     return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
                            "OrtModel has no model_path set and no output location was specified. "
                            "Please either: (1) set the model_path on the OrtGraph before adding to OrtModel, "
-                           "(2) call SetOutputModelPath/SetOutputModelBuffer to specify an output location, "
-                           "(3) call SetEpContextEmbedMode(true) to embed EP context in the model, or "
-                           "(4) call SetEpContextBinaryInformation to specify the binary output directory.");
+                           "(2) call SetOutputModelPath/SetOutputModelBuffer to specify an output location, or "
+                           "(3) call SetEpContextEmbedMode(true) to embed EP context in the model.");
   }
 
   if (has_no_output_model_location && can_derive_output_path) {
@@ -360,7 +380,13 @@ Status ModelCompilationOptions::Check() const {
 }
 
 std::string ModelCompilationOptions::GetInputSourceForTelemetry() const {
-  return InputModelComesFromFile() ? "file" : "buffer";
+  if (InputModelComesFromFile()) {
+    return "file";
+  }
+  if (InputModelComesFromOrtModel()) {
+    return "ort_model";
+  }
+  return "buffer";
 }
 
 std::string ModelCompilationOptions::GetOutputTargetForTelemetry() const {

onnxruntime/core/session/model_compilation_options.h

@@ -183,9 +183,9 @@ class ModelCompilationOptions {
   // Telemetry helper methods
 
   /// <summary>
-  /// Returns a string describing the input source type: "file" or "buffer".
+  /// Returns a string describing the input source type: "file", "buffer", or "ort_model".
   /// </summary>
-  /// <returns>"file" or "buffer"</returns>
+  /// <returns>"file", "buffer", or "ort_model"</returns>
   std::string GetInputSourceForTelemetry() const;
 
   /// <summary>

onnxruntime/core/session/utils.cc

@@ -301,10 +301,6 @@ static OrtStatus* CreateSessionAndLoadModelImpl(_In_ const OrtSessionOptions* op
     return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "OrtModel pointer is null");
   }
 
-  const Env& os_env = Env::Default();
-  bool load_config_from_model =
-      os_env.GetEnvironmentVar(inference_session_utils::kOrtLoadConfigFromModelEnvVar) == "1";
-
   // Check EPContext model generation options - OrtModel has no file path by default,
   // so we need explicit output location or embedded model path.
   if (options) {
@@ -325,22 +321,14 @@ static OrtStatus* CreateSessionAndLoadModelImpl(_In_ const OrtSessionOptions* op
           (!ep_ctx_gen_options.HasOutputModelLocation() ||
            (output_model_path != nullptr && output_model_path->empty()))) {
         return OrtApis::CreateStatus(ORT_FAIL,
-                                     "OrtModel loaded without a model_path was configured with EPContext "
-                                     "model generation enabled but without a valid location (e.g., file or buffer) "
-                                     "for the output model. Please specify a valid output location via "
+                                     "OrtModel has no model_path set and no valid output location was specified "
+                                     "for EPContext model generation. "
                                      "SetOutputModelPath/SetOutputModelBuffer, or set the model_path on the "
-                                     "OrtGraph before adding to OrtModel.");
+                                     "OrtGraph before adding it to OrtModel.");
       }
     }
   }
 
-  // Note: load_config_from_model is not applicable for OrtModel since there's no serialized
-  // model to load config from. We treat this as a regular load.
-  if (load_config_from_model) {
-    // Log a warning but proceed - OrtModel doesn't support loading config from model
-    // as it's already in-memory and constructed programmatically
-  }
-
   sess = std::make_unique<onnxruntime::InferenceSession>(
       options == nullptr ? onnxruntime::SessionOptions() : options->value,
       env);
@@ -352,6 +340,30 @@ static OrtStatus* CreateSessionAndLoadModelImpl(_In_ const OrtSessionOptions* op
   }
 #endif
 
+  // Add custom domains for all OrtEpDevice instances to inference session.
+  // The custom domains should be registered before model load for ORT to validate the custom ops.
+  // This mirrors the same block in the file/buffer overload to maintain load-path parity.
+  if (options != nullptr &&
+      options->provider_factories.empty() &&
+      options->value.ep_selection_policy.enable) {
+    InlinedVector<OrtCustomOpDomain*> all_ep_custom_op_domains;
+
+    for (const OrtEpDevice* ep_device : env.GetOrtEpDevices()) {
+      InlinedVector<OrtCustomOpDomain*> domains;
+      ORT_API_RETURN_IF_STATUS_NOT_OK(GetCustomOpDomainsFromEpDevice(*ep_device, domains));
+
+      for (auto domain : domains) {
+        if (ShouldAddDomain(domain, options->custom_op_domains_)) {
+          all_ep_custom_op_domains.push_back(domain);
+        }
+      }
+    }
+
+    if (!all_ep_custom_op_domains.empty()) {
+      ORT_API_RETURN_IF_STATUS_NOT_OK(sess->AddCustomOpDomains(all_ep_custom_op_domains));
+    }
+  }
+
   // Load from OrtModel
   ORT_API_RETURN_IF_STATUS_NOT_OK(sess->Load(*model));
 

onnxruntime/test/shared_lib/test_model_builder_api.cc

@@ -18,6 +18,7 @@
 
 #include "test/shared_lib/test_fixture.h"
 #include "test/shared_lib/utils.h"
+#include "test/util/include/scoped_env_vars.h"
 #include "test/util/include/test_allocator.h"
 
 #include "onnxruntime_config.h"  // generated file in build output dir
@@ -1014,3 +1015,79 @@ TEST(ModelEditorCompileAPITest, CompileFromOrtModelToFile) {
   // Cleanup
   std::filesystem::remove(output_path);
 }
+
+// Test: ORT_LOAD_CONFIG_FROM_MODEL=1 with OrtModel input fails fast with a clear,
+// actionable error message.
+TEST(ModelEditorCompileAPITest, LoadConfigFromModelEnvVarFailsForOrtModel) {
+  // RAII helper saves the current env var value and restores it when the scope exits.
+  onnxruntime::test::ScopedEnvironmentVariables scoped_env(
+      {{"ORT_LOAD_CONFIG_FROM_MODEL", "1"}});
+
+  std::vector<std::unique_ptr<std::vector<float>>> weights;
+  auto model = CreateSimpleGemmModel(weights);
+
+  Ort::SessionOptions session_options;
+  Ort::ModelCompilationOptions compile_options(*ort_env, session_options);
+  compile_options.SetInputModel(static_cast<const OrtModel*>(model));
+  compile_options.SetEpContextEmbedMode(true);
+
+  std::unique_ptr<MockedOrtAllocator> allocator = std::make_unique<MockedOrtAllocator>();
+  void* output_buffer = nullptr;
+  size_t output_size = 0;
+  compile_options.SetOutputModelBuffer(allocator.get(), &output_buffer, &output_size);
+
+  // Should fail with a clear error about the unsupported env-var/input combination.
+  Ort::Status status = Ort::CompileModel(*ort_env, compile_options);
+  EXPECT_FALSE(status.IsOK());
+  EXPECT_THAT(status.GetErrorMessage(), ::testing::HasSubstr("ORT_LOAD_CONFIG_FROM_MODEL=1"));
+  EXPECT_THAT(status.GetErrorMessage(), ::testing::HasSubstr("in-memory OrtModel input"));
+  EXPECT_THAT(status.GetErrorMessage(), ::testing::HasSubstr("Unset ORT_LOAD_CONFIG_FROM_MODEL"));
+
+  if (output_buffer != nullptr) {
+    allocator->Free(output_buffer);
+  }
+}
+
+// Test: Validation error for OrtModel with no model_path, no output location, and no embed mode.
+// Verifies the error message contains the expected remediation guidance.
+TEST(ModelEditorCompileAPITest, NoOutputLocationNoModelPathFails) {
+  std::vector<std::unique_ptr<std::vector<float>>> weights;
+  auto model = CreateSimpleGemmModel(weights);
+
+  Ort::SessionOptions session_options;
+  Ort::ModelCompilationOptions compile_options(*ort_env, session_options);
+  compile_options.SetInputModel(static_cast<const OrtModel*>(model));
+  // Intentionally do NOT call SetEpContextEmbedMode, SetOutputModelPath, or SetOutputModelBuffer
+
+  Ort::Status status = Ort::CompileModel(*ort_env, compile_options);
+  EXPECT_FALSE(status.IsOK());
+  // Should suggest setting output location or model_path, including embed mode as an option
+  EXPECT_THAT(status.GetErrorMessage(), ::testing::HasSubstr("SetOutputModelPath"));
+  EXPECT_THAT(status.GetErrorMessage(), ::testing::HasSubstr("SetEpContextEmbedMode(true)"));
+  // Should NOT suggest SetEpContextBinaryInformation (that alone is not sufficient)
+  EXPECT_THAT(status.GetErrorMessage(), ::testing::Not(::testing::HasSubstr("SetEpContextBinaryInformation")));
+}
+
+// Test: Setting embed mode with buffer output satisfies the output location requirement
+// for OrtModel with no model_path.
+TEST(ModelEditorCompileAPITest, EmbedModeWithBufferOutputSatisfiesValidation) {
+  std::vector<std::unique_ptr<std::vector<float>>> weights;
+  auto model = CreateSimpleGemmModel(weights);
+
+  Ort::SessionOptions session_options;
+  Ort::ModelCompilationOptions compile_options(*ort_env, session_options);
+  compile_options.SetInputModel(static_cast<const OrtModel*>(model));
+  compile_options.SetEpContextEmbedMode(true);
+
+  std::unique_ptr<MockedOrtAllocator> allocator = std::make_unique<MockedOrtAllocator>();
+  void* output_buffer = nullptr;
+  size_t output_size = 0;
+  compile_options.SetOutputModelBuffer(allocator.get(), &output_buffer, &output_size);
+
+  Ort::Status status = Ort::CompileModel(*ort_env, compile_options);
+  EXPECT_TRUE(status.IsOK()) << "CompileModel failed: " << status.GetErrorMessage();
+
+  if (output_buffer != nullptr) {
+    allocator->Free(output_buffer);
+  }
+}