Re-route asyncio plasma code path through raylet instead of direct plasma connection (#7234)

python/ray/_raylet.pyx

@@ -635,9 +635,12 @@ cdef class CoreWorker:
     def set_actor_title(self, title):
         self.core_worker.get().SetActorTitle(title)
 
-    def subscribe_to_plasma(self, plasma_event_handler):
+    def set_plasma_added_callback(self, plasma_event_handler):
         self.plasma_event_handler = plasma_event_handler
-        self.core_worker.get().SubscribeToAsyncPlasma(async_plasma_callback)
+        self.core_worker.get().SetPlasmaAddedCallback(async_plasma_callback)
+
+    def subscribe_to_plasma_object(self, ObjectID object_id):
+        self.core_worker.get().SubscribeToPlasmaAdd(object_id.native())
 
     def get_plasma_event_handler(self):
         return self.plasma_event_handler

python/ray/experimental/async_api.py

@@ -16,7 +16,7 @@ async def _async_init():
         worker = ray.worker.global_worker
         loop = asyncio.get_event_loop()
         handler = PlasmaEventHandler(loop, worker)
-        worker.core_worker.subscribe_to_plasma(handler)
+        worker.core_worker.set_plasma_added_callback(handler)
         logger.debug("AsyncPlasma Connection Created!")
 
 

python/ray/experimental/async_plasma.py

@@ -71,6 +71,8 @@ class PlasmaEventHandler:
 
         future = PlasmaObjectFuture(loop=self._loop)
         self._waiting_dict[object_id].append(future)
-        self.check_immediately(object_id)
-
+        if not self.check_immediately(object_id) and len(
+                self._waiting_dict[object_id]) == 1:
+            # Only subscribe once
+            self._worker.core_worker.subscribe_to_plasma_object(object_id)
         return future

python/ray/includes/libcoreworker.pxd

@@ -197,4 +197,6 @@ cdef extern from "ray/core_worker/core_worker.h" nogil:
                                const double capacity,
                                const CClientID &client_Id)
 
-        void SubscribeToAsyncPlasma(plasma_callback_function callback)
+        void SetPlasmaAddedCallback(plasma_callback_function callback)
+
+        void SubscribeToPlasmaAdd(const CObjectID &object_id)