Return RayObjects to core worker (#6052)

python/ray/_raylet.pyx

@@ -85,6 +85,7 @@ from ray.exceptions import (
     RayTaskError,
     ObjectStoreFullError
 )
+from ray.experimental.no_return import NoReturn
 from ray.function_manager import FunctionDescriptor
 from ray.utils import decode
 from ray.ray_constants import (
@@ -115,6 +116,8 @@ include "includes/libcoreworker.pxi"
 
 logger = logging.getLogger(__name__)
 
+MEMCOPY_THREADS = 12
+
 
 if cpython.PY_MAJOR_VERSION >= 3:
     import pickle
@@ -456,39 +459,6 @@ cdef deserialize_args(
     return ray.signature.recover_args(args)
 
 
-cdef _store_task_outputs(
-        worker, return_ids, outputs,
-        c_bool return_outputs_directly,
-        c_vector[shared_ptr[CRayObject]] *returns):
-
-    # Direct actor call returns are not placed in the object store directly,
-    # but returned to the core worker.
-    if return_outputs_directly:
-        return_buffer = []
-    else:
-        return_buffer = None
-
-    for i in range(len(return_ids)):
-        return_id, output = return_ids[i], outputs[i]
-        if isinstance(output, ray.actor.ActorHandle):
-            raise Exception("Returning an actor handle from a remote "
-                            "function is not allowed).")
-        if output is ray.experimental.no_return.NoReturn:
-            if not worker.core_worker.object_exists(return_id):
-                raise RuntimeError(
-                    "Attempting to return 'ray.experimental.NoReturn' "
-                    "from a remote function, but the corresponding "
-                    "ObjectID does not exist in the local object store.")
-        else:
-            worker.put_object(
-                output, object_id=return_id, return_buffer=return_buffer)
-
-    if return_outputs_directly:
-        assert len(return_ids) == len(return_buffer), \
-            (return_ids, return_buffer)
-        push_objects_into_return_vector(return_buffer, returns)
-
-
 cdef execute_task(
         CTaskType task_type,
         const CRayFunction &ray_function,
@@ -496,7 +466,6 @@ cdef execute_task(
         const c_vector[shared_ptr[CRayObject]] &c_args,
         const c_vector[CObjectID] &c_arg_reference_ids,
         const c_vector[CObjectID] &c_return_ids,
-        c_bool return_outputs_directly,
         c_vector[shared_ptr[CRayObject]] *returns):
 
     worker = ray.worker.global_worker
@@ -562,7 +531,6 @@ cdef execute_task(
         def function_executor(*arguments, **kwarguments):
             return execution_info.function(actor, *arguments, **kwarguments)
 
-    return_ids = VectorToObjectIDs(c_return_ids)
     with core_worker.profile_event(b"task", extra_data=extra_data):
         try:
             task_exception = False
@@ -580,14 +548,13 @@ cdef execute_task(
                     task_exception = True
                     outputs = function_executor(*args, **kwargs)
                     task_exception = False
-                    if len(return_ids) == 1:
+                    if c_return_ids.size() == 1:
                         outputs = (outputs,)
 
             # Store the outputs in the object store.
             with core_worker.profile_event(b"task:store_outputs"):
-                _store_task_outputs(
-                    worker, return_ids, outputs, return_outputs_directly,
-                    returns)
+                core_worker.store_task_outputs(
+                    worker, outputs, c_return_ids, returns)
         except Exception as error:
             if (<int>task_type == <int>TASK_TYPE_ACTOR_CREATION_TASK):
                 worker.mark_actor_init_failed(error)
@@ -601,9 +568,11 @@ cdef execute_task(
             else:
                 failure_object = RayTaskError(function_name, backtrace,
                                               error.__class__)
-            _store_task_outputs(
-                worker, return_ids, [failure_object] * len(return_ids),
-                return_outputs_directly, returns)
+            errors = []
+            for _ in range(c_return_ids.size()):
+                errors.append(failure_object)
+            core_worker.store_task_outputs(
+                worker, errors, c_return_ids, returns)
             ray.utils.push_error_to_driver(
                 worker,
                 ray_constants.TASK_PUSH_ERROR,
@@ -643,7 +612,6 @@ cdef CRayStatus task_execution_handler(
         const c_vector[shared_ptr[CRayObject]] &c_args,
         const c_vector[CObjectID] &c_arg_reference_ids,
         const c_vector[CObjectID] &c_return_ids,
-        c_bool return_results_directly,
         c_vector[shared_ptr[CRayObject]] *returns) nogil:
 
     with gil:
@@ -652,8 +620,7 @@ cdef CRayStatus task_execution_handler(
                 # The call to execute_task should never raise an exception. If
                 # it does, that indicates that there was an internal error.
                 execute_task(task_type, ray_function, c_resources, c_args,
-                             c_arg_reference_ids, c_return_ids,
-                             return_results_directly, returns)
+                             c_arg_reference_ids, c_return_ids, returns)
             except Exception:
                 traceback_str = traceback.format_exc() + (
                     "An unexpected internal error occurred while the worker "
@@ -665,13 +632,9 @@ cdef CRayStatus task_execution_handler(
                     job_id=None)
                 sys.exit(1)
         except SystemExit:
-            if isinstance(threading.current_thread(), threading._MainThread):
-                raise
-            else:
-                # We cannot exit from a non-main thread, so return a special
-                # status that tells the core worker to call sys.exit() on the
-                # main thread instead. This only applies to direct actor calls.
-                return CRayStatus.SystemExit()
+            # Tell the core worker to exit as soon as the result objects
+            # are processed.
+            return CRayStatus.SystemExit()
 
     return CRayStatus.OK()
 
@@ -690,45 +653,6 @@ cdef void exit_handler() nogil:
         sys.exit(0)
 
 
-cdef void push_objects_into_return_vector(
-        py_objects,
-        c_vector[shared_ptr[CRayObject]] *returns):
-
-    cdef:
-        c_string metadata_str = RAW_BUFFER_METADATA
-        c_string raw_data_str
-        shared_ptr[CBuffer] data
-        shared_ptr[CBuffer] metadata
-        shared_ptr[CRayObject] ray_object
-        int64_t data_size
-
-    for serialized_object in py_objects:
-        if isinstance(serialized_object, bytes):
-            data_size = len(serialized_object)
-            raw_data_str = serialized_object
-            data = dynamic_pointer_cast[
-                CBuffer, LocalMemoryBuffer](
-                    make_shared[LocalMemoryBuffer](
-                        <uint8_t*>(raw_data_str.data()), raw_data_str.size()))
-            metadata = dynamic_pointer_cast[
-                CBuffer, LocalMemoryBuffer](
-                    make_shared[LocalMemoryBuffer](
-                        <uint8_t*>(metadata_str.data()), metadata_str.size()))
-            ray_object = make_shared[CRayObject](data, metadata, True)
-            returns.push_back(ray_object)
-        else:
-            data_size = serialized_object.total_bytes
-            data = dynamic_pointer_cast[
-                CBuffer, LocalMemoryBuffer](
-                    make_shared[LocalMemoryBuffer](data_size))
-            metadata.reset()
-            stream = pyarrow.FixedSizeBufferWriter(
-                pyarrow.py_buffer(Buffer.make(data)))
-            serialized_object.write_to(stream)
-            ray_object = make_shared[CRayObject](data, metadata)
-            returns.push_back(ray_object)
-
-
 cdef class CoreWorker:
     cdef unique_ptr[CCoreWorker] core_worker
 
@@ -821,8 +745,8 @@ cdef class CoreWorker:
         # and deal with it here.
         return data.get() == NULL
 
-    def put_serialized_object(self, serialized_object, ObjectID object_id=None,
-                              int memcopy_threads=6):
+    def put_serialized_object(self, serialized_object,
+                              ObjectID object_id=None):
         cdef:
             CObjectID c_object_id
             shared_ptr[CBuffer] data
@@ -834,7 +758,7 @@ cdef class CoreWorker:
         if not object_already_exists:
             stream = pyarrow.FixedSizeBufferWriter(
                 pyarrow.py_buffer(Buffer.make(data)))
-            stream.set_memcopy_threads(memcopy_threads)
+            stream.set_memcopy_threads(MEMCOPY_THREADS)
             serialized_object.write_to(stream)
 
             with nogil:
@@ -843,8 +767,7 @@ cdef class CoreWorker:
 
         return ObjectID(c_object_id.Binary())
 
-    def put_raw_buffer(self, c_string value, ObjectID object_id=None,
-                       int memcopy_threads=6):
+    def put_raw_buffer(self, c_string value, ObjectID object_id=None):
         cdef:
             c_string metadata_str = RAW_BUFFER_METADATA
             CObjectID c_object_id
@@ -859,7 +782,7 @@ cdef class CoreWorker:
         if not object_already_exists:
             stream = pyarrow.FixedSizeBufferWriter(
                 pyarrow.py_buffer(Buffer.make(data)))
-            stream.set_memcopy_threads(memcopy_threads)
+            stream.set_memcopy_threads(MEMCOPY_THREADS)
             stream.write(pyarrow.py_buffer(value))
 
             with nogil:
@@ -869,8 +792,7 @@ cdef class CoreWorker:
         return ObjectID(c_object_id.Binary())
 
     def put_pickle5_buffers(self, c_string inband,
-                            Pickle5Writer writer, ObjectID object_id=None,
-                            int memcopy_threads=6):
+                            Pickle5Writer writer, ObjectID object_id=None):
         cdef:
             CObjectID c_object_id
             c_string metadata_str = PICKLE5_BUFFER_METADATA
@@ -884,7 +806,7 @@ cdef class CoreWorker:
             metadata, writer.get_total_bytes(inband),
             object_id, &c_object_id, &data)
         if not object_already_exists:
-            writer.write_to(inband, data, memcopy_threads)
+            writer.write_to(inband, data, MEMCOPY_THREADS)
             with nogil:
                 check_status(
                     self.core_worker.get().Seal(c_object_id))
@@ -1089,3 +1011,78 @@ cdef class CoreWorker:
             CObjectID c_object_id = object_id.native()
         # Note: faster to not release GIL for short-running op.
         self.core_worker.get().RemoveActiveObjectID(c_object_id)
+
+    # TODO: handle noreturn better
+    cdef store_task_outputs(
+            self, worker, outputs, const c_vector[CObjectID] return_ids,
+            c_vector[shared_ptr[CRayObject]] *returns):
+        cdef:
+            c_vector[size_t] data_sizes
+            c_string metadata_str
+            shared_ptr[CBuffer] empty_metadata
+            c_vector[shared_ptr[CBuffer]] metadatas
+
+        if return_ids.size() == 0:
+            return
+
+        serialized_objects = []
+        for i in range(len(outputs)):
+            return_id, output = return_ids[i], outputs[i]
+            if isinstance(output, ray.actor.ActorHandle):
+                raise Exception("Returning an actor handle from a remote "
+                                "function is not allowed).")
+            elif output is NoReturn:
+                serialized_objects.append(output)
+                data_sizes.push_back(0)
+                metadatas.push_back(empty_metadata)
+            elif isinstance(output, bytes):
+                serialized_objects.append(output)
+                data_sizes.push_back(len(output))
+                metadata_str = RAW_BUFFER_METADATA
+                metadatas.push_back(dynamic_pointer_cast[
+                    CBuffer, LocalMemoryBuffer](
+                    make_shared[LocalMemoryBuffer](
+                        <uint8_t*>(metadata_str.data()),
+                        metadata_str.size(), True)))
+            elif worker.use_pickle:
+                inband, writer = worker._serialize_with_pickle5(output)
+                serialized_objects.append((inband, writer))
+                data_sizes.push_back(writer.get_total_bytes(inband))
+                metadata_str = PICKLE5_BUFFER_METADATA
+                metadatas.push_back(dynamic_pointer_cast[
+                    CBuffer, LocalMemoryBuffer](
+                    make_shared[LocalMemoryBuffer](
+                        <uint8_t*>(metadata_str.data()),
+                        metadata_str.size(), True)))
+            else:
+                serialized_object = worker._serialize_with_pyarrow(output)
+                serialized_objects.append(serialized_object)
+                data_sizes.push_back(serialized_object.total_bytes)
+                metadatas.push_back(empty_metadata)
+
+        check_status(self.core_worker.get().AllocateReturnObjects(
+            return_ids, data_sizes, metadatas, returns))
+
+        for i, serialized_object in enumerate(serialized_objects):
+            # A nullptr is returned if the object already exists.
+            if returns[0][i].get() == NULL:
+                continue
+
+            if serialized_object is NoReturn:
+                returns[0][i].reset()
+            elif isinstance(serialized_object, bytes):
+                buffer = Buffer.make(returns[0][i].get().GetData())
+                stream = pyarrow.FixedSizeBufferWriter(
+                    pyarrow.py_buffer(buffer))
+                stream.set_memcopy_threads(MEMCOPY_THREADS)
+                stream.write(pyarrow.py_buffer(serialized_object))
+            elif worker.use_pickle:
+                inband, writer = serialized_object
+                (<Pickle5Writer>writer).write_to(
+                    inband, returns[0][i].get().GetData(), MEMCOPY_THREADS)
+            else:
+                buffer = Buffer.make(returns[0][i].get().GetData())
+                stream = pyarrow.FixedSizeBufferWriter(
+                    pyarrow.py_buffer(buffer))
+                stream.set_memcopy_threads(MEMCOPY_THREADS)
+                serialized_object.write_to(stream)

python/ray/includes/libcoreworker.pxd

@@ -62,7 +62,6 @@ cdef extern from "ray/core_worker/core_worker.h" nogil:
                         const c_vector[shared_ptr[CRayObject]] &args,
                         const c_vector[CObjectID] &arg_reference_ids,
                         const c_vector[CObjectID] &return_ids,
-                        c_bool is_direct_call,
                         c_vector[shared_ptr[CRayObject]] *returns) nogil,
                     CRayStatus() nogil,
                     void () nogil)
@@ -85,6 +84,11 @@ cdef extern from "ray/core_worker/core_worker.h" nogil:
 
         unique_ptr[CProfileEvent] CreateProfileEvent(
             const c_string &event_type)
+        CRayStatus AllocateReturnObjects(
+            const c_vector[CObjectID] &object_ids,
+            const c_vector[size_t] &data_sizes,
+            const c_vector[shared_ptr[CBuffer]] &metadatas,
+            c_vector[shared_ptr[CRayObject]] *return_objects)
 
         # TODO(edoakes): remove this once the raylet client is no longer used
         # directly.

python/ray/worker.py

@@ -25,7 +25,6 @@ import random
 import pyarrow
 import pyarrow.plasma as plasma
 import ray.cloudpickle as pickle
-import ray.experimental.no_return
 import ray.gcs_utils
 import ray.memory_monitor as memory_monitor
 import ray.node
@@ -128,9 +127,6 @@ class Worker(object):
         # Information used to maintain actor checkpoints.
         self.actor_checkpoint_info = {}
         self.actor_task_counter = 0
-        # The number of threads Plasma should use when putting an object in the
-        # object store.
-        self.memcopy_threads = 12
         # When the worker is constructed. Record the original value of the
         # CUDA_VISIBLE_DEVICES environment variable.
         self.original_gpu_ids = ray.utils.get_cuda_visible_devices()
@@ -251,7 +247,7 @@ class Worker(object):
         """
         self.mode = mode
 
-    def put_object(self, value, object_id=None, return_buffer=None):
+    def put_object(self, value, object_id=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,8 +261,6 @@ 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.
@@ -286,25 +280,15 @@ class Worker(object):
                 "call 'put' on it (or return it).")
 
         if isinstance(value, bytes):
-            if return_buffer is not None:
-                return_buffer.append(value)
-                return
             # 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.
-            return self.core_worker.put_raw_buffer(
-                value,
-                object_id=object_id,
-                memcopy_threads=self.memcopy_threads)
+            return self.core_worker.put_raw_buffer(value, object_id=object_id)
 
         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_buffer=return_buffer)
+            return self._serialize_and_put_pyarrow(value, object_id=object_id)
 
     def _serialize_and_put_pickle5(self, value, object_id=None):
         """Serialize an object using pickle5 and store it in the object store.
@@ -318,33 +302,34 @@ class Worker(object):
             Exception: An exception is raised if the attempt to store the
                 object fails. This can happen if the object store is full.
         """
+        inband, writer = self._serialize_with_pickle5(value)
+        return self.core_worker.put_pickle5_buffers(
+            inband, writer, object_id=object_id)
+
+    def _serialize_with_pickle5(self, value):
         writer = Pickle5Writer()
         if ray.cloudpickle.FAST_CLOUDPICKLE_USED:
             inband = pickle.dumps(
                 value, protocol=5, buffer_callback=writer.buffer_callback)
         else:
             inband = pickle.dumps(value)
-        return self.core_worker.put_pickle5_buffers(
-            inband,
-            writer,
-            object_id=object_id,
-            memcopy_threads=self.memcopy_threads)
+        return inband, writer
 
-    def _serialize_and_put_pyarrow(self,
-                                   value,
-                                   object_id=None,
-                                   return_buffer=None):
+    def _serialize_and_put_pyarrow(self, value, object_id=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.
         """
+        serialized_value = self._serialize_with_pyarrow(value)
+        return self.core_worker.put_serialized_object(
+            serialized_value, object_id=object_id)
+
+    def _serialize_with_pyarrow(self, value):
         try:
-            serialized_value = self._serialize_with_pyarrow(value)
+            serialized_value = self._store_and_register_pyarrow(value)
         except TypeError:
             # TypeError can happen because one of the members of the object
             # may not be serializable for cloudpickle. So we need
@@ -353,17 +338,11 @@ class Worker(object):
             _register_custom_serializer(type(value), use_pickle=True)
             logger.warning("WARNING: Serializing the class {} failed, "
                            "falling back to cloudpickle.".format(type(value)))
-            serialized_value = self._serialize_with_pyarrow(value)
+            serialized_value = self._store_and_register_pyarrow(value)
 
-        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)
+        return serialized_value
 
-    def _serialize_with_pyarrow(self, value, depth=100):
+    def _store_and_register_pyarrow(self, value, depth=100):
         """Store an object and attempt to register its class if needed.
 
         Args: