Fix release test -- client remote put

python/ray/_private/ray_client_microbenchmark.py

@@ -5,7 +5,6 @@
 from ray.util.client.ray_client_helpers import ray_start_client_server
 
 from ray._private.ray_microbenchmark_helpers import timeit
-from ray._private.ray_microbenchmark_helpers import ray_setup_and_teardown
 
 
 def benchmark_get_calls(ray):
@@ -69,14 +68,26 @@ def actor_concurrent():
 
 
 def main():
-    system_config = {"put_small_object_in_memory_store": True}
-    with ray_setup_and_teardown(
-            logging_level=logging.WARNING, _system_config=system_config):
-        for name, obj in inspect.getmembers(sys.modules[__name__]):
-            if not name.startswith("benchmark_"):
-                continue
-            with ray_start_client_server() as ray:
-                obj(ray)
+    ray_config = {
+        "_system_config": {
+            "put_small_object_in_memory_store": True
+        },
+        "logging_level": logging.WARNING
+    }
+
+    def ray_connect_handler(job_config=None):
+        import ray as real_ray
+        from ray._private.client_mode_hook import disable_client_hook
+        with disable_client_hook():
+            if not real_ray.is_initialized():
+                real_ray.init(**ray_config)
+
+    for name, obj in inspect.getmembers(sys.modules[__name__]):
+        if not name.startswith("benchmark_"):
+            continue
+        with ray_start_client_server(
+                ray_connect_handler=ray_connect_handler) as ray:
+            obj(ray)
 
 
 if __name__ == "__main__":

python/ray/util/client/ray_client_helpers.py

@@ -6,16 +6,19 @@
 
 
 @contextmanager
-def ray_start_client_server(metadata=None):
-    with ray_start_client_server_pair(metadata=metadata) as pair:
+def ray_start_client_server(metadata=None, ray_connect_handler=None):
+    with ray_start_client_server_pair(
+            metadata=metadata,
+            ray_connect_handler=ray_connect_handler) as pair:
         client, server = pair
         yield client
 
 
 @contextmanager
-def ray_start_client_server_pair(metadata=None):
+def ray_start_client_server_pair(metadata=None, ray_connect_handler=None):
     ray._inside_client_test = True
-    server = ray_client_server.serve("localhost:50051")
+    server = ray_client_server.serve(
+        "localhost:50051", ray_connect_handler=ray_connect_handler)
     ray.connect("localhost:50051", metadata=metadata)
     try:
         yield ray, server

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

@@ -82,7 +82,6 @@ def Datapath(self, request_iterator, context):
         finally:
             logger.debug(f"Lost data connection from client {client_id}")
             self.basic_service.release_all(client_id)
-
             with self.clients_lock:
                 if accepted_connection:
                     # Could fail before client accounting happens
@@ -91,8 +90,8 @@ def Datapath(self, request_iterator, context):
 
                 # It's important to keep the Ray shutdown
                 # within this locked context or else Ray could hang.
-                with disable_client_hook():
-                    if self.num_clients == 0:
+                if self.num_clients == 0:
+                    with disable_client_hook():
                         logger.debug("Shutting down ray.")
                         ray.shutdown()