[Core] Run Plasma Store as a Raylet thread (with a feature flag) (#8897)

* integrate plasma store as a thread (C++)

* integrate plasma store as a thread (Python)

* fix config issues

* remove plasma component fail tests

* without forcefully kill the plasma store thread

src/ray/common/ray_config_def.h

@@ -297,6 +297,9 @@ RAY_CONFIG(int64_t, ping_gcs_rpc_server_interval_milliseconds, 1000)
 /// Maximum number of times to retry ping gcs rpc server when gcs server restarts.
 RAY_CONFIG(int32_t, ping_gcs_rpc_server_max_retries, 600)
 
+// Whether start the Plasma Store as a Raylet thread.
+RAY_CONFIG(bool, plasma_store_as_thread, false)
+
 /// Whether to enable gcs service.
 /// RAY_GCS_SERVICE_ENABLED is an env variable which only set in ci job.
 /// If the value of RAY_GCS_SERVICE_ENABLED is false, we will disable gcs service,

src/ray/object_manager/object_manager.cc

@@ -14,6 +14,8 @@
 
 #include "ray/object_manager/object_manager.h"
 
+#include <chrono>
+
 #include "ray/common/common_protocol.h"
 #include "ray/stats/stats.h"
 #include "ray/util/util.h"
@@ -24,12 +26,34 @@ namespace object_manager_protocol = ray::object_manager::protocol;
 
 namespace ray {
 
+ObjectStoreRunner::ObjectStoreRunner(const ObjectManagerConfig &config) {
+  if (config.object_store_memory > 0) {
+    plasma_store_.reset(new plasma::PlasmaStoreRunner(
+        config.store_socket_name, config.object_store_memory, config.huge_pages,
+        config.plasma_directory, ""));
+    // Initialize object store.
+    store_thread_ = std::thread(&plasma::PlasmaStoreRunner::Start, plasma_store_.get());
+    // Sleep for sometime until the store is working. This can suppress some
+    // connection warnings.
+    std::this_thread::sleep_for(std::chrono::microseconds(500));
+  }
+}
+
+ObjectStoreRunner::~ObjectStoreRunner() {
+  if (plasma_store_ != nullptr) {
+    plasma_store_->Stop();
+    store_thread_.join();
+    plasma_store_.reset();
+  }
+}
+
 ObjectManager::ObjectManager(asio::io_service &main_service, const ClientID &self_node_id,
                              const ObjectManagerConfig &config,
                              std::shared_ptr<ObjectDirectoryInterface> object_directory)
     : self_node_id_(self_node_id),
       config_(config),
       object_directory_(std::move(object_directory)),
+      object_store_internal_(config),
       store_notification_(main_service, config_.store_socket_name),
       buffer_pool_(config_.store_socket_name, config_.object_chunk_size),
       rpc_work_(rpc_service_),

src/ray/object_manager/object_manager.h

@@ -37,6 +37,7 @@
 #include "ray/object_manager/object_directory.h"
 #include "ray/object_manager/object_store_notification_manager.h"
 #include "ray/object_manager/plasma/client.h"
+#include "ray/object_manager/plasma/store.h"
 #include "ray/rpc/object_manager/object_manager_client.h"
 #include "ray/rpc/object_manager/object_manager_server.h"
 
@@ -62,6 +63,12 @@ struct ObjectManagerConfig {
   /// Number of threads of rpc service
   /// Send and receive request in these threads
   int rpc_service_threads_number;
+  /// Initial memory allocation for store.
+  int64_t object_store_memory = -1;
+  /// The directory for shared memory files.
+  std::string plasma_directory;
+  /// Enable huge pages.
+  bool huge_pages;
 };
 
 struct LocalObjectInfo {
@@ -72,6 +79,16 @@ struct LocalObjectInfo {
   std::unordered_map<ClientID, int64_t> recent_pushes;
 };
 
+class ObjectStoreRunner {
+ public:
+  ObjectStoreRunner(const ObjectManagerConfig &config);
+  ~ObjectStoreRunner();
+
+ private:
+  std::unique_ptr<plasma::PlasmaStoreRunner> plasma_store_;
+  std::thread store_thread_;
+};
+
 class ObjectManagerInterface {
  public:
   virtual ray::Status Pull(const ObjectID &object_id) = 0;
@@ -370,6 +387,8 @@ class ObjectManager : public ObjectManagerInterface,
   ClientID self_node_id_;
   const ObjectManagerConfig config_;
   std::shared_ptr<ObjectDirectoryInterface> object_directory_;
+  // Object store runner.
+  ObjectStoreRunner object_store_internal_;
   ObjectStoreNotificationManager store_notification_;
   ObjectBufferPool buffer_pool_;
 

src/ray/plasma/store_exec.cc

@@ -1,6 +1,9 @@
 #include <getopt.h>
 #include <stdio.h>
 
+#include <chrono>
+#include <thread>
+
 #include "ray/object_manager/plasma/store.h"
 // TODO(pcm): Convert getopt and sscanf in the store to use more idiomatic C++
 // and get rid of the next three lines:
@@ -14,35 +17,50 @@ int main(int argc, char *argv[]) {
   std::string plasma_directory;
   std::string external_store_endpoint;
   bool hugepages_enabled = false;
+  bool keep_idle = false;
   int64_t system_memory = -1;
   int c;
-  while ((c = getopt(argc, argv, "s:m:d:e:h")) != -1) {
+  while ((c = getopt(argc, argv, "s:m:d:e:h:z")) != -1) {
     switch (c) {
-      case 'd':
-        plasma_directory = std::string(optarg);
-        break;
-      case 'e':
-        external_store_endpoint = std::string(optarg);
-        break;
-      case 'h':
-        hugepages_enabled = true;
-        break;
-      case 's':
-        socket_name = std::string(optarg);
-        break;
-      case 'm': {
-        char extra;
-        int scanned = sscanf(optarg, "%" SCNd64 "%c", &system_memory, &extra);
-        ARROW_CHECK(scanned == 1);
-        break;
-      }
-      default:
-        exit(-1);
+    case 'd':
+      plasma_directory = std::string(optarg);
+      break;
+    case 'e':
+      external_store_endpoint = std::string(optarg);
+      break;
+    case 'h':
+      hugepages_enabled = true;
+      break;
+    case 's':
+      socket_name = std::string(optarg);
+      break;
+    case 'm': {
+      char extra;
+      int scanned = sscanf(optarg, "%" SCNd64 "%c", &system_memory, &extra);
+      ARROW_CHECK(scanned == 1);
+      break;
+    }
+    case 'z': {
+      keep_idle = true;
+      break;
+    }
+    default:
+      exit(-1);
+    }
+  }
+
+  if (!keep_idle) {
+    plasma::PlasmaStoreRunner runner(socket_name, system_memory, hugepages_enabled,
+                                     plasma_directory, external_store_endpoint);
+    runner.Start();
+  } else {
+    printf(
+        "The Plasma Store is started with the '-z' flag, "
+        "and it will run idle as a placeholder.");
+    while (true) {
+      std::this_thread::sleep_for(std::chrono::hours(1000));
     }
   }
 
-  plasma::PlasmaStoreRunner runner(socket_name, system_memory, hugepages_enabled,
-                                   plasma_directory, external_store_endpoint);
-  runner.Start();
   return 0;
 }

src/ray/raylet/main.cc

@@ -48,6 +48,10 @@ DEFINE_bool(disable_stats, false, "Whether disable the stats.");
 DEFINE_string(stat_address, "127.0.0.1:8888", "The address that we report metrics to.");
 DEFINE_bool(enable_stdout_exporter, false,
             "Whether enable the stdout exporter for stats.");
+// store options
+DEFINE_int64(object_store_memory, -1, "The initial memory of the object store.");
+DEFINE_string(plasma_directory, "", "The shared memory directory of the object store.");
+DEFINE_bool(huge_pages, false, "Whether enable huge pages");
 
 #ifndef RAYLET_TEST
 
@@ -81,6 +85,9 @@ int main(int argc, char *argv[]) {
   const bool disable_stats = FLAGS_disable_stats;
   const std::string stat_address = FLAGS_stat_address;
   const bool enable_stdout_exporter = FLAGS_enable_stdout_exporter;
+  const int64_t object_store_memory = FLAGS_object_store_memory;
+  const std::string plasma_directory = FLAGS_plasma_directory;
+  const bool huge_pages = FLAGS_huge_pages;
   gflags::ShutDownCommandLineFlags();
 
   // Initialize stats.
@@ -166,6 +173,9 @@ int main(int argc, char *argv[]) {
       RayConfig::instance().object_manager_pull_timeout_ms();
   object_manager_config.push_timeout_ms =
       RayConfig::instance().object_manager_push_timeout_ms();
+  object_manager_config.object_store_memory = object_store_memory;
+  object_manager_config.plasma_directory = plasma_directory;
+  object_manager_config.huge_pages = huge_pages;
 
   int num_cpus = static_cast<int>(static_resource_conf["CPU"]);
   object_manager_config.rpc_service_threads_number =