Implement concurrent memcpy for building Python objects into vineyard (…

…#1646) Remove the problematic `.buffer` property (as it cannot bind the lifetime of the underlying blob to the memoryview object) and add concurrent support for memcpy for faster object building. Fixes #1631 Signed-off-by: Tao He <[email protected]>
v6d-io · Dec 13, 2023 · 89530bf · 89530bf
1 parent a219eae
commit 89530bf
Show file tree

Hide file tree

Showing 12 changed files with 282 additions and 123 deletions.
diff --git a/python/core.cc b/python/core.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include "client/ds/object_meta.h"
 #include "client/ds/remote_blob.h"
 #include "client/rpc_client.h"
+#include "common/memory/memcpy.h"
 #include "common/util/json.h"
 #include "common/util/status.h"
 #include "common/util/uuid.h"
@@ -481,18 +482,6 @@ void bind_blobs(py::module& mod) {
           "address",
           [](Blob* self) { return reinterpret_cast<uintptr_t>(self->data()); },
           doc::Blob_address)
-      .def_property_readonly(
-          "buffer",
-          [](Blob& blob) -> py::object {
-            auto buffer = blob.Buffer();
-            if (buffer == nullptr) {
-              return py::none();
-            } else {
-              return py::memoryview::from_memory(
-                  const_cast<uint8_t*>(buffer->data()), buffer->size(), true);
-            }
-          },
-          doc::Blob_buffer)
       .def_buffer([](Blob& blob) -> py::buffer_info {
         return py::buffer_info(const_cast<char*>(blob.data()), sizeof(int8_t),
                                py::format_descriptor<int8_t>::format(), 1,
@@ -550,21 +539,29 @@ void bind_blobs(py::module& mod) {
       .def(
           "copy",
           [](BlobWriter* self, size_t const offset, uintptr_t ptr,
-             size_t const size) {
-            std::memcpy(self->data() + offset, reinterpret_cast<void*>(ptr),
-                        size);
-          },
-          "offset"_a, "address"_a, "size"_a, doc::BlobBuilder_copy)
+             size_t const size,
+             size_t const concurrency = memory::default_memcpy_concurrency) {
+            memory::concurrent_memcpy(self->data() + offset,
+                                      reinterpret_cast<void*>(ptr), size,
+                                      concurrency);
+          },
+          "offset"_a, "address"_a, "size"_a,
+          py::arg("concurrency") = memory::default_memcpy_concurrency,
+          doc::BlobBuilder_copy)
       .def(
           "copy",
-          [](BlobWriter* self, size_t offset, py::buffer const& buffer) {
+          [](BlobWriter* self, size_t offset, py::buffer const& buffer,
+             size_t const concurrency = memory::default_memcpy_concurrency) {
             throw_on_error(copy_memoryview(buffer.ptr(), self->data(),
-                                           self->size(), offset));
+                                           self->size(), offset, concurrency));
           },
-          "offset"_a, "buffer"_a)
+          "offset"_a, "buffer"_a,
+          py::arg("concurrency") = memory::default_memcpy_concurrency,
+          doc::BlobBuilder_copy)
       .def(
           "copy",
-          [](BlobWriter* self, size_t offset, py::bytes const& bs) {
+          [](BlobWriter* self, size_t offset, py::bytes const& bs,
+             size_t const concurrency = memory::default_memcpy_concurrency) {
             char* buffer = nullptr;
             ssize_t length = 0;
             if (PYBIND11_BYTES_AS_STRING_AND_SIZE(bs.ptr(), &buffer, &length)) {
@@ -577,27 +574,18 @@ void bind_blobs(py::module& mod) {
                   "', but the buffer size is '" + std::to_string(length) +
                   "'"));
             }
-            std::memcpy(self->data() + offset, buffer, length);
+            memory::concurrent_memcpy(self->data() + offset, buffer, length,
+                                      concurrency);
           },
-          "offset"_a, "bytes"_a)
+          "offset"_a, "bytes"_a,
+          py::arg("concurrency") = memory::default_memcpy_concurrency,
+          doc::BlobBuilder_copy)
       .def_property_readonly(
           "address",
           [](BlobWriter* self) {
             return reinterpret_cast<uintptr_t>(self->data());
           },
           doc::BlobBuilder_address)
-      .def_property_readonly(
-          "buffer",
-          [](BlobWriter& blob) -> py::object {
-            auto buffer = blob.Buffer();
-            if (buffer == nullptr) {
-              return py::none();
-            } else {
-              return py::memoryview::from_memory(buffer->mutable_data(),
-                                                 buffer->size(), false);
-            }
-          },
-          doc::BlobBuilder_buffer)
       .def_buffer([](BlobWriter& blob) -> py::buffer_info {
         return py::buffer_info(blob.data(), sizeof(int8_t),
                                py::format_descriptor<int8_t>::format(), 1,
@@ -641,18 +629,6 @@ void bind_blobs(py::module& mod) {
             return reinterpret_cast<uintptr_t>(self->data());
           },
           doc::RemoteBlob_address)
-      .def_property_readonly(
-          "buffer",
-          [](RemoteBlob& blob) -> py::object {
-            auto buffer = blob.Buffer();
-            if (buffer == nullptr) {
-              return py::none();
-            } else {
-              return py::memoryview::from_memory(
-                  const_cast<uint8_t*>(buffer->data()), buffer->size(), true);
-            }
-          },
-          doc::RemoteBlob_buffer)
       .def_buffer([](RemoteBlob& blob) -> py::buffer_info {
         return py::buffer_info(const_cast<char*>(blob.data()), sizeof(int8_t),
                                py::format_descriptor<int8_t>::format(), 1,
@@ -733,21 +709,29 @@ void bind_blobs(py::module& mod) {
       .def(
           "copy",
           [](RemoteBlobWriter* self, size_t const offset, uintptr_t ptr,
-             size_t const size) {
-            std::memcpy(self->data() + offset, reinterpret_cast<void*>(ptr),
-                        size);
-          },
-          "offset"_a, "address"_a, "size"_a, doc::RemoteBlobBuilder_copy)
+             size_t const size,
+             size_t const concurrency = memory::default_memcpy_concurrency) {
+            memory::concurrent_memcpy(self->data() + offset,
+                                      reinterpret_cast<void*>(ptr), size,
+                                      concurrency);
+          },
+          "offset"_a, "address"_a, "size"_a,
+          py::arg("concurrency") = memory::default_memcpy_concurrency,
+          doc::RemoteBlobBuilder_copy)
       .def(
           "copy",
-          [](RemoteBlobWriter* self, size_t offset, py::buffer const& buffer) {
+          [](RemoteBlobWriter* self, size_t offset, py::buffer const& buffer,
+             size_t const concurrency = memory::default_memcpy_concurrency) {
             throw_on_error(copy_memoryview(buffer.ptr(), self->data(),
-                                           self->size(), offset));
+                                           self->size(), offset, concurrency));
           },
-          "offset"_a, "buffer"_a)
+          "offset"_a, "buffer"_a,
+          py::arg("concurrency") = memory::default_memcpy_concurrency,
+          doc::RemoteBlobBuilder_copy)
       .def(
           "copy",
-          [](RemoteBlobWriter* self, size_t offset, py::bytes const& bs) {
+          [](RemoteBlobWriter* self, size_t offset, py::bytes const& bs,
+             size_t const concurrency = memory::default_memcpy_concurrency) {
             char* buffer = nullptr;
             ssize_t length = 0;
             if (PYBIND11_BYTES_AS_STRING_AND_SIZE(bs.ptr(), &buffer, &length)) {
@@ -760,27 +744,18 @@ void bind_blobs(py::module& mod) {
                   "', but the buffer size is '" + std::to_string(length) +
                   "'"));
             }
-            std::memcpy(self->data() + offset, buffer, length);
+            memory::concurrent_memcpy(self->data() + offset, buffer, length,
+                                      concurrency);
           },
-          "offset"_a, "bytes"_a)
+          "offset"_a, "bytes"_a,
+          py::arg("concurrency") = memory::default_memcpy_concurrency,
+          doc::RemoteBlobBuilder_copy)
       .def_property_readonly(
           "address",
           [](RemoteBlobWriter* self) {
             return reinterpret_cast<uintptr_t>(self->data());
           },
           doc::RemoteBlobBuilder_address)
-      .def_property_readonly(
-          "buffer",
-          [](RemoteBlobWriter& blob) -> py::object {
-            auto buffer = blob.Buffer();
-            if (buffer == nullptr) {
-              return py::none();
-            } else {
-              return py::memoryview::from_memory(buffer->mutable_data(),
-                                                 buffer->size(), false);
-            }
-          },
-          doc::RemoteBlobBuilder_buffer)
       .def_buffer([](RemoteBlobWriter& blob) -> py::buffer_info {
         return py::buffer_info(blob.data(), sizeof(int8_t),
                                py::format_descriptor<int8_t>::format(), 1,

diff --git a/python/pybind11_docs.cc b/python/pybind11_docs.cc
@@ -339,11 +339,6 @@ const char* Blob_address = R"doc(
 The memory address value of this blob.
 )doc";
 
-const char* Blob_buffer = R"doc(
-The readonly buffer behind this blob. The result buffer has type
-:code:`memoryview`.
-)doc";
-
 const char* BlobBuilder = R"doc(
 :class:`BlobBuilder` is the builder for creating a finally immutable blob in
 vineyard server.
@@ -373,7 +368,7 @@ Shrink the blob builder to the given size if it is not sealed yet.
 )doc";
 
 const char* BlobBuilder_copy = R"doc(
-.. method:: copy(self, offset: int, ptr: int, size: int)
+.. method:: copy(self, offset: int, ptr: int, size: int, concurrency: int = 6)
     :noindex:
 
 Copy the given address to the given offset.
@@ -383,11 +378,6 @@ const char* BlobBuilder_address = R"doc(
 The memory address value of this blob builder.
 )doc";
 
-const char* BlobBuilder_buffer = R"doc(
-The writeable buffer behind this blob builder. The result buffer has type
-:code:`memoryview`, and it is a mutable one.
-)doc";
-
 const char* RemoteBlob = R"doc(
 :class:`RemoteBlob` is a holder for :class:`Blob` in cases like the :class:`Blob`
 doesn't exist in the local vineyardd instance but the client still want th access
@@ -414,11 +404,6 @@ const char* RemoteBlob_address = R"doc(
 The memory address value of this blob.
 )doc";
 
-const char* RemoteBlob_buffer = R"doc(
-The readonly buffer behind this blob. The result buffer has type
-:code:`memoryview`.
-)doc";
-
 const char* RemoteBlobBuilder = R"doc(
 :class:`RemoteBlobBuilder` is the builder for creating a finally immutable blob in
 vineyard server over a RPC client.
@@ -476,7 +461,7 @@ Abort the blob builder if it is not sealed yet.
 )doc";
 
 const char* RemoteBlobBuilder_copy = R"doc(
-.. method:: copy(self, offset: int, ptr: int, size: int)
+.. method:: copy(self, offset: int, ptr: int, size: int, concurrency: int = 6)
     :noindex:
 
 Copy the given address to the given offset.
@@ -486,11 +471,6 @@ const char* RemoteBlobBuilder_address = R"doc(
 The memory address value of this blob builder.
 )doc";
 
-const char* RemoteBlobBuilder_buffer = R"doc(
-The writeable buffer behind this blob builder. The result buffer has type
-:code:`memoryview`, and it is a mutable one.
-)doc";
-
 const char* InstanceStatus = R"doc(
 :class:`InstanceStatus` represents the status of connected vineyard instance, including
 the instance identity, memory statistics and workloads on this instance.

diff --git a/python/pybind11_docs.h b/python/pybind11_docs.h
@@ -58,7 +58,6 @@ extern const char* Blob_is_empty;
 extern const char* Blob_empty;
 extern const char* Blob__len__;
 extern const char* Blob_address;
-extern const char* Blob_buffer;
 
 extern const char* BlobBuilder;
 extern const char* BlobBuilder_id;
@@ -67,22 +66,19 @@ extern const char* BlobBuilder_abort;
 extern const char* BlobBuilder_shrink;
 extern const char* BlobBuilder_copy;
 extern const char* BlobBuilder_address;
-extern const char* BlobBuilder_buffer;
 
 extern const char* RemoteBlob;
 extern const char* RemoteBlob_id;
 extern const char* RemoteBlob_instance_id;
 extern const char* RemoteBlob_is_empty;
 extern const char* RemoteBlob__len__;
 extern const char* RemoteBlob_address;
-extern const char* RemoteBlob_buffer;
 
 extern const char* RemoteBlobBuilder;
 extern const char* RemoteBlobBuilder_size;
 extern const char* RemoteBlobBuilder_abort;
 extern const char* RemoteBlobBuilder_copy;
 extern const char* RemoteBlobBuilder_address;
-extern const char* RemoteBlobBuilder_buffer;
 
 extern const char* InstanceStatus;
 extern const char* InstanceStatus_instance_id;

diff --git a/python/pybind11_utils.cc b/python/pybind11_utils.cc
@@ -18,7 +18,6 @@ limitations under the License.
 #include <cstdlib>
 #include <cstring>
 #include <string>
-#include <unordered_map>
 #include <utility>
 
 #include "common/memory/memcpy.h"
@@ -133,11 +132,13 @@ void bind_utils(py::module& mod) {
   mod.def(
       "memory_copy",
       [](py::buffer const dst, size_t offset, py::buffer const src,
-         size_t const size) {
-        throw_on_error(
-            copy_memoryview_to_memoryview(src.ptr(), dst.ptr(), size, offset));
+         size_t const size,
+         size_t const concurrency = memory::default_memcpy_concurrency) {
+        throw_on_error(copy_memoryview_to_memoryview(src.ptr(), dst.ptr(), size,
+                                                     offset, concurrency));
       },
-      "dst"_a, "offset"_a, "src"_a, py::arg("size") = 0 /* not checked */);
+      "dst"_a, "offset"_a, "src"_a, py::arg("size") = 0 /* not checked */,
+      py::arg("concurrency") = memory::default_memcpy_concurrency);
 
   PyModule_AddFunctions(mod.ptr(), vineyard_utils_methods);
 }
@@ -177,7 +178,7 @@ class PyBufferGetter {
 }  // namespace detail
 
 Status copy_memoryview(PyObject* src, void* dst, size_t const size,
-                       size_t const offset) {
+                       size_t const offset, size_t const concurrency) {
   detail::PyBufferGetter src_buffer(src);
   if (!src_buffer.has_buffer()) {
     return Status::AssertionFailed(
@@ -201,15 +202,17 @@ Status copy_memoryview(PyObject* src, void* dst, size_t const size,
   {
     py::gil_scoped_release release;
     // memcpy
-    memory::inline_memcpy(reinterpret_cast<uint8_t*>(dst) + offset,
-                          src_buffer.data(), src_buffer.size());
+    memory::concurrent_memcpy(reinterpret_cast<uint8_t*>(dst) + offset,
+                              src_buffer.data(), src_buffer.size(),
+                              concurrency);
   }
 
   return Status::OK();
 }
 
 Status copy_memoryview_to_memoryview(PyObject* src, PyObject* dst,
-                                     size_t const size, size_t const offset) {
+                                     size_t const size, size_t const offset,
+                                     size_t const concurrency) {
   detail::PyBufferGetter src_buffer(src);
   if (!src_buffer.has_buffer()) {
     return Status::AssertionFailed(
@@ -256,8 +259,8 @@ Status copy_memoryview_to_memoryview(PyObject* src, PyObject* dst,
   {
     py::gil_scoped_release release;
     // memcpy
-    memory::inline_memcpy(dst_buffer.data() + offset, src_buffer.data(),
-                          src_buffer.size());
+    memory::concurrent_memcpy(dst_buffer.data() + offset, src_buffer.data(),
+                              src_buffer.size(), concurrency);
   }
 
   return Status::OK();

diff --git a/python/pybind11_utils.h b/python/pybind11_utils.h
@@ -22,6 +22,7 @@ limitations under the License.
 #include <string>
 #include <utility>
 
+#include "common/memory/memcpy.h"
 #include "common/util/json.h"
 #include "common/util/status.h"
 #include "common/util/uuid.h"
@@ -107,17 +108,18 @@ void throw_on_error(Status const& status);
  *
  * @size: capacity of the dst memory block.
  */
-Status copy_memoryview(PyObject* src, void* dst, size_t const size,
-                       size_t const offset = 0);
+Status copy_memoryview(
+    PyObject* src, void* dst, size_t const size, size_t const offset = 0,
+    size_t const concurrency = memory::default_memcpy_concurrency);
 
 /**
  * Copy a memoryview/buffer to a dst pointer.
  *
  * @size: capacity of the dst memoryview.
  */
-Status copy_memoryview_to_memoryview(PyObject* src, PyObject* dst,
-                                     size_t const size,
-                                     size_t const offset = 0);
+Status copy_memoryview_to_memoryview(
+    PyObject* src, PyObject* dst, size_t const size, size_t const offset = 0,
+    size_t const concurrency = memory::default_memcpy_concurrency);
 
 namespace detail {
 py::object from_json(const json& value);