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:

src/ray/common/buffer.h

@@ -24,6 +24,8 @@ class Buffer {
   /// Whether this buffer owns the data.
   virtual bool OwnsData() const = 0;
 
+  virtual bool IsPlasmaBuffer() const = 0;
+
   virtual ~Buffer(){};
 
   bool operator==(const Buffer &rhs) const {
@@ -75,6 +77,8 @@ class LocalMemoryBuffer : public Buffer {
 
   bool OwnsData() const override { return has_data_copy_; }
 
+  bool IsPlasmaBuffer() const override { return false; }
+
   ~LocalMemoryBuffer() {}
 
  private:
@@ -105,6 +109,8 @@ class PlasmaBuffer : public Buffer {
 
   bool OwnsData() const override { return true; }
 
+  bool IsPlasmaBuffer() const override { return true; }
+
  private:
   /// shared_ptr to arrow buffer which can potentially hold a reference
   /// for the object (when it's a plasma::PlasmaBuffer).

src/ray/core_worker/core_worker.cc

@@ -114,7 +114,7 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language,
     raylet_task_receiver_ =
         std::unique_ptr<CoreWorkerRayletTaskReceiver>(new CoreWorkerRayletTaskReceiver(
             worker_context_, raylet_client_, task_execution_service_, worker_server_,
-            execute_task));
+            execute_task, exit_handler));
     direct_actor_task_receiver_ = std::unique_ptr<CoreWorkerDirectActorTaskReceiver>(
         new CoreWorkerDirectActorTaskReceiver(worker_context_, task_execution_service_,
                                               worker_server_, execute_task,
@@ -612,9 +612,40 @@ std::unique_ptr<worker::ProfileEvent> CoreWorker::CreateProfileEvent(
 
 void CoreWorker::StartExecutingTasks() { task_execution_service_.run(); }
 
+Status CoreWorker::AllocateReturnObjects(
+    const std::vector<ObjectID> &object_ids, const std::vector<size_t> &data_sizes,
+    const std::vector<std::shared_ptr<Buffer>> &metadatas,
+    std::vector<std::shared_ptr<RayObject>> *return_objects) {
+  RAY_CHECK(object_ids.size() == metadatas.size());
+  RAY_CHECK(object_ids.size() == data_sizes.size());
+  return_objects->resize(object_ids.size(), nullptr);
+
+  for (size_t i = 0; i < object_ids.size(); i++) {
+    bool object_already_exists = false;
+    std::shared_ptr<Buffer> data_buffer;
+    if (data_sizes[i] > 0) {
+      if (!worker_context_.CurrentActorUseDirectCall()) {
+        RAY_RETURN_NOT_OK(
+            Create(metadatas[i], data_sizes[i], object_ids[i], &data_buffer));
+        object_already_exists = !data_buffer;
+      } else {
+        data_buffer = std::make_shared<LocalMemoryBuffer>(data_sizes[i]);
+      }
+    }
+    // Leave the return object as a nullptr if there is no data or metadata.
+    // This allows the caller to prevent the core worker from storing an output
+    // (e.g., to support ray.experimental.no_return.NoReturn).
+    if (!object_already_exists && (data_buffer || metadatas[i])) {
+      return_objects->at(i) = std::make_shared<RayObject>(data_buffer, metadatas[i]);
+    }
+  }
+
+  return Status::OK();
+}
+
 Status CoreWorker::ExecuteTask(const TaskSpecification &task_spec,
                                const ResourceMappingType &resource_ids,
-                               std::vector<std::shared_ptr<RayObject>> *results) {
+                               std::vector<std::shared_ptr<RayObject>> *return_by_value) {
   resource_ids_ = resource_ids;
   worker_context_.SetCurrentTask(task_spec);
   SetCurrentTaskId(task_spec.TaskId());
@@ -642,31 +673,34 @@ Status CoreWorker::ExecuteTask(const TaskSpecification &task_spec,
     return_ids.pop_back();
     task_type = TaskType::ACTOR_TASK;
   }
-  bool direct_call = worker_context_.CurrentActorUseDirectCall();
-  status = task_execution_callback_(
-      task_type, func, task_spec.GetRequiredResources().GetResourceMap(), args,
-      arg_reference_ids, return_ids, direct_call, results);
+
+  std::vector<std::shared_ptr<RayObject>> return_objects;
+  status = task_execution_callback_(task_type, func,
+                                    task_spec.GetRequiredResources().GetResourceMap(),
+                                    args, arg_reference_ids, return_ids, &return_objects);
+
+  for (size_t i = 0; i < return_objects.size(); i++) {
+    // The object is nullptr if it already existed in the object store.
+    if (!return_objects[i]) {
+      continue;
+    }
+    if (return_objects[i]->GetData()->IsPlasmaBuffer()) {
+      if (!Seal(return_ids[i]).ok()) {
+        RAY_LOG(ERROR) << "Task " << task_spec.TaskId() << " failed to seal object "
+                       << return_ids[i] << " in store: " << status.message();
+      }
+    } else if (!worker_context_.CurrentActorUseDirectCall()) {
+      if (!Put(*return_objects[i], return_ids[i]).ok()) {
+        RAY_LOG(ERROR) << "Task " << task_spec.TaskId() << " failed to seal object "
+                       << return_ids[i] << " in store: " << status.message();
+      }
+    } else {
+      return_by_value->push_back(return_objects[i]);
+    }
+  }
 
   SetCurrentTaskId(TaskID::Nil());
   worker_context_.ResetCurrentTask(task_spec);
-
-  // TODO(edoakes): this is only used by java.
-  if (results->size() != 0 && !direct_call) {
-    for (size_t i = 0; i < results->size(); i++) {
-      ObjectID id = ObjectID::ForTaskReturn(
-          task_spec.TaskId(), /*index=*/i + 1,
-          /*transport_type=*/static_cast<int>(TaskTransportType::RAYLET));
-      if (!Put(*results->at(i), id).ok()) {
-        // NOTE(hchen): `PlasmaObjectExists` error is already ignored inside
-        // Put`, we treat other error types as fatal here.
-        RAY_LOG(FATAL) << "Task " << task_spec.TaskId() << " failed to put object " << id
-                       << " in store: " << status.message();
-      } else {
-        RAY_LOG(DEBUG) << "Task " << task_spec.TaskId() << " put object " << id
-                       << " in store.";
-      }
-    }
-  }
   return status;
 }
 

src/ray/core_worker/core_worker.h

@@ -33,7 +33,7 @@ class CoreWorker {
       const std::unordered_map<std::string, double> &required_resources,
       const std::vector<std::shared_ptr<RayObject>> &args,
       const std::vector<ObjectID> &arg_reference_ids,
-      const std::vector<ObjectID> &return_ids, const bool return_results_directly,
+      const std::vector<ObjectID> &return_ids,
       std::vector<std::shared_ptr<RayObject>> *results)>;
 
  public:
@@ -282,6 +282,19 @@ class CoreWorker {
   /// \return void.
   void StartExecutingTasks();
 
+  /// Allocate the return objects for an executing task. The caller should write into the
+  /// data buffers of the allocated buffers.
+  ///
+  /// \param[in] object_ids Object IDs of the return values.
+  /// \param[in] data_sizes Sizes of the return values.
+  /// \param[in] metadatas Metadata buffers of the return values.
+  /// \param[out] return_objects RayObjects containing buffers to write results into.
+  /// \return Status.
+  Status AllocateReturnObjects(const std::vector<ObjectID> &object_ids,
+                               const std::vector<size_t> &data_sizes,
+                               const std::vector<std::shared_ptr<Buffer>> &metadatas,
+                               std::vector<std::shared_ptr<RayObject>> *return_objects);
+
  private:
   /// Run the io_service_ event loop. This should be called in a background thread.
   void RunIOService();
@@ -321,11 +334,12 @@ class CoreWorker {
   ///
   /// \param spec[in] Task specification.
   /// \param spec[in] Resource IDs of resources assigned to this worker.
-  /// \param results[out] Results for task execution.
+  /// \param results[out] Result objects that should be returned by value (not via
+  ///                     plasma).
   /// \return Status.
   Status ExecuteTask(const TaskSpecification &task_spec,
                      const ResourceMappingType &resource_ids,
-                     std::vector<std::shared_ptr<RayObject>> *results);
+                     std::vector<std::shared_ptr<RayObject>> *return_by_value);
 
   /// Build arguments for task executor. This would loop through all the arguments
   /// in task spec, and for each of them that's passed by reference (ObjectID),

src/ray/core_worker/test/mock_worker.cc

@@ -20,11 +20,10 @@ class MockWorker {
  public:
   MockWorker(const std::string &store_socket, const std::string &raylet_socket,
              const gcs::GcsClientOptions &gcs_options)
-      : worker_(
-            WorkerType::WORKER, Language::PYTHON, store_socket, raylet_socket,
-            JobID::FromInt(1), gcs_options, /*log_dir=*/"",
-            /*node_id_address=*/"127.0.0.1",
-            std::bind(&MockWorker::ExecuteTask, this, _1, _2, _3, _4, _5, _6, _7, _8)) {}
+      : worker_(WorkerType::WORKER, Language::PYTHON, store_socket, raylet_socket,
+                JobID::FromInt(1), gcs_options, /*log_dir=*/"",
+                /*node_id_address=*/"127.0.0.1",
+                std::bind(&MockWorker::ExecuteTask, this, _1, _2, _3, _4, _5, _6, _7)) {}
 
   void StartExecutingTasks() { worker_.StartExecutingTasks(); }
 
@@ -34,7 +33,6 @@ class MockWorker {
                      const std::vector<std::shared_ptr<RayObject>> &args,
                      const std::vector<ObjectID> &arg_reference_ids,
                      const std::vector<ObjectID> &return_ids,
-                     const bool return_results_directly,
                      std::vector<std::shared_ptr<RayObject>> *results) {
     // Note that this doesn't include dummy object id.
     RAY_CHECK(return_ids.size() >= 0);

src/ray/core_worker/transport/direct_actor_transport.cc

@@ -266,8 +266,8 @@ void CoreWorkerDirectActorTaskReceiver::HandlePushTask(
         // TODO(edoakes): resource IDs are currently kept track of in the raylet,
         // need to come up with a solution for this.
         ResourceMappingType resource_ids;
-        std::vector<std::shared_ptr<RayObject>> results;
-        auto status = task_handler_(task_spec, resource_ids, &results);
+        std::vector<std::shared_ptr<RayObject>> return_by_value;
+        auto status = task_handler_(task_spec, resource_ids, &return_by_value);
         if (status.IsSystemExit()) {
           // In Python, SystemExit cannot be raised except on the main thread. To work
           // around this when we are executing tasks on worker threads, we re-post the
@@ -275,15 +275,16 @@ void CoreWorkerDirectActorTaskReceiver::HandlePushTask(
           task_main_io_service_.post([this]() { exit_handler_(); });
           return;
         }
-        RAY_CHECK(results.size() == num_returns) << results.size() << "  " << num_returns;
+        RAY_CHECK(return_by_value.size() == num_returns)
+            << return_by_value.size() << "  " << num_returns;
 
-        for (size_t i = 0; i < results.size(); i++) {
+        for (size_t i = 0; i < return_by_value.size(); i++) {
           auto return_object = reply->add_return_objects();
           ObjectID id = ObjectID::ForTaskReturn(
               task_spec.TaskId(), /*index=*/i + 1,
               /*transport_type=*/static_cast<int>(TaskTransportType::DIRECT_ACTOR));
           return_object->set_object_id(id.Binary());
-          const auto &result = results[i];
+          const auto &result = return_by_value[i];
           if (result->GetData() != nullptr) {
             return_object->set_data(result->GetData()->Data(), result->GetData()->Size());
           }

src/ray/core_worker/transport/direct_actor_transport.h

@@ -331,7 +331,7 @@ class CoreWorkerDirectActorTaskReceiver : public rpc::DirectActorHandler {
  public:
   using TaskHandler = std::function<Status(
       const TaskSpecification &task_spec, const ResourceMappingType &resource_ids,
-      std::vector<std::shared_ptr<RayObject>> *results)>;
+      std::vector<std::shared_ptr<RayObject>> *return_by_value)>;
 
   CoreWorkerDirectActorTaskReceiver(WorkerContext &worker_context,
                                     boost::asio::io_service &main_io_service,

src/ray/core_worker/transport/raylet_transport.cc

@@ -8,11 +8,12 @@ namespace ray {
 CoreWorkerRayletTaskReceiver::CoreWorkerRayletTaskReceiver(
     WorkerContext &worker_context, std::unique_ptr<RayletClient> &raylet_client,
     boost::asio::io_service &io_service, rpc::GrpcServer &server,
-    const TaskHandler &task_handler)
+    const TaskHandler &task_handler, const std::function<void()> &exit_handler)
     : worker_context_(worker_context),
       raylet_client_(raylet_client),
       task_service_(io_service, *this),
-      task_handler_(task_handler) {
+      task_handler_(task_handler),
+      exit_handler_(exit_handler) {
   server.RegisterService(task_service_);
 }
 
@@ -56,16 +57,14 @@ void CoreWorkerRayletTaskReceiver::HandleAssignTask(
 
   std::vector<std::shared_ptr<RayObject>> results;
   auto status = task_handler_(task_spec, resource_ids, &results);
-
-  auto num_returns = task_spec.NumReturns();
-  if (task_spec.IsActorCreationTask() || task_spec.IsActorTask()) {
-    RAY_CHECK(num_returns > 0);
-    // Decrease to account for the dummy object id.
-    num_returns--;
+  if (status.IsSystemExit()) {
+    exit_handler_();
+    return;
   }
+  // Raylet transport doesn't currently support returning objects inline.
+  RAY_CHECK(results.size() == 0);
 
-  RAY_LOG(DEBUG) << "Assigned task " << task_spec.TaskId()
-                 << " finished execution. num_returns: " << num_returns;
+  RAY_LOG(DEBUG) << "Assigned task " << task_spec.TaskId() << " finished execution.";
 
   // Notify raylet that current task is done via a `TaskDone` message. This is to
   // ensure that the task is marked as finished by raylet only after previous

src/ray/core_worker/transport/raylet_transport.h

@@ -14,12 +14,13 @@ class CoreWorkerRayletTaskReceiver : public rpc::WorkerTaskHandler {
  public:
   using TaskHandler = std::function<Status(
       const TaskSpecification &task_spec, const ResourceMappingType &resource_ids,
-      std::vector<std::shared_ptr<RayObject>> *results)>;
+      std::vector<std::shared_ptr<RayObject>> *return_by_value)>;
 
   CoreWorkerRayletTaskReceiver(WorkerContext &worker_context,
                                std::unique_ptr<RayletClient> &raylet_client,
                                boost::asio::io_service &io_service,
-                               rpc::GrpcServer &server, const TaskHandler &task_handler);
+                               rpc::GrpcServer &server, const TaskHandler &task_handler,
+                               const std::function<void()> &exit_handler);
 
   /// Handle a `AssignTask` request.
   /// The implementation can handle this request asynchronously. When handling is done,
@@ -41,6 +42,8 @@ class CoreWorkerRayletTaskReceiver : public rpc::WorkerTaskHandler {
   rpc::WorkerTaskGrpcService task_service_;
   /// The callback function to process a task.
   TaskHandler task_handler_;
+  /// The callback function to exit the worker.
+  std::function<void()> exit_handler_;
   /// The callback to process arg wait complete.
   std::function<void(int64_t)> on_wait_complete_;
 };