diff --git a/src/ray/raylet/worker_pool.cc b/src/ray/raylet/worker_pool.cc index 606b46edb..878259ffe 100644 --- a/src/ray/raylet/worker_pool.cc +++ b/src/ray/raylet/worker_pool.cc @@ -38,10 +38,9 @@ namespace ray { namespace raylet { -/// A constructor that initializes a worker pool with -/// (num_worker_processes * states_by_lang_[language].num_workers_per_process) workers for +/// A constructor that initializes a worker pool with num_workers workers for /// each language. -WorkerPool::WorkerPool(int num_worker_processes, int maximum_startup_concurrency, +WorkerPool::WorkerPool(int num_workers, int maximum_startup_concurrency, std::shared_ptr gcs_client, const WorkerCommandMap &worker_commands) : maximum_startup_concurrency_(maximum_startup_concurrency), @@ -70,12 +69,20 @@ WorkerPool::WorkerPool(int num_worker_processes, int maximum_startup_concurrency << "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; + std::max(state.num_workers_per_process, + std::max(num_workers, maximum_startup_concurrency)); // Set worker command for this language. state.worker_command = entry.second; RAY_CHECK(!state.worker_command.empty()) << "Worker command must not be empty."; - // Force-start num_workers worker processes for this language. + } + Start(num_workers); +} + +void WorkerPool::Start(int num_workers) { + for (auto &entry : states_by_lang_) { + auto &state = entry.second; + int num_worker_processes = static_cast( + std::ceil(static_cast(num_workers) / state.num_workers_per_process)); for (int i = 0; i < num_worker_processes; i++) { StartWorkerProcess(entry.first); } @@ -120,11 +127,14 @@ int WorkerPool::StartWorkerProcess(const Language &language, auto &state = GetStateForLanguage(language); // If we are already starting up too many workers, then return without starting // more. - if (static_cast(state.starting_worker_processes.size()) >= - maximum_startup_concurrency_) { + int starting_workers = 0; + for (auto &entry : state.starting_worker_processes) { + starting_workers += entry.second; + } + if (starting_workers >= maximum_startup_concurrency_) { // Workers have been started, but not registered. Force start disabled -- returning. - RAY_LOG(DEBUG) << "Worker not started, " << state.starting_worker_processes.size() - << " worker processes of language type " << static_cast(language) + RAY_LOG(DEBUG) << "Worker not started, " << starting_workers + << " workers of language type " << static_cast(language) << " pending registration"; return -1; } diff --git a/src/ray/raylet/worker_pool.h b/src/ray/raylet/worker_pool.h index ec3cc80ac..8af5b4f15 100644 --- a/src/ray/raylet/worker_pool.h +++ b/src/ray/raylet/worker_pool.h @@ -5,6 +5,7 @@ #include #include #include +#include "gtest/gtest.h" #include "ray/common/client_connection.h" #include "ray/common/task/task.h" @@ -27,18 +28,19 @@ class Worker; /// is a container for a unit of work. class WorkerPool { public: - /// Create a pool and asynchronously start the specified number of worker processes. - /// Once each worker process has registered with an external server, - /// the process should create and register the specified number of workers, - /// and add them to the pool. + /// Create a pool and asynchronously start at least the specified number of workers per + /// language. + /// Once each worker process has registered with an external server, the + /// process should create and register the specified number of workers, and add them to + /// the pool. /// - /// \param num_worker_processes The number of worker processes to start, per language. + /// \param num_workers The number of workers to start, per language. /// \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 maximum_startup_concurrency, + WorkerPool(int num_workers, int maximum_startup_concurrency, std::shared_ptr gcs_client, const WorkerCommandMap &worker_commands); @@ -205,6 +207,12 @@ class WorkerPool { std::unordered_map> states_by_lang_; private: + /// Force-start at least num_workers workers for this language. Used for internal and + /// test purpose only. + /// + /// \param num_workers The number of workers to start, per language. + void Start(int num_workers); + /// A helper function that returns the reference of the pool state /// for a given language. State &GetStateForLanguage(const Language &language); @@ -213,6 +221,8 @@ class WorkerPool { int maximum_startup_concurrency_; /// A client connection to the GCS. std::shared_ptr gcs_client_; + + FRIEND_TEST(WorkerPoolTest, InitialWorkerProcessCount); }; } // namespace raylet diff --git a/src/ray/raylet/worker_pool_test.cc b/src/ray/raylet/worker_pool_test.cc index aae9714ea..e5c2b7346 100644 --- a/src/ray/raylet/worker_pool_test.cc +++ b/src/ray/raylet/worker_pool_test.cc @@ -12,6 +12,8 @@ namespace raylet { int NUM_WORKERS_PER_PROCESS = 3; int MAXIMUM_STARTUP_CONCURRENCY = 5; +std::vector LANGUAGES = {Language::PYTHON, Language::JAVA}; + class WorkerPoolMock : public WorkerPool { public: WorkerPoolMock() @@ -54,6 +56,16 @@ class WorkerPoolMock : public WorkerPool { return worker_commands_by_pid[pid]; } + int NumWorkersStarting() const { + int total = 0; + for (auto &state_entry : states_by_lang_) { + for (auto &process_entry : state_entry.second.starting_worker_processes) { + total += process_entry.second; + } + } + return total; + } + int NumWorkerProcessesStarting() const { int total = 0; for (auto &entry : states_by_lang_) { @@ -156,19 +168,23 @@ TEST_F(WorkerPoolTest, HandleWorkerRegistration) { } } -TEST_F(WorkerPoolTest, StartupWorkerCount) { +TEST_F(WorkerPoolTest, StartupWorkerProcessCount) { std::string num_workers_arg = std::string("--foo=") + std::to_string(NUM_WORKERS_PER_PROCESS); - std::vector languages = {Language::PYTHON, Language::JAVA}; std::vector> 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(); + int desired_initial_worker_process_count_per_language = 100; + int expected_worker_process_count = + static_cast(std::ceil(static_cast(MAXIMUM_STARTUP_CONCURRENCY) / + NUM_WORKERS_PER_PROCESS * LANGUAGES.size())); + ASSERT_TRUE(expected_worker_process_count < + static_cast(desired_initial_worker_process_count_per_language * + 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]); + 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()) { @@ -179,16 +195,24 @@ TEST_F(WorkerPoolTest, StartupWorkerCount) { } else { ASSERT_TRUE(worker_pool_.NumWorkerProcessesStarting() == expected_worker_process_count); - ASSERT_TRUE(static_cast(i * languages.size() + j) >= + ASSERT_TRUE(static_cast(i * LANGUAGES.size() + j) >= expected_worker_process_count); } } } - // Check number of starting worker processes + // Check number of starting workers ASSERT_EQ(worker_pool_.NumWorkerProcessesStarting(), expected_worker_process_count); - ASSERT_TRUE(worker_pool_.NumWorkerProcessesStarting() < - static_cast(desired_initial_worker_process_count_per_language * - languages.size())); +} + +TEST_F(WorkerPoolTest, InitialWorkerProcessCount) { + worker_pool_.Start(1); + // Here we try to start only 1 worker for each worker language. But since each worker + // process contains exactly NUM_WORKERS_PER_PROCESS (3) workers here, it's expected to + // see 3 workers for each worker language, instead of 1. + ASSERT_NE(worker_pool_.NumWorkersStarting(), 1 * LANGUAGES.size()); + ASSERT_EQ(worker_pool_.NumWorkersStarting(), + NUM_WORKERS_PER_PROCESS * LANGUAGES.size()); + ASSERT_EQ(worker_pool_.NumWorkerProcessesStarting(), LANGUAGES.size()); } TEST_F(WorkerPoolTest, HandleWorkerPushPop) {