Fix Experimental Async API (#7391)

python/ray/experimental/async_plasma.py

@@ -10,6 +10,23 @@ class PlasmaObjectFuture(asyncio.Future):
     pass
 
 
+def _complete_future(event_handler, ray_object_id):
+    # TODO(ilr): Consider race condition between popping from the
+    # waiting_dict and as_future appending to the waiting_dict's list.
+    logger.debug(
+        "Completing plasma futures for object id {}".format(ray_object_id))
+    obj = event_handler._worker.get_objects([ray_object_id], timeout=0)[0]
+    futures = event_handler._waiting_dict.pop(ray_object_id)
+    for fut in futures:
+        try:
+            fut.set_result(obj)
+        except asyncio.InvalidStateError:
+            # Avoid issues where process_notifications
+            # and check_immediately both get executed
+            logger.debug("Failed to set result for future {}."
+                         "Most likely already set.".format(fut))
+
+
 class PlasmaEventHandler:
     """This class is an event handler for Plasma."""
 
@@ -23,7 +40,10 @@ class PlasmaEventHandler:
         """Process notifications."""
         for object_id, object_size, metadata_size in messages:
             if object_size > 0 and object_id in self._waiting_dict:
-                self._complete_future(object_id)
+                # This must be asynchronous to allow objects to be locally
+                # received
+                self._loop.call_soon_threadsafe(_complete_future, self,
+                                                object_id)
 
     def close(self):
         """Clean up this handler."""
@@ -31,32 +51,10 @@ class PlasmaEventHandler:
             for fut in futures:
                 fut.cancel()
 
-    def _complete_future(self, ray_object_id):
-        # TODO(ilr): Consider race condition between popping from the
-        # waiting_dict and as_future appending to the waiting_dict's list.
-        logger.debug(
-            "Completing plasma futures for object id {}".format(ray_object_id))
-
-        obj = self._worker.get_objects([ray_object_id])[0]
-        futures = self._waiting_dict.pop(ray_object_id)
-        for fut in futures:
-            loop = fut._loop
-
-            def complete_closure():
-                try:
-                    fut.set_result(obj)
-                except asyncio.InvalidStateError:
-                    # Avoid issues where process_notifications
-                    # and check_ready both get executed
-                    logger.debug("Failed to set result for future {}."
-                                 "Most likely already set.".format(fut))
-
-            loop.call_soon_threadsafe(complete_closure)
-
     def check_immediately(self, object_id):
         ready, _ = ray.wait([object_id], timeout=0)
         if ready:
-            self._complete_future(object_id)
+            _complete_future(self, object_id)
 
     def as_future(self, object_id, check_ready=True):
         """Turn an object_id into a Future object.