[Core] Make some grpcio imports lazy (ray-project#35705)

Part of ray-project#35472
mohanasudhan · May 24, 2023 · c3232c5 · c3232c5
1 parent f035000
commit c3232c5
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 13 deletions.
diff --git a/python/ray/_private/utils.py b/python/ray/_private/utils.py
@@ -35,8 +35,6 @@
     List,
 )
 
-import grpc
-
 # Import psutil after ray so the packaged version is used.
 import psutil
 from google.protobuf import json_format
@@ -51,11 +49,6 @@
 if TYPE_CHECKING:
     from ray.runtime_env import RuntimeEnv
 
-try:
-    from grpc import aio as aiogrpc
-except ImportError:
-    from grpc.experimental import aio as aiogrpc
-
 
 pwd = None
 if sys.platform != "win32":
@@ -1301,6 +1294,13 @@ def init_grpc_channel(
     options: Optional[Sequence[Tuple[str, Any]]] = None,
     asynchronous: bool = False,
 ):
+    import grpc
+
+    try:
+        from grpc import aio as aiogrpc
+    except ImportError:
+        from grpc.experimental import aio as aiogrpc
+
     grpc_module = aiogrpc if asynchronous else grpc
 
     options = options or []
@@ -1355,8 +1355,8 @@ def internal_kv_list_with_retry(gcs_client, prefix, namespace, num_retries=20):
             result = gcs_client.internal_kv_keys(prefix, namespace)
         except Exception as e:
             if isinstance(e, ray.exceptions.RpcError) and e.rpc_code in (
-                grpc.StatusCode.UNAVAILABLE.value[0],
-                grpc.StatusCode.UNKNOWN.value[0],
+                ray._raylet.GRPC_STATUS_CODE_UNAVAILABLE,
+                ray._raylet.GRPC_STATUS_CODE_UNKNOWN,
             ):
                 logger.warning(
                     f"Unable to connect to GCS at {gcs_client.address}. "
@@ -1389,8 +1389,8 @@ def internal_kv_get_with_retry(gcs_client, key, namespace, num_retries=20):
             result = gcs_client.internal_kv_get(key, namespace)
         except Exception as e:
             if isinstance(e, ray.exceptions.RpcError) and e.rpc_code in (
-                grpc.StatusCode.UNAVAILABLE.value[0],
-                grpc.StatusCode.UNKNOWN.value[0],
+                ray._raylet.GRPC_STATUS_CODE_UNAVAILABLE,
+                ray._raylet.GRPC_STATUS_CODE_UNKNOWN,
             ):
                 logger.warning(
                     f"Unable to connect to GCS at {gcs_client.address}. "
@@ -1447,8 +1447,8 @@ def internal_kv_put_with_retry(gcs_client, key, value, namespace, num_retries=20
             )
         except ray.exceptions.RpcError as e:
             if e.rpc_code in (
-                grpc.StatusCode.UNAVAILABLE.value[0],
-                grpc.StatusCode.UNKNOWN.value[0],
+                ray._raylet.GRPC_STATUS_CODE_UNAVAILABLE,
+                ray._raylet.GRPC_STATUS_CODE_UNKNOWN,
             ):
                 logger.warning(
                     f"Unable to connect to GCS at {gcs_client.address}. "

diff --git a/python/ray/_raylet.pyx b/python/ray/_raylet.pyx
@@ -80,6 +80,7 @@ from ray.includes.common cimport (
     CWorkerType,
     CJobConfig,
     CConcurrencyGroup,
+    CGrpcStatusCode,
     move,
     LANGUAGE_CPP,
     LANGUAGE_JAVA,
@@ -172,6 +173,9 @@ include "includes/metric.pxi"
 # whether C++ optimizations were enabled during compilation.
 OPTIMIZED = __OPTIMIZE__
 
+GRPC_STATUS_CODE_UNAVAILABLE = CGrpcStatusCode.UNAVAILABLE
+GRPC_STATUS_CODE_UNKNOWN = CGrpcStatusCode.UNKNOWN
+
 logger = logging.getLogger(__name__)
 
 # The currently executing task, if any. These are used to synchronize task

diff --git a/python/ray/includes/common.pxd b/python/ray/includes/common.pxd
@@ -316,6 +316,10 @@ cdef extern from "ray/core_worker/common.h" nogil:
         const CNodeID &GetSpilledNodeID() const
 
 cdef extern from "ray/gcs/gcs_client/gcs_client.h" nogil:
+    cdef enum CGrpcStatusCode "grpc::StatusCode":
+        UNAVAILABLE "grpc::StatusCode::UNAVAILABLE",
+        UNKNOWN "grpc::StatusCode::UNKNOWN",
+
     cdef cppclass CGcsClientOptions "ray::gcs::GcsClientOptions":
         CGcsClientOptions(const c_string &gcs_address)