Enable a memory_trim option in DelData (#1775)

Signed-off-by: Tao He <[email protected]>

.gitignore

@@ -25,6 +25,9 @@ cmake-build-debug
 *.pyc
 *.bin
 *.whl
+*.tar.gz
+*.zip
+*.xz
 
 # hive sql work directory
 /java/hive/docker/dependency/mysql/conf/

.pre-commit-config.yaml

@@ -5,7 +5,3 @@ repos:
       - id: gitleaks
         args:
         - '--verbose'
-  - repo: https://github.com/crate-ci/typos
-    rev: v1.18.2
-    hooks:
-      - id: typos

modules/graph/writer/util.cc

@@ -14,15 +14,17 @@
  * limitations under the License.
  */
 
+#include "graph/writer/util.h"
+
+#ifdef ENABLE_GAR
+
 #include <unordered_map>
 #include <vector>
 
 #include "gar/util/adj_list_type.h"
 #include "gar/util/data_type.h"
 #include "gar/util/file_type.h"
 
-#include "graph/writer/util.h"
-
 namespace GAR = GraphArchive;
 
 namespace vineyard {
@@ -71,3 +73,5 @@ std::shared_ptr<GraphArchive::GraphInfo> generate_graph_info_with_schema(
                               nullptr, extra_info);
 }
 }  // namespace vineyard
+
+#endif

modules/graph/writer/util.h

@@ -16,6 +16,8 @@
 #ifndef MODULES_GRAPH_WRITER_UTIL_H_
 #define MODULES_GRAPH_WRITER_UTIL_H_
 
+#ifdef ENABLE_GAR
+
 #include <memory>
 #include <string>
 
@@ -37,4 +39,6 @@ std::shared_ptr<GraphArchive::GraphInfo> generate_graph_info_with_schema(
 
 }  // namespace vineyard
 
+#endif
+
 #endif  // MODULES_GRAPH_WRITER_UTIL_H_

python/client.cc

@@ -257,36 +257,42 @@ void bind_client(py::module& mod) {
       .def(
           "delete",
           [](ClientBase* self, const ObjectIDWrapper object_id,
-             const bool force, const bool deep) {
-            throw_on_error(self->DelData(object_id, force, deep));
+             const bool force, const bool deep, const bool memory_trim) {
+            throw_on_error(self->DelData(object_id, force, deep, memory_trim));
           },
           "object_id"_a, py::arg("force") = false, py::arg("deep") = true,
-          doc::ClientBase_delete)
+          py::arg("memory_trim") = false, doc::ClientBase_delete)
       .def(
           "delete",
           [](ClientBase* self, const std::vector<ObjectIDWrapper>& object_ids,
-             const bool force, const bool deep) {
+             const bool force, const bool deep, const bool memory_trim) {
             std::vector<ObjectID> unwrapped_object_ids(object_ids.size());
             for (size_t idx = 0; idx < object_ids.size(); ++idx) {
               unwrapped_object_ids[idx] = object_ids[idx];
             }
-            throw_on_error(self->DelData(unwrapped_object_ids, force, deep));
+            throw_on_error(
+                self->DelData(unwrapped_object_ids, force, deep, memory_trim));
           },
-          "object_ids"_a, py::arg("force") = false, py::arg("deep") = true)
+          "object_ids"_a, py::arg("force") = false, py::arg("deep") = true,
+          py::arg("memory_trim") = false, doc::ClientBase_delete)
       .def(
           "delete",
           [](ClientBase* self, const ObjectMeta& meta, const bool force,
-             const bool deep) {
-            throw_on_error(self->DelData(meta.GetId(), force, deep));
+             const bool deep, const bool memory_trim) {
+            throw_on_error(
+                self->DelData(meta.GetId(), force, deep, memory_trim));
           },
-          "object_meta"_a, py::arg("force") = false, py::arg("deep") = true)
+          "object_meta"_a, py::arg("force") = false, py::arg("deep") = true,
+          py::arg("memory_trim") = false, doc::ClientBase_delete)
       .def(
           "delete",
           [](ClientBase* self, const Object* object, const bool force,
-             const bool deep) {
-            throw_on_error(self->DelData(object->id(), force, deep));
+             const bool deep, const bool memory_trim) {
+            throw_on_error(
+                self->DelData(object->id(), force, deep, memory_trim));
           },
-          "object"_a, py::arg("force") = false, py::arg("deep") = true)
+          "object"_a, py::arg("force") = false, py::arg("deep") = true,
+          py::arg("memory_trim") = false, doc::ClientBase_delete)
       .def(
           "create_stream",
           [](ClientBase* self, ObjectID const id) {

python/pybind11_docs.cc

@@ -571,7 +571,7 @@ Create metadata in vineyardd with a specified instance id.
 
 const char* ClientBase_delete = R"doc(
 .. method:: delete(object_id: ObjectID or List[ObjectID], force: bool = false,
-                   deep: bool = true) -> None
+                   deep: bool = true, memory_trim: bool = False) -> None
     :noindex:
 
 Delete the specific vineyard object.
@@ -590,17 +590,25 @@ Delete the specific vineyard object.
 
         Note that when deleting objects which have *direct* blob members, the
         processing on those blobs yields a "deep" behavior.
+    memory_trim: bool
+        Whether to trim the memory pool inside the shared memory allocator to
+        return the unused physical memory back to the OS kernel, like the
+        `malloc_trim` API from glibc. Default value is False.
 
-.. method:: delete(object_meta: ObjectMeta, force: bool = false, deep: bool = true)
-                   -> None
+        Note that the memory trimming operation is a best-effort operation and
+        may not release any memory at all.
+
+.. method:: delete(object_meta: ObjectMeta, force: bool = false, deep: bool = true,
+                   memory_trim: bool = False) -> None
     :noindex:
 
 Delete the specific vineyard object.
 
 Parameters:
     object_meta: The corresponding object meta to delete.
 
-.. method:: delete(object: Object, force: bool = false, deep: bool = true) -> None
+.. method:: delete(object: Object, force: bool = false, deep: bool = true,
+                   memory_trim: bool = False) -> None
     :noindex:
 
 Delete the specific vineyard object.

python/vineyard/_C.pyi

@@ -234,6 +234,7 @@ class ClientBase:
         object: Union[ObjectID, Object, ObjectMeta, List[ObjectID]],
         force: bool = False,
         deep: bool = True,
+        memory_trim: bool = False,
     ) -> None: ...
     def create_stream(self, id: ObjectID) -> None: ...
     def open_stream(self, id: ObjectID, mode: str) -> None: ...

python/vineyard/core/client.py

@@ -265,8 +265,9 @@ def delete(
         object: Union[ObjectID, Object, ObjectMeta, List[ObjectID]],
         force: bool = False,
         deep: bool = True,
+        memory_trim: bool = False,
     ) -> None:
-        return self.default_client().delete(object, force, deep)
+        return self.default_client().delete(object, force, deep, memory_trim)
 
     @_apply_docstring(IPCClient.create_stream)
     def create_stream(self, id: ObjectID) -> None:

python/vineyard/core/tests/test_client.py

@@ -313,11 +313,7 @@ def test_memory_trim(vineyard_client):
         assert current_memory_usage >= original_memory_usage + i * data_kbytes
 
     for r in rs:
-        vineyard_client.delete(r)
+        vineyard_client.delete(r, memory_trim=True)
 
-    current_memory_usage = parse_shared_memory_usage()
-    assert current_memory_usage >= original_memory_usage + (8 - 1) * data_kbytes
-
-    vineyard_client.memory_trim()
     # there might be some fragmentation overhead
     assert parse_shared_memory_usage() <= original_memory_usage + 2 * data_kbytes

src/client/client.cc

@@ -918,19 +918,29 @@ Status Client::Release(ObjectID const& id) {
 }
 
 Status Client::DelData(const ObjectID id, const bool force, const bool deep) {
-  return this->DelData(std::vector<ObjectID>{id}, force, deep);
+  return this->DelData(id, force, deep, false);
+}
+
+Status Client::DelData(const ObjectID id, const bool force, const bool deep,
+                       const bool memory_trim) {
+  return this->DelData(std::vector<ObjectID>{id}, force, deep, memory_trim);
 }
 
 Status Client::DelData(const std::vector<ObjectID>& ids, const bool force,
                        const bool deep) {
+  return this->DelData(ids, force, deep, false);
+}
+
+Status Client::DelData(const std::vector<ObjectID>& ids, const bool force,
+                       const bool deep, const bool memory_trim) {
   ENSURE_CONNECTED(this);
   for (auto id : ids) {
     // May contain duplicated blob ids.
     VINEYARD_DISCARD(Release(id));
   }
   std::string message_out;
-  WriteDelDataWithFeedbacksRequest(ids, force, deep, /*fastpath=*/false,
-                                   message_out);
+  WriteDelDataWithFeedbacksRequest(ids, force, deep, memory_trim,
+                                   /*fastpath=*/false, message_out);
   RETURN_ON_ERROR(doWrite(message_out));
   json message_in;
   std::vector<ObjectID> deleted_bids;

src/client/client.h

@@ -781,6 +781,10 @@ class Client final : public BasicIPCClient,
    */
   Status DelData(const ObjectID id, const bool force = false,
                  const bool deep = true);
+
+  Status DelData(const ObjectID id, const bool force, const bool deep,
+                 const bool memory_trim);
+
   /**
    * @brief Delete multiple metadatas in vineyard.
    *
@@ -797,6 +801,9 @@ class Client final : public BasicIPCClient,
   Status DelData(const std::vector<ObjectID>& ids, const bool force = false,
                  const bool deep = true);
 
+  Status DelData(const std::vector<ObjectID>& ids, const bool force,
+                 const bool deep, const bool memory_trim);
+
   /**
    * @brief Create a GPU buffer on vineyard server. See also `CreateBuffer`.
    *

src/client/client_base.cc

@@ -198,14 +198,24 @@ Status ClientBase::SyncMetaData() {
 
 Status ClientBase::DelData(const ObjectID id, const bool force,
                            const bool deep) {
-  return DelData(std::vector<ObjectID>{id}, force, deep);
+  return DelData(id, force, deep, false);
+}
+
+Status ClientBase::DelData(const ObjectID id, const bool force, const bool deep,
+                           const bool memory_trim) {
+  return DelData(std::vector<ObjectID>{id}, force, deep, memory_trim);
 }
 
 Status ClientBase::DelData(const std::vector<ObjectID>& ids, const bool force,
                            const bool deep) {
+  return DelData(ids, force, deep, false);
+}
+
+Status ClientBase::DelData(const std::vector<ObjectID>& ids, const bool force,
+                           const bool deep, const bool memory_trim) {
   ENSURE_CONNECTED(this);
   std::string message_out;
-  WriteDelDataRequest(ids, force, deep, false, message_out);
+  WriteDelDataRequest(ids, force, deep, memory_trim, false, message_out);
   RETURN_ON_ERROR(doWrite(message_out));
   json message_in;
   RETURN_ON_ERROR(doRead(message_in));

src/client/client_base.h

@@ -170,11 +170,17 @@ class ClientBase {
    * @param deep Whether to delete the member of this object. Default is true.
    *        Note that when deleting object which has *direct* blob members, the
    *        processing on those blobs yields a "deep" behavior.
+   * @param memory_trim Whether to trim the memory pool inside the shared memory
+   *        allocator to return the unused physical memory back to the OS.
    *
    * @return Status that indicates whether the delete action has succeeded.
    */
   Status DelData(const ObjectID id, const bool force = false,
                  const bool deep = true);
+
+  Status DelData(const ObjectID id, const bool force, const bool deep,
+                 const bool memory_trim);
+
   /**
    * @brief Delete multiple metadatas in vineyard.
    *
@@ -185,12 +191,17 @@ class ClientBase {
    * @param deep Whether to delete the member of this object. Default is true.
    *        Note that when deleting objects which have *direct* blob members,
    *        the processing on those blobs yields a "deep" behavior.
+   * @param memory_trim Whether to trim the memory pool inside the shared memory
+   *        allocator to return the unused physical memory back to the OS.
    *
    * @return Status that indicates whether the delete action has succeeded.
    */
   Status DelData(const std::vector<ObjectID>& ids, const bool force = false,
                  const bool deep = true);
 
+  Status DelData(const std::vector<ObjectID>& ids, const bool force,
+                 const bool deep, const bool memory_trim);
+
   /**
    * @brief List objectmetas in vineyard, using the given typename patterns.
    *

src/common/util/protocols.cc

@@ -858,24 +858,28 @@ Status ReadReleaseReply(json const& root) {
 
 void WriteDelDataWithFeedbacksRequest(const std::vector<ObjectID>& id,
                                       const bool force, const bool deep,
+                                      const bool memory_trim,
                                       const bool fastpath, std::string& msg) {
   json root;
   root["type"] = command_t::DEL_DATA_WITH_FEEDBACKS_REQUEST;
   root["id"] = std::vector<ObjectID>{id};
   root["force"] = force;
   root["deep"] = deep;
+  root["memory_trim"] = memory_trim;
   root["fastpath"] = fastpath;
 
   encode_msg(root, msg);
 }
 
 Status ReadDelDataWithFeedbacksRequest(json const& root,
                                        std::vector<ObjectID>& ids, bool& force,
-                                       bool& deep, bool& fastpath) {
+                                       bool& deep, bool& memory_trim,
+                                       bool& fastpath) {
   CHECK_IPC_ERROR(root, command_t::DEL_DATA_WITH_FEEDBACKS_REQUEST);
   root["id"].get_to(ids);
   force = root.value("force", false);
   deep = root.value("deep", false);
+  memory_trim = root.value("memory_trim", false);
   fastpath = root.value("fastpath", false);
   return Status::OK();
 }
@@ -1213,37 +1217,42 @@ Status ReadListDataRequest(const json& root, std::string& pattern, bool& regex,
 }
 
 void WriteDelDataRequest(const ObjectID id, const bool force, const bool deep,
-                         const bool fastpath, std::string& msg) {
+                         const bool memory_trim, const bool fastpath,
+                         std::string& msg) {
   json root;
   root["type"] = command_t::DELETE_DATA_REQUEST;
   root["id"] = std::vector<ObjectID>{id};
   root["force"] = force;
   root["deep"] = deep;
   root["fastpath"] = fastpath;
+  root["memory_trim"] = memory_trim;
 
   encode_msg(root, msg);
 }
 
 void WriteDelDataRequest(const std::vector<ObjectID>& ids, const bool force,
-                         const bool deep, const bool fastpath,
-                         std::string& msg) {
+                         const bool deep, const bool memory_trim,
+                         const bool fastpath, std::string& msg) {
   json root;
   root["type"] = command_t::DELETE_DATA_REQUEST;
   root["id"] = ids;
   root["force"] = force;
   root["deep"] = deep;
   root["fastpath"] = fastpath;
+  root["memory_trim"] = memory_trim;
 
   encode_msg(root, msg);
 }
 
 Status ReadDelDataRequest(const json& root, std::vector<ObjectID>& ids,
-                          bool& force, bool& deep, bool& fastpath) {
+                          bool& force, bool& deep, bool& memory_trim,
+                          bool& fastpath) {
   CHECK_IPC_ERROR(root, command_t::DELETE_DATA_REQUEST);
   root["id"].get_to(ids);
   force = root.value("force", false);
   deep = root.value("deep", false);
   fastpath = root.value("fastpath", false);
+  memory_trim = root.value("memory_trim", false);
   return Status::OK();
 }