Basic direct actor call support in Python (#5991)

python/ray/worker.py

@@ -251,7 +251,7 @@ class Worker(object):
         """
         self.mode = mode
 
-    def put_object(self, value, object_id=None):
+    def put_object(self, value, object_id=None, return_buffer=None):
         """Put value in the local object store with object id `objectid`.
 
         This assumes that the value for `objectid` has not yet been placed in
@@ -265,6 +265,8 @@ class Worker(object):
             value: The value to put in the object store.
             object_id (object_id.ObjectID): The object ID of the value to be
                 put. If None, one will be generated.
+            return_buffer: If specified, append returns to this list instead
+                of storing directly in the object store.
 
         Returns:
             object_id.ObjectID: The object ID the object was put under.
@@ -284,6 +286,9 @@ class Worker(object):
                 "call 'put' on it (or return it).")
 
         if isinstance(value, bytes):
+            if return_buffer is not None:
+                raise NotImplementedError(
+                    "returning raw buffers from direct actor calls")
             # If the object is a byte array, skip serializing it and
             # use a special metadata to indicate it's raw binary. So
             # that this object can also be read by Java.
@@ -293,9 +298,13 @@ class Worker(object):
                 memcopy_threads=self.memcopy_threads)
 
         if self.use_pickle:
+            if return_buffer is not None:
+                raise NotImplementedError(
+                    "pickle5 serialization with direct actor calls")
             return self._serialize_and_put_pickle5(value, object_id=object_id)
         else:
-            return self._serialize_and_put_pyarrow(value, object_id=object_id)
+            return self._serialize_and_put_pyarrow(
+                value, object_id=object_id, return_buffer=return_buffer)
 
     def _serialize_and_put_pickle5(self, value, object_id=None):
         """Serialize an object using pickle5 and store it in the object store.
@@ -321,13 +330,18 @@ class Worker(object):
             object_id=object_id,
             memcopy_threads=self.memcopy_threads)
 
-    def _serialize_and_put_pyarrow(self, value, object_id=None):
+    def _serialize_and_put_pyarrow(self,
+                                   value,
+                                   object_id=None,
+                                   return_buffer=None):
         """Wraps `store_and_register` with cases for existence and pickling.
 
         Args:
             object_id (object_id.ObjectID): The object ID of the value to be
                 put.
             value: The value to put in the object store.
+            return_buffer: If specified, append returns to this list instead
+                of storing directly in the object store.
         """
         try:
             serialized_value = self._serialize_with_pyarrow(value)
@@ -341,10 +355,13 @@ class Worker(object):
                            "falling back to cloudpickle.".format(type(value)))
             serialized_value = self._serialize_with_pyarrow(value)
 
-        return self.core_worker.put_serialized_object(
-            serialized_value,
-            object_id=object_id,
-            memcopy_threads=self.memcopy_threads)
+        if return_buffer is not None:
+            return_buffer.append(serialized_value)
+        else:
+            return self.core_worker.put_serialized_object(
+                serialized_value,
+                object_id=object_id,
+                memcopy_threads=self.memcopy_threads)
 
     def _serialize_with_pyarrow(self, value, depth=100):
         """Store an object and attempt to register its class if needed.
@@ -721,11 +738,22 @@ def _initialize_serialization(job_id, worker=global_worker):
     serialization_context.set_pickle(pickle.dumps, pickle.loads)
     pyarrow.register_torch_serialization_handlers(serialization_context)
 
+    def id_serializer(obj):
+        if isinstance(obj, ray.ObjectID) and obj.is_direct_actor_type():
+            raise NotImplementedError(
+                "Objects produced by direct actor calls cannot be "
+                "passed to other tasks as arguments.")
+        return pickle.dumps(obj)
+
+    def id_deserializer(serialized_obj):
+        return pickle.loads(serialized_obj)
+
     for id_type in ray._raylet._ID_TYPES:
         serialization_context.register_type(
             id_type,
             "{}.{}".format(id_type.__module__, id_type.__name__),
-            pickle=True)
+            custom_serializer=id_serializer,
+            custom_deserializer=id_deserializer)
 
     def actor_handle_serializer(obj):
         return obj._serialization_helper(True)