[ray_client] close ray connection upon client deactivation

python/ray/util/client/server/dataservicer.py

@@ -3,7 +3,7 @@
 import grpc
 import sys
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable
 from threading import Lock
 
 import ray.core.generated.ray_client_pb2 as ray_client_pb2
@@ -17,10 +17,12 @@
 
 
 class DataServicer(ray_client_pb2_grpc.RayletDataStreamerServicer):
-    def __init__(self, basic_service: "RayletServicer"):
+    def __init__(self, basic_service: "RayletServicer",
+                 ray_connect_handler: Callable):
         self.basic_service = basic_service
         self._clients_lock = Lock()
         self._num_clients = 0  # guarded by self._clients_lock
+        self.ray_connect_handler = ray_connect_handler
 
     def Datapath(self, request_iterator, context):
         metadata = {k: v for k, v in context.invocation_metadata()}
@@ -31,6 +33,8 @@ def Datapath(self, request_iterator, context):
         logger.info(f"New data connection from client {client_id}")
         try:
             with self._clients_lock:
+                if self._num_clients == 0 and not ray.is_initialized():
+                    self.ray_connect_handler()
                 self._num_clients += 1
             for req in request_iterator:
                 resp = None
@@ -63,9 +67,13 @@ def Datapath(self, request_iterator, context):
         finally:
             logger.info(f"Lost data connection from client {client_id}")
             self.basic_service.release_all(client_id)
+
             with self._clients_lock:
                 self._num_clients -= 1
 
+            if self._num_clients == 0:
+                ray.shutdown()
+
     def _build_connection_response(self):
         with self._clients_lock:
             cur_num_clients = self._num_clients

python/ray/util/client/server/server.py

@@ -422,10 +422,11 @@ def __getattr__(self, attr):
         return getattr(self.grpc_server, attr)
 
 
-def serve(connection_str):
+def serve(connection_str, ray_connect_handler):
     server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
     task_servicer = RayletServicer()
-    data_servicer = DataServicer(task_servicer)
+    data_servicer = DataServicer(
+        task_servicer, ray_connect_handler=ray_connect_handler)
     logs_servicer = LogstreamServicer()
     ray_client_pb2_grpc.add_RayletDriverServicer_to_server(
         task_servicer, server)
@@ -477,18 +478,21 @@ def main():
         help="Password for connecting to Redis")
     args = parser.parse_args()
     logging.basicConfig(level="INFO")
-    if args.redis_address:
-        if args.redis_password:
-            ray.init(
-                address=args.redis_address,
-                _redis_password=args.redis_password)
+
+    def ray_connect_handler():
+        if args.redis_address:
+            if args.redis_password:
+                ray.init(
+                    address=args.redis_address,
+                    _redis_password=args.redis_password)
+            else:
+                ray.init(address=args.redis_address)
         else:
-            ray.init(address=args.redis_address)
-    else:
-        ray.init()
+            ray.init()
+
     hostport = "%s:%d" % (args.host, args.port)
     logger.info(f"Starting Ray Client server on {hostport}")
-    server = serve(hostport)
+    server = serve(hostport, ray_connect_handler)
     try:
         while True:
             time.sleep(1000)

python/ray/util/client/worker.py

@@ -94,16 +94,16 @@ def __init__(self,
                 # RayletDriverStub, allowing for unary requests.
                 self.server = ray_client_pb2_grpc.RayletDriverStub(
                     self.channel)
-                # Now the HTTP2 channel is ready, or proxied, but the
-                # servicer may not be ready. Call is_initialized() and if
-                # it throws, the servicer is not ready. On success, the
-                # `ray_ready` result is checked.
-                ray_ready = self.is_initialized()
-                if ray_ready:
-                    # Ray is ready! Break out of the retry loop
-                    break
-                # Ray is not ready yet, wait a timeout
-                time.sleep(timeout)
+                # # Now the HTTP2 channel is ready, or proxied, but the
+                # # servicer may not be ready. Call is_initialized() and if
+                # # it throws, the servicer is not ready. On success, the
+                # # `ray_ready` result is checked.
+                # ray_ready = self.is_initialized()
+                # if ray_ready:
+                #     # Ray is ready! Break out of the retry loop
+                #     break
+                # # Ray is not ready yet, wait a timeout
+                # time.sleep(timeout)
             except grpc.FutureTimeoutError:
                 logger.info(
                     f"Couldn't connect channel in {timeout} seconds, retrying")
@@ -120,10 +120,10 @@ def __init__(self,
                         f"retry in {timeout}s...")
             timeout = backoff(timeout)
 
-        # If we made it through the loop without ray_ready it means we've used
-        # up our retries and should error back to the user.
-        if not ray_ready:
-            raise ConnectionError("ray client connection timeout")
+        # # If we made it through the loop without ray_ready it means we've used
+        # # up our retries and should error back to the user.
+        # if not ray_ready:
+        #     raise ConnectionError("ray client connection timeout")
 
         # Initialize the streams to finish protocol negotiation.
         self.data_client = DataClient(self.channel, self._client_id,