[Core] Multi-tenancy: Worker capping (#10500)

src/ray/raylet/main.cc

@@ -161,6 +161,10 @@ int main(int argc, char *argv[]) {
           // about this?
           static_resource_conf[resource_name] = std::stod(resource_quantity);
         }
+        auto num_cpus_it = static_resource_conf.find("CPU");
+        int num_cpus = num_cpus_it != static_resource_conf.end()
+                           ? static_cast<int>(num_cpus_it->second)
+                           : 0;
 
         node_manager_config.raylet_config = raylet_config;
         node_manager_config.resource_config =
@@ -170,6 +174,7 @@ int main(int argc, char *argv[]) {
         node_manager_config.node_manager_address = node_ip_address;
         node_manager_config.node_manager_port = node_manager_port;
         node_manager_config.num_initial_workers = num_initial_workers;
+        node_manager_config.num_workers_soft_limit = num_cpus;
         node_manager_config.num_initial_python_workers_for_first_job =
             num_initial_python_workers_for_first_job;
         node_manager_config.maximum_startup_concurrency = maximum_startup_concurrency;
@@ -225,7 +230,6 @@ int main(int argc, char *argv[]) {
         object_manager_config.plasma_directory = plasma_directory;
         object_manager_config.huge_pages = huge_pages;
 
-        int num_cpus = static_cast<int>(static_resource_conf["CPU"]);
         object_manager_config.rpc_service_threads_number =
             std::min(std::max(2, num_cpus / 4), 8);
         object_manager_config.object_chunk_size =

src/ray/raylet/node_manager.cc

@@ -139,7 +139,7 @@ NodeManager::NodeManager(boost::asio::io_service &io_service,
       initial_config_(config),
       local_available_resources_(config.resource_config),
       worker_pool_(
-          io_service, config.num_initial_workers,
+          io_service, config.num_initial_workers, config.num_workers_soft_limit,
           config.num_initial_python_workers_for_first_job,
           config.maximum_startup_concurrency, config.min_worker_port,
           config.max_worker_port, gcs_client_, config.worker_commands,
@@ -1362,6 +1362,11 @@ void NodeManager::HandleWorkerAvailable(const std::shared_ptr<WorkerInterface> &
     // Call task dispatch to assign work to the new worker.
     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);
+  }
 }
 
 void NodeManager::ProcessDisconnectClientMessage(

src/ray/raylet/node_manager.h

@@ -67,6 +67,8 @@ struct NodeManagerConfig {
   int max_worker_port;
   /// The initial number of workers to create.
   int num_initial_workers;
+  /// The soft limit of the number of workers.
+  int num_workers_soft_limit;
   /// Number of initial Python workers for the first job.
   int num_initial_python_workers_for_first_job;
   /// The maximum number of workers that can be started concurrently by a

src/ray/raylet/worker_pool.cc

@@ -55,6 +55,7 @@ namespace ray {
 namespace raylet {
 
 WorkerPool::WorkerPool(boost::asio::io_service &io_service, int num_workers,
+                       int num_workers_soft_limit,
                        int num_initial_python_workers_for_first_job,
                        int maximum_startup_concurrency, int min_worker_port,
                        int max_worker_port, std::shared_ptr<gcs::GcsClient> gcs_client,
@@ -62,6 +63,7 @@ WorkerPool::WorkerPool(boost::asio::io_service &io_service, int num_workers,
                        const std::unordered_map<std::string, std::string> &raylet_config,
                        std::function<void()> starting_worker_timeout_callback)
     : io_service_(&io_service),
+      num_workers_soft_limit_(num_workers_soft_limit),
       maximum_startup_concurrency_(maximum_startup_concurrency),
       gcs_client_(std::move(gcs_client)),
       raylet_config_(raylet_config),
@@ -591,6 +593,64 @@ void WorkerPool::PushWorker(const std::shared_ptr<WorkerInterface> &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);
+      }
+    }
+  }
+
+  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);
+  }
+  worker->GetProcess().Kill();
+}
+
 std::shared_ptr<WorkerInterface> WorkerPool::PopWorker(
     const TaskSpecification &task_spec) {
   auto &state = GetStateForLanguage(task_spec.GetLanguage());
@@ -671,6 +731,7 @@ 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())},

src/ray/raylet/worker_pool.h

@@ -73,6 +73,7 @@ class WorkerPool : public WorkerPoolInterface {
   /// the pool.
   ///
   /// \param num_workers The number of workers to start, per language.
+  /// \param num_workers_soft_limit The soft limit of the number of workers.
   /// \param num_initial_python_workers_for_first_job The number of initial Python
   /// workers for the first job.
   /// \param maximum_startup_concurrency The maximum number of worker processes
@@ -88,7 +89,7 @@ class WorkerPool : public WorkerPoolInterface {
   /// \param starting_worker_timeout_callback The callback that will be triggered once
   /// it times out to start a worker.
   WorkerPool(boost::asio::io_service &io_service, int num_workers,
-             int num_initial_python_workers_for_first_job,
+             int num_workers_soft_limit, int num_initial_python_workers_for_first_job,
              int maximum_startup_concurrency, int min_worker_port, int max_worker_port,
              std::shared_ptr<gcs::GcsClient> gcs_client,
              const WorkerCommandMap &worker_commands,
@@ -180,6 +181,11 @@ 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);
+
   /// 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.
   ///
@@ -292,6 +298,10 @@ 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;
@@ -354,6 +364,8 @@ class WorkerPool : public WorkerPoolInterface {
 
   /// 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.
+  int num_workers_soft_limit_;
   /// The maximum number of worker processes that can be started concurrently.
   int maximum_startup_concurrency_;
   /// Keeps track of unused ports that newly-created workers can bind on.

src/ray/raylet/worker_pool_test.cc

@@ -34,7 +34,7 @@ class WorkerPoolMock : public WorkerPool {
  public:
   explicit WorkerPoolMock(boost::asio::io_service &io_service,
                           const WorkerCommandMap &worker_commands)
-      : WorkerPool(io_service, 0, 0, MAXIMUM_STARTUP_CONCURRENCY, 0, 0, nullptr,
+      : WorkerPool(io_service, 0, 0, 0, MAXIMUM_STARTUP_CONCURRENCY, 0, 0, nullptr,
                    worker_commands, {}, []() {}),
         last_worker_process_() {
     states_by_lang_[ray::Language::JAVA].num_workers_per_process =