diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index b3c3352fd..9b73cbd76 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -3231,9 +3231,10 @@ void NodeManager::HandleGetNodeStats(const rpc::GetNodeStatsRequest &node_stats_ // and return the information from HandleNodesStatsRequest. The caller of // HandleGetNodeStats should set a timeout so that the rpc finishes even if not all // workers have replied. - auto all_workers = worker_pool_.GetAllRegisteredWorkers(); + auto all_workers = worker_pool_.GetAllRegisteredWorkers(/* filter_dead_worker */ true); absl::flat_hash_set driver_ids; - for (auto driver : worker_pool_.GetAllRegisteredDrivers()) { + for (auto driver : + worker_pool_.GetAllRegisteredDrivers(/* filter_dead_driver */ true)) { all_workers.push_back(driver); driver_ids.insert(driver->WorkerId()); } @@ -3242,6 +3243,9 @@ void NodeManager::HandleGetNodeStats(const rpc::GetNodeStatsRequest &node_stats_ return; } for (const auto &worker : all_workers) { + if (worker->IsDead()) { + continue; + } rpc::GetCoreWorkerStatsRequest request; request.set_intended_worker_id(worker->WorkerId().Binary()); request.set_include_memory_info(node_stats_request.include_memory_info()); diff --git a/src/ray/raylet/worker_pool.cc b/src/ray/raylet/worker_pool.cc index 8a7e57044..31a9cef74 100644 --- a/src/ray/raylet/worker_pool.cc +++ b/src/ray/raylet/worker_pool.cc @@ -903,30 +903,40 @@ std::vector> WorkerPool::GetWorkersRunningTasks return workers; } -const std::vector> WorkerPool::GetAllRegisteredWorkers() - const { +const std::vector> WorkerPool::GetAllRegisteredWorkers( + bool filter_dead_workers) const { std::vector> workers; for (const auto &entry : states_by_lang_) { for (const auto &worker : entry.second.registered_workers) { - if (worker->IsRegistered()) { - workers.push_back(worker); + if (!worker->IsRegistered()) { + continue; } + + if (filter_dead_workers && worker->IsDead()) { + continue; + } + workers.push_back(worker); } } return workers; } -const std::vector> WorkerPool::GetAllRegisteredDrivers() - const { +const std::vector> WorkerPool::GetAllRegisteredDrivers( + bool filter_dead_drivers) const { std::vector> drivers; for (const auto &entry : states_by_lang_) { for (const auto &driver : entry.second.registered_drivers) { - if (driver->IsRegistered()) { - drivers.push_back(driver); + if (!driver->IsRegistered()) { + continue; } + + if (filter_dead_drivers && driver->IsDead()) { + continue; + } + drivers.push_back(driver); } } diff --git a/src/ray/raylet/worker_pool.h b/src/ray/raylet/worker_pool.h index b78d6cdd5..fce840368 100644 --- a/src/ray/raylet/worker_pool.h +++ b/src/ray/raylet/worker_pool.h @@ -211,13 +211,17 @@ class WorkerPool : public WorkerPoolInterface { /// Get all the registered workers. /// - /// \return A list containing all the workers. - const std::vector> GetAllRegisteredWorkers() const; + /// \param filter_dead_workers whether or not if this method will filter dead workers + /// that are still registered. \return A list containing all the workers. + const std::vector> GetAllRegisteredWorkers( + bool filter_dead_workers = false) const; /// Get all the registered drivers. /// - /// \return A list containing all the drivers. - const std::vector> GetAllRegisteredDrivers() const; + /// \param filter_dead_drivers whether or not if this method will filter dead drivers + /// that are still registered. \return A list containing all the drivers. + const std::vector> GetAllRegisteredDrivers( + bool filter_dead_drivers = false) const; /// Whether there is a pending worker for the given task. /// Note that, this is only used for actor creation task with dynamic options.