[core] add cython wrapper for raylet client

Signed-off-by: tianyi-ge <tianyig@outlook.com>

Author: tianyi-ge

python/ray/_raylet.pyx

@@ -199,6 +199,7 @@ include "includes/libcoreworker.pxi"
 include "includes/global_state_accessor.pxi"
 include "includes/metric.pxi"
 include "includes/setproctitle.pxi"
+include "includes/raylet_client.pxi"
 
 import ray
 from ray.exceptions import (

python/ray/dashboard/consts.py

@@ -103,4 +103,4 @@
     "RAY_DASHBOARD_SUBPROCESS_MODULE_WAIT_READY_TIMEOUT", 30.0
 )
 
-NODE_MANAGER_RPC_TIMEOUT_SECONDS = 1
+RAYLET_RPC_TIMEOUT_SECONDS = 1

python/ray/dashboard/modules/reporter/reporter_agent.py

@@ -25,7 +25,7 @@
 import ray._private.prometheus_exporter as prometheus_exporter
 import ray.dashboard.modules.reporter.reporter_consts as reporter_consts
 import ray.dashboard.utils as dashboard_utils
-from ray._common.network_utils import build_address, parse_address
+from ray._common.network_utils import parse_address
 from ray._common.utils import (
     get_or_create_event_loop,
     get_user_temp_dir,
@@ -34,18 +34,15 @@
 from ray._private.metrics_agent import Gauge, MetricsAgent, Record
 from ray._private.ray_constants import (
     DEBUG_AUTOSCALING_STATUS,
-    GLOBAL_GRPC_OPTIONS,
     RAY_ENABLE_OPEN_TELEMETRY,
     env_integer,
 )
 from ray._private.telemetry.open_telemetry_metric_recorder import (
     OpenTelemetryMetricRecorder,
 )
 from ray._private.utils import get_system_memory
-from ray._raylet import GCS_PID_KEY, WorkerID
+from ray._raylet import GCS_PID_KEY, RayletClient, WorkerID
 from ray.core.generated import (
-    node_manager_pb2,
-    node_manager_pb2_grpc,
     reporter_pb2,
     reporter_pb2_grpc,
 )
@@ -56,8 +53,8 @@
     COMPONENT_METRICS_TAG_KEYS,
     GCS_RPC_TIMEOUT_SECONDS,
     GPU_TAG_KEYS,
-    NODE_MANAGER_RPC_TIMEOUT_SECONDS,
     NODE_TAG_KEYS,
+    RAYLET_RPC_TIMEOUT_SECONDS,
     TPU_TAG_KEYS,
 )
 from ray.dashboard.modules.reporter.gpu_profile_manager import GpuProfilingManager
@@ -490,10 +487,6 @@ def __init__(self, dashboard_agent):
         # Create GPU metric provider instance
         self._gpu_metric_provider = GpuMetricProvider()
 
-        self._node_manager_address = build_address(
-            self._ip, self._dashboard_agent.node_manager_port
-        )
-
     async def GetTraceback(self, request, context):
         pid = request.pid
         native = request.native
@@ -896,17 +889,14 @@ def _get_disk_io_stats():
                 stats.write_count,
             )
 
-    async def _get_worker_pids_from_raylet(self):
-        channel = ray._private.utils.init_grpc_channel(
-            self._node_manager_address, GLOBAL_GRPC_OPTIONS, asynchronous=True
+    def _get_worker_pids_from_raylet(self) -> Optional[List[int]]:
+        # Get worker pids from raylet via gRPC.
+        timeout = RAYLET_RPC_TIMEOUT_SECONDS * 1000  # in milliseconds
+        raylet_client = RayletClient(
+            ip_address=self._ip, port=self._dashboard_agent.node_manager_port
         )
-        timeout = NODE_MANAGER_RPC_TIMEOUT_SECONDS
-        stub = node_manager_pb2_grpc.NodeManagerServiceStub(channel)
         try:
-            reply = await stub.GetDriverAndWorkerPids(
-                node_manager_pb2.GetDriverAndWorkerPidsRequest(), timeout=timeout
-            )
-            return reply.pids
+            return raylet_client.get_worker_pids(timeout=timeout)
         except Exception as e:
             logger.debug(f"Failed to get worker pids from raylet via gRPC: {e}")
             return None
@@ -920,7 +910,7 @@ def _generate_worker_key(self, proc: psutil.Process) -> Tuple[int, float]:
         return (proc.pid, proc.create_time())
 
     def _get_worker_processes(self):
-        pids = asyncio.run(self._get_worker_pids_from_raylet())
+        pids = self._get_worker_pids_from_raylet()
         if pids is not None:
             workers = {}
             for pid in pids:
@@ -931,25 +921,6 @@ def _get_worker_processes(self):
                     continue
             return workers
 
-        logger.debug("fallback to get worker processes from raylet children")
-        raylet_proc = self._get_raylet_proc()
-        if raylet_proc is None:
-            return []
-        else:
-            workers = {}
-            if sys.platform == "win32":
-                # windows, get the child process not the runner
-                for child in raylet_proc.children():
-                    if child.children():
-                        child = child.children()[0]
-                    workers[self._generate_worker_key(child)] = child
-            else:
-                workers = {
-                    self._generate_worker_key(proc): proc
-                    for proc in raylet_proc.children()
-                }
-            return workers
-
     def _get_workers(self, gpus: Optional[List[GpuUtilizationInfo]] = None):
         workers = self._get_worker_processes()
         if not workers:

python/ray/includes/common.pxd

@@ -763,6 +763,11 @@ cdef extern from "src/ray/protobuf/autoscaler.pb.h" nogil:
         void ParseFromString(const c_string &serialized)
         const c_string &SerializeAsString() const
 
+cdef extern from "ray/raylet_rpc_client/raylet_client.h" nogil:
+    cdef cppclass CRayletClient "ray::rpc::RayletClient":
+        CRayletClient(const c_string &ip_address, int port)
+        CRayStatus GetWorkerPIDs(c_vector[int32_t] &worker_pids, int64_t timeout_ms)
+
 cdef extern from "ray/common/task/task_spec.h" nogil:
     cdef cppclass CConcurrencyGroup "ray::ConcurrencyGroup":
         CConcurrencyGroup(

python/ray/includes/raylet_client.pxi

@@ -0,0 +1,27 @@
+from libcpp.vector cimport vector as c_vector
+from libcpp.string cimport string as c_string
+from libc.stdint cimport int32_t as c_int32_t
+from libcpp.memory cimport unique_ptr, make_unique
+from ray.includes.common cimport CRayletClient, CRayStatus, CAddress
+
+cdef class RayletClient:
+    cdef:
+        unique_ptr[CRayletClient] inner
+
+    def __cinit__(self, ip_address: str, port: int):
+        cdef:
+            c_string c_ip_address
+            c_int32_t c_port
+        c_ip_address = ip_address.encode('utf-8')
+        c_port = <int32_t>port
+        self.inner = make_unique[CRayletClient](c_ip_address, c_port)
+
+    cdef list get_worker_pids(self, timeout_ms: int):
+        cdef:
+            c_vector[c_int32_t] pids
+            CRayStatus status
+        assert self.inner.get() is not NULL
+        status = self.inner.get().GetWorkerPIDs(pids, timeout_ms)
+        if not status.ok():
+            raise RuntimeError("Failed to get worker PIDs from raylet: " + status.message())
+        return [pid for pid in pids]

src/ray/protobuf/node_manager.proto

@@ -399,11 +399,11 @@ message IsLocalWorkerDeadReply {
   bool is_dead = 1;
 }
 
-message GetDriverAndWorkerPidsRequest {}
+message GetWorkerPIDsRequest {}
 
-message GetDriverAndWorkerPidsReply {
+message GetWorkerPIDsReply {
   // PIDs of all drivers and workers managed by the local raylet.
-  repeated uint32 pids = 1;
+  repeated int32 pids = 1;
 }
 
 // Service for inter-node-manager communication.
@@ -520,8 +520,8 @@ service NodeManagerService {
   // Failure: Is currently only used when grpc channel is unavailable for retryable core
   // worker clients. The unavailable callback will eventually be retried so if this fails.
   rpc IsLocalWorkerDead(IsLocalWorkerDeadRequest) returns (IsLocalWorkerDeadReply);
-  // Get the worker managed by local raylet.
-  // Failure: Sends to local raylet, so should never fail.
-  rpc GetDriverAndWorkerPids(GetDriverAndWorkerPidsRequest)
-      returns (GetDriverAndWorkerPidsReply);
+  // Get the PIDs of all workers currently alive that are managed by the local Raylet.
+  // This includes connected driver processes.
+  // Failure: Will retry on failure with logging
+  rpc GetWorkerPIDs(GetWorkerPIDsRequest) returns (GetWorkerPIDsReply);
 }

src/ray/raylet/node_manager.cc

@@ -2723,14 +2723,15 @@ void NodeManager::TriggerGlobalGC() {
   should_local_gc_ = true;
 }
 
-void NodeManager::HandleGetDriverAndWorkerPids(
-    rpc::GetDriverAndWorkerPidsRequest request,
-    rpc::GetDriverAndWorkerPidsReply *reply,
-    rpc::SendReplyCallback send_reply_callback) {
-  auto all_workers = worker_pool_.GetAllRegisteredWorkers(/* filter_dead_worker */ true);
-  const auto &drivers =
-      worker_pool_.GetAllRegisteredDrivers(/* filter_dead_driver */ true);
-  all_workers.insert(all_workers.end(), drivers.begin(), drivers.end());
+void NodeManager::HandleGetWorkerPIDs(rpc::GetWorkerPIDsRequest request,
+                                      rpc::GetWorkerPIDsReply *reply,
+                                      rpc::SendReplyCallback send_reply_callback) {
+  auto all_workers = worker_pool_.GetAllRegisteredWorkers(/* filter_dead_worker */ true,
+                                                          /* filter_io_workers */ true);
+  auto drivers = worker_pool_.GetAllRegisteredDrivers(/* filter_dead_driver */ true);
+  all_workers.insert(all_workers.end(),
+                     std::make_move_iterator(drivers.begin()),
+                     std::make_move_iterator(drivers.end()));
   for (const auto &worker : all_workers) {
     reply->add_pids(worker->GetProcess().GetId());
   }

src/ray/raylet/node_manager.h

@@ -636,6 +636,12 @@ class NodeManager : public rpc::NodeManagerServiceHandler,
                               rpc::NotifyGCSRestartReply *reply,
                               rpc::SendReplyCallback send_reply_callback) override;
 
+  /// Handle a `GetWorkerPIDs` request.
+  void NodeManager::HandleGetWorkerPIDs(
+      rpc::GetWorkerPIDsRequest request,
+      rpc::GetWorkerPIDsReply *reply,
+      rpc::SendReplyCallback send_reply_callback) override;
+
   /// Checks the local socket connection for all registered workers and drivers.
   /// If any of them have disconnected unexpectedly (i.e., we receive a SIGHUP),
   /// we disconnect and kill the worker process.

src/ray/raylet_rpc_client/raylet_client.cc

@@ -47,6 +47,32 @@ RayletClient::RayletClient(const rpc::Address &address,
           std::move(raylet_unavailable_timeout_callback),
           /*server_name=*/std::string("Raylet ") + address.ip_address())) {}
 
+RayletClient::RayletClient(const std::string &ip_address, int port) {
+  io_service_ = std::make_unique<instrumented_io_context>();
+  client_call_manager_ =
+      std::make_unique<rpc::ClientCallManager>(*io_service_, /*record_stats=*/false);
+  grpc_client_ = std::make_unique<rpc::GrpcClient<rpc::NodeManagerService>>(
+      ip_address, port, *client_call_manager_);
+  auto raylet_unavailable_timeout_callback = []() {
+    RAY_LOG(WARNING) << "Raylet is unavailable for "
+                     << ::RayConfig::instance().raylet_rpc_server_reconnect_timeout_s()
+                     << "s";
+  };
+  retryable_grpc_client_ = rpc::RetryableGrpcClient::Create(
+      grpc_client_->Channel(),
+      client_call_manager_->GetMainService(),
+      /*max_pending_requests_bytes=*/
+      std::numeric_limits<uint64_t>::max(),
+      /*check_channel_status_interval_milliseconds=*/
+      ::RayConfig::instance().grpc_client_check_connection_status_interval_milliseconds(),
+      /*server_unavailable_timeout_seconds=*/
+      ::RayConfig::instance().raylet_rpc_server_reconnect_timeout_s(),
+      /*server_unavailable_timeout_callback=*/
+      raylet_unavailable_timeout_callback,
+      /*server_name=*/
+      std::string("Raylet ") + ip_address);
+}
+
 void RayletClient::RequestWorkerLease(
     const rpc::LeaseSpec &lease_spec,
     bool grant_or_reject,
@@ -465,5 +491,30 @@ void RayletClient::GetNodeStats(
                   /*method_timeout_ms*/ -1);
 }
 
+Status RayletClient::GetWorkerPIDs(std::vector<int32_t> &worker_pids,
+                                   int64_t timeout_ms) {
+  rpc::GetWorkerPIDsRequest request;
+  std::promise<Status> promise;
+  auto future = promise.get_future();
+  auto callback = [&promise, &worker_pids](const Status &status,
+                                           rpc::GetWorkerPIDsReply &&reply) {
+    if (status.ok()) {
+      worker_pids = std::vector<int32_t>(reply.pids().begin(), reply.pids().end());
+    }
+    promise.set_value(status);
+  };
+  INVOKE_RPC_CALL(NodeManagerService,
+                  GetWorkerPIDs,
+                  request,
+                  callback,
+                  grpc_client_,
+                  /*method_timeout_ms*/ timeout_ms);
+  if (future.wait_for(std::chrono::milliseconds(timeout_ms)) ==
+      std::future_status::timeout) {
+    return Status::TimedOut("Timed out getting worker PIDs from raylet");
+  }
+  return future.get();
+}
+
 }  // namespace rpc
 }  // namespace ray

src/ray/raylet_rpc_client/raylet_client.h

@@ -43,13 +43,20 @@ class RayletClient : public RayletClientInterface {
   /// Connect to the raylet.
   ///
   /// \param address The IP address of the worker.
-  /// \param port The port that the worker should listen on for gRPC requests. If
-  /// 0, the worker should choose a random port.
   /// \param client_call_manager The client call manager to use for the grpc connection.
+  /// \param raylet_unavailable_timeout_callback callback to be called when the raylet is
+  /// unavailable for a certain period of time.
   explicit RayletClient(const rpc::Address &address,
                         rpc::ClientCallManager &client_call_manager,
                         std::function<void()> raylet_unavailable_timeout_callback);
 
+  /// Connect to the raylet. only used for cython wrapper `CRayletClient`
+  /// `client_call_manager` will be created inside.
+  ///
+  /// \param ip_address The IP address of the worker.
+  /// \param port The port of the worker.
+  explicit RayletClient(const std::string &ip_address, int port);
+
   std::shared_ptr<grpc::Channel> GetChannel() const override;
 
   void RequestWorkerLease(
@@ -163,6 +170,8 @@ class RayletClient : public RayletClientInterface {
   void GetNodeStats(const rpc::GetNodeStatsRequest &request,
                     const rpc::ClientCallback<rpc::GetNodeStatsReply> &callback) override;
 
+  Status GetWorkerPIDs(std::vector<int32_t> &worker_pids, int64_t timeout_ms);
+
  private:
   /// gRPC client to the NodeManagerService.
   std::shared_ptr<rpc::GrpcClient<rpc::NodeManagerService>> grpc_client_;
@@ -177,6 +186,12 @@ class RayletClient : public RayletClientInterface {
 
   /// The number of object ID pin RPCs currently in flight.
   std::atomic<int64_t> pins_in_flight_ = 0;
+
+ private:
+  /// if io context and client call manager are created inside the raylet client, they
+  /// should be kept active during the whole lifetime of client.
+  std::unique_ptr<instrumented_io_context> io_service_;
+  std::unique_ptr<rpc::ClientCallManager> client_call_manager_;
 };
 
 }  // namespace rpc