[xray] Adds a driver table. (#2289)

This PR adds a driver table for the new GCS, which enables cleanup functionality associated with monitoring driver death.

Some testing in `monitor_test.py` is restored, but redis sharding for xray is needed to enable remaining tests.

src/common/lib/python/common_extension.cc

@@ -907,3 +907,12 @@ PyObject *check_simple_value(PyObject *self, PyObject *args) {
   }
   Py_RETURN_FALSE;
 }
+
+PyObject *compute_task_id(PyObject *self, PyObject *args) {
+  ObjectID object_id;
+  if (!PyArg_ParseTuple(args, "O&", &PyObjectToUniqueID, &object_id)) {
+    return NULL;
+  }
+  TaskID task_id = ray::ComputeTaskId(object_id);
+  return PyObjectID_make(task_id);
+}

src/common/lib/python/common_extension.h

@@ -56,6 +56,7 @@ int PyObjectToUniqueID(PyObject *object, ray::ObjectID *object_id);
 PyObject *PyObjectID_make(ray::ObjectID object_id);
 
 PyObject *check_simple_value(PyObject *self, PyObject *args);
+PyObject *compute_task_id(PyObject *self, PyObject *args);
 
 PyObject *PyTask_to_string(PyObject *, PyObject *args);
 PyObject *PyTask_from_string(PyObject *, PyObject *args);

src/local_scheduler/lib/python/local_scheduler_extension.cc

@@ -493,6 +493,8 @@ static PyTypeObject PyLocalSchedulerClientType = {
 static PyMethodDef local_scheduler_methods[] = {
     {"check_simple_value", check_simple_value, METH_VARARGS,
      "Should the object be passed by value?"},
+    {"compute_task_id", compute_task_id, METH_VARARGS,
+     "Return the task ID of an object ID."},
     {"task_from_string", PyTask_from_string, METH_VARARGS,
      "Creates a Python PyTask object from a string representation of "
      "TaskSpec."},

src/ray/gcs/client.cc

@@ -17,6 +17,7 @@ AsyncGcsClient::AsyncGcsClient(const ClientID &client_id, CommandType command_ty
   task_reconstruction_log_.reset(new TaskReconstructionLog(context_, this));
   task_lease_table_.reset(new TaskLeaseTable(context_, this));
   heartbeat_table_.reset(new HeartbeatTable(context_, this));
+  driver_table_.reset(new DriverTable(primary_context_, this));
   error_table_.reset(new ErrorTable(primary_context_, this));
   profile_table_.reset(new ProfileTable(context_, this));
   command_type_ = command_type;
@@ -88,6 +89,8 @@ HeartbeatTable &AsyncGcsClient::heartbeat_table() { return *heartbeat_table_; }
 
 ErrorTable &AsyncGcsClient::error_table() { return *error_table_; }
 
+DriverTable &AsyncGcsClient::driver_table() { return *driver_table_; }
+
 ProfileTable &AsyncGcsClient::profile_table() { return *profile_table_; }
 
 }  // namespace gcs

src/ray/gcs/client.h

@@ -60,6 +60,7 @@ class RAY_EXPORT AsyncGcsClient {
   ClientTable &client_table();
   HeartbeatTable &heartbeat_table();
   ErrorTable &error_table();
+  DriverTable &driver_table();
   ProfileTable &profile_table();
 
   // We also need something to export generic code to run on workers from the
@@ -92,6 +93,7 @@ class RAY_EXPORT AsyncGcsClient {
   std::unique_ptr<RedisAsioClient> asio_subscribe_client_;
   // The following context writes everything to the primary shard
   std::shared_ptr<RedisContext> primary_context_;
+  std::unique_ptr<DriverTable> driver_table_;
   std::unique_ptr<RedisAsioClient> asio_async_auxiliary_client_;
   std::unique_ptr<RedisAsioClient> asio_subscribe_auxiliary_client_;
   CommandType command_type_;

src/ray/gcs/format/gcs.fbs

@@ -4,6 +4,7 @@ enum Language:int {
   JAVA = 2
 }
 
+// These indexes are mapped to strings in ray_redis_module.cc.
 enum TablePrefix:int {
   UNUSED = 0,
   TASK,
@@ -15,6 +16,7 @@ enum TablePrefix:int {
   TASK_RECONSTRUCTION,
   HEARTBEAT,
   ERROR_INFO,
+  DRIVER,
   PROFILE,
   TASK_LEASE,
 }
@@ -30,6 +32,7 @@ enum TablePubsub:int {
   HEARTBEAT,
   ERROR_INFO,
   TASK_LEASE,
+  DRIVER,
 }
 
 table GcsTableEntry {
@@ -202,3 +205,10 @@ table TaskLeaseData {
   // The period that the lease is active for.
   timeout: long;
 }
+
+table DriverTableData {
+  // The driver ID.
+  driver_id: string;
+  // Whether it's dead.
+  is_dead: bool;
+}

src/ray/gcs/tables.cc

@@ -266,6 +266,17 @@ Status ProfileTable::AddProfileEventBatch(const ProfileTableData &profile_events
                 });
 }
 
+Status DriverTable::AppendDriverData(const JobID &driver_id, bool is_dead) {
+  auto data = std::make_shared<DriverTableDataT>();
+  data->driver_id = driver_id.binary();
+  data->is_dead = is_dead;
+  return Append(driver_id, driver_id, data,
+                [](ray::gcs::AsyncGcsClient *client, const JobID &id,
+                   const DriverTableDataT &data) {
+                  RAY_LOG(DEBUG) << "Driver entry added callback";
+                });
+}
+
 void ClientTable::RegisterClientAddedCallback(const ClientTableCallback &callback) {
   client_added_callback_ = callback;
   // Call the callback for any added clients that are cached.
@@ -425,6 +436,7 @@ template class Table<TaskID, TaskLeaseData>;
 template class Table<ClientID, HeartbeatTableData>;
 template class Log<JobID, ErrorTableData>;
 template class Log<UniqueID, ClientTableData>;
+template class Log<JobID, DriverTableData>;
 template class Log<UniqueID, ProfileTableData>;
 
 }  // namespace gcs

src/ray/gcs/tables.h

@@ -317,6 +317,23 @@ class HeartbeatTable : public Table<ClientID, HeartbeatTableData> {
   virtual ~HeartbeatTable() {}
 };
 
+class DriverTable : public Log<JobID, DriverTableData> {
+ public:
+  DriverTable(const std::shared_ptr<RedisContext> &context, AsyncGcsClient *client)
+      : Log(context, client) {
+    pubsub_channel_ = TablePubsub::DRIVER;
+    prefix_ = TablePrefix::DRIVER;
+  };
+  virtual ~DriverTable() {}
+
+  /// Appends driver data to the driver table.
+  ///
+  /// \param driver_id The driver id.
+  /// \param is_dead Whether the driver is dead.
+  /// \return The return status.
+  Status AppendDriverData(const JobID &driver_id, bool is_dead);
+};
+
 class FunctionTable : public Table<ObjectID, FunctionTableData> {
  public:
   FunctionTable(const std::shared_ptr<RedisContext> &context, AsyncGcsClient *client)

src/ray/object_manager/object_manager.cc

@@ -25,8 +25,10 @@ ObjectManager::ObjectManager(asio::io_service &main_service,
   RAY_CHECK(config_.max_sends > 0);
   RAY_CHECK(config_.max_receives > 0);
   main_service_ = &main_service;
-  store_notification_.SubscribeObjAdded(
-      [this](const ObjectInfoT &object_info) { NotifyDirectoryObjectAdd(object_info); });
+  store_notification_.SubscribeObjAdded([this](const ObjectInfoT &object_info) {
+    NotifyDirectoryObjectAdd(object_info);
+    HandleUnfulfilledPushRequests(object_info);
+  });
   store_notification_.SubscribeObjDeleted(
       [this](const ObjectID &oid) { NotifyDirectoryObjectDeleted(oid); });
   StartIOService();
@@ -49,8 +51,10 @@ ObjectManager::ObjectManager(asio::io_service &main_service,
   RAY_CHECK(config_.max_receives > 0);
   // TODO(hme) Client ID is never set with this constructor.
   main_service_ = &main_service;
-  store_notification_.SubscribeObjAdded(
-      [this](const ObjectInfoT &object_info) { NotifyDirectoryObjectAdd(object_info); });
+  store_notification_.SubscribeObjAdded([this](const ObjectInfoT &object_info) {
+    NotifyDirectoryObjectAdd(object_info);
+    HandleUnfulfilledPushRequests(object_info);
+  });
   store_notification_.SubscribeObjDeleted(
       [this](const ObjectID &oid) { NotifyDirectoryObjectDeleted(oid); });
   StartIOService();
@@ -89,6 +93,10 @@ void ObjectManager::NotifyDirectoryObjectAdd(const ObjectInfoT &object_info) {
   local_objects_[object_id] = object_info;
   ray::Status status =
       object_directory_->ReportObjectAdded(object_id, client_id_, object_info);
+}
+
+void ObjectManager::HandleUnfulfilledPushRequests(const ObjectInfoT &object_info) {
+  ObjectID object_id = ObjectID::from_binary(object_info.object_id);
   // Handle the unfulfilled_push_requests_ which contains the push request that is not
   // completed due to unsatisfied local objects.
   auto iter = unfulfilled_push_requests_.find(object_id);

src/ray/object_manager/object_manager.h

@@ -266,6 +266,10 @@ class ObjectManager : public ObjectManagerInterface {
   /// Register object remove with directory.
   void NotifyDirectoryObjectDeleted(const ObjectID &object_id);
 
+  /// Handle any push requests that were made before an object was available.
+  /// This is invoked when an "object added" notification is received from the store.
+  void HandleUnfulfilledPushRequests(const ObjectInfoT &object_info);
+
   /// Part of an asynchronous sequence of Pull methods.
   /// Uses an existing connection or creates a connection to ClientID.
   /// Executes on main_service_ thread.

src/ray/raylet/node_manager.cc

@@ -193,12 +193,33 @@ ray::Status NodeManager::RegisterGcs() {
         RAY_LOG(DEBUG) << "heartbeat table subscription done callback called.";
       }));
 
+  // Subscribe to driver table updates.
+  const auto driver_table_handler = [this](
+      gcs::AsyncGcsClient *client, const ClientID &client_id,
+      const std::vector<DriverTableDataT> &driver_data) {
+    HandleDriverTableUpdate(client_id, driver_data);
+  };
+  RAY_RETURN_NOT_OK(gcs_client_->driver_table().Subscribe(JobID::nil(), UniqueID::nil(),
+                                                          driver_table_handler, nullptr));
+
   // Start sending heartbeats to the GCS.
   Heartbeat();
 
   return ray::Status::OK();
 }
 
+void NodeManager::HandleDriverTableUpdate(
+    const ClientID &id, const std::vector<DriverTableDataT> &driver_data) {
+  for (const auto &entry : driver_data) {
+    RAY_LOG(DEBUG) << "HandleDriverTableUpdate " << UniqueID::from_binary(entry.driver_id)
+                   << " " << entry.is_dead;
+    if (entry.is_dead) {
+      // TODO: Implement cleanup on driver death. For reference,
+      // see handle_driver_removed_callback in local_scheduler.cc
+    }
+  }
+}
+
 void NodeManager::Heartbeat() {
   RAY_LOG(DEBUG) << "[Heartbeat] sending heartbeat.";
   auto &heartbeat_table = gcs_client_->heartbeat_table();
@@ -449,6 +470,7 @@ void NodeManager::ProcessClientMessage(
   switch (static_cast<protocol::MessageType>(message_type)) {
   case protocol::MessageType::RegisterClientRequest: {
     auto message = flatbuffers::GetRoot<protocol::RegisterClientRequest>(message_data);
+    client->SetClientID(from_flatbuf(*message->client_id()));
     auto worker = std::make_shared<Worker>(message->worker_pid(), client);
     if (message->is_worker()) {
       // Register the new worker.
@@ -543,6 +565,8 @@ void NodeManager::ProcessClientMessage(
       DispatchTasks();
     } else {
       // The client is a driver.
+      RAY_CHECK_OK(gcs_client_->driver_table().AppendDriverData(client->GetClientID(),
+                                                                /*is_dead=*/true));
       const std::shared_ptr<Worker> driver = worker_pool_.GetRegisteredDriver(client);
       RAY_CHECK(driver);
       auto driver_id = driver->GetAssignedTaskId();

src/ray/raylet/node_manager.h

@@ -144,6 +144,10 @@ class NodeManager {
   /// accounting, but does not write to any global accounting in the GCS.
   void HandleObjectMissing(const ObjectID &object_id);
 
+  /// Handles updates to driver table.
+  void HandleDriverTableUpdate(const ClientID &id,
+                               const std::vector<DriverTableDataT> &driver_data);
+
   boost::asio::io_service &io_service_;
   ObjectManager &object_manager_;
   /// A Plasma object store client. This is used exclusively for creating new