XRay Task Forwarding Milestone (ray-project#1785)

Summary:
Able to run 1000 tasks with object dependencies on a set of distributed Raylets.

Raylet Changes:

Finalized ClientConnection class.
Task forwarding.
NM-to-NM heartbeats.
NM resource accounting for tasks.
Simple scheduling policy with task forwarding.
Creating and maintaining NM 2 NM long-lived connections and reusing them for task forwarding.
LineageCache Changes:

LineageCache without cleanup of tasks committed by remote nodes.
Lineage cache writeback and cleanup implementation.
ObjectManager Changes:

Object manager event loop/ClientConnection refactor.
Multithreaded object manager (disabled in this PR).
Testing Changes:

Integration tests for task submission on multiple Raylets.
Stress tests for object manager (with GCS and object store integration).


Co-authored-by: Stephanie Wang <swang@cs.berkeley.edu>
Co-authored-by: Alexey Tumanov <atumanov@gmail.com>

.travis.yml

@@ -102,6 +102,13 @@ install:
 
   - cd python/ray/core
   - bash ../../../src/ray/test/run_gcs_tests.sh
+  # Raylet tests.
+  - bash ../../../src/ray/test/run_object_manager_tests.sh
+  - bash ../../../src/ray/test/run_task_test.sh
+  - ./src/ray/raylet/task_test
+  - ./src/ray/raylet/worker_pool_test
+  - ./src/ray/raylet/lineage_cache_test
+
   - bash ../../../src/common/test/run_tests.sh
   - bash ../../../src/plasma/test/run_tests.sh
   - bash ../../../src/local_scheduler/test/run_tests.sh

src/common/redis_module/ray_redis_module.cc

@@ -50,7 +50,14 @@
   }
 
 static const char *table_prefixes[] = {
-    NULL, "TASK:", "TASK:", "CLIENT:", "OBJECT:", "FUNCTION:",
+    NULL,
+    "TASK:",
+    "TASK:",
+    "CLIENT:",
+    "OBJECT:",
+    "FUNCTION:",
+    "TASK_RECONSTRUCTION:",
+    "HEARTBEAT:",
 };
 
 /// Parse a Redis string into a TablePubsub channel.
@@ -811,8 +818,9 @@ int TableRequestNotifications_RedisCommand(RedisModuleCtx *ctx,
     // notifications.
     flatbuffers::FlatBufferBuilder fbb;
     TableEntryToFlatbuf(table_key, id, fbb);
-    RedisModule_Call(ctx, "PUBLISH", "sb", client_channel, reinterpret_cast<const char *>(fbb.GetBufferPointer()),
-        fbb.GetSize());
+    RedisModule_Call(ctx, "PUBLISH", "sb", client_channel,
+                     reinterpret_cast<const char *>(fbb.GetBufferPointer()),
+                     fbb.GetSize());
   }
   RedisModule_CloseKey(table_key);
 

src/global_scheduler/global_scheduler.cc

@@ -141,13 +141,8 @@ GlobalSchedulerState *GlobalSchedulerState_init(event_loop *loop,
                          std::vector<std::string>());
   db_attach(state->db, loop, false);
 
-  ClientTableDataT client_info;
-  client_info.client_id = get_db_client_id(state->db).binary();
-  client_info.node_manager_address = std::string(node_ip_address);
-  client_info.local_scheduler_port = 0;
-  client_info.object_manager_port = 0;
   RAY_CHECK_OK(state->gcs_client.Connect(std::string(redis_primary_addr),
-                                         redis_primary_port, client_info));
+                                         redis_primary_port));
   RAY_CHECK_OK(state->gcs_client.context()->AttachToEventLoop(loop));
   state->policy_state = GlobalSchedulerPolicyState_init();
   return state;

src/local_scheduler/local_scheduler.cc

@@ -355,13 +355,8 @@ LocalSchedulerState *LocalSchedulerState_init(
                            "local_scheduler", node_ip_address, db_connect_args);
     db_attach(state->db, loop, false);
 
-    ClientTableDataT client_info;
-    client_info.client_id = get_db_client_id(state->db).binary();
-    client_info.node_manager_address = std::string(node_ip_address);
-    client_info.local_scheduler_port = 0;
-    client_info.object_manager_port = 0;
     RAY_CHECK_OK(state->gcs_client.Connect(std::string(redis_primary_addr),
-                                           redis_primary_port, client_info));
+                                           redis_primary_port));
     RAY_CHECK_OK(state->gcs_client.context()->AttachToEventLoop(loop));
   } else {
     state->db = NULL;

src/plasma/plasma_manager.cc

@@ -487,13 +487,8 @@ PlasmaManagerState *PlasmaManagerState_init(const char *store_socket_name,
                            "plasma_manager", manager_addr, db_connect_args);
     db_attach(state->db, state->loop, false);
 
-    ClientTableDataT client_info;
-    client_info.client_id = get_db_client_id(state->db).binary();
-    client_info.node_manager_address = std::string(manager_addr);
-    client_info.local_scheduler_port = 0;
-    client_info.object_manager_port = manager_port;
     RAY_CHECK_OK(state->gcs_client.Connect(std::string(redis_primary_addr),
-                                           redis_primary_port, client_info));
+                                           redis_primary_port));
     RAY_CHECK_OK(state->gcs_client.context()->AttachToEventLoop(state->loop));
   } else {
     state->db = NULL;

src/ray/CMakeLists.txt

@@ -38,8 +38,11 @@ set(RAY_SRCS
   gcs/asio.cc
   common/client_connection.cc
   object_manager/object_manager_client_connection.cc
-  object_manager/object_store_client.cc
+  object_manager/connection_pool.cc
+  object_manager/object_store_client_pool.cc
+  object_manager/object_store_notification_manager.cc
   object_manager/object_directory.cc
+  object_manager/transfer_queue.cc
   object_manager/object_manager.cc
   raylet/mock_gcs_client.cc
   raylet/task.cc

src/ray/common/client_connection.cc

@@ -7,20 +7,82 @@
 
 namespace ray {
 
+ray::Status TcpConnect(boost::asio::ip::tcp::socket &socket,
+                       const std::string &ip_address_string, int port) {
+  boost::asio::ip::address ip_address =
+      boost::asio::ip::address::from_string(ip_address_string);
+  boost::asio::ip::tcp::endpoint endpoint(ip_address, port);
+  boost::system::error_code error;
+  socket.connect(endpoint, error);
+  if (error) {
+    return ray::Status::IOError(error.message());
+  } else {
+    return ray::Status::OK();
+  }
+}
+
+template <class T>
+ServerConnection<T>::ServerConnection(boost::asio::basic_stream_socket<T> &&socket)
+    : socket_(std::move(socket)) {}
+
+template <class T>
+void ServerConnection<T>::WriteBuffer(
+    const std::vector<boost::asio::const_buffer> &buffer, boost::system::error_code &ec) {
+  boost::asio::write(socket_, buffer, ec);
+}
+
+template <class T>
+void ServerConnection<T>::ReadBuffer(
+    const std::vector<boost::asio::mutable_buffer> &buffer,
+    boost::system::error_code &ec) {
+  boost::asio::read(socket_, buffer, ec);
+}
+
+template <class T>
+ray::Status ServerConnection<T>::WriteMessage(int64_t type, int64_t length,
+                                              const uint8_t *message) {
+  std::vector<boost::asio::const_buffer> message_buffers;
+  auto write_version = RayConfig::instance().ray_protocol_version();
+  message_buffers.push_back(boost::asio::buffer(&write_version, sizeof(write_version)));
+  message_buffers.push_back(boost::asio::buffer(&type, sizeof(type)));
+  message_buffers.push_back(boost::asio::buffer(&length, sizeof(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;
+  boost::asio::write(socket_, message_buffers, error);
+  if (error) {
+    return ray::Status::IOError(error.message());
+  } else {
+    return ray::Status::OK();
+  }
+}
+
 template <class T>
 std::shared_ptr<ClientConnection<T>> ClientConnection<T>::Create(
-    ClientManager<T> &manager, boost::asio::basic_stream_socket<T> &&socket) {
+    ClientHandler<T> &client_handler, MessageHandler<T> &message_handler,
+    boost::asio::basic_stream_socket<T> &&socket) {
   std::shared_ptr<ClientConnection<T>> self(
-      new ClientConnection(manager, std::move(socket)));
+      new ClientConnection(message_handler, std::move(socket)));
   // Let our manager process our new connection.
-  self->manager_.ProcessNewClient(self);
+  client_handler(self);
   return self;
 }
 
 template <class T>
-ClientConnection<T>::ClientConnection(ClientManager<T> &manager,
+ClientConnection<T>::ClientConnection(MessageHandler<T> &message_handler,
                                       boost::asio::basic_stream_socket<T> &&socket)
-    : socket_(std::move(socket)), manager_(manager) {}
+    : ServerConnection<T>(std::move(socket)), message_handler_(message_handler) {}
+
+template <class T>
+const ClientID &ClientConnection<T>::GetClientID() {
+  return client_id_;
+}
+
+template <class T>
+void ClientConnection<T>::SetClientID(const ClientID &client_id) {
+  client_id_ = client_id;
+}
 
 template <class T>
 void ClientConnection<T>::ProcessMessages() {
@@ -31,7 +93,7 @@ void ClientConnection<T>::ProcessMessages() {
   header.push_back(boost::asio::buffer(&read_type_, sizeof(read_type_)));
   header.push_back(boost::asio::buffer(&read_length_, sizeof(read_length_)));
   boost::asio::async_read(
-      socket_, header,
+      ServerConnection<T>::socket_, header,
       boost::bind(&ClientConnection<T>::ProcessMessageHeader, this->shared_from_this(),
                   boost::asio::placeholders::error));
 }
@@ -52,57 +114,23 @@ void ClientConnection<T>::ProcessMessageHeader(const boost::system::error_code &
   read_message_.resize(read_length_);
   // Wait for the message to be read.
   boost::asio::async_read(
-      socket_, boost::asio::buffer(read_message_),
+      ServerConnection<T>::socket_, boost::asio::buffer(read_message_),
       boost::bind(&ClientConnection<T>::ProcessMessage, this->shared_from_this(),
                   boost::asio::placeholders::error));
 }
 
-template <class T>
-void ClientConnection<T>::WriteMessage(int64_t type, size_t length,
-                                       const uint8_t *message) {
-  std::vector<boost::asio::const_buffer> message_buffers;
-  write_version_ = RayConfig::instance().ray_protocol_version();
-  write_type_ = type;
-  write_length_ = length;
-  write_message_.assign(message, message + length);
-  message_buffers.push_back(boost::asio::buffer(&write_version_, sizeof(write_version_)));
-  message_buffers.push_back(boost::asio::buffer(&write_type_, sizeof(write_type_)));
-  message_buffers.push_back(boost::asio::buffer(&write_length_, sizeof(write_length_)));
-  message_buffers.push_back(boost::asio::buffer(write_message_));
-  boost::system::error_code error;
-  // Write the message and then wait for more messages.
-  boost::asio::async_write(
-      socket_, message_buffers,
-      boost::bind(&ClientConnection<T>::ProcessMessages, this->shared_from_this(),
-                  boost::asio::placeholders::error));
-}
-
 template <class T>
 void ClientConnection<T>::ProcessMessage(const boost::system::error_code &error) {
   if (error) {
     // TODO(hme): Disconnect differently & remove dependency on node_manager_generated.h
     read_type_ = protocol::MessageType_DisconnectClient;
   }
-  manager_.ProcessClientMessage(this->shared_from_this(), read_type_,
-                                read_message_.data());
-}
-
-template <class T>
-void ClientConnection<T>::ProcessMessages(const boost::system::error_code &error) {
-  if (error) {
-    ProcessMessage(error);
-  } else {
-    ProcessMessages();
-  }
+  message_handler_(this->shared_from_this(), read_type_, read_message_.data());
 }
 
+template class ServerConnection<boost::asio::local::stream_protocol>;
+template class ServerConnection<boost::asio::ip::tcp>;
 template class ClientConnection<boost::asio::local::stream_protocol>;
 template class ClientConnection<boost::asio::ip::tcp>;
 
-template <class T>
-ClientManager<T>::~ClientManager<T>() {}
-
-template class ClientManager<boost::asio::local::stream_protocol>;
-template class ClientManager<boost::asio::ip::tcp>;
-
 }  // namespace ray