diff --git a/ci/travis/install-dependencies.sh b/ci/travis/install-dependencies.sh index 9cee9ff58..4657b42aa 100755 --- a/ci/travis/install-dependencies.sh +++ b/ci/travis/install-dependencies.sh @@ -36,7 +36,8 @@ elif [[ "$PYTHON" == "3.5" ]] && [[ "$platform" == "linux" ]]; then export PATH="$HOME/miniconda/bin:$PATH" pip install -q scipy tensorflow cython==0.29.0 gym opencv-python-headless pyyaml pandas==0.24.2 requests \ feather-format lxml openpyxl xlrd py-spy setproctitle pytest-timeout networkx tabulate psutil aiohttp \ - uvicorn dataclasses pygments werkzeug kubernetes flask grpcio pytest-sugar pytest-rerunfailures blist + uvicorn dataclasses pygments werkzeug kubernetes flask grpcio pytest-sugar pytest-rerunfailures pytest-asyncio \ + blist elif [[ "$PYTHON" == "2.7" ]] && [[ "$platform" == "macosx" ]]; then # Install miniconda. wget -q https://repo.continuum.io/miniconda/Miniconda2-4.5.4-MacOSX-x86_64.sh -O miniconda.sh -nv @@ -52,7 +53,8 @@ elif [[ "$PYTHON" == "3.5" ]] && [[ "$platform" == "macosx" ]]; then export PATH="$HOME/miniconda/bin:$PATH" pip install -q cython==0.29.0 tensorflow gym opencv-python-headless pyyaml pandas==0.24.2 requests \ feather-format lxml openpyxl xlrd py-spy setproctitle pytest-timeout networkx tabulate psutil aiohttp \ - uvicorn dataclasses pygments werkzeug kubernetes flask grpcio pytest-sugar pytest-rerunfailures blist + uvicorn dataclasses pygments werkzeug kubernetes flask grpcio pytest-sugar pytest-rerunfailures pytest-asyncio \ + blist elif [[ "$LINT" == "1" ]]; then sudo apt-get update sudo apt-get install -y build-essential curl unzip diff --git a/python/ray/_raylet.pyx b/python/ray/_raylet.pyx index 71a083e13..dc964d94c 100644 --- a/python/ray/_raylet.pyx +++ b/python/ray/_raylet.pyx @@ -138,7 +138,7 @@ else: import cPickle as pickle if PY3: - from ray.async_compat import sync_to_async + from ray.async_compat import sync_to_async, AsyncGetResponse def set_internal_config(dict options): @@ -1172,3 +1172,41 @@ cdef class CoreWorker: def current_actor_is_asyncio(self): return self.core_worker.get().GetWorkerContext().CurrentActorIsAsync() + + def in_memory_store_get_async(self, ObjectID object_id, future): + self.core_worker.get().GetAsync( + object_id.native(), + async_set_result_callback, + async_retry_with_plasma_callback, + future) + +cdef void async_set_result_callback(shared_ptr[CRayObject] obj, + CObjectID object_id, + void *future) with gil: + cdef: + c_vector[shared_ptr[CRayObject]] objects_to_deserialize + + py_future = (future) + loop = py_future._loop + + # Object is retrieved from in memory store. + # Here we go through the code path used to deserialize objects. + objects_to_deserialize.push_back(obj) + data_metadata_pairs = RayObjectsToDataMetadataPairs( + objects_to_deserialize) + ids_to_deserialize = [ObjectID(object_id.Binary())] + objects = ray.worker.global_worker.deserialize_objects( + data_metadata_pairs, ids_to_deserialize) + loop.call_soon_threadsafe(lambda: py_future.set_result( + AsyncGetResponse( + plasma_fallback_id=None, result=objects[0]))) + +cdef void async_retry_with_plasma_callback(shared_ptr[CRayObject] obj, + CObjectID object_id, + void *future) with gil: + py_future = (future) + loop = py_future._loop + loop.call_soon_threadsafe(lambda: py_future.set_result( + AsyncGetResponse( + plasma_fallback_id=ObjectID(object_id.Binary()), + result=None))) diff --git a/python/ray/async_compat.py b/python/ray/async_compat.py index 79fdb7d78..2b839c77b 100644 --- a/python/ray/async_compat.py +++ b/python/ray/async_compat.py @@ -2,6 +2,11 @@ This file should only be imported from Python 3. It will raise SyntaxError when importing from Python 2. """ +import asyncio +from collections import namedtuple +import time + +import ray def sync_to_async(func): @@ -11,3 +16,81 @@ def sync_to_async(func): return func(*args, **kwargs) return wrapper + + +# Class encapsulate the get result from direct actor. +# Case 1: plasma_fallback_id=None, result= +# Case 2: plasma_fallback_id=ObjectID, result=None +AsyncGetResponse = namedtuple("AsyncGetResponse", + ["plasma_fallback_id", "result"]) + + +def get_async(object_id): + """Asyncio compatible version of ray.get""" + # Delayed import because raylet import this file and + # it creates circular imports. + from ray.experimental.async_api import init as async_api_init, as_future + from ray.experimental.async_plasma import PlasmaObjectFuture + + assert isinstance(object_id, ray.ObjectID), "Batched get is not supported." + + # Setup + async_api_init() + loop = asyncio.get_event_loop() + core_worker = ray.worker.global_worker.core_worker + + # Here's the callback used to implement async get logic. + # What we want: + # - If direct call, first try to get it from in memory store. + # If the object if promoted to plasma, retry it from plasma API. + # - If not direct call, directly use plasma API to get it. + user_future = loop.create_future() + + # We have three future objects here. + # user_future is directly returned to the user from this function. + # and it will be eventually fulfilled by the final result. + # inner_future is the first attempt to retrieve the object. It can be + # fulfilled by either core_worker.get_async or plasma_api.as_future. + # When inner_future completes, done_callback will be invoked. This + # callback set the final object in user_future if the object hasn't + # been promoted by plasma, otherwise it will retry from plasma. + # retry_plasma_future is only created when we are getting objects that's + # promoted to plasma. It will also invoke the done_callback when it's + # fulfilled. + + def done_callback(future): + result = future.result() + # Result from async plasma, transparently pass it to user future + if isinstance(future, PlasmaObjectFuture): + if isinstance(result, ray.exceptions.RayTaskError): + ray.worker.last_task_error_raise_time = time.time() + user_future.set_exception(result.as_instanceof_cause()) + else: + user_future.set_result(result) + else: + # Result from direct call. + assert isinstance(result, AsyncGetResponse), result + if result.plasma_fallback_id is None: + if isinstance(result.result, ray.exceptions.RayTaskError): + ray.worker.last_task_error_raise_time = time.time() + user_future.set_exception( + result.result.as_instanceof_cause()) + else: + user_future.set_result(result.result) + else: + # Schedule plasma to async get, use the the same callback. + retry_plasma_future = as_future(result.plasma_fallback_id) + retry_plasma_future.add_done_callback(done_callback) + # A hack to keep reference to the future so it doesn't get GC. + user_future.retry_plasma_future = retry_plasma_future + + if object_id.is_direct_call_type(): + inner_future = loop.create_future() + core_worker.in_memory_store_get_async(object_id, inner_future) + else: + inner_future = as_future(object_id) + inner_future.add_done_callback(done_callback) + # A hack to keep reference to inner_future so it doesn't get GC. + user_future.inner_future = inner_future + + return user_future diff --git a/python/ray/experimental/async_api.py b/python/ray/experimental/async_api.py index 8f1eb7246..2a3a6fad2 100644 --- a/python/ray/experimental/async_api.py +++ b/python/ray/experimental/async_api.py @@ -88,9 +88,19 @@ def init(): """ assert ray.is_initialized(), "Please call ray.init before async_api.init" + # Noop when handler is set. + if handler is not None: + return + loop = asyncio.get_event_loop() if loop.is_running(): - asyncio.ensure_future(_async_init()) + assert loop._thread_id != threading.get_ident(), ( + "You are using async_api inside a running event loop. " + "Please call `await _async_init()` to initialize inside " + "asynchrounous context.") + # If the loop is runing outside current thread, we actually need + # to do this to make sure the context is initialized. + asyncio.run_coroutine_threadsafe(_async_init(), loop=loop) else: asyncio.get_event_loop().run_until_complete(_async_init()) diff --git a/python/ray/includes/common.pxd b/python/ray/includes/common.pxd index 17349b23d..ade09ea2f 100644 --- a/python/ray/includes/common.pxd +++ b/python/ray/includes/common.pxd @@ -192,6 +192,7 @@ cdef extern from "ray/common/ray_object.h" nogil: const size_t DataSize() const const shared_ptr[CBuffer] &GetData() const shared_ptr[CBuffer] &GetMetadata() const + c_bool IsInPlasmaError() const cdef extern from "ray/core_worker/common.h" nogil: cdef cppclass CRayFunction "ray::RayFunction": diff --git a/python/ray/includes/libcoreworker.pxd b/python/ray/includes/libcoreworker.pxd index 1ab49cf86..b8153fe13 100644 --- a/python/ray/includes/libcoreworker.pxd +++ b/python/ray/includes/libcoreworker.pxd @@ -36,6 +36,10 @@ from ray.includes.libraylet cimport CRayletClient ctypedef unordered_map[c_string, c_vector[pair[int64_t, double]]] \ ResourceMappingType +ctypedef void (*ray_callback_function) \ + (shared_ptr[CRayObject] result_object, + CObjectID object_id, void* user_data) + cdef extern from "ray/core_worker/profiling.h" nogil: cdef cppclass CProfiler "ray::worker::Profiler": void Start() @@ -145,3 +149,7 @@ cdef extern from "ray/core_worker/core_worker.h" nogil: CWorkerContext &GetWorkerContext() void YieldCurrentFiber(CFiberEvent &coroutine_done) + void GetAsync(const CObjectID &object_id, + ray_callback_function successs_callback, + ray_callback_function fallback_callback, + void* python_future) diff --git a/python/ray/includes/unique_ids.pxi b/python/ray/includes/unique_ids.pxi index 872571b42..5aaf65d6d 100644 --- a/python/ray/includes/unique_ids.pxi +++ b/python/ray/includes/unique_ids.pxi @@ -196,6 +196,15 @@ cdef class ObjectID(BaseID): def from_random(cls): return cls(CObjectID.FromRandom().Binary()) + def __await__(self): + # Delayed import because this can only be imported in py3. + from ray.async_compat import get_async + return get_async(self).__await__() + + def as_future(self): + # Delayed import because this can only be imported in py3. + from ray.async_compat import get_async + return get_async(self) cdef class TaskID(BaseID): cdef CTaskID data diff --git a/python/ray/tests/conftest.py b/python/ray/tests/conftest.py index 3c5ce095b..fa5fb0bda 100644 --- a/python/ray/tests/conftest.py +++ b/python/ray/tests/conftest.py @@ -63,6 +63,13 @@ def ray_start_regular(request): yield res +@pytest.fixture(scope="session") +def ray_start_regular_shared(request): + param = getattr(request, "param", {}) + with _ray_start(**param) as res: + yield res + + @pytest.fixture def ray_start_2_cpus(request): param = getattr(request, "param", {}) diff --git a/python/ray/tests/py3_test.py b/python/ray/tests/py3_test.py index e171db5e4..abc547b72 100644 --- a/python/ray/tests/py3_test.py +++ b/python/ray/tests/py3_test.py @@ -98,7 +98,7 @@ def test_args_intertwined(ray_start_regular): ray.get(remote_test_function.remote(local_method, actor_method)) -def test_asyncio_actor(ray_start_regular): +def test_asyncio_actor(ray_start_regular_shared): @ray.remote class AsyncBatcher(object): def __init__(self): @@ -124,7 +124,7 @@ def test_asyncio_actor(ray_start_regular): assert r1 == r2 == r3 -def test_asyncio_actor_same_thread(ray_start_regular): +def test_asyncio_actor_same_thread(ray_start_regular_shared): @ray.remote class Actor: def sync_thread_id(self): @@ -140,7 +140,7 @@ def test_asyncio_actor_same_thread(ray_start_regular): assert sync_id == async_id -def test_asyncio_actor_concurrency(ray_start_regular): +def test_asyncio_actor_concurrency(ray_start_regular_shared): @ray.remote class RecordOrder: def __init__(self): @@ -170,3 +170,56 @@ def test_asyncio_actor_concurrency(ray_start_regular): answer.append(status) assert history == answer + + +@pytest.mark.asyncio +async def test_asyncio_get(ray_start_regular_shared, event_loop): + loop = event_loop + asyncio.set_event_loop(loop) + loop.set_debug(True) + + # This is needed for async plasma + from ray.experimental.async_api import _async_init + await _async_init() + + # Test Async Plasma + @ray.remote + def task(): + return 1 + + assert await ray.async_compat.get_async(task.remote()) == 1 + + @ray.remote + def task_throws(): + 1 / 0 + + with pytest.raises(ray.exceptions.RayTaskError): + await ray.async_compat.get_async(task_throws.remote()) + + # Test Direct Actor Call + str_len = 200 * 1024 + + @ray.remote + class DirectActor: + def echo(self, i): + return i + + def big_object(self): + # 100Kb is the limit for direct call + return "a" * (str_len) + + def throw_error(self): + 1 / 0 + + direct = DirectActor.options(is_direct_call=True).remote() + + direct_actor_call_future = ray.async_compat.get_async( + direct.echo.remote(2)) + assert await direct_actor_call_future == 2 + + promoted_to_plasma_future = ray.async_compat.get_async( + direct.big_object.remote()) + assert await promoted_to_plasma_future == "a" * str_len + + with pytest.raises(ray.exceptions.RayTaskError): + await ray.async_compat.get_async(direct.throw_error.remote()) diff --git a/python/ray/worker.py b/python/ray/worker.py index 044858967..b0fa9b43c 100644 --- a/python/ray/worker.py +++ b/python/ray/worker.py @@ -61,6 +61,8 @@ LOCAL_MODE = 2 ERROR_KEY_PREFIX = b"Error:" +PY3 = sys.version_info.major >= 3 + # Logger for this module. It should be configured at the entry point # into the program using Ray. Ray provides a default configuration at # entry/init points. @@ -1436,6 +1438,14 @@ def get(object_ids, timeout=None): """ worker = global_worker worker.check_connected() + + if PY3 and hasattr( + worker, + "core_worker") and worker.core_worker.current_actor_is_asyncio(): + raise RayError("Using blocking ray.get inside async actor. " + "This blocks the event loop. Please " + "use `await` on object id with asyncio.gather.") + with profiling.profile("ray.get"): is_individual_id = isinstance(object_ids, ray.ObjectID) if is_individual_id: @@ -1548,6 +1558,13 @@ def wait(object_ids, num_returns=1, timeout=None): """ worker = global_worker + if PY3 and hasattr( + worker, + "core_worker") and worker.core_worker.current_actor_is_asyncio(): + raise RayError("Using blocking ray.wait inside async method. " + "This blocks the event loop. Please use `await` " + "on object id with asyncio.wait. ") + if isinstance(object_ids, ObjectID): raise TypeError( "wait() expected a list of ray.ObjectID, got a single ray.ObjectID" diff --git a/src/ray/common/ray_config_def.h b/src/ray/common/ray_config_def.h index fe64479d0..962b4fa98 100644 --- a/src/ray/common/ray_config_def.h +++ b/src/ray/common/ray_config_def.h @@ -46,7 +46,7 @@ RAY_CONFIG(bool, fair_queueing_enabled, true) RAY_CONFIG(bool, new_scheduler_enabled, false) // The max allowed size in bytes of a return object from direct actor calls. -// Objects larger than this size will be spilled to plasma. +// Objects larger than this size will be spilled/promoted to plasma. RAY_CONFIG(int64_t, max_direct_call_object_size, 100 * 1024) /// The initial period for a task execution lease. The lease will expire this diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc index f024eda81..ed2abadd7 100644 --- a/src/ray/core_worker/core_worker.cc +++ b/src/ray/core_worker/core_worker.cc @@ -1071,4 +1071,17 @@ void CoreWorker::YieldCurrentFiber(FiberEvent &event) { event.Wait(); } +void CoreWorker::GetAsync(const ObjectID &object_id, SetResultCallback success_callback, + SetResultCallback fallback_callback, void *python_future) { + RAY_CHECK(object_id.IsDirectCallType()); + memory_store_->GetAsync(object_id, [python_future, success_callback, fallback_callback, + object_id](std::shared_ptr ray_object) { + if (ray_object->IsInPlasmaError()) { + fallback_callback(ray_object, object_id.WithPlasmaTransportType(), python_future); + } else { + success_callback(ray_object, object_id, python_future); + } + }); +} + } // namespace ray diff --git a/src/ray/core_worker/core_worker.h b/src/ray/core_worker/core_worker.h index 127412db3..9eda952c5 100644 --- a/src/ray/core_worker/core_worker.h +++ b/src/ray/core_worker/core_worker.h @@ -402,6 +402,20 @@ class CoreWorker { /// Block current fiber until event is triggered. void YieldCurrentFiber(FiberEvent &event); + /// The callback expected to be implemented by the client. + using SetResultCallback = + std::function, ObjectID object_id, void *)>; + + /// Perform async get from in-memory store. + /// + /// \param[in] object_id The id to call get on. Assumes object_id.IsDirectCallType(). + /// \param[in] success_callback The callback to use the result object. + /// \param[in] fallback_callback The callback to use when failed to get result. + /// \param[in] python_future the void* object to be passed to SetResultCallback + /// \return void + void GetAsync(const ObjectID &object_id, SetResultCallback success_callback, + SetResultCallback fallback_callback, void *python_future); + private: /// Run the io_service_ event loop. This should be called in a background thread. void RunIOService();