[Java] Support multiple workers in Java worker process (#5505)

src/ray/common/constants.h

@@ -21,6 +21,9 @@ constexpr char kObjectTablePrefix[] = "ObjectTable";
 /// Prefix for the task table keys in redis.
 constexpr char kTaskTablePrefix[] = "TaskTable";
 
-constexpr char kWorkerDynamicOptionPlaceholderPrefix[] = "RAY_WORKER_OPTION_";
+constexpr char kWorkerDynamicOptionPlaceholderPrefix[] =
+    "RAY_WORKER_DYNAMIC_OPTION_PLACEHOLDER_";
+
+constexpr char kWorkerNumWorkersPlaceholder[] = "RAY_WORKER_NUM_WORKERS_PLACEHOLDER";
 
 #endif  // RAY_CONSTANTS_H_

src/ray/common/ray_config_def.h

@@ -137,8 +137,11 @@ RAY_CONFIG(int, object_manager_repeated_push_delay_ms, 60000)
 /// chunks exceeds the number of available sending threads.
 RAY_CONFIG(uint64_t, object_manager_default_chunk_size, 1000000)
 
-/// Number of workers per process
-RAY_CONFIG(int, num_workers_per_process, 1)
+/// Number of workers per Python worker process
+RAY_CONFIG(int, num_workers_per_process_python, 1)
+
+/// Number of workers per Java worker process
+RAY_CONFIG(int, num_workers_per_process_java, 10)
 
 /// Maximum timeout in milliseconds within which a task lease must be renewed.
 RAY_CONFIG(int64_t, max_task_lease_timeout_ms, 60000)

src/ray/core_worker/lib/java/org_ray_runtime_task_NativeTaskSubmitter.cc

@@ -84,10 +84,12 @@ inline ray::ActorCreationOptions ToActorCreationOptions(JNIEnv *env,
     jobject java_resources =
         env->GetObjectField(actorCreationOptions, java_base_task_options_resources);
     resources = ToResources(env, java_resources);
-    std::string jvm_options = JavaStringToNativeString(
-        env, (jstring)env->GetObjectField(actorCreationOptions,
-                                          java_actor_creation_options_jvm_options));
-    dynamic_worker_options.emplace_back(jvm_options);
+    jstring java_jvm_options = (jstring)env->GetObjectField(
+        actorCreationOptions, java_actor_creation_options_jvm_options);
+    if (java_jvm_options) {
+      std::string jvm_options = JavaStringToNativeString(env, java_jvm_options);
+      dynamic_worker_options.emplace_back(jvm_options);
+    }
   }
 
   ray::ActorCreationOptions action_creation_options{

src/ray/protobuf/common.proto

@@ -85,8 +85,8 @@ message ActorCreationTaskSpec {
   uint64 max_actor_reconstructions = 3;
   // The dynamic options used in the worker command when starting a worker process for
   // an actor creation task. If the list isn't empty, the options will be used to replace
-  // the placeholder strings (`RAY_WORKER_OPTION_0`, `RAY_WORKER_OPTION_1`, etc) in the
-  // worker command.
+  // the placeholder strings (`RAY_WORKER_DYNAMIC_OPTION_PLACEHOLDER_0`,
+  // `RAY_WORKER_DYNAMIC_OPTION_PLACEHOLDER_1`, etc) in the worker command.
   repeated string dynamic_worker_options = 4;
 }
 

src/ray/raylet/main.cc

@@ -104,8 +104,6 @@ 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_per_process =
-      RayConfig::instance().num_workers_per_process();
   node_manager_config.maximum_startup_concurrency = maximum_startup_concurrency;
 
   if (!python_worker_command.empty()) {

src/ray/raylet/node_manager.cc

@@ -82,9 +82,8 @@ NodeManager::NodeManager(boost::asio::io_service &io_service,
       object_manager_profile_timer_(io_service),
       initial_config_(config),
       local_available_resources_(config.resource_config),
-      worker_pool_(config.num_initial_workers, config.num_workers_per_process,
-                   config.maximum_startup_concurrency, gcs_client_,
-                   config.worker_commands),
+      worker_pool_(config.num_initial_workers, config.maximum_startup_concurrency,
+                   gcs_client_, config.worker_commands),
       scheduling_policy_(local_queues_),
       reconstruction_policy_(
           io_service_,

src/ray/raylet/node_manager.h

@@ -43,8 +43,6 @@ struct NodeManagerConfig {
   int node_manager_port;
   /// The initial number of workers to create.
   int num_initial_workers;
-  /// The number of workers per process.
-  int num_workers_per_process;
   /// The maximum number of workers that can be started concurrently by a
   /// worker pool.
   int maximum_startup_concurrency;

src/ray/raylet/object_manager_integration_test.cc

@@ -39,7 +39,6 @@ class TestObjectManagerBase : public ::testing::Test {
     node_manager_config.resource_config =
         ray::raylet::ResourceSet(std::move(static_resource_conf));
     node_manager_config.num_initial_workers = 0;
-    node_manager_config.num_workers_per_process = 1;
     // Use a default worker that can execute empty tasks with dependencies.
     std::vector<std::string> py_worker_command;
     py_worker_command.push_back("python");

src/ray/raylet/worker_pool.cc

@@ -40,17 +40,13 @@ namespace ray {
 namespace raylet {
 
 /// A constructor that initializes a worker pool with
-/// (num_worker_processes * num_workers_per_process) workers for each language.
-WorkerPool::WorkerPool(int num_worker_processes, int num_workers_per_process,
-                       int maximum_startup_concurrency,
+/// (num_worker_processes * states_by_lang_[language].num_workers_per_process) workers for
+/// each language.
+WorkerPool::WorkerPool(int num_worker_processes, int maximum_startup_concurrency,
                        std::shared_ptr<gcs::RedisGcsClient> gcs_client,
                        const WorkerCommandMap &worker_commands)
-    : num_workers_per_process_(num_workers_per_process),
-      multiple_for_warning_(std::max(num_worker_processes, maximum_startup_concurrency)),
-      maximum_startup_concurrency_(maximum_startup_concurrency),
-      last_warning_multiple_(0),
+    : maximum_startup_concurrency_(maximum_startup_concurrency),
       gcs_client_(std::move(gcs_client)) {
-  RAY_CHECK(num_workers_per_process > 0) << "num_workers_per_process must be positive.";
   RAY_CHECK(maximum_startup_concurrency > 0);
   // Ignore SIGCHLD signals. If we don't do this, then worker processes will
   // become zombies instead of dying gracefully.
@@ -58,6 +54,25 @@ WorkerPool::WorkerPool(int num_worker_processes, int num_workers_per_process,
   for (const auto &entry : worker_commands) {
     // Initialize the pool state for this language.
     auto &state = states_by_lang_[entry.first];
+    switch (entry.first) {
+    case Language::PYTHON:
+      state.num_workers_per_process =
+          RayConfig::instance().num_workers_per_process_python();
+      break;
+    case Language::JAVA:
+      state.num_workers_per_process =
+          RayConfig::instance().num_workers_per_process_java();
+      break;
+    default:
+      RAY_LOG(FATAL) << "The number of workers per process for "
+                     << Language_Name(entry.first) << " worker is not set.";
+    }
+    RAY_CHECK(state.num_workers_per_process > 0)
+        << "Number of workers per process of language " << Language_Name(entry.first)
+        << " must be positive.";
+    state.multiple_for_warning =
+        std::max(num_worker_processes, maximum_startup_concurrency) *
+        state.num_workers_per_process;
     // Set worker command for this language.
     state.worker_command = entry.second;
     RAY_CHECK(!state.worker_command.empty()) << "Worker command must not be empty.";
@@ -119,9 +134,17 @@ int WorkerPool::StartWorkerProcess(const Language &language,
                  << state.idle_actor.size() << " actor workers, and " << state.idle.size()
                  << " non-actor workers";
 
+  int workers_to_start;
+  if (dynamic_options.empty()) {
+    workers_to_start = state.num_workers_per_process;
+  } else {
+    workers_to_start = 1;
+  }
+
   // Extract pointers from the worker command to pass into execvp.
   std::vector<std::string> worker_command_args;
   size_t dynamic_option_index = 0;
+  bool num_workers_arg_replaced = false;
   for (auto const &token : state.worker_command) {
     const auto option_placeholder =
         kWorkerDynamicOptionPlaceholderPrefix + std::to_string(dynamic_option_index);
@@ -135,9 +158,22 @@ int WorkerPool::StartWorkerProcess(const Language &language,
         ++dynamic_option_index;
       }
     } else {
-      worker_command_args.push_back(token);
+      size_t num_workers_index = token.find(kWorkerNumWorkersPlaceholder);
+      if (num_workers_index != std::string::npos) {
+        std::string arg = token;
+        worker_command_args.push_back(arg.replace(num_workers_index,
+                                                  strlen(kWorkerNumWorkersPlaceholder),
+                                                  std::to_string(workers_to_start)));
+        num_workers_arg_replaced = true;
+      } else {
+        worker_command_args.push_back(token);
+      }
     }
   }
+  RAY_CHECK(num_workers_arg_replaced || state.num_workers_per_process == 1)
+      << "Expect to start " << state.num_workers_per_process << " workers per "
+      << Language_Name(language) << " worker process. But the "
+      << kWorkerNumWorkersPlaceholder << "placeholder is not found in worker command.";
 
   pid_t pid = StartProcess(worker_command_args);
   if (pid < 0) {
@@ -145,9 +181,9 @@ int WorkerPool::StartWorkerProcess(const Language &language,
     RAY_LOG(FATAL) << "Failed to fork worker process: " << strerror(errno);
   } else if (pid > 0) {
     // Parent process case.
-    RAY_LOG(DEBUG) << "Started worker process with pid " << pid;
-    state.starting_worker_processes.emplace(
-        std::make_pair(pid, num_workers_per_process_));
+    RAY_LOG(DEBUG) << "Started worker process of " << workers_to_start
+                   << " worker(s) with pid " << pid;
+    state.starting_worker_processes.emplace(pid, workers_to_start);
     return pid;
   }
   return -1;
@@ -355,27 +391,30 @@ std::vector<std::shared_ptr<Worker>> WorkerPool::GetWorkersRunningTasksForJob(
 }
 
 void WorkerPool::WarnAboutSize() {
-  int64_t num_workers_started_or_registered = 0;
   for (const auto &entry : states_by_lang_) {
+    auto state = entry.second;
+    int64_t num_workers_started_or_registered = 0;
     num_workers_started_or_registered +=
-        static_cast<int64_t>(entry.second.registered_workers.size());
-    num_workers_started_or_registered +=
-        static_cast<int64_t>(entry.second.starting_worker_processes.size());
-  }
-  int64_t multiple = num_workers_started_or_registered / multiple_for_warning_;
-  std::stringstream warning_message;
-  if (multiple >= 3 && multiple > last_warning_multiple_) {
-    // Push an error message to the user if the worker pool tells us that it is
-    // getting too big.
-    last_warning_multiple_ = multiple;
-    warning_message << "WARNING: " << num_workers_started_or_registered
-                    << " workers have been started. This could be a result of using "
-                    << "a large number of actors, or it could be a consequence of "
-                    << "using nested tasks "
-                    << "(see https://github.com/ray-project/ray/issues/3644) for "
-                    << "some a discussion of workarounds.";
-    RAY_CHECK_OK(gcs_client_->error_table().PushErrorToDriver(
-        JobID::Nil(), "worker_pool_large", warning_message.str(), current_time_ms()));
+        static_cast<int64_t>(state.registered_workers.size());
+    for (const auto &starting_process : state.starting_worker_processes) {
+      num_workers_started_or_registered += starting_process.second;
+    }
+    int64_t multiple = num_workers_started_or_registered / state.multiple_for_warning;
+    std::stringstream warning_message;
+    if (multiple >= 3 && multiple > state.last_warning_multiple) {
+      // Push an error message to the user if the worker pool tells us that it is
+      // getting too big.
+      state.last_warning_multiple = multiple;
+      warning_message << "WARNING: " << num_workers_started_or_registered << " "
+                      << Language_Name(entry.first)
+                      << " workers have been started. This could be a result of using "
+                      << "a large number of actors, or it could be a consequence of "
+                      << "using nested tasks "
+                      << "(see https://github.com/ray-project/ray/issues/3644) for "
+                      << "some a discussion of workarounds.";
+      RAY_CHECK_OK(gcs_client_->error_table().PushErrorToDriver(
+          JobID::Nil(), "worker_pool_large", warning_message.str(), current_time_ms()));
+    }
   }
 }
 
@@ -390,8 +429,10 @@ std::string WorkerPool::DebugString() const {
   std::stringstream result;
   result << "WorkerPool:";
   for (const auto &entry : states_by_lang_) {
-    result << "\n- num workers: " << entry.second.registered_workers.size();
-    result << "\n- num drivers: " << entry.second.registered_drivers.size();
+    result << "\n- num " << Language_Name(entry.first)
+           << " workers: " << entry.second.registered_workers.size();
+    result << "\n- num " << Language_Name(entry.first)
+           << " drivers: " << entry.second.registered_drivers.size();
   }
   return result.str();
 }

src/ray/raylet/worker_pool.h

@@ -33,14 +33,12 @@ class WorkerPool {
   /// and add them to the pool.
   ///
   /// \param num_worker_processes The number of worker processes to start, per language.
-  /// \param num_workers_per_process The number of workers per process.
   /// \param maximum_startup_concurrency The maximum number of worker processes
   /// that can be started in parallel (typically this should be set to the number of CPU
   /// resources on the machine).
   /// \param worker_commands The commands used to start the worker process, grouped by
   /// language.
-  WorkerPool(int num_worker_processes, int num_workers_per_process,
-             int maximum_startup_concurrency,
+  WorkerPool(int num_worker_processes, int maximum_startup_concurrency,
              std::shared_ptr<gcs::RedisGcsClient> gcs_client,
              const WorkerCommandMap &worker_commands);
 
@@ -134,7 +132,7 @@ class WorkerPool {
  protected:
   /// Asynchronously start a new worker process. Once the worker process has
   /// registered with an external server, the process should create and
-  /// register num_workers_per_process_ workers, then add them to the pool.
+  /// register num_workers_per_process workers, then add them to the pool.
   /// Failure to start the worker process is a fatal error. If too many workers
   /// are already being started, then this function will return without starting
   /// any workers.
@@ -159,6 +157,8 @@ class WorkerPool {
   struct State {
     /// The commands and arguments used to start the worker process
     std::vector<std::string> worker_command;
+    /// The number of workers per process.
+    int num_workers_per_process;
     /// The pool of dedicated workers for actor creation tasks
     /// with prefix or suffix worker command.
     std::unordered_map<TaskID, std::shared_ptr<Worker>> idle_dedicated_workers;
@@ -179,10 +179,14 @@ class WorkerPool {
     std::unordered_map<pid_t, TaskID> dedicated_workers_to_tasks;
     /// A map for speeding up looking up the pending worker for the given task.
     std::unordered_map<TaskID, pid_t> tasks_to_dedicated_workers;
+    /// We'll push a warning to the user every time a multiple of this many
+    /// worker processes has been started.
+    int multiple_for_warning;
+    /// The last size at which a warning about the number of registered workers
+    /// was generated.
+    int64_t last_warning_multiple;
   };
 
-  /// The number of workers per process.
-  int num_workers_per_process_;
   /// Pool states per language.
   std::unordered_map<Language, State, std::hash<int>> states_by_lang_;
 
@@ -191,14 +195,8 @@ class WorkerPool {
   /// for a given language.
   State &GetStateForLanguage(const Language &language);
 
-  /// We'll push a warning to the user every time a multiple of this many
-  /// workers has been started.
-  int multiple_for_warning_;
-  /// The maximum number of workers that can be started concurrently.
+  /// The maximum number of worker processes that can be started concurrently.
   int maximum_startup_concurrency_;
-  /// The last size at which a warning about the number of registered workers
-  /// was generated.
-  int64_t last_warning_multiple_;
   /// A client connection to the GCS.
   std::shared_ptr<gcs::RedisGcsClient> gcs_client_;
 };

src/ray/raylet/worker_pool_test.cc

@@ -15,13 +15,20 @@ int MAXIMUM_STARTUP_CONCURRENCY = 5;
 class WorkerPoolMock : public WorkerPool {
  public:
   WorkerPoolMock()
-      : WorkerPoolMock({{Language::PYTHON, {"dummy_py_worker_command"}},
-                        {Language::JAVA, {"dummy_java_worker_command"}}}) {}
+      : WorkerPoolMock(
+            {{Language::PYTHON,
+              {"dummy_py_worker_command", "--foo=RAY_WORKER_NUM_WORKERS_PLACEHOLDER"}},
+             {Language::JAVA,
+              {"dummy_java_worker_command",
+               "--foo=RAY_WORKER_NUM_WORKERS_PLACEHOLDER"}}}) {}
 
   explicit WorkerPoolMock(const WorkerCommandMap &worker_commands)
-      : WorkerPool(0, NUM_WORKERS_PER_PROCESS, MAXIMUM_STARTUP_CONCURRENCY, nullptr,
-                   worker_commands),
-        last_worker_pid_(0) {}
+      : WorkerPool(0, MAXIMUM_STARTUP_CONCURRENCY, nullptr, worker_commands),
+        last_worker_pid_(0) {
+    for (auto &entry : states_by_lang_) {
+      entry.second.num_workers_per_process = NUM_WORKERS_PER_PROCESS;
+    }
+  }
 
   ~WorkerPoolMock() {
     // Avoid killing real processes
@@ -150,16 +157,38 @@ TEST_F(WorkerPoolTest, HandleWorkerRegistration) {
 }
 
 TEST_F(WorkerPoolTest, StartupWorkerCount) {
-  int desired_initial_worker_count_per_language = 20;
-  for (int i = 0; i < desired_initial_worker_count_per_language; i++) {
-    worker_pool_.StartWorkerProcess(Language::PYTHON);
-    worker_pool_.StartWorkerProcess(Language::JAVA);
+  std::string num_workers_arg =
+      std::string("--foo=") + std::to_string(NUM_WORKERS_PER_PROCESS);
+  std::vector<Language> languages = {Language::PYTHON, Language::JAVA};
+  std::vector<std::vector<std::string>> worker_commands = {
+      {{"dummy_py_worker_command", num_workers_arg},
+       {"dummy_java_worker_command", num_workers_arg}}};
+  int desired_initial_worker_process_count_per_language = MAXIMUM_STARTUP_CONCURRENCY + 1;
+  int expected_worker_process_count = MAXIMUM_STARTUP_CONCURRENCY * languages.size();
+  pid_t last_started_worker_process = 0;
+  for (int i = 0; i < desired_initial_worker_process_count_per_language; i++) {
+    for (size_t j = 0; j < languages.size(); j++) {
+      worker_pool_.StartWorkerProcess(languages[j]);
+      ASSERT_TRUE(worker_pool_.NumWorkerProcessesStarting() <=
+                  expected_worker_process_count);
+      if (last_started_worker_process != worker_pool_.LastStartedWorkerProcess()) {
+        last_started_worker_process = worker_pool_.LastStartedWorkerProcess();
+        const auto &real_command =
+            worker_pool_.GetWorkerCommand(worker_pool_.LastStartedWorkerProcess());
+        ASSERT_EQ(real_command, worker_commands[j]);
+      } else {
+        ASSERT_TRUE(worker_pool_.NumWorkerProcessesStarting() ==
+                    expected_worker_process_count);
+        ASSERT_TRUE(static_cast<int>(i * languages.size() + j) >=
+                    expected_worker_process_count);
+      }
+    }
   }
-  // Check that number of starting worker processes equals to
-  // maximum_startup_concurrency_ * 2. (because we started both python and java workers)
-  ASSERT_EQ(
-      worker_pool_.NumWorkerProcessesStarting(),
-      /* Provided in constructor of WorkerPoolMock */ MAXIMUM_STARTUP_CONCURRENCY * 2);
+  // Check number of starting worker processes
+  ASSERT_EQ(worker_pool_.NumWorkerProcessesStarting(), expected_worker_process_count);
+  ASSERT_TRUE(worker_pool_.NumWorkerProcessesStarting() <
+              static_cast<int>(desired_initial_worker_process_count_per_language *
+                               languages.size()));
 }
 
 TEST_F(WorkerPoolTest, HandleWorkerPushPop) {
@@ -232,7 +261,9 @@ TEST_F(WorkerPoolTest, PopWorkersOfMultipleLanguages) {
 
 TEST_F(WorkerPoolTest, StartWorkerWithDynamicOptionsCommand) {
   const std::vector<std::string> java_worker_command = {
-      "RAY_WORKER_OPTION_0", "dummy_java_worker_command", "RAY_WORKER_OPTION_1"};
+      "RAY_WORKER_DYNAMIC_OPTION_PLACEHOLDER_0", "dummy_java_worker_command",
+      "--foo=RAY_WORKER_NUM_WORKERS_PLACEHOLDER",
+      "RAY_WORKER_DYNAMIC_OPTION_PLACEHOLDER_1"};
   SetWorkerCommands({{Language::PYTHON, {"dummy_py_worker_command"}},
                      {Language::JAVA, java_worker_command}});
 
@@ -243,8 +274,9 @@ TEST_F(WorkerPoolTest, StartWorkerWithDynamicOptionsCommand) {
   worker_pool_.StartWorkerProcess(Language::JAVA, task_spec.DynamicWorkerOptions());
   const auto real_command =
       worker_pool_.GetWorkerCommand(worker_pool_.LastStartedWorkerProcess());
-  ASSERT_EQ(real_command, std::vector<std::string>(
-                              {"test_op_0", "dummy_java_worker_command", "test_op_1"}));
+  ASSERT_EQ(real_command,
+            std::vector<std::string>(
+                {"test_op_0", "dummy_java_worker_command", "--foo=1", "test_op_1"}));
 }
 
 }  // namespace raylet