[xray] Start actor workers in parallel (#2168)

src/ray/raylet/node_manager.cc

@@ -70,7 +70,9 @@ NodeManager::NodeManager(boost::asio::io_service &io_service,
       heartbeat_timer_(io_service),
       heartbeat_period_ms_(config.heartbeat_period_ms),
       local_resources_(config.resource_config),
-      worker_pool_(config.num_initial_workers, config.worker_command),
+      worker_pool_(config.num_initial_workers,
+                   static_cast<int>(config.resource_config.GetNumCpus()),
+                   config.worker_command),
       local_queues_(SchedulingQueue()),
       scheduling_policy_(local_queues_),
       reconstruction_policy_([this](const TaskID &task_id) { ResubmitTask(task_id); }),
@@ -400,8 +402,7 @@ void NodeManager::ProcessClientMessage(
       const auto &task = tasks.front();
       // Get the CPU resources required by the running task.
       const auto required_resources = task.GetTaskSpecification().GetRequiredResources();
-      double required_cpus = 0;
-      RAY_CHECK(required_resources.GetResource(kCPU_ResourceLabel, &required_cpus));
+      double required_cpus = required_resources.GetNumCpus();
       const std::unordered_map<std::string, double> cpu_resources = {
           {kCPU_ResourceLabel, required_cpus}};
       // Release the CPU resources.
@@ -429,8 +430,7 @@ void NodeManager::ProcessClientMessage(
       const auto &task = tasks.front();
       // Get the CPU resources required by the running task.
       const auto required_resources = task.GetTaskSpecification().GetRequiredResources();
-      double required_cpus = 0;
-      RAY_CHECK(required_resources.GetResource(kCPU_ResourceLabel, &required_cpus));
+      double required_cpus = required_resources.GetNumCpus();
       const std::unordered_map<std::string, double> cpu_resources = {
           {kCPU_ResourceLabel, required_cpus}};
       // Acquire the CPU resources.

src/ray/raylet/scheduling_resources.cc

@@ -116,6 +116,12 @@ bool ResourceSet::GetResource(const std::string &resource_name, double *value) c
   return true;
 }
 
+double ResourceSet::GetNumCpus() const {
+  double num_cpus;
+  RAY_CHECK(GetResource(kCPU_ResourceLabel, &num_cpus));
+  return num_cpus;
+}
+
 const std::string ResourceSet::ToString() const {
   std::string return_string = "";
   for (const auto &resource_pair : this->resource_capacity_) {

src/ray/raylet/scheduling_resources.h

@@ -99,6 +99,11 @@ class ResourceSet {
   ///         False otherwise.
   bool GetResource(const std::string &resource_name, double *value) const;
 
+  /// Return the number of CPUs.
+  ///
+  /// \return Number of CPUs.
+  double GetNumCpus() const;
+
   /// Return true if the resource set is empty. False otherwise.
   ///
   /// \return True if the resource capacity is zero. False otherwise.

src/ray/raylet/worker_pool.cc

@@ -10,8 +10,9 @@ namespace ray {
 namespace raylet {
 
 /// A constructor that initializes a worker pool with num_workers workers.
-WorkerPool::WorkerPool(int num_workers, const std::vector<std::string> &worker_command)
-    : worker_command_(worker_command) {
+WorkerPool::WorkerPool(int num_workers, int num_cpus,
+                       const std::vector<std::string> &worker_command)
+    : num_cpus_(num_cpus), worker_command_(worker_command) {
   // Ignore SIGCHLD signals. If we don't do this, then worker processes will
   // become zombies instead of dying gracefully.
   signal(SIGCHLD, SIG_IGN);
@@ -52,7 +53,9 @@ uint32_t WorkerPool::Size() const {
 
 void WorkerPool::StartWorker(bool force_start) {
   RAY_CHECK(!worker_command_.empty()) << "No worker command provided";
-  if (!started_worker_pids_.empty() && !force_start) {
+  // The first condition makes sure that we are always starting up to
+  // num_cpus_ number of processes in parallel.
+  if (NumWorkersStarting() > num_cpus_ && !force_start) {
     // Workers have been started, but not registered. Force start disabled -- returning.
     RAY_LOG(DEBUG) << started_worker_pids_.size() << " workers pending registration";
     return;
@@ -155,7 +158,7 @@ bool WorkerPool::DisconnectWorker(std::shared_ptr<Worker> worker) {
 // Protected WorkerPool methods.
 void WorkerPool::AddStartedWorker(pid_t pid) { started_worker_pids_.insert(pid); }
 
-uint32_t WorkerPool::NumStartedWorkers() const { return started_worker_pids_.size(); }
+int WorkerPool::NumWorkersStarting() const { return started_worker_pids_.size(); }
 
 }  // namespace raylet
 

src/ray/raylet/worker_pool.h

@@ -28,7 +28,8 @@ class WorkerPool {
   ///
   /// \param num_workers The number of workers to start.
   /// \param worker_command The command used to start the worker process.
-  WorkerPool(int num_workers, const std::vector<std::string> &worker_command);
+  WorkerPool(int num_workers, int num_cpus,
+             const std::vector<std::string> &worker_command);
 
   /// Create a pool with zero workers.
   ///
@@ -42,7 +43,8 @@ class WorkerPool {
   /// Asynchronously start a new worker process. Once the worker process has
   /// registered with an external server, the process should create and
   /// register a new Worker, then add itself to the pool. Failure to start
-  /// the worker process is a fatal error.
+  /// the worker process is a fatal error. This function will start up to
+  /// num_cpus many workers in parallel if it is called multiple times.
   ///
   /// \param force_start Controls whether to force starting a worker regardless of any
   /// workers that have already been started but not yet registered.
@@ -93,12 +95,15 @@ class WorkerPool {
   /// \param pid A process identifier for the worker being started.
   void AddStartedWorker(pid_t pid);
 
-  /// Return a number of workers currently started but not registered.
+  /// Return a number of workers currently starting but not registered.
   ///
   /// \return The number of worker PIDs stored for started workers.
-  uint32_t NumStartedWorkers() const;
+  int NumWorkersStarting() const;
 
  private:
+  /// The number of CPUs this Raylet has available.
+  int num_cpus_;
+  /// The command and arguments used to start the worker.
   std::vector<std::string> worker_command_;
   /// The pool of idle workers.
   std::list<std::shared_ptr<Worker>> pool_;

src/ray/raylet/worker_pool_test.cc

@@ -14,9 +14,9 @@ class WorkerPoolMock : public WorkerPool {
       : WorkerPool(worker_command) {}
 
   void StartWorker(pid_t pid, bool force_start = false) {
-    if (NumStartedWorkers() > 0 && !force_start) {
+    if (NumWorkersStarting() > 0 && !force_start) {
       // Workers have been started, but not registered. Force start disabled -- returning.
-      RAY_LOG(DEBUG) << NumStartedWorkers() << " workers pending registration";
+      RAY_LOG(DEBUG) << NumWorkersStarting() << " workers pending registration";
       return;
     }
     // Either no workers are pending registration or the worker start is being forced.