Prototype distributed actor handles (#1137)

* Add actor handle ID to the task spec

* Local scheduler dispatches actor tasks according to a task counter per handle

* Fix python test

* Allow passing actor handles into tasks. Not completely working yet. Also this is very messy.

* Fixes, should be roughly working now.

* Refactor actor handle wrapper

* Fix __init__ tests

* Terminate actor when the original handle goes out of scope

* TODO and a couple test cases

* Make tests for unsupported cases

* Fix Python mode tests

* Linting.

* Cache actor definitions that occur before ray.init() is called.

* Fix export actor class

* Deterministically compute actor handle ID

* Fix __getattribute__

* Fix string encoding for python3

* doc

* Add comment and assertion.

src/common/format/common.fbs

@@ -35,6 +35,9 @@ table TaskInfo {
   // Actor ID of the task. This is the actor that this task is executed on
   // or NIL_ACTOR_ID if the task is just a normal task.
   actor_id: string;
+  // The ID of the handle that was used to submit the task. This should be
+  // unique across handles with the same actor_id.
+  actor_handle_id: string;
   // Number of tasks that have been submitted to this actor so far.
   actor_counter: int;
   // True if this task is an actor checkpoint task and false otherwise.

src/common/lib/python/common_extension.cc

@@ -271,6 +271,8 @@ static int PyTask_init(PyTask *self, PyObject *args, PyObject *kwds) {
   UniqueID driver_id;
   /* ID of the actor this task should run on. */
   UniqueID actor_id = NIL_ACTOR_ID;
+  /* ID of the actor handle used to submit this task. */
+  UniqueID actor_handle_id = NIL_ACTOR_ID;
   /* How many tasks have been launched on the actor so far? */
   int actor_counter = 0;
   /* True if this is an actor checkpoint task and false otherwise. */
@@ -287,12 +289,13 @@ static int PyTask_init(PyTask *self, PyObject *args, PyObject *kwds) {
   int parent_counter;
   /* Resource vector of the required resources to execute this task. */
   PyObject *resource_vector = NULL;
-  if (!PyArg_ParseTuple(args, "O&O&OiO&i|O&iOO", &PyObjectToUniqueID,
+  if (!PyArg_ParseTuple(args, "O&O&OiO&i|O&O&iOO", &PyObjectToUniqueID,
                         &driver_id, &PyObjectToUniqueID, &function_id,
                         &arguments, &num_returns, &PyObjectToUniqueID,
                         &parent_task_id, &parent_counter, &PyObjectToUniqueID,
-                        &actor_id, &actor_counter,
-                        &is_actor_checkpoint_method_object, &resource_vector)) {
+                        &actor_id, &PyObjectToUniqueID, &actor_handle_id,
+                        &actor_counter, &is_actor_checkpoint_method_object,
+                        &resource_vector)) {
     return -1;
   }
 
@@ -304,9 +307,10 @@ static int PyTask_init(PyTask *self, PyObject *args, PyObject *kwds) {
 
   Py_ssize_t size = PyList_Size(arguments);
   /* Construct the task specification. */
-  TaskSpec_start_construct(
-      g_task_builder, driver_id, parent_task_id, parent_counter, actor_id,
-      actor_counter, is_actor_checkpoint_method, function_id, num_returns);
+  TaskSpec_start_construct(g_task_builder, driver_id, parent_task_id,
+                           parent_counter, actor_id, actor_handle_id,
+                           actor_counter, is_actor_checkpoint_method,
+                           function_id, num_returns);
   /* Add the task arguments. */
   for (Py_ssize_t i = 0; i < size; ++i) {
     PyObject *arg = PyList_GetItem(arguments, i);

src/common/task.cc

@@ -38,6 +38,7 @@ class TaskBuilder {
              TaskID parent_task_id,
              int64_t parent_counter,
              ActorID actor_id,
+             ActorID actor_handle_id,
              int64_t actor_counter,
              bool is_actor_checkpoint_method,
              FunctionID function_id,
@@ -46,6 +47,7 @@ class TaskBuilder {
     parent_task_id_ = parent_task_id;
     parent_counter_ = parent_counter;
     actor_id_ = actor_id;
+    actor_handle_id_ = actor_handle_id;
     actor_counter_ = actor_counter;
     is_actor_checkpoint_method_ = is_actor_checkpoint_method;
     function_id_ = function_id;
@@ -107,7 +109,8 @@ class TaskBuilder {
     auto message = CreateTaskInfo(
         fbb, to_flatbuf(fbb, driver_id_), to_flatbuf(fbb, task_id),
         to_flatbuf(fbb, parent_task_id_), parent_counter_,
-        to_flatbuf(fbb, actor_id_), actor_counter_, is_actor_checkpoint_method_,
+        to_flatbuf(fbb, actor_id_), to_flatbuf(fbb, actor_handle_id_),
+        actor_counter_, is_actor_checkpoint_method_,
         to_flatbuf(fbb, function_id_), arguments, fbb.CreateVector(returns),
         fbb.CreateVector(resource_vector_));
     /* Finish the TaskInfo. */
@@ -130,6 +133,7 @@ class TaskBuilder {
   TaskID parent_task_id_;
   int64_t parent_counter_;
   ActorID actor_id_;
+  ActorID actor_handle_id_;
   int64_t actor_counter_;
   bool is_actor_checkpoint_method_;
   FunctionID function_id_;
@@ -172,13 +176,14 @@ void TaskSpec_start_construct(TaskBuilder *builder,
                               TaskID parent_task_id,
                               int64_t parent_counter,
                               ActorID actor_id,
+                              ActorID actor_handle_id,
                               int64_t actor_counter,
                               bool is_actor_checkpoint_method,
                               FunctionID function_id,
                               int64_t num_returns) {
   builder->Start(driver_id, parent_task_id, parent_counter, actor_id,
-                 actor_counter, is_actor_checkpoint_method, function_id,
-                 num_returns);
+                 actor_handle_id, actor_counter, is_actor_checkpoint_method,
+                 function_id, num_returns);
 }
 
 uint8_t *TaskSpec_finish_construct(TaskBuilder *builder, int64_t *size) {
@@ -221,6 +226,12 @@ ActorID TaskSpec_actor_id(TaskSpec *spec) {
   return from_flatbuf(message->actor_id());
 }
 
+ActorID TaskSpec_actor_handle_id(TaskSpec *spec) {
+  CHECK(spec);
+  auto message = flatbuffers::GetRoot<TaskInfo>(spec);
+  return from_flatbuf(message->actor_handle_id());
+}
+
 bool TaskSpec_is_actor_task(TaskSpec *spec) {
   return !ActorID_equal(TaskSpec_actor_id(spec), NIL_ACTOR_ID);
 }

src/common/task.h

@@ -86,6 +86,9 @@ void free_task_builder(TaskBuilder *builder);
  *        the parent task prior to this one.
  * @param actor_id The ID of the actor that this task is for. If it is not an
  *        actor task, then this if NIL_ACTOR_ID.
+ * @param actor_handle_id The ID of the actor handle that this task was
+ *        submitted through. If it is not an actor task, or if this is the
+ *        original handle, then this is NIL_ACTOR_ID.
  * @param actor_counter A counter indicating how many tasks have been submitted
  *        to the same actor before this one.
  * @param is_actor_checkpoint_method True if this is an actor checkpoint method
@@ -102,6 +105,7 @@ void TaskSpec_start_construct(TaskBuilder *B,
                               TaskID parent_task_id,
                               int64_t parent_counter,
                               UniqueID actor_id,
+                              UniqueID actor_handle_id,
                               int64_t actor_counter,
                               bool is_actor_checkpoint_method,
                               FunctionID function_id,
@@ -133,6 +137,14 @@ FunctionID TaskSpec_function(TaskSpec *spec);
  */
 UniqueID TaskSpec_actor_id(TaskSpec *spec);
 
+/**
+ * Return the actor handle ID of the task.
+ *
+ * @param spec The task_spec in question.
+ * @return The ID of the actor handle that the task was submitted through.
+ */
+UniqueID TaskSpec_actor_handle_id(TaskSpec *spec);
+
 /**
  * Return whether this task is for an actor.
  *

src/common/test/example_task.h

@@ -14,7 +14,8 @@ static inline TaskSpec *example_task_spec_with_args(int64_t num_args,
   TaskID parent_task_id = globally_unique_id();
   FunctionID func_id = globally_unique_id();
   TaskSpec_start_construct(g_task_builder, NIL_ID, parent_task_id, 0,
-                           NIL_ACTOR_ID, 0, false, func_id, num_returns);
+                           NIL_ACTOR_ID, NIL_ACTOR_ID, 0, false, func_id,
+                           num_returns);
   for (int64_t i = 0; i < num_args; ++i) {
     ObjectID arg_id;
     if (arg_ids == NULL) {

src/common/test/task_tests.cc

@@ -15,8 +15,8 @@ TEST task_test(void) {
   TaskID parent_task_id = globally_unique_id();
   FunctionID func_id = globally_unique_id();
   TaskBuilder *builder = make_task_builder();
-  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0,
-                           false, func_id, 2);
+  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 0, NIL_ACTOR_ID,
+                           NIL_ACTOR_ID, 0, false, func_id, 2);
 
   UniqueID arg1 = globally_unique_id();
   TaskSpec_args_add_ref(builder, arg1);
@@ -54,16 +54,16 @@ TEST deterministic_ids_test(void) {
   uint8_t *arg2 = (uint8_t *) "hello world";
 
   /* Construct a first task. */
-  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0,
-                           false, func_id, 3);
+  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 0, NIL_ACTOR_ID,
+                           NIL_ACTOR_ID, 0, false, func_id, 3);
   TaskSpec_args_add_ref(builder, arg1);
   TaskSpec_args_add_val(builder, arg2, 11);
   int64_t size1;
   TaskSpec *spec1 = TaskSpec_finish_construct(builder, &size1);
 
   /* Construct a second identical task. */
-  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0,
-                           false, func_id, 3);
+  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 0, NIL_ACTOR_ID,
+                           NIL_ACTOR_ID, 0, false, func_id, 3);
   TaskSpec_args_add_ref(builder, arg1);
   TaskSpec_args_add_val(builder, arg2, 11);
   int64_t size2;
@@ -83,39 +83,39 @@ TEST deterministic_ids_test(void) {
 
   /* Construct a task with a different parent task ID. */
   TaskSpec_start_construct(builder, NIL_ID, globally_unique_id(), 0,
-                           NIL_ACTOR_ID, 0, false, func_id, 3);
+                           NIL_ACTOR_ID, NIL_ACTOR_ID, 0, false, func_id, 3);
   TaskSpec_args_add_ref(builder, arg1);
   TaskSpec_args_add_val(builder, arg2, 11);
   int64_t size3;
   TaskSpec *spec3 = TaskSpec_finish_construct(builder, &size3);
 
   /* Construct a task with a different parent counter. */
-  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 1, NIL_ACTOR_ID, 0,
-                           false, func_id, 3);
+  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 1, NIL_ACTOR_ID,
+                           NIL_ACTOR_ID, 0, false, func_id, 3);
   TaskSpec_args_add_ref(builder, arg1);
   TaskSpec_args_add_val(builder, arg2, 11);
   int64_t size4;
   TaskSpec *spec4 = TaskSpec_finish_construct(builder, &size4);
 
   /* Construct a task with a different function ID. */
-  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0,
-                           false, globally_unique_id(), 3);
+  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 0, NIL_ACTOR_ID,
+                           NIL_ACTOR_ID, 0, false, globally_unique_id(), 3);
   TaskSpec_args_add_ref(builder, arg1);
   TaskSpec_args_add_val(builder, arg2, 11);
   int64_t size5;
   TaskSpec *spec5 = TaskSpec_finish_construct(builder, &size5);
 
   /* Construct a task with a different object ID argument. */
-  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0,
-                           false, func_id, 3);
+  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 0, NIL_ACTOR_ID,
+                           NIL_ACTOR_ID, 0, false, func_id, 3);
   TaskSpec_args_add_ref(builder, globally_unique_id());
   TaskSpec_args_add_val(builder, arg2, 11);
   int64_t size6;
   TaskSpec *spec6 = TaskSpec_finish_construct(builder, &size6);
 
   /* Construct a task with a different value argument. */
-  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0,
-                           false, func_id, 3);
+  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 0, NIL_ACTOR_ID,
+                           NIL_ACTOR_ID, 0, false, func_id, 3);
   TaskSpec_args_add_ref(builder, arg1);
   TaskSpec_args_add_val(builder, (uint8_t *) "hello_world", 11);
   int64_t size7;
@@ -159,8 +159,8 @@ TEST send_task(void) {
   TaskBuilder *builder = make_task_builder();
   TaskID parent_task_id = globally_unique_id();
   FunctionID func_id = globally_unique_id();
-  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0,
-                           false, func_id, 2);
+  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 0, NIL_ACTOR_ID,
+                           NIL_ACTOR_ID, 0, false, func_id, 2);
   TaskSpec_args_add_ref(builder, globally_unique_id());
   TaskSpec_args_add_val(builder, (uint8_t *) "Hello", 5);
   TaskSpec_args_add_val(builder, (uint8_t *) "World", 5);

src/local_scheduler/local_scheduler_algorithm.cc

@@ -51,16 +51,21 @@ struct ObjectEntry {
 /** This struct contains information about a specific actor. This struct will be
  *  used inside of a hash table. */
 typedef struct {
-  /** The number of tasks that have been executed on this actor so far. This is
-   *  used to guarantee the in-order execution of tasks on actors (in the order
-   *  that the tasks were submitted). This is currently meaningful because we
-   *  restrict the submission of tasks on actors to the process that created the
-   *  actor. */
-  int64_t task_counter;
+  /** The number of tasks that have been executed on this actor so far, per
+   *  handle. This is used to guarantee execution of tasks on actors in the
+   *  order that the tasks were submitted, per handle. Tasks from different
+   *  handles to the same actor may be interleaved. */
+  std::unordered_map<ActorID, int64_t, UniqueIDHasher> task_counters;
   /** 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. */
+   *  assigned task execution, then the corresponding task_counter should be
+   *  updated to this value. */
   int64_t assigned_task_counter;
+  /** The handle that the currently assigned task was submitted by. This field
+   *  is only valid if assigned_task_counter is set. If the actor process
+   *  reports back success for the assigned task execution, then the
+   *  task_counter corresponding to this handle should be updated. */
+  ActorID assigned_task_handle_id;
   /** Whether the actor process has loaded yet. The actor counts as loaded once
    *  it has either executed its first task or successfully resumed from a
    *  checkpoint. Before the actor has loaded, we may dispatch the first task
@@ -247,8 +252,9 @@ void create_actor(SchedulingAlgorithmState *algorithm_state,
                   ActorID actor_id,
                   LocalSchedulerClient *worker) {
   LocalActorInfo entry;
-  entry.task_counter = 0;
+  entry.task_counters[NIL_ACTOR_ID] = 0;
   entry.assigned_task_counter = -1;
+  entry.assigned_task_handle_id = NIL_ACTOR_ID;
   entry.task_queue = new std::list<TaskQueueEntry>();
   entry.worker = worker;
   entry.worker_available = false;
@@ -333,16 +339,17 @@ bool dispatch_actor_task(LocalSchedulerState *state,
   /* Check whether we can execute the first task in the queue. */
   auto task = entry.task_queue->begin();
   int64_t next_task_counter = TaskSpec_actor_counter(task->spec);
+  ActorID next_task_handle_id = TaskSpec_actor_handle_id(task->spec);
   if (entry.loaded) {
     /* Once the actor has loaded, we can only execute tasks in order of
      * task_counter. */
-    if (next_task_counter != entry.task_counter) {
+    if (next_task_counter != entry.task_counters[next_task_handle_id]) {
       return false;
     }
   } else {
     /* If the actor has not yet loaded, we can only execute the task that
      * matches task_counter (the first task), or a checkpoint task. */
-    if (next_task_counter != entry.task_counter) {
+    if (next_task_counter != entry.task_counters[next_task_handle_id]) {
       /* No other task should be first in the queue. */
       CHECK(TaskSpec_is_actor_checkpoint_method(task->spec));
     }
@@ -361,6 +368,7 @@ bool dispatch_actor_task(LocalSchedulerState *state,
    * as unavailable. */
   assign_task_to_worker(state, task->spec, task->task_spec_size, entry.worker);
   entry.assigned_task_counter = next_task_counter;
+  entry.assigned_task_handle_id = next_task_handle_id;
   entry.worker_available = false;
   /* Free the task queue entry. */
   TaskQueueEntry_free(&(*task));
@@ -407,6 +415,8 @@ void insert_actor_task_queue(LocalSchedulerState *state,
                              TaskQueueEntry task_entry) {
   /* Get the local actor entry for this actor. */
   ActorID actor_id = TaskSpec_actor_id(task_entry.spec);
+  ActorID task_handle_id = TaskSpec_actor_handle_id(task_entry.spec);
+  int64_t task_counter = TaskSpec_actor_counter(task_entry.spec);
 
   /* Handle the case in which there is no LocalActorInfo struct yet. */
   if (algorithm_state->local_actor_infos.count(actor_id) == 0) {
@@ -418,40 +428,43 @@ void insert_actor_task_queue(LocalSchedulerState *state,
   }
   LocalActorInfo &entry =
       algorithm_state->local_actor_infos.find(actor_id)->second;
+  if (entry.task_counters.count(task_handle_id) == 0) {
+    entry.task_counters[task_handle_id] = 0;
+  }
 
-  int64_t task_counter = TaskSpec_actor_counter(task_entry.spec);
   /* As a sanity check, the counter of the new task should be greater than the
    * number of tasks that have executed on this actor so far (since we are
    * 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. */
-  if (task_counter < entry.task_counter) {
+  if (task_counter < entry.task_counters[task_handle_id]) {
     LOG_INFO(
         "A task that has already been executed has been resubmitted, so we "
         "are ignoring it. This should only happen during reconstruction.");
     return;
   }
 
-  /* Add the task spec to the actor's task queue in a manner that preserves the
-   * order of the actor task counters. Iterate from the beginning of the queue
-   * to find the right place to insert the task queue entry. TODO(pcm): This
-   * makes submitting multiple actor tasks take quadratic time, which needs to
-   * be optimized. */
+  /* Insert the task spec to the actor's task queue in sorted order, per actor
+   * handle ID. Find the first task in the queue with a counter greater than
+   * the submitted task's and the same handle ID. */
   auto it = entry.task_queue->begin();
-  while (it != entry.task_queue->end() &&
-         (task_counter > TaskSpec_actor_counter(it->spec))) {
-    ++it;
+  for (; it != entry.task_queue->end(); it++) {
+    /* Skip tasks submitted by a different handle. */
+    if (!ActorID_equal(task_handle_id, TaskSpec_actor_handle_id(it->spec))) {
+      continue;
+    }
+    /* A duplicate task submitted by the same handle. */
+    if (task_counter == TaskSpec_actor_counter(it->spec)) {
+      LOG_INFO(
+          "A task was resubmitted, so we are ignoring it. This should only "
+          "happen during reconstruction.");
+      return;
+    }
+    /* We found a task with the same handle ID and a greater task counter. */
+    if (task_counter < TaskSpec_actor_counter(it->spec)) {
+      break;
+    }
   }
-  if (it != entry.task_queue->end() &&
-      task_counter == TaskSpec_actor_counter(it->spec)) {
-    LOG_INFO(
-        "A task was resubmitted, so we are ignoring it. This should only "
-        "happen during reconstruction.");
-    return;
-  }
-
-  /* The task has a counter that has not been executed or submitted before. Add
-   * it to the actor queue. */
   entry.task_queue->insert(it, task_entry);
 
   /* Record the fact that this actor has a task waiting to execute. */
@@ -1266,7 +1279,8 @@ void handle_actor_worker_available(LocalSchedulerState *state,
    * 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.task_counters[entry.assigned_task_handle_id] =
+        entry.assigned_task_counter + 1;
     /* If a task was assigned to this actor and there was no checkpoint
      * failure, then it is now loaded. */
     if (entry.assigned_task_counter > -1) {
@@ -1274,6 +1288,7 @@ void handle_actor_worker_available(LocalSchedulerState *state,
     }
   }
   entry.assigned_task_counter = -1;
+  entry.assigned_task_handle_id = NIL_ACTOR_ID;
   entry.worker_available = true;
   /* Assign new tasks if possible. */
   dispatch_all_tasks(state, algorithm_state);