Fix the reading of epctx blob using stream

preetha-intel · preetha-intel · commit 91b627b962ca · 2025-03-14T06:22:39.000-07:00
diff --git a/onnxruntime/core/providers/openvino/backend_manager.cc b/onnxruntime/core/providers/openvino/backend_manager.cc
@@ -76,7 +76,13 @@ BackendManager::BackendManager(SessionContext& session_context,
   ptr_stream_t model_stream;
   std::unique_ptr<onnx::ModelProto> model_proto;
   if (subgraph_context_.is_ep_ctx_graph) {
-    model_stream = ep_ctx_handle_.GetModelBlobStream(session_context_.so_context_file_path, subgraph);
+    std::cout << " inside is_ep_ctx_graph " << std::endl;
+    std::string model_name = onnxruntime::openvino_ep::BackendManager::stripAfterFirstDot(session_context_.onnx_model_path_name.filename().string());
+    auto subgraph_name = model_name + "_" + subgraph_context_.subgraph_name;
+    model_stream = ep_ctx_handle_.GetModelBlobStream(shared_context_,
+                                                     session_context_.so_context_file_path,
+                                                     subgraph_name,
+                                                     subgraph);
   } else {
     model_proto = GetModelProtoFromFusedNode(fused_node, subgraph, logger);
   }
@@ -96,7 +102,9 @@ BackendManager::BackendManager(SessionContext& session_context,
       if (!sw.mapped_weights) {
         sw.mapped_weights = std::make_unique<SharedContext::SharedWeights::WeightsFile>(weight_filename);
       }
+      std::cout << " Call createOVTensors in backend_manager.cc" << std::endl;
       backend_utils::CreateOVTensors(session_context_.device_type, sw.metadata, *sw.mapped_weights);
+      std::cout << " create OVTensors successful " << std::endl;
     }
   }
 
@@ -197,6 +205,14 @@ BackendManager::BackendManager(SessionContext& session_context,
   }
 }
 
+std::string BackendManager::stripAfterFirstDot(std::string filename) {
+  size_t dotPos = filename.find('.');  // Find first dot
+  if (dotPos == std::string::npos) {
+    return filename;  // No dot found, return full filename
+  }
+  return filename.substr(0, dotPos);  // Return everything before first dot
+}
+
 // Call EPContext model exporter here if the provider option for exporting
 // precompiled blob is set. If that's the case:
 // By default, create model in embed mode where the blob stream is exported as data within
@@ -210,27 +226,33 @@ Status BackendManager::ExportCompiledBlobAsEPCtxNode(const onnxruntime::GraphVie
     ORT_THROW(exception_str);
   }
 
+  std::cout << " inside export compiled model " << std::endl;
+
   // If embed_mode, then pass on the serialized blob
   // If not embed_mode, dump the blob here and only pass on the path to the blob
   std::string model_blob_str;
   auto compiled_model = concrete_backend_->GetOVCompiledModel();
-  if (session_context_.so_share_ep_contexts){
-    // std::ostringstream model_blob_stream;
-    // compiled_model.export_model(model_blob_stream);
+  if (session_context_.so_share_ep_contexts) {
+    std::ostringstream model_blob_stream;
+    compiled_model.export_model(model_blob_stream);
+    std::cout << " inside export compiled model - share ep contexts" << std::endl;
 
     // std::ofstream file(metadata_filename, std::ios::app| std::ios::binary);
     // std::cout << " write to metadata bin - " << metadata_filename << std::endl;
     auto& subgraph_metadata = shared_context_.shared_weights.subgraph_metadata;
-
-    sw::SubgraphMetadata::Map::key_type key{subgraph_context_.subgraph_name};
+    std::string model_name = onnxruntime::openvino_ep::BackendManager::stripAfterFirstDot(session_context_.onnx_model_path_name.filename().string());
+    auto subgraph_name = model_name + "_" + subgraph_context_.subgraph_name;
+    sw::SubgraphMetadata::Map::key_type key{subgraph_name};
     sw::SubgraphMetadata::Map::mapped_type value{};
 
     auto& bin_file = shared_context_.shared_weights.shared_bin_file.bin_file_;
-    if (bin_file.is_open()) {
+    std::cout << " subgraph name " << subgraph_name << "key = " << key.name << " For bin write " << std::endl;
+    if (!subgraph_metadata.contains(key) && bin_file.is_open()) {
       // std::cout << "Current offset before "<< subgraph_context_.subgraph_name << "  = " << bin_file.tellp() << std::endl;
       value.epctx_offset = bin_file.tellp();
-      // bin_file << model_blob_stream.str();
-      compiled_model.export_model(bin_file);
+      std::cout << " bin file location for writing subgraph = " << bin_file.tellp() << std::endl;
+      bin_file << model_blob_stream.str();
+      // compiled_model.export_model(bin_file);
       // std::cout << "Current offset after "<< subgraph_context_.subgraph_name << "  = " << bin_file.tellp() << std::endl;
       value.epctx_length = static_cast<size_t>(static_cast<std::streamoff>(bin_file.tellp()) - value.epctx_offset);
       // std::cout << "Key = " << key.name << " Offset = " << value.epctx_offset << " , length = " << value.epctx_length << std::endl;
diff --git a/onnxruntime/core/providers/openvino/backend_manager.h b/onnxruntime/core/providers/openvino/backend_manager.h
@@ -47,6 +47,8 @@ class BackendManager {
   ReWriteInputShapeInfo(const ONNX_NAMESPACE::ModelProto& model_proto,
                         const std::vector<std::vector<int64_t>>& input_shapes);
 
+  std::string stripAfterFirstDot(std::string filename);
+
   std::unique_ptr<ONNX_NAMESPACE::ModelProto> model_proto_;
   std::shared_ptr<IBackend> concrete_backend_;
   std::map<std::string, std::shared_ptr<IBackend>> backend_map_;
diff --git a/onnxruntime/core/providers/openvino/backend_utils.cc b/onnxruntime/core/providers/openvino/backend_utils.cc
@@ -303,22 +303,33 @@ void CreateOVTensors(const std::string& device_name,
                      SharedContext::SharedWeights::Metadata::Map& metadata_map,
                      SharedContext::SharedWeights::WeightsFile& weights) {
   for (auto& [key, value] : metadata_map) {
-    if (value.tensor) continue;
+    // std::cout << " Key = " << key.name << std::endl;
+    if (value.tensor) {
+      // std::cout << " Value already present for key = " << key.name << std::endl;
+      continue;
+    }
 
     // Get element data type
+    // std::cout << " value element type = " << value.element_type << std::endl;
     auto onnx_element_type = (ONNX_NAMESPACE::TensorProto_DataType)value.element_type;
 
     ov::element::Type ov_elementType = GetOpenVINOElementType(onnx_element_type);  // Map to OpenVINO data type
-
+    // std::cout << "value dimensions = " << std::endl;
+    // for (auto dim:value.dimensions){
+    //   std::cout << dim << std::endl;
+    // }
     // Create OpenVINO Tensor
     if (device_name == "NPU") {
       // Use remote tensors
       auto npu_context = OVCore::Get()->core.get_default_context("NPU").as<ov::intel_npu::level_zero::ZeroContext>();
       auto&& remote_tensor = npu_context.create_l0_host_tensor(ov_elementType, value.dimensions, ov::intel_npu::TensorType::INPUT);
-
+      // std::cout << " Remote tensor created " << std::endl;
       // Copy data to remote tensor
+      // std::cout << " value size = " << value.size << std::endl;
       weights.load_weights(value.data_offset, remote_tensor.get(), value.size);
       value.tensor = std::make_shared<ov::Tensor>(remote_tensor);
+      // std::cout << " value tensor created " << std::endl;
+
     } else {
       // Use vanilla tensors
       value.tensor = std::make_shared<ov::Tensor>(ov_elementType, value.dimensions);
diff --git a/onnxruntime/core/providers/openvino/backends/basic_backend.cc b/onnxruntime/core/providers/openvino/backends/basic_backend.cc
@@ -114,6 +114,7 @@ BasicBackend::BasicBackend(std::unique_ptr<ONNX_NAMESPACE::ModelProto>& model_pr
           ov_model, hw_target, device_config, subgraph_context_.subgraph_name);
     }
 #endif
+    std::cout << " loaded model to the plugin " << std::endl;
     LOGS_DEFAULT(INFO) << log_tag << "Loaded model to the plugin";
   } catch (const char* msg) {
     ORT_THROW(msg);
@@ -125,11 +126,14 @@ BasicBackend::BasicBackend(std::unique_ptr<ONNX_NAMESPACE::ModelProto>& model_pr
   if (session_context_.so_share_ep_contexts) {
     initializer = [&metadata](OVInferRequestPtr ir_ptr) {
       const auto input_count = ir_ptr->GetNumInputs();
+      std::cout << " ov ir input count = " << input_count << std::endl;
       for (auto i = 0u; i < input_count; i++) {
         using Key = SharedContext::SharedWeights::Metadata::Key;
         const auto tensor_key = Key{ir_ptr->GetInputTensorName(i)};
         if (metadata.contains(tensor_key)) {
           auto& value = metadata.at(tensor_key);
+          // ORT_ENFORCE(value.tensor->get_byte_size() == value.size, "Unexpected tensor size mismatch");
+          std::cout << " value tensor is set with shape = " << value.tensor->get_byte_size() << " input size from metadata = " << value.size << std::endl;
           ir_ptr->SetTensor(tensor_key.name, value.tensor);
         }
       }
diff --git a/onnxruntime/core/providers/openvino/contexts.h b/onnxruntime/core/providers/openvino/contexts.h
@@ -25,15 +25,15 @@ class SharedContext : public WeakSingleton<SharedContext> {
   SharedContext() : OVCore_(OVCore::Get()) {}
   struct SharedWeights {
     struct Header {
-      uint32_t bin_version=1;
-      long footer_offset=0;
-    }header_;
+      uint32_t bin_version = 1;
+      long footer_offset = 0;
+    } header_;
     struct Footer {
       long subgraph_offset;
       size_t subgraph_length;
       long metadata_offset;
       size_t metadata_length;
-    }footer_;
+    } footer_;
 
     struct Metadata {
       struct Key {
@@ -56,8 +56,6 @@ class SharedContext : public WeakSingleton<SharedContext> {
       using Map = std::unordered_map<Key, Value, Hash>;
       void writeMetadataToBinaryFile(SharedContext& shared_context, const Metadata::Map& metadata);
       void readMetadataFromBinaryFile(SharedContext& shared_context, Metadata::Map& metadata);
-      // friend std::ostream& operator<<(std::ostream& right, const Metadata::Map& metadata);
-      // friend std::istream& operator>>(std::istream& right, Metadata::Map& metadata);
     };
 
     struct SubgraphMetadata {
@@ -78,9 +76,7 @@ class SharedContext : public WeakSingleton<SharedContext> {
       void writeSubgraphDataToBinaryFile(SharedContext& shared_context,
                                          const SubgraphMetadata::Map& subgraph_metadata);
       void readSubgraphDataFromBinaryFile(SharedContext& shared_context,
-                                         SubgraphMetadata::Map& subgraph_metadata);
-      // friend std::ostream& operator<<(std::ostream& right, const SubgraphMetadata::Map& subgraph_metadata);
-      // friend std::istream& operator>>(std::istream& right, SubgraphMetadata::Map& subgraph_metadata);
+                                          SubgraphMetadata::Map& subgraph_metadata);
     };
 
     struct WeightsFile {
@@ -103,56 +99,49 @@ class SharedContext : public WeakSingleton<SharedContext> {
       SharedBinFile() = default;  // Default constructor
       ~SharedBinFile() {
         if (bin_file_.is_open()) {
-            bin_file_.close();  // Close file when object is destroyed
+          bin_file_.close();  // Close file when object is destroyed
         }
       }
 
       void openBinFile(const fs::path shared_bin_filename) {
-      // Check if the file exists before trying to open
+        // Check if the file exists before trying to open
         if (!fs::exists(shared_bin_filename)) {
-            std::cerr << "Error: The file does not exist at path: " << shared_bin_filename << std::endl;
-            std::ofstream createFile(shared_bin_filename, std::ios::binary);  // Create an empty binary file
-            if (!createFile) {
-                throw std::runtime_error("Failed to create the file!");
-            }
-            createFile.close();
-            // throw std::runtime_error("Failed to open log file! File does not exist.");
+          std::ofstream createFile(shared_bin_filename, std::ios::binary);  // Create an empty binary file
+          if (!createFile) {
+            throw std::runtime_error("Failed to create the file!");
+          }
+          createFile.close();
         }
 
         // Check if the file is accessible for reading and writing
         fs::perms file_perms = fs::status(shared_bin_filename).permissions();
 
         if ((file_perms & fs::perms::owner_read) == fs::perms::none ||
             (file_perms & fs::perms::owner_write) == fs::perms::none) {
-            std::cerr << "Error: Insufficient permissions for file: " << shared_bin_filename << std::endl;
-            throw std::runtime_error("Failed to open log file! Insufficient permissions.");
+          std::cerr << "Error: Insufficient permissions for file: " << shared_bin_filename << std::endl;
+          throw std::runtime_error("Failed to open log file! Insufficient permissions.");
         }
 
-
         if (!bin_file_.is_open()) {  // Prevent reopening
-          std::cout << " Bin file is not open " << std::endl;
           bin_file_.open(shared_bin_filename, std::ios::in | std::ios::out | std::ios::binary);
-          std::cout << " bin file opened " << std::endl;
           bin_size_ = bin_file_.seekg(0, std::ios::end).tellg();
-
-          std::cout << " bin size = " << bin_size_ << std::endl;
           bin_file_.seekg(0, std::ios::beg);  // Reset to the beginning of the file
 
-
           if (!bin_file_) {
-              throw std::runtime_error("Failed to open log file!");
+            throw std::runtime_error("Failed to open log file!");
           }
         }
       }
-    }shared_bin_file;
+      void readBinFile(SharedContext& shared_context_);
+    } shared_bin_file;
 
     fs::path external_weight_filename;
     std::unique_ptr<WeightsFile> mapped_weights;
     Metadata metadata_;
     Metadata::Map metadata;
     SubgraphMetadata subgraph_metadata_;
     SubgraphMetadata::Map subgraph_metadata;
-  }shared_weights;
+  } shared_weights;
 };
 
 using config_t = std::map<std::string, ov::AnyMap>;
diff --git a/onnxruntime/core/providers/openvino/onnx_ctx_model_helper.cc b/onnxruntime/core/providers/openvino/onnx_ctx_model_helper.cc
@@ -99,7 +99,10 @@ Status EPCtxHandler::AddOVEPCtxNodeToGraph(const GraphViewer& graph_viewer,
   return Status::OK();
 }
 
-std::unique_ptr<std::istream> EPCtxHandler::GetModelBlobStream(const std::filesystem::path& so_context_file_path, const GraphViewer& graph_viewer) const {
+std::unique_ptr<std::istream> EPCtxHandler::GetModelBlobStream(SharedContext& shared_context_,
+                                                               const std::filesystem::path& so_context_file_path,
+                                                               const std::string& subgraph_name,
+                                                               const GraphViewer& graph_viewer) const {
   auto first_index = *graph_viewer.GetNodesInTopologicalOrder().begin();
   auto node = graph_viewer.GetNode(first_index);
   ORT_ENFORCE(node != nullptr);
@@ -121,7 +124,40 @@ std::unique_ptr<std::istream> EPCtxHandler::GetModelBlobStream(const std::filesy
     }
     blob_filepath = blob_filepath.parent_path() / ep_cache_context;
     ORT_ENFORCE(std::filesystem::exists(blob_filepath), "Blob file not found: ", blob_filepath.string());
-    result.reset((std::istream*)new std::ifstream(blob_filepath, std::ios_base::binary | std::ios_base::in));
+    std::cout << " blob_filepath " << blob_filepath.filename().string() << std::endl;
+    std::cout << " shared bin filename = " << shared_context_.shared_weights.shared_bin_file.shared_bin_filename.filename().string() << std::endl;
+    if (blob_filepath == shared_context_.shared_weights.shared_bin_file.shared_bin_filename) {
+      LOGS_DEFAULT(VERBOSE) << "[OpenVINO EP] Read blob from Shared bin file - " << blob_filepath;
+      auto& sb = shared_context_.shared_weights.shared_bin_file;
+      // check if size of bin file is greater than the header as it gets written at the begining
+      ORT_ENFORCE(sb.bin_size_ > 8, " Bin file is empty. Regenerate the epctx model. Bin file path : ", blob_filepath.string());
+      auto subgraph_metadata = shared_context_.shared_weights.subgraph_metadata;
+      using Key = SharedContext::SharedWeights::SubgraphMetadata::Key;
+      std::cout << " subgraph name = " << subgraph_name << std::endl;
+      const auto subgraph_key = Key{subgraph_name};
+      auto it = subgraph_metadata.find(subgraph_key);
+      if (it != subgraph_metadata.end()) {
+        auto& value = it->second;
+        std::cout << " value.epctx_offset =  " << value.epctx_offset << std::endl;
+        std::cout << " value.epctx_length =  " << value.epctx_length << std::endl;
+        std::cout << " sb.bin_size_ =  " << sb.bin_size_ << std::endl;
+
+        if (value.epctx_offset < sb.bin_size_ && value.epctx_length <= sb.bin_size_ &&
+            (value.epctx_offset <= sb.bin_size_ - value.epctx_length)) {
+          sb.bin_file_.seekg(value.epctx_offset);             // Move to the specified offset
+          std::string buffer(value.epctx_length, '\0');       // preallocate space
+          sb.bin_file_.read(&buffer[0], value.epctx_length);  // Read the specified length
+          // Adjust string size in case of a short read
+          buffer.resize(sb.bin_file_.gcount());
+          std::cout << " Read epctx into stream " << std::endl;
+          result.reset((std::istream*)new std::istringstream(buffer));
+        }
+      }
+      ORT_ENFORCE(result != nullptr, " Epctx blob is not read. Check bin file correctness from Bin path: ",
+                  blob_filepath.string());
+    } else {
+      result.reset((std::istream*)new std::ifstream(blob_filepath, std::ios_base::binary | std::ios_base::in));
+    }
   }
   LOGS_DEFAULT(VERBOSE) << "[OpenVINO EP] Read blob from EPContext Node";
   return result;
diff --git a/onnxruntime/core/providers/openvino/onnx_ctx_model_helper.h b/onnxruntime/core/providers/openvino/onnx_ctx_model_helper.h
@@ -6,8 +6,10 @@
 #include <sstream>
 #include <string>
 #include <memory>
+#include <streambuf>
 
 #include "core/providers/shared_library/provider_api.h"
+#include "core/providers/openvino/contexts.h"
 
 namespace onnxruntime {
 namespace openvino_ep {
@@ -31,7 +33,10 @@ class EPCtxHandler {
                                const std::string& graph_name,
                                const bool embed_mode,
                                std::string&& model_blob_str) const;
-  std::unique_ptr<std::istream> GetModelBlobStream(const std::filesystem::path& so_context_file_path, const GraphViewer& graph_viewer) const;
+  std::unique_ptr<std::istream> GetModelBlobStream(SharedContext& shared_context_,
+                                                   const std::filesystem::path& so_context_file_path,
+                                                   const std::string& subgraph_name,
+                                                   const GraphViewer& graph_viewer) const;
   InlinedVector<const Node*> GetEPCtxNodes() const;
 
  private:
@@ -40,5 +45,26 @@ class EPCtxHandler {
   const logging::Logger& logger_;
 };
 
+// class LimitedFileStreambuf : public std::streambuf {
+// private:
+//     std::fstream& file; // Reference to the existing file stream
+//     long start, end; // Start and end positions
+
+// protected:
+//     int_type underflow() override {
+//         if (file.tellg() >= end || file.eof())
+//             return traits_type::eof(); // Stop reading if we reach the limit
+
+//         return file.get(); // Read next character directly from the file
+//     }
+
+// public:
+//     LimitedFileStreambuf(std::fstream& bin_file_, long start, long end)
+//         : file(bin_file_), start(start), end(end) {
+//         file.clear(); // Clear error flags in case of previous reads
+//         file.seekg(start); // Move file pointer to the start position
+//     }
+// };
+
 }  // namespace openvino_ep
 }  // namespace onnxruntime
diff --git a/onnxruntime/core/providers/openvino/openvino_execution_provider.cc b/onnxruntime/core/providers/openvino/openvino_execution_provider.cc
diff --git a/onnxruntime/core/providers/openvino/openvino_file_handling.cc b/onnxruntime/core/providers/openvino/openvino_file_handling.cc
diff --git a/onnxruntime/core/providers/openvino/ov_interface.cc b/onnxruntime/core/providers/openvino/ov_interface.cc
