[Core] Exit the Core Worker Early Error Received from Plasma Store (#53679)

What's the issue:
- During node shutdown, when the raylet is killed before its core
workers, and the tasks on the core workers read/write objects from the
plasma store, a broken pipe error will be obtained and the tasks will
fail due to the ray task error with reason broken pipe and thus the
whole job failed.
- This is not the desired behavior because the task failure due to node
shutdown should be seen as the system failure and the core worker
shouldn't continue executing tasks when the raylet is down.

The PR made the change to mitigate the above issue: 
- In the plasmas store client, add the logic to do core worker quick
exit when error happens during read/write buffer and the plasma store
client is on the core worker side

Test the logic manually to verify the behavior:
- With the following test code:
```
ray.init()

@ray.remote(max_retries=2)
def test_task(obj_ref):
    time.sleep(1)
    raylet_pid = int(os.environ["RAY_RAYLET_PID"])
    os.kill(raylet_pid, signal.SIGKILL)
    ray.put(obj_ref)

a = ray.put([0] * 250000)
ray.get(test_task.remote(a))
```
- Without the change:
```
ray.exceptions.RayTaskError(OSError): ray::test_task() (pid=30681, ip=127.0.0.1)
  File "/Users/myan/ray-core-quickstart/test-tasks/test-tasks.py", line 18, in test_task
    ray.get(test_ref)
  File "python/ray/includes/common.pxi", line 93, in ray._raylet.check_status
    raise IOError(message)
OSError: Failed to read data from the socket: End of file
```
- With the change in the PR:
```
ray.exceptions.LocalRayletDiedError: The task's local raylet died. Check raylet.out for more information.
```
---------

Signed-off-by: Mengjin Yan <mengjinyan3@gmail.com>
Co-authored-by: Ibrahim Rabbani <israbbani@gmail.com>

Author: MengjinYan

python/ray/tests/test_failure_3.py

@@ -47,6 +47,31 @@ def get_raylet_pid(self):
     wait_for_pid_to_exit(worker_pid)
 
 
+def test_plasma_store_operation_after_raylet_dies(ray_start_cluster):
+    """
+    Test that the operation on the plasma store after the raylet dies will not fail the
+    task with an application level error (RayTaskError) but a system level error
+    (RayletDiedError).
+    """
+    cluster = ray_start_cluster
+    cluster.add_node(num_cpus=1)
+    cluster.wait_for_nodes()
+
+    ray.init(address=cluster.address)
+
+    @ray.remote
+    def get_after_raylet_dies():
+        raylet_pid = int(os.environ["RAY_RAYLET_PID"])
+        os.kill(raylet_pid, SIGKILL)
+        wait_for_pid_to_exit(raylet_pid)
+        ray.put([0] * 100000)
+
+    try:
+        ray.get(get_after_raylet_dies.remote(), timeout=10)
+    except Exception as e:
+        assert isinstance(e, ray.exceptions.LocalRayletDiedError)
+
+
 @pytest.mark.parametrize(
     "ray_start_cluster_head",
     [

src/ray/common/client_connection.h

@@ -99,7 +99,7 @@ class ServerConnection : public std::enable_shared_from_this<ServerConnection> {
   ///
   /// \param buffer The buffer.
   /// \return Status.
-  Status WriteBuffer(const std::vector<boost::asio::const_buffer> &buffer);
+  virtual Status WriteBuffer(const std::vector<boost::asio::const_buffer> &buffer);
 
   /// Write a buffer to this connection asynchronously.
   ///
@@ -113,7 +113,7 @@ class ServerConnection : public std::enable_shared_from_this<ServerConnection> {
   ///
   /// \param buffer The buffer.
   /// \return Status.
-  Status ReadBuffer(const std::vector<boost::asio::mutable_buffer> &buffer);
+  virtual Status ReadBuffer(const std::vector<boost::asio::mutable_buffer> &buffer);
 
   /// Read a buffer from this connection asynchronously.
   ///

src/ray/core_worker/store_provider/plasma_store_provider.cc

@@ -65,7 +65,13 @@ CoreWorkerPlasmaStoreProvider::CoreWorkerPlasmaStoreProvider(
     bool warmup,
     std::function<std::string()> get_current_call_site)
     : raylet_client_(raylet_client),
-      store_client_(std::make_shared<plasma::PlasmaClient>()),
+      // We can turn on exit_on_connection_failure on for the core worker plasma
+      // client to early exit core worker after the raylet's death because on the
+      // raylet side, we never proactively close the plasma store connection even
+      // during shutdown. So any error from the raylet side should be a sign of raylet
+      // death.
+      store_client_(
+          std::make_shared<plasma::PlasmaClient>(/*exit_on_connection_failure*/ true)),
       reference_counter_(reference_counter),
       check_signals_(std::move(check_signals)) {
   if (get_current_call_site != nullptr) {

src/ray/object_manager/plasma/client.cc

@@ -97,6 +97,7 @@ struct ObjectInUseEntry {
 class PlasmaClient::Impl : public std::enable_shared_from_this<PlasmaClient::Impl> {
  public:
   Impl();
+  explicit Impl(bool exit_on_connection_failure);
   ~Impl();
 
   // PlasmaClient method implementations
@@ -235,11 +236,17 @@ class PlasmaClient::Impl : public std::enable_shared_from_this<PlasmaClient::Imp
   std::unordered_set<ObjectID> deletion_cache_;
   /// A mutex which protects this class.
   std::recursive_mutex client_mutex_;
+  /// Whether the current process should exit when read or write to the connection fails.
+  /// It should only be turned on when the plasma client is in a core worker.
+  bool exit_on_connection_failure_;
 };
 
 PlasmaBuffer::~PlasmaBuffer() { RAY_UNUSED(client_->Release(object_id_)); }
 
-PlasmaClient::Impl::Impl() : store_capacity_(0) {}
+PlasmaClient::Impl::Impl() : store_capacity_(0), exit_on_connection_failure_(false) {}
+
+PlasmaClient::Impl::Impl(bool exit_on_connection_failure)
+    : store_capacity_(0), exit_on_connection_failure_(exit_on_connection_failure) {}
 
 PlasmaClient::Impl::~Impl() {}
 
@@ -868,7 +875,7 @@ Status PlasmaClient::Impl::Connect(const std::string &store_socket_name,
   /// The local stream socket that connects to store.
   ray::local_stream_socket socket(main_service_);
   RAY_RETURN_NOT_OK(ray::ConnectSocketRetry(socket, store_socket_name));
-  store_conn_.reset(new StoreConn(std::move(socket)));
+  store_conn_.reset(new StoreConn(std::move(socket), exit_on_connection_failure_));
   // Send a ConnectRequest to the store to get its memory capacity.
   RAY_RETURN_NOT_OK(SendConnectRequest(store_conn_));
   std::vector<uint8_t> buffer;
@@ -912,6 +919,9 @@ std::string PlasmaClient::Impl::DebugString() {
 
 PlasmaClient::PlasmaClient() : impl_(std::make_shared<PlasmaClient::Impl>()) {}
 
+PlasmaClient::PlasmaClient(bool exit_on_connection_failure)
+    : impl_(std::make_shared<PlasmaClient::Impl>(exit_on_connection_failure)) {}
+
 Status PlasmaClient::Connect(const std::string &store_socket_name,
                              const std::string &manager_socket_name,
                              int num_retries) {

src/ray/object_manager/plasma/client.h

@@ -242,6 +242,8 @@ class PlasmaClient : public PlasmaClientInterface {
  public:
   PlasmaClient();
 
+  explicit PlasmaClient(bool exit_on_connection_failure);
+
   Status Connect(const std::string &store_socket_name,
                  const std::string &manager_socket_name = "",
                  int num_retries = -1) override;

src/ray/object_manager/plasma/connection.cc

@@ -170,7 +170,11 @@ Status Client::SendFd(MEMFD_TYPE fd) {
 }
 
 StoreConn::StoreConn(ray::local_stream_socket &&socket)
-    : ray::ServerConnection(std::move(socket)) {}
+    : ray::ServerConnection(std::move(socket)), exit_on_connection_failure_(false) {}
+
+StoreConn::StoreConn(ray::local_stream_socket &&socket, bool exit_on_connection_failure)
+    : ray::ServerConnection(std::move(socket)),
+      exit_on_connection_failure_(exit_on_connection_failure) {}
 
 Status StoreConn::RecvFd(MEMFD_TYPE_NON_UNIQUE *fd) {
 #ifdef _WIN32
@@ -192,4 +196,28 @@ Status StoreConn::RecvFd(MEMFD_TYPE_NON_UNIQUE *fd) {
   return Status::OK();
 }
 
+ray::Status StoreConn::WriteBuffer(const std::vector<boost::asio::const_buffer> &buffer) {
+  auto status = ray::ServerConnection::WriteBuffer(buffer);
+  ExitIfErrorStatus(status);
+  return status;
+}
+
+ray::Status StoreConn::ReadBuffer(
+    const std::vector<boost::asio::mutable_buffer> &buffer) {
+  auto status = ray::ServerConnection::ReadBuffer(buffer);
+  ExitIfErrorStatus(status);
+  return status;
+}
+
+void StoreConn::ExitIfErrorStatus(const ray::Status &status) {
+  if (!status.ok() && exit_on_connection_failure_) {
+    RAY_LOG(WARNING) << "The connection to the plasma store is failed. Terminate the "
+                     << "process. Status: " << status;
+    ray::QuickExit();
+    RAY_LOG(FATAL)
+        << "Accessing unreachable code. This line should never be reached "
+        << "after quick process exit due to plasma store connection failure. Please "
+           "create a github issue at https://github.com/ray-project/ray.";
+  }
+}
 }  // namespace plasma

src/ray/object_manager/plasma/connection.h

@@ -164,10 +164,28 @@ class StoreConn : public ray::ServerConnection {
  public:
   explicit StoreConn(ray::local_stream_socket &&socket);
 
+  explicit StoreConn(ray::local_stream_socket &&socket, bool exit_on_connection_failure);
+
   /// Receive a file descriptor for the store.
   ///
   /// \return A file descriptor.
   ray::Status RecvFd(MEMFD_TYPE_NON_UNIQUE *fd);
+
+  ray::Status WriteBuffer(const std::vector<boost::asio::const_buffer> &buffer) override;
+
+  ray::Status ReadBuffer(const std::vector<boost::asio::mutable_buffer> &buffer) override;
+
+ private:
+  // Whether the current process should exit when WriteBuffer or ReadBuffer fails.
+  // Currently it is only turned on when the plasma client is in a core worker.
+  // TODO(myan): The better way is to handle the failure outside of the plasma client
+  // and inside the core worker's logic and propogate the correct exception to the user.
+  bool exit_on_connection_failure_ = false;
+
+  // Shutdown the current process if the passed in status is not OK and the client is
+  // configured to exit on failure.
+  // @param status: The status to check.
+  void ExitIfErrorStatus(const ray::Status &status);
 };
 
 std::ostream &operator<<(std::ostream &os, const std::shared_ptr<StoreConn> &store_conn);