Use Plasma with LRU refreshing integrated (#6050)

bazel/BUILD.plasma

@@ -9,19 +9,28 @@ cc_library(
         "cpp/src/arrow/io/interfaces.cc",
         "cpp/src/arrow/memory_pool.cc",
         "cpp/src/arrow/status.cc",
-        "cpp/src/arrow/util/io-util.cc",
+        "cpp/src/arrow/util/io_util.cc",
         "cpp/src/arrow/util/logging.cc",
         "cpp/src/arrow/util/memory.cc",
+        "cpp/src/arrow/util/string.cc",
         "cpp/src/arrow/util/string_builder.cc",
-        "cpp/src/arrow/util/thread-pool.cc",
+        "cpp/src/arrow/util/thread_pool.cc",
     ],
     hdrs = [
         "cpp/src/arrow/buffer.h",
+        "cpp/src/arrow/io/concurrency.h",
         "cpp/src/arrow/io/interfaces.h",
+        "cpp/src/arrow/io/util_internal.h",
         "cpp/src/arrow/memory_pool.h",
+        "cpp/src/arrow/result.h",
         "cpp/src/arrow/status.h",
-        "cpp/src/arrow/util/bit-util.h",
-        "cpp/src/arrow/util/io-util.h",
+        "cpp/src/arrow/type_fwd.h",
+        "cpp/src/arrow/util/bit_util.h",
+        "cpp/src/arrow/util/checked_cast.h",
+        "cpp/src/arrow/util/compare.h",
+        "cpp/src/arrow/util/functional.h",
+        "cpp/src/arrow/util/io_util.h",
+        "cpp/src/arrow/util/iterator.h",
         "cpp/src/arrow/util/logging.h",
         "cpp/src/arrow/util/macros.h",
         "cpp/src/arrow/util/memory.h",
@@ -29,12 +38,16 @@ cc_library(
         "cpp/src/arrow/util/string.h",
         "cpp/src/arrow/util/string_builder.h",
         "cpp/src/arrow/util/string_view.h",
-        "cpp/src/arrow/util/thread-pool.h",
+        "cpp/src/arrow/util/thread_pool.h",
         "cpp/src/arrow/util/type_traits.h",
         "cpp/src/arrow/util/ubsan.h",
+        "cpp/src/arrow/util/variant.h",
         "cpp/src/arrow/util/visibility.h",
         "cpp/src/arrow/util/windows_compatibility.h",
         "cpp/src/arrow/vendored/string_view.hpp",
+        "cpp/src/arrow/vendored/variant.hpp",
+        "cpp/src/arrow/vendored/xxhash.h",
+        "cpp/src/arrow/vendored/xxhash/xxh3.h",
         "cpp/src/arrow/vendored/xxhash/xxhash.c",
         "cpp/src/arrow/vendored/xxhash/xxhash.h",
     ],
@@ -176,17 +189,17 @@ FLATC_ARGS = [
 
 flatbuffer_cc_library(
     name = "common_fbs",
-    srcs = ["cpp/src/plasma/format/common.fbs"],
+    srcs = ["cpp/src/plasma/common.fbs"],
     flatc_args = FLATC_ARGS,
     out_prefix = "cpp/src/plasma/",
 )
 
 flatbuffer_cc_library(
     name = "plasma_fbs",
-    srcs = ["cpp/src/plasma/format/plasma.fbs"],
+    srcs = ["cpp/src/plasma/plasma.fbs"],
     flatc_args = FLATC_ARGS,
-    includes = ["cpp/src/plasma/format/common.fbs"],
+    includes = ["cpp/src/plasma/common.fbs"],
     out_prefix = "cpp/src/plasma/",
 )
 
-exports_files(["cpp/src/plasma/format/common.fbs"])
+exports_files(["cpp/src/plasma/common.fbs"])

bazel/ray_deps_build_all.bzl

@@ -17,4 +17,3 @@ def ray_deps_build_all():
   grpc_deps()
   java_proto_compile()
   python_proto_compile()
-

bazel/ray_deps_setup.bzl

@@ -74,7 +74,7 @@ def ray_deps_setup():
     new_git_repository(
         name = "plasma",
         build_file = "@//bazel:BUILD.plasma",
-        commit = "141a213a54f4979ab0b94b94928739359a2ee9ad",
+        commit = "86f34aa07e611787d9cc98c6a33b0a0a536dce57",
         remote = "https://github.com/apache/arrow",
     )
 

build.sh

@@ -99,7 +99,7 @@ pushd "$BUILD_DIR"
 if [ -z "$SKIP_PYARROW_INSTALL" ]; then
     "$PYTHON_EXECUTABLE" -m pip install -q \
         --target="$ROOT_DIR/python/ray/pyarrow_files" pyarrow==0.14.0.RAY \
-        --find-links https://s3-us-west-2.amazonaws.com/arrow-wheels/516e15028091b5e287200b5df77d77f72d9a6c9a/index.html
+        --find-links https://s3-us-west-2.amazonaws.com/arrow-wheels/3a11193d9530fe8ec7fdb98057f853b708f6f6ae/index.html
 fi
 export PYTHON_BIN_PATH="$PYTHON_EXECUTABLE"
 

python/ray/experimental/async_api.py

@@ -72,7 +72,7 @@ async def _async_init():
     if handler is None:
         worker = ray.worker.global_worker
         plasma_client = thread_safe_client(
-            plasma.connect(worker.node.plasma_store_socket_name, None, 0, 300))
+            plasma.connect(worker.node.plasma_store_socket_name, 300))
         loop = asyncio.get_event_loop()
         plasma_client.subscribe()
         rsock = plasma_client.get_notification_socket()

python/ray/experimental/async_plasma.py

@@ -39,7 +39,12 @@ class PlasmaProtocol(asyncio.Protocol):
             i += INT64_SIZE
             segment = self._buffer[i:i + msg_len]
             i += msg_len
-            messages.append(self.plasma_client.decode_notification(segment))
+            (object_ids, object_sizes,
+             metadata_sizes) = self.plasma_client.decode_notifications(segment)
+            assert len(object_ids) == len(object_sizes) == len(metadata_sizes)
+            for j in range(len(object_ids)):
+                messages.append((object_ids[j], object_sizes[j],
+                                 metadata_sizes[j]))
 
         self._buffer = self._buffer[i:]
         self.plasma_event_handler.process_notifications(messages)

src/ray/object_manager/format/object_manager.fbs

@@ -3,7 +3,7 @@ namespace ray.object_manager.protocol;
 
 // Object information data structure.
 // NOTE(pcm): This structure is replicated in
-// https://github.com/apache/arrow/blob/master/cpp/src/plasma/format/common.fbs,
+// https://github.com/apache/arrow/blob/master/cpp/src/plasma/common.fbs,
 // so if you modify it, you should also modify that one.
 table ObjectInfo {
   // Object ID of this object.
@@ -24,3 +24,10 @@ table ObjectInfo {
   // Specifies if this object was deleted or added.
   is_deletion: bool;
 }
+
+// NOTE(pcm): This structure is replicated in
+// https://github.com/apache/arrow/blob/master/cpp/src/plasma/plasma.fbs
+// so if you modify it, you should also modify that one.
+table PlasmaNotification {
+  object_info: [ObjectInfo];
+}

src/ray/object_manager/object_store_notification_manager.cc

@@ -61,16 +61,20 @@ void ObjectStoreNotificationManager::ProcessStoreNotification(
         << "dmesg for previous errors: " << boost_to_ray_status(error).ToString();
   }
 
-  const auto &object_info =
-      flatbuffers::GetRoot<object_manager::protocol::ObjectInfo>(notification_.data());
-  const ObjectID object_id =
-      ObjectID::FromPlasmaIdBinary(object_info->object_id()->str());
-  if (object_info->is_deletion()) {
-    ProcessStoreRemove(object_id);
-  } else {
-    object_manager::protocol::ObjectInfoT result;
-    object_info->UnPackTo(&result);
-    ProcessStoreAdd(result);
+  const auto &object_notification =
+      flatbuffers::GetRoot<object_manager::protocol::PlasmaNotification>(
+          notification_.data());
+  for (size_t i = 0; i < object_notification->object_info()->size(); ++i) {
+    auto object_info = object_notification->object_info()->Get(i);
+    const ObjectID object_id =
+        ObjectID::FromPlasmaIdBinary(object_info->object_id()->str());
+    if (object_info->is_deletion()) {
+      ProcessStoreRemove(object_id);
+    } else {
+      object_manager::protocol::ObjectInfoT result;
+      object_info->UnPackTo(&result);
+      ProcessStoreAdd(result);
+    }
   }
   NotificationWait();
 }