make grpc optional to tls_utils

add grpc deps to default and remove from minimal

add DrainNode to python gcs_client, and remove grpc from autoscaler.py

Signed-off-by: Ruiyang Wang <[email protected]>

dashboard/optional_deps.py

@@ -16,3 +16,4 @@
 from aiohttp.typedefs import PathLike # noqa: F401
 from aiohttp.web import RouteDef # noqa: F401
 import pydantic # noqa: F401
+import grpc # noqa: F401

python/ray/_private/tls_utils.py

@@ -2,8 +2,6 @@
 import os
 import socket
 
-import grpc
-
 
 def generate_self_signed_tls_certs():
  """Create self-signed key/cert pair for testing.
@@ -68,6 +66,7 @@ def generate_self_signed_tls_certs():
 
 
 def add_port_to_grpc_server(server, address):
+ import grpc
  if os.environ.get("RAY_USE_TLS", "0").lower() in ("1", "true"):
  server_cert_chain, private_key, ca_cert = load_certs_from_env()
  credentials = grpc.ssl_server_credentials(

python/ray/_raylet.pyx

@@ -2421,6 +2421,22 @@ cdef class GcsClient:
 
  return num_added
 
+ @_auto_reconnect
+ def drain_node(self, node_ids, timeout=None):
+ cdef:
+ c_vector[c_string] c_node_ids
+ int64_t timeout_ms = round(1000 * timeout) if timeout else -1
+ c_vector[c_string] c_drained_node_ids
+ for node_id in node_ids:
+ c_node_ids.push_back(node_id)
+ with nogil:
+ check_status(self.inner.get().DrainNode(
+ c_node_ids, timeout_ms, c_drained_node_ids))
+ result = []
+ for drain_node_id in c_drained_node_ids:
+ result.append(drain_node_id)
+ return result
+
  @_auto_reconnect
  def internal_kv_del(self, c_string key, c_bool del_by_prefix,
  namespace=None, timeout=None):

python/ray/autoscaler/_private/autoscaler.py

@@ -12,7 +12,6 @@
 from enum import Enum
 from typing import Any, Callable, Dict, FrozenSet, List, Optional, Set, Tuple, Union
 
-import grpc
 import yaml
 
 from ray.autoscaler._private.constants import (
@@ -76,7 +75,8 @@
  TAG_RAY_RUNTIME_CONFIG,
  TAG_RAY_USER_NODE_TYPE,
 )
-from ray.core.generated import gcs_service_pb2, gcs_service_pb2_grpc
+from ray._raylet import GcsClient
+from ray.exceptions import RpcError
 
 logger = logging.getLogger(__name__)
 
@@ -184,7 +184,7 @@ def __init__(
  # TODO(ekl): require config reader to be a callable always.
  config_reader: Union[str, Callable[[], dict]],
  load_metrics: LoadMetrics,
- gcs_node_info_stub: gcs_service_pb2_grpc.NodeInfoGcsServiceStub,
+ gcs_client: GcsClient,
  session_name: Optional[str] = None,
  max_launch_batch: int = AUTOSCALER_MAX_LAUNCH_BATCH,
  max_concurrent_launches: int = AUTOSCALER_MAX_CONCURRENT_LAUNCHES,
@@ -214,8 +214,8 @@ def __init__(
  prefix_cluster_info: Whether to add the cluster name to info strs.
  event_summarizer: Utility to consolidate duplicated messages.
  prom_metrics: Prometheus metrics for autoscaler-related operations.
- gcs_node_info_stub: Stub for interactions with Ray nodes via gRPC
- request to the GCS. Used to drain nodes before termination.
+ gcs_client: client for interactions with the GCS. Used to drain nodes
+ before termination.
  """
 
  if isinstance(config_reader, str):
@@ -355,7 +355,7 @@ def read_fn():
  for remote, local in self.config["file_mounts"].items()
  }
 
- self.gcs_node_info_stub = gcs_node_info_stub
+ self.gcs_client = gcs_client
 
  for local_path in self.config["file_mounts"].values():
  assert os.path.exists(local_path)
@@ -675,39 +675,22 @@ def drain_nodes_via_gcs(self, provider_node_ids_to_drain: List[NodeID]):
 
  logger.info(f"Draining {len(raylet_ids_to_drain)} raylet(s).")
  try:
- request = gcs_service_pb2.DrainNodeRequest(
- drain_node_data=[
- gcs_service_pb2.DrainNodeData(node_id=raylet_id)
- for raylet_id in raylet_ids_to_drain
- ]
- )
-
  # A successful response indicates that the GCS has marked the
  # desired nodes as "drained." The cloud provider can then terminate
  # the nodes without the GCS printing an error.
- response = self.gcs_node_info_stub.DrainNode(request, timeout=5)
-
  # Check if we succeeded in draining all of the intended nodes by
  # looking at the RPC response.
- drained_raylet_ids = {
- status_item.node_id for status_item in response.drain_node_status
- }
+ drained_raylet_ids = set(self.gcs_client.drain_node(
+ raylet_ids_to_drain, timeout=5))
  failed_to_drain = raylet_ids_to_drain - drained_raylet_ids
  if failed_to_drain:
  self.prom_metrics.drain_node_exceptions.inc()
  logger.error(f"Failed to drain {len(failed_to_drain)} raylet(s).")
-
- # If we get a gRPC error with an UNIMPLEMENTED code, fail silently.
- # This error indicates that the GCS is using Ray version < 1.8.0,
- # for which DrainNode is not implemented.
- except grpc.RpcError as e:
- # If the code is UNIMPLEMENTED, pass.
- if e.code() == grpc.StatusCode.UNIMPLEMENTED:
- pass
- # Otherwise, it's a plane old gRPC error and we should log it.
- else:
- self.prom_metrics.drain_node_exceptions.inc()
- logger.exception("Failed to drain Ray nodes. Traceback follows.")
+ except RpcError as e:
+ # Otherwise, it's a plane old RPC error and we should log it.
+ self.prom_metrics.drain_node_exceptions.inc()
+ logger.exception("Failed to drain Ray nodes. Traceback follows.")
+ logger.exception("Error: ", e)
  except Exception:
  # We don't need to interrupt the autoscaler update with an
  # exception, but we should log what went wrong and record the

python/ray/includes/common.pxd

@@ -391,6 +391,10 @@ cdef extern from "ray/gcs/gcs_client/gcs_client.h" nogil:
  CRayStatus InternalKVExists(
  const c_string &ns, const c_string &key,
  int64_t timeout_ms, c_bool &exists)
+ CRayStatus DrainNode(
+ const c_vector[c_string]& node_ids,
+ int64_t timeout_ms,
+ c_vector[c_string]& drained_node_ids)
 
  CRayStatus PinRuntimeEnvUri(
  const c_string &uri, int expiration_s, int64_t timeout_ms)

python/setup.py

@@ -256,6 +256,8 @@ def get_packages(self):
  "py-spy >= 0.2.0",
  "requests",
  "gpustat >= 1.0.0", # for windows
+ "grpcio >= 1.32.0; python_version < '3.10'", # noqa:E501
+ "grpcio >= 1.42.0; python_version >= '3.10'", # noqa:E501
  "opencensus",
  "pydantic < 2", # 2.0.0 brings breaking changes
  "prometheus_client >= 0.7.1",

src/ray/gcs/gcs_client/gcs_client.cc

@@ -353,6 +353,34 @@ Status PythonGcsClient::InternalKVExists(const std::string &ns,
  return Status::RpcError(status.error_message(), status.error_code());
 }
 
+Status PythonGcsClient::DrainNode(const std::vector<std::string> &node_ids,
+ int64_t timeout_ms,
+ std::vector<std::string> &drained_node_ids) {
+ grpc::ClientContext context;
+ GrpcClientContextWithTimeoutMs(context, timeout_ms);
+
+ rpc::DrainNodeRequest request;
+ for (const std::string &node_id : node_ids) {
+ request.add_drain_node_data()->set_node_id(node_id);
+ }
+
+ rpc::DrainNodeReply reply;
+
+ grpc::Status status = node_info_stub_->DrainNode(&context, request, &reply);
+ if (status.ok()) {
+ if (reply.status().code() == static_cast<int>(StatusCode::OK)) {
+ drained_node_ids.clear();
+ drained_node_ids.reserve(reply.drain_node_status().size());
+ for (const auto &node_status : reply.drain_node_status()) {
+ drained_node_ids.push_back(node_status.node_id());
+ }
+ return Status::OK();
+ }
+ return HandleGcsError(reply.status());
+ }
+ return Status::RpcError(status.error_message(), status.error_code());
+}
+
 Status PythonGcsClient::PinRuntimeEnvUri(const std::string &uri,
  int expiration_s,
  int64_t timeout_ms) {

src/ray/gcs/gcs_client/gcs_client.h

@@ -224,6 +224,9 @@ class RAY_EXPORT PythonGcsClient {
  const std::string &key,
  int64_t timeout_ms,
  bool &exists);
+ Status DrainNode(const std::vector<std::string> &node_ids,
+ int64_t timeout_ms,
+ std::vector<std::string> &drained_node_ids);
 
  Status PinRuntimeEnvUri(const std::string &uri, int expiration_s, int64_t timeout_ms);
  Status GetAllNodeInfo(int64_t timeout_ms, std::vector<rpc::GcsNodeInfo> &result);