From 02766adecad20891150e7810f3381055d320d2be Mon Sep 17 00:00:00 2001 From: Kai Yang Date: Wed, 30 Jan 2019 13:43:12 +0800 Subject: [PATCH] Limit maximum starting workers per language (#3852) --- src/ray/raylet/worker_pool.cc | 71 ++++++++++++++++++------------ src/ray/raylet/worker_pool.h | 22 +++++---- src/ray/raylet/worker_pool_test.cc | 54 ++++++++++++++++------- 3 files changed, 94 insertions(+), 53 deletions(-) diff --git a/src/ray/raylet/worker_pool.cc b/src/ray/raylet/worker_pool.cc index 8387656e0..12c753bdc 100644 --- a/src/ray/raylet/worker_pool.cc +++ b/src/ray/raylet/worker_pool.cc @@ -71,10 +71,10 @@ WorkerPool::~WorkerPool() { for (const auto &worker : entry.second.registered_workers) { pids_to_kill.insert(worker->Pid()); } - } - // Kill all the workers that have been started but not registered. - for (const auto &entry : starting_worker_processes_) { - pids_to_kill.insert(entry.first); + // Kill all the workers that have been started but not registered. + for (const auto &starting_worker : entry.second.starting_worker_processes) { + pids_to_kill.insert(starting_worker.first); + } } for (const auto &pid : pids_to_kill) { RAY_CHECK(pid > 0); @@ -97,38 +97,22 @@ uint32_t WorkerPool::Size(const Language &language) const { } void 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(starting_worker_processes_.size()) >= + if (static_cast(state.starting_worker_processes.size()) >= maximum_startup_concurrency_) { // Workers have been started, but not registered. Force start disabled -- returning. - RAY_LOG(DEBUG) << "Worker not started, " << starting_worker_processes_.size() - << " worker processes pending registration"; + RAY_LOG(DEBUG) << "Worker not started, " << state.starting_worker_processes.size() + << " worker processes of language type " << static_cast(language) + << " pending registration"; return; } - auto &state = GetStateForLanguage(language); // Either there are no workers pending registration or the worker start is being forced. RAY_LOG(DEBUG) << "Starting new worker process, current pool has " << state.idle_actor.size() << " actor workers, and " << state.idle.size() << " non-actor workers"; - // Launch the process to create the worker. - pid_t pid = fork(); - if (pid < 0) { - // Failure case. - RAY_LOG(FATAL) << "Failed to fork worker process: " << strerror(errno); - return; - } else if (pid > 0) { - // Parent process case. - RAY_LOG(DEBUG) << "Started worker process with pid " << pid; - starting_worker_processes_.emplace(std::make_pair(pid, num_workers_per_process_)); - return; - } - - // Child process case. - // Reset the SIGCHLD handler for the worker. - signal(SIGCHLD, SIG_DFL); - // Extract pointers from the worker command to pass into execvp. std::vector worker_command_args; for (auto const &token : state.worker_command) { @@ -136,12 +120,39 @@ void WorkerPool::StartWorkerProcess(const Language &language) { } worker_command_args.push_back(nullptr); + pid_t pid = StartProcess(worker_command_args); + if (pid < 0) { + // Failure case. + RAY_LOG(FATAL) << "Failed to fork worker process: " << strerror(errno); + return; + } 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_)); + return; + } +} + +pid_t WorkerPool::StartProcess(const std::vector &worker_command_args) { + // Launch the process to create the worker. + pid_t pid = fork(); + + if (pid != 0) { + return pid; + } + + // Child process case. + // Reset the SIGCHLD handler for the worker. + signal(SIGCHLD, SIG_DFL); + // Try to execute the worker command. int rv = execvp(worker_command_args[0], const_cast(worker_command_args.data())); // The worker failed to start. This is a fatal error. RAY_LOG(FATAL) << "Failed to start worker with return value " << rv << ": " << strerror(errno); + return 0; } void WorkerPool::RegisterWorker(const std::shared_ptr &worker) { @@ -150,11 +161,11 @@ void WorkerPool::RegisterWorker(const std::shared_ptr &worker) { auto &state = GetStateForLanguage(worker->GetLanguage()); state.registered_workers.insert(std::move(worker)); - auto it = starting_worker_processes_.find(pid); - RAY_CHECK(it != starting_worker_processes_.end()); + auto it = state.starting_worker_processes.find(pid); + RAY_CHECK(it != state.starting_worker_processes.end()); it->second--; if (it->second == 0) { - starting_worker_processes_.erase(it); + state.starting_worker_processes.erase(it); } } @@ -251,10 +262,12 @@ std::vector> WorkerPool::GetWorkersRunningTasksForDriver } std::string WorkerPool::WarningAboutSize() { - int64_t num_workers_started_or_registered = starting_worker_processes_.size(); + int64_t num_workers_started_or_registered = 0; for (const auto &entry : states_by_lang_) { num_workers_started_or_registered += static_cast(entry.second.registered_workers.size()); + num_workers_started_or_registered += + static_cast(entry.second.starting_worker_processes.size()); } int64_t multiple = num_workers_started_or_registered / multiple_for_warning_; std::stringstream warning_message; diff --git a/src/ray/raylet/worker_pool.h b/src/ray/raylet/worker_pool.h index ab9ca04ef..ccad659f5 100644 --- a/src/ray/raylet/worker_pool.h +++ b/src/ray/raylet/worker_pool.h @@ -131,13 +131,12 @@ class WorkerPool { std::string WarningAboutSize(); protected: - /// A map from the pids of starting worker processes - /// to the number of their unregistered workers. - std::unordered_map starting_worker_processes_; - /// The number of workers per process. - int num_workers_per_process_; + /// The implementation of how to start a new worker process with command arguments. + /// + /// \param worker_command_args The command arguments of new worker process. + /// \return The process ID of started worker process. + virtual pid_t StartProcess(const std::vector &worker_command_args); - private: /// An internal data structure that maintains the pool state per language. struct State { /// The commands and arguments used to start the worker process @@ -151,8 +150,17 @@ class WorkerPool { std::unordered_set> registered_workers; /// All drivers that have registered and are still connected. std::unordered_set> registered_drivers; + /// A map from the pids of starting worker processes + /// to the number of their unregistered workers. + std::unordered_map starting_worker_processes; }; + /// The number of workers per process. + int num_workers_per_process_; + /// Pool states per language. + std::unordered_map states_by_lang_; + + private: /// A helper function that returns the reference of the pool state /// for a given language. inline State &GetStateForLanguage(const Language &language); @@ -162,8 +170,6 @@ class WorkerPool { int multiple_for_warning_; /// The maximum number of workers that can be started concurrently. int maximum_startup_concurrency_; - /// Pool states per language. - std::unordered_map states_by_lang_; /// The last size at which a warning about the number of registered workers /// was generated. int64_t last_warning_multiple_; diff --git a/src/ray/raylet/worker_pool_test.cc b/src/ray/raylet/worker_pool_test.cc index e03c61bb5..4a7f71ea8 100644 --- a/src/ray/raylet/worker_pool_test.cc +++ b/src/ray/raylet/worker_pool_test.cc @@ -9,27 +9,36 @@ namespace ray { namespace raylet { int NUM_WORKERS_PER_PROCESS = 3; +int MAXIMUM_STARTUP_CONCURRENCY = 5; class WorkerPoolMock : public WorkerPool { public: WorkerPoolMock() - : WorkerPool(0, NUM_WORKERS_PER_PROCESS, 1, + : WorkerPool(0, NUM_WORKERS_PER_PROCESS, MAXIMUM_STARTUP_CONCURRENCY, {{Language::PYTHON, {"dummy_py_worker_command"}}, - {Language::JAVA, {"dummy_java_worker_command"}}}) {} - - void StartWorkerProcess(pid_t pid, const Language &language = Language::PYTHON) { - if (starting_worker_processes_.size() > 0) { - // Workers have been started, but not registered. Force start disabled -- returning. - RAY_LOG(DEBUG) << starting_worker_processes_.size() - << " worker processes pending registration"; - return; - } - // Either no workers are pending registration or the worker start is being forced. - RAY_LOG(DEBUG) << "starting new worker process, worker pool size " << Size(language); - starting_worker_processes_.emplace(std::make_pair(pid, num_workers_per_process_)); + {Language::JAVA, {"dummy_java_worker_command"}}}), + last_worker_pid_(0) {} + ~WorkerPoolMock() { + // Avoid killing real processes + states_by_lang_.clear(); } - int NumWorkerProcessesStarting() const { return starting_worker_processes_.size(); } + pid_t StartProcess(const std::vector &worker_command_args) override { + return ++last_worker_pid_; + } + + pid_t LastStartedWorkerProcess() const { return last_worker_pid_; } + + int NumWorkerProcessesStarting() const { + int total = 0; + for (auto &entry : states_by_lang_) { + total += entry.second.starting_worker_processes.size(); + } + return total; + } + + private: + int last_worker_pid_; }; class WorkerPoolTest : public ::testing::Test { @@ -72,8 +81,8 @@ static inline TaskSpecification ExampleTaskSpec( } TEST_F(WorkerPoolTest, HandleWorkerRegistration) { - pid_t pid = 1234; - worker_pool_.StartWorkerProcess(pid); + worker_pool_.StartWorkerProcess(Language::PYTHON); + pid_t pid = worker_pool_.LastStartedWorkerProcess(); std::vector> workers; for (int i = 0; i < NUM_WORKERS_PER_PROCESS; i++) { workers.push_back(CreateWorker(pid)); @@ -97,6 +106,19 @@ 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); + } + // 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); +} + TEST_F(WorkerPoolTest, HandleWorkerPushPop) { // Try to pop a worker from the empty pool and make sure we don't get one. std::shared_ptr popped_worker;