From 043d1f40949bd0d16b02c7f38931a534638a4744 Mon Sep 17 00:00:00 2001 From: Edward Oakes Date: Mon, 4 Nov 2019 20:27:57 -0800 Subject: [PATCH] Return RayObjects to core worker (#6052) --- python/ray/_raylet.pyx | 197 +++++++++--------- python/ray/includes/libcoreworker.pxd | 6 +- python/ray/worker.py | 59 ++---- src/ray/common/buffer.h | 6 + src/ray/core_worker/core_worker.cc | 82 +++++--- src/ray/core_worker/core_worker.h | 20 +- src/ray/core_worker/test/mock_worker.cc | 10 +- .../transport/direct_actor_transport.cc | 11 +- .../transport/direct_actor_transport.h | 2 +- .../core_worker/transport/raylet_transport.cc | 19 +- .../core_worker/transport/raylet_transport.h | 7 +- 11 files changed, 227 insertions(+), 192 deletions(-) diff --git a/python/ray/_raylet.pyx b/python/ray/_raylet.pyx index e3fe76352..2b31761e4 100644 --- a/python/ray/_raylet.pyx +++ b/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 (task_type == 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]( - (raw_data_str.data()), raw_data_str.size())) - metadata = dynamic_pointer_cast[ - CBuffer, LocalMemoryBuffer]( - make_shared[LocalMemoryBuffer]( - (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]( + (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]( + (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 + (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) diff --git a/python/ray/includes/libcoreworker.pxd b/python/ray/includes/libcoreworker.pxd index cbde7a278..19c8dda53 100644 --- a/python/ray/includes/libcoreworker.pxd +++ b/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. diff --git a/python/ray/worker.py b/python/ray/worker.py index a23d96163..8b95fa431 100644 --- a/python/ray/worker.py +++ b/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: diff --git a/src/ray/common/buffer.h b/src/ray/common/buffer.h index da9ceadab..5aba5c9d3 100644 --- a/src/ray/common/buffer.h +++ b/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). diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc index aaf30e59f..a31aa335d 100644 --- a/src/ray/core_worker/core_worker.cc +++ b/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(new CoreWorkerRayletTaskReceiver( worker_context_, raylet_client_, task_execution_service_, worker_server_, - execute_task)); + execute_task, exit_handler)); direct_actor_task_receiver_ = std::unique_ptr( new CoreWorkerDirectActorTaskReceiver(worker_context_, task_execution_service_, worker_server_, execute_task, @@ -612,9 +612,40 @@ std::unique_ptr CoreWorker::CreateProfileEvent( void CoreWorker::StartExecutingTasks() { task_execution_service_.run(); } +Status CoreWorker::AllocateReturnObjects( + const std::vector &object_ids, const std::vector &data_sizes, + const std::vector> &metadatas, + std::vector> *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 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(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(data_buffer, metadatas[i]); + } + } + + return Status::OK(); +} + Status CoreWorker::ExecuteTask(const TaskSpecification &task_spec, const ResourceMappingType &resource_ids, - std::vector> *results) { + std::vector> *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> 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(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; } diff --git a/src/ray/core_worker/core_worker.h b/src/ray/core_worker/core_worker.h index e4d58b689..40b099c93 100644 --- a/src/ray/core_worker/core_worker.h +++ b/src/ray/core_worker/core_worker.h @@ -33,7 +33,7 @@ class CoreWorker { const std::unordered_map &required_resources, const std::vector> &args, const std::vector &arg_reference_ids, - const std::vector &return_ids, const bool return_results_directly, + const std::vector &return_ids, std::vector> *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 &object_ids, + const std::vector &data_sizes, + const std::vector> &metadatas, + std::vector> *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> *results); + std::vector> *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), diff --git a/src/ray/core_worker/test/mock_worker.cc b/src/ray/core_worker/test/mock_worker.cc index 0d8e26541..c50d4187c 100644 --- a/src/ray/core_worker/test/mock_worker.cc +++ b/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> &args, const std::vector &arg_reference_ids, const std::vector &return_ids, - const bool return_results_directly, std::vector> *results) { // Note that this doesn't include dummy object id. RAY_CHECK(return_ids.size() >= 0); diff --git a/src/ray/core_worker/transport/direct_actor_transport.cc b/src/ray/core_worker/transport/direct_actor_transport.cc index 07d9f2d66..b0e75a2c8 100644 --- a/src/ray/core_worker/transport/direct_actor_transport.cc +++ b/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> results; - auto status = task_handler_(task_spec, resource_ids, &results); + std::vector> 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(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()); } diff --git a/src/ray/core_worker/transport/direct_actor_transport.h b/src/ray/core_worker/transport/direct_actor_transport.h index 20ea36dd5..9fb604ebb 100644 --- a/src/ray/core_worker/transport/direct_actor_transport.h +++ b/src/ray/core_worker/transport/direct_actor_transport.h @@ -331,7 +331,7 @@ class CoreWorkerDirectActorTaskReceiver : public rpc::DirectActorHandler { public: using TaskHandler = std::function> *results)>; + std::vector> *return_by_value)>; CoreWorkerDirectActorTaskReceiver(WorkerContext &worker_context, boost::asio::io_service &main_io_service, diff --git a/src/ray/core_worker/transport/raylet_transport.cc b/src/ray/core_worker/transport/raylet_transport.cc index 21b955d1b..339df1c14 100644 --- a/src/ray/core_worker/transport/raylet_transport.cc +++ b/src/ray/core_worker/transport/raylet_transport.cc @@ -8,11 +8,12 @@ namespace ray { CoreWorkerRayletTaskReceiver::CoreWorkerRayletTaskReceiver( WorkerContext &worker_context, std::unique_ptr &raylet_client, boost::asio::io_service &io_service, rpc::GrpcServer &server, - const TaskHandler &task_handler) + const TaskHandler &task_handler, const std::function &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> 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 diff --git a/src/ray/core_worker/transport/raylet_transport.h b/src/ray/core_worker/transport/raylet_transport.h index 99735128c..958dea4dc 100644 --- a/src/ray/core_worker/transport/raylet_transport.h +++ b/src/ray/core_worker/transport/raylet_transport.h @@ -14,12 +14,13 @@ class CoreWorkerRayletTaskReceiver : public rpc::WorkerTaskHandler { public: using TaskHandler = std::function> *results)>; + std::vector> *return_by_value)>; CoreWorkerRayletTaskReceiver(WorkerContext &worker_context, std::unique_ptr &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 &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 exit_handler_; /// The callback to process arg wait complete. std::function on_wait_complete_; };