[xray] Track ray.get calls as task dependencies (#2362)

java/runtime-common/src/main/java/org/ray/core/WorkerContext.java

@@ -2,6 +2,7 @@ package org.ray.core;
 
 import org.ray.api.UniqueID;
 import org.ray.core.model.RayParameters;
+import org.ray.core.model.WorkerMode;
 import org.ray.spi.model.TaskSpec;
 
 public class WorkerContext {
@@ -35,7 +36,11 @@ public class WorkerContext {
 
     TaskSpec dummy = new TaskSpec();
     dummy.parentTaskId = UniqueID.nil;
-    dummy.taskId = UniqueID.nil;
+    if (params.worker_mode == WorkerMode.DRIVER) {
+      dummy.taskId = UniqueID.randomId();
+    } else {
+      dummy.taskId = UniqueID.nil;
+    }
     dummy.actorId = UniqueID.nil;
     dummy.driverId = params.driver_id;
     prepare(dummy, null);

java/runtime-native/src/main/java/org/ray/core/impl/RayNativeRuntime.java

@@ -109,6 +109,7 @@ public class RayNativeRuntime extends RayRuntime {
           WorkerContext.currentWorkerId(),
           UniqueID.nil,
           isWorker,
+          WorkerContext.currentTask().taskId,
           0
       );
 
@@ -237,4 +238,4 @@ public class RayNativeRuntime extends RayRuntime {
       throw new TaskExecutionException(log, e);
     }
   }
-}
+}

java/runtime-native/src/main/java/org/ray/spi/impl/DefaultLocalSchedulerClient.java

@@ -26,13 +26,13 @@ public class DefaultLocalSchedulerClient implements LocalSchedulerLink {
   private long client = 0;
 
   public DefaultLocalSchedulerClient(String schedulerSockName, UniqueID clientId, UniqueID actorId,
-                                     boolean isWorker, long numGpus) {
+                                     boolean isWorker, UniqueID driverId, long numGpus) {
     client = _init(schedulerSockName, clientId.getBytes(), actorId.getBytes(), isWorker,
-        numGpus);
+        driverId.getBytes(), numGpus);
   }
 
   private static native long _init(String localSchedulerSocket, byte[] workerId, byte[] actorId,
-                                   boolean isWorker, long numGpus);
+                                   boolean isWorker, byte[] driverTaskId, long numGpus);
 
   private static native byte[] _computePutId(long client, byte[] taskId, int putIndex);
 

python/ray/global_scheduler/test/test.py

@@ -96,7 +96,8 @@ class TestGlobalScheduler(unittest.TestCase):
                 static_resources={"CPU": 10})
             # Connect to the scheduler.
             local_scheduler_client = local_scheduler.LocalSchedulerClient(
-                local_scheduler_name, NIL_WORKER_ID, False, False)
+                local_scheduler_name, NIL_WORKER_ID, False, random_task_id(),
+                False)
             self.local_scheduler_clients.append(local_scheduler_client)
             self.local_scheduler_pids.append(p4)
 

python/ray/local_scheduler/test/test.py

@@ -46,7 +46,7 @@ class TestLocalSchedulerClient(unittest.TestCase):
             plasma_store_name, use_valgrind=USE_VALGRIND)
         # Connect to the scheduler.
         self.local_scheduler_client = local_scheduler.LocalSchedulerClient(
-            scheduler_name, NIL_WORKER_ID, False, False)
+            scheduler_name, NIL_WORKER_ID, False, random_task_id(), False)
 
     def tearDown(self):
         # Check that the processes are still alive.

python/ray/worker.py

@@ -503,15 +503,6 @@ class Worker(object):
                 # get them until at least get_timeout_milliseconds
                 # milliseconds passes, then repeat.
                 while len(unready_ids) > 0:
-                    for unready_id in unready_ids:
-                        if not self.use_raylet:
-                            self.local_scheduler_client.reconstruct_objects(
-                                [ray.ObjectID(unready_id)], False)
-                    # Do another fetch for objects that aren't available
-                    # locally yet, in case they were evicted since the last
-                    # fetch. We divide the fetch into smaller fetches so as
-                    # to not block the manager for a prolonged period of time
-                    # in a single call.
                     object_ids_to_fetch = [
                         plasma.ObjectID(unready_id)
                         for unready_id in unready_ids.keys()
@@ -525,6 +516,18 @@ class Worker(object):
                     for i in range(0, len(object_ids_to_fetch),
                                    fetch_request_size):
                         if not self.use_raylet:
+                            for unready_id in ray_object_ids_to_fetch[i:(
+                                    i + fetch_request_size)]:
+                                (self.local_scheduler_client.
+                                 reconstruct_objects([unready_id], False))
+                            # Do another fetch for objects that aren't
+                            # available locally yet, in case they were evicted
+                            # since the last fetch. We divide the fetch into
+                            # smaller fetches so as to not block the manager
+                            # for a prolonged period of time in a single call.
+                            # This is only necessary for legacy ray since
+                            # reconstruction and fetch are implemented by
+                            # different processes.
                             self.plasma_client.fetch(object_ids_to_fetch[i:(
                                 i + fetch_request_size)])
                         else:
@@ -2162,9 +2165,6 @@ def connect(info,
     else:
         local_scheduler_socket = info["raylet_socket_name"]
 
-    worker.local_scheduler_client = ray.local_scheduler.LocalSchedulerClient(
-        local_scheduler_socket, worker.worker_id, is_worker, worker.use_raylet)
-
     # If this is a driver, set the current task ID, the task driver ID, and set
     # the task index to 0.
     if mode in [SCRIPT_MODE, SILENT_MODE]:
@@ -2219,6 +2219,13 @@ def connect(info,
         # Set the driver's current task ID to the task ID assigned to the
         # driver task.
         worker.current_task_id = driver_task.task_id()
+    else:
+        # A non-driver worker begins without an assigned task.
+        worker.current_task_id = ray.ObjectID(NIL_ID)
+
+    worker.local_scheduler_client = ray.local_scheduler.LocalSchedulerClient(
+        local_scheduler_socket, worker.worker_id, is_worker,
+        worker.current_task_id, worker.use_raylet)
 
     # Start the import thread
     import_thread.ImportThread(worker, mode).start()

src/local_scheduler/format/local_scheduler.fbs

@@ -79,6 +79,8 @@ table RegisterClientRequest {
   client_id: string;
   // The process ID of this worker.
   worker_pid: long;
+  // The driver ID. This is non-nil if the client is a driver.
+  driver_id: string;
 }
 
 table DisconnectClient {

src/local_scheduler/lib/java/org_ray_spi_impl_DefaultLocalSchedulerClient.cc

@@ -42,14 +42,16 @@ Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1init(JNIEnv *env,
                                                          jbyteArray wid,
                                                          jbyteArray actorId,
                                                          jboolean isWorker,
+                                                         jbyteArray driverId,
                                                          jlong numGpus) {
   // 	native private static long _init(String localSchedulerSocket,
   //     byte[] workerId, byte[] actorId, boolean isWorker, long numGpus);
   UniqueIdFromJByteArray worker_id(env, wid);
+  UniqueIdFromJByteArray driver_id(env, driverId);
   const char *nativeString = env->GetStringUTFChars(sockName, JNI_FALSE);
   bool use_raylet = false;
-  auto client = LocalSchedulerConnection_init(nativeString, *worker_id.PID,
-                                              isWorker, use_raylet);
+  auto client = LocalSchedulerConnection_init(
+      nativeString, *worker_id.PID, isWorker, *driver_id.PID, use_raylet);
   env->ReleaseStringUTFChars(sockName, nativeString);
   return reinterpret_cast<jlong>(client);
 }

src/local_scheduler/lib/java/org_ray_spi_impl_DefaultLocalSchedulerClient.h

@@ -19,6 +19,7 @@ Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1init(JNIEnv *,
                                                          jbyteArray,
                                                          jbyteArray,
                                                          jboolean,
+                                                         jbyteArray,
                                                          jlong);
 
 /*

src/local_scheduler/lib/python/local_scheduler_extension.cc

@@ -20,16 +20,18 @@ static int PyLocalSchedulerClient_init(PyLocalSchedulerClient *self,
   char *socket_name;
   UniqueID client_id;
   PyObject *is_worker;
+  JobID driver_id;
   PyObject *use_raylet;
-  if (!PyArg_ParseTuple(args, "sO&OO", &socket_name, PyStringToUniqueID,
-                        &client_id, &is_worker, &use_raylet)) {
+  if (!PyArg_ParseTuple(args, "sO&OO&O", &socket_name, PyStringToUniqueID,
+                        &client_id, &is_worker, &PyObjectToUniqueID, &driver_id,
+                        &use_raylet)) {
     self->local_scheduler_connection = NULL;
     return -1;
   }
   /* Connect to the local scheduler. */
   self->local_scheduler_connection = LocalSchedulerConnection_init(
       socket_name, client_id, static_cast<bool>(PyObject_IsTrue(is_worker)),
-      static_cast<bool>(PyObject_IsTrue(use_raylet)));
+      driver_id, static_cast<bool>(PyObject_IsTrue(use_raylet)));
   return 0;
 }
 

src/local_scheduler/local_scheduler_client.cc

@@ -14,8 +14,9 @@ using MessageType = ray::local_scheduler::protocol::MessageType;
 
 LocalSchedulerConnection *LocalSchedulerConnection_init(
     const char *local_scheduler_socket,
-    UniqueID client_id,
+    const UniqueID &client_id,
     bool is_worker,
+    const JobID &driver_id,
     bool use_raylet) {
   LocalSchedulerConnection *result = new LocalSchedulerConnection();
   result->use_raylet = use_raylet;
@@ -26,7 +27,8 @@ LocalSchedulerConnection *LocalSchedulerConnection_init(
    * worker, we will get killed. */
   flatbuffers::FlatBufferBuilder fbb;
   auto message = ray::local_scheduler::protocol::CreateRegisterClientRequest(
-      fbb, is_worker, to_flatbuf(fbb, client_id), getpid());
+      fbb, is_worker, to_flatbuf(fbb, client_id), getpid(),
+      to_flatbuf(fbb, driver_id));
   fbb.Finish(message);
   /* Register the process ID with the local scheduler. */
   int success = write_message(

src/local_scheduler/local_scheduler_client.h

@@ -32,16 +32,20 @@ struct LocalSchedulerConnection {
  *
  * @param local_scheduler_socket The name of the socket to use to connect to the
  *        local scheduler.
+ * @param worker_id A unique ID to represent the worker.
  * @param is_worker Whether this client is a worker. If it is a worker, an
  *        additional message will be sent to register as one.
+ * @param driver_id The ID of the driver. This is non-nil if the client is a
+ *        driver.
  * @param use_raylet True if we should use the raylet code path and false
  *        otherwise.
  * @return The connection information.
  */
 LocalSchedulerConnection *LocalSchedulerConnection_init(
     const char *local_scheduler_socket,
-    UniqueID worker_id,
+    const UniqueID &worker_id,
     bool is_worker,
+    const JobID &driver_id,
     bool use_raylet);
 
 /**

src/local_scheduler/test/local_scheduler_tests.cc

@@ -125,7 +125,8 @@ LocalSchedulerMock *LocalSchedulerMock_init(int num_workers,
 
   for (int i = 0; i < num_mock_workers; ++i) {
     mock->conns[i] = LocalSchedulerConnection_init(
-        local_scheduler_socket_name.c_str(), WorkerID::nil(), true, false);
+        local_scheduler_socket_name.c_str(), WorkerID::nil(), true,
+        JobID::nil(), false);
   }
 
   background_thread.join();

src/ray/raylet/format/node_manager.fbs

@@ -119,6 +119,8 @@ table RegisterClientRequest {
   client_id: string;
   // The process ID of this worker.
   worker_pid: long;
+  // The driver ID. This is non-nil if the client is a driver.
+  driver_id: string;
 }
 
 table RegisterClientReply {

src/ray/raylet/node_manager.cc

@@ -381,9 +381,8 @@ void NodeManager::DispatchTasks() {
     }
     // We have enough resources for this task. Assign task.
     // TODO(atumanov): perform the task state/queue transition inside AssignTask.
-    auto dispatched_task =
-        local_queues_.RemoveTasks({task.GetTaskSpecification().TaskId()});
-    AssignTask(dispatched_task.front());
+    auto dispatched_task = local_queues_.RemoveTask(task.GetTaskSpecification().TaskId());
+    AssignTask(dispatched_task);
   }
 }
 
@@ -395,11 +394,16 @@ void NodeManager::ProcessClientMessage(
   switch (static_cast<protocol::MessageType>(message_type)) {
   case protocol::MessageType::RegisterClientRequest: {
     auto message = flatbuffers::GetRoot<protocol::RegisterClientRequest>(message_data);
+    auto worker = std::make_shared<Worker>(message->worker_pid(), client);
     if (message->is_worker()) {
-      // Create a new worker from the registration request.
-      auto worker = std::make_shared<Worker>(message->worker_pid(), client);
       // Register the new worker.
       worker_pool_.RegisterWorker(std::move(worker));
+    } else {
+      // Register the new driver.
+      JobID job_id = from_flatbuf(*message->driver_id());
+      worker->AssignTaskId(job_id);
+      worker_pool_.RegisterDriver(std::move(worker));
+      local_queues_.AddDriverTaskId(job_id);
     }
   } break;
   case protocol::MessageType::GetTask: {
@@ -419,10 +423,10 @@ void NodeManager::ProcessClientMessage(
     // Remove the dead worker from the pool and stop listening for messages.
     const std::shared_ptr<Worker> worker = worker_pool_.GetRegisteredWorker(client);
 
-    // This if statement distinguishes workers from drivers.
     if (worker) {
-      // 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.
+      // The client is a worker. 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();
@@ -460,6 +464,14 @@ void NodeManager::ProcessClientMessage(
 
       // Since some resources may have been released, we can try to dispatch more tasks.
       DispatchTasks();
+    } else {
+      // The client is a driver.
+      const std::shared_ptr<Worker> driver = worker_pool_.GetRegisteredDriver(client);
+      RAY_CHECK(driver);
+      auto driver_id = driver->GetAssignedTaskId();
+      RAY_CHECK(!driver_id.is_nil());
+      local_queues_.RemoveDriverTaskId(driver_id);
+      worker_pool_.DisconnectDriver(driver);
     }
     return;
   } break;
@@ -475,95 +487,77 @@ void NodeManager::ProcessClientMessage(
     SubmitTask(task, Lineage());
   } break;
   case protocol::MessageType::ReconstructObjects: {
-    // TODO(hme): handle multiple object ids.
     auto message = flatbuffers::GetRoot<protocol::ReconstructObjects>(message_data);
+    std::vector<ObjectID> required_object_ids;
     for (size_t i = 0; i < message->object_ids()->size(); ++i) {
       ObjectID object_id = from_flatbuf(*message->object_ids()->Get(i));
-      RAY_LOG(DEBUG) << "reconstructing object " << object_id;
       if (!task_dependency_manager_.CheckObjectLocal(object_id)) {
-        // TODO(swang): Instead of calling Pull on the object directly, record the
-        // fact that the blocked task is dependent on this object_id in the task
-        // dependency manager.
-        RAY_CHECK_OK(object_manager_.Pull(object_id));
-      }
-
-      if (!message->fetch_only()) {
-        // If the blocked client is a worker, and the worker isn't already blocked,
-        // then release any CPU resources that it acquired for its assigned task
-        // while it is blocked. The resources will be acquired again once the
-        // worker is unblocked.
-        std::shared_ptr<Worker> worker = worker_pool_.GetRegisteredWorker(client);
-        if (worker && !worker->IsBlocked()) {
-          RAY_CHECK(!worker->GetAssignedTaskId().is_nil());
-          auto tasks = local_queues_.RemoveTasks({worker->GetAssignedTaskId()});
-          const auto &task = tasks.front();
-          // Get the CPU resources required by the running task.
-          const auto required_resources =
-              task.GetTaskSpecification().GetRequiredResources();
-          double required_cpus = required_resources.GetNumCpus();
-          const std::unordered_map<std::string, double> cpu_resources = {
-              {kCPU_ResourceLabel, required_cpus}};
-
-          // Release the CPU resources.
-          auto const cpu_resource_ids = worker->ReleaseTaskCpuResources();
-          local_available_resources_.Release(cpu_resource_ids);
-          RAY_CHECK(cluster_resource_map_[gcs_client_->client_table().GetLocalClientId()]
-                        .Release(ResourceSet(cpu_resources)));
-
-          // Mark the task as blocked.
-          local_queues_.QueueBlockedTasks(tasks);
-          worker->MarkBlocked();
-
-          // Try to dispatch more tasks since the blocked worker released some
-          // resources.
-          DispatchTasks();
+        if (message->fetch_only()) {
+          // If only a fetch is required, then do not subscribe to the
+          // dependencies to the task dependency manager.
+          RAY_CHECK_OK(object_manager_.Pull(object_id));
+        } else {
+          // If reconstruction is also required, then add any missing objects
+          // to the list to subscribe to in the task dependency manager. These
+          // objects will be pulled from remote node managers and reconstructed
+          // if necessary.
+          required_object_ids.push_back(object_id);
         }
       }
     }
+
+    if (!required_object_ids.empty()) {
+      std::shared_ptr<Worker> worker = worker_pool_.GetRegisteredWorker(client);
+      if (worker) {
+        // The client is a worker.  Mark the worker as blocked. This
+        // temporarily releases any resources that the worker holds while it is
+        // blocked.
+        HandleWorkerBlocked(worker);
+      } else {
+        // The client is a driver. Drivers do not hold resources, so we simply
+        // mark the driver as blocked.
+        worker = worker_pool_.GetRegisteredDriver(client);
+        RAY_CHECK(worker);
+        worker->MarkBlocked();
+      }
+      const TaskID current_task_id = worker->GetAssignedTaskId();
+      RAY_CHECK(!current_task_id.is_nil());
+      // Subscribe to the objects required by the ray.get. These objects will
+      // be fetched and/or reconstructed as necessary, until the objects become
+      // local or are unsubscribed.
+      task_dependency_manager_.SubscribeDependencies(current_task_id,
+                                                     required_object_ids);
+    }
   } break;
   case protocol::MessageType::NotifyUnblocked: {
     std::shared_ptr<Worker> worker = worker_pool_.GetRegisteredWorker(client);
+
     // Re-acquire the CPU resources for the task that was assigned to the
     // unblocked worker.
+    // TODO(swang): Because the object dependencies are tracked in the task
+    // dependency manager, we could actually remove this message entirely and
+    // instead unblock the worker once all the objects become available.
+    bool was_blocked;
     if (worker) {
-      RAY_CHECK(worker->IsBlocked());
-      RAY_CHECK(!worker->GetAssignedTaskId().is_nil());
-
-      auto tasks = local_queues_.RemoveTasks({worker->GetAssignedTaskId()});
-      const auto &task = tasks.front();
-      // Get the CPU resources required by the running task.
-      const auto required_resources = task.GetTaskSpecification().GetRequiredResources();
-      double required_cpus = required_resources.GetNumCpus();
-      const ResourceSet cpu_resources(
-          std::unordered_map<std::string, double>({{kCPU_ResourceLabel, required_cpus}}));
-
-      // Check if we can reacquire the CPU resources.
-      bool oversubscribed = !local_available_resources_.Contains(cpu_resources);
-
-      if (!oversubscribed) {
-        // Reacquire the CPU resources for the worker. Note that care needs to be
-        // taken if the user is using the specific CPU IDs since the IDs that we
-        // reacquire here may be different from the ones that the task started with.
-        auto const resource_ids = local_available_resources_.Acquire(cpu_resources);
-        worker->AcquireTaskCpuResources(resource_ids);
-        RAY_CHECK(
-            cluster_resource_map_[gcs_client_->client_table().GetLocalClientId()].Acquire(
-                cpu_resources));
-      } else {
-        // In this case, we simply don't reacquire the CPU resources for the worker.
-        // The worker can keep running and when the task finishes, it will simply
-        // not have any CPU resources to release.
-        RAY_LOG(WARNING)
-            << "Resources oversubscribed: "
-            << cluster_resource_map_[gcs_client_->client_table().GetLocalClientId()]
-                   .GetAvailableResources()
-                   .ToString();
-      }
-
-      // Mark the task as running again.
-      local_queues_.QueueRunningTasks(tasks);
+      was_blocked = worker->IsBlocked();
+      // Mark the worker as unblocked. This returns the temporarily released
+      // resources to the worker.
+      HandleWorkerUnblocked(worker);
+    } else {
+      // The client is a driver. Drivers do not hold resources, so we simply
+      // mark the driver as unblocked.
+      worker = worker_pool_.GetRegisteredDriver(client);
+      RAY_CHECK(worker);
+      was_blocked = worker->IsBlocked();
       worker->MarkUnblocked();
     }
+    // Unsubscribe to the objects. Any fetch or reconstruction operations to
+    // make the objects local are canceled.
+    if (was_blocked) {
+      const TaskID current_task_id = worker->GetAssignedTaskId();
+      RAY_CHECK(!current_task_id.is_nil());
+      task_dependency_manager_.UnsubscribeDependencies(current_task_id);
+    }
   } break;
   case protocol::MessageType::WaitRequest: {
     // Read the data.
@@ -664,9 +658,7 @@ void NodeManager::ScheduleTasks() {
       local_task_ids.insert(task_id);
     } else {
       // TODO(atumanov): need a better interface for task exit on forward.
-      auto tasks = local_queues_.RemoveTasks({task_id});
-      RAY_CHECK(1 == tasks.size());
-      Task &task = tasks.front();
+      const auto task = local_queues_.RemoveTask(task_id);
       // TODO(swang): Handle forward task failure.
       RAY_CHECK_OK(ForwardTask(task, client_id));
     }
@@ -742,6 +734,78 @@ void NodeManager::SubmitTask(const Task &task, const Lineage &uncommitted_lineag
   }
 }
 
+void NodeManager::HandleWorkerBlocked(std::shared_ptr<Worker> worker) {
+  RAY_CHECK(worker);
+  if (worker->IsBlocked()) {
+    return;
+  }
+  // If the worker isn't already blocked, then release any CPU resources that
+  // it acquired for its assigned task while it is blocked. The resources will
+  // be acquired again once the worker is unblocked.
+  RAY_CHECK(!worker->GetAssignedTaskId().is_nil());
+  const auto task = local_queues_.RemoveTask(worker->GetAssignedTaskId());
+  // Get the CPU resources required by the running task.
+  const auto required_resources = task.GetTaskSpecification().GetRequiredResources();
+  double required_cpus = required_resources.GetNumCpus();
+  const std::unordered_map<std::string, double> cpu_resources = {
+      {kCPU_ResourceLabel, required_cpus}};
+
+  // Release the CPU resources.
+  auto const cpu_resource_ids = worker->ReleaseTaskCpuResources();
+  local_available_resources_.Release(cpu_resource_ids);
+  RAY_CHECK(cluster_resource_map_[gcs_client_->client_table().GetLocalClientId()].Release(
+      ResourceSet(cpu_resources)));
+
+  // Mark the task as blocked.
+  local_queues_.QueueBlockedTasks({task});
+  worker->MarkBlocked();
+
+  // Try to dispatch more tasks since the blocked worker released some
+  // resources.
+  DispatchTasks();
+}
+
+void NodeManager::HandleWorkerUnblocked(std::shared_ptr<Worker> worker) {
+  RAY_CHECK(worker);
+  if (!worker->IsBlocked()) {
+    return;
+  }
+
+  const auto task = local_queues_.RemoveTask(worker->GetAssignedTaskId());
+  // Get the CPU resources required by the running task.
+  const auto required_resources = task.GetTaskSpecification().GetRequiredResources();
+  double required_cpus = required_resources.GetNumCpus();
+  const ResourceSet cpu_resources(
+      std::unordered_map<std::string, double>({{kCPU_ResourceLabel, required_cpus}}));
+
+  // Check if we can reacquire the CPU resources.
+  bool oversubscribed = !local_available_resources_.Contains(cpu_resources);
+
+  if (!oversubscribed) {
+    // Reacquire the CPU resources for the worker. Note that care needs to be
+    // taken if the user is using the specific CPU IDs since the IDs that we
+    // reacquire here may be different from the ones that the task started with.
+    auto const resource_ids = local_available_resources_.Acquire(cpu_resources);
+    worker->AcquireTaskCpuResources(resource_ids);
+    RAY_CHECK(
+        cluster_resource_map_[gcs_client_->client_table().GetLocalClientId()].Acquire(
+            cpu_resources));
+  } else {
+    // In this case, we simply don't reacquire the CPU resources for the worker.
+    // The worker can keep running and when the task finishes, it will simply
+    // not have any CPU resources to release.
+    RAY_LOG(WARNING)
+        << "Resources oversubscribed: "
+        << cluster_resource_map_[gcs_client_->client_table().GetLocalClientId()]
+               .GetAvailableResources()
+               .ToString();
+  }
+
+  // Mark the task as running again.
+  local_queues_.QueueRunningTasks({task});
+  worker->MarkUnblocked();
+}
+
 void NodeManager::HandleRemoteDependencyRequired(const ObjectID &dependency_id) {
   // Try to fetch the object from the object manager.
   RAY_CHECK_OK(object_manager_.Pull(dependency_id));
@@ -872,8 +936,7 @@ void NodeManager::AssignTask(Task &task) {
 void NodeManager::FinishAssignedTask(Worker &worker) {
   TaskID task_id = worker.GetAssignedTaskId();
   RAY_LOG(DEBUG) << "Finished task " << task_id;
-  auto tasks = local_queues_.RemoveTasks({task_id});
-  auto task = *tasks.begin();
+  const auto task = local_queues_.RemoveTask(task_id);
 
   if (task.GetTaskSpecification().IsActorCreationTask()) {
     // If this was an actor creation task, then convert the worker to an actor.
@@ -936,9 +999,15 @@ void NodeManager::HandleObjectLocal(const ObjectID &object_id) {
     std::unordered_set<TaskID> ready_task_id_set(ready_task_ids.begin(),
                                                  ready_task_ids.end());
     // Transition tasks from waiting to scheduled.
-    local_queues_.MoveTasks(ready_task_id_set, WAITING, READY);
+    local_queues_.MoveTasks(ready_task_id_set, TaskState::WAITING, TaskState::READY);
     // New scheduled tasks appeared in the queue, try to dispatch them.
     DispatchTasks();
+
+    // Check that remaining tasks that could not be transitioned are blocked
+    // workers or drivers.
+    local_queues_.FilterState(ready_task_id_set, TaskState::BLOCKED);
+    local_queues_.FilterState(ready_task_id_set, TaskState::DRIVER);
+    RAY_CHECK(ready_task_id_set.empty());
   }
 }
 
@@ -952,8 +1021,15 @@ void NodeManager::HandleObjectMissing(const ObjectID &object_id) {
     // runnable once the deleted object becomes available again.
     std::unordered_set<TaskID> waiting_task_id_set(waiting_task_ids.begin(),
                                                    waiting_task_ids.end());
-    auto waiting_tasks = local_queues_.RemoveTasks(waiting_task_id_set);
-    local_queues_.QueueWaitingTasks(std::vector<Task>(waiting_tasks));
+    local_queues_.MoveTasks(waiting_task_id_set, TaskState::READY, TaskState::WAITING);
+
+    // Check that remaining tasks that could not be transitioned are running
+    // workers or drivers, now blocked in a get.
+    local_queues_.FilterState(waiting_task_id_set, TaskState::RUNNING);
+    if (!waiting_task_id_set.empty()) {
+      RAY_CHECK(waiting_task_id_set.size() == 1);
+      RAY_CHECK(waiting_task_id_set.begin()->is_nil());
+    }
   }
 }
 

src/ray/raylet/node_manager.h

@@ -90,6 +90,10 @@ class NodeManager {
   /// Dispatch locally scheduled tasks. This attempts the transition from "scheduled" to
   /// "running" task state.
   void DispatchTasks();
+  /// Handle a worker becoming blocked in a `ray.get`.
+  void HandleWorkerBlocked(std::shared_ptr<Worker> worker);
+  /// Handle a worker exiting a `ray.get`.
+  void HandleWorkerUnblocked(std::shared_ptr<Worker> worker);
 
   /// Methods for actor scheduling.
   /// Handler for the creation of an actor, possibly on a remote node.

src/ray/raylet/scheduling_queue.cc

@@ -2,6 +2,45 @@
 
 #include "ray/status.h"
 
+namespace {
+
+// Helper function to remove tasks in the given set of task_ids from a
+// queue, and append them to the given vector removed_tasks.
+void RemoveTasksFromQueue(std::list<ray::raylet::Task> &queue,
+                          std::unordered_set<ray::TaskID> &task_ids,
+                          std::vector<ray::raylet::Task> &removed_tasks) {
+  for (auto it = queue.begin(); it != queue.end();) {
+    auto task_id = task_ids.find(it->GetTaskSpecification().TaskId());
+    if (task_id != task_ids.end()) {
+      task_ids.erase(task_id);
+      removed_tasks.push_back(std::move(*it));
+      it = queue.erase(it);
+    } else {
+      it++;
+    }
+  }
+}
+
+// Helper function to queue the given tasks to the given queue.
+inline void QueueTasks(std::list<ray::raylet::Task> &queue,
+                       const std::vector<ray::raylet::Task> &tasks) {
+  queue.insert(queue.end(), tasks.begin(), tasks.end());
+}
+
+// Helper function to filter out tasks of a given state.
+inline void FilterStateFromQueue(const std::list<ray::raylet::Task> &queue,
+                                 std::unordered_set<ray::TaskID> &task_ids,
+                                 ray::raylet::TaskState filter_state) {
+  for (auto it = queue.begin(); it != queue.end(); it++) {
+    auto task_id = task_ids.find(it->GetTaskSpecification().TaskId());
+    if (task_id != task_ids.end()) {
+      task_ids.erase(task_id);
+    }
+  }
+}
+
+}  // namespace
+
 namespace ray {
 
 namespace raylet {
@@ -30,106 +69,147 @@ const std::list<Task> &SchedulingQueue::GetBlockedTasks() const {
   return this->blocked_tasks_;
 }
 
-// Helper function to remove tasks in the given set of task_ids from a
-// queue, and append them to the given vector removed_tasks.
-void removeTasksFromQueue(std::list<Task> &queue, std::unordered_set<TaskID> &task_ids,
-                          std::vector<Task> &removed_tasks) {
-  for (auto it = queue.begin(); it != queue.end();) {
-    auto task_id = task_ids.find(it->GetTaskSpecification().TaskId());
-    if (task_id != task_ids.end()) {
-      task_ids.erase(task_id);
-      removed_tasks.push_back(std::move(*it));
-      it = queue.erase(it);
-    } else {
-      it++;
+void SchedulingQueue::FilterState(std::unordered_set<TaskID> &task_ids,
+                                  TaskState filter_state) const {
+  switch (filter_state) {
+  case TaskState::PLACEABLE:
+    FilterStateFromQueue(placeable_tasks_, task_ids, filter_state);
+    break;
+  case TaskState::WAITING:
+    FilterStateFromQueue(waiting_tasks_, task_ids, filter_state);
+    break;
+  case TaskState::READY:
+    FilterStateFromQueue(ready_tasks_, task_ids, filter_state);
+    break;
+  case TaskState::RUNNING:
+    FilterStateFromQueue(running_tasks_, task_ids, filter_state);
+    break;
+  case TaskState::BLOCKED:
+    FilterStateFromQueue(blocked_tasks_, task_ids, filter_state);
+    break;
+  case TaskState::DRIVER: {
+    const auto driver_ids = GetDriverTaskIds();
+    for (auto it = task_ids.begin(); it != task_ids.end();) {
+      if (driver_ids.count(*it) == 1) {
+        it = task_ids.erase(it);
+      } else {
+        it++;
+      }
     }
+  } break;
+  default:
+    RAY_LOG(FATAL) << "Attempting to filter tasks on unrecognized state "
+                   << static_cast<std::underlying_type<TaskState>::type>(filter_state);
   }
 }
 
-// Helper function to queue the given tasks to the given queue.
-inline void queueTasks(std::list<Task> &queue, const std::vector<Task> &tasks) {
-  queue.insert(queue.end(), tasks.begin(), tasks.end());
-}
-
-std::vector<Task> SchedulingQueue::RemoveTasks(std::unordered_set<TaskID> task_ids) {
+std::vector<Task> SchedulingQueue::RemoveTasks(std::unordered_set<TaskID> &task_ids) {
   // List of removed tasks to be returned.
   std::vector<Task> removed_tasks;
 
-  // Try to find the tasks to remove from the waiting tasks.
-  removeTasksFromQueue(uncreated_actor_methods_, task_ids, removed_tasks);
-  removeTasksFromQueue(waiting_tasks_, task_ids, removed_tasks);
-  removeTasksFromQueue(placeable_tasks_, task_ids, removed_tasks);
-  removeTasksFromQueue(ready_tasks_, task_ids, removed_tasks);
-  removeTasksFromQueue(running_tasks_, task_ids, removed_tasks);
-  removeTasksFromQueue(blocked_tasks_, task_ids, removed_tasks);
-  // TODO(swang): Remove from running methods.
+  // Try to find the tasks to remove from the queues.
+  RemoveTasksFromQueue(uncreated_actor_methods_, task_ids, removed_tasks);
+  RemoveTasksFromQueue(waiting_tasks_, task_ids, removed_tasks);
+  RemoveTasksFromQueue(placeable_tasks_, task_ids, removed_tasks);
+  RemoveTasksFromQueue(ready_tasks_, task_ids, removed_tasks);
+  RemoveTasksFromQueue(running_tasks_, task_ids, removed_tasks);
+  RemoveTasksFromQueue(blocked_tasks_, task_ids, removed_tasks);
 
   RAY_CHECK(task_ids.size() == 0);
   return removed_tasks;
 }
 
-void SchedulingQueue::MoveTasks(std::unordered_set<TaskID> task_ids, TaskState src_state,
+Task SchedulingQueue::RemoveTask(const TaskID &task_id) {
+  std::unordered_set<TaskID> task_id_set = {task_id};
+  auto task = RemoveTasks(task_id_set).front();
+  RAY_CHECK(task.GetTaskSpecification().TaskId() == task_id);
+  return task;
+}
+
+void SchedulingQueue::MoveTasks(std::unordered_set<TaskID> &task_ids, TaskState src_state,
                                 TaskState dst_state) {
   // TODO(atumanov): check the states first to ensure the move is transactional.
   std::vector<Task> removed_tasks;
   // Remove the tasks from the specified source queue.
   switch (src_state) {
-  case PLACEABLE:
-    removeTasksFromQueue(placeable_tasks_, task_ids, removed_tasks);
+  case TaskState::PLACEABLE:
+    RemoveTasksFromQueue(placeable_tasks_, task_ids, removed_tasks);
     break;
-  case WAITING:
-    removeTasksFromQueue(waiting_tasks_, task_ids, removed_tasks);
+  case TaskState::WAITING:
+    RemoveTasksFromQueue(waiting_tasks_, task_ids, removed_tasks);
     break;
-  case READY:
-    removeTasksFromQueue(ready_tasks_, task_ids, removed_tasks);
+  case TaskState::READY:
+    RemoveTasksFromQueue(ready_tasks_, task_ids, removed_tasks);
     break;
-  case RUNNING:
-    removeTasksFromQueue(running_tasks_, task_ids, removed_tasks);
+  case TaskState::RUNNING:
+    RemoveTasksFromQueue(running_tasks_, task_ids, removed_tasks);
+    break;
+  case TaskState::BLOCKED:
+    RemoveTasksFromQueue(blocked_tasks_, task_ids, removed_tasks);
     break;
   default:
-    RAY_LOG(ERROR) << "Attempting to move tasks from unrecognized state " << src_state;
+    RAY_LOG(FATAL) << "Attempting to move tasks from unrecognized state "
+                   << static_cast<std::underlying_type<TaskState>::type>(src_state);
   }
   // Add the tasks to the specified destination queue.
   switch (dst_state) {
-  case PLACEABLE:
-    queueTasks(placeable_tasks_, removed_tasks);
+  case TaskState::PLACEABLE:
+    QueueTasks(placeable_tasks_, removed_tasks);
     break;
-  case WAITING:
-    queueTasks(waiting_tasks_, removed_tasks);
+  case TaskState::WAITING:
+    QueueTasks(waiting_tasks_, removed_tasks);
     break;
-  case READY:
-    queueTasks(ready_tasks_, removed_tasks);
+  case TaskState::READY:
+    QueueTasks(ready_tasks_, removed_tasks);
     break;
-  case RUNNING:
-    queueTasks(running_tasks_, removed_tasks);
+  case TaskState::RUNNING:
+    QueueTasks(running_tasks_, removed_tasks);
+    break;
+  case TaskState::BLOCKED:
+    QueueTasks(blocked_tasks_, removed_tasks);
     break;
   default:
-    RAY_LOG(ERROR) << "Attempting to move tasks to unrecognized state " << dst_state;
+    RAY_LOG(FATAL) << "Attempting to move tasks to unrecognized state "
+                   << static_cast<std::underlying_type<TaskState>::type>(dst_state);
   }
 }
 
 void SchedulingQueue::QueueUncreatedActorMethods(const std::vector<Task> &tasks) {
-  queueTasks(uncreated_actor_methods_, tasks);
+  QueueTasks(uncreated_actor_methods_, tasks);
 }
 
 void SchedulingQueue::QueueWaitingTasks(const std::vector<Task> &tasks) {
-  queueTasks(waiting_tasks_, tasks);
+  QueueTasks(waiting_tasks_, tasks);
 }
 
 void SchedulingQueue::QueuePlaceableTasks(const std::vector<Task> &tasks) {
-  queueTasks(placeable_tasks_, tasks);
+  QueueTasks(placeable_tasks_, tasks);
 }
 
 void SchedulingQueue::QueueReadyTasks(const std::vector<Task> &tasks) {
-  queueTasks(ready_tasks_, tasks);
+  QueueTasks(ready_tasks_, tasks);
 }
 
 void SchedulingQueue::QueueRunningTasks(const std::vector<Task> &tasks) {
-  queueTasks(running_tasks_, tasks);
+  QueueTasks(running_tasks_, tasks);
 }
 
 void SchedulingQueue::QueueBlockedTasks(const std::vector<Task> &tasks) {
-  queueTasks(blocked_tasks_, tasks);
+  QueueTasks(blocked_tasks_, tasks);
+}
+
+void SchedulingQueue::AddDriverTaskId(const TaskID &driver_id) {
+  auto inserted = driver_task_ids_.insert(driver_id);
+  RAY_CHECK(inserted.second);
+}
+
+void SchedulingQueue::RemoveDriverTaskId(const TaskID &driver_id) {
+  auto erased = driver_task_ids_.erase(driver_id);
+  RAY_CHECK(erased == 1);
+}
+
+const std::unordered_set<TaskID> &SchedulingQueue::GetDriverTaskIds() const {
+  return driver_task_ids_;
 }
 
 }  // namespace raylet

src/ray/raylet/scheduling_queue.h

@@ -12,7 +12,8 @@ namespace ray {
 
 namespace raylet {
 
-enum TaskState { INIT, PLACEABLE, WAITING, READY, RUNNING };
+enum class TaskState { INIT, PLACEABLE, WAITING, READY, RUNNING, BLOCKED, DRIVER };
+
 /// \class SchedulingQueue
 ///
 /// Encapsulates task queues. Each queue represents a scheduling state for a
@@ -67,12 +68,31 @@ class SchedulingQueue {
   /// at runtime.
   const std::list<Task> &GetBlockedTasks() const;
 
+  /// Get the set of driver task IDs.
+  ///
+  /// \return A const reference to the set of driver task IDs. These are empty
+  /// tasks used to represent drivers.
+  const std::unordered_set<TaskID> &GetDriverTaskIds() const;
+
   /// Remove tasks from the task queue.
   ///
   /// \param tasks The set of task IDs to remove from the queue. The
-  ///        corresponding tasks must be contained in the queue.
+  /// corresponding tasks must be contained in the queue. The IDs of removed
+  /// tasks will be erased from the set.
   /// \return A vector of the tasks that were removed.
-  std::vector<Task> RemoveTasks(std::unordered_set<TaskID> tasks);
+  std::vector<Task> RemoveTasks(std::unordered_set<TaskID> &tasks);
+
+  /// Remove a task from the task queue.
+  ///
+  /// \param task_id The task ID to remove from the queue. The corresponding
+  /// task must be contained in the queue.
+  /// \return The task that was removed.
+  Task RemoveTask(const TaskID &task_id);
+
+  /// Remove a driver task ID. This is an empty task used to represent a driver.
+  ///
+  /// \param The driver task ID to remove.
+  void RemoveDriverTaskId(const TaskID &task_id);
 
   /// Queue tasks that are destined for actors that have not yet been created.
   ///
@@ -107,15 +127,30 @@ class SchedulingQueue {
   /// \param tasks The tasks to queue.
   void QueueBlockedTasks(const std::vector<Task> &tasks);
 
-  /// \brief Move the specified tasks from the source state to the destination state.
+  /// Add a driver task ID. This is an empty task used to represent a driver.
   ///
-  /// \param tasks The set of task IDs to move.
-  /// \param src_state Source state, which corresponds to one of the internal task queues.
-  /// \param dst_state Destination state, corresponding to one of the internal task
-  /// queues.
-  void MoveTasks(std::unordered_set<TaskID> tasks, TaskState src_state,
+  /// \param The driver task ID to add.
+  void AddDriverTaskId(const TaskID &task_id);
+
+  /// \brief Move the specified tasks from the source state to the destination
+  /// state.
+  ///
+  /// \param tasks The set of task IDs to move. The IDs of successfully moved
+  /// tasks will be erased from the set.
+  /// \param src_state Source state, which corresponds to one of the internal
+  /// task queues.
+  /// \param dst_state Destination state, corresponding to one of the internal
+  /// task queues.
+  void MoveTasks(std::unordered_set<TaskID> &tasks, TaskState src_state,
                  TaskState dst_state);
 
+  /// \brief Filter out task IDs based on their scheduling state.
+  ///
+  /// \param task_ids The set of task IDs to filter. All tasks that have the
+  /// given filter_state will be removed from this set.
+  /// \param filter_state The task state to filter out.
+  void FilterState(std::unordered_set<TaskID> &task_ids, TaskState filter_state) const;
+
  private:
   /// Tasks that are destined for actors that have not yet been created.
   std::list<Task> uncreated_actor_methods_;
@@ -131,6 +166,9 @@ class SchedulingQueue {
   /// Tasks that were dispatched to a worker but are blocked on a data
   /// dependency that was missing at runtime.
   std::list<Task> blocked_tasks_;
+  /// The set of currently running driver tasks. These are empty tasks that are
+  /// started by a driver process on initialization.
+  std::unordered_set<TaskID> driver_task_ids_;
 };
 
 }  // namespace raylet

src/ray/raylet/worker_pool.cc

@@ -5,6 +5,36 @@
 #include "ray/status.h"
 #include "ray/util/logging.h"
 
+namespace {
+
+// A helper function to remove a worker from a list. Returns true if the worker
+// was found and removed.
+std::shared_ptr<ray::raylet::Worker> GetWorker(
+    const std::list<std::shared_ptr<ray::raylet::Worker>> &worker_pool,
+    const std::shared_ptr<ray::LocalClientConnection> &connection) {
+  for (auto it = worker_pool.begin(); it != worker_pool.end(); it++) {
+    if ((*it)->Connection() == connection) {
+      return (*it);
+    }
+  }
+  return nullptr;
+}
+
+// A helper function to remove a worker from a list. Returns true if the worker
+// was found and removed.
+bool RemoveWorker(std::list<std::shared_ptr<ray::raylet::Worker>> &worker_pool,
+                  const std::shared_ptr<ray::raylet::Worker> &worker) {
+  for (auto it = worker_pool.begin(); it != worker_pool.end(); it++) {
+    if (*it == worker) {
+      worker_pool.erase(it);
+      return true;
+    }
+  }
+  return false;
+}
+
+}  // namespace
+
 namespace ray {
 
 namespace raylet {
@@ -99,14 +129,19 @@ void WorkerPool::RegisterWorker(std::shared_ptr<Worker> worker) {
   }
 }
 
+void WorkerPool::RegisterDriver(std::shared_ptr<Worker> driver) {
+  RAY_CHECK(!driver->GetAssignedTaskId().is_nil());
+  registered_drivers_.push_back(driver);
+}
+
 std::shared_ptr<Worker> WorkerPool::GetRegisteredWorker(
     const std::shared_ptr<LocalClientConnection> &connection) const {
-  for (auto it = registered_workers_.begin(); it != registered_workers_.end(); it++) {
-    if ((*it)->Connection() == connection) {
-      return (*it);
-    }
-  }
-  return nullptr;
+  return GetWorker(registered_workers_, connection);
+}
+
+std::shared_ptr<Worker> WorkerPool::GetRegisteredDriver(
+    const std::shared_ptr<LocalClientConnection> &connection) const {
+  return GetWorker(registered_drivers_, connection);
 }
 
 void WorkerPool::PushWorker(std::shared_ptr<Worker> worker) {
@@ -138,22 +173,13 @@ std::shared_ptr<Worker> WorkerPool::PopWorker(const ActorID &actor_id) {
   return worker;
 }
 
-// A helper function to remove a worker from a list. Returns true if the worker
-// was found and removed.
-bool removeWorker(std::list<std::shared_ptr<Worker>> &worker_pool,
-                  const std::shared_ptr<Worker> &worker) {
-  for (auto it = worker_pool.begin(); it != worker_pool.end(); it++) {
-    if (*it == worker) {
-      worker_pool.erase(it);
-      return true;
-    }
-  }
-  return false;
+bool WorkerPool::DisconnectWorker(std::shared_ptr<Worker> worker) {
+  RAY_CHECK(RemoveWorker(registered_workers_, worker));
+  return RemoveWorker(pool_, worker);
 }
 
-bool WorkerPool::DisconnectWorker(std::shared_ptr<Worker> worker) {
-  RAY_CHECK(removeWorker(registered_workers_, worker));
-  return removeWorker(pool_, worker);
+void WorkerPool::DisconnectDriver(std::shared_ptr<Worker> driver) {
+  RAY_CHECK(RemoveWorker(registered_drivers_, driver));
 }
 
 }  // namespace raylet

src/ray/raylet/worker_pool.h

@@ -52,6 +52,11 @@ class WorkerPool {
   /// \param The Worker to be registered.
   void RegisterWorker(std::shared_ptr<Worker> worker);
 
+  /// Register a new driver.
+  ///
+  /// \param The driver to be registered.
+  void RegisterDriver(const std::shared_ptr<Worker> worker);
+
   /// Get the client connection's registered worker.
   ///
   /// \param The client connection owned by a registered worker.
@@ -60,12 +65,25 @@ class WorkerPool {
   std::shared_ptr<Worker> GetRegisteredWorker(
       const std::shared_ptr<LocalClientConnection> &connection) const;
 
+  /// Get the client connection's registered driver.
+  ///
+  /// \param The client connection owned by a registered driver.
+  /// \return The Worker that owns the given client connection. Returns nullptr
+  /// if the client has not registered a driver.
+  std::shared_ptr<Worker> GetRegisteredDriver(
+      const std::shared_ptr<LocalClientConnection> &connection) const;
+
   /// Disconnect a registered worker.
   ///
   /// \param The worker to disconnect. The worker must be registered.
   /// \return Whether the given worker was in the pool of idle workers.
   bool DisconnectWorker(std::shared_ptr<Worker> worker);
 
+  /// Disconnect a registered driver.
+  ///
+  /// \param The driver to disconnect. The driver must be registered.
+  void DisconnectDriver(std::shared_ptr<Worker> driver);
+
   /// Add an idle worker to the pool.
   ///
   /// \param The idle worker to add.
@@ -105,6 +123,8 @@ class WorkerPool {
   /// idle and executing.
   // TODO(swang): Make this a map to make GetRegisteredWorker faster.
   std::list<std::shared_ptr<Worker>> registered_workers_;
+  /// All drivers that have registered and are still connected.
+  std::list<std::shared_ptr<Worker>> registered_drivers_;
 };
 
 }  // namespace raylet