Stop vendoring pyarrow (#7233)

src/ray/core_worker/core_worker.cc

@@ -103,7 +103,6 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
     RayLog::InstallFailureSignalHandler();
   }
   RAY_LOG(INFO) << "Initializing worker " << worker_context_.GetWorkerID();
-
   // Initialize gcs client.
   gcs_client_ = std::make_shared<gcs::RedisGcsClient>(gcs_options);
   RAY_CHECK_OK(gcs_client_->Connect(io_service_));
@@ -247,9 +246,15 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
   if (direct_task_receiver_ != nullptr) {
     direct_task_receiver_->Init(client_factory, rpc_address_);
   }
+  plasma_notifier_.reset(new ObjectStoreNotificationManager(io_service_, store_socket,
+                                                            /*exit_on_error*/ false));
 }
 
 CoreWorker::~CoreWorker() {
+  // ObjectStoreNotificationManager depends on io_service_ so we need to shut it down
+  // first.
+  plasma_notifier_->Shutdown();
+
   io_service_.stop();
   io_thread_.join();
   if (log_dir_ != "") {
@@ -1342,6 +1347,14 @@ void CoreWorker::GetAsync(const ObjectID &object_id, SetResultCallback success_c
   });
 }
 
+void CoreWorker::SubscribeToAsyncPlasma(PlasmaSubscriptionCallback subscribe_callback) {
+  plasma_notifier_->SubscribeObjAdded(
+      [subscribe_callback](const object_manager::protocol::ObjectInfoT &info) {
+        subscribe_callback(ObjectID::FromPlasmaIdBinary(info.object_id), info.data_size,
+                           info.metadata_size);
+      });
+}
+
 void CoreWorker::SetActorId(const ActorID &actor_id) {
   absl::MutexLock lock(&mutex_);
   RAY_CHECK(actor_id_.IsNil());

src/ray/core_worker/core_worker.h

@@ -18,6 +18,7 @@
 #include "ray/core_worker/transport/raylet_transport.h"
 #include "ray/gcs/redis_gcs_client.h"
 #include "ray/gcs/subscription_executor.h"
+#include "ray/object_manager/object_store_notification_manager.h"
 #include "ray/raylet/raylet_client.h"
 #include "ray/rpc/node_manager/node_manager_client.h"
 #include "ray/rpc/worker/core_worker_client.h"
@@ -507,6 +508,15 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
   void GetAsync(const ObjectID &object_id, SetResultCallback success_callback,
                 SetResultCallback fallback_callback, void *python_future);
 
+  /// Connect to plasma store for async futures
+  using PlasmaSubscriptionCallback = std::function<void(ray::ObjectID, int64_t, int64_t)>;
+
+  /// Subscribe to plasma store
+  ///
+  /// \param[in] subscribe_callback The callback when an item is added to plasma.
+  /// \return void
+  void SubscribeToAsyncPlasma(PlasmaSubscriptionCallback subscribe_callback);
+
  private:
   /// Run the io_service_ event loop. This should be called in a background thread.
   void RunIOService();
@@ -766,6 +776,9 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
   // Queue of tasks to resubmit when the specified time passes.
   std::deque<std::pair<int64_t, TaskSpecification>> to_resubmit_ GUARDED_BY(mutex_);
 
+  // Plasma notification manager
+  std::unique_ptr<ObjectStoreNotificationManager> plasma_notifier_;
+
   friend class CoreWorkerTest;
 };
 

src/ray/object_manager/object_store_notification_manager.cc

@@ -14,12 +14,14 @@
 namespace ray {
 
 ObjectStoreNotificationManager::ObjectStoreNotificationManager(
-    boost::asio::io_service &io_service, const std::string &store_socket_name)
+    boost::asio::io_service &io_service, const std::string &store_socket_name,
+    bool exit_on_error)
     : store_client_(),
       length_(0),
       num_adds_processed_(0),
       num_removes_processed_(0),
-      socket_(io_service) {
+      socket_(io_service),
+      exit_on_error_(exit_on_error) {
   RAY_ARROW_CHECK_OK(store_client_.Connect(store_socket_name.c_str(), "", 0, 300));
 
   int c_socket;  // TODO(mehrdadn): This should be type SOCKET for Windows
@@ -57,6 +59,10 @@ ObjectStoreNotificationManager::~ObjectStoreNotificationManager() {
   RAY_ARROW_CHECK_OK(store_client_.Disconnect());
 }
 
+void ObjectStoreNotificationManager::Shutdown() {
+  RAY_ARROW_CHECK_OK(store_client_.Disconnect());
+}
+
 void ObjectStoreNotificationManager::NotificationWait() {
   boost::asio::async_read(socket_, boost::asio::buffer(&length_, sizeof(length_)),
                           boost::bind(&ObjectStoreNotificationManager::ProcessStoreLength,
@@ -67,15 +73,26 @@ void ObjectStoreNotificationManager::ProcessStoreLength(
     const boost::system::error_code &error) {
   notification_.resize(length_);
   if (error) {
-    // When shutting down a cluster, it's possible that the plasma store is killed
-    // earlier than raylet, in this case we don't want raylet to crash, we instead
-    // log an error message and exit.
-    RAY_LOG(ERROR) << "Failed to process store length: "
-                   << boost_to_ray_status(error).ToString()
-                   << ", most likely plasma store is down, raylet will exit";
-    // Exit raylet process.
-    _exit(kRayletStoreErrorExitCode);
+    if (exit_on_error_) {
+      // When shutting down a cluster, it's possible that the plasma store is killed
+      // earlier than raylet. In this case we don't want raylet to crash, we instead
+      // log an error message and exit.
+      RAY_LOG(ERROR) << "Failed to process store length: "
+                     << boost_to_ray_status(error).ToString()
+                     << ", most likely plasma store is down, raylet will exit";
+      // Exit raylet process.
+      _exit(kRayletStoreErrorExitCode);
+    } else {
+      // The log level is set to debug so user don't see it on ctrl+c exit.
+      RAY_LOG(DEBUG) << "Failed to process store length: "
+                     << boost_to_ray_status(error).ToString()
+                     << ", most likely plasma store is down. "
+                     << "The error is silenced because exit_on_error_ "
+                     << "flag is set.";
+      return;
+    }
   }
+
   boost::asio::async_read(
       socket_, boost::asio::buffer(notification_),
       boost::bind(&ObjectStoreNotificationManager::ProcessStoreNotification, this,
@@ -85,9 +102,19 @@ void ObjectStoreNotificationManager::ProcessStoreLength(
 void ObjectStoreNotificationManager::ProcessStoreNotification(
     const boost::system::error_code &error) {
   if (error) {
-    RAY_LOG(FATAL)
-        << "Problem communicating with the object store from raylet, check logs or "
-        << "dmesg for previous errors: " << boost_to_ray_status(error).ToString();
+    if (exit_on_error_) {
+      RAY_LOG(FATAL)
+          << "Problem communicating with the object store from raylet, check logs or "
+          << "dmesg for previous errors: " << boost_to_ray_status(error).ToString();
+    } else {
+      // The log level is set to debug so user don't see it on ctrl+c exit.
+      RAY_LOG(DEBUG)
+          << "Problem communicating with the object store from raylet, check logs or "
+          << "dmesg for previous errors: " << boost_to_ray_status(error).ToString()
+          << " The error is silenced because exit_on_error_ "
+          << "flag is set.";
+      return;
+    }
   }
 
   const auto &object_notification =

src/ray/object_manager/object_store_notification_manager.h

@@ -28,8 +28,11 @@ class ObjectStoreNotificationManager {
   ///
   /// \param io_service The asio service to be used.
   /// \param store_socket_name The store socket to connect to.
+  /// \param exit_on_error The manager will exit with error when it fails
+  ///                      to process messages from socket.
   ObjectStoreNotificationManager(boost::asio::io_service &io_service,
-                                 const std::string &store_socket_name);
+                                 const std::string &store_socket_name,
+                                 bool exit_on_error = true);
 
   ~ObjectStoreNotificationManager();
 
@@ -46,6 +49,9 @@ class ObjectStoreNotificationManager {
   /// \param callback A callback expecting an ObjectID.
   void SubscribeObjDeleted(std::function<void(const ray::ObjectID &)> callback);
 
+  /// Explicitly shutdown the manager.
+  void Shutdown();
+
   /// Returns debug string for class.
   ///
   /// \return string.
@@ -71,6 +77,13 @@ class ObjectStoreNotificationManager {
   int64_t num_removes_processed_;
   std::vector<uint8_t> notification_;
   local_stream_protocol::socket socket_;
+
+  /// Flag to indicate whether or not to exit the process when received socket
+  /// error. When it is false, socket error will be ignored. This flag is needed
+  /// when running object store notification manager in core worker. On core worker
+  /// exit, plasma store will be killed before deconstructor of this manager. So we
+  /// we have to silence the errors.
+  bool exit_on_error_;
 };
 
 }  // namespace ray