[xray] Cleanup Raylet processes on exit (#1839)

* Add raylet monitor script to timeout Raylet heartbeats

* Unit test for removing a different client from the client table

* Set node manager heartbeat according to global config

* Doc and fixes

* Add regression test for client table disconnect, refactor client table

* Convert 'Terminate' methods to destructors

* Destroy the Raylet on a SIGTERM

* Clean up workers on a SIGTERM

src/ray/object_manager/object_directory.cc

@@ -101,6 +101,4 @@ ray::Status ObjectDirectory::Cancel(const ObjectID &object_id) {
   return ray::Status::OK();
 };
 
-ray::Status ObjectDirectory::Terminate() { return ray::Status::OK(); };
-
 }  // namespace ray

src/ray/object_manager/object_directory.h

@@ -82,11 +82,6 @@ class ObjectDirectoryInterface {
   /// \return Status of whether this method succeeded.
   virtual ray::Status ReportObjectRemoved(const ObjectID &object_id,
                                           const ClientID &client_id) = 0;
-
-  /// Terminate this object.
-  ///
-  /// \return Status of whether termination succeeded.
-  virtual ray::Status Terminate() = 0;
 };
 
 /// Ray ObjectDirectory declaration.
@@ -102,7 +97,6 @@ class ObjectDirectory : public ObjectDirectoryInterface {
                            const OnLocationsSuccess &success_callback,
                            const OnLocationsFailure &fail_callback) override;
   ray::Status Cancel(const ObjectID &object_id) override;
-  ray::Status Terminate() override;
   ray::Status ReportObjectAdded(const ObjectID &object_id, const ClientID &client_id,
                                 const ObjectInfoT &object_info) override;
   ray::Status ReportObjectRemoved(const ObjectID &object_id,

src/ray/object_manager/object_manager.cc

@@ -53,6 +53,13 @@ ObjectManager::ObjectManager(asio::io_service &main_service,
   StartIOService();
 }
 
+ObjectManager::~ObjectManager() {
+  object_manager_service_->stop();
+  for (int i = 0; i < config_.num_threads; ++i) {
+    io_threads_[i].join();
+  }
+}
+
 void ObjectManager::StartIOService() {
   for (int i = 0; i < config_.num_threads; ++i) {
     io_threads_.emplace_back(std::thread(&ObjectManager::IOServiceLoop, this));
@@ -61,13 +68,6 @@ void ObjectManager::StartIOService() {
 
 void ObjectManager::IOServiceLoop() { object_manager_service_->run(); }
 
-void ObjectManager::StopIOService() {
-  object_manager_service_->stop();
-  for (int i = 0; i < config_.num_threads; ++i) {
-    io_threads_[i].join();
-  }
-}
-
 void ObjectManager::NotifyDirectoryObjectAdd(const ObjectInfoT &object_info) {
   ObjectID object_id = ObjectID::from_binary(object_info.object_id);
   local_objects_[object_id] = object_info;
@@ -80,15 +80,6 @@ void ObjectManager::NotifyDirectoryObjectDeleted(const ObjectID &object_id) {
   ray::Status status = object_directory_->ReportObjectRemoved(object_id, client_id_);
 }
 
-ray::Status ObjectManager::Terminate() {
-  StopIOService();
-  ray::Status status_code = object_directory_->Terminate();
-  // TODO: evaluate store client termination status.
-  store_notification_.Terminate();
-  store_pool_.Terminate();
-  return status_code;
-}
-
 ray::Status ObjectManager::SubscribeObjAdded(
     std::function<void(const ObjectInfoT &)> callback) {
   store_notification_.SubscribeObjAdded(callback);

src/ray/object_manager/object_manager.h

@@ -71,6 +71,8 @@ class ObjectManager {
                          const ObjectManagerConfig &config,
                          std::unique_ptr<ObjectDirectoryInterface> od);
 
+  ~ObjectManager();
+
   /// Subscribe to notifications of objects added to local store.
   /// Upon subscribing, the callback will be invoked for all objects that
   ///
@@ -148,9 +150,6 @@ class ObjectManager {
   ray::Status Wait(const std::vector<ObjectID> &object_ids, uint64_t timeout_ms,
                    int num_ready_objects, const WaitCallback &callback);
 
-  /// \return Whether this object was successfully terminated.
-  ray::Status Terminate();
-
  private:
   ClientID client_id_;
   ObjectManagerConfig config_;

src/ray/object_manager/object_store_client_pool.cc

@@ -5,6 +5,12 @@ namespace ray {
 ObjectStoreClientPool::ObjectStoreClientPool(const std::string &store_socket_name)
     : store_socket_name_(store_socket_name) {}
 
+ObjectStoreClientPool::~ObjectStoreClientPool() {
+  for (const auto &client : clients) {
+    ARROW_CHECK_OK(client->Disconnect());
+  }
+}
+
 std::shared_ptr<plasma::PlasmaClient> ObjectStoreClientPool::GetObjectStore() {
   std::lock_guard<std::mutex> lock(pool_mutex);
   if (available_clients.empty()) {
@@ -21,14 +27,6 @@ void ObjectStoreClientPool::ReleaseObjectStore(
   available_clients.push_back(client);
 }
 
-void ObjectStoreClientPool::Terminate() {
-  for (const auto &client : clients) {
-    ARROW_CHECK_OK(client->Disconnect());
-  }
-  available_clients.clear();
-  clients.clear();
-}
-
 void ObjectStoreClientPool::Add() {
   clients.emplace_back(new plasma::PlasmaClient());
   ARROW_CHECK_OK(clients.back()->Connect(store_socket_name_.c_str(), "",

src/ray/object_manager/object_store_client_pool.h

@@ -31,6 +31,8 @@ class ObjectStoreClientPool {
   /// \param store_socket_name The object store socket name.
   ObjectStoreClientPool(const std::string &store_socket_name);
 
+  ~ObjectStoreClientPool();
+
   /// This object cannot be copied due to pool_mutex.
   RAY_DISALLOW_COPY_AND_ASSIGN(ObjectStoreClientPool);
 
@@ -47,9 +49,6 @@ class ObjectStoreClientPool {
   /// \param client
   void ReleaseObjectStore(std::shared_ptr<plasma::PlasmaClient> client);
 
-  /// Terminates this object.
-  void Terminate();
-
  private:
   /// Adds a client to the client pool and mark it as available.
   void Add();

src/ray/object_manager/object_store_notification_manager.cc

@@ -25,7 +25,7 @@ ObjectStoreNotificationManager::ObjectStoreNotificationManager(
   NotificationWait();
 }
 
-void ObjectStoreNotificationManager::Terminate() {
+ObjectStoreNotificationManager::~ObjectStoreNotificationManager() {
   ARROW_CHECK_OK(store_client_.Disconnect());
 }
 

src/ray/object_manager/object_store_notification_manager.h

@@ -32,6 +32,8 @@ class ObjectStoreNotificationManager {
   ObjectStoreNotificationManager(boost::asio::io_service &io_service,
                                  const std::string &store_socket_name);
 
+  ~ObjectStoreNotificationManager();
+
   /// Subscribe to notifications of objects added to local store.
   /// Upon subscribing, the callback will be invoked for all objects that
   /// already exist in the local store
@@ -44,9 +46,6 @@ class ObjectStoreNotificationManager {
   /// \param callback A callback expecting an ObjectID.
   void SubscribeObjDeleted(std::function<void(const ray::ObjectID &)> callback);
 
-  /// Terminate this object.
-  void Terminate();
-
  private:
   /// Async loop for handling object store notifications.
   void NotificationWait();

src/ray/object_manager/test/object_manager_stress_test.cc

@@ -43,7 +43,6 @@ class MockServer {
 
   ~MockServer() {
     RAY_CHECK_OK(gcs_client_->client_table().Disconnect());
-    RAY_CHECK_OK(object_manager_.Terminate());
   }
 
  private:

src/ray/object_manager/test/object_manager_test.cc

@@ -34,7 +34,6 @@ class MockServer {
 
   ~MockServer() {
     RAY_CHECK_OK(gcs_client_->client_table().Disconnect());
-    RAY_CHECK_OK(object_manager_.Terminate());
   }
 
  private:

src/ray/raylet/format/node_manager.fbs

@@ -91,12 +91,8 @@ table RegisterClientRequest {
   is_worker: bool;
   // The ID of the worker or driver.
   client_id: string;
-  // The ID of the actor. This is NIL_ACTOR_ID if the worker is not an actor.
-  actor_id: string;
   // The process ID of this worker.
   worker_pid: long;
-  // The number of GPUs required by this actor.
-  num_gpus: long;
 }
 
 table RegisterClientReply {

src/ray/raylet/main.cc

@@ -50,6 +50,15 @@ int main(int argc, char *argv[]) {
                              raylet_socket_name, node_ip_address, redis_address,
                              redis_port, node_manager_config, object_manager_config,
                              gcs_client);
+
+  // Destroy the Raylet on a SIGTERM. The pointer to main_service is
+  // guaranteed to be valid since this function will run the event loop
+  // instead of returning immediately.
+  auto handler = [&main_service](const boost::system::error_code &error,
+                                 int signal_number) { main_service.stop(); };
+  boost::asio::signal_set signals(main_service, SIGTERM);
+  signals.async_wait(handler);
+
   main_service.run();
 }
 #endif

src/ray/raylet/raylet.cc

@@ -18,18 +18,18 @@ Raylet::Raylet(boost::asio::io_service &main_service,
                const NodeManagerConfig &node_manager_config,
                const ObjectManagerConfig &object_manager_config,
                std::shared_ptr<gcs::AsyncGcsClient> gcs_client)
-    : acceptor_(main_service, boost::asio::local::stream_protocol::endpoint(socket_name)),
+    : gcs_client_(gcs_client),
+      object_manager_(main_service, std::move(object_manager_service),
+                      object_manager_config, gcs_client),
+      node_manager_(main_service, node_manager_config, object_manager_, gcs_client_),
+      acceptor_(main_service, boost::asio::local::stream_protocol::endpoint(socket_name)),
       socket_(main_service),
       object_manager_acceptor_(
           main_service, boost::asio::ip::tcp::endpoint(boost::asio::ip::tcp::v4(), 0)),
       object_manager_socket_(main_service),
       node_manager_acceptor_(
           main_service, boost::asio::ip::tcp::endpoint(boost::asio::ip::tcp::v4(), 0)),
-      node_manager_socket_(main_service),
-      gcs_client_(gcs_client),
-      object_manager_(main_service, std::move(object_manager_service),
-                      object_manager_config, gcs_client),
-      node_manager_(main_service, node_manager_config, object_manager_, gcs_client_) {
+      node_manager_socket_(main_service) {
   // Start listening for clients.
   DoAccept();
   DoAcceptObjectManager();
@@ -43,7 +43,6 @@ Raylet::Raylet(boost::asio::io_service &main_service,
 
 Raylet::~Raylet() {
   RAY_CHECK_OK(gcs_client_->client_table().Disconnect());
-  RAY_CHECK_OK(object_manager_.Terminate());
 }
 
 ray::Status Raylet::RegisterPeriodicTimer(boost::asio::io_service &io_service) {

src/ray/raylet/raylet.h

@@ -65,6 +65,13 @@ class Raylet {
 
   friend class TestObjectManagerIntegration;
 
+  /// A client connection to the GCS.
+  std::shared_ptr<gcs::AsyncGcsClient> gcs_client_;
+  /// Manages client requests for object transfers and availability.
+  ObjectManager object_manager_;
+  /// Manages client requests for task submission and execution.
+  NodeManager node_manager_;
+
   /// An acceptor for new clients.
   boost::asio::local::stream_protocol::acceptor acceptor_;
   /// The socket to listen on for new clients.
@@ -77,13 +84,6 @@ class Raylet {
   boost::asio::ip::tcp::acceptor node_manager_acceptor_;
   /// The socket to listen on for new tcp clients.
   boost::asio::ip::tcp::socket node_manager_socket_;
-
-  /// A client connection to the GCS.
-  std::shared_ptr<gcs::AsyncGcsClient> gcs_client_;
-  /// Manages client requests for object transfers and availability.
-  ObjectManager object_manager_;
-  /// Manages client requests for task submission and execution.
-  NodeManager node_manager_;
 };
 
 }  // namespace raylet

src/ray/raylet/worker_pool.cc

@@ -1,5 +1,7 @@
 #include "ray/raylet/worker_pool.h"
 
+#include <sys/wait.h>
+
 #include "ray/status.h"
 #include "ray/util/logging.h"
 
@@ -19,9 +21,13 @@ WorkerPool::WorkerPool(int num_workers, const std::vector<std::string> &worker_c
 }
 
 WorkerPool::~WorkerPool() {
-  // TODO(swang): Kill registered workers.
-  pool_.clear();
-  registered_workers_.clear();
+  // Kill all registered workers. NOTE(swang): This assumes that the registered
+  // workers were started by the pool.
+  for (const auto &worker : registered_workers_) {
+    RAY_CHECK(worker->Pid() > 0);
+    kill(worker->Pid(), SIGKILL);
+    waitpid(worker->Pid(), NULL, 0);
+  }
 }
 
 void WorkerPool::StartWorker() {