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

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

@@ -552,28 +552,30 @@ void regclass_Core(py::module m) {
                                      "`model_stream` must be an io.BytesIO object but " +
                                      (std::string)(py::repr(model_stream)) + "` provided");
             }
-            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 _stream(filename, std::ios::out | std::ios::binary);
+            py::buffer_info info = py::buffer(model_stream.attr("getbuffer")()).request();
+            constexpr auto one_gigabyte = static_cast<double>(1024) * 1024 * 1024;
+            const double size_in_gb = (info.size * info.itemsize) / one_gigabyte;
             model_stream.attr("seek")(0);  // Always rewind stream!
-            if (_stream.is_open()) {
+            ov::CompiledModel result;
+            // std::stringstream cannot handle streams > 2GB, in that case use std::fstream
+            if (size_in_gb > 2) {
+                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 _stream(filename, std::ios::out | std::ios::binary);
+                OPENVINO_ASSERT(_stream.is_open(), "Failed to open temporary file for model stream");
                 const py::bytes data = model_stream.attr("read")();
+
                 // convert the Python bytes object to C++ string
                 char* buffer;
                 Py_ssize_t length;
                 PYBIND11_BYTES_AS_STRING_AND_SIZE(data.ptr(), &buffer, &length);
                 _stream.write(buffer, length);
                 _stream.close();
-            } else {
-                OPENVINO_THROW("Failed to open temporary file for model stream");
-            }
 
-            ov::CompiledModel result;
-            std::fstream _fstream(filename, std::ios::in | std::ios::binary);
-            if (_fstream.is_open()) {
-                py::gil_scoped_release release;
+                std::fstream _fstream(filename, std::ios::in | std::ios::binary);
+                OPENVINO_ASSERT(_fstream.is_open(), "Failed to open temporary file for model stream");
                 result = self.import_model(_fstream, device_name, _properties);
                 _fstream.close();
                 if (std::remove(filename.c_str()) != 0) {
@@ -583,9 +585,11 @@ void regclass_Core(py::module m) {
                     PyErr_WarnEx(PyExc_RuntimeWarning, warning_message.c_str(), 1);
                 }
             } else {
-                OPENVINO_THROW("Failed to open temporary file for model stream");
+                std::stringstream _stream;
+                _stream << model_stream.attr("read")().cast<std::string>();
+                py::gil_scoped_release release;
+                result = self.import_model(_stream, device_name, _properties);
             }
-
             return result;
         },
         py::arg("model_stream"),

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

@@ -11,8 +11,9 @@
     generate_image,
     generate_model_and_image,
     generate_relu_compiled_model,
+    generate_model_with_memory,
     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.properties as props
@@ -76,6 +77,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_model_with_memory([6, 10000, 9000], 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([6, 10000, 9000])
+    res = new_compiled.infer_new_request({"input_data": img})
+    assert np.argmax(res[new_compiled.outputs[0]]) == 14970
+
+
 @pytest.mark.parametrize("input_arguments", [[0], ["data"], []])
 def test_get_input(device, input_arguments):
     compiled_model = generate_relu_compiled_model(device)