[xray] Make sure raylet does not crash if remote raylet dies (#2619)

* Log a warning on remote object manager failures

* Mark a task that was failed to be forwarded as pending

* Raylet component failure test and make it harder

* Turn on component failure test for xray

* Remove return status from ReleaseSender

* lint

.travis.yml

@@ -145,7 +145,7 @@ matrix:
         - python -m pytest test/failure_test.py
         - python -m pytest test/microbenchmarks.py
         - python -m pytest test/stress_tests.py
-        # - pytest test/component_failures_test.py
+        - pytest test/component_failures_test.py
         - python test/multi_node_test.py
         - python -m pytest test/recursion_test.py
         - pytest test/monitor_test.py

src/ray/common/client_connection.cc

@@ -23,8 +23,9 @@ ServerConnection<T>::ServerConnection(boost::asio::basic_stream_socket<T> &&sock
     : socket_(std::move(socket)) {}
 
 template <class T>
-void ServerConnection<T>::WriteBuffer(
-    const std::vector<boost::asio::const_buffer> &buffer, boost::system::error_code &ec) {
+Status ServerConnection<T>::WriteBuffer(
+    const std::vector<boost::asio::const_buffer> &buffer) {
+  boost::system::error_code error;
   // Loop until all bytes are written while handling interrupts.
   // When profiling with pprof, unhandled interrupts were being sent by the profiler to
   // the raylet process, which was causing synchronous reads and writes to fail.
@@ -33,16 +34,17 @@ void ServerConnection<T>::WriteBuffer(
     uint64_t position = 0;
     while (bytes_remaining != 0) {
       size_t bytes_written =
-          socket_.write_some(boost::asio::buffer(b + position, bytes_remaining), ec);
+          socket_.write_some(boost::asio::buffer(b + position, bytes_remaining), error);
       position += bytes_written;
       bytes_remaining -= bytes_written;
-      if (ec.value() == EINTR) {
+      if (error.value() == EINTR) {
         continue;
-      } else if (ec.value() != boost::system::errc::errc_t::success) {
-        return;
+      } else if (error.value() != boost::system::errc::errc_t::success) {
+        return boost_to_ray_status(error);
       }
     }
   }
+  return ray::Status::OK();
 }
 
 template <class T>
@@ -78,9 +80,7 @@ ray::Status ServerConnection<T>::WriteMessage(int64_t type, int64_t length,
   message_buffers.push_back(boost::asio::buffer(message, length));
   // Write the message and then wait for more messages.
   // TODO(swang): Does this need to be an async write?
-  boost::system::error_code error;
-  WriteBuffer(message_buffers, error);
-  return boost_to_ray_status(error);
+  return WriteBuffer(message_buffers);
 }
 
 template <class T>

src/ray/common/client_connection.h

@@ -43,8 +43,7 @@ class ServerConnection {
   ///
   /// \param buffer The buffer.
   /// \param ec The error code object in which to store error codes.
-  void WriteBuffer(const std::vector<boost::asio::const_buffer> &buffer,
-                   boost::system::error_code &ec);
+  Status WriteBuffer(const std::vector<boost::asio::const_buffer> &buffer);
 
   /// Read a buffer from this connection.
   ///

src/ray/object_manager/connection_pool.cc

@@ -38,8 +38,8 @@ void ConnectionPool::RegisterSender(ConnectionType type, const ClientID &client_
   // Don't add to available connections. It will become available once it is released.
 }
 
-ray::Status ConnectionPool::GetSender(ConnectionType type, const ClientID &client_id,
-                                      std::shared_ptr<SenderConnection> *conn) {
+void ConnectionPool::GetSender(ConnectionType type, const ClientID &client_id,
+                               std::shared_ptr<SenderConnection> *conn) {
   std::unique_lock<std::mutex> guard(connection_mutex);
   SenderMapType &avail_conn_map = (type == ConnectionType::MESSAGE)
                                       ? available_message_send_connections_
@@ -49,17 +49,15 @@ ray::Status ConnectionPool::GetSender(ConnectionType type, const ClientID &clien
   } else {
     *conn = nullptr;
   }
-  return ray::Status::OK();
 }
 
-ray::Status ConnectionPool::ReleaseSender(ConnectionType type,
-                                          std::shared_ptr<SenderConnection> &conn) {
+void ConnectionPool::ReleaseSender(ConnectionType type,
+                                   std::shared_ptr<SenderConnection> &conn) {
   std::unique_lock<std::mutex> guard(connection_mutex);
   SenderMapType &conn_map = (type == ConnectionType::MESSAGE)
                                 ? available_message_send_connections_
                                 : available_transfer_send_connections_;
   Return(conn_map, conn->GetClientID(), conn);
-  return ray::Status::OK();
 }
 
 void ConnectionPool::Add(ReceiverMapType &conn_map, const ClientID &client_id,

src/ray/object_manager/connection_pool.h

@@ -65,16 +65,16 @@ class ConnectionPool {
   /// \param[in] type The type of connection.
   /// \param[in] client_id The ClientID of the remote object manager.
   /// \param[out] conn An empty pointer to a shared pointer.
-  /// \return Status of invoking this method.
-  ray::Status GetSender(ConnectionType type, const ClientID &client_id,
-                        std::shared_ptr<SenderConnection> *conn);
+  /// \return Void.
+  void GetSender(ConnectionType type, const ClientID &client_id,
+                 std::shared_ptr<SenderConnection> *conn);
 
   /// Releases a sender connection, allowing it to be used by another operation.
   ///
   /// \param type The type of connection.
   /// \param conn The actual connection.
-  /// \return Status of invoking this method.
-  ray::Status ReleaseSender(ConnectionType type, std::shared_ptr<SenderConnection> &conn);
+  /// \return Void.
+  void ReleaseSender(ConnectionType type, std::shared_ptr<SenderConnection> &conn);
 
   // TODO(hme): Implement with error handling.
   /// Remove a sender connection. This is invoked if the connection is no longer

src/ray/object_manager/object_directory.cc

@@ -107,7 +107,7 @@ ray::Status ObjectDirectory::GetInformation(const ClientID &client_id,
   const ClientTableDataT &data = gcs_client_->client_table().GetClient(client_id);
   ClientID result_client_id = ClientID::from_binary(data.client_id);
   if (result_client_id == ClientID::nil() || !data.is_insertion) {
-    fail_callback(ray::Status::RedisError("ClientID not found."));
+    fail_callback();
   } else {
     const auto &info = RemoteConnectionInfo(client_id, data.node_manager_address,
                                             (uint16_t)data.object_manager_port);

src/ray/object_manager/object_directory.h

@@ -31,7 +31,7 @@ class ObjectDirectoryInterface {
 
   /// Callbacks for GetInformation.
   using InfoSuccessCallback = std::function<void(const ray::RemoteConnectionInfo &info)>;
-  using InfoFailureCallback = std::function<void(ray::Status status)>;
+  using InfoFailureCallback = std::function<void()>;
 
   virtual void RegisterBackend() = 0;
 

src/ray/object_manager/object_manager.cc

@@ -5,6 +5,15 @@ namespace asio = boost::asio;
 
 namespace object_manager_protocol = ray::object_manager::protocol;
 
+namespace {
+
+void CheckIOError(ray::Status &status, const std::string &operation) {
+  RAY_CHECK(status.IsIOError());
+  RAY_LOG(ERROR) << "Failed to contact remote object manager during " << operation;
+}
+
+}  // namespace
+
 namespace ray {
 
 ObjectManager::ObjectManager(asio::io_service &main_service,
@@ -152,37 +161,31 @@ void ObjectManager::GetLocationsSuccess(const std::vector<ray::ClientID> &client
     // Only pull objects that aren't local.
     RAY_CHECK(!client_ids.empty());
     ClientID client_id = client_ids.front();
-    ray::Status status_code = Pull(object_id, client_id);
-    RAY_CHECK_OK(status_code);
+    Pull(object_id, client_id);
   }
 }
 
-ray::Status ObjectManager::Pull(const ObjectID &object_id, const ClientID &client_id) {
+void ObjectManager::Pull(const ObjectID &object_id, const ClientID &client_id) {
   // Check if object is already local.
   if (local_objects_.count(object_id) != 0) {
     RAY_LOG(ERROR) << object_id << " attempted to pull an object that's already local.";
-    return ray::Status::OK();
+    return;
   }
   // Check if we're pulling from self.
   if (client_id == client_id_) {
     RAY_LOG(ERROR) << client_id_ << " attempted to pull an object from itself.";
-    return ray::Status::Invalid("A node cannot pull an object from itself.");
+    return;
   }
-  return PullEstablishConnection(object_id, client_id);
+  PullEstablishConnection(object_id, client_id);
 };
 
-ray::Status ObjectManager::PullEstablishConnection(const ObjectID &object_id,
-                                                   const ClientID &client_id) {
+void ObjectManager::PullEstablishConnection(const ObjectID &object_id,
+                                            const ClientID &client_id) {
   // Acquire a message connection and send pull request.
   ray::Status status;
   std::shared_ptr<SenderConnection> conn;
   // TODO(hme): There is no cap on the number of pull request connections.
-  status = connection_pool_.GetSender(ConnectionPool::ConnectionType::MESSAGE, client_id,
-                                      &conn);
-  // Currently, acquiring a connection should not fail.
-  // No status from GetSender is returned which can be
-  // handled without failing.
-  RAY_CHECK_OK(status);
+  connection_pool_.GetSender(ConnectionPool::ConnectionType::MESSAGE, client_id, &conn);
 
   if (conn == nullptr) {
     status = object_directory_->GetInformation(
@@ -190,19 +193,25 @@ ray::Status ObjectManager::PullEstablishConnection(const ObjectID &object_id,
         [this, object_id, client_id](const RemoteConnectionInfo &connection_info) {
           std::shared_ptr<SenderConnection> async_conn = CreateSenderConnection(
               ConnectionPool::ConnectionType::MESSAGE, connection_info);
+          if (async_conn == nullptr) {
+            return;
+          }
           connection_pool_.RegisterSender(ConnectionPool::ConnectionType::MESSAGE,
                                           client_id, async_conn);
           Status pull_send_status = PullSendRequest(object_id, async_conn);
-          RAY_CHECK_OK(pull_send_status);
+          if (!pull_send_status.ok()) {
+            CheckIOError(pull_send_status, "Pull");
+          }
         },
-        [](const Status &status) {
+        []() {
           RAY_LOG(ERROR) << "Failed to establish connection with remote object manager.";
-          RAY_CHECK_OK(status);
         });
   } else {
     status = PullSendRequest(object_id, conn);
+    if (!status.ok()) {
+      CheckIOError(status, "Pull");
+    }
   }
-  return status;
 }
 
 ray::Status ObjectManager::PullSendRequest(const ObjectID &object_id,
@@ -211,12 +220,13 @@ ray::Status ObjectManager::PullSendRequest(const ObjectID &object_id,
   auto message = object_manager_protocol::CreatePullRequestMessage(
       fbb, fbb.CreateString(client_id_.binary()), fbb.CreateString(object_id.binary()));
   fbb.Finish(message);
-  RAY_CHECK_OK(conn->WriteMessage(
+  Status status = conn->WriteMessage(
       static_cast<int64_t>(object_manager_protocol::MessageType::PullRequest),
-      fbb.GetSize(), fbb.GetBufferPointer()));
-  RAY_CHECK_OK(
-      connection_pool_.ReleaseSender(ConnectionPool::ConnectionType::MESSAGE, conn));
-  return ray::Status::OK();
+      fbb.GetSize(), fbb.GetBufferPointer());
+  if (status.ok()) {
+    connection_pool_.ReleaseSender(ConnectionPool::ConnectionType::MESSAGE, conn);
+  }
+  return status;
 }
 
 void ObjectManager::HandlePushTaskTimeout(const ObjectID &object_id,
@@ -283,8 +293,10 @@ void ObjectManager::Push(const ObjectID &object_id, const ClientID &client_id) {
           });
         }
       },
-      [](const Status &status) {
+      []() {
         // Push is best effort, so do nothing here.
+        RAY_LOG(ERROR)
+            << "Failed to establish connection for Push with remote object manager.";
       }));
 }
 
@@ -296,16 +308,20 @@ void ObjectManager::ExecuteSendObject(const ClientID &client_id,
                  << chunk_index;
   ray::Status status;
   std::shared_ptr<SenderConnection> conn;
-  status = connection_pool_.GetSender(ConnectionPool::ConnectionType::TRANSFER, client_id,
-                                      &conn);
+  connection_pool_.GetSender(ConnectionPool::ConnectionType::TRANSFER, client_id, &conn);
   if (conn == nullptr) {
     conn =
         CreateSenderConnection(ConnectionPool::ConnectionType::TRANSFER, connection_info);
     connection_pool_.RegisterSender(ConnectionPool::ConnectionType::TRANSFER, client_id,
                                     conn);
+    if (conn == nullptr) {
+      return;
+    }
   }
   status = SendObjectHeaders(object_id, data_size, metadata_size, chunk_index, conn);
-  RAY_CHECK_OK(status);
+  if (!status.ok()) {
+    CheckIOError(status, "Push");
+  }
 }
 
 ray::Status ObjectManager::SendObjectHeaders(const ObjectID &object_id,
@@ -329,30 +345,28 @@ ray::Status ObjectManager::SendObjectHeaders(const ObjectID &object_id,
   ray::Status status = conn->WriteMessage(
       static_cast<int64_t>(object_manager_protocol::MessageType::PushRequest),
       fbb.GetSize(), fbb.GetBufferPointer());
-  RAY_CHECK_OK(status);
+  if (!status.ok()) {
+    return status;
+  }
   return SendObjectData(object_id, chunk_info, conn);
 }
 
 ray::Status ObjectManager::SendObjectData(const ObjectID &object_id,
                                           const ObjectBufferPool::ChunkInfo &chunk_info,
                                           std::shared_ptr<SenderConnection> &conn) {
-  boost::system::error_code ec;
+  boost::system::error_code error;
   std::vector<asio::const_buffer> buffer;
   buffer.push_back(asio::buffer(chunk_info.data, chunk_info.buffer_length));
-  conn->WriteBuffer(buffer, ec);
-
-  ray::Status status = boost_to_ray_status(ec);
-  if (ec.value() != boost::system::errc::success) {
-    // Push failed. Deal with partial objects on the receiving end.
-    // TODO(hme): Try to invoke disconnect on sender connection, then remove it.
-  }
+  Status status = conn->WriteBuffer(buffer);
 
   // Do this regardless of whether it failed or succeeded.
   buffer_pool_.ReleaseGetChunk(object_id, chunk_info.chunk_index);
-  RAY_CHECK_OK(
-      connection_pool_.ReleaseSender(ConnectionPool::ConnectionType::TRANSFER, conn));
-  RAY_LOG(DEBUG) << "SendCompleted " << client_id_ << " " << object_id << " "
-                 << config_.max_sends;
+
+  if (status.ok()) {
+    connection_pool_.ReleaseSender(ConnectionPool::ConnectionType::TRANSFER, conn);
+    RAY_LOG(DEBUG) << "SendCompleted " << client_id_ << " " << object_id << " "
+                   << config_.max_sends;
+  }
   return status;
 }
 
@@ -528,6 +542,10 @@ std::shared_ptr<SenderConnection> ObjectManager::CreateSenderConnection(
     ConnectionPool::ConnectionType type, RemoteConnectionInfo info) {
   std::shared_ptr<SenderConnection> conn =
       SenderConnection::Create(*main_service_, info.client_id, info.ip, info.port);
+  if (conn == nullptr) {
+    RAY_LOG(ERROR) << "Failed to connect to remote object manager.";
+    return conn;
+  }
   // Prepare client connection info buffer
   flatbuffers::FlatBufferBuilder fbb;
   bool is_transfer = (type == ConnectionPool::ConnectionType::TRANSFER);

src/ray/object_manager/object_manager.h

@@ -116,8 +116,8 @@ class ObjectManager : public ObjectManagerInterface {
   ///
   /// \param object_id The object's object id.
   /// \param client_id The remote node's client id.
-  /// \return Status of whether the pull request successfully initiated.
-  ray::Status Pull(const ObjectID &object_id, const ClientID &client_id);
+  /// \return Void.
+  void Pull(const ObjectID &object_id, const ClientID &client_id);
 
   /// Add a connection to a remote object manager.
   /// This is invoked by an external server.
@@ -273,8 +273,7 @@ class ObjectManager : public ObjectManagerInterface {
   /// Part of an asynchronous sequence of Pull methods.
   /// Uses an existing connection or creates a connection to ClientID.
   /// Executes on main_service_ thread.
-  ray::Status PullEstablishConnection(const ObjectID &object_id,
-                                      const ClientID &client_id);
+  void PullEstablishConnection(const ObjectID &object_id, const ClientID &client_id);
 
   /// Private callback implementation for success on get location. Called from
   /// ObjectDirectory.

src/ray/object_manager/object_manager_client_connection.cc

@@ -8,10 +8,14 @@ std::shared_ptr<SenderConnection> SenderConnection::Create(
     boost::asio::io_service &io_service, const ClientID &client_id, const std::string &ip,
     uint16_t port) {
   boost::asio::ip::tcp::socket socket(io_service);
-  RAY_CHECK_OK(TcpConnect(socket, ip, port));
-  std::shared_ptr<TcpServerConnection> conn =
-      std::make_shared<TcpServerConnection>(std::move(socket));
-  return std::make_shared<SenderConnection>(std::move(conn), client_id);
+  Status status = TcpConnect(socket, ip, port);
+  if (status.ok()) {
+    std::shared_ptr<TcpServerConnection> conn =
+        std::make_shared<TcpServerConnection>(std::move(socket));
+    return std::make_shared<SenderConnection>(std::move(conn), client_id);
+  } else {
+    return nullptr;
+  }
 };
 
 SenderConnection::SenderConnection(std::shared_ptr<TcpServerConnection> conn,

src/ray/object_manager/object_manager_client_connection.h

@@ -24,7 +24,8 @@ class SenderConnection : public boost::enable_shared_from_this<SenderConnection>
   /// \param client_id The ClientID of the remote node.
   /// \param ip The ip address of the remote node server.
   /// \param port The port of the remote node server.
-  /// \return A connection to the remote object manager.
+  /// \return A connection to the remote object manager. This is null if the
+  /// connection was unsuccessful.
   static std::shared_ptr<SenderConnection> Create(boost::asio::io_service &io_service,
                                                   const ClientID &client_id,
                                                   const std::string &ip, uint16_t port);
@@ -47,9 +48,8 @@ class SenderConnection : public boost::enable_shared_from_this<SenderConnection>
   ///
   /// \param buffer The buffer.
   /// \param ec The error code object in which to store error codes.
-  void WriteBuffer(const std::vector<boost::asio::const_buffer> &buffer,
-                   boost::system::error_code &ec) {
-    return conn_->WriteBuffer(buffer, ec);
+  Status WriteBuffer(const std::vector<boost::asio::const_buffer> &buffer) {
+    return conn_->WriteBuffer(buffer);
   }
 
   /// Read a buffer from this connection.

src/ray/raylet/node_manager.cc

@@ -1274,6 +1274,11 @@ void NodeManager::ForwardTaskOrResubmit(const Task &task,
   if (!ForwardTask(task, node_manager_id).ok()) {
     RAY_LOG(INFO) << "Failed to forward task " << task_id << " to node manager "
                   << node_manager_id;
+    // Mark the failed task as pending to let other raylets know that we still
+    // have the task. Once the task is successfully retried, it will be
+    // canceled. TaskDependencyManager::TaskPending() is assumed to be
+    // idempotent.
+    task_dependency_manager_.TaskPending(task);
 
     // Create a timer to resubmit the task in a little bit. TODO(rkn): Really
     // this should be a unique_ptr instead of a shared_ptr. However, it's a

test/component_failures_test.py

@@ -16,6 +16,9 @@ class ComponentFailureTest(unittest.TestCase):
 
     # This test checks that when a worker dies in the middle of a get, the
     # plasma store and manager will not die.
+    @unittest.skipIf(
+        os.environ.get('RAY_USE_XRAY', False),
+        "Workers are all started by Raylet, so cannot be killed from Python.")
     @unittest.skipIf(
         os.environ.get('RAY_USE_NEW_GCS', False),
         "Not working with new GCS API.")
@@ -55,6 +58,9 @@ class ComponentFailureTest(unittest.TestCase):
 
     # This test checks that when a worker dies in the middle of a wait, the
     # plasma store and manager will not die.
+    @unittest.skipIf(
+        os.environ.get('RAY_USE_XRAY', False),
+        "Workers are all started by Raylet, so cannot be killed from Python.")
     @unittest.skipIf(
         os.environ.get('RAY_USE_NEW_GCS', False),
         "Not working with new GCS API.")
@@ -133,11 +139,6 @@ class ComponentFailureTest(unittest.TestCase):
     def _testComponentFailed(self, component_type):
         """Kill a component on all worker nodes and check workload succeeds."""
 
-        @ray.remote
-        def f(x, j):
-            time.sleep(0.2)
-            return x
-
         # Start with 4 workers and 4 cores.
         num_local_schedulers = 4
         num_workers_per_scheduler = 8
@@ -148,14 +149,24 @@ class ComponentFailureTest(unittest.TestCase):
             num_cpus=[num_workers_per_scheduler] * num_local_schedulers,
             redirect_output=True)
 
-        # Submit more tasks than there are workers so that all workers and
-        # cores are utilized.
-        object_ids = [
-            f.remote(i, 0)
-            for i in range(num_workers_per_scheduler * num_local_schedulers)
-        ]
-        object_ids += [f.remote(object_id, 1) for object_id in object_ids]
-        object_ids += [f.remote(object_id, 2) for object_id in object_ids]
+        # Submit many tasks with many dependencies.
+        @ray.remote
+        def f(x):
+            return x
+
+        x = 1
+        for _ in range(1000):
+            x = f.remote(x)
+        ray.get(x)
+
+        @ray.remote
+        def g(*xs):
+            return 1
+
+        xs = [g.remote(1)]
+        for _ in range(100):
+            xs.append(g.remote(*xs))
+            xs.append(g.remote(1))
 
         # Kill the component on all nodes except the head node as the tasks
         # execute.
@@ -172,10 +183,7 @@ class ComponentFailureTest(unittest.TestCase):
 
         # Make sure that we can still get the objects after the executing tasks
         # died.
-        results = ray.get(object_ids)
-        expected_results = 4 * list(
-            range(num_workers_per_scheduler * num_local_schedulers))
-        assert results == expected_results
+        ray.get(xs)
 
     def check_components_alive(self, component_type, check_component_alive):
         """Check that a given component type is alive on all worker nodes.
@@ -192,6 +200,20 @@ class ComponentFailureTest(unittest.TestCase):
                       str(component.pid) + "to terminate")
                 assert not component.poll() is None
 
+    @unittest.skipIf(not os.environ.get('RAY_USE_XRAY', False),
+                     "Only tests Raylet failure.")
+    def testRayletFailed(self):
+        # Kill all local schedulers on worker nodes.
+        self._testComponentFailed(ray.services.PROCESS_TYPE_RAYLET)
+
+        # The plasma stores and plasma managers should still be alive on the
+        # worker nodes.
+        self.check_components_alive(ray.services.PROCESS_TYPE_PLASMA_STORE,
+                                    True)
+
+    @unittest.skipIf(
+        os.environ.get('RAY_USE_XRAY', False),
+        "Raylet codepath does not have this component")
     @unittest.skipIf(
         os.environ.get('RAY_USE_NEW_GCS', False), "Hanging with new GCS API.")
     def testLocalSchedulerFailed(self):
@@ -207,6 +229,9 @@ class ComponentFailureTest(unittest.TestCase):
         self.check_components_alive(ray.services.PROCESS_TYPE_LOCAL_SCHEDULER,
                                     False)
 
+    @unittest.skipIf(
+        os.environ.get('RAY_USE_XRAY', False),
+        "Raylet codepath does not have this component")
     @unittest.skipIf(
         os.environ.get('RAY_USE_NEW_GCS', False), "Hanging with new GCS API.")
     def testPlasmaManagerFailed(self):
@@ -235,6 +260,7 @@ class ComponentFailureTest(unittest.TestCase):
                                     False)
         self.check_components_alive(ray.services.PROCESS_TYPE_LOCAL_SCHEDULER,
                                     False)
+        self.check_components_alive(ray.services.PROCESS_TYPE_RAYLET, False)
 
     @unittest.skipIf(
         os.environ.get('RAY_USE_NEW_GCS', False),
@@ -242,12 +268,11 @@ class ComponentFailureTest(unittest.TestCase):
     def testDriverLivesSequential(self):
         ray.worker.init(redirect_output=True)
         all_processes = ray.services.all_processes
-        processes = [
-            all_processes[ray.services.PROCESS_TYPE_PLASMA_STORE][0],
-            all_processes[ray.services.PROCESS_TYPE_PLASMA_MANAGER][0],
-            all_processes[ray.services.PROCESS_TYPE_LOCAL_SCHEDULER][0],
-            all_processes[ray.services.PROCESS_TYPE_GLOBAL_SCHEDULER][0]
-        ]
+        processes = (all_processes[ray.services.PROCESS_TYPE_PLASMA_STORE] +
+                     all_processes[ray.services.PROCESS_TYPE_PLASMA_MANAGER] +
+                     all_processes[ray.services.PROCESS_TYPE_LOCAL_SCHEDULER] +
+                     all_processes[ray.services.PROCESS_TYPE_GLOBAL_SCHEDULER]
+                     + all_processes[ray.services.PROCESS_TYPE_RAYLET])
 
         # Kill all the components sequentially.
         for process in processes:
@@ -264,12 +289,11 @@ class ComponentFailureTest(unittest.TestCase):
     def testDriverLivesParallel(self):
         ray.worker.init(redirect_output=True)
         all_processes = ray.services.all_processes
-        processes = [
-            all_processes[ray.services.PROCESS_TYPE_PLASMA_STORE][0],
-            all_processes[ray.services.PROCESS_TYPE_PLASMA_MANAGER][0],
-            all_processes[ray.services.PROCESS_TYPE_LOCAL_SCHEDULER][0],
-            all_processes[ray.services.PROCESS_TYPE_GLOBAL_SCHEDULER][0]
-        ]
+        processes = (all_processes[ray.services.PROCESS_TYPE_PLASMA_STORE] +
+                     all_processes[ray.services.PROCESS_TYPE_PLASMA_MANAGER] +
+                     all_processes[ray.services.PROCESS_TYPE_LOCAL_SCHEDULER] +
+                     all_processes[ray.services.PROCESS_TYPE_GLOBAL_SCHEDULER]
+                     + all_processes[ray.services.PROCESS_TYPE_RAYLET])
 
         # Kill all the components in parallel.
         for process in processes: