Implement async get for direct actor call (#6339)

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,
+            <void*>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 = <object>(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 = <object>(future)
+    loop = py_future._loop
+    loop.call_soon_threadsafe(lambda: py_future.set_result(
+                AsyncGetResponse(
+                    plasma_fallback_id=ObjectID(object_id.Binary()),
+                    result=None)))

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=<Object>
+# 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

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())
 

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":

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)

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

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", {})

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())

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"