[PyOV] Branch import_model and export_model implementations based on the model size

src/bindings/python/src/pyopenvino/core/compiled_model.cpp

@@ -7,6 +7,9 @@
 #include <pybind11/iostream.h>
 #include <pybind11/stl.h>
 
+#include <fstream>
+#include <random>
+
 #include "common.hpp"
 #include "pyopenvino/core/compiled_model.hpp"
 #include "pyopenvino/core/infer_request.hpp"
@@ -81,13 +84,42 @@ void regclass_CompiledModel(py::module m) {
                                      "`model_stream` must be an io.BytesIO object but " +
                                      (std::string)(py::repr(model_stream)) + "` provided");
             }
-            std::stringstream _stream;
-            {
-                py::gil_scoped_release release;
-                self.export_model(_stream);
+            model_stream.attr("seek")(0);  // Always rewind stream!
+            // std::stringstream cannot handle streams > 2GB, in that case we use std::fstream
+            // TODO: Determine model size after serialization
+            if (true) {
+                std::random_device rd;
+                std::mt19937 gen(rd());
+                std::uniform_int_distribution<> distr(1000, 9999);
+                std::string filename = "model_stream_" + std::to_string(distr(gen)) + ".txt";
+
+                std::fstream _fstream(filename, std::ios::out | std::ios::binary);
+                OPENVINO_ASSERT(_fstream.is_open(), "Failed to open temporary file for model stream");
+
+                // TODO: Assert all plugins' export_model functions work correctly for std::fstream and models > 2GB
+                self.export_model(_fstream);
+                _fstream.seekg(0, std::ios::beg);
+                {
+                    std::string str((std::istreambuf_iterator<char>(_fstream)), std::istreambuf_iterator<char>());
+                    _fstream.close();
+                    model_stream.attr("flush")();
+                    model_stream.attr("write")(py::bytes(str));
+                }
+                if (std::remove(filename.c_str()) != 0) {
+                    const std::string abs_path =
+                        py::module_::import("os").attr("getcwd")().cast<std::string>() + "/" + filename;
+                    const std::string warning_message = "Temporary file " + abs_path + " failed to delete!";
+                    PyErr_WarnEx(PyExc_RuntimeWarning, warning_message.c_str(), 1);
+                }
+            } else {
+                std::stringstream _stream;
+                {
+                    py::gil_scoped_release release;
+                    self.export_model(_stream);
+                }
+                model_stream.attr("flush")();
+                model_stream.attr("write")(py::bytes(_stream.str()));
             }
-            model_stream.attr("flush")();
-            model_stream.attr("write")(py::bytes(_stream.str()));
             model_stream.attr("seek")(0);  // Always rewind stream!
         },
         py::arg("model_stream"),

src/bindings/python/tests/test_runtime/test_compiled_model.py

@@ -12,12 +12,13 @@
     generate_model_and_image,
     generate_concat_compiled_model,
     generate_relu_compiled_model,
+    generate_big_model_with_tile,
     generate_relu_compiled_model_with_config,
     encrypt_base64,
     decrypt_base64,
     create_filenames_for_ir,
     create_filename_for_test)
-from openvino import Model, Shape, Core, Tensor, serialize
+from openvino import Model, Shape, Core, Tensor, serialize, Type
 from openvino.runtime import ConstOutput
 
 import openvino.runtime.opset13 as ops
@@ -99,6 +100,24 @@ def test_export_import_advanced(device):
     assert np.argmax(res[new_compiled.outputs[0]]) == 531
 
 
+def test_export_import_large_model(device):
+    import io
+    core = Core()
+    if props.device.Capability.EXPORT_IMPORT not in core.get_property(device, props.device.capabilities):
+        pytest.skip(f"{core.get_property(device, props.device.full_name)} plugin due-to export, import model API isn't implemented.")
+
+    # model of size of roughly 2.01GB
+    model = generate_big_model_with_tile([1, 10, 9], [6, 10000, 9000], [6, 1000, 1000], Type.f32)
+    core = Core()
+    compiled_model = core.compile_model(model, device, {})
+    user_stream = io.BytesIO()
+    compiled_model.export_model(user_stream)
+    new_compiled = core.import_model(user_stream, device)
+    img = generate_image([1, 10, 9])
+    res = new_compiled.infer_new_request({"input_data": img})
+    assert np.argmax(res[new_compiled.outputs[0]]) == 63006
+
+
 # request - https://docs.pytest.org/en/7.1.x/reference/reference.html#request
 @pytest.fixture
 def prepare_blob_path(request, tmp_path):

src/bindings/python/tests/utils/helpers.py

@@ -256,6 +256,14 @@ def generate_model_with_memory(input_shape, data_type) -> openvino._pyopenvino.M
     return model
 
 
+def generate_big_model_with_tile(input_shape, constant_shape, tile_shape, data_type) -> openvino._pyopenvino.Model:
+    input_data = ops.parameter(input_shape, name="input_data", dtype=data_type)
+    init_val = ops.constant(np.ones(constant_shape, np.float32), data_type)
+    tile = ops.tile(input_data, tile_shape)
+    add = ops.add(init_val, tile, name="MemoryAdd")
+    return Model(add, [input_data], "TestModel")
+
+
 def generate_concat_compiled_model(device, input_shape: List[int] = None, ov_type=Type.f32, numpy_dtype=np.float32):
     if input_shape is None:
         input_shape = [5]