Don't unsubscribe dependencies for infeasible tasks. (#3338)

* Make scheduling queues RemoveTasks return task states as well.

* Add test

* Don't unsubscribe for infeasible tasks when spilling over.

* Linting

* Address comments.

.travis.yml

@@ -151,6 +151,7 @@ script:
   - python -m pytest -v test/monitor_test.py
   - python -m pytest -v test/cython_test.py
   - python -m pytest -v test/credis_test.py
+  - python -m pytest -v test/node_manager_test.py
 
   # ray tune tests
   - python python/ray/tune/test/dependency_test.py

src/ray/raylet/node_manager.cc

@@ -433,10 +433,15 @@ void NodeManager::HeartbeatAdded(gcs::AsyncGcsClient *client, const ClientID &cl
   std::unordered_set<TaskID> local_task_ids;
   for (const auto &task_id : decision) {
     // (See design_docs/task_states.rst for the state transition diagram.)
-    const auto task = local_queues_.RemoveTask(task_id);
+    TaskState state;
+    const auto task = local_queues_.RemoveTask(task_id, &state);
     // Since we are spilling back from the ready and waiting queues, we need
     // to unsubscribe the dependencies.
-    task_dependency_manager_.UnsubscribeDependencies(task_id);
+    if (state != TaskState::INFEASIBLE) {
+      // Don't unsubscribe for infeasible tasks because we never subscribed in
+      // the first place.
+      task_dependency_manager_.UnsubscribeDependencies(task_id);
+    }
     // Attempt to forward the task. If this fails to forward the task,
     // the task will be resubmit locally.
     ForwardTaskOrResubmit(task, client_id);

src/ray/raylet/scheduling_queue.cc

@@ -8,12 +8,17 @@ 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(ray::raylet::SchedulingQueue::TaskQueue &queue,
+void RemoveTasksFromQueue(ray::raylet::TaskState task_state,
+                          ray::raylet::SchedulingQueue::TaskQueue &queue,
                           std::unordered_set<ray::TaskID> &task_ids,
-                          std::vector<ray::raylet::Task> &removed_tasks) {
+                          std::vector<ray::raylet::Task> &removed_tasks,
+                          std::vector<ray::raylet::TaskState> *task_states = nullptr) {
   for (auto it = task_ids.begin(); it != task_ids.end();) {
     if (queue.RemoveTask(*it, &removed_tasks)) {
       it = task_ids.erase(it);
+      if (task_states != nullptr) {
+        task_states->push_back(task_state);
+      }
     } else {
       it++;
     }
@@ -197,25 +202,43 @@ void SchedulingQueue::FilterState(std::unordered_set<TaskID> &task_ids,
   }
 }
 
-std::vector<Task> SchedulingQueue::RemoveTasks(std::unordered_set<TaskID> &task_ids) {
+std::vector<Task> SchedulingQueue::RemoveTasks(std::unordered_set<TaskID> &task_ids,
+                                               std::vector<TaskState> *task_states) {
   // List of removed tasks to be returned.
   std::vector<Task> removed_tasks;
 
   // Try to find the tasks to remove from the queues.
-  RemoveTasksFromQueue(methods_waiting_for_actor_creation_, 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(infeasible_tasks_, task_ids, removed_tasks);
+
+  RemoveTasksFromQueue(TaskState::WAITING_FOR_ACTOR, methods_waiting_for_actor_creation_,
+                       task_ids, removed_tasks, task_states);
+  RemoveTasksFromQueue(TaskState::WAITING, waiting_tasks_, task_ids, removed_tasks,
+                       task_states);
+  RemoveTasksFromQueue(TaskState::PLACEABLE, placeable_tasks_, task_ids, removed_tasks,
+                       task_states);
+  RemoveTasksFromQueue(TaskState::READY, ready_tasks_, task_ids, removed_tasks,
+                       task_states);
+  RemoveTasksFromQueue(TaskState::RUNNING, running_tasks_, task_ids, removed_tasks,
+                       task_states);
+  RemoveTasksFromQueue(TaskState::INFEASIBLE, infeasible_tasks_, task_ids, removed_tasks,
+                       task_states);
 
   RAY_CHECK(task_ids.size() == 0);
+  if (task_states != nullptr) {
+    RAY_CHECK(removed_tasks.size() == task_states->size());
+  }
   return removed_tasks;
 }
 
-Task SchedulingQueue::RemoveTask(const TaskID &task_id) {
+Task SchedulingQueue::RemoveTask(const TaskID &task_id, TaskState *task_state) {
   std::unordered_set<TaskID> task_id_set = {task_id};
-  auto task = RemoveTasks(task_id_set).front();
+  std::vector<TaskState> task_state_vector;
+  auto const task = RemoveTasks(task_id_set, &task_state_vector).front();
+
+  RAY_CHECK(task_state_vector.size() == 1);
+  if (task_state != nullptr) {
+    *task_state = task_state_vector[0];
+  }
+
   RAY_CHECK(task.GetTaskSpecification().TaskId() == task_id);
   return task;
 }
@@ -224,22 +247,24 @@ void SchedulingQueue::MoveTasks(std::unordered_set<TaskID> &task_ids, TaskState
                                 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 TaskState::PLACEABLE:
-    RemoveTasksFromQueue(placeable_tasks_, task_ids, removed_tasks);
+    RemoveTasksFromQueue(TaskState::PLACEABLE, placeable_tasks_, task_ids, removed_tasks);
     break;
   case TaskState::WAITING:
-    RemoveTasksFromQueue(waiting_tasks_, task_ids, removed_tasks);
+    RemoveTasksFromQueue(TaskState::WAITING, waiting_tasks_, task_ids, removed_tasks);
     break;
   case TaskState::READY:
-    RemoveTasksFromQueue(ready_tasks_, task_ids, removed_tasks);
+    RemoveTasksFromQueue(TaskState::READY, ready_tasks_, task_ids, removed_tasks);
     break;
   case TaskState::RUNNING:
-    RemoveTasksFromQueue(running_tasks_, task_ids, removed_tasks);
+    RemoveTasksFromQueue(TaskState::RUNNING, running_tasks_, task_ids, removed_tasks);
     break;
   case TaskState::INFEASIBLE:
-    RemoveTasksFromQueue(infeasible_tasks_, task_ids, removed_tasks);
+    RemoveTasksFromQueue(TaskState::INFEASIBLE, infeasible_tasks_, task_ids,
+                         removed_tasks);
     break;
   default:
     RAY_LOG(FATAL) << "Attempting to move tasks from unrecognized state "

src/ray/raylet/scheduling_queue.h

@@ -34,7 +34,10 @@ enum class TaskState {
   DRIVER,
   // The task has resources that cannot be satisfied by any node, as far as we
   // know.
-  INFEASIBLE
+  INFEASIBLE,
+  // The task is an actor method and is waiting to learn where the actor was
+  // created.
+  WAITING_FOR_ACTOR,
 };
 
 /// \class SchedulingQueue
@@ -118,15 +121,20 @@ class SchedulingQueue {
   /// \param tasks The set of task IDs to remove from the queue. The
   /// corresponding tasks must be contained in the queue. The IDs of removed
   /// tasks will be erased from the set.
+  /// \param task_states If this is not nullptr, then, the states of the removed
+  /// tasks will be appended to this vector.
   /// \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> &task_ids,
+                                std::vector<TaskState> *task_states = nullptr);
 
   /// 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.
+  /// \param task_state If this is not nullptr, then the state of the removed
+  /// task will be written here.
   /// \return The task that was removed.
-  Task RemoveTask(const TaskID &task_id);
+  Task RemoveTask(const TaskID &task_id, TaskState *task_state = nullptr);
 
   /// Remove a driver task ID. This is an empty task used to represent a driver.
   ///

test/node_manager_test.py

@@ -0,0 +1,63 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import pytest
+
+import ray
+from ray.test.cluster_utils import Cluster
+from ray.test.test_utils import run_string_as_driver
+
+
+@pytest.fixture()
+def ray_start_empty_cluster():
+    cluster = Cluster()
+    yield cluster
+
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+    cluster.shutdown()
+
+
+# This tests the queue transitions for infeasible tasks. This has been an issue
+# in the past, e.g., https://github.com/ray-project/ray/issues/3275.
+def test_infeasible_tasks(ray_start_empty_cluster):
+    cluster = ray_start_empty_cluster
+
+    @ray.remote
+    def f():
+        return
+
+    cluster.add_node(resources={str(0): 100})
+    ray.init(redis_address=cluster.redis_address)
+
+    # Submit an infeasible task.
+    x_id = f._submit(args=[], kwargs={}, resources={str(1): 1})
+
+    # Add a node that makes the task feasible and make sure we can get the
+    # result.
+    cluster.add_node(resources={str(1): 100})
+    ray.get(x_id)
+
+    # Start a driver that submits an infeasible task and then let it exit.
+    driver_script = """
+import ray
+
+ray.init(redis_address="{}")
+
+@ray.remote(resources={})
+def f():
+{}pass  # This is a weird hack to insert some blank space.
+
+f.remote()
+""".format(cluster.redis_address, "{str(2): 1}", "    ")
+
+    run_string_as_driver(driver_script)
+
+    # Now add a new node that makes the task feasible.
+    cluster.add_node(resources={str(2): 100})
+
+    # Make sure we can still run tasks on all nodes.
+    ray.get([
+        f._submit(args=[], kwargs={}, resources={str(i): 1}) for i in range(3)
+    ])