[xray] Sophisticated task dependency management (#2035)

.travis.yml

@@ -162,6 +162,7 @@ install:
   - ./src/ray/raylet/task_test
   - ./src/ray/raylet/worker_pool_test
   - ./src/ray/raylet/lineage_cache_test
+  - ./src/ray/raylet/task_dependency_manager_test
 
   - bash ../../../src/common/test/run_tests.sh
   - bash ../../../src/plasma/test/run_tests.sh

src/ray/object_manager/object_directory.cc

@@ -9,24 +9,34 @@ ObjectDirectory::ObjectDirectory(std::shared_ptr<gcs::AsyncGcsClient> gcs_client
 ray::Status ObjectDirectory::ReportObjectAdded(const ObjectID &object_id,
                                                const ClientID &client_id,
                                                const ObjectInfoT &object_info) {
-  // TODO(hme): Determine whether we need to do lookup to append.
-  JobID job_id = JobID::from_random();
+  // Append the addition entry to the object table.
+  JobID job_id = JobID::nil();
   auto data = std::make_shared<ObjectTableDataT>();
   data->manager = client_id.binary();
   data->is_eviction = false;
+  data->num_evictions = object_evictions_[object_id];
   data->object_size = object_info.data_size;
-  ray::Status status = gcs_client_->object_table().Append(
-      job_id, object_id, data,
-      [](gcs::AsyncGcsClient *client, const UniqueID &id, const ObjectTableDataT &data) {
-        // Do nothing.
-      });
+  ray::Status status =
+      gcs_client_->object_table().Append(job_id, object_id, data, nullptr);
   return status;
 };
 
 ray::Status ObjectDirectory::ReportObjectRemoved(const ObjectID &object_id,
                                                  const ClientID &client_id) {
-  // TODO(hme): Need corresponding remove method in GCS.
-  return ray::Status::NotImplemented("ObjectTable.Remove is not implemented");
+  // Append the eviction entry to the object table.
+  JobID job_id = JobID::nil();
+  auto data = std::make_shared<ObjectTableDataT>();
+  data->manager = client_id.binary();
+  data->is_eviction = true;
+  data->num_evictions = object_evictions_[object_id];
+  ray::Status status =
+      gcs_client_->object_table().Append(job_id, object_id, data, nullptr);
+  // Increment the number of times we've evicted this object. NOTE(swang): This
+  // is only necessary because the Ray redis module expects unique entries in a
+  // log. We track the number of evictions so that the next eviction, if there
+  // is one, is unique.
+  object_evictions_[object_id]++;
+  return status;
 };
 
 ray::Status ObjectDirectory::GetInformation(const ClientID &client_id,

src/ray/object_manager/object_directory.h

@@ -126,6 +126,9 @@ class ObjectDirectory : public ObjectDirectoryInterface {
   std::unordered_map<ObjectID, ODCallbacks> existing_requests_;
   /// Reference to the gcs client.
   std::shared_ptr<gcs::AsyncGcsClient> gcs_client_;
+  /// Map from object ID to the number of times it's been evicted on this
+  /// node before.
+  std::unordered_map<ObjectID, int> object_evictions_;
 };
 
 }  // namespace ray

src/ray/object_manager/object_manager.h

@@ -43,8 +43,15 @@ struct ObjectManagerConfig {
   std::string store_socket_name;
 };
 
+class ObjectManagerInterface {
+ public:
+  virtual ray::Status Pull(const ObjectID &object_id) = 0;
+  virtual ray::Status Cancel(const ObjectID &object_id) = 0;
+  virtual ~ObjectManagerInterface(){};
+};
+
 // TODO(hme): Add success/failure callbacks for push and pull.
-class ObjectManager {
+class ObjectManager : public ObjectManagerInterface {
  public:
   /// Implicitly instantiates Ray implementation of ObjectDirectory.
   ///

src/ray/raylet/CMakeLists.txt

@@ -25,6 +25,7 @@ ADD_RAY_TEST(worker_pool_test STATIC_LINK_LIBS ray_static ${PLASMA_STATIC_LIB} $
 
 ADD_RAY_TEST(task_test STATIC_LINK_LIBS ray_static gtest gtest_main gmock_main pthread ${Boost_SYSTEM_LIBRARY})
 ADD_RAY_TEST(lineage_cache_test STATIC_LINK_LIBS ray_static gtest gtest_main gmock_main pthread ${Boost_SYSTEM_LIBRARY})
+ADD_RAY_TEST(task_dependency_manager_test STATIC_LINK_LIBS ray_static gtest gtest_main gmock_main pthread ${Boost_SYSTEM_LIBRARY})
 
 add_library(rayletlib raylet.cc ${NODE_MANAGER_FBS_OUTPUT_FILES})
 target_link_libraries(rayletlib ray_static ${Boost_SYSTEM_LIBRARY})

src/ray/raylet/node_manager.cc

@@ -74,10 +74,7 @@ NodeManager::NodeManager(boost::asio::io_service &io_service,
       local_queues_(SchedulingQueue()),
       scheduling_policy_(local_queues_),
       reconstruction_policy_([this](const TaskID &task_id) { ResubmitTask(task_id); }),
-      task_dependency_manager_(
-          object_manager,
-          // reconstruction_policy_,
-          [this](const TaskID &task_id) { HandleWaitingTaskReady(task_id); }),
+      task_dependency_manager_(object_manager),
       lineage_cache_(gcs_client_->client_table().GetLocalClientId(),
                      gcs_client->raylet_task_table(), gcs_client->raylet_task_table()),
       remote_clients_(),
@@ -88,6 +85,13 @@ NodeManager::NodeManager(boost::asio::io_service &io_service,
   ClientID local_client_id = gcs_client_->client_table().GetLocalClientId();
   cluster_resource_map_.emplace(local_client_id,
                                 SchedulingResources(config.resource_config));
+
+  RAY_CHECK_OK(object_manager_.SubscribeObjAdded([this](const ObjectInfoT &object_info) {
+    ObjectID object_id = ObjectID::from_binary(object_info.object_id);
+    HandleObjectLocal(object_id);
+  }));
+  RAY_CHECK_OK(object_manager_.SubscribeObjDeleted(
+      [this](const ObjectID &object_id) { HandleObjectMissing(object_id); }));
 }
 
 ray::Status NodeManager::RegisterGcs() {
@@ -374,6 +378,9 @@ void NodeManager::ProcessClientMessage(
     auto message = flatbuffers::GetRoot<protocol::ReconstructObject>(message_data);
     ObjectID object_id = from_flatbuf(*message->object_id());
     RAY_LOG(DEBUG) << "reconstructing object " << 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 the blocked client is a worker, and the worker isn't already blocked,
@@ -468,13 +475,6 @@ void NodeManager::ProcessNodeManagerMessage(TcpClientConnection &node_manager_cl
   node_manager_client.ProcessMessages();
 }
 
-void NodeManager::HandleWaitingTaskReady(const TaskID &task_id) {
-  auto ready_tasks = local_queues_.RemoveTasks({task_id});
-  local_queues_.QueueReadyTasks(std::vector<Task>(ready_tasks));
-  // Schedule the newly ready tasks if possible.
-  ScheduleTasks();
-}
-
 void NodeManager::ScheduleTasks() {
   // This method performs the transition of tasks from PENDING to SCHEDULED.
   auto policy_decision = scheduling_policy_.Schedule(
@@ -500,23 +500,32 @@ void NodeManager::ScheduleTasks() {
       RAY_CHECK(1 == tasks.size());
       Task &task = tasks.front();
       // TODO(swang): Handle forward task failure.
-      // TODO(swang): Unsubscribe this task in the task dependency manager.
       RAY_CHECK_OK(ForwardTask(task, client_id));
     }
+    // Notify the task dependency manager that we no longer need this task's
+    // object dependencies.
+    // NOTE(swang): For local tasks, the scheduled task's dependencies may get
+    // evicted before it can be assigned to a worker.
+    task_dependency_manager_.UnsubscribeDependencies(task_id);
   }
 
   // Transition locally scheduled tasks to SCHEDULED and dispatch scheduled tasks.
-  std::vector<Task> tasks = local_queues_.RemoveTasks(local_task_ids);
-  local_queues_.QueueScheduledTasks(tasks);
-  DispatchTasks();
+  if (local_task_ids.size() > 0) {
+    std::vector<Task> tasks = local_queues_.RemoveTasks(local_task_ids);
+    local_queues_.QueueScheduledTasks(tasks);
+    DispatchTasks();
+  }
 }
 
 void NodeManager::SubmitTask(const Task &task, const Lineage &uncommitted_lineage) {
-  const TaskSpecification &spec = task.GetTaskSpecification();
-
   // Add the task and its uncommitted lineage to the lineage cache.
   lineage_cache_.AddWaitingTask(task, uncommitted_lineage);
+  // Mark the task as pending. Once the task has finished execution, or once it
+  // has been forwarded to another node, the task must be marked as canceled in
+  // the TaskDependencyManager.
+  task_dependency_manager_.TaskPending(task);
 
+  const TaskSpecification &spec = task.GetTaskSpecification();
   if (spec.IsActorTask()) {
     // Check whether we know the location of the actor.
     const auto actor_entry = actor_registry_.find(spec.ActorId());
@@ -561,14 +570,31 @@ void NodeManager::SubmitTask(const Task &task, const Lineage &uncommitted_lineag
   }
 }
 
+void NodeManager::HandleRemoteDependencyRequired(const ObjectID &dependency_id) {
+  // Try to fetch the object from the object manager.
+  RAY_CHECK_OK(object_manager_.Pull(dependency_id));
+  // TODO(swang): Request reconstruction of the object, possibly after a
+  // timeout.
+}
+
+void NodeManager::HandleRemoteDependencyCanceled(const ObjectID &dependency_id) {
+  // Cancel the fetch request from the object manager.
+  RAY_CHECK_OK(object_manager_.Cancel(dependency_id));
+  // TODO(swang): Cancel reconstruction of the object.
+}
+
 void NodeManager::QueueTask(const Task &task) {
-  // Queue the task depending on the availability of its arguments.
-  if (task_dependency_manager_.TaskReady(task)) {
-    local_queues_.QueueReadyTasks(std::vector<Task>({task}));
+  // Subscribe to the task's dependencies.
+  bool ready = task_dependency_manager_.SubscribeDependencies(
+      task.GetTaskSpecification().TaskId(), task.GetDependencies());
+  // Queue the task. If all dependencies are available, then the task is queued
+  // in the READY state, else the WAITING.
+  if (ready) {
+    local_queues_.QueueReadyTasks({task});
+    // Try to schedule the newly ready task.
     ScheduleTasks();
   } else {
-    local_queues_.QueueWaitingTasks(std::vector<Task>({task}));
-    task_dependency_manager_.SubscribeTaskReady(task);
+    local_queues_.QueueWaitingTasks({task});
   }
 }
 
@@ -687,9 +713,12 @@ void NodeManager::FinishAssignedTask(Worker &worker) {
   if (task.GetTaskSpecification().IsActorCreationTask() ||
       task.GetTaskSpecification().IsActorTask()) {
     auto dummy_object = task.GetTaskSpecification().ActorDummyObject();
-    task_dependency_manager_.MarkDependencyReady(dummy_object);
+    HandleObjectLocal(dummy_object);
   }
 
+  // Notify the task dependency manager that this task has finished execution.
+  task_dependency_manager_.TaskCanceled(task_id);
+
   // Unset the worker's assigned task.
   worker.AssignTaskId(TaskID::nil());
 }
@@ -698,6 +727,36 @@ void NodeManager::ResubmitTask(const TaskID &task_id) {
   throw std::runtime_error("Method not implemented");
 }
 
+void NodeManager::HandleObjectLocal(const ObjectID &object_id) {
+  // Notify the task dependency manager that this object is local.
+  const auto ready_task_ids = task_dependency_manager_.HandleObjectLocal(object_id);
+  // Transition the tasks whose dependencies are now fulfilled to the ready
+  // state.
+  if (ready_task_ids.size() > 0) {
+    std::unordered_set<TaskID> ready_task_id_set(ready_task_ids.begin(),
+                                                 ready_task_ids.end());
+    auto ready_tasks = local_queues_.RemoveTasks(ready_task_id_set);
+    local_queues_.QueueReadyTasks(std::vector<Task>(ready_tasks));
+    // Schedule the newly ready tasks.
+    ScheduleTasks();
+  }
+}
+
+void NodeManager::HandleObjectMissing(const ObjectID &object_id) {
+  // Notify the task dependency manager that this object is no longer local.
+  const auto waiting_task_ids = task_dependency_manager_.HandleObjectMissing(object_id);
+  // Transition any tasks that were in the runnable state and are dependent on
+  // this object to the waiting state.
+  if (!waiting_task_ids.empty()) {
+    // Transition the tasks back to the waiting state. They will be made
+    // 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));
+  }
+}
+
 ray::Status NodeManager::ForwardTask(const Task &task, const ClientID &node_id) {
   const auto &spec = task.GetTaskSpecification();
   auto task_id = spec.TaskId();
@@ -733,7 +792,9 @@ ray::Status NodeManager::ForwardTask(const Task &task, const ClientID &node_id)
     // lineage cache since the receiving node is now responsible for writing
     // the task to the GCS.
     lineage_cache_.RemoveWaitingTask(task_id);
-
+    // Notify the task dependency manager that we are no longer responsible
+    // for executing this task.
+    task_dependency_manager_.TaskCanceled(task_id);
     // Preemptively push any local arguments to the receiving node. For now, we
     // only do this with actor tasks, since actor tasks must be executed by a
     // specific process and therefore have affinity to the receiving node.

src/ray/raylet/node_manager.h

@@ -57,11 +57,16 @@ class NodeManager {
   ray::Status RegisterGcs();
 
  private:
-  // Handler for the addition of a new GCS client.
+  /// Methods for handling clients.
+  /// Handler for the addition of a new GCS client.
   void ClientAdded(const ClientTableDataT &data);
-  // Handler for the creation of an actor, possibly on a remote node.
-  void HandleActorCreation(const ActorID &actor_id,
-                           const std::vector<ActorTableDataT> &data);
+  /// Send heartbeats to the GCS.
+  void Heartbeat();
+  /// Handler for a heartbeat notification from the GCS.
+  void HeartbeatAdded(gcs::AsyncGcsClient *client, const ClientID &id,
+                      const HeartbeatTableDataT &data);
+
+  /// Methods for task scheduling.
   // Queue a task for local execution.
   void QueueTask(const Task &task);
   /// Submit a task to this node.
@@ -72,25 +77,35 @@ class NodeManager {
   void FinishAssignedTask(Worker &worker);
   /// Schedule tasks.
   void ScheduleTasks();
-  /// Handle a task whose local dependencies were missing and are now available.
-  void HandleWaitingTaskReady(const TaskID &task_id);
   /// Resubmit a task whose return value needs to be reconstructed.
   void ResubmitTask(const TaskID &task_id);
   /// Forward a task to another node to execute. The task is assumed to not be
   /// queued in local_queues_.
   ray::Status ForwardTask(const Task &task, const ClientID &node_id);
-  /// Send heartbeats to the GCS.
-  void Heartbeat();
-  /// Handler for a notification about a new client from the GCS.
-  void ClientAdded(gcs::AsyncGcsClient *client, const UniqueID &id,
-                   const ClientTableDataT &data);
-  /// Handler for a heartbeat notification from the GCS.
-  void HeartbeatAdded(gcs::AsyncGcsClient *client, const ClientID &id,
-                      const HeartbeatTableDataT &data);
   /// Dispatch locally scheduled tasks. This attempts the transition from "scheduled" to
   /// "running" task state.
   void DispatchTasks();
 
+  /// Methods for actor scheduling.
+  /// Handler for the creation of an actor, possibly on a remote node.
+  void HandleActorCreation(const ActorID &actor_id,
+                           const std::vector<ActorTableDataT> &data);
+
+  /// Methods for managing object dependencies.
+  /// Handle a dependency required by a queued task that is missing locally.
+  /// The dependency is (1) on a remote node, (2) pending creation on a remote
+  /// node, or (3) missing from all nodes and requires reconstruction.
+  void HandleRemoteDependencyRequired(const ObjectID &dependency_id);
+  /// Handle a dependency that was previously required by a queued task that is
+  /// no longer required.
+  void HandleRemoteDependencyCanceled(const ObjectID &dependency_id);
+  /// Handle an object becoming local. This updates any local accounting, but
+  /// does not write to any global accounting in the GCS.
+  void HandleObjectLocal(const ObjectID &object_id);
+  /// Handle an object that is no longer local. This updates any local
+  /// accounting, but does not write to any global accounting in the GCS.
+  void HandleObjectMissing(const ObjectID &object_id);
+
   boost::asio::io_service &io_service_;
   ObjectManager &object_manager_;
   /// A client connection to the GCS.

src/ray/raylet/task_dependency_manager.cc

@@ -4,114 +4,225 @@ namespace ray {
 
 namespace raylet {
 
-TaskDependencyManager::TaskDependencyManager(
-    ObjectManager &object_manager,
-    // ReconstructionPolicy &reconstruction_policy,
-    std::function<void(const TaskID &)> handler)
-    : object_manager_(object_manager),
-      // reconstruction_policy_(reconstruction_policy),
-      task_ready_callback_(handler) {
-  // TODO(swang): Check return status.
-  ray::Status status =
-      object_manager_.SubscribeObjAdded([this](const ObjectInfoT &object_info) {
-        handleObjectReady(ObjectID::from_binary(object_info.object_id));
-      });
-  // TODO(swang): Subscribe to object removed notifications.
-}
+TaskDependencyManager::TaskDependencyManager(ObjectManagerInterface &object_manager)
+    : object_manager_(object_manager) {}
 
 bool TaskDependencyManager::CheckObjectLocal(const ObjectID &object_id) const {
   return local_objects_.count(object_id) == 1;
 }
 
-bool TaskDependencyManager::argumentsReady(const std::vector<ObjectID> arguments) const {
-  for (auto &argument : arguments) {
-    // Check if any argument is missing.
-    if (local_objects_.count(argument) == 0) {
-      return false;
-    }
+bool TaskDependencyManager::CheckObjectRequired(const ObjectID &object_id) const {
+  const TaskID task_id = ComputeTaskId(object_id);
+  auto task_entry = required_tasks_.find(task_id);
+  // If there are no subscribed tasks that are dependent on the object, then do
+  // nothing.
+  if (task_entry == required_tasks_.end()) {
+    return false;
+  }
+  if (task_entry->second.count(object_id) == 0) {
+    return false;
+  }
+  // If the object is already local, then the dependency is fulfilled. Do
+  // nothing.
+  if (local_objects_.count(object_id) == 1) {
+    return false;
+  }
+  // If the task that creates the object is pending execution, then the
+  // dependency will be fulfilled locally. Do nothing.
+  if (pending_tasks_.count(task_id) == 1) {
+    return false;
   }
-  // All arguments are ready.
   return true;
 }
 
-void TaskDependencyManager::handleObjectReady(const ray::ObjectID &object_id) {
-  RAY_LOG(DEBUG) << "object ready " << object_id;
-  // Add the object to the table of locally available objects.
-  RAY_CHECK(local_objects_.count(object_id) == 0);
-  local_objects_.insert(object_id);
+void TaskDependencyManager::HandleRemoteDependencyRequired(const ObjectID &object_id) {
+  bool required = CheckObjectRequired(object_id);
+  // If the object is required, then try to make the object available locally.
+  if (required) {
+    auto inserted = required_objects_.insert(object_id);
+    if (inserted.second) {
+      // If we haven't already, request the object manager to pull it from a
+      // remote node.
+      RAY_CHECK_OK(object_manager_.Pull(object_id));
+    }
+  }
+}
 
-  // Handle any tasks that were dependent on the newly available object.
+void TaskDependencyManager::HandleRemoteDependencyCanceled(const ObjectID &object_id) {
+  bool required = CheckObjectRequired(object_id);
+  // If the object is no longer required, then cancel the object.
+  if (!required) {
+    auto it = required_objects_.find(object_id);
+    if (it != required_objects_.end()) {
+      RAY_CHECK_OK(object_manager_.Cancel(object_id));
+      required_objects_.erase(it);
+    }
+  }
+}
+
+std::vector<TaskID> TaskDependencyManager::HandleObjectLocal(
+    const ray::ObjectID &object_id) {
+  RAY_LOG(DEBUG) << "object ready " << object_id.hex();
+  // Add the object to the table of locally available objects.
+  auto inserted = local_objects_.insert(object_id);
+  RAY_CHECK(inserted.second);
+
+  // Find any tasks that are dependent on the newly available object.
   std::vector<TaskID> ready_task_ids;
-  auto dependent_tasks = remote_object_dependencies_.find(object_id);
-  if (dependent_tasks != remote_object_dependencies_.end()) {
-    for (auto &dependent_task_id : dependent_tasks->second) {
-      // If the dependent task now has all of its arguments ready, it's ready
-      // to run.
-      if (argumentsReady(task_dependencies_[dependent_task_id])) {
-        ready_task_ids.push_back(dependent_task_id);
+  auto creating_task_entry = required_tasks_.find(ComputeTaskId(object_id));
+  if (creating_task_entry != required_tasks_.end()) {
+    auto object_entry = creating_task_entry->second.find(object_id);
+    if (object_entry != creating_task_entry->second.end()) {
+      for (auto &dependent_task_id : object_entry->second) {
+        auto &task_entry = task_dependencies_[dependent_task_id];
+        task_entry.num_missing_dependencies--;
+        // If the dependent task now has all of its arguments ready, it's ready
+        // to run.
+        if (task_entry.num_missing_dependencies == 0) {
+          ready_task_ids.push_back(dependent_task_id);
+        }
       }
     }
-    remote_object_dependencies_.erase(dependent_tasks);
   }
+
+  // The object is now local, so cancel any in-progress operations to make the
+  // object local.
+  HandleRemoteDependencyCanceled(object_id);
+
+  return ready_task_ids;
+}
+
+std::vector<TaskID> TaskDependencyManager::HandleObjectMissing(
+    const ray::ObjectID &object_id) {
+  // Add the object to the table of locally available objects.
+  auto erased = local_objects_.erase(object_id);
+  RAY_CHECK(erased == 1);
+
+  // Find any tasks that are dependent on the missing object.
+  std::vector<TaskID> waiting_task_ids;
+  TaskID creating_task_id = ComputeTaskId(object_id);
+  auto creating_task_entry = required_tasks_.find(creating_task_id);
+  if (creating_task_entry != required_tasks_.end()) {
+    auto object_entry = creating_task_entry->second.find(object_id);
+    if (object_entry != creating_task_entry->second.end()) {
+      for (auto &dependent_task_id : object_entry->second) {
+        auto &task_entry = task_dependencies_[dependent_task_id];
+        // If the dependent task had all of its arguments ready, it was ready to
+        // run but must be switched to waiting since one of its arguments is now
+        // missing.
+        if (task_entry.num_missing_dependencies == 0) {
+          waiting_task_ids.push_back(dependent_task_id);
+        }
+        task_entry.num_missing_dependencies++;
+      }
+    }
+  }
+  // The object is no longer local. Try to make the object local if necessary.
+  HandleRemoteDependencyRequired(object_id);
   // Process callbacks for all of the tasks dependent on the object that are
   // now ready to run.
-  for (auto &ready_task_id : ready_task_ids) {
-    UnsubscribeTaskReady(ready_task_id);
-    task_ready_callback_(ready_task_id);
-  }
+  return waiting_task_ids;
 }
 
-bool TaskDependencyManager::TaskReady(const Task &task) const {
-  const std::vector<ObjectID> arguments = task.GetDependencies();
-  return argumentsReady(arguments);
-}
+bool TaskDependencyManager::SubscribeDependencies(
+    const TaskID &task_id, const std::vector<ObjectID> &required_objects) {
+  auto &task_entry = task_dependencies_[task_id];
 
-void TaskDependencyManager::SubscribeTaskReady(const Task &task) {
-  // TODO(swang): Don't pull arguments that are going to be created by a queued
-  // or running task.
-  TaskID task_id = task.GetTaskSpecification().TaskId();
-  const std::vector<ObjectID> arguments = task.GetDependencies();
-  // Add the task's arguments to the table of subscribed tasks.
-  task_dependencies_[task_id] = arguments;
-  // Add the task's remote arguments to the table of remote objects.
-  int num_missing_arguments = 0;
-  for (auto &argument : arguments) {
-    if (local_objects_.count(argument) == 0) {
-      remote_object_dependencies_[argument].push_back(task_id);
-      num_missing_arguments++;
-      // TODO(swang): Check return status.
-      // TODO(swang): Handle Pull failure (if object manager does not retry).
-      // TODO(atumanov): pull return status should be propagated back to the caller.
-      ray::Status status = object_manager_.Pull(argument);
+  // Record the task's dependencies.
+  for (const auto &object_id : required_objects) {
+    auto inserted = task_entry.object_dependencies.insert(object_id);
+    if (inserted.second) {
+      // Get the ID of the task that creates the dependency.
+      TaskID creating_task_id = ComputeTaskId(object_id);
+      // Determine whether the dependency can be fulfilled by the local node.
+      if (local_objects_.count(object_id) == 0) {
+        // The object is not local.
+        task_entry.num_missing_dependencies++;
+      }
+      // Add the subscribed task to the mapping from object ID to list of
+      // dependent tasks.
+      required_tasks_[creating_task_id][object_id].push_back(task_id);
     }
   }
-  // Check that the task has some missing arguments.
-  RAY_CHECK(num_missing_arguments > 0);
+
+  // These dependencies are required by the given task. Try to make them local
+  // if necessary.
+  for (const auto &object_id : required_objects) {
+    HandleRemoteDependencyRequired(object_id);
+  }
+
+  // Return whether all dependencies are local.
+  return (task_entry.num_missing_dependencies == 0);
 }
 
-void TaskDependencyManager::UnsubscribeTaskReady(const TaskID &task_id) {
-  const std::vector<ObjectID> arguments = task_dependencies_[task_id];
+void TaskDependencyManager::UnsubscribeDependencies(const TaskID &task_id) {
   // Remove the task from the table of subscribed tasks.
-  task_dependencies_.erase(task_id);
-  // Remove the task from the table of remote objects to dependent tasks.
-  for (auto &argument : arguments) {
-    if (local_objects_.count(argument) == 1) {
-      continue;
-    }
-    // The argument is not local. Remove the task from the list of tasks that
-    // are dependent on this object.
-    std::vector<TaskID> &dependent_tasks = remote_object_dependencies_[task_id];
-    for (auto it = dependent_tasks.begin(); it != dependent_tasks.end(); it++) {
-      if (*it == task_id) {
-        it = dependent_tasks.erase(it);
-        break;
+  auto it = task_dependencies_.find(task_id);
+  RAY_CHECK(it != task_dependencies_.end());
+  const TaskDependencies task_entry = std::move(it->second);
+  task_dependencies_.erase(it);
+
+  // Remove the task's dependencies.
+  for (const auto &object_id : task_entry.object_dependencies) {
+    // Remove the task from the list of tasks that are dependent on this
+    // object.
+    // Get the ID of the task that creates the dependency.
+    TaskID creating_task_id = ComputeTaskId(object_id);
+    auto creating_task_entry = required_tasks_.find(creating_task_id);
+    std::vector<TaskID> &dependent_tasks = creating_task_entry->second[object_id];
+    auto it = std::find(dependent_tasks.begin(), dependent_tasks.end(), task_id);
+    RAY_CHECK(it != dependent_tasks.end());
+    dependent_tasks.erase(it);
+    // If the unsubscribed task was the only task dependent on the object, then
+    // erase the object entry.
+    if (dependent_tasks.empty()) {
+      creating_task_entry->second.erase(object_id);
+      // Remove the task that creates this object if there are no more object
+      // dependencies created by the task.
+      if (creating_task_entry->second.empty()) {
+        required_tasks_.erase(creating_task_entry);
       }
     }
   }
+
+  // These dependencies are no longer required by the given task. Cancel any
+  // in-progress operations to make them local.
+  for (const auto &object_id : task_entry.object_dependencies) {
+    HandleRemoteDependencyCanceled(object_id);
+  }
 }
 
-void TaskDependencyManager::MarkDependencyReady(const ObjectID &object) {
-  handleObjectReady(object);
+void TaskDependencyManager::TaskPending(const Task &task) {
+  TaskID task_id = task.GetTaskSpecification().TaskId();
+
+  // Record that the task is pending execution.
+  pending_tasks_.insert(task_id);
+  // Find any subscribed tasks that are dependent on objects created by the
+  // pending task.
+  auto remote_task_entry = required_tasks_.find(task_id);
+  if (remote_task_entry != required_tasks_.end()) {
+    for (const auto &object_entry : remote_task_entry->second) {
+      // This object created by the pending task will appear locally once the
+      // task completes execution. Cancel any in-progress operations to make
+      // the object local.
+      HandleRemoteDependencyCanceled(object_entry.first);
+    }
+  }
+}
+
+void TaskDependencyManager::TaskCanceled(const TaskID &task_id) {
+  // Record that the task is no longer pending execution.
+  pending_tasks_.erase(task_id);
+  // Find any subscribed tasks that are dependent on objects created by the
+  // canceled task.
+  auto remote_task_entry = required_tasks_.find(task_id);
+  if (remote_task_entry != required_tasks_.end()) {
+    for (const auto &object_entry : remote_task_entry->second) {
+      // This object created by the task will no longer appear locally since
+      // the task is canceled.  Try to make the object local if necessary.
+      HandleRemoteDependencyRequired(object_entry.first);
+    }
+  }
 }
 
 }  // namespace raylet

src/ray/raylet/task_dependency_manager.h

@@ -16,20 +16,18 @@ class ReconstructionPolicy;
 
 /// \class TaskDependencyManager
 ///
-/// Responsible for managing task dependencies. Tasks that have object
-/// dependencies that are missing locally should contact the manager to try to
-/// make their dependencies available.
+/// Responsible for managing object dependencies for tasks.  The caller can
+/// subscribe to object dependencies for a task. The task manager will
+/// determine which object dependencies are remote. These are the objects that
+/// are neither in the local object store, nor will they be created by a
+/// locally queued task. The task manager will request that these objects be
+/// made available locally, either by object transfer from a remote node or
+/// reconstruction. The task manager will also cancel these objects if they are
+/// no longer needed by any task.
 class TaskDependencyManager {
  public:
   /// Create a task dependency manager.
-  ///
-  /// \param object_manager A reference to the object manager so that the task
-  /// dependency manager can issue requests to transfer objects.
-  /// \param handler The handler to call for subscribed tasks whose
-  /// dependencies have become available locally.
-  TaskDependencyManager(ObjectManager &object_manager,
-                        // ReconstructionPolicy &reconstruction_policy,
-                        std::function<void(const TaskID &)> handler);
+  TaskDependencyManager(ObjectManagerInterface &object_manager);
 
   /// Check whether an object is locally available.
   ///
@@ -37,50 +35,109 @@ class TaskDependencyManager {
   /// \return Whether the object is local.
   bool CheckObjectLocal(const ObjectID &object_id) const;
 
-  /// Check whether a task's object dependencies are locally available.
+  /// Subscribe to object depedencies required by the task and check whether
+  /// all dependencies are fulfilled. This will track this task's dependencies
+  /// until UnsubscribeDependencies is called on the same task ID. If any
+  /// dependencies are remote, then they will be requested. When the last
+  /// remote dependency later appears locally via a call to HandleObjectLocal,
+  /// the subscribed task will be returned by the HandleObjectLocal call,
+  /// signifying that it is ready to run. This method may be called multiple
+  /// times per task.
   ///
-  /// \param task The task whose object dependencies will be checked.
-  /// \return Whether the task's object dependencies are ready.
-  bool TaskReady(const Task &task) const;
+  /// \param task_id The ID of the task whose dependencies to subscribe to.
+  /// \param required_objects The objects required by the task.
+  /// \return Whether all of the given dependencies for the given task are
+  /// local.
+  bool SubscribeDependencies(const TaskID &task_id,
+                             const std::vector<ObjectID> &required_objects);
 
-  /// Subscribe to a task that has missing dependencies. The manager will
-  /// attempt to make any missing dependencies available locally by transfer or
-  /// by reconstruction. The registered handler will be called when the task's
-  /// dependencies become locally available.
+  /// Unsubscribe from the object dependencies required by this task. If the
+  /// objects were remote and are no longer required by any subscribed task,
+  /// then they will be canceled.
   ///
-  /// \param task The task with missing dependencies.
-  void SubscribeTaskReady(const Task &task);
+  /// \param task_id The ID of the task whose dependencies to unsubscribe from.
+  void UnsubscribeDependencies(const TaskID &task_id);
 
-  /// Stop waiting for a task's dependencies to become available.
+  /// Mark that the given task is pending execution. Any objects that it creates
+  /// are now considered to be pending creation. If there are any subscribed
+  /// tasks that depend on these objects, then the objects will be canceled.
   ///
-  /// \param task_id The task ID of the task with missing dependencies.
-  void UnsubscribeTaskReady(const TaskID &task_id);
+  /// \param task The task that is pending execution.
+  void TaskPending(const Task &task);
 
-  /// Mark an object as locally available. This is used for objects that do not
-  /// have a stored value (e.g., actor execution dependencies).
+  /// Mark that the given task is no longer pending execution. Any objects that
+  /// it creates that are not already local are now considered to be remote. If
+  /// there are any subscribed tasks that depend on these objects, then the
+  /// objects will be requested.
+  ///
+  /// \param task_id The ID of the task to cancel.
+  void TaskCanceled(const TaskID &task_id);
+
+  /// Handle an object becoming locally available. If there are any subscribed
+  /// tasks that depend on this object, then the object will be canceled.
   ///
   /// \param object_id The object ID of the object to mark as locally
   /// available.
-  void MarkDependencyReady(const ObjectID &object_id);
+  /// \return A list of task IDs. This contains all subscribed tasks that now
+  /// have all of their dependencies fulfilled, once this object was made
+  /// local.
+  std::vector<TaskID> HandleObjectLocal(const ray::ObjectID &object_id);
+
+  /// Handle an object that is no longer locally available. If there are any
+  /// subscribed tasks that depend on this object, then the object will be
+  /// requested.
+  ///
+  /// \param object_id The object ID of the object that was previously locally
+  /// available.
+  /// \return A list of task IDs. This contains all subscribed tasks that
+  /// previously had all of their dependencies fulfilled, but are now missing
+  /// this object dependency.
+  std::vector<TaskID> HandleObjectMissing(const ray::ObjectID &object_id);
 
  private:
-  /// Check whether the given list of objects are ready.
-  bool argumentsReady(const std::vector<ObjectID> arguments) const;
-  /// Handle an object added to the object store.
-  void handleObjectReady(const ray::ObjectID &object_id);
-  /// A reference to the object manager so that we can issue Pull requests of
-  /// missing objects.
-  ObjectManager &object_manager_;
-  /// A mapping from task ID of each subscribed task to its list of
+  using ObjectDependencyMap = std::unordered_map<ray::ObjectID, std::vector<ray::TaskID>>;
+
+  /// A struct to represent the object dependencies of a task.
+  struct TaskDependencies {
+    /// The objects that the task is dependent on. These must be local before
+    /// the task is ready to execute.
+    std::unordered_set<ObjectID> object_dependencies;
+    /// The number of object arguments that are not available locally. This
+    /// must be zero before the task is ready to execute.
+    int64_t num_missing_dependencies;
+  };
+
+  /// Check whether the given object needs to be made available through object
+  /// transfer or reconstruction. These are objects for which: (1) there is a
+  /// subscribed task dependent on it, (2) the object is not local, and (3) the
+  /// task that creates the object is not pending execution locally.
+  bool CheckObjectRequired(const ObjectID &object_id) const;
+  /// If the given object is required, then request that the object be made
+  /// available through object transfer or reconstruction.
+  void HandleRemoteDependencyRequired(const ObjectID &object_id);
+  /// If the given object is no longer required, then cancel any in-progress
+  /// operations to make the object available through object transfer or
+  /// reconstruction.
+  void HandleRemoteDependencyCanceled(const ObjectID &object_id);
+
+  ObjectManagerInterface &object_manager_;
+  /// A mapping from task ID of each subscribed task to its list of object
   /// dependencies.
-  std::unordered_map<ray::TaskID, std::vector<ray::ObjectID>> task_dependencies_;
-  // A mapping from object ID of each object that is not locally available to
-  // the list of subscribed tasks that are dependent on it.
-  std::unordered_map<ray::ObjectID, std::vector<ray::TaskID>> remote_object_dependencies_;
-  // The set of locally available objects.
+  std::unordered_map<ray::TaskID, TaskDependencies> task_dependencies_;
+  /// All tasks whose outputs are required by a subscribed task. This is a
+  /// mapping from task ID to information about the objects that the task
+  /// creates, either by return value or by `ray.put`. For each object, we
+  /// store the IDs of the subscribed tasks that are dependent on the object.
+  std::unordered_map<ray::TaskID, ObjectDependencyMap> required_tasks_;
+  /// Objects that are required by a subscribed task, are not local, and are
+  /// not created by a pending task. For these objects, there are pending
+  /// operations to make the object available.
+  std::unordered_set<ray::ObjectID> required_objects_;
+  /// The set of locally available objects.
   std::unordered_set<ray::ObjectID> local_objects_;
-  // The callback to call when a subscribed task becomes ready.
-  std::function<void(const TaskID &)> task_ready_callback_;
+  /// The set of tasks that are pending execution. Any objects created by these
+  /// tasks that are not already local are pending creation.
+  std::unordered_set<ray::TaskID> pending_tasks_;
 };
 
 }  // namespace raylet

src/ray/raylet/task_dependency_manager_test.cc

@@ -0,0 +1,347 @@
+#include <list>
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+#include <boost/asio.hpp>
+
+#include "ray/raylet/task_dependency_manager.h"
+
+namespace ray {
+
+namespace raylet {
+
+using ::testing::_;
+
+class MockObjectManager : public ObjectManagerInterface {
+ public:
+  MOCK_METHOD1(Pull, ray::Status(const ObjectID &object_id));
+  MOCK_METHOD1(Cancel, ray::Status(const ObjectID &object_id));
+};
+
+class TaskDependencyManagerTest : public ::testing::Test {
+ public:
+  TaskDependencyManagerTest()
+      : object_manager_mock_(), task_dependency_manager_(object_manager_mock_) {}
+
+ protected:
+  MockObjectManager object_manager_mock_;
+  TaskDependencyManager task_dependency_manager_;
+};
+
+static inline Task ExampleTask(const std::vector<ObjectID> &arguments,
+                               int64_t num_returns) {
+  std::unordered_map<std::string, double> required_resources;
+  std::vector<std::shared_ptr<TaskArgument>> task_arguments;
+  for (auto &argument : arguments) {
+    std::vector<ObjectID> references = {argument};
+    task_arguments.emplace_back(std::make_shared<TaskArgumentByReference>(references));
+  }
+  auto spec = TaskSpecification(UniqueID::nil(), UniqueID::from_random(), 0,
+                                UniqueID::from_random(), task_arguments, num_returns,
+                                required_resources);
+  auto execution_spec = TaskExecutionSpecification(std::vector<ObjectID>());
+  execution_spec.IncrementNumForwards();
+  Task task = Task(execution_spec, spec);
+  return task;
+}
+
+std::vector<Task> MakeTaskChain(int chain_size,
+                                const std::vector<ObjectID> &initial_arguments,
+                                int64_t num_returns) {
+  std::vector<Task> task_chain;
+  std::vector<ObjectID> arguments = initial_arguments;
+  for (int i = 0; i < chain_size; i++) {
+    auto task = ExampleTask(arguments, num_returns);
+    task_chain.push_back(task);
+    arguments.clear();
+    for (int j = 0; j < task.GetTaskSpecification().NumReturns(); j++) {
+      arguments.push_back(task.GetTaskSpecification().ReturnId(j));
+    }
+  }
+  return task_chain;
+}
+
+TEST_F(TaskDependencyManagerTest, TestSimpleTask) {
+  // Create a task with 3 arguments.
+  int num_arguments = 3;
+  std::vector<ObjectID> arguments;
+  for (int i = 0; i < num_arguments; i++) {
+    arguments.push_back(ObjectID::from_random());
+  }
+  TaskID task_id = TaskID::from_random();
+  // No objects have been registered in the task dependency manager, so all
+  // arguments should be remote.
+  for (const auto &argument_id : arguments) {
+    EXPECT_CALL(object_manager_mock_, Pull(argument_id));
+  }
+  // Subscribe to the task's dependencies.
+  bool ready = task_dependency_manager_.SubscribeDependencies(task_id, arguments);
+  ASSERT_FALSE(ready);
+
+  // All arguments should be canceled as they become available locally.
+  for (const auto &argument_id : arguments) {
+    EXPECT_CALL(object_manager_mock_, Cancel(argument_id));
+  }
+  // For each argument except the last, tell the task dependency manager that
+  // the argument is local.
+  int i = 0;
+  for (; i < num_arguments - 1; i++) {
+    auto ready_task_ids = task_dependency_manager_.HandleObjectLocal(arguments[i]);
+    ASSERT_TRUE(ready_task_ids.empty());
+  }
+  // Tell the task dependency manager that the last argument is local. Now the
+  // task should be ready to run.
+  auto ready_task_ids = task_dependency_manager_.HandleObjectLocal(arguments[i]);
+  ASSERT_EQ(ready_task_ids.size(), 1);
+  ASSERT_EQ(ready_task_ids.front(), task_id);
+}
+
+TEST_F(TaskDependencyManagerTest, TestDuplicateSubscribe) {
+  // Create a task with 3 arguments.
+  TaskID task_id = TaskID::from_random();
+  int num_arguments = 3;
+  std::vector<ObjectID> arguments;
+  for (int i = 0; i < num_arguments; i++) {
+    // Add the new argument to the list of dependencies to subscribe to.
+    ObjectID argument_id = ObjectID::from_random();
+    arguments.push_back(argument_id);
+    // Subscribe to the task's dependencies. All arguments except the last are
+    // duplicates of previous subscription calls. Each argument should only be
+    // requested from the node manager once.
+    EXPECT_CALL(object_manager_mock_, Pull(argument_id));
+    bool ready = task_dependency_manager_.SubscribeDependencies(task_id, arguments);
+    ASSERT_FALSE(ready);
+  }
+
+  // All arguments should be canceled as they become available locally.
+  for (const auto &argument_id : arguments) {
+    EXPECT_CALL(object_manager_mock_, Cancel(argument_id));
+  }
+  // For each argument except the last, tell the task dependency manager that
+  // the argument is local.
+  int i = 0;
+  for (; i < num_arguments - 1; i++) {
+    auto ready_task_ids = task_dependency_manager_.HandleObjectLocal(arguments[i]);
+    ASSERT_TRUE(ready_task_ids.empty());
+  }
+  // Tell the task dependency manager that the last argument is local. Now the
+  // task should be ready to run.
+  auto ready_task_ids = task_dependency_manager_.HandleObjectLocal(arguments[i]);
+  ASSERT_EQ(ready_task_ids.size(), 1);
+  ASSERT_EQ(ready_task_ids.front(), task_id);
+}
+
+TEST_F(TaskDependencyManagerTest, TestMultipleTasks) {
+  // Create 3 tasks that are dependent on the same object.
+  ObjectID argument_id = ObjectID::from_random();
+  std::vector<TaskID> dependent_tasks;
+  int num_dependent_tasks = 3;
+  // The object should only be requested from the object manager once for all
+  // three tasks.
+  EXPECT_CALL(object_manager_mock_, Pull(argument_id));
+  for (int i = 0; i < num_dependent_tasks; i++) {
+    TaskID task_id = TaskID::from_random();
+    dependent_tasks.push_back(task_id);
+    // Subscribe to each of the task's dependencies.
+    bool ready = task_dependency_manager_.SubscribeDependencies(task_id, {argument_id});
+    ASSERT_FALSE(ready);
+  }
+
+  // Tell the task dependency manager that the object is local.
+  EXPECT_CALL(object_manager_mock_, Cancel(argument_id));
+  auto ready_task_ids = task_dependency_manager_.HandleObjectLocal(argument_id);
+  // Check that all tasks are now ready to run.
+  ASSERT_EQ(ready_task_ids.size(), dependent_tasks.size());
+  for (const auto &task_id : ready_task_ids) {
+    ASSERT_NE(std::find(dependent_tasks.begin(), dependent_tasks.end(), task_id),
+              dependent_tasks.end());
+  }
+}
+
+TEST_F(TaskDependencyManagerTest, TestTaskChain) {
+  // Create 3 tasks, each dependent on the previous. The first task has no
+  // arguments.
+  int num_tasks = 3;
+  auto tasks = MakeTaskChain(num_tasks, {}, 1);
+  int num_ready_tasks = 1;
+  int i = 0;
+  // No objects should be remote or canceled since each task depends on a
+  // locally queued task.
+  EXPECT_CALL(object_manager_mock_, Pull(_)).Times(0);
+  EXPECT_CALL(object_manager_mock_, Cancel(_)).Times(0);
+  for (const auto &task : tasks) {
+    // Subscribe to each of the tasks' arguments.
+    auto arguments = task.GetDependencies();
+    bool ready = task_dependency_manager_.SubscribeDependencies(
+        task.GetTaskSpecification().TaskId(), arguments);
+    if (i < num_ready_tasks) {
+      // The first task should be ready to run since it has no arguments.
+      ASSERT_TRUE(ready);
+    } else {
+      // All remaining tasks depend on the previous task.
+      ASSERT_FALSE(ready);
+    }
+
+    // Mark each task as pending.
+    task_dependency_manager_.TaskPending(task);
+
+    i++;
+  }
+
+  // Simulate executing each task. Each task's completion should make the next
+  // task runnable.
+  while (!tasks.empty()) {
+    auto task = tasks.front();
+    tasks.erase(tasks.begin());
+    TaskID task_id = task.GetTaskSpecification().TaskId();
+    auto return_id = task.GetTaskSpecification().ReturnId(0);
+
+    task_dependency_manager_.UnsubscribeDependencies(task_id);
+    // Simulate the object notifications for the task's return values.
+    auto ready_tasks = task_dependency_manager_.HandleObjectLocal(return_id);
+    if (tasks.empty()) {
+      // If there are no more tasks, then there should be no more tasks that
+      // become ready to run.
+      ASSERT_TRUE(ready_tasks.empty());
+    } else {
+      // If there are more tasks to run, then the next task in the chain should
+      // now be ready to run.
+      ASSERT_EQ(ready_tasks.size(), 1);
+      ASSERT_EQ(ready_tasks.front(), tasks.front().GetTaskSpecification().TaskId());
+    }
+    // Simulate the task finishing execution.
+    task_dependency_manager_.TaskCanceled(task_id);
+  }
+}
+
+TEST_F(TaskDependencyManagerTest, TestDependentPut) {
+  // Create a task with 3 arguments.
+  auto task1 = ExampleTask({}, 0);
+  ObjectID put_id = ComputePutId(task1.GetTaskSpecification().TaskId(), 1);
+  auto task2 = ExampleTask({put_id}, 0);
+
+  // No objects have been registered in the task dependency manager, so the put
+  // object should be remote.
+  EXPECT_CALL(object_manager_mock_, Pull(put_id));
+  // Subscribe to the task's dependencies.
+  bool ready = task_dependency_manager_.SubscribeDependencies(
+      task2.GetTaskSpecification().TaskId(), {put_id});
+  ASSERT_FALSE(ready);
+
+  // The put object should be considered local as soon as the task that creates
+  // it is pending execution.
+  EXPECT_CALL(object_manager_mock_, Cancel(put_id));
+  task_dependency_manager_.TaskPending(task1);
+}
+
+TEST_F(TaskDependencyManagerTest, TestTaskForwarding) {
+  // Create 2 tasks, one dependent on the other. The first has no arguments.
+  int num_tasks = 2;
+  auto tasks = MakeTaskChain(num_tasks, {}, 1);
+  for (const auto &task : tasks) {
+    // Subscribe to each of the tasks' arguments.
+    auto arguments = task.GetDependencies();
+    static_cast<void>(task_dependency_manager_.SubscribeDependencies(
+        task.GetTaskSpecification().TaskId(), arguments));
+    task_dependency_manager_.TaskPending(task);
+  }
+
+  // Get the first task.
+  const auto task = tasks.front();
+  TaskID task_id = task.GetTaskSpecification().TaskId();
+  ObjectID return_id = task.GetTaskSpecification().ReturnId(0);
+  // Simulate forwarding the first task to a remote node.
+  task_dependency_manager_.UnsubscribeDependencies(task_id);
+  // The object returned by the first task should be considered remote once we
+  // cancel the forwarded task, since the second task depends on it.
+  EXPECT_CALL(object_manager_mock_, Pull(return_id));
+  task_dependency_manager_.TaskCanceled(task_id);
+
+  // Simulate the task executing on a remote node and its return value
+  // appearing locally.
+  EXPECT_CALL(object_manager_mock_, Cancel(return_id));
+  auto ready_tasks = task_dependency_manager_.HandleObjectLocal(return_id);
+  // Check that the task that we kept is now ready to run.
+  ASSERT_EQ(ready_tasks.size(), 1);
+  ASSERT_EQ(ready_tasks.front(), tasks.back().GetTaskSpecification().TaskId());
+}
+
+TEST_F(TaskDependencyManagerTest, TestEviction) {
+  // Create a task with 3 arguments.
+  int num_arguments = 3;
+  std::vector<ObjectID> arguments;
+  for (int i = 0; i < num_arguments; i++) {
+    arguments.push_back(ObjectID::from_random());
+  }
+  TaskID task_id = TaskID::from_random();
+  // No objects have been registered in the task dependency manager, so all
+  // arguments should be remote.
+  for (const auto &argument_id : arguments) {
+    EXPECT_CALL(object_manager_mock_, Pull(argument_id));
+  }
+  // Subscribe to the task's dependencies.
+  bool ready = task_dependency_manager_.SubscribeDependencies(task_id, arguments);
+  ASSERT_FALSE(ready);
+
+  // Tell the task dependency manager that each of the arguments is now
+  // available.
+  for (const auto &argument_id : arguments) {
+    EXPECT_CALL(object_manager_mock_, Cancel(argument_id));
+  }
+  for (size_t i = 0; i < arguments.size(); i++) {
+    std::vector<TaskID> ready_tasks;
+    ready_tasks = task_dependency_manager_.HandleObjectLocal(arguments[i]);
+    if (i == arguments.size() - 1) {
+      ASSERT_EQ(ready_tasks.size(), 1);
+      ASSERT_EQ(ready_tasks.front(), task_id);
+    } else {
+      ASSERT_TRUE(ready_tasks.empty());
+    }
+  }
+
+  // Simulate each of the arguments getting evicted. Each object should now be
+  // considered remote.
+  for (const auto &argument_id : arguments) {
+    EXPECT_CALL(object_manager_mock_, Pull(argument_id));
+  }
+  for (size_t i = 0; i < arguments.size(); i++) {
+    std::vector<TaskID> waiting_tasks;
+    waiting_tasks = task_dependency_manager_.HandleObjectMissing(arguments[i]);
+    if (i == 0) {
+      // The first eviction should cause the task to go back to the waiting
+      // state.
+      ASSERT_EQ(waiting_tasks.size(), 1);
+      ASSERT_EQ(waiting_tasks.front(), task_id);
+    } else {
+      // The subsequent evictions shouldn't cause any more tasks to go back to
+      // the waiting state.
+      ASSERT_TRUE(waiting_tasks.empty());
+    }
+  }
+
+  // Tell the task dependency manager that each of the arguments is available
+  // again.
+  for (const auto &argument_id : arguments) {
+    EXPECT_CALL(object_manager_mock_, Cancel(argument_id));
+  }
+  for (size_t i = 0; i < arguments.size(); i++) {
+    std::vector<TaskID> ready_tasks;
+    ready_tasks = task_dependency_manager_.HandleObjectLocal(arguments[i]);
+    if (i == arguments.size() - 1) {
+      ASSERT_EQ(ready_tasks.size(), 1);
+      ASSERT_EQ(ready_tasks.front(), task_id);
+    } else {
+      ASSERT_TRUE(ready_tasks.empty());
+    }
+  }
+}
+
+}  // namespace raylet
+
+}  // namespace ray
+
+int main(int argc, char **argv) {
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}