Implement load cancellation flag.

Author: yuslepukhin

Diff for: include/onnxruntime/core/session/onnxruntime_c_api.h

@@ -4898,6 +4898,18 @@ struct OrtApi {
                   _In_ const int64_t* shape, size_t shape_len,
                   ONNXTensorElementDataType type,
                   _Outptr_ OrtValue** out);
+
+  /** \brief changes the setting of the session to cancel loading
+   *
+   * \param[in] options options instance that was passed to the session at creation time.
+   * \param[in] is_cancel setting this to true after model loading process was initiated will
+   *            cancel the loading process within some reasonable time frame.
+   *
+   * \snippet{doc} snippets.dox OrtStatus Return Value
+   *
+   */
+  ORT_API2_STATUS(SessionOptionsSetLoadCancellationFlag, _Inout_ OrtSessionOptions* options,
+                  _In_ bool is_cancel);
 };
 
 /*

Diff for: include/onnxruntime/core/session/onnxruntime_cxx_api.h

@@ -928,6 +928,8 @@ struct SessionOptionsImpl : ConstSessionOptionsImpl<T> {
 
   SessionOptionsImpl& SetExecutionMode(ExecutionMode execution_mode);  ///< Wraps OrtApi::SetSessionExecutionMode
 
+  SessionOptionsImpl& SetLoadCancellationFlag(bool value);  ///< Wraps OrtApi::SessionOptionsSetLoadCancellationFlag
+
   SessionOptionsImpl& SetLogId(const char* logid);     ///< Wraps OrtApi::SetSessionLogId
   SessionOptionsImpl& SetLogSeverityLevel(int level);  ///< Wraps OrtApi::SetSessionLogSeverityLevel
 

Diff for: include/onnxruntime/core/session/onnxruntime_cxx_inline.h

@@ -747,6 +747,12 @@ inline SessionOptionsImpl<T>& SessionOptionsImpl<T>::SetExecutionMode(ExecutionM
   return *this;
 }
 
+template <typename T>
+inline SessionOptionsImpl<T>& SessionOptionsImpl<T>::SetLoadCancellationFlag(bool value) {
+  ThrowOnError(GetApi().SessionOptionsSetLoadCancellationFlag(p_, value));
+  return *this;
+}
+
 template <typename T>
 inline SessionOptionsImpl<T>& SessionOptionsImpl<T>::SetLogId(const char* logid) {
   ThrowOnError(GetApi().SetSessionLogId(this->p_, logid));

Diff for: onnxruntime/core/framework/graph_partitioner.cc

@@ -413,6 +413,8 @@ static Status PartitionOnnxFormatModelImpl(Graph& graph, FuncManager& func_mgr,
     return Status::OK();
   }
 
+  const volatile auto& load_cancellaton_flag = graph_optimizer_registry.GetLoadCancellationFlagRef();
+
   // recurse into nested graphs first to partition bottom up.
   for (auto& node : graph.Nodes()) {
     for (auto& entry : node.GetAttributeNameToMutableSubgraphMap()) {
@@ -506,6 +508,7 @@ static Status PartitionOnnxFormatModelImpl(Graph& graph, FuncManager& func_mgr,
         capabilities_to_complete_fuse.push_back(std::move(capability));
       }
     }
+    ORT_RETURN_IF(load_cancellaton_flag, "Graph partitioning is canceled due to user request.");
   }
 
   // NOTE: if mode_ is kAssignOnly, nodes_to_compile will be empty at this point due to logic in PlaceNode
@@ -561,6 +564,7 @@ static Status PartitionOnnxFormatModelImpl(Graph& graph, FuncManager& func_mgr,
         graph.FinalizeFuseSubGraph(indexed_sub_graph, *node);
       }
     }
+    ORT_RETURN_IF(load_cancellaton_flag, "Graph partitioning is canceled due to user request.");
   }
 
   if (!nodes_to_complete_fuse.empty()) {
@@ -641,6 +645,8 @@ static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_provide
     return Status::OK();
   }
 
+  const volatile bool& load_cancellation_flag = graph_optimizer_registry.GetLoadCancellationFlagRef();
+
   for (auto& node : graph.Nodes()) {
     for (auto& entry : node.GetAttributeNameToMutableSubgraphMap()) {
       Graph* subgraph = entry.second;
@@ -689,6 +695,7 @@ static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_provide
         }
       }
     }
+    ORT_RETURN_IF(load_cancellation_flag, "AOT inlining is canceled due to user request.");
   }
 
   // TODO: Insert version check. We need to collect all the versions
@@ -701,6 +708,7 @@ static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_provide
       if (claimed_by_ep.count(node_index) == 0) {
         ORT_RETURN_IF_ERROR(graph.InlineFunction(*node));
         ++inlined_count;
+        ORT_RETURN_IF(load_cancellation_flag, "AOT inlining is canceled due to user request.");
       } else {
         // OpType is the same as function name.
         auto function_id = function_utils::GetFunctionIdentifier(node->Domain(), node->OpType());
@@ -1032,6 +1040,8 @@ Status GraphPartitioner::InlineFunctionsAOT(Model& model,
     return Status::OK();
   }
 
+  ORT_RETURN_IF(IsLoadCancellationFlagSet(), "AOT inlining is canceled due to user request.");
+
   auto& graph = model.MainGraph();
   InlinedHashSet<std::string> not_inlined;
   do {
@@ -1048,6 +1058,7 @@ Status GraphPartitioner::InlineFunctionsAOT(Model& model,
       break;
     }
 
+    ORT_RETURN_IF(IsLoadCancellationFlagSet(), "AOT inlining is canceled due to user request.");
     ORT_RETURN_IF_ERROR(graph.Resolve());
   } while (true);
 
@@ -1082,6 +1093,9 @@ Status GraphPartitioner::Partition(Graph& graph, FuncManager& func_mgr,
     return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "No provider specified.");
   }
 
+  const volatile auto& load_cancellaton_flag = graph_optimizer_registry_->GetLoadCancellationFlagRef();
+  ORT_RETURN_IF(load_cancellaton_flag, "Graph partitioning is canceled due user request.");
+
 #if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD)
   // fused_kernel_registry is preparing the kernels created on the fly for fused sub graph.
   // It is only visible for current session.

Diff for: onnxruntime/core/framework/graph_partitioner.h

@@ -41,6 +41,10 @@ class GraphPartitioner {
                    Mode mode = Mode::kNormal,
                    const layout_transformation::DebugGraphFn& debug_graph_fn = {}) const;
 
+  bool IsLoadCancellationFlagSet() const noexcept {
+    return graph_optimizer_registry_->GetSessionOptions().IsLoadCancellationFlagSet();
+  }
+
 #ifndef ORT_MINIMAL_BUILD
   /// <summary>
   // Ahead of Time Function inlining. The main purpose of the function is to inline as many

Diff for: onnxruntime/core/framework/session_options.h

@@ -181,6 +181,17 @@ struct SessionOptions {
   void AddCustomOpLibraryHandle(PathString library_name, void* library_handle);
 #endif
 
+  // Load cancellation flag is necessary to be within shared memory as session_options are
+  std::shared_ptr<bool> load_cancellation_flag = std::make_shared<bool>(false);
+
+  bool IsLoadCancellationFlagSet() const noexcept {
+    return *load_cancellation_flag;
+  }
+
+  const bool& GetLoadCancellationFlagRef() const noexcept {
+    return *load_cancellation_flag;
+  }
+
   // User specified logging func and param
   OrtLoggingFunction user_logging_function = nullptr;
   void* user_logging_param = nullptr;

Diff for: onnxruntime/core/framework/session_state.cc

@@ -585,6 +585,8 @@ Status SessionState::PrepackConstantInitializedTensors(
                     constant_initialized_tensors.erase(ort_value_idx);
                   }
                 }
+                ORT_RETURN_IF(sess_options_.IsLoadCancellationFlagSet(),
+                              "Weight pre-packing was canceled due to user request.");
               }
               // stop searching in 2 cases:
               // 1. value is not from OuterScope
@@ -1528,6 +1530,9 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string<PATH_CHAR_
                                               p_seq_exec_plan_);
   ORT_RETURN_IF_ERROR(status);
 
+  ORT_RETURN_IF(session_options.IsLoadCancellationFlagSet(),
+                "SessionState finalize is canceled due to user request");
+
   // Record the allocation plan
 
   // Uncomment the below to dump the allocation plan to std::cout

Diff for: onnxruntime/core/framework/session_state_utils.cc

@@ -378,6 +378,10 @@ common::Status SaveInitializedTensors(
 
   // 3. create weight tensors based on weights buffer
   for (const auto& entry : id_to_initialized_tensor) {
+    // We check for cancelleation for every initializer since mapping from disk can be costly
+    ORT_RETURN_IF(session_options.IsLoadCancellationFlagSet(),
+                  "Saving session state weights is canceled due to user request.");
+
     int ort_value_index = entry.first;
     const std::string& name = entry.second->name();
 

Diff for: onnxruntime/core/graph/graph.cc

@@ -1268,6 +1268,10 @@ Graph::Graph(const Model& owning_model,
 #endif
   }
 
+  if (owning_model_.IsCancellationFalgSet()) {
+    ORT_THROW("Graph loading canceled due to user request.");
+  }
+
   // Remove constant nodes as they're replaced with initializers above.
   const gsl::not_null<RepeatedPtrField<NodeProto>*> graph_mutable_nodes{graph_proto_->mutable_node()};
   graph_mutable_nodes->erase(
@@ -1300,6 +1304,9 @@ Graph::Graph(const Model& owning_model,
       delete graph_proto_->mutable_sparse_initializer()->ReleaseCleared();
     }
 #endif
+    if (owning_model_.IsCancellationFalgSet()) {
+      ORT_THROW("Graph loading canceled due to user request.");
+    }
   }
 #endif
 
@@ -1365,6 +1372,10 @@ Graph::Graph(const Model& owning_model,
     }
   }
 
+  if (owning_model_.IsCancellationFalgSet()) {
+    ORT_THROW("Graph loading canceled due to user request.");
+  }
+
   for (auto& graph_output : graph_proto_->output()) {
     if (utils::HasName(graph_output) && utils::HasType(graph_output)) {
       auto& name = graph_output.name();

Diff for: onnxruntime/core/graph/model.cc

@@ -82,7 +82,7 @@ Model::Model(const std::string& graph_name,
              const std::vector<ONNX_NAMESPACE::FunctionProto>& model_local_functions,
              const logging::Logger& logger,
              const ModelOptions& options)
-    : model_path_(model_path) {
+    : model_path_(model_path), load_cancellation_flag_(options.load_cancellation_flag) {
   model_proto_.set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION);
   model_proto_.mutable_graph()->set_name(graph_name);
   model_metadata_ = model_metadata;
@@ -161,7 +161,7 @@ Model::Model(const ModelProto& model_proto, const PathString& model_path,
 Model::Model(ModelProto&& model_proto, const PathString& model_path,
              const IOnnxRuntimeOpSchemaRegistryList* local_registries,
              const logging::Logger& logger, const ModelOptions& options)
-    : model_path_(model_path) {
+    : model_path_(model_path), load_cancellation_flag_(options.load_cancellation_flag) {
   if (!utils::HasGraph(model_proto)) {
     ORT_THROW("ModelProto does not have a graph.");
   }

Diff for: onnxruntime/core/graph/model.h

@@ -38,9 +38,13 @@ struct ModelOptions {
   // be returned.
   bool strict_shape_type_inference;
 
-  ModelOptions(bool allow_released_opsets_only, bool strict_shape_type_inference)
+  const bool* load_cancellation_flag = nullptr;
+
+  ModelOptions(bool allow_released_opsets_only, bool strict_shape_type_inference,
+               const bool* load_cancellation_flag = nullptr)
       : allow_released_opsets_only(allow_released_opsets_only),
-        strict_shape_type_inference(strict_shape_type_inference) {}
+        strict_shape_type_inference(strict_shape_type_inference),
+        load_cancellation_flag(load_cancellation_flag) {}
 
   ModelOptions() : ModelOptions(true, false) {}
 };
@@ -143,6 +147,11 @@ class Model {
 
   const NodeHashMap<std::string, std::unique_ptr<FunctionTemplate>>& GetModelLocalFunctionTemplates() const;
 
+  // Check for load cancellation.
+  bool IsCancellationFalgSet() const noexcept {
+    return load_cancellation_flag_ && *load_cancellation_flag_;
+  }
+
 #else
   // Get model's IR version.
   // Return <kNoVersion> if not specified.
@@ -343,5 +352,7 @@ class Model {
 
   // Main graph of the model.
   std::unique_ptr<Graph> graph_;
+
+  const bool* load_cancellation_flag_ = nullptr;
 };
 }  // namespace onnxruntime

Diff for: onnxruntime/core/optimizer/graph_optimizer_registry.h

@@ -55,6 +55,16 @@ class GraphOptimizerRegistry {
    */
   const onnxruntime::SessionOptions& GetSessionOptions() const { return *session_options_; }
 
+  /**
+   * Get Load Cancellation Flag.
+   * This flag is used to cancel the loading of the model if it takes too long.
+   * It is set to true when the user cancels the loading process.
+   * The flag is shared across all threads, so it is read-only
+   */
+  const bool& GetLoadCancellationFlagRef() const {
+    return session_options_->GetLoadCancellationFlagRef();
+  }
+
   /**
    * Get Logger.
    */

Diff for: onnxruntime/core/optimizer/graph_transformer_mgr.cc

@@ -35,6 +35,8 @@ common::Status GraphTransformerManager::ApplyTransformers(Graph& graph, Transfor
       bool modified = false;
       ORT_RETURN_IF_ERROR(transformer->Apply(graph, modified, logger));
       graph_changed = graph_changed || modified;
+      ORT_RETURN_IF(IsLoadCancellationFlagSet(),
+                    "Graph transformation canceled due to user request.");
     }
     if (!graph_changed) {
       break;

Diff for: onnxruntime/core/optimizer/graph_transformer_mgr.h

@@ -24,6 +24,16 @@ class GraphTransformerManager {
   // Get the maximum number of graph transformation steps
   common::Status GetSteps(unsigned& steps) const;
 
+  // Set the cancellation flag ptr from session_options
+  void SetLoadCancellationFlagRef(const bool& load_cancellation_flag) {
+    load_cancellation_flag_ = &load_cancellation_flag;
+  }
+
+  // Get the cancellation flag ptr
+  bool IsLoadCancellationFlagSet() const noexcept {
+    return load_cancellation_flag_ && *load_cancellation_flag_;
+  }
+
   // Register a transformer with a level.
   common::Status Register(std::unique_ptr<GraphTransformer> transformer, TransformerLevel level);
 
@@ -38,5 +48,6 @@ class GraphTransformerManager {
 
   InlinedHashMap<TransformerLevel, InlinedVector<std::unique_ptr<GraphTransformer>>> level_to_transformer_map_;
   InlinedHashMap<std::string, GraphTransformer*> transformers_info_;
+  const bool* load_cancellation_flag_ = nullptr;
 };
 }  // namespace onnxruntime

Diff for: onnxruntime/core/session/abi_session_options.cc

@@ -340,3 +340,9 @@ ORT_API_STATUS_IMPL(OrtApis::SetDeterministicCompute, _Inout_ OrtSessionOptions*
   return nullptr;
   API_IMPL_END
 }
+
+ORT_API_STATUS_IMPL(OrtApis::SessionOptionsSetLoadCancellationFlag, _Inout_ OrtSessionOptions* options,
+                    _In_ bool is_cancel) {
+  *options->value.load_cancellation_flag = is_cancel;
+  return nullptr;
+}