Actor checkpointing with object lineage reconstruction (#1004)

* Worker reports error in previous task, actor task counter is incremented after task is successful

* Refactor actor task execution

- Return new task counter in GetTaskRequest
- Update worker state for actor tasks inside of the actor method
  executor

* Manually invoked checkpoint method

* Scheduling for actor checkpoint methods

* Fix python bugs in checkpointing

* Return task success from worker to local scheduler instead of actor counter

* Kill local schedulers halfway through actor execution instead of waiting for all tasks to execute once

* Remove redundant actor tasks during dispatch, reconstruct missing dependencies for actor tasks

* Make executor for temporary actor methods

* doc

* Set default argument for whether the previous task was a success

* Refactor actor method call

* Simplify checkpoint task submission

* lint

* fix philipp's comments

* Add missing line

* Make actor reconstruction tests run faster

* Unimportant whitespace.

* Unimportant whitespace.

* Update checkpoint method signature

* Documentation and handle exceptions during checkpoint save/resume

* Rename get_task message field to actor_checkpoint_failed

* Fix bug.

* Remove debugging check, redirect test output

src/common/task.cc

@@ -217,7 +217,14 @@ ActorID TaskSpec_actor_id(TaskSpec *spec) {
 int64_t TaskSpec_actor_counter(TaskSpec *spec) {
   CHECK(spec);
   auto message = flatbuffers::GetRoot<TaskInfo>(spec);
-  return message->actor_counter();
+  return std::abs(message->actor_counter());
+}
+
+bool TaskSpec_actor_is_checkpoint_method(TaskSpec *spec) {
+  CHECK(spec);
+  auto message = flatbuffers::GetRoot<TaskInfo>(spec);
+  int64_t actor_counter = message->actor_counter();
+  return actor_counter < 0;
 }
 
 UniqueID TaskSpec_driver_id(TaskSpec *spec) {

src/common/task.h

@@ -135,6 +135,8 @@ UniqueID TaskSpec_actor_id(TaskSpec *spec);
  */
 int64_t TaskSpec_actor_counter(TaskSpec *spec);
 
+bool TaskSpec_actor_is_checkpoint_method(TaskSpec *spec);
+
 /**
  * Return the driver ID of the task.
  *

src/local_scheduler/format/local_scheduler.fbs

@@ -36,6 +36,14 @@ enum MessageType:int {
   PutObject
 }
 
+// This message is sent from a worker to a local scheduler.
+table GetTaskRequest {
+  // Whether the previously assigned task was a checkpoint task that failed.
+  // If true, then the local scheduler will not update the actor's task
+  // counter to match the assigned checkpoint index.
+  actor_checkpoint_failed: bool;
+}
+
 // This message is sent from the local scheduler to a worker.
 table GetTaskReply {
   // A string of bytes representing the task specification.

src/local_scheduler/local_scheduler.cc

@@ -568,7 +568,9 @@ void assign_task_to_worker(LocalSchedulerState *state,
   }
 }
 
-void finish_task(LocalSchedulerState *state, LocalSchedulerClient *worker) {
+void finish_task(LocalSchedulerState *state,
+                 LocalSchedulerClient *worker,
+                 bool actor_checkpoint_failed) {
   if (worker->task_in_progress != NULL) {
     TaskSpec *spec = Task_task_spec(worker->task_in_progress);
     /* Return dynamic resources back for the task in progress. */
@@ -589,8 +591,11 @@ void finish_task(LocalSchedulerState *state, LocalSchedulerClient *worker) {
     }
     /* If we're connected to Redis, update tables. */
     if (state->db != NULL) {
-      /* Update control state tables. */
-      Task_set_state(worker->task_in_progress, TASK_STATUS_DONE);
+      /* Update control state tables. If there was an error while executing a *
+       * checkpoint task, report the task as lost. Else, the task succeeded. */
+      int task_state =
+          actor_checkpoint_failed ? TASK_STATUS_LOST : TASK_STATUS_DONE;
+      Task_set_state(worker->task_in_progress, task_state);
       task_table_update(state->db, worker->task_in_progress, NULL, NULL, NULL);
       /* The call to task_table_update takes ownership of the
        * task_in_progress, so we set the pointer to NULL so it is not used. */
@@ -734,18 +739,18 @@ void reconstruct_object_lookup_callback(
    * object table entry is up-to-date. */
   LocalSchedulerState *state = (LocalSchedulerState *) user_context;
   /* Look up the task that created the object in the result table. */
-  if (!never_created && manager_vector.size() == 0) {
-    /* If the object was created and later evicted, we reconstruct the object
-     * if and only if there are no other instances of the task running. */
-    result_table_lookup(state->db, reconstruct_object_id, NULL,
-                        reconstruct_evicted_result_lookup_callback,
-                        (void *) state);
-  } else if (never_created) {
+  if (never_created) {
     /* If the object has not been created yet, we reconstruct the object if and
      * only if the task that created the object failed to complete. */
     result_table_lookup(state->db, reconstruct_object_id, NULL,
                         reconstruct_failed_result_lookup_callback,
                         (void *) state);
+  } else if (manager_vector.size() == 0) {
+    /* If the object was created and later evicted, we reconstruct the object
+     * if and only if there are no other instances of the task running. */
+    result_table_lookup(state->db, reconstruct_object_id, NULL,
+                        reconstruct_evicted_result_lookup_callback,
+                        (void *) state);
   }
 }
 
@@ -951,7 +956,7 @@ void process_message(event_loop *loop,
   case MessageType_TaskDone: {
   } break;
   case MessageType_DisconnectClient: {
-    finish_task(state, worker);
+    finish_task(state, worker, false);
     CHECK(!worker->disconnected);
     worker->disconnected = true;
     /* If the disconnected worker was not an actor, start a new worker to make
@@ -977,13 +982,16 @@ void process_message(event_loop *loop,
   } break;
   case MessageType_GetTask: {
     /* If this worker reports a completed task, account for resources. */
-    finish_task(state, worker);
+    auto message = flatbuffers::GetRoot<GetTaskRequest>(input);
+    bool actor_checkpoint_failed = message->actor_checkpoint_failed();
+    finish_task(state, worker, actor_checkpoint_failed);
     /* Let the scheduling algorithm process the fact that there is an available
      * worker. */
     if (ActorID_equal(worker->actor_id, NIL_ACTOR_ID)) {
       handle_worker_available(state, state->algorithm_state, worker);
     } else {
-      handle_actor_worker_available(state, state->algorithm_state, worker);
+      handle_actor_worker_available(state, state->algorithm_state, worker,
+                                    actor_checkpoint_failed);
     }
   } break;
   case MessageType_ReconstructObject: {

src/local_scheduler/local_scheduler.h

@@ -56,9 +56,13 @@ void assign_task_to_worker(LocalSchedulerState *state,
  *
  * @param state The local scheduler state.
  * @param worker The worker that finished the task.
+ * @param actor_checkpoint_failed If the last task assigned was a checkpoint
+ *        task that failed.
  * @return Void.
  */
-void finish_task(LocalSchedulerState *state, LocalSchedulerClient *worker);
+void finish_task(LocalSchedulerState *state,
+                 LocalSchedulerClient *worker,
+                 bool actor_checkpoint_failed);
 
 /**
  * This is the callback that is used to process a notification from the Plasma

src/local_scheduler/local_scheduler_algorithm.cc

@@ -53,6 +53,10 @@ typedef struct {
    *  restrict the submission of tasks on actors to the process that created the
    *  actor. */
   int64_t task_counter;
+  /** The index of the task assigned to this actor. Set to -1 if no task is
+   *  currently assigned. If the actor process reports back success for the
+   *  assigned task execution, task_counter should be set to this value. */
+  int64_t assigned_task_counter;
   /** A queue of tasks to be executed on this actor. The tasks will be sorted by
    *  the order of their actor counters. */
   std::list<TaskQueueEntry> *task_queue;
@@ -234,6 +238,7 @@ void create_actor(SchedulingAlgorithmState *algorithm_state,
                   LocalSchedulerClient *worker) {
   LocalActorInfo entry;
   entry.task_counter = 0;
+  entry.assigned_task_counter = -1;
   entry.task_queue = new std::list<TaskQueueEntry>();
   entry.worker = worker;
   entry.worker_available = false;
@@ -309,39 +314,56 @@ bool dispatch_actor_task(LocalSchedulerState *state,
 
   /* There should be some queued tasks for this actor. */
   CHECK(!entry.task_queue->empty());
-
-  TaskQueueEntry first_task = entry.task_queue->front();
-  int64_t next_task_counter = TaskSpec_actor_counter(first_task.spec);
-  if (next_task_counter != entry.task_counter) {
-    /* We cannot execute the next task on this actor without violating the
-     * in-order execution guarantee for actor tasks. */
-    CHECK(next_task_counter > entry.task_counter);
-    return false;
-  }
   /* If the worker is not available, we cannot assign a task to it. */
   if (!entry.worker_available) {
     return false;
   }
+
+  /* Find the first task that either matches the task counter or that is a
+   * checkpoint method. Remove any tasks that we have already executed past
+   * (e.g., by executing a more recent checkpoint method). */
+  auto task = entry.task_queue->begin();
+  int64_t next_task_counter = TaskSpec_actor_counter(task->spec);
+  while (next_task_counter != entry.task_counter) {
+    if (next_task_counter < entry.task_counter) {
+      /* A task that we have already executed past. Remove it. */
+      task = entry.task_queue->erase(task);
+      /* If there are no more tasks in the queue, wait. */
+      if (task == entry.task_queue->end()) {
+        algorithm_state->actors_with_pending_tasks.erase(actor_id);
+        return false;
+      }
+      /* Move on to the next task. */
+      next_task_counter = TaskSpec_actor_counter(task->spec);
+    } else if (TaskSpec_actor_is_checkpoint_method(task->spec)) {
+      /* A later task that is a checkpoint method. Checkpoint methods can
+       * always be executed. */
+      break;
+    } else {
+      /* A later task that is not a checkpoint. Wait for the preceding tasks to
+       * execute. */
+      return false;
+    }
+  }
+
   /* If there are not enough resources available, we cannot assign the task. */
-  CHECK(0 ==
-        TaskSpec_get_required_resource(first_task.spec, ResourceIndex_GPU));
+  CHECK(0 == TaskSpec_get_required_resource(task->spec, ResourceIndex_GPU));
   if (!check_dynamic_resources(
-          state,
-          TaskSpec_get_required_resource(first_task.spec, ResourceIndex_CPU), 0,
-          TaskSpec_get_required_resource(first_task.spec,
-                                         ResourceIndex_CustomResource))) {
+          state, TaskSpec_get_required_resource(task->spec, ResourceIndex_CPU),
+          0, TaskSpec_get_required_resource(task->spec,
+                                            ResourceIndex_CustomResource))) {
     return false;
   }
+
   /* Assign the first task in the task queue to the worker and mark the worker
    * as unavailable. */
-  entry.task_counter += 1;
-  assign_task_to_worker(state, first_task.spec, first_task.task_spec_size,
-                        entry.worker);
+  assign_task_to_worker(state, task->spec, task->task_spec_size, entry.worker);
+  entry.assigned_task_counter = next_task_counter;
   entry.worker_available = false;
   /* Free the task queue entry. */
-  TaskQueueEntry_free(&first_task);
+  TaskQueueEntry_free(&(*task));
   /* Remove the task from the actor's task queue. */
-  entry.task_queue->pop_front();
+  entry.task_queue->erase(task);
 
   /* If there are no more tasks in the queue, then indicate that the actor has
    * no tasks. */
@@ -414,12 +436,11 @@ void add_task_to_actor_queue(LocalSchedulerState *state,
    * guaranteeing in-order execution of the tasks on the actor). TODO(rkn): This
    * check will fail if the fault-tolerance mechanism resubmits a task on an
    * actor. */
-  bool task_is_redundant = false;
   if (task_counter < entry.task_counter) {
     LOG_INFO(
         "A task that has already been executed has been resubmitted, so we "
         "are ignoring it. This should only happen during reconstruction.");
-    task_is_redundant = true;
+    return;
   }
 
   /* Create a new task queue entry. */
@@ -437,32 +458,51 @@ void add_task_to_actor_queue(LocalSchedulerState *state,
   if (it != entry.task_queue->end() &&
       task_counter == TaskSpec_actor_counter(it->spec)) {
     LOG_INFO(
-        "A task that has already been executed has been resubmitted, so we "
-        "are ignoring it. This should only happen during reconstruction.");
-    task_is_redundant = true;
+        "A task was resubmitted, so we are ignoring it. This should only "
+        "happen during reconstruction.");
+    return;
   }
 
-  if (!task_is_redundant) {
-    entry.task_queue->insert(it, elt);
+  /* The task has a counter that has not been executed or submitted before. Add
+   * it to the actor queue. */
+  entry.task_queue->insert(it, elt);
 
-    /* Update the task table. */
-    if (state->db != NULL) {
-      Task *task = Task_alloc(spec, task_spec_size, TASK_STATUS_QUEUED,
-                              get_db_client_id(state->db));
-      if (from_global_scheduler) {
-        /* If the task is from the global scheduler, it's already been added to
-         * the task table, so just update the entry. */
-        task_table_update(state->db, task, NULL, NULL, NULL);
-      } else {
-        /* Otherwise, this is the first time the task has been seen in the
-         * system (unless it's a resubmission of a previous task), so add the
-         * entry. */
-        task_table_add_task(state->db, task, NULL, NULL, NULL);
+  /* Update the task table. */
+  if (state->db != NULL) {
+    Task *task = Task_alloc(spec, task_spec_size, TASK_STATUS_QUEUED,
+                            get_db_client_id(state->db));
+    if (from_global_scheduler) {
+      /* If the task is from the global scheduler, it's already been added to
+       * the task table, so just update the entry. */
+      task_table_update(state->db, task, NULL, NULL, NULL);
+    } else {
+      /* Otherwise, this is the first time the task has been seen in the
+       * system (unless it's a resubmission of a previous task), so add the
+       * entry. */
+      task_table_add_task(state->db, task, NULL, NULL, NULL);
+    }
+  }
+
+  /* Record the fact that this actor has a task waiting to execute. */
+  algorithm_state->actors_with_pending_tasks.insert(actor_id);
+
+  /* Register a missing dependency on the preceding task. TODO(swang): Unify
+   * with `fetch_missing_dependencies` for non-actor tasks. */
+  if (entry.task_counter != task_counter) {
+    int64_t num_args = TaskSpec_num_args(spec);
+    /* The last argument represents dependency on a preceding task. If it is by
+     * reference, then it is an explicit dependency. */
+    if (TaskSpec_arg_by_ref(spec, num_args - 1)) {
+      ObjectID dummy_object_id = TaskSpec_arg_id(spec, num_args - 1);
+      if (algorithm_state->local_objects.count(dummy_object_id) == 0) {
+        ObjectEntry entry;
+        /* TODO(swang): Objects in `remote_objects` will get fetched from
+         * remote plasma managers. Do not fetch actor dummy objects. Otherwise,
+         * if the plasma manager associated with the dead local scheduler is
+         * still alive, reconstruction will never complete. */
+        state->algorithm_state->remote_objects[dummy_object_id] = entry;
       }
     }
-
-    /* Record the fact that this actor has a task waiting to execute. */
-    algorithm_state->actors_with_pending_tasks.insert(actor_id);
   }
 }
 
@@ -1202,7 +1242,8 @@ void handle_actor_worker_disconnect(LocalSchedulerState *state,
 
 void handle_actor_worker_available(LocalSchedulerState *state,
                                    SchedulingAlgorithmState *algorithm_state,
-                                   LocalSchedulerClient *worker) {
+                                   LocalSchedulerClient *worker,
+                                   bool actor_checkpoint_failed) {
   ActorID actor_id = worker->actor_id;
   CHECK(!ActorID_equal(actor_id, NIL_ACTOR_ID));
   /* Get the actor info for this worker. */
@@ -1212,6 +1253,13 @@ void handle_actor_worker_available(LocalSchedulerState *state,
 
   CHECK(worker == entry.worker);
   CHECK(!entry.worker_available);
+  /* If the assigned task was not a checkpoint task, or if it was but it
+   * loaded the checkpoint successfully, then we update the actor's counter
+   * to the assigned counter. */
+  if (!actor_checkpoint_failed) {
+    entry.task_counter = entry.assigned_task_counter + 1;
+  }
+  entry.assigned_task_counter = -1;
   entry.worker_available = true;
   /* Assign new tasks if possible. */
   dispatch_all_tasks(state, algorithm_state);

src/local_scheduler/local_scheduler_algorithm.h

@@ -178,12 +178,15 @@ void handle_worker_removed(LocalSchedulerState *state,
  *
  * @param state The state of the local scheduler.
  * @param algorithm_state State maintained by the scheduling algorithm.
- * @param wi Information about the worker that is available.
+ * @param worker The worker that is available.
+ * @param actor_checkpoint_failed If the last task assigned was a checkpoint
+ *        task that failed.
  * @return Void.
  */
 void handle_actor_worker_available(LocalSchedulerState *state,
                                    SchedulingAlgorithmState *algorithm_state,
-                                   LocalSchedulerClient *worker);
+                                   LocalSchedulerClient *worker,
+                                   bool actor_checkpoint_failed);
 
 /**
  * Handle the fact that a new worker is available for running an actor.

src/local_scheduler/local_scheduler_client.cc

@@ -95,14 +95,19 @@ void local_scheduler_submit(LocalSchedulerConnection *conn,
 }
 
 TaskSpec *local_scheduler_get_task(LocalSchedulerConnection *conn,
-                                   int64_t *task_size) {
-  write_message(conn->conn, MessageType_GetTask, 0, NULL);
+                                   int64_t *task_size,
+                                   bool actor_checkpoint_failed) {
+  flatbuffers::FlatBufferBuilder fbb;
+  auto message = CreateGetTaskRequest(fbb, actor_checkpoint_failed);
+  fbb.Finish(message);
+  write_message(conn->conn, MessageType_GetTask, fbb.GetSize(),
+                fbb.GetBufferPointer());
   int64_t type;
-  int64_t message_size;
-  uint8_t *message;
+  int64_t reply_size;
+  uint8_t *reply;
   /* Receive a task from the local scheduler. This will block until the local
    * scheduler gives this client a task. */
-  read_message(conn->conn, &type, &message_size, &message);
+  read_message(conn->conn, &type, &reply_size, &reply);
   if (type == DISCONNECT_CLIENT) {
     LOG_WARN("Exiting because local scheduler closed connection.");
     exit(1);
@@ -110,7 +115,7 @@ TaskSpec *local_scheduler_get_task(LocalSchedulerConnection *conn,
   CHECK(type == MessageType_ExecuteTask);
 
   /* Parse the flatbuffer object. */
-  auto reply_message = flatbuffers::GetRoot<GetTaskReply>(message);
+  auto reply_message = flatbuffers::GetRoot<GetTaskReply>(reply);
 
   /* Set the GPU IDs for this task. We only do this for non-actor tasks because
    * for actors the GPUs are associated with the actor itself and not with the
@@ -127,7 +132,7 @@ TaskSpec *local_scheduler_get_task(LocalSchedulerConnection *conn,
   TaskSpec *data = (TaskSpec *) reply_message->task_spec()->data();
   TaskSpec *spec = TaskSpec_copy(data, *task_size);
   /* Free the original message from the local scheduler. */
-  free(message);
+  free(reply);
   /* Return the copy of the task spec and pass ownership to the caller. */
   return spec;
 }

src/local_scheduler/local_scheduler_client.h

@@ -94,10 +94,14 @@ void local_scheduler_log_event(LocalSchedulerConnection *conn,
  * @todo When does this actually get freed?
  *
  * @param conn The connection information.
+ * @param task_size A pointer to fill out with the task size.
+ * @param actor_checkpoint_failed If the last task assigned was a checkpoint
+ *        task that failed.
  * @return The address of the assigned task.
  */
 TaskSpec *local_scheduler_get_task(LocalSchedulerConnection *conn,
-                                   int64_t *task_size);
+                                   int64_t *task_size,
+                                   bool actor_checkpoint_failed);
 
 /**
  * Tell the local scheduler that the client has finished executing a task.

src/local_scheduler/local_scheduler_extension.cc

@@ -59,14 +59,25 @@ static PyObject *PyLocalSchedulerClient_submit(PyObject *self, PyObject *args) {
 }
 
 // clang-format off
-static PyObject *PyLocalSchedulerClient_get_task(PyObject *self) {
+static PyObject *PyLocalSchedulerClient_get_task(PyObject *self, PyObject *args) {
+  PyObject *py_actor_checkpoint_failed = NULL;
+  if (!PyArg_ParseTuple(args, "|O", &py_actor_checkpoint_failed)) {
+    return NULL;
+  }
   TaskSpec *task_spec;
+  int64_t task_size;
+  /* If no argument for actor_checkpoint_failed was provided, default to false,
+   * since we assume that there was no previous task. */
+  bool actor_checkpoint_failed = false;
+  if (py_actor_checkpoint_failed != NULL) {
+    actor_checkpoint_failed = (bool) PyObject_IsTrue(py_actor_checkpoint_failed);
+  }
   /* Drop the global interpreter lock while we get a task because
    * local_scheduler_get_task may block for a long time. */
-  int64_t task_size;
   Py_BEGIN_ALLOW_THREADS
   task_spec = local_scheduler_get_task(
-      ((PyLocalSchedulerClient *) self)->local_scheduler_connection, &task_size);
+      ((PyLocalSchedulerClient *) self)->local_scheduler_connection,
+      &task_size, actor_checkpoint_failed);
   Py_END_ALLOW_THREADS
   return PyTask_make(task_spec, task_size);
 }
@@ -138,7 +149,7 @@ static PyMethodDef PyLocalSchedulerClient_methods[] = {
      "Notify the local scheduler that this client is exiting gracefully."},
     {"submit", (PyCFunction) PyLocalSchedulerClient_submit, METH_VARARGS,
      "Submit a task to the local scheduler."},
-    {"get_task", (PyCFunction) PyLocalSchedulerClient_get_task, METH_NOARGS,
+    {"get_task", (PyCFunction) PyLocalSchedulerClient_get_task, METH_VARARGS,
      "Get a task from the local scheduler."},
     {"reconstruct_object",
      (PyCFunction) PyLocalSchedulerClient_reconstruct_object, METH_VARARGS,

src/local_scheduler/test/local_scheduler_tests.cc

@@ -210,12 +210,12 @@ TEST object_reconstruction_test(void) {
     int64_t task_assigned_size;
     local_scheduler_submit(worker, spec, task_size);
     TaskSpec *task_assigned =
-        local_scheduler_get_task(worker, &task_assigned_size);
+        local_scheduler_get_task(worker, &task_assigned_size, true);
     ASSERT_EQ(memcmp(task_assigned, spec, task_size), 0);
     ASSERT_EQ(task_assigned_size, task_size);
     int64_t reconstruct_task_size;
     TaskSpec *reconstruct_task =
-        local_scheduler_get_task(worker, &reconstruct_task_size);
+        local_scheduler_get_task(worker, &reconstruct_task_size, true);
     ASSERT_EQ(memcmp(reconstruct_task, spec, task_size), 0);
     ASSERT_EQ(reconstruct_task_size, task_size);
     /* Clean up. */
@@ -315,7 +315,8 @@ TEST object_reconstruction_recursive_test(void) {
     /* Make sure we receive each task from the initial submission. */
     for (int i = 0; i < NUM_TASKS; ++i) {
       int64_t task_size;
-      TaskSpec *task_assigned = local_scheduler_get_task(worker, &task_size);
+      TaskSpec *task_assigned =
+          local_scheduler_get_task(worker, &task_size, true);
       ASSERT_EQ(memcmp(task_assigned, specs[i], task_sizes[i]), 0);
       ASSERT_EQ(task_size, task_sizes[i]);
       free(task_assigned);
@@ -325,7 +326,7 @@ TEST object_reconstruction_recursive_test(void) {
     for (int i = 0; i < NUM_TASKS; ++i) {
       int64_t task_assigned_size;
       TaskSpec *task_assigned =
-          local_scheduler_get_task(worker, &task_assigned_size);
+          local_scheduler_get_task(worker, &task_assigned_size, true);
       bool found = false;
       for (int j = 0; j < NUM_TASKS; ++j) {
         if (specs[j] == NULL) {
@@ -410,7 +411,7 @@ TEST object_reconstruction_suppression_test(void) {
      * object_table_add callback completes. */
     int64_t task_assigned_size;
     TaskSpec *task_assigned =
-        local_scheduler_get_task(worker, &task_assigned_size);
+        local_scheduler_get_task(worker, &task_assigned_size, true);
     ASSERT_EQ(memcmp(task_assigned, object_reconstruction_suppression_spec,
                      object_reconstruction_suppression_size),
               0);