From 2fe1321c3f137ea62ef71e7ad5117cb3f6594499 Mon Sep 17 00:00:00 2001 From: Barak Michener Date: Wed, 18 Nov 2020 22:42:02 -0800 Subject: [PATCH] [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 | 89 ++++++++----------- python/ray/experimental/client/client_app.py | 2 +- python/ray/experimental/client/common.py | 4 +- .../client/{ => server}/core_ray_api.py | 5 -- .../client/{ => server}/server.py | 2 - python/ray/experimental/client/util.py | 26 ------ python/ray/tests/test_experimental_client.py | 4 +- 7 files changed, 42 insertions(+), 90 deletions(-) rename python/ray/experimental/client/{ => server}/core_ray_api.py (89%) rename python/ray/experimental/client/{ => server}/server.py (97%) delete mode 100644 python/ray/experimental/client/util.py diff --git a/python/ray/experimental/client/__init__.py b/python/ray/experimental/client/__init__.py index 738c5356b..76a2ea91f 100644 --- a/python/ray/experimental/client/__init__.py +++ b/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. diff --git a/python/ray/experimental/client/client_app.py b/python/ray/experimental/client/client_app.py index f6eedeee3..a87223d13 100644 --- a/python/ray/experimental/client/client_app.py +++ b/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") diff --git a/python/ray/experimental/client/common.py b/python/ray/experimental/client/common.py index 17456a4be..eb721fbbd 100644 --- a/python/ray/experimental/client/common.py +++ b/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) diff --git a/python/ray/experimental/client/core_ray_api.py b/python/ray/experimental/client/server/core_ray_api.py similarity index 89% rename from python/ray/experimental/client/core_ray_api.py rename to python/ray/experimental/client/server/core_ray_api.py index 2738e7a4c..4a58af49d 100644 --- a/python/ray/experimental/client/core_ray_api.py +++ b/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) diff --git a/python/ray/experimental/client/server.py b/python/ray/experimental/client/server/server.py similarity index 97% rename from python/ray/experimental/client/server.py rename to python/ray/experimental/client/server/server.py index f70e86e3c..89284ecf8 100644 --- a/python/ray/experimental/client/server.py +++ b/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: diff --git a/python/ray/experimental/client/util.py b/python/ray/experimental/client/util.py deleted file mode 100644 index f571b4a6d..000000000 --- a/python/ray/experimental/client/util.py +++ /dev/null @@ -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 diff --git a/python/ray/tests/test_experimental_client.py b/python/ray/tests/test_experimental_client.py index 1a2ac0b60..647b97578 100644 --- a/python/ray/tests/test_experimental_client.py +++ b/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