Push an error to the driver when the workload hangs on ray.put reconstruction (#382)

* Fix worker blocked bug

* tmp

* Push an error to the driver on ray.put for non-driver tasks

* Fix result table tests

* Fix test, logging

* Address comments

* Fix suppression bug

* Fix redis module test

* Edit error message

* Get values in chunks during reconstruction

* Test case for driver ray.put errors

* Error for evicting ray.put objects from the driver

* Fix tests

* Reduce verbosity

* Documentation

src/common/state/error_table.h

@@ -18,14 +18,19 @@ typedef enum {
   /** An object was added with a different hash from the existing
    *  one. */
   OBJECT_HASH_MISMATCH_ERROR_INDEX = 0,
+  /** An object that was created through a ray.put is lost. */
+  PUT_RECONSTRUCTION_ERROR_INDEX,
   /** The total number of error types. */
   MAX_ERROR_INDEX
 } error_index;
 
 /** Information about the error to be displayed to the user. */
-static const char *error_types[] = {"object_hash_mismatch"};
+static const char *error_types[] = {"object_hash_mismatch",
+                                    "put_reconstruction"};
 static const char *error_messages[] = {
-    "A nondeterministic task was reexecuted."};
+    "A nondeterministic task was reexecuted.",
+    "An object created by ray.put was evicted and could not be reconstructed. "
+    "The driver may need to be restarted."};
 
 /**
  * Push an error to the given Python driver.

src/common/state/object_table.cc

@@ -104,13 +104,16 @@ void object_info_subscribe(DBHandle *db_handle,
 
 void result_table_add(DBHandle *db_handle,
                       ObjectID object_id,
-                      TaskID task_id_arg,
+                      TaskID task_id,
+                      bool is_put,
                       RetryInfo *retry,
                       result_table_done_callback done_callback,
                       void *user_context) {
-  TaskID *task_id_copy = (TaskID *) malloc(sizeof(TaskID));
-  memcpy(task_id_copy, task_id_arg.id, sizeof(*task_id_copy));
-  init_table_callback(db_handle, object_id, __func__, task_id_copy, retry,
+  ResultTableAddInfo *info =
+      (ResultTableAddInfo *) malloc(sizeof(ResultTableAddInfo));
+  info->task_id = task_id;
+  info->is_put = is_put;
+  init_table_callback(db_handle, object_id, __func__, info, retry,
                       (table_done_callback) done_callback,
                       redis_result_table_add, user_context);
 }

src/common/state/object_table.h

@@ -224,6 +224,15 @@ typedef struct {
 typedef void (*result_table_done_callback)(ObjectID object_id,
                                            void *user_context);
 
+/** Information about a result table entry to add. */
+typedef struct {
+  /** The task ID of the task that created the requested object. */
+  TaskID task_id;
+  /** True if the object was created through a put, and false if created by
+   *  return value. */
+  bool is_put;
+} ResultTableAddInfo;
+
 /**
  * Add information about a new object to the object table. This
  * is immutable information like the ID of the task that
@@ -232,6 +241,8 @@ typedef void (*result_table_done_callback)(ObjectID object_id,
  * @param db_handle Handle to object_table database.
  * @param object_id ID of the object to add.
  * @param task_id ID of the task that creates this object.
+ * @param is_put A boolean that is true if the object was created through a
+ *        ray.put, and false if the object was created by return value.
  * @param retry Information about retrying the request to the database.
  * @param done_callback Function to be called when database returns result.
  * @param user_context Context passed by the caller.
@@ -240,6 +251,7 @@ typedef void (*result_table_done_callback)(ObjectID object_id,
 void result_table_add(DBHandle *db_handle,
                       ObjectID object_id,
                       TaskID task_id,
+                      bool is_put,
                       RetryInfo *retry,
                       result_table_done_callback done_callback,
                       void *user_context);
@@ -247,6 +259,7 @@ void result_table_add(DBHandle *db_handle,
 /** Callback called when the result table lookup completes. */
 typedef void (*result_table_lookup_callback)(ObjectID object_id,
                                              TaskID task_id,
+                                             bool is_put,
                                              void *user_context);
 
 /**

src/common/state/redis.cc

@@ -342,12 +342,14 @@ void redis_result_table_add(TableCallbackData *callback_data) {
   CHECK(callback_data);
   DBHandle *db = callback_data->db_handle;
   ObjectID id = callback_data->id;
-  TaskID *result_task_id = (TaskID *) callback_data->data;
+  ResultTableAddInfo *info = (ResultTableAddInfo *) callback_data->data;
+  int is_put = info->is_put ? 1 : 0;
+
   /* Add the result entry to the result table. */
   int status = redisAsyncCommand(
       db->context, redis_result_table_add_callback,
-      (void *) callback_data->timer_id, "RAY.RESULT_TABLE_ADD %b %b", id.id,
-      sizeof(id.id), result_task_id->id, sizeof(result_task_id->id));
+      (void *) callback_data->timer_id, "RAY.RESULT_TABLE_ADD %b %b %d", id.id,
+      sizeof(id.id), info->task_id.id, sizeof(info->task_id.id), is_put);
   if ((status == REDIS_ERR) || db->context->err) {
     LOG_REDIS_DEBUG(db->context, "Error in result table add");
   }
@@ -386,16 +388,19 @@ void redis_result_table_lookup_callback(redisAsyncContext *c,
          reply->type);
   /* Parse the task from the reply. */
   TaskID result_id = NIL_TASK_ID;
+  bool is_put = false;
   if (reply->type == REDIS_REPLY_STRING) {
-    CHECK(reply->len == sizeof(result_id));
-    memcpy(&result_id, reply->str, reply->len);
+    auto message = flatbuffers::GetRoot<ResultTableReply>(reply->str);
+    result_id = from_flatbuf(message->task_id());
+    is_put = message->is_put();
   }
 
   /* Call the done callback if there is one. */
   result_table_lookup_callback done_callback =
       (result_table_lookup_callback) callback_data->done_callback;
   if (done_callback != NULL) {
-    done_callback(callback_data->id, result_id, callback_data->user_context);
+    done_callback(callback_data->id, result_id, is_put,
+                  callback_data->user_context);
   }
   /* Clean up timer and callback. */
   destroy_timer_callback(db->loop, callback_data);
@@ -761,11 +766,15 @@ void redis_task_table_test_and_update_callback(redisAsyncContext *c,
   redisReply *reply = (redisReply *) r;
   /* Parse the task from the reply. */
   Task *task = parse_and_construct_task_from_redis_reply(reply);
+  /* Determine whether the update happened. */
+  auto message = flatbuffers::GetRoot<TaskReply>(reply->str);
+  bool updated = message->updated();
+
   /* Call the done callback if there is one. */
-  task_table_get_callback done_callback =
-      (task_table_get_callback) callback_data->done_callback;
+  task_table_test_and_update_callback done_callback =
+      (task_table_test_and_update_callback) callback_data->done_callback;
   if (done_callback != NULL) {
-    done_callback(task, callback_data->user_context);
+    done_callback(task, callback_data->user_context, updated);
   }
   /* Free the task if it is not NULL. */
   if (task != NULL) {

src/common/state/task_table.cc

@@ -33,13 +33,14 @@ void task_table_update(DBHandle *db_handle,
                       redis_task_table_update, user_context);
 }
 
-void task_table_test_and_update(DBHandle *db_handle,
-                                TaskID task_id,
-                                int test_state_bitmask,
-                                int update_state,
-                                RetryInfo *retry,
-                                task_table_get_callback done_callback,
-                                void *user_context) {
+void task_table_test_and_update(
+    DBHandle *db_handle,
+    TaskID task_id,
+    int test_state_bitmask,
+    int update_state,
+    RetryInfo *retry,
+    task_table_test_and_update_callback done_callback,
+    void *user_context) {
   TaskTableTestAndUpdateData *update_data =
       (TaskTableTestAndUpdateData *) malloc(sizeof(TaskTableTestAndUpdateData));
   update_data->test_state_bitmask = test_state_bitmask;

src/common/state/task_table.h

@@ -28,6 +28,13 @@ typedef void (*task_table_done_callback)(TaskID task_id, void *user_context);
  * was not in the task table, then the task pointer will be NULL. */
 typedef void (*task_table_get_callback)(Task *task, void *user_context);
 
+/* Callback called when a task table test-and-update operation completes. If
+ * the task ID was not in the task table, then the task pointer will be NULL.
+ * If the update succeeded, the updated field will be set to true. */
+typedef void (*task_table_test_and_update_callback)(Task *task,
+                                                    void *user_context,
+                                                    bool updated);
+
 /**
  * Get a task's entry from the task table.
  *
@@ -107,13 +114,14 @@ void task_table_update(DBHandle *db_handle,
  *        fail_callback.
  * @return Void.
  */
-void task_table_test_and_update(DBHandle *db_handle,
-                                TaskID task_id,
-                                int test_state_bitmask,
-                                int update_state,
-                                RetryInfo *retry,
-                                task_table_get_callback done_callback,
-                                void *user_context);
+void task_table_test_and_update(
+    DBHandle *db_handle,
+    TaskID task_id,
+    int test_state_bitmask,
+    int update_state,
+    RetryInfo *retry,
+    task_table_test_and_update_callback done_callback,
+    void *user_context);
 
 /* Data that is needed to test and set the task's scheduling state. */
 typedef struct {