Add gRPC endpoint to raylet to expose metrics (#6005)

src/ray/protobuf/node_manager.proto

@@ -15,8 +15,24 @@ message ForwardTaskRequest {
 message ForwardTaskReply {
 }
 
+message NodeStatsRequest {
+}
+
+message WorkerStats {
+  // PID of the worker process.
+  uint32 pid = 1;
+  // Whether this is a driver.
+  bool is_driver = 2;
+}
+
+message NodeStatsReply {
+  repeated WorkerStats workers_stats = 1;
+}
+
 // Service for inter-node-manager communication.
 service NodeManagerService {
   // Forward a task and its uncommitted lineage to the remote node manager.
   rpc ForwardTask(ForwardTaskRequest) returns (ForwardTaskReply);
+  // Get the current node stats.
+  rpc GetNodeStats(NodeStatsRequest) returns (NodeStatsReply);
 }

src/ray/raylet/node_manager.cc

@@ -2537,6 +2537,22 @@ std::string NodeManager::DebugString() const {
   return result.str();
 }
 
+void NodeManager::HandleNodeStatsRequest(const rpc::NodeStatsRequest &request,
+                                         rpc::NodeStatsReply *reply,
+                                         rpc::SendReplyCallback send_reply_callback) {
+  for (const auto &worker : worker_pool_.GetAllWorkers()) {
+    auto worker_stats = reply->add_workers_stats();
+    worker_stats->set_pid(worker->Pid());
+    worker_stats->set_is_driver(false);
+  }
+  for (const auto &driver : worker_pool_.GetAllDrivers()) {
+    auto worker_stats = reply->add_workers_stats();
+    worker_stats->set_pid(driver->Pid());
+    worker_stats->set_is_driver(true);
+  }
+  send_reply_callback(Status::OK(), nullptr, nullptr);
+}
+
 void NodeManager::RecordMetrics() const {
   if (stats::StatsConfig::instance().IsStatsDisabled()) {
     return;

src/ray/raylet/node_manager.h

@@ -492,6 +492,11 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
                          rpc::ForwardTaskReply *reply,
                          rpc::SendReplyCallback send_reply_callback) override;
 
+  /// Handle a `NodeStats` request.
+  void HandleNodeStatsRequest(const rpc::NodeStatsRequest &request,
+                              rpc::NodeStatsReply *reply,
+                              rpc::SendReplyCallback send_reply_callback) override;
+
   /// Push an error to the driver if this node is full of actors and so we are
   /// unable to schedule new tasks or actors at all.
   void WarnResourceDeadlock();

src/ray/raylet/worker_pool.cc

@@ -389,6 +389,30 @@ std::vector<std::shared_ptr<Worker>> WorkerPool::GetWorkersRunningTasksForJob(
   return workers;
 }
 
+const std::vector<std::shared_ptr<Worker>> WorkerPool::GetAllWorkers() const {
+  std::vector<std::shared_ptr<Worker>> workers;
+
+  for (const auto &entry : states_by_lang_) {
+    for (const auto &worker : entry.second.registered_workers) {
+      workers.push_back(worker);
+    }
+  }
+
+  return workers;
+}
+
+const std::vector<std::shared_ptr<Worker>> WorkerPool::GetAllDrivers() const {
+  std::vector<std::shared_ptr<Worker>> drivers;
+
+  for (const auto &entry : states_by_lang_) {
+    for (const auto &driver : entry.second.registered_drivers) {
+      drivers.push_back(driver);
+    }
+  }
+
+  return drivers;
+}
+
 void WorkerPool::WarnAboutSize() {
   for (const auto &entry : states_by_lang_) {
     auto state = entry.second;

src/ray/raylet/worker_pool.h

@@ -112,6 +112,16 @@ class WorkerPool {
   std::vector<std::shared_ptr<Worker>> GetWorkersRunningTasksForJob(
       const JobID &job_id) const;
 
+  /// Get all the workers.
+  ///
+  /// \return A list containing all the workers.
+  const std::vector<std::shared_ptr<Worker>> GetAllWorkers() const;
+
+  /// Get all the drivers.
+  ///
+  /// \return A list containing all the drivers.
+  const std::vector<std::shared_ptr<Worker>> GetAllDrivers() const;
+
   /// Whether there is a pending worker for the given task.
   /// Note that, this is only used for actor creation task with dynamic options.
   /// And if the worker registered but isn't assigned a task,

src/ray/rpc/node_manager/node_manager_client.h

@@ -42,6 +42,13 @@ class NodeManagerClient {
             callback);
   }
 
+  /// Get current node stats.
+  void GetNodeStats(const ClientCallback<NodeStatsReply> &callback) {
+    NodeStatsRequest request;
+    client_call_manager_.CreateCall<NodeManagerService, NodeStatsRequest, NodeStatsReply>(
+        *stub_, &NodeManagerService::Stub::PrepareAsyncGetNodeStats, request, callback);
+  }
+
  private:
   /// The gRPC-generated stub.
   std::unique_ptr<NodeManagerService::Stub> stub_;

src/ray/rpc/node_manager/node_manager_server.h

@@ -23,6 +23,17 @@ class NodeManagerServiceHandler {
   virtual void HandleForwardTask(const ForwardTaskRequest &request,
                                  ForwardTaskReply *reply,
                                  SendReplyCallback send_reply_callback) = 0;
+
+  /// Handle a `GetNodeStats` request.
+  /// The implementation can handle this request asynchronously. When handling is done,
+  /// the `send_reply_callback` should be called.
+  ///
+  /// \param[in] request The request message.
+  /// \param[out] reply The reply message.
+  /// \param[in] send_reply_callback The callback to be called when the request is done.
+  virtual void HandleNodeStatsRequest(const NodeStatsRequest &request,
+                                      NodeStatsReply *reply,
+                                      SendReplyCallback send_reply_callback) = 0;
 };
 
 /// The `GrpcService` for `NodeManagerService`.
@@ -43,7 +54,7 @@ class NodeManagerGrpcService : public GrpcService {
       const std::unique_ptr<grpc::ServerCompletionQueue> &cq,
       std::vector<std::pair<std::unique_ptr<ServerCallFactory>, int>>
           *server_call_factories_and_concurrencies) override {
-    // Initialize the factory for `ForwardTask` requests.
+    // Initialize the factory for requests.
     std::unique_ptr<ServerCallFactory> forward_task_call_factory(
         new ServerCallFactoryImpl<NodeManagerService, NodeManagerServiceHandler,
                                   ForwardTaskRequest, ForwardTaskReply>(
@@ -51,9 +62,18 @@ class NodeManagerGrpcService : public GrpcService {
             service_handler_, &NodeManagerServiceHandler::HandleForwardTask, cq,
             main_service_));
 
-    // Set `ForwardTask`'s accept concurrency to 100.
+    std::unique_ptr<ServerCallFactory> node_stats_call_factory(
+        new ServerCallFactoryImpl<NodeManagerService, NodeManagerServiceHandler,
+                                  NodeStatsRequest, NodeStatsReply>(
+            service_, &NodeManagerService::AsyncService::RequestGetNodeStats,
+            service_handler_, &NodeManagerServiceHandler::HandleNodeStatsRequest, cq,
+            main_service_));
+
+    // Set accept concurrency.
     server_call_factories_and_concurrencies->emplace_back(
         std::move(forward_task_call_factory), 100);
+    server_call_factories_and_concurrencies->emplace_back(
+        std::move(node_stats_call_factory), 1);
   }
 
  private: