Fix the reading of epctx blob using stream

preetha-intel · preetha-intel · commit 98ef5838b1b2 · 2025-03-04T06:48:07.000-08:00
diff --git a/onnxruntime/core/providers/openvino/backend_manager.cc b/onnxruntime/core/providers/openvino/backend_manager.cc
@@ -76,7 +76,12 @@ 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(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_,
+                                                    subgraph_name,
+                                                    subgraph);
   } else {
     model_proto = GetModelProtoFromFusedNode(fused_node, subgraph, logger);
   }
@@ -96,7 +101,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 +204,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 +225,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);
+    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
@@ -302,23 +302,35 @@ ov::element::Type GetOpenVINOElementType(ONNX_NAMESPACE::TensorProto_DataType dt
 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().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
@@ -88,10 +88,12 @@ BasicBackend::BasicBackend(std::unique_ptr<ONNX_NAMESPACE::ModelProto>& model_pr
     if (subgraph_context_.is_ep_ctx_graph) {
       // If the blob is held in an EPContext node, then skip FE+Compile
       // and directly move on to creating a backend with the executable blob
+      std::cout << " before ovcore import model " << std::endl;
       exe_network_ = OVCore::ImportModel(*model_stream,
                                          hw_target,
                                          device_config,
                                          subgraph_context_.subgraph_name);
+      std::cout << " import model is successful " << std::endl;
       model_stream.reset();  // Delete stream after it is no longer needed
     } else if (!session_context_.has_external_weights &&
                !subgraph_context_.has_dynamic_input_shape &&
@@ -120,6 +122,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);
@@ -131,11 +134,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
@@ -51,8 +51,6 @@ struct 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 {
@@ -74,8 +72,6 @@ struct SharedContext {
                                          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);
     };
 
     struct WeightsFile {
@@ -105,13 +101,11 @@ struct SharedContext {
       void openBinFile(const fs::path shared_bin_filename) {
       // 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.");
         }
 
         // Check if the file is accessible for reading and writing
@@ -125,20 +119,16 @@ struct SharedContext {
 
 
         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!");
           }
         }
       }
+      void readBinFile(SharedContext& shared_context_);
     }shared_bin_file;
 
     fs::path external_weight_filename;
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,9 @@ Status EPCtxHandler::AddOVEPCtxNodeToGraph(const GraphViewer& graph_viewer,
   return Status::OK();
 }
 
-std::unique_ptr<std::istream> EPCtxHandler::GetModelBlobStream(const GraphViewer& graph_viewer) const {
+std::unique_ptr<std::istream> EPCtxHandler::GetModelBlobStream(SharedContext& shared_context_,
+                                                              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);
@@ -117,7 +119,52 @@ std::unique_ptr<std::istream> EPCtxHandler::GetModelBlobStream(const GraphViewer
   } else {
     const auto& blob_filepath = graph_viewer.ModelPath().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());
+              // long end = value.epctx_offset + value.epctx_length ;
+              // LimitedFileStreambuf limitedBuffer(sb.bin_file_, value.epctx_offset, end);
+              // std::cout << " string stream read from " << value.epctx_offset << " to " << end << std::endl;
+
+
+              // sb.bin_file_.seekg(value.epctx_offset, std::ios::beg);
+              // // Read exactly 'length' bytes into a string
+              // std::string data(value.epctx_length, '\0');  // Allocate a string with 'length' bytes
+              // sb.bin_file_.read(&data[0], value.epctx_length);     // Read directly into the string
+              // result = std::make_unique<std::stringstream>(std::move(data));
+              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,9 @@ class EPCtxHandler {
                                const std::string& graph_name,
                                const bool embed_mode,
                                std::string&& model_blob_str) const;
-  std::unique_ptr<std::istream> GetModelBlobStream(const GraphViewer& graph_viewer) const;
+  std::unique_ptr<std::istream> GetModelBlobStream(SharedContext& shared_context_,
+                                                   const std::string &subgraph_name,
+                                                   const GraphViewer& graph_viewer) const;
   InlinedVector<const Node*> GetEPCtxNodes() const;
 
  private:
@@ -40,5 +44,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
