From a6f91664c1ec3e38b7837c756f2275b8cc5e393a Mon Sep 17 00:00:00 2001
From: Alex Wu <itswu.alex@gmail.com>
Date: Thu, 8 Oct 2020 13:19:01 -0700
Subject: [PATCH] [New Scheduler] Multi tenancy edge case  (#11164)

* .

* refactor

* .

* .

* done?

* .

* .

* .

* lint

* no light heartbeat, no tests, fields 2,3

* .

* manually clang format :(

* .

* .

* test

* .

* .

* task manager heartbeat

* lint

* .

* add reminder

* CR

* CR

* cleanup

* CR

* comment

* lint

* .

* .
---
 src/ray/raylet/node_manager.cc | 26 ++++++++++++++++----------
 1 file changed, 16 insertions(+), 10 deletions(-)

diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc
index 28573b4e7..52b23357e 100644
--- a/src/ray/raylet/node_manager.cc
+++ b/src/ray/raylet/node_manager.cc
@@ -341,7 +341,11 @@ void NodeManager::HandleJobStarted(const JobID &job_id, const JobTableData &job_
     // Tasks of this job may already arrived but failed to pop a worker because the job
     // config is not local yet. So we trigger dispatching again here to try to
     // reschedule these tasks.
-    DispatchTasks(local_queues_.GetReadyTasksByClass());
+    if (new_scheduler_enabled_) {
+      ScheduleAndDispatch();
+    } else {
+      DispatchTasks(local_queues_.GetReadyTasksByClass());
+    }
   }
 }
 
@@ -365,15 +369,17 @@ void NodeManager::HandleJobFinished(const JobID &job_id, const JobTableData &job
     }
   }
 
-  // Remove all tasks for this job from the scheduling queues, mark
-  // the results for these tasks as not required, cancel any attempts
-  // at reconstruction. Note that at this time the workers are likely
-  // alive because of the delay in killing workers.
-  auto tasks_to_remove = local_queues_.GetTaskIdsForJob(job_id);
-  task_dependency_manager_.RemoveTasksAndRelatedObjects(tasks_to_remove);
-  // NOTE(swang): SchedulingQueue::RemoveTasks modifies its argument so we must
-  // call it last.
-  local_queues_.RemoveTasks(tasks_to_remove);
+  if (!new_scheduler_enabled_) {
+    // Remove all tasks for this job from the scheduling queues, mark
+    // the results for these tasks as not required, cancel any attempts
+    // at reconstruction. Note that at this time the workers are likely
+    // alive because of the delay in killing workers.
+    auto tasks_to_remove = local_queues_.GetTaskIdsForJob(job_id);
+    task_dependency_manager_.RemoveTasksAndRelatedObjects(tasks_to_remove);
+    // NOTE(swang): SchedulingQueue::RemoveTasks modifies its argument so we must
+    // call it last.
+    local_queues_.RemoveTasks(tasks_to_remove);
+  }
 }
 
 void NodeManager::Heartbeat() {