[Ray] Support new ray cluster through ray client (#2981)

Author: fyrestone

mars/deploy/oscar/tests/test_ray.py

@@ -15,6 +15,10 @@
 import asyncio
 import copy
 import os
+import subprocess
+import sys
+import tempfile
+import threading
 import time
 
 import numpy as np
@@ -236,11 +240,55 @@ def new_ray_session_test():
 
 @require_ray
 def test_ray_client(ray_start_regular):
-    from ray.util.client.ray_client_helpers import ray_start_client_server
-    from ray._private.client_mode_hook import enable_client_mode
-
-    with ray_start_client_server(), enable_client_mode():
-        new_ray_session_test()
+    server_code = """import time
+import ray.util.client.server.server as ray_client_server
+
+server = ray_client_server.init_and_serve("{address}", num_cpus=20)
+print("OK", flush=True)
+while True:
+    time.sleep(1)
+"""
+
+    address = "127.0.0.1:50051"
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
+        f.write(server_code.format(address=address))
+        f.flush()
+
+        proc = subprocess.Popen([sys.executable, "-u", f.name], stdout=subprocess.PIPE)
+
+        try:
+
+            def _check_ready(expect_exit=False):
+                while True:
+                    line = proc.stdout.readline()
+                    if proc.returncode is not None:
+                        if expect_exit:
+                            break
+                        raise Exception(
+                            f"Failed to start ray server at {address}, "
+                            f"the return code is {proc.returncode}."
+                        )
+                    if b"OK" in line:
+                        break
+
+            # Avoid ray.init timeout.
+            _check_ready()
+
+            # Avoid blocking the subprocess when the stdout pipe is full.
+            t = threading.Thread(target=_check_ready, args=(True,))
+            t.start()
+
+            ray.init(f"ray://{address}")
+            ray._inside_client_test = True
+            try:
+                new_ray_session_test()
+            finally:
+                ray._inside_client_test = False
+                ray.shutdown()
+        finally:
+            proc.kill()
+            proc.wait()
 
 
 @require_ray

mars/oscar/backends/ray/communication.py

@@ -14,6 +14,7 @@
 
 import asyncio
 import concurrent.futures as futures
+import functools
 import itertools
 import logging
 import time
@@ -25,6 +26,7 @@
 
 from ....oscar.profiling import ProfilingData
 from ....serialization import serialize, deserialize
+from ....serialization.ray import register_ray_serializers
 from ....metrics import Metrics
 from ....utils import lazy_import, implements, classproperty, Timer
 from ...debug import debug_async_timeout
@@ -79,7 +81,11 @@ def msg_to_simple_str(msg):  # pragma: no cover
     return str(type(msg))
 
 
+_register_ray_serializers_once = functools.lru_cache(1)(register_ray_serializers)
+
+
 def _argwrapper_unpickler(serialized_message):
+    _register_ray_serializers_once()
     return _ArgWrapper(deserialize(*serialized_message))
 
 
@@ -91,6 +97,7 @@ def __init__(self, message):
         self.message = message
 
     def __reduce__(self):
+        _register_ray_serializers_once()
         return _argwrapper_unpickler, (serialize(self.message),)
 
 
@@ -273,9 +280,15 @@ async def handle_task(message: Any, object_ref: "ray.ObjectRef"):
                 try:
                     result = await object_ref
                 except Exception as e:  # pragma: no cover
+                    # The error ClientObjectRef can't be formatted, so
+                    # we give it a string `ClientObjectRef` instead.
+                    try:
+                        object_ref_str = str(object_ref)
+                    except Exception:
+                        object_ref_str = "ClientObjectRef"
                     logger.exception(
                         "Get object %s from %s failed, got exception %s.",
-                        object_ref,
+                        object_ref_str,
                         self.dest_address,
                         e,
                     )