[ray_client] __getattr__ for the API Import interface (#12089)

* move all things that import real-ray into the server folder

* change the import line and have a __getattr__-able API stub

* formatting

* remove unused (duplicated) util file

* Remove module methods (but leave comment on why)

python/ray/experimental/client/__init__.py

@@ -6,61 +6,46 @@ import logging
 
 logger = logging.getLogger(__name__)
 
+# _client_api has to be external to the API stub, below.
+# Otherwise, ray.remote() that contains ray.remote()
+# contains a reference to the RayAPIStub, therefore a
+# reference to the _client_api, and then tries to pickle
+# the thing.
 _client_api: Optional[APIImpl] = None
 
 
-def check_client_api():
-    global _client_api
-    if _client_api is None:
-        logger.info(
-            "No client API initialized: probably a worker, using core ray")
-        from ray.experimental.client.core_ray_api import set_client_api_as_ray
-        set_client_api_as_ray()
+class RayAPIStub:
+    def connect(self, conn_str):
+        global _client_api
+        from ray.experimental.client.worker import Worker
+        _client_worker = Worker(conn_str)
+        _client_api = ClientAPI(_client_worker)
+
+    def disconnect(self):
+        global _client_api
+        if _client_api is not None:
+            _client_api.close()
+        _client_api = None
+
+    def __getattr__(self, key: str):
+        global _client_api
+        self.__check_client_api()
+        return _client_api.__getattribute__(key)
+
+    def __check_client_api(self):
+        global _client_api
+        if _client_api is None:
+            from ray.experimental.client.server.core_ray_api import CoreRayAPI
+            _client_api = CoreRayAPI()
 
 
-def get(*args, **kwargs):
-    global _client_api
-    check_client_api()
-    return _client_api.get(*args, **kwargs)
+ray = RayAPIStub()
 
-
-def put(*args, **kwargs):
-    global _client_api
-    check_client_api()
-    return _client_api.put(*args, **kwargs)
-
-
-def wait(*args, **kwargs):
-    global _client_api
-    check_client_api()
-    return _client_api.wait(*args, **kwargs)
-
-
-def remote(*args, **kwargs):
-    global _client_api
-    check_client_api()
-    return _client_api.remote(*args, **kwargs)
-
-
-def call_remote(f, *args, **kwargs):
-    global _client_api
-    check_client_api()
-    return _client_api.call_remote(f, *args, **kwargs)
-
-
-def connect(conn_str):
-    global _client_api
-    from ray.experimental.client.worker import Worker
-    _client_worker = Worker(conn_str)
-    _client_api = ClientAPI(_client_worker)
-
-
-def disconnect():
-    global _client_api
-    _client_api.close()
-    _client_api = None
-
-
-def _set_client_api(api: Optional[APIImpl]):
-    global _client_api
-    _client_api = api
+# Someday we might add methods in this module so that someone who
+# tries to `import ray_client as ray` -- as a module, instead of
+# `from ray_client import ray` -- as the API stub
+# still gets expected functionality. This is the way the ray package
+# worked in the past.
+#
+# This really calls for PEP 562: https://www.python.org/dev/peps/pep-0562/
+# But until Python 3.6 is EOL, here we are.

python/ray/experimental/client/client_app.py

@@ -1,4 +1,4 @@
-import ray.experimental.client as ray
+from ray.experimental.client import ray
 
 ray.connect("localhost:50051")
 

python/ray/experimental/client/common.py

@@ -1,5 +1,5 @@
 import ray.core.generated.ray_client_pb2 as ray_client_pb2
-from ray.experimental.client import call_remote
+from ray.experimental.client import ray
 from typing import Any
 from ray import cloudpickle
 
@@ -29,7 +29,7 @@ class ClientRemoteFunc:
     def remote(self, *args, **kwargs):
         if self._raylet_remote_func is not None:
             return self._raylet_remote_func.remote(*args, **kwargs)
-        return call_remote(self, *args, **kwargs)
+        return ray.call_remote(self, *args, **kwargs)
 
     def __repr__(self):
         return "ClientRemoteFunc(%s, %s)" % (self._name, self.id)

python/ray/experimental/client/server/core_ray_api.py

@@ -30,8 +30,3 @@ class CoreRayAPI(APIImpl):
 
     def close(self, *args, **kwargs):
         return None
-
-
-def set_client_api_as_ray():
-    ray_api = CoreRayAPI()
-    ray.experimental.client._set_client_api(ray_api)

python/ray/experimental/client/server/server.py

@@ -6,7 +6,6 @@ import ray
 import ray.core.generated.ray_client_pb2 as ray_client_pb2
 import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc
 import time
-from ray.experimental.client.core_ray_api import set_client_api_as_ray
 from ray.experimental.client.common import convert_from_arg
 from ray.experimental.client.common import ClientObjectRef
 from ray.experimental.client.common import ClientRemoteFunc
@@ -109,7 +108,6 @@ if __name__ == "__main__":
     logging.basicConfig(level="INFO")
     # TODO(barakmich): Perhaps wrap ray init
     ray.init()
-    set_client_api_as_ray()
     server = serve("0.0.0.0:50051")
     try:
         while True:

python/ray/experimental/client/util.py

@@ -1,26 +0,0 @@
-from ray import cloudpickle
-from ray.experimental.client.worker import ClientObjectRef
-
-import ray
-import ray.core.generated.ray_client_pb2 as ray_client_pb2
-
-
-def dump_args_proto(arg):
-    if arg.local == ray_client_pb2.Arg.Locality.INTERNED:
-        return cloudpickle.loads(arg.data)
-    else:
-        # TODO(barakmich): This is a dirty hack that assumes the
-        # server maintains a reference to the ID we've been given
-        ref = ray.ObjectRef(arg.reference_id)
-        return ray.get(ref)
-
-
-def load_args_proto(thing):
-    arg = ray_client_pb2.Arg()
-    if isinstance(thing, ClientObjectRef):
-        arg.local = ray_client_pb2.Arg.Locality.REFERENCE
-        arg.reference_id = thing.id
-    else:
-        arg.local = ray_client_pb2.Arg.Locality.INTERNED
-        arg.data = cloudpickle.dumps(thing)
-    return arg

python/ray/tests/test_experimental_client.py

@@ -1,6 +1,6 @@
 import pytest
-import ray.experimental.client.server as ray_client_server
-import ray.experimental.client as ray
+import ray.experimental.client.server.server as ray_client_server
+from ray.experimental.client import ray
 from ray.experimental.client.common import ClientObjectRef