[Core] Multi-tenancy: Kill idle workers in FIFO order (#10597)

* Kill idle workers in FIFO order

* Update test

* minor update

* Address comments

* fix after merge

* fix worker_pool_test

src/ray/raylet/node_manager.cc

@@ -1370,9 +1370,9 @@ void NodeManager::HandleWorkerAvailable(const std::shared_ptr<WorkerInterface> &
     DispatchTasks(local_queues_.GetReadyTasksByClass());
   }
   if (RayConfig::instance().enable_multi_tenancy()) {
-    // If the worker remains idle after scheduling, we may kill it to ensure the
-    // registered workers are in a reasonable size.
-    worker_pool_.TryKillingIdleWorker(worker);
+    // We trigger killing here instead of inside `Worker::PushWorker` because we
+    // only kill an idle worker if it remains idle after scheduling.
+    worker_pool_.TryKillingIdleWorkers();
   }
 }
 

src/ray/raylet/worker_pool.cc

@@ -165,16 +165,6 @@ WorkerPool::~WorkerPool() {
   }
 }
 
-uint32_t WorkerPool::Size(const Language &language) const {
-  const auto state = states_by_lang_.find(language);
-  if (state == states_by_lang_.end()) {
-    return 0;
-  } else {
-    return static_cast<uint32_t>(state->second.idle.size() +
-                                 state->second.idle_actor.size());
-  }
-}
-
 Process WorkerPool::StartWorkerProcess(const Language &language,
                                        const rpc::WorkerType worker_type,
                                        const JobID &job_id,
@@ -635,68 +625,83 @@ void WorkerPool::PushWorker(const std::shared_ptr<WorkerInterface> &worker) {
     // Put the worker to the corresponding idle pool.
     if (worker->GetActorId().IsNil()) {
       state.idle.insert(worker);
+      if (RayConfig::instance().enable_multi_tenancy()) {
+        idle_of_all_languages.push_back(worker);
+      }
     } else {
       state.idle_actor[worker->GetActorId()] = worker;
     }
   }
 }
 
-void WorkerPool::TryKillingIdleWorker(std::shared_ptr<WorkerInterface> worker) {
-  auto &worker_state = GetStateForLanguage(worker->GetLanguage());
-  if (worker_state.pending_unregistration_workers.count(worker) > 0) {
-    // This worker has already been killed.
-    // This is possible because a Java worker process may hold multiple workers.
-    return;
-  }
-
-  auto running_size = GetAllRegisteredWorkers().size();
-  for (const auto &entry : states_by_lang_) {
-    running_size -= entry.second.pending_unregistration_workers.size();
-  }
-  if (running_size <= static_cast<size_t>(num_workers_soft_limit_)) {
-    return;
-  }
-
-  auto worker_id = worker->WorkerId();
-  const auto pid = worker->GetProcess().GetId();
-  if (worker_state.idle.count(worker) == 0) {
-    return;
-  }
-  if (worker_state.starting_worker_processes.count(worker->GetProcess()) > 0) {
-    // A Java worker process may hold multiple workers.
-    RAY_LOG(DEBUG) << "Some workers of pid " << pid
-                   << " are pending registration. Skip killing worker " << worker_id;
-    return;
-  }
-
-  // Make sure all workers in this worker process are idle.
-  // This block of code is needed by Java workers.
-  std::unordered_set<std::shared_ptr<WorkerInterface>> workers_in_the_same_process;
-  for (const auto &worker_in_the_same_process : worker_state.registered_workers) {
-    if (worker_in_the_same_process->GetProcess().GetId() == pid) {
-      if (worker_state.idle.count(worker_in_the_same_process) == 0) {
-        // Another worker in this process isn't idle, so this process can't be killed.
-        return;
-      } else {
-        workers_in_the_same_process.insert(worker_in_the_same_process);
-      }
+void WorkerPool::TryKillingIdleWorkers() {
+  size_t running_size = 0;
+  for (const auto worker : GetAllRegisteredWorkers()) {
+    if (!worker->IsDead()) {
+      running_size++;
     }
   }
 
-  for (auto worker_it = workers_in_the_same_process.begin();
-       worker_it != workers_in_the_same_process.end(); worker_it++) {
-    RAY_LOG(INFO) << "The worker pool has " << running_size
-                  << " registered workers which exceeds the soft limit of "
-                  << num_workers_soft_limit_ << ", and worker "
-                  << (*worker_it)->WorkerId() << " with pid " << pid
-                  << " is idle. Kill it.";
-    // Remove the worker from the idle pool so it can't be popped anymore. However, we
-    // don't remove it from the registered pool because we want the worker to go through
-    // the normal disconnection logic in Node Manager.
-    RemoveWorker(worker_state.idle, *worker_it);
-    worker_state.pending_unregistration_workers.insert(*worker_it);
+  // Kill idle workers in FIFO order.
+  for (auto it = idle_of_all_languages.begin();
+       it != idle_of_all_languages.end() &&
+       running_size > static_cast<size_t>(num_workers_soft_limit_);
+       it++) {
+    if ((*it)->IsDead()) {
+      // This worker has already been killed.
+      // This is possible because a Java worker process may hold multiple workers.
+      continue;
+    }
+    auto process = (*it)->GetProcess();
+
+    auto &worker_state = GetStateForLanguage((*it)->GetLanguage());
+
+    if (worker_state.starting_worker_processes.count(process) > 0) {
+      // A Java worker process may hold multiple workers.
+      // Some workers of this process are pending registration. Skip killing this worker.
+      continue;
+    }
+
+    // Make sure all workers in this worker process are idle.
+    // This block of code is needed by Java workers.
+    auto workers_in_the_same_process = GetWorkersByProcess(process);
+    bool can_be_killed = true;
+    for (const auto &worker : workers_in_the_same_process) {
+      if (worker_state.idle.count(worker) == 0) {
+        // Another worker in this process isn't idle, so this process can't be killed.
+        can_be_killed = false;
+        break;
+      }
+    }
+    if (!can_be_killed) {
+      continue;
+    }
+
+    for (auto worker_it = workers_in_the_same_process.begin();
+         worker_it != workers_in_the_same_process.end(); worker_it++) {
+      RAY_LOG(INFO) << "The worker pool has " << running_size
+                    << " registered workers which exceeds the soft limit of "
+                    << num_workers_soft_limit_ << ", and worker "
+                    << (*worker_it)->WorkerId() << " with pid " << process.GetId()
+                    << " is idle. Kill it.";
+      // Remove the worker from the idle pool so it can't be popped anymore.
+      RemoveWorker(worker_state.idle, *worker_it);
+      if (!(*worker_it)->IsDead()) {
+        (*worker_it)->MarkDead();
+        running_size--;
+      }
+    }
+    process.Kill();
   }
-  worker->GetProcess().Kill();
+
+  std::list<std::shared_ptr<WorkerInterface>> new_idle_of_all_languages;
+  for (auto it = idle_of_all_languages.begin(); it != idle_of_all_languages.end(); it++) {
+    if (!(*it)->IsDead()) {
+      new_idle_of_all_languages.push_back(*it);
+    }
+  }
+
+  idle_of_all_languages = std::move(new_idle_of_all_languages);
 }
 
 std::shared_ptr<WorkerInterface> WorkerPool::PopWorker(
@@ -741,12 +746,19 @@ std::shared_ptr<WorkerInterface> WorkerPool::PopWorker(
       }
     } else {
       // Find an available worker which is already assigned to this job.
-      for (auto it = state.idle.begin(); it != state.idle.end(); it++) {
-        if ((*it)->GetAssignedJobId() != task_spec.JobId()) {
+      // Try to pop the most recently pushed worker.
+      for (auto it = idle_of_all_languages.rbegin(); it != idle_of_all_languages.rend();
+           it++) {
+        if (task_spec.GetLanguage() != (*it)->GetLanguage() ||
+            (*it)->GetAssignedJobId() != task_spec.JobId()) {
           continue;
         }
-        worker = std::move(*it);
-        state.idle.erase(it);
+        state.idle.erase(*it);
+        // We can't erase a reverse_iterator.
+        auto lit = it.base();
+        lit--;
+        worker = std::move(*lit);
+        idle_of_all_languages.erase(lit);
         break;
       }
       if (worker == nullptr) {
@@ -779,7 +791,6 @@ std::shared_ptr<WorkerInterface> WorkerPool::PopWorker(
 bool WorkerPool::DisconnectWorker(const std::shared_ptr<WorkerInterface> &worker) {
   auto &state = GetStateForLanguage(worker->GetLanguage());
   RAY_CHECK(RemoveWorker(state.registered_workers, worker));
-  RemoveWorker(state.pending_unregistration_workers, worker);
 
   stats::CurrentWorker().Record(
       0, {{stats::LanguageKey, Language_Name(worker->GetLanguage())},
@@ -911,6 +922,20 @@ void WorkerPool::TryStartIOWorkers(const Language &language, State &state) {
   }
 }
 
+std::unordered_set<std::shared_ptr<WorkerInterface>> WorkerPool::GetWorkersByProcess(
+    const Process &process) {
+  std::unordered_set<std::shared_ptr<WorkerInterface>> workers_of_process;
+  for (auto &entry : states_by_lang_) {
+    auto &worker_state = entry.second;
+    for (const auto &worker : worker_state.registered_workers) {
+      if (worker->GetProcess().GetId() == process.GetId()) {
+        workers_of_process.insert(worker);
+      }
+    }
+  }
+  return workers_of_process;
+}
+
 std::string WorkerPool::DebugString() const {
   std::stringstream result;
   result << "WorkerPool:";
@@ -920,6 +945,7 @@ std::string WorkerPool::DebugString() const {
     result << "\n- num " << Language_Name(entry.first)
            << " drivers: " << entry.second.registered_drivers.size();
   }
+  result << "- num idle workers: " << idle_of_all_languages.size();
   return result.str();
 }
 

src/ray/raylet/worker_pool.h

@@ -181,10 +181,9 @@ class WorkerPool : public WorkerPoolInterface {
   /// \param The idle worker to add.
   void PushWorker(const std::shared_ptr<WorkerInterface> &worker);
 
-  /// Try to kill the worker if it's idle.
-  ///
-  /// \param worker The worker to be killed.
-  void TryKillingIdleWorker(std::shared_ptr<WorkerInterface> worker);
+  /// Try killing idle workers to ensure the running workers are in a
+  /// reasonable size.
+  void TryKillingIdleWorkers();
 
   /// Pop an idle worker from the pool. The caller is responsible for pushing
   /// the worker back onto the pool once the worker has completed its work.
@@ -298,10 +297,6 @@ class WorkerPool : public WorkerPoolInterface {
     std::unordered_set<std::shared_ptr<WorkerInterface>> registered_workers;
     /// All drivers that have registered and are still connected.
     std::unordered_set<std::shared_ptr<WorkerInterface>> registered_drivers;
-    /// All workers that have been killed but been unregistered yet.
-    /// This field is used to calculate the size of running workers when trying to kill an
-    /// idle worker.
-    std::unordered_set<std::shared_ptr<WorkerInterface>> pending_unregistration_workers;
     /// A map from the pids of starting worker processes
     /// to the number of their unregistered workers.
     std::unordered_map<Process, int> starting_worker_processes;
@@ -362,6 +357,13 @@ class WorkerPool : public WorkerPoolInterface {
   /// started.
   void TryStartIOWorkers(const Language &language, State &state);
 
+  /// Get all workers of the given process.
+  ///
+  /// \param process The process of workers.
+  /// \return The workers of the given process.
+  std::unordered_set<std::shared_ptr<WorkerInterface>> GetWorkersByProcess(
+      const Process &process);
+
   /// For Process class for managing subprocesses (e.g. reaping zombies).
   boost::asio::io_service *io_service_;
   /// The soft limit of the number of registered workers.
@@ -397,6 +399,10 @@ class WorkerPool : public WorkerPoolInterface {
 
   /// This map tracks the latest infos of unfinished jobs.
   absl::flat_hash_map<JobID, rpc::JobConfig> unfinished_jobs_;
+
+  /// The pool of idle non-actor workers of all languages. This is used to kill idle
+  /// workers in FIFO order.
+  std::list<std::shared_ptr<WorkerInterface>> idle_of_all_languages;
 };
 
 }  // namespace raylet