Clean up when a driver disconnects. (#462)

* Clean up state when drivers exit.

* Remove unnecessary field in ActorMapEntry struct.

* Have monitor release GPU resources in Redis when driver exits.

* Enable multiple drivers in multi-node tests and test driver cleanup.

* Make redis GPU allocation a redis transaction and small cleanups.

* Fix multi-node test.

* Small cleanups.

* Make global scheduler take node_ip_address so it appears in the right place in the client table.

* Cleanups.

* Fix linting and cleanups in local scheduler.

* Fix removed_driver_test.

* Fix bug related to vector -> list.

* Fix linting.

* Cleanup.

* Fix multi node tests.

* Fix jenkins tests.

* Add another multi node test with many drivers.

* Fix linting.

* Make the actor creation notification a flatbuffer message.

* Revert "Make the actor creation notification a flatbuffer message."

This reverts commit af99099c8084dbf9177fb4e34c0c9b1a12c78f39.

* Add comment explaining flatbuffer problems.

src/common/CMakeLists.txt

@@ -61,6 +61,7 @@ add_library(common STATIC
   state/object_table.cc
   state/task_table.cc
   state/db_client_table.cc
+  state/driver_table.cc
   state/actor_notification_table.cc
   state/local_scheduler_table.cc
   state/error_table.cc

src/common/common.cc

@@ -47,6 +47,10 @@ bool DBClientID_equal(DBClientID first_id, DBClientID second_id) {
   return UNIQUE_ID_EQ(first_id, second_id);
 }
 
+bool WorkerID_equal(WorkerID first_id, WorkerID second_id) {
+  return UNIQUE_ID_EQ(first_id, second_id);
+}
+
 char *ObjectID_to_string(ObjectID obj_id, char *id_string, int id_length) {
   CHECK(id_length >= ID_STRING_SIZE);
   static const char hex[] = "0123456789abcdef";

src/common/common.h

@@ -153,7 +153,9 @@ extern const UniqueID NIL_ID;
 UniqueID globally_unique_id(void);
 
 #define NIL_OBJECT_ID NIL_ID
+#define NIL_WORKER_ID NIL_ID
 
+/** The object ID is the type used to identify objects. */
 typedef UniqueID ObjectID;
 
 #ifdef __cplusplus
@@ -202,6 +204,18 @@ bool ObjectID_equal(ObjectID first_id, ObjectID second_id);
  */
 bool ObjectID_is_nil(ObjectID id);
 
+/** The worker ID is the ID of a worker or driver. */
+typedef UniqueID WorkerID;
+
+/**
+ * Compare two worker IDs.
+ *
+ * @param first_id The first worker ID to compare.
+ * @param second_id The first worker ID to compare.
+ * @return True if the worker IDs are the same and false otherwise.
+ */
+bool WorkerID_equal(WorkerID first_id, WorkerID second_id);
+
 typedef UniqueID DBClientID;
 
 /**

src/common/format/common.fbs

@@ -139,3 +139,8 @@ table ResultTableReply {
 }
 
 root_type ResultTableReply;
+
+table DriverTableMessage {
+  // The driver ID of the driver that died.
+  driver_id: string;
+}

src/common/state/actor_notification_table.h

@@ -5,20 +5,16 @@
 #include "db.h"
 #include "table.h"
 
-typedef struct {
-  /** The ID of the actor. */
-  ActorID actor_id;
-  /** The ID of the local scheduler that is responsible for the actor. */
-  DBClientID local_scheduler_id;
-} ActorInfo;
-
 /*
  *  ==== Subscribing to the actor notification table ====
  */
 
 /* Callback for subscribing to the local scheduler table. */
-typedef void (*actor_notification_table_subscribe_callback)(ActorInfo info,
-                                                            void *user_context);
+typedef void (*actor_notification_table_subscribe_callback)(
+    ActorID actor_id,
+    WorkerID driver_id,
+    DBClientID local_scheduler_id,
+    void *user_context);
 
 /**
  * Register a callback to process actor notification events.

src/common/state/driver_table.cc

@@ -0,0 +1,21 @@
+#include "driver_table.h"
+#include "redis.h"
+
+void driver_table_subscribe(DBHandle *db_handle,
+                            driver_table_subscribe_callback subscribe_callback,
+                            void *subscribe_context,
+                            RetryInfo *retry) {
+  DriverTableSubscribeData *sub_data =
+      (DriverTableSubscribeData *) malloc(sizeof(DriverTableSubscribeData));
+  sub_data->subscribe_callback = subscribe_callback;
+  sub_data->subscribe_context = subscribe_context;
+  init_table_callback(db_handle, NIL_ID, __func__, sub_data, retry, NULL,
+                      redis_driver_table_subscribe, NULL);
+}
+
+void driver_table_send_driver_death(DBHandle *db_handle,
+                                    WorkerID driver_id,
+                                    RetryInfo *retry) {
+  init_table_callback(db_handle, driver_id, __func__, NULL, retry, NULL,
+                      redis_driver_table_send_driver_death, NULL);
+}

src/common/state/driver_table.h

@@ -0,0 +1,50 @@
+#ifndef DRIVER_TABLE_H
+#define DRIVER_TABLE_H
+
+#include "db.h"
+#include "table.h"
+#include "task.h"
+
+/*
+ *  ==== Subscribing to the driver table ====
+ */
+
+/* Callback for subscribing to the driver table. */
+typedef void (*driver_table_subscribe_callback)(WorkerID driver_id,
+                                                void *user_context);
+
+/**
+ * Register a callback for a driver table event.
+ *
+ * @param db_handle Database handle.
+ * @param subscribe_callback Callback that will be called when the driver event
+ *        happens.
+ * @param subscribe_context Context that will be passed into the
+ *        subscribe_callback.
+ * @param retry Information about retrying the request to the database.
+ * @return Void.
+ */
+void driver_table_subscribe(DBHandle *db_handle,
+                            driver_table_subscribe_callback subscribe_callback,
+                            void *subscribe_context,
+                            RetryInfo *retry);
+
+/* Data that is needed to register driver table subscribe callbacks with the
+ * state database. */
+typedef struct {
+  driver_table_subscribe_callback subscribe_callback;
+  void *subscribe_context;
+} DriverTableSubscribeData;
+
+/**
+ * Send driver death update to all subscribers.
+ *
+ * @param db_handle Database handle.
+ * @param driver_id The ID of the driver that died.
+ * @param retry Information about retrying the request to the database.
+ */
+void driver_table_send_driver_death(DBHandle *db_handle,
+                                    WorkerID driver_id,
+                                    RetryInfo *retry);
+
+#endif /* DRIVER_TABLE_H */

src/common/state/redis.cc

@@ -17,6 +17,7 @@ extern "C" {
 #include "db.h"
 #include "db_client_table.h"
 #include "actor_notification_table.h"
+#include "driver_table.h"
 #include "local_scheduler_table.h"
 #include "object_table.h"
 #include "task.h"
@@ -1089,6 +1090,90 @@ void redis_local_scheduler_table_send_info(TableCallbackData *callback_data) {
   }
 }
 
+void redis_driver_table_subscribe_callback(redisAsyncContext *c,
+                                           void *r,
+                                           void *privdata) {
+  REDIS_CALLBACK_HEADER(db, callback_data, r);
+
+  redisReply *reply = (redisReply *) r;
+  CHECK(reply->type == REDIS_REPLY_ARRAY);
+  CHECK(reply->elements == 3);
+  redisReply *message_type = reply->element[0];
+  LOG_DEBUG("Driver table subscribe callback, message %s", message_type->str);
+
+  if (strcmp(message_type->str, "message") == 0) {
+    /* Handle a driver heartbeat. Parse the payload and call the subscribe
+     * callback. */
+    auto message =
+        flatbuffers::GetRoot<DriverTableMessage>(reply->element[2]->str);
+    /* Extract the client ID. */
+    WorkerID driver_id = from_flatbuf(message->driver_id());
+
+    /* Call the subscribe callback. */
+    DriverTableSubscribeData *data =
+        (DriverTableSubscribeData *) callback_data->data;
+    if (data->subscribe_callback) {
+      data->subscribe_callback(driver_id, data->subscribe_context);
+    }
+  } else if (strcmp(message_type->str, "subscribe") == 0) {
+    /* The reply for the initial SUBSCRIBE command. */
+    CHECK(callback_data->done_callback == NULL);
+    /* If the initial SUBSCRIBE was successful, clean up the timer, but don't
+     * destroy the callback data. */
+    event_loop_remove_timer(db->loop, callback_data->timer_id);
+
+  } else {
+    LOG_FATAL("Unexpected reply type from driver subscribe.");
+  }
+}
+
+void redis_driver_table_subscribe(TableCallbackData *callback_data) {
+  DBHandle *db = callback_data->db_handle;
+  int status = redisAsyncCommand(
+      db->sub_context, redis_driver_table_subscribe_callback,
+      (void *) callback_data->timer_id, "SUBSCRIBE driver_deaths");
+  if ((status == REDIS_ERR) || db->sub_context->err) {
+    LOG_REDIS_DEBUG(db->sub_context, "error in redis_driver_table_subscribe");
+  }
+}
+
+void redis_driver_table_send_driver_death_callback(redisAsyncContext *c,
+                                                   void *r,
+                                                   void *privdata) {
+  REDIS_CALLBACK_HEADER(db, callback_data, r);
+
+  redisReply *reply = (redisReply *) r;
+  CHECK(reply->type == REDIS_REPLY_INTEGER);
+  LOG_DEBUG("%" PRId64 " subscribers received this publish.\n", reply->integer);
+  /* At the very least, the local scheduler that publishes this message should
+   * also receive it. */
+  CHECK(reply->integer >= 1);
+
+  CHECK(callback_data->done_callback == NULL);
+  /* Clean up the timer and callback. */
+  destroy_timer_callback(db->loop, callback_data);
+}
+
+void redis_driver_table_send_driver_death(TableCallbackData *callback_data) {
+  DBHandle *db = callback_data->db_handle;
+  WorkerID driver_id = callback_data->id;
+
+  /* Create a flatbuffer object to serialize and publish. */
+  flatbuffers::FlatBufferBuilder fbb;
+  /* Create the flatbuffers message. */
+  auto message = CreateDriverTableMessage(fbb, to_flatbuf(fbb, driver_id));
+  fbb.Finish(message);
+
+  int status = redisAsyncCommand(
+      db->context, redis_driver_table_send_driver_death_callback,
+      (void *) callback_data->timer_id, "PUBLISH driver_deaths %b",
+      fbb.GetBufferPointer(), fbb.GetSize());
+  if ((status == REDIS_ERR) || db->context->err) {
+    LOG_REDIS_DEBUG(db->context,
+                    "error in redis_driver_table_send_driver_death");
+  }
+}
+
 void redis_plasma_manager_send_heartbeat(TableCallbackData *callback_data) {
   DBHandle *db = callback_data->db_handle;
   /* NOTE(swang): We purposefully do not provide a callback, leaving the table
@@ -1124,15 +1209,20 @@ void redis_actor_notification_table_subscribe_callback(redisAsyncContext *c,
     redisReply *payload = reply->element[2];
     ActorNotificationTableSubscribeData *data =
         (ActorNotificationTableSubscribeData *) callback_data->data;
-    ActorInfo info;
-    /* The payload should be the concatenation of these two structs. */
-    CHECK(sizeof(info.actor_id) + sizeof(info.local_scheduler_id) ==
+    /* The payload should be the concatenation of three IDs. */
+    ActorID actor_id;
+    WorkerID driver_id;
+    DBClientID local_scheduler_id;
+    CHECK(sizeof(actor_id) + sizeof(driver_id) + sizeof(local_scheduler_id) ==
           payload->len);
-    memcpy(&info.actor_id, payload->str, sizeof(info.actor_id));
-    memcpy(&info.local_scheduler_id, payload->str + sizeof(info.actor_id),
-           sizeof(info.local_scheduler_id));
+    memcpy(&actor_id, payload->str, sizeof(actor_id));
+    memcpy(&driver_id, payload->str + sizeof(actor_id), sizeof(driver_id));
+    memcpy(&local_scheduler_id,
+           payload->str + sizeof(actor_id) + sizeof(driver_id),
+           sizeof(local_scheduler_id));
     if (data->subscribe_callback) {
-      data->subscribe_callback(info, data->subscribe_context);
+      data->subscribe_callback(actor_id, driver_id, local_scheduler_id,
+                               data->subscribe_context);
     }
   } else if (strcmp(message_type->str, "subscribe") == 0) {
     /* The reply for the initial SUBSCRIBE command. */

src/common/state/redis.h

@@ -253,6 +253,24 @@ void redis_local_scheduler_table_subscribe(TableCallbackData *callback_data);
  */
 void redis_local_scheduler_table_send_info(TableCallbackData *callback_data);
 
+/**
+ * Subscribe to updates from the driver table.
+ *
+ * @param callback_data Data structure containing redis connection and timeout
+ *        information.
+ * @return Void.
+ */
+void redis_driver_table_subscribe(TableCallbackData *callback_data);
+
+/**
+ * Publish an update to the driver table.
+ *
+ * @param callback_data Data structure containing redis connection and timeout
+ *        information.
+ * @return Void.
+ */
+void redis_driver_table_send_driver_death(TableCallbackData *callback_data);
+
 void redis_plasma_manager_send_heartbeat(TableCallbackData *callback_data);
 
 /**

src/common/task.h

@@ -16,7 +16,6 @@ struct TaskBuilder;
 #define NIL_TASK_ID NIL_ID
 #define NIL_ACTOR_ID NIL_ID
 #define NIL_FUNCTION_ID NIL_ID
-#define NIL_WORKER_ID NIL_ID
 
 typedef UniqueID FunctionID;
 

src/global_scheduler/global_scheduler.cc

@@ -61,14 +61,15 @@ void assign_task_to_local_scheduler(GlobalSchedulerState *state,
 }
 
 GlobalSchedulerState *GlobalSchedulerState_init(event_loop *loop,
+                                                const char *node_ip_address,
                                                 const char *redis_addr,
                                                 int redis_port) {
   GlobalSchedulerState *state =
       (GlobalSchedulerState *) malloc(sizeof(GlobalSchedulerState));
   /* Must initialize state to 0. Sets hashmap head(s) to NULL. */
   memset(state, 0, sizeof(GlobalSchedulerState));
-  state->db =
-      db_connect(redis_addr, redis_port, "global_scheduler", ":", 0, NULL);
+  state->db = db_connect(redis_addr, redis_port, "global_scheduler",
+                         node_ip_address, 0, NULL);
   db_attach(state->db, loop, false);
   utarray_new(state->local_schedulers, &local_scheduler_icd);
   state->policy_state = GlobalSchedulerPolicyState_init();
@@ -416,9 +417,12 @@ int heartbeat_timeout_handler(event_loop *loop, timer_id id, void *context) {
   return HEARTBEAT_TIMEOUT_MILLISECONDS;
 }
 
-void start_server(const char *redis_addr, int redis_port) {
+void start_server(const char *node_ip_address,
+                  const char *redis_addr,
+                  int redis_port) {
   event_loop *loop = event_loop_create();
-  g_state = GlobalSchedulerState_init(loop, redis_addr, redis_port);
+  g_state =
+      GlobalSchedulerState_init(loop, node_ip_address, redis_addr, redis_port);
   /* TODO(rkn): subscribe to notifications from the object table. */
   /* Subscribe to notifications about new local schedulers. TODO(rkn): this
    * needs to also get all of the clients that registered with the database
@@ -456,12 +460,17 @@ int main(int argc, char *argv[]) {
   signal(SIGTERM, signal_handler);
   /* IP address and port of redis. */
   char *redis_addr_port = NULL;
+  /* The IP address of the node that this global scheduler is running on. */
+  char *node_ip_address = NULL;
   int c;
-  while ((c = getopt(argc, argv, "s:m:h:p:r:")) != -1) {
+  while ((c = getopt(argc, argv, "h:r:")) != -1) {
     switch (c) {
     case 'r':
       redis_addr_port = optarg;
       break;
+    case 'h':
+      node_ip_address = optarg;
+      break;
     default:
       LOG_ERROR("unknown option %c", c);
       exit(-1);
@@ -472,8 +481,11 @@ int main(int argc, char *argv[]) {
   if (!redis_addr_port ||
       parse_ip_addr_port(redis_addr_port, redis_addr, &redis_port) == -1) {
     LOG_ERROR(
-        "need to specify redis address like 127.0.0.1:6379 with -r switch");
+        "specify the redis address like 127.0.0.1:6379 with the -r switch");
     exit(-1);
   }
-  start_server(redis_addr, redis_port);
+  if (!node_ip_address) {
+    LOG_FATAL("specify the node IP address with the -h switch");
+  }
+  start_server(node_ip_address, redis_addr, redis_port);
 }

src/local_scheduler/local_scheduler.cc

@@ -18,6 +18,7 @@
 #include "local_scheduler_algorithm.h"
 #include "state/actor_notification_table.h"
 #include "state/db.h"
+#include "state/driver_table.h"
 #include "state/task_table.h"
 #include "state/object_table.h"
 #include "state/error_table.h"
@@ -67,7 +68,8 @@ int force_kill_worker(event_loop *loop, timer_id id, void *context) {
 
 /**
  * Kill a worker, if it is a child process, and clean up all of its associated
- * state.
+ * state. Note that this function is also called on drivers, but it should not
+ * actually send a kill signal to drivers.
  *
  * @param worker A pointer to the worker we want to kill.
  * @param cleanup A bool representing whether we're cleaning up the entire local
@@ -89,7 +91,13 @@ void kill_worker(LocalSchedulerState *state,
   CHECK(it == state->workers.end());
 
   /* Erase the algorithm state's reference to the worker. */
-  handle_worker_removed(state, state->algorithm_state, worker);
+  if (ActorID_equal(worker->actor_id, NIL_ACTOR_ID)) {
+    handle_worker_removed(state, state->algorithm_state, worker);
+  } else {
+    /* Let the scheduling algorithm process the absence of this worker. */
+    handle_actor_worker_disconnect(state, state->algorithm_state,
+                                   worker->actor_id);
+  }
 
   /* Remove the client socket from the event loop so that we don't process the
    * SIGPIPE when the worker is killed. */
@@ -99,6 +107,8 @@ void kill_worker(LocalSchedulerState *state,
    * process, use it to send a kill signal. */
   bool free_worker = true;
   if (worker->is_child && worker->pid != 0) {
+    /* If worker is a driver, we should not enter this condition because
+     * worker->pid should be 0. */
     if (cleanup) {
       /* If we're exiting the local scheduler anyway, it's okay to force kill
        * the worker immediately. Wait for the process to exit. */
@@ -425,10 +435,18 @@ void update_dynamic_resources(LocalSchedulerState *state,
   print_resource_info(state, spec);
 }
 
+bool is_driver_alive(LocalSchedulerState *state, WorkerID driver_id) {
+  return state->removed_drivers.count(driver_id) == 0;
+}
+
 void assign_task_to_worker(LocalSchedulerState *state,
                            TaskSpec *spec,
                            int64_t task_spec_size,
                            LocalSchedulerClient *worker) {
+  /* Make sure the driver for this task is still alive. */
+  WorkerID driver_id = TaskSpec_driver_id(spec);
+  CHECK(is_driver_alive(state, driver_id));
+
   /* Construct a flatbuffer object to send to the worker. */
   flatbuffers::FlatBufferBuilder fbb;
   auto message =
@@ -649,8 +667,15 @@ void send_client_register_reply(LocalSchedulerState *state,
 void handle_client_register(LocalSchedulerState *state,
                             LocalSchedulerClient *worker,
                             const RegisterClientRequest *message) {
+  /* Make sure this worker hasn't already registered. */
+  CHECK(!worker->registered);
+  worker->registered = true;
+  worker->is_worker = message->is_worker();
+  CHECK(WorkerID_equal(worker->client_id, NIL_WORKER_ID));
+  worker->client_id = from_flatbuf(message->client_id());
+
   /* Register the worker or driver. */
-  if (message->is_worker()) {
+  if (worker->is_worker) {
     /* Update the actor mapping with the actor ID of the worker (if an actor is
      * running on the worker). */
     worker->pid = message->worker_pid();
@@ -683,12 +708,74 @@ void handle_client_register(LocalSchedulerState *state,
       state->child_pids.erase(it);
       LOG_DEBUG("Found matching child pid %d", worker->pid);
     }
+
+    /* If the worker is an actor that corresponds to a driver that has been
+     * removed, then kill the worker. */
+    if (!ActorID_equal(actor_id, NIL_ACTOR_ID)) {
+      WorkerID driver_id = state->actor_mapping[actor_id].driver_id;
+      if (state->removed_drivers.count(driver_id) == 1) {
+        kill_worker(state, worker, false);
+      }
+    }
   } else {
     /* Register the driver. Currently we don't do anything here. */
   }
 }
 
-/* End of the cleanup code. */
+void handle_driver_removed_callback(WorkerID driver_id, void *user_context) {
+  LocalSchedulerState *state = (LocalSchedulerState *) user_context;
+
+  /* Kill any actors that were created by the removed driver, and kill any
+   * workers that are currently running tasks from the dead driver. */
+  auto it = state->workers.begin();
+  while (it != state->workers.end()) {
+    /* Increment the iterator by one before calling kill_worker, because
+     * kill_worker will invalidate the iterator. Note that this requires
+     * knowledge of the particular container that we are iterating over (in this
+     * case it is a list). */
+    auto next_it = it;
+    next_it++;
+
+    ActorID actor_id = (*it)->actor_id;
+    Task *task = (*it)->task_in_progress;
+
+    if (!ActorID_equal(actor_id, NIL_ACTOR_ID)) {
+      /* This is an actor. */
+      CHECK(state->actor_mapping.count(actor_id) == 1);
+      if (WorkerID_equal(state->actor_mapping[actor_id].driver_id, driver_id)) {
+        /* This actor was created by the removed driver, so kill the actor. */
+        LOG_DEBUG("Killing an actor for a removed driver.");
+        kill_worker(state, *it, false);
+        break;
+      }
+    } else if (task != NULL) {
+      if (WorkerID_equal(TaskSpec_driver_id(Task_task_spec(task)), driver_id)) {
+        LOG_DEBUG("Killing a worker executing a task for a removed driver.");
+        kill_worker(state, *it, false);
+        break;
+      }
+    }
+
+    it = next_it;
+  }
+
+  /* Add the driver to a list of dead drivers. */
+  state->removed_drivers.insert(driver_id);
+
+  /* Notify the scheduling algorithm that the driver has been removed. It should
+   * remove tasks for that driver from its data structures. */
+  handle_driver_removed(state, state->algorithm_state, driver_id);
+}
+
+void handle_client_disconnect(LocalSchedulerState *state,
+                              LocalSchedulerClient *worker) {
+  if (!worker->registered || worker->is_worker) {
+  } else {
+    /* In this case, a driver is disconecting. */
+    driver_table_send_driver_death(state->db, worker->client_id, NULL);
+  }
+  kill_worker(state, worker, false);
+}
 
 void process_message(event_loop *loop,
                      int client_sock,
@@ -786,12 +873,7 @@ void process_message(event_loop *loop,
   } break;
   case DISCONNECT_CLIENT: {
     LOG_INFO("Disconnecting client on fd %d", client_sock);
-    kill_worker(state, worker, false);
-    if (!ActorID_equal(worker->actor_id, NIL_ACTOR_ID)) {
-      /* Let the scheduling algorithm process the absence of this worker. */
-      handle_actor_worker_disconnect(state, state->algorithm_state,
-                                     worker->actor_id);
-    }
+    handle_client_disconnect(state, worker);
   } break;
   case MessageType_NotifyUnblocked: {
     if (worker->task_in_progress != NULL) {
@@ -827,6 +909,11 @@ void new_client_connection(event_loop *loop,
    * scheduler state. */
   LocalSchedulerClient *worker = new LocalSchedulerClient();
   worker->sock = new_socket;
+  worker->registered = false;
+  /* We don't know whether this is a worker or not, so just initialize is_worker
+   * to false. */
+  worker->is_worker = true;
+  worker->client_id = NIL_WORKER_ID;
   worker->task_in_progress = NULL;
   worker->is_blocked = false;
   worker->pid = 0;
@@ -857,10 +944,21 @@ void signal_handler(int signal) {
   }
 }
 
+/* End of the cleanup code. */
+
 void handle_task_scheduled_callback(Task *original_task,
                                     void *subscribe_context) {
   LocalSchedulerState *state = (LocalSchedulerState *) subscribe_context;
   TaskSpec *spec = Task_task_spec(original_task);
+
+  /* If the driver for this task has been removed, then don't bother telling the
+   * scheduling algorithm. */
+  WorkerID driver_id = TaskSpec_driver_id(spec);
+  if (!is_driver_alive(state, driver_id)) {
+    LOG_DEBUG("Ignoring scheduled task for removed driver.")
+    return;
+  }
+
   if (ActorID_equal(TaskSpec_actor_id(spec), NIL_ACTOR_ID)) {
     /* This task does not involve an actor. Handle it normally. */
     handle_task_scheduled(state, state->algorithm_state, spec,
@@ -884,10 +982,17 @@ void handle_task_scheduled_callback(Task *original_task,
  *        for creating the actor.
  * @return Void.
  */
-void handle_actor_creation_callback(ActorInfo info, void *context) {
-  ActorID actor_id = info.actor_id;
-  DBClientID local_scheduler_id = info.local_scheduler_id;
+void handle_actor_creation_callback(ActorID actor_id,
+                                    WorkerID driver_id,
+                                    DBClientID local_scheduler_id,
+                                    void *context) {
   LocalSchedulerState *state = (LocalSchedulerState *) context;
+
+  /* If the driver has been removed, don't bother doing anything. */
+  if (state->removed_drivers.count(driver_id) == 1) {
+    return;
+  }
+
   /* Make sure the actor entry is not already present in the actor map table.
    * TODO(rkn): We will need to remove this check to handle the case where the
    * corresponding publish is retried and the case in which a task that creates
@@ -897,9 +1002,10 @@ void handle_actor_creation_callback(ActorInfo info, void *context) {
    * Currently this is never removed (except when the local scheduler state is
    * deleted). */
   ActorMapEntry entry;
-  entry.actor_id = actor_id;
   entry.local_scheduler_id = local_scheduler_id;
+  entry.driver_id = driver_id;
   state->actor_mapping[actor_id] = entry;
+
   /* If this local scheduler is responsible for the actor, then start a new
    * worker for the actor. */
   if (DBClientID_equal(local_scheduler_id, get_db_client_id(state->db))) {
@@ -960,6 +1066,11 @@ void start_server(const char *node_ip_address,
     actor_notification_table_subscribe(
         g_state->db, handle_actor_creation_callback, g_state, NULL);
   }
+  /* Subscribe to notifications about removed drivers. */
+  if (g_state->db != NULL) {
+    driver_table_subscribe(g_state->db, handle_driver_removed_callback, g_state,
+                           NULL);
+  }
   /* Create a timer for publishing information about the load on the local
    * scheduler to the local scheduler table. This message also serves as a
    * heartbeat. */

src/local_scheduler/local_scheduler.h

@@ -26,6 +26,14 @@ void new_client_connection(event_loop *loop,
                            void *context,
                            int events);
 
+/**
+ * Check if a driver is still alive.
+ *
+ * @param driver_id The ID of the driver.
+ * @return True if the driver is still alive and false otherwise.
+ */
+bool is_driver_alive(WorkerID driver_id);
+
 /**
  * This function can be called by the scheduling algorithm to assign a task
  * to a worker.

src/local_scheduler/local_scheduler_algorithm.cc

@@ -964,6 +964,9 @@ void handle_worker_available(LocalSchedulerState *state,
 void handle_worker_removed(LocalSchedulerState *state,
                            SchedulingAlgorithmState *algorithm_state,
                            LocalSchedulerClient *worker) {
+  /* Make sure this is not an actor. */
+  CHECK(ActorID_equal(worker->actor_id, NIL_ACTOR_ID));
+
   /* Make sure that we remove the worker at most once. */
   int num_times_removed = 0;
 
@@ -1141,6 +1144,59 @@ void handle_object_removed(LocalSchedulerState *state,
   }
 }
 
+void handle_driver_removed(LocalSchedulerState *state,
+                           SchedulingAlgorithmState *algorithm_state,
+                           WorkerID driver_id) {
+  /* Loop over fetch requests. This must be done before we clean up the waiting
+   * task queue and the dispatch task queue because this map contains iterators
+   * for those lists, which will be invalidated when we clean up those lists.*/
+  for (auto it = algorithm_state->remote_objects.begin();
+       it != algorithm_state->remote_objects.end();) {
+    /* Loop over the tasks that are waiting for this object and remove the tasks
+     * for the removed driver. */
+    auto task_it_it = it->second.dependent_tasks.begin();
+    while (task_it_it != it->second.dependent_tasks.end()) {
+      /* If the dependent task was a task for the removed driver, remove it from
+       * this vector. */
+      TaskSpec *spec = (*task_it_it)->spec;
+      if (WorkerID_equal(TaskSpec_driver_id(spec), driver_id)) {
+        task_it_it = it->second.dependent_tasks.erase(task_it_it);
+      } else {
+        task_it_it++;
+      }
+    }
+    /* If there are no more dependent tasks for this object, then remove the
+     * ObjectEntry. */
+    if (it->second.dependent_tasks.size() == 0) {
+      it = algorithm_state->remote_objects.erase(it);
+    } else {
+      it++;
+    }
+  }
+
+  /* Remove this driver's tasks from the waiting task queue. */
+  auto it = algorithm_state->waiting_task_queue->begin();
+  while (it != algorithm_state->waiting_task_queue->end()) {
+    if (WorkerID_equal(TaskSpec_driver_id(it->spec), driver_id)) {
+      it = algorithm_state->waiting_task_queue->erase(it);
+    } else {
+      it++;
+    }
+  }
+
+  /* Remove this driver's tasks from the dispatch task queue. */
+  it = algorithm_state->dispatch_task_queue->begin();
+  while (it != algorithm_state->dispatch_task_queue->end()) {
+    if (WorkerID_equal(TaskSpec_driver_id(it->spec), driver_id)) {
+      it = algorithm_state->dispatch_task_queue->erase(it);
+    } else {
+      it++;
+    }
+  }
+
+  /* TODO(rkn): Should we clean up the actor data structures? */
+}
+
 int num_waiting_tasks(SchedulingAlgorithmState *algorithm_state) {
   return algorithm_state->waiting_task_queue->size();
 }

src/local_scheduler/local_scheduler_algorithm.h

@@ -230,6 +230,20 @@ void handle_worker_unblocked(LocalSchedulerState *state,
                              SchedulingAlgorithmState *algorithm_state,
                              LocalSchedulerClient *worker);
 
+/**
+ * Process the fact that a driver has been removed. This will remove all of the
+ * tasks for that driver from the scheduling algorithm's internal data
+ * structures.
+ *
+ * @param state The state of the local scheduler.
+ * @param algorithm_state State maintained by the scheduling algorithm.
+ * @param driver_id The ID of the driver that was removed.
+ * @return Void.
+ */
+void handle_driver_removed(LocalSchedulerState *state,
+                           SchedulingAlgorithmState *algorithm_state,
+                           WorkerID driver_id);
+
 /**
  * This function fetches queued task's missing object dependencies. It is
  * called every LOCAL_SCHEDULER_FETCH_TIMEOUT_MILLISECONDS.

src/local_scheduler/local_scheduler_shared.h

@@ -8,7 +8,9 @@
 #include "utarray.h"
 #include "uthash.h"
 
+#include <list>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 /* These are needed to define the UT_arrays. */
@@ -17,8 +19,8 @@ extern UT_icd task_ptr_icd;
 /** This struct is used to maintain a mapping from actor IDs to the ID of the
  *  local scheduler that is responsible for the actor. */
 struct ActorMapEntry {
-  /** The ID of the actor. This is used as a key in the hash table. */
-  ActorID actor_id;
+  /** The ID of the driver that created the actor. */
+  WorkerID driver_id;
   /** The ID of the local scheduler that is responsible for the actor. */
   DBClientID local_scheduler_id;
 };
@@ -47,7 +49,10 @@ struct LocalSchedulerState {
   /** List of workers available to this node. This is used to free the worker
    *  structs when we free the scheduler state and also to access the worker
    *  structs in the tests. */
-  std::vector<LocalSchedulerClient *> workers;
+  std::list<LocalSchedulerClient *> workers;
+  /** A set of driver IDs corresponding to drivers that have been removed. This
+   *  is used to make sure we don't execute any tasks belong to dead drivers. */
+  std::unordered_set<WorkerID, UniqueIDHasher> removed_drivers;
   /** List of the process IDs for child processes (workers) started by the
    *  local scheduler that have not sent a REGISTER_PID message yet. */
   std::vector<pid_t> child_pids;
@@ -75,6 +80,13 @@ struct LocalSchedulerState {
 struct LocalSchedulerClient {
   /** The socket used to communicate with the client. */
   int sock;
+  /** True if the client has registered and false otherwise. */
+  bool registered;
+  /** True if the client is a worker and false if it is a driver. */
+  bool is_worker;
+  /** The worker ID if the client is a worker and the driver ID if the client is
+   *  a driver. */
+  WorkerID client_id;
   /** A pointer to the task object that is currently running on this client. If
    *  no task is running on the worker, this will be NULL. This is used to
    *  update the task table. */

src/local_scheduler/test/local_scheduler_tests.cc

@@ -64,9 +64,7 @@ static void register_clients(int num_mock_workers, LocalSchedulerMock *mock) {
   for (int i = 0; i < num_mock_workers; ++i) {
     new_client_connection(mock->loop, mock->local_scheduler_fd,
                           (void *) mock->local_scheduler_state, 0);
-
-    LocalSchedulerClient *worker = mock->local_scheduler_state->workers[i];
-
+    LocalSchedulerClient *worker = mock->local_scheduler_state->workers.back();
     process_message(mock->local_scheduler_state->loop, worker->sock, worker, 0);
   }
 }
@@ -644,7 +642,7 @@ TEST start_kill_workers_test(void) {
   ASSERT_EQ(local_scheduler->local_scheduler_state->workers.size(),
             num_workers);
   /* Make sure that the new worker registers its process ID. */
-  worker = local_scheduler->local_scheduler_state->workers[num_workers - 1];
+  worker = local_scheduler->local_scheduler_state->workers.back();
   process_message(local_scheduler->local_scheduler_state->loop, worker->sock,
                   worker, 0);
   ASSERT_EQ(local_scheduler->local_scheduler_state->child_pids.size(), 0);