[xray] Add error table and push error messages to driver through node manager. (#2256)

* Fix documentation indentation.

* Add error table to GCS and push error messages through node manager.

* Add type to error data.

* Linting

* Fix failure_test bug.

* Linting.

* Enable one more test.

* Attempt to fix doc building.

* Restructuring

* Fixes

* More fixes.

* Move current_time_ms function into util.h.

src/common/redis_module/ray_redis_module.cc

@@ -61,14 +61,6 @@ extern RedisChainModule module;
     return RedisModule_ReplyWithError(ctx, (MESSAGE)); \
   }
 
-// NOTE(swang): The order of prefixes here must match the TablePrefix enum
-// defined in src/ray/gcs/format/gcs.fbs.
-static const char *table_prefixes[] = {
-    NULL,         "TASK:",  "TASK:",     "CLIENT:",
-    "OBJECT:",    "ACTOR:", "FUNCTION:", "TASK_RECONSTRUCTION:",
-    "HEARTBEAT:",
-};
-
 /// Parse a Redis string into a TablePubsub channel.
 TablePubsub ParseTablePubsub(const RedisModuleString *pubsub_channel_str) {
   long long pubsub_channel_long;
@@ -128,8 +120,8 @@ RedisModuleKey *OpenPrefixedKey(RedisModuleCtx *ctx,
       << "This table has no prefix registered";
   RAY_CHECK(prefix >= TablePrefix::MIN && prefix <= TablePrefix::MAX)
       << "Prefix must be a valid TablePrefix";
-  return OpenPrefixedKey(ctx, table_prefixes[static_cast<long long>(prefix)],
-                         keyname, mode, mutated_key_str);
+  return OpenPrefixedKey(ctx, EnumNameTablePrefix(prefix), keyname, mode,
+                         mutated_key_str);
 }
 
 RedisModuleKey *OpenPrefixedKey(RedisModuleCtx *ctx,

src/local_scheduler/lib/python/local_scheduler_extension.cc

@@ -286,6 +286,29 @@ static PyObject *PyLocalSchedulerClient_wait(PyObject *self, PyObject *args) {
   return Py_BuildValue("(OO)", py_found, py_remaining);
 }
 
+static PyObject *PyLocalSchedulerClient_push_error(PyObject *self,
+                                                   PyObject *args) {
+  JobID job_id;
+  const char *type;
+  int type_length;
+  const char *error_message;
+  int error_message_length;
+  double timestamp;
+  if (!PyArg_ParseTuple(args, "O&s#s#d", &PyObjectToUniqueID, &job_id, &type,
+                        &type_length, &error_message, &error_message_length,
+                        &timestamp)) {
+    return NULL;
+  }
+
+  local_scheduler_push_error(reinterpret_cast<PyLocalSchedulerClient *>(self)
+                                 ->local_scheduler_connection,
+                             job_id, std::string(type, type_length),
+                             std::string(error_message, error_message_length),
+                             timestamp);
+
+  Py_RETURN_NONE;
+}
+
 static PyMethodDef PyLocalSchedulerClient_methods[] = {
     {"disconnect", (PyCFunction) PyLocalSchedulerClient_disconnect, METH_NOARGS,
      "Notify the local scheduler that this client is exiting gracefully."},
@@ -313,6 +336,8 @@ static PyMethodDef PyLocalSchedulerClient_methods[] = {
      (PyCFunction) PyLocalSchedulerClient_set_actor_frontier, METH_VARARGS, ""},
     {"wait", (PyCFunction) PyLocalSchedulerClient_wait, METH_VARARGS,
      "Wait for a list of objects to be created."},
+    {"push_error", (PyCFunction) PyLocalSchedulerClient_push_error,
+     METH_VARARGS, "Push an error message to the relevant driver."},
     {NULL} /* Sentinel */
 };
 

src/local_scheduler/local_scheduler_client.cc

@@ -306,3 +306,19 @@ std::pair<std::vector<ObjectID>, std::vector<ObjectID>> local_scheduler_wait(
   free(reply);
   return result;
 }
+
+void local_scheduler_push_error(LocalSchedulerConnection *conn,
+                                const JobID &job_id,
+                                const std::string &type,
+                                const std::string &error_message,
+                                double timestamp) {
+  flatbuffers::FlatBufferBuilder fbb;
+  auto message = ray::protocol::CreatePushErrorRequest(
+      fbb, to_flatbuf(fbb, job_id), fbb.CreateString(type),
+      fbb.CreateString(error_message), timestamp);
+  fbb.Finish(message);
+
+  write_message(conn->conn, static_cast<int64_t>(
+                                ray::protocol::MessageType::PushErrorRequest),
+                fbb.GetSize(), fbb.GetBufferPointer());
+}

src/local_scheduler/local_scheduler_client.h

@@ -211,4 +211,18 @@ std::pair<std::vector<ObjectID>, std::vector<ObjectID>> local_scheduler_wait(
     int64_t timeout_milliseconds,
     bool wait_local);
 
+/// Push an error to the relevant driver.
+///
+/// \param conn The connection information.
+/// \param The ID of the job that the error is for.
+/// \param The type of the error.
+/// \param The error message.
+/// \param The timestamp of the error.
+/// \return Void.
+void local_scheduler_push_error(LocalSchedulerConnection *conn,
+                                const JobID &job_id,
+                                const std::string &type,
+                                const std::string &error_message,
+                                double timestamp);
+
 #endif

src/ray/gcs/client.cc

@@ -15,6 +15,7 @@ AsyncGcsClient::AsyncGcsClient(const ClientID &client_id, CommandType command_ty
   raylet_task_table_.reset(new raylet::TaskTable(context_, this, command_type));
   task_reconstruction_log_.reset(new TaskReconstructionLog(context_, this));
   heartbeat_table_.reset(new HeartbeatTable(context_, this));
+  error_table_.reset(new ErrorTable(context_, this));
   command_type_ = command_type;
 }
 
@@ -74,6 +75,9 @@ FunctionTable &AsyncGcsClient::function_table() { return *function_table_; }
 ClassTable &AsyncGcsClient::class_table() { return *class_table_; }
 
 HeartbeatTable &AsyncGcsClient::heartbeat_table() { return *heartbeat_table_; }
+
+ErrorTable &AsyncGcsClient::error_table() { return *error_table_; }
+
 }  // namespace gcs
 
 }  // namespace ray

src/ray/gcs/client.h

@@ -57,7 +57,7 @@ class RAY_EXPORT AsyncGcsClient {
   TaskReconstructionLog &task_reconstruction_log();
   ClientTable &client_table();
   HeartbeatTable &heartbeat_table();
-  inline ErrorTable &error_table();
+  ErrorTable &error_table();
 
   // We also need something to export generic code to run on workers from the
   // driver (to set the PYTHONPATH)
@@ -78,6 +78,7 @@ class RAY_EXPORT AsyncGcsClient {
   std::unique_ptr<ActorTable> actor_table_;
   std::unique_ptr<TaskReconstructionLog> task_reconstruction_log_;
   std::unique_ptr<HeartbeatTable> heartbeat_table_;
+  std::unique_ptr<ErrorTable> error_table_;
   std::unique_ptr<ClientTable> client_table_;
   std::shared_ptr<RedisContext> context_;
   std::unique_ptr<RedisAsioClient> asio_async_client_;

src/ray/gcs/format/gcs.fbs

@@ -14,6 +14,7 @@ enum TablePrefix:int {
   FUNCTION,
   TASK_RECONSTRUCTION,
   HEARTBEAT,
+  ERROR_INFO,
 }
 
 // The channel that Add operations to the Table should be published on, if any.
@@ -24,7 +25,8 @@ enum TablePubsub:int {
   CLIENT,
   OBJECT,
   ACTOR,
-  HEARTBEAT
+  HEARTBEAT,
+  ERROR_INFO,
 }
 
 table GcsTableEntry {
@@ -103,6 +105,14 @@ table ActorTableData {
 }
 
 table ErrorTableData {
+  // The ID of the job that the error is for.
+  job_id: string;
+  // The type of the error.
+  type: string;
+  // The error message.
+  error_message: string;
+  // The timestamp of the error message.
+  timestamp: double;
 }
 
 table CustomSerializerData {

src/ray/gcs/tables.cc

@@ -183,6 +183,19 @@ Status Table<ID, Data>::Subscribe(const JobID &job_id, const ClientID &client_id
       done);
 }
 
+Status ErrorTable::PushErrorToDriver(const JobID &job_id, const std::string &type,
+                                     const std::string &error_message, double timestamp) {
+  auto data = std::make_shared<ErrorTableDataT>();
+  data->job_id = job_id.binary();
+  data->type = type;
+  data->error_message = error_message;
+  data->timestamp = timestamp;
+  return Append(job_id, job_id, data, [](ray::gcs::AsyncGcsClient *client,
+                                         const JobID &id, const ErrorTableDataT &data) {
+    RAY_LOG(DEBUG) << "Error message pushed callback";
+  });
+}
+
 void ClientTable::RegisterClientAddedCallback(const ClientTableCallback &callback) {
   client_added_callback_ = callback;
   // Call the callback for any added clients that are cached.
@@ -333,6 +346,7 @@ template class Table<TaskID, TaskTableData>;
 template class Log<ActorID, ActorTableData>;
 template class Log<TaskID, TaskReconstructionData>;
 template class Table<ClientID, HeartbeatTableData>;
+template class Log<JobID, ErrorTableData>;
 template class Log<UniqueID, ClientTableData>;
 
 }  // namespace gcs

src/ray/gcs/tables.h

@@ -95,7 +95,7 @@ class Log : virtual public PubsubInterface<ID> {
   /// \param id The ID of the data that is added to the GCS.
   /// \param data Data to append to the log.
   /// \param done Callback that is called once the data has been written to the
-  ///        GCS.
+  /// GCS.
   /// \return Status
   Status Append(const JobID &job_id, const ID &id, std::shared_ptr<DataT> &data,
                 const WriteCallback &done);
@@ -108,10 +108,9 @@ class Log : virtual public PubsubInterface<ID> {
   /// \param data Data to append to the log.
   /// \param done Callback that is called if the data was appended to the log.
   /// \param failure Callback that is called if the data was not appended to
-  ///        the log because the log length did not match the given
-  ///        `log_length`.
+  /// the log because the log length did not match the given `log_length`.
   /// \param log_length The number of entries that the log must have for the
-  ///        append to succeed.
+  /// append to succeed.
   /// \return Status
   Status AppendAt(const JobID &job_id, const ID &id, std::shared_ptr<DataT> &data,
                   const WriteCallback &done, const WriteCallback &failure,
@@ -122,7 +121,7 @@ class Log : virtual public PubsubInterface<ID> {
   /// \param job_id The ID of the job (= driver).
   /// \param id The ID of the data that is looked up in the GCS.
   /// \param lookup Callback that is called after lookup. If the callback is
-  ///        called with an empty vector, then there was no data at the key.
+  /// called with an empty vector, then there was no data at the key.
   /// \return Status
   Status Lookup(const JobID &job_id, const ID &id, const Callback &lookup);
 
@@ -133,15 +132,14 @@ class Log : virtual public PubsubInterface<ID> {
   ///
   /// \param job_id The ID of the job (= driver).
   /// \param client_id The type of update to listen to. If this is nil, then a
-  ///        message for each Add to the table will be received. Else, only
-  ///        messages for the given client will be received. In the latter
-  ///        case, the client may request notifications on specific keys in the
-  ///        table via `RequestNotifications`.
+  /// message for each Add to the table will be received. Else, only
+  /// messages for the given client will be received. In the latter
+  /// case, the client may request notifications on specific keys in the
+  /// table via `RequestNotifications`.
   /// \param subscribe Callback that is called on each received message. If the
-  ///        callback is called with an empty vector, then there was no data at
-  ///        the key.
+  /// callback is called with an empty vector, then there was no data at the key.
   /// \param done Callback that is called when subscription is complete and we
-  ///        are ready to receive messages.
+  /// are ready to receive messages.
   /// \return Status
   Status Subscribe(const JobID &job_id, const ClientID &client_id,
                    const Callback &subscribe, const SubscriptionCallback &done);
@@ -158,8 +156,8 @@ class Log : virtual public PubsubInterface<ID> {
   /// \param job_id The ID of the job (= driver).
   /// \param id The ID of the key to request notifications for.
   /// \param client_id The client who is requesting notifications. Before
-  ///        notifications can be requested, a call to `Subscribe` to this
-  ///        table with the same `client_id` must complete successfully.
+  /// notifications can be requested, a call to `Subscribe` to this
+  /// table with the same `client_id` must complete successfully.
   /// \return Status
   Status RequestNotifications(const JobID &job_id, const ID &id,
                               const ClientID &client_id);
@@ -241,7 +239,7 @@ class Table : private Log<ID, Data>,
   /// \param id The ID of the data that is added to the GCS.
   /// \param data Data that is added to the GCS.
   /// \param done Callback that is called once the data has been written to the
-  ///        GCS.
+  /// GCS.
   /// \return Status
   Status Add(const JobID &job_id, const ID &id, std::shared_ptr<DataT> &data,
              const WriteCallback &done);
@@ -251,9 +249,9 @@ class Table : private Log<ID, Data>,
   /// \param job_id The ID of the job (= driver).
   /// \param id The ID of the data that is looked up in the GCS.
   /// \param lookup Callback that is called after lookup if there was data the
-  ///        key.
+  /// key.
   /// \param failure Callback that is called after lookup if there was no data
-  ///        at the key.
+  /// at the key.
   /// \return Status
   Status Lookup(const JobID &job_id, const ID &id, const Callback &lookup,
                 const FailureCallback &failure);
@@ -366,10 +364,10 @@ class TaskTable : public Table<TaskID, TaskTableData> {
   ///
   /// \param task_id The task ID of the task entry to update.
   /// \param test_state_bitmask The bitmask to apply to the task entry's current
-  ///        scheduling state.  The update happens if and only if the current
-  ///        scheduling state AND-ed with the bitmask is greater than 0.
+  /// scheduling state.  The update happens if and only if the current
+  /// scheduling state AND-ed with the bitmask is greater than 0.
   /// \param update_state The value to update the task entry's scheduling state
-  ///        with, if the current state matches test_state_bitmask.
+  /// with, if the current state matches test_state_bitmask.
   /// \param callback Function to be called when database returns result.
   /// \return Status
   Status TestAndUpdate(const JobID &job_id, const TaskID &id,
@@ -397,16 +395,14 @@ class TaskTable : public Table<TaskID, TaskTableData> {
   /// task's local scheduler ID.
   ///
   /// \param local_scheduler_id The db_client_id of the local scheduler whose
-  ///        events we want to listen to. If you want to subscribe to updates
-  ///        from
-  ///        all local schedulers, pass in NIL_ID.
+  /// events we want to listen to. If you want to subscribe to updates from
+  /// all local schedulers, pass in NIL_ID.
   /// \param subscribe_callback Callback that will be called when the task table
-  /// is
-  ///        updated.
+  /// is updated.
   /// \param state_filter Events we want to listen to. Can have values from the
-  ///        enum "scheduling_state" in task.h.
-  ///        TODO(pcm): Make it possible to combine these using flags like
-  ///        TASK_STATUS_WAITING | TASK_STATUS_SCHEDULED.
+  /// enum "scheduling_state" in task.h.
+  /// TODO(pcm): Make it possible to combine these using flags like
+  /// TASK_STATUS_WAITING | TASK_STATUS_SCHEDULED.
   /// \param callback Function to be called when database returns result.
   /// \return Status
   Status SubscribeToTask(const JobID &job_id, const ClientID &local_scheduler_id,
@@ -422,7 +418,28 @@ Status TaskTableTestAndUpdate(AsyncGcsClient *gcs_client, const TaskID &task_id,
                               SchedulingState update_state,
                               const TaskTable::TestAndUpdateCallback &callback);
 
-using ErrorTable = Table<TaskID, ErrorTableData>;
+class ErrorTable : private Log<JobID, ErrorTableData> {
+ public:
+  ErrorTable(const std::shared_ptr<RedisContext> &context, AsyncGcsClient *client)
+      : Log(context, client) {
+    pubsub_channel_ = TablePubsub::ERROR_INFO;
+    prefix_ = TablePrefix::ERROR_INFO;
+  };
+
+  /// Push an error message for a specific job.
+  ///
+  /// TODO(rkn): We need to make sure that the errors are unique because
+  /// duplicate messages currently cause failures (the GCS doesn't allow it).
+  ///
+  /// \param job_id The ID of the job that generated the error. If the error
+  /// should be pushed to all jobs, then this should be nil.
+  /// \param type The type of the error.
+  /// \param error_message The error message to push.
+  /// \param timestamp The timestamp of the error.
+  /// \return Status.
+  Status PushErrorToDriver(const JobID &job_id, const std::string &type,
+                           const std::string &error_message, double timestamp);
+};
 
 using CustomSerializerTable = Table<ClassID, CustomSerializerData>;
 
@@ -467,7 +484,7 @@ class ClientTable : private Log<UniqueID, ClientTableData> {
   /// and begins subscription to client table notifications.
   ///
   /// \param Information about the connecting client. This must have the
-  ///        same client_id as the one set in the client table.
+  /// same client_id as the one set in the client table.
   /// \return Status
   ray::Status Connect(const ClientTableDataT &local_client);
 
@@ -499,7 +516,7 @@ class ClientTable : private Log<UniqueID, ClientTableData> {
   ///
   /// \param client The client to get information about.
   /// \return A reference to the requested client. If the client is not in the
-  ///         cache, then an entry with a nil ClientID will be returned.
+  /// cache, then an entry with a nil ClientID will be returned.
   const ClientTableDataT &GetClient(const ClientID &client);
 
   /// Get the local client's ID.

src/ray/raylet/format/node_manager.fbs

@@ -58,7 +58,10 @@ enum MessageType:int {
   WaitRequest,
   // The response message to WaitRequest; replies with the objects found and objects
   // remaining.
-  WaitReply
+  WaitReply,
+  // Push an error to the relevant driver. This is sent from a worker to the
+  // node manager.
+  PushErrorRequest,
 }
 
 table TaskExecutionSpecification {
@@ -154,3 +157,15 @@ table WaitReply {
   // List of object ids not found.
   remaining: [string];
 }
+
+// This struct is the same as ErrorTableData.
+table PushErrorRequest {
+  // The ID of the job that the error is for.
+  job_id: string;
+  // The type of the error.
+  type: string;
+  // The error message.
+  error_message: string;
+  // The timestamp of the error message.
+  timestamp: double;
+}

src/ray/raylet/node_manager.cc

@@ -3,6 +3,7 @@
 #include "common_protocol.h"
 #include "local_scheduler/format/local_scheduler_generated.h"
 #include "ray/raylet/format/node_manager_generated.h"
+#include "ray/util/util.h"
 
 namespace {
 
@@ -372,11 +373,28 @@ void NodeManager::ProcessClientMessage(
 
     // This if statement distinguishes workers from drivers.
     if (worker) {
-      // TODO(swang): Handle the case where the worker is killed while
-      // executing a task. Clean up the assigned task's resources, return an
-      // error to the driver.
-      // RAY_CHECK(worker->GetAssignedTaskId().is_nil())
-      //    << "Worker died while executing task: " << worker->GetAssignedTaskId();
+      // Handle the case where the worker is killed while executing a task.
+      // Clean up the assigned task's resources, push an error to the driver.
+      const TaskID &task_id = worker->GetAssignedTaskId();
+      if (!task_id.is_nil()) {
+        auto const &running_tasks = local_queues_.GetRunningTasks();
+        // TODO(rkn): This is too heavyweight just to get the task's driver ID.
+        auto const it = std::find_if(
+            running_tasks.begin(), running_tasks.end(), [task_id](const Task &task) {
+              return task.GetTaskSpecification().TaskId() == task_id;
+            });
+        RAY_CHECK(running_tasks.size() != 0);
+        RAY_CHECK(it != running_tasks.end());
+        JobID job_id = it->GetTaskSpecification().DriverId();
+        // TODO(rkn): Define this constant somewhere else.
+        std::string type = "worker_died";
+        std::ostringstream error_message;
+        error_message << "A worker died or was killed while executing task " << task_id
+                      << ".";
+        RAY_CHECK_OK(gcs_client_->error_table().PushErrorToDriver(
+            job_id, type, error_message.str(), current_time_ms()));
+      }
+
       worker_pool_.DisconnectWorker(worker);
 
       const ClientID &client_id = gcs_client_->client_table().GetLocalClientId();
@@ -521,6 +539,17 @@ void NodeManager::ProcessClientMessage(
         });
     RAY_CHECK_OK(status);
   } break;
+  case protocol::MessageType::PushErrorRequest: {
+    auto message = flatbuffers::GetRoot<protocol::PushErrorRequest>(message_data);
+
+    JobID job_id = from_flatbuf(*message->job_id());
+    auto const &type = string_from_flatbuf(*message->type());
+    auto const &error_message = string_from_flatbuf(*message->error_message());
+    double timestamp = message->timestamp();
+
+    RAY_CHECK_OK(gcs_client_->error_table().PushErrorToDriver(job_id, type, error_message,
+                                                              timestamp));
+  } break;
 
   default:
     RAY_LOG(FATAL) << "Received unexpected message type " << message_type;

src/ray/util/CMakeLists.txt

@@ -1,6 +1,7 @@
 install(FILES
   logging.h
   macros.h
+  util.h
   visibility.h
   DESTINATION "${CMAKE_INSTALL_INCLUDEDIR}/ray/util"
 )

src/ray/util/util.h

@@ -0,0 +1,19 @@
+#ifndef RAY_UTIL_UTIL_H
+#define RAY_UTIL_UTIL_H
+
+#include <chrono>
+
+/// Return the number of milliseconds since the Unix epoch.
+///
+/// TODO(rkn): This function appears in multiple places. It should be
+/// deduplicated.
+///
+/// \return The number of milliseconds since the Unix epoch.
+int64_t current_time_ms() {
+  std::chrono::milliseconds ms_since_epoch =
+      std::chrono::duration_cast<std::chrono::milliseconds>(
+          std::chrono::steady_clock::now().time_since_epoch());
+  return ms_since_epoch.count();
+}
+
+#endif  // RAY_UTIL_UTIL_H