From 2797c11b6983e55081733c8c498eea4290a155ea Mon Sep 17 00:00:00 2001 From: Stephanie Wang Date: Wed, 27 Nov 2019 15:31:44 -0800 Subject: [PATCH] [direct task] For serialized object IDs, check with owner before declaring object unreconstructable (#6286) * Track borrowed vs owned objects * Serialize owner address with object ID * serialize owner task id * Deserialize object IDs * Pass direct task ID instead of plasma ID * it works * Fix ref count test * Add unit test * update warning * we own ray.put objects * missing file * doc * Fix unit test * comments * Fix py2 * lint * update --- python/ray/_raylet.pyx | 25 +++++- python/ray/includes/common.pxd | 4 + python/ray/includes/libcoreworker.pxd | 7 ++ python/ray/includes/task.pxd | 1 - python/ray/serialization.py | 85 +++++++++++++++--- python/ray/tests/test_failure.py | 29 ++++-- src/ray/common/ray_object.cc | 11 +++ src/ray/common/ray_object.h | 2 + src/ray/core_worker/core_worker.cc | 53 ++++++++++- src/ray/core_worker/core_worker.h | 57 +++++++++--- src/ray/core_worker/future_resolver.cc | 54 +++++++++++ src/ray/core_worker/future_resolver.h | 71 +++++++++++++++ src/ray/core_worker/reference_count.cc | 93 +++++++++++-------- src/ray/core_worker/reference_count.h | 90 +++++++++++++++---- src/ray/core_worker/reference_count_test.cc | 99 ++++++++++++++------- src/ray/core_worker/task_manager.cc | 13 +-- src/ray/protobuf/core_worker.proto | 16 ++++ src/ray/rpc/worker/core_worker_client.h | 15 ++++ 18 files changed, 594 insertions(+), 131 deletions(-) create mode 100644 src/ray/core_worker/future_resolver.cc create mode 100644 src/ray/core_worker/future_resolver.h diff --git a/python/ray/_raylet.pyx b/python/ray/_raylet.pyx index 2d5b51933..382e190b7 100644 --- a/python/ray/_raylet.pyx +++ b/python/ray/_raylet.pyx @@ -41,6 +41,7 @@ from libcpp.vector cimport vector as c_vector from cython.operator import dereference, postincrement from ray.includes.common cimport ( + CAddress, CLanguage, CRayObject, CRayStatus, @@ -1052,11 +1053,29 @@ cdef class CoreWorker: # Note: faster to not release GIL for short-running op. self.core_worker.get().RemoveObjectIDReference(c_object_id) - def promote_object_to_plasma(self, ObjectID object_id): + def serialize_and_promote_object_id(self, ObjectID object_id): cdef: CObjectID c_object_id = object_id.native() - self.core_worker.get().PromoteObjectToPlasma(c_object_id) - return object_id.with_plasma_transport_type() + CTaskID c_owner_id = CTaskID.Nil() + CAddress c_owner_address = CAddress() + self.core_worker.get().PromoteToPlasmaAndGetOwnershipInfo( + c_object_id, &c_owner_id, &c_owner_address) + return (object_id, + TaskID(c_owner_id.Binary()), + c_owner_address.SerializeAsString()) + + def deserialize_and_register_object_id( + self, const c_string &object_id_binary, const c_string + &owner_id_binary, const c_string &serialized_owner_address): + cdef: + CObjectID c_object_id = CObjectID.FromBinary(object_id_binary) + CTaskID c_owner_id = CTaskID.FromBinary(owner_id_binary) + CAddress c_owner_address = CAddress() + c_owner_address.ParseFromString(serialized_owner_address) + self.core_worker.get().RegisterOwnershipInfoAndResolveFuture( + c_object_id, + c_owner_id, + c_owner_address) # TODO: handle noreturn better cdef store_task_outputs( diff --git a/python/ray/includes/common.pxd b/python/ray/includes/common.pxd index 76e34bbc2..b228477bc 100644 --- a/python/ray/includes/common.pxd +++ b/python/ray/includes/common.pxd @@ -132,6 +132,10 @@ cdef extern from "ray/protobuf/common.pb.h" nogil: pass cdef cppclass CTaskType "ray::TaskType": pass + cdef cppclass CAddress "ray::rpc::Address": + CAddress() + const c_string &SerializeAsString() + void ParseFromString(const c_string &serialized) # This is a workaround for C++ enum class since Cython has no corresponding diff --git a/python/ray/includes/libcoreworker.pxd b/python/ray/includes/libcoreworker.pxd index 412729a9d..9bd3844dc 100644 --- a/python/ray/includes/libcoreworker.pxd +++ b/python/ray/includes/libcoreworker.pxd @@ -17,6 +17,7 @@ from ray.includes.unique_ids cimport ( CObjectID, ) from ray.includes.common cimport ( + CAddress, CActorCreationOptions, CBuffer, CRayFunction, @@ -116,6 +117,12 @@ cdef extern from "ray/core_worker/core_worker.h" nogil: void AddObjectIDReference(const CObjectID &object_id) void RemoveObjectIDReference(const CObjectID &object_id) void PromoteObjectToPlasma(const CObjectID &object_id) + void PromoteToPlasmaAndGetOwnershipInfo(const CObjectID &object_id, + CTaskID *owner_id, + CAddress *owner_address) + void RegisterOwnershipInfoAndResolveFuture( + const CObjectID &object_id, const CTaskID &owner_id, const + CAddress &owner_address) CRayStatus SetClientOptions(c_string client_name, int64_t limit) CRayStatus Put(const CRayObject &object, CObjectID *object_id) diff --git a/python/ray/includes/task.pxd b/python/ray/includes/task.pxd index 4b742c3f4..a46ed2162 100644 --- a/python/ray/includes/task.pxd +++ b/python/ray/includes/task.pxd @@ -27,7 +27,6 @@ cdef extern from "ray/protobuf/common.pb.h" nogil: cdef cppclass RpcTask "ray::rpc::Task": RpcTaskSpec *mutable_task_spec() - cdef extern from "ray/protobuf/gcs.pb.h" nogil: cdef cppclass TaskTableData "ray::rpc::TaskTableData": RpcTask *mutable_task() diff --git a/python/ray/serialization.py b/python/ray/serialization.py index 39663e4e4..04530ff2d 100644 --- a/python/ray/serialization.py +++ b/python/ray/serialization.py @@ -157,19 +157,52 @@ class SerializationContext(object): serialization_context) def id_serializer(obj): - if isinstance(obj, ray.ObjectID) and obj.is_direct_call_type(): - obj = self.worker.core_worker.promote_object_to_plasma(obj) return pickle.dumps(obj) def id_deserializer(serialized_obj): return pickle.loads(serialized_obj) + def object_id_serializer(obj): + owner_id = "" + owner_address = "" + if obj.is_direct_call_type(): + worker = ray.worker.get_global_worker() + worker.check_connected() + obj, owner_id, owner_address = ( + worker.core_worker.serialize_and_promote_object_id(obj) + ) + obj = obj.__reduce__() + owner_id = owner_id.__reduce__() if owner_id else owner_id + return pickle.dumps((obj, owner_id, owner_address)) + + def object_id_deserializer(serialized_obj): + obj_id, owner_id, owner_address = pickle.loads(serialized_obj) + # Must deserialize the object in the core worker before we + # create the ObjectID to ensure that the reference is added + # before we increment its count to 1. + if owner_id: + worker = ray.worker.get_global_worker() + worker.check_connected() + # UniqueIDs are serialized as + # (class name, (unique bytes,)). + worker.core_worker.deserialize_and_register_object_id( + obj_id[1][0], owner_id[1][0], owner_address) + obj_id = obj_id[0](obj_id[1][0]) + return obj_id + for id_type in ray._raylet._ID_TYPES: - serialization_context.register_type( - id_type, - "{}.{}".format(id_type.__module__, id_type.__name__), - custom_serializer=id_serializer, - custom_deserializer=id_deserializer) + if id_type == ray._raylet.ObjectID: + serialization_context.register_type( + id_type, + "{}.{}".format(id_type.__module__, id_type.__name__), + custom_serializer=object_id_serializer, + custom_deserializer=object_id_deserializer) + else: + serialization_context.register_type( + id_type, + "{}.{}".format(id_type.__module__, id_type.__name__), + custom_serializer=id_serializer, + custom_deserializer=id_deserializer) # We register this serializer on each worker instead of calling # _register_custom_serializer from the driver so that isinstance @@ -188,16 +221,46 @@ class SerializationContext(object): custom_deserializer=actor_handle_deserializer) def id_serializer(obj): - if isinstance(obj, ray.ObjectID) and obj.is_direct_call_type(): - obj = self.worker.core_worker.promote_object_to_plasma(obj) return obj.__reduce__() def id_deserializer(serialized_obj): return serialized_obj[0](*serialized_obj[1]) + def object_id_serializer(obj): + owner_id = "" + owner_address = "" + if obj.is_direct_call_type(): + worker = ray.worker.get_global_worker() + worker.check_connected() + obj, owner_id, owner_address = ( + worker.core_worker.serialize_and_promote_object_id(obj) + ) + obj = id_serializer(obj) + owner_id = id_serializer(owner_id) if owner_id else owner_id + return (obj, owner_id, owner_address) + + def object_id_deserializer(serialized_obj): + obj_id, owner_id, owner_address = serialized_obj + # Must deserialize the object in the core worker before we + # create the ObjectID to ensure that the reference is added + # before we increment its count to 1. + if owner_id: + worker = ray.worker.get_global_worker() + worker.check_connected() + # UniqueIDs are serialized as + # (class name, (unique bytes,)). + worker.core_worker.deserialize_and_register_object_id( + obj_id[1][0], owner_id[1][0], owner_address) + obj_id = id_deserializer(obj_id) + return obj_id + for id_type in ray._raylet._ID_TYPES: - self._register_cloudpickle_serializer(id_type, id_serializer, - id_deserializer) + if id_type == ray._raylet.ObjectID: + self._register_cloudpickle_serializer( + id_type, object_id_serializer, object_id_deserializer) + else: + self._register_cloudpickle_serializer( + id_type, id_serializer, id_deserializer) def initialize(self): """ Register custom serializers """ diff --git a/python/ray/tests/test_failure.py b/python/ray/tests/test_failure.py index dd05c59c3..3d7466e32 100644 --- a/python/ray/tests/test_failure.py +++ b/python/ray/tests/test_failure.py @@ -951,15 +951,13 @@ def test_direct_call_serialized_id_eviction(ray_start_cluster): ray.get(get.remote([obj])) -@pytest.mark.skip( - "Uncomment once eviction errors for serialized IDs are implemented") @pytest.mark.parametrize( "ray_start_cluster", [{ "num_nodes": 2, - "num_cpus": 10, + "num_cpus": 1, }, { "num_nodes": 1, - "num_cpus": 20, + "num_cpus": 2, }], indirect=True) def test_direct_call_serialized_id(ray_start_cluster): @@ -971,16 +969,33 @@ def test_direct_call_serialized_id(ray_start_cluster): return 1 @ray.remote - def get(obj_ids): + def dependent_task(x): + return x + + @ray.remote + def get(obj_ids, test_dependent_task): print("get", obj_ids) obj_id = obj_ids[0] - assert ray.get(obj_id) == 1 + if test_dependent_task: + assert ray.get(dependent_task.remote(obj_id)) == 1 + else: + assert ray.get(obj_id) == 1 small_object = small_object.options(is_direct_call=True) + dependent_task = dependent_task.options(is_direct_call=True) get = get.options(is_direct_call=True) obj = small_object.remote() - ray.get(get.remote([obj])) + ray.get(get.remote([obj], False)) + + obj = small_object.remote() + ray.get(get.remote([obj], True)) + + obj = ray.put(1) + ray.get(get.remote([obj], False)) + + obj = ray.put(1) + ray.get(get.remote([obj], True)) if __name__ == "__main__": diff --git a/src/ray/common/ray_object.cc b/src/ray/common/ray_object.cc index faabaaee5..ecc652b8f 100644 --- a/src/ray/common/ray_object.cc +++ b/src/ray/common/ray_object.cc @@ -2,6 +2,17 @@ namespace ray { +std::shared_ptr MakeErrorMetadataBuffer(rpc::ErrorType error_type) { + std::string meta = std::to_string(static_cast(error_type)); + auto metadata = const_cast(reinterpret_cast(meta.data())); + auto meta_buffer = + std::make_shared(metadata, meta.size(), /*copy_data=*/true); + return meta_buffer; +} + +RayObject::RayObject(rpc::ErrorType error_type) + : RayObject(nullptr, MakeErrorMetadataBuffer(error_type)) {} + bool RayObject::IsException(rpc::ErrorType *error_type) const { if (metadata_ == nullptr) { return false; diff --git a/src/ray/common/ray_object.h b/src/ray/common/ray_object.h index 6c82a5d6f..4f8a5e003 100644 --- a/src/ray/common/ray_object.h +++ b/src/ray/common/ray_object.h @@ -42,6 +42,8 @@ class RayObject { RAY_CHECK(data_ || metadata_) << "Data and metadata cannot both be empty."; } + RayObject(rpc::ErrorType error_type); + /// Return the data of the ray object. const std::shared_ptr &GetData() const { return data_; }; diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc index db9d7d3e9..80ac7e608 100644 --- a/src/ray/core_worker/core_worker.cc +++ b/src/ray/core_worker/core_worker.cc @@ -164,6 +164,7 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language, RAY_CHECK_OK(plasma_store_provider_->Put(obj, obj_id)); }, ref_counting_enabled ? reference_counter_ : nullptr, raylet_client_)); + task_manager_.reset(new TaskManager(memory_store_)); resolver_.reset(new LocalDependencyResolver(memory_store_)); @@ -208,6 +209,7 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language, }, memory_store_, task_manager_, RayConfig::instance().worker_lease_timeout_milliseconds())); + future_resolver_.reset(new FutureResolver(memory_store_, client_factory, io_service_)); } CoreWorker::~CoreWorker() { @@ -281,13 +283,36 @@ void CoreWorker::ReportActiveObjectIDs() { heartbeat_timer_.async_wait(boost::bind(&CoreWorker::ReportActiveObjectIDs, this)); } -void CoreWorker::PromoteObjectToPlasma(const ObjectID &object_id) { +void CoreWorker::PromoteToPlasmaAndGetOwnershipInfo(const ObjectID &object_id, + TaskID *owner_id, + rpc::Address *owner_address) { RAY_CHECK(object_id.IsDirectCallType()); auto value = memory_store_->GetOrPromoteToPlasma(object_id); if (value != nullptr) { RAY_CHECK_OK( plasma_store_provider_->Put(*value, object_id.WithPlasmaTransportType())); } + + auto has_owner = reference_counter_->GetOwner(object_id, owner_id, owner_address); + RAY_CHECK(has_owner) + << "Object IDs generated randomly (ObjectID.from_random()) or out-of-band " + "(ObjectID.from_binary(...)) cannot be serialized because Ray does not know " + "which task will create them. " + "If this was not how your object ID was generated, please file an issue " + "at https://github.com/ray-project/ray/issues/"; +} + +void CoreWorker::RegisterOwnershipInfoAndResolveFuture( + const ObjectID &object_id, const TaskID &owner_id, + const rpc::Address &owner_address) { + // Add the object's owner to the local metadata in case it gets serialized + // again. + reference_counter_->AddBorrowedObject(object_id, owner_id, owner_address); + + RAY_CHECK(!owner_id.IsNil()); + // We will ask the owner about the object until the object is + // created or we can no longer reach the owner. + future_resolver_->ResolveFutureAsync(object_id, owner_id, owner_address); } Status CoreWorker::SetClientOptions(std::string name, int64_t limit_bytes) { @@ -299,6 +324,8 @@ Status CoreWorker::Put(const RayObject &object, ObjectID *object_id) { *object_id = ObjectID::ForPut(worker_context_.GetCurrentTaskID(), worker_context_.GetNextPutIndex(), static_cast(TaskTransportType::RAYLET)); + reference_counter_->AddOwnedObject(*object_id, GetCallerId(), rpc_address_, + std::make_shared>()); return Put(object, *object_id); } @@ -557,7 +584,8 @@ void CoreWorker::PinObjectReferences(const TaskSpecification &task_spec, // Note that we call this even if task_deps.size() == 0, in order to pin the return id. for (size_t i = 0; i < num_returns; i++) { - reference_counter_->SetDependencies(task_spec.ReturnId(i, transport_type), task_deps); + reference_counter_->AddOwnedObject(task_spec.ReturnId(i, transport_type), + GetCallerId(), rpc_address_, task_deps); } } @@ -928,6 +956,27 @@ void CoreWorker::HandleDirectActorCallArgWaitComplete( }); } +void CoreWorker::HandleGetObjectStatus(const rpc::GetObjectStatusRequest &request, + rpc::GetObjectStatusReply *reply, + rpc::SendReplyCallback send_reply_callback) { + ObjectID object_id = ObjectID::FromBinary(request.object_id()); + TaskID owner_id = TaskID::FromBinary(request.owner_id()); + if (owner_id != GetCallerId()) { + // We may have owned this object in the past, but we are now executing some + // other task or actor. + reply->set_status(rpc::GetObjectStatusReply::WRONG_OWNER); + } else { + if (task_manager_->IsTaskPending(object_id.TaskId())) { + reply->set_status(rpc::GetObjectStatusReply::PENDING); + } else { + // TODO: We could probably just send the object value if it is small + // enough and we have it local. + reply->set_status(rpc::GetObjectStatusReply::CREATED); + } + } + send_reply_callback(Status::OK(), nullptr, nullptr); +} + void CoreWorker::YieldCurrentFiber(FiberEvent &event) { RAY_CHECK(worker_context_.CurrentActorIsAsync()); boost::this_fiber::yield(); diff --git a/src/ray/core_worker/core_worker.h b/src/ray/core_worker/core_worker.h index 2746a9f7c..b2590ab72 100644 --- a/src/ray/core_worker/core_worker.h +++ b/src/ray/core_worker/core_worker.h @@ -7,6 +7,7 @@ #include "ray/core_worker/actor_handle.h" #include "ray/core_worker/common.h" #include "ray/core_worker/context.h" +#include "ray/core_worker/future_resolver.h" #include "ray/core_worker/profiling.h" #include "ray/core_worker/reference_count.h" #include "ray/core_worker/store_provider/memory_store/memory_store.h" @@ -25,10 +26,11 @@ /// 1) Add the rpc to the CoreWorkerService in core_worker.proto, e.g., "ExampleCall" /// 2) Add a new handler to the macro below: "RAY_CORE_WORKER_RPC_HANDLER(ExampleCall, 1)" /// 3) Add a method to the CoreWorker class below: "CoreWorker::HandleExampleCall" -#define RAY_CORE_WORKER_RPC_HANDLERS \ - RAY_CORE_WORKER_RPC_HANDLER(AssignTask, 5) \ - RAY_CORE_WORKER_RPC_HANDLER(PushTask, 9999) \ - RAY_CORE_WORKER_RPC_HANDLER(DirectActorCallArgWaitComplete, 100) +#define RAY_CORE_WORKER_RPC_HANDLERS \ + RAY_CORE_WORKER_RPC_HANDLER(AssignTask, 5) \ + RAY_CORE_WORKER_RPC_HANDLER(PushTask, 9999) \ + RAY_CORE_WORKER_RPC_HANDLER(DirectActorCallArgWaitComplete, 100) \ + RAY_CORE_WORKER_RPC_HANDLER(GetObjectStatus, 100) namespace ray { @@ -104,7 +106,7 @@ class CoreWorker { /// /// \param[in] object_id The object ID to increase the reference count for. void AddObjectIDReference(const ObjectID &object_id) { - reference_counter_->AddReference(object_id); + reference_counter_->AddLocalReference(object_id); } /// Decrease the reference count for this object ID. @@ -112,20 +114,46 @@ class CoreWorker { /// \param[in] object_id The object ID to decrease the reference count for. void RemoveObjectIDReference(const ObjectID &object_id) { std::vector deleted; - reference_counter_->RemoveReference(object_id, &deleted); + reference_counter_->RemoveLocalReference(object_id, &deleted); if (ref_counting_enabled_) { memory_store_->Delete(deleted); } } - /// Promote an object to plasma. If it already exists locally, it will be - /// put into the plasma store. If it doesn't yet exist, it will be spilled to - /// plasma once available. + /// Promote an object to plasma and get its owner information. This should be + /// called when serializing an object ID, and the returned information should + /// be stored with the serialized object ID. For plasma promotion, if the + /// object already exists locally, it will be put into the plasma store. If + /// it doesn't yet exist, it will be spilled to plasma once available. + /// + /// This can only be called on object IDs that we created via task + /// submission, ray.put, or object IDs that we deserialized. It cannot be + /// called on object IDs that were created randomly, e.g., + /// ObjectID::FromRandom. /// /// Postcondition: Get(object_id.WithPlasmaTransportType()) is valid. /// - /// \param[in] object_id The object ID to promote to plasma. - void PromoteObjectToPlasma(const ObjectID &object_id); + /// \param[in] object_id The object ID to serialize. + /// \param[out] owner_id The ID of the object's owner. This should be + /// appended to the serialized object ID. + /// \param[out] owner_address The address of the object's owner. This should + /// be appended to the serialized object ID. + void PromoteToPlasmaAndGetOwnershipInfo(const ObjectID &object_id, TaskID *owner_id, + rpc::Address *owner_address); + + /// Add a reference to an ObjectID that was deserialized by the language + /// frontend. This will also start the process to resolve the future. + /// Specifically, we will periodically contact the owner, until we learn that + /// the object has been created or the owner is no longer reachable. This + /// will then unblock any Gets or submissions of tasks dependent on the + /// object. + /// + /// \param[in] object_id The object ID to deserialize. + /// \param[out] owner_id The ID of the object's owner. + /// \param[out] owner_address The address of the object's owner. + void RegisterOwnershipInfoAndResolveFuture(const ObjectID &object_id, + const TaskID &owner_id, + const rpc::Address &owner_address); /// /// Public methods related to storing and retrieving objects. @@ -354,6 +382,11 @@ class CoreWorker { rpc::DirectActorCallArgWaitCompleteReply *reply, rpc::SendReplyCallback send_reply_callback); + /// Implements gRPC server handler. + void HandleGetObjectStatus(const rpc::GetObjectStatusRequest &request, + rpc::GetObjectStatusReply *reply, + rpc::SendReplyCallback send_reply_callback); + /// /// Public methods related to async actor call. This should only be used when /// the actor is (1) direct actor and (2) using asyncio mode. @@ -509,6 +542,8 @@ class CoreWorker { /// Plasma store interface. std::shared_ptr plasma_store_provider_; + std::unique_ptr future_resolver_; + /// /// Fields related to task submission. /// diff --git a/src/ray/core_worker/future_resolver.cc b/src/ray/core_worker/future_resolver.cc new file mode 100644 index 000000000..2ce70f810 --- /dev/null +++ b/src/ray/core_worker/future_resolver.cc @@ -0,0 +1,54 @@ +#include "ray/core_worker/future_resolver.h" + +namespace ray { + +void FutureResolver::ResolveFutureAsync(const ObjectID &object_id, const TaskID &owner_id, + const rpc::Address &owner_address) { + RAY_CHECK(object_id.IsDirectCallType()); + absl::MutexLock lock(&mu_); + auto it = owner_clients_.find(owner_id); + if (it == owner_clients_.end()) { + auto client = std::shared_ptr( + client_factory_({owner_address.ip_address(), owner_address.port()})); + owner_clients_.emplace(owner_id, std::move(client)); + } + + // This timer will get deallocated once the future has been resolved. + auto timer = std::make_shared(io_service_); + AttemptFutureResolution(object_id, owner_id, std::move(timer)); +} + +void FutureResolver::AttemptFutureResolution( + const ObjectID &object_id, const TaskID &owner_id, + std::shared_ptr timer) { + auto &owner_client = owner_clients_[owner_id]; + rpc::GetObjectStatusRequest request; + request.set_object_id(object_id.Binary()); + request.set_owner_id(owner_id.Binary()); + auto status = owner_client->GetObjectStatus( + request, [this, object_id, owner_id, timer]( + const Status &status, const rpc::GetObjectStatusReply &reply) { + if (!status.ok() || reply.status() != rpc::GetObjectStatusReply::PENDING) { + // Either the owner is gone or the owner replied that the object has + // been created. In both cases, we can now try to fetch the object via + // plasma. + RAY_CHECK_OK(in_memory_store_->Put(RayObject(rpc::ErrorType::OBJECT_IN_PLASMA), + object_id)); + } else { + // Try again later. + timer->expires_from_now( + boost::posix_time::milliseconds(wait_future_resolution_milliseconds_)); + timer->async_wait( + [this, object_id, owner_id, timer](const boost::system::error_code &error) { + absl::MutexLock lock(&mu_); + AttemptFutureResolution(object_id, owner_id, std::move(timer)); + }); + } + }); + if (!status.ok()) { + RAY_CHECK_OK( + in_memory_store_->Put(RayObject(rpc::ErrorType::OBJECT_IN_PLASMA), object_id)); + } +} + +} // namespace ray diff --git a/src/ray/core_worker/future_resolver.h b/src/ray/core_worker/future_resolver.h new file mode 100644 index 000000000..c67c4d88e --- /dev/null +++ b/src/ray/core_worker/future_resolver.h @@ -0,0 +1,71 @@ +#ifndef RAY_CORE_WORKER_FUTURE_RESOLVER_H +#define RAY_CORE_WORKER_FUTURE_RESOLVER_H + +#include + +#include "ray/common/id.h" +#include "ray/core_worker/store_provider/memory_store/memory_store.h" +#include "ray/protobuf/core_worker.pb.h" +#include "ray/rpc/worker/core_worker_client.h" + +namespace ray { + +/// Max time between requests to the owner to check whether the object is still +/// being computed. +const int kWaitObjectEvictionMilliseconds = 100; + +// Resolve values for futures that were given to us before the value +// was available. This class is thread-safe. +class FutureResolver { + public: + FutureResolver( + std::shared_ptr store, rpc::ClientFactoryFn client_factory, + boost::asio::io_service &io_service, + int wait_future_resolution_milliseconds = kWaitObjectEvictionMilliseconds) + : in_memory_store_(store), + client_factory_(client_factory), + io_service_(io_service), + wait_future_resolution_milliseconds_(wait_future_resolution_milliseconds) {} + + /// Resolve the value for a future. This will periodically contact the given + /// owner until the owner dies or the owner has finished creating the object. + /// In either case, this will put an OBJECT_IN_PLASMA error as the future's + /// value. + /// + /// \param[in] object_id The ID of the future to resolve. + /// \param[in] owner_id The ID of the task or actor that owns the future. + /// \param[in] owner_address The address of the task or actor that owns the + /// future. + void ResolveFutureAsync(const ObjectID &object_id, const TaskID &owner_id, + const rpc::Address &owner_address); + + private: + // Attempt to contact the owner to ask about the future's current status. + void AttemptFutureResolution(const ObjectID &object_id, const TaskID &owner_id, + std::shared_ptr timer) + EXCLUSIVE_LOCKS_REQUIRED(mu_); + + /// Used to set timers. + boost::asio::io_service &io_service_; + + /// Used to store values of resolved futures. + std::shared_ptr in_memory_store_; + + /// Factory for producing new core worker clients. + const rpc::ClientFactoryFn client_factory_; + + /// The amount of time to wait between requests to a future's owner to get + /// the object's current status. + const int wait_future_resolution_milliseconds_; + + /// Protects against concurrent access to internal state. + absl::Mutex mu_; + + /// Cache of gRPC clients to the objects' owners. + absl::flat_hash_map> + owner_clients_ GUARDED_BY(mu_); +}; + +} // namespace ray + +#endif // RAY_CORE_WORKER_FUTURE_RESOLVER_H diff --git a/src/ray/core_worker/reference_count.cc b/src/ray/core_worker/reference_count.cc index 6e8b4fcd0..bae2f3d4d 100644 --- a/src/ray/core_worker/reference_count.cc +++ b/src/ray/core_worker/reference_count.cc @@ -2,42 +2,48 @@ namespace ray { -void ReferenceCounter::AddReference(const ObjectID &object_id) { +void ReferenceCounter::AddBorrowedObject(const ObjectID &object_id, + const TaskID &owner_id, + const rpc::Address &owner_address) { absl::MutexLock lock(&mutex_); - AddReferenceInternal(object_id); + RAY_CHECK( + object_id_refs_.emplace(object_id, Reference(owner_id, owner_address)).second); } -void ReferenceCounter::AddReferenceInternal(const ObjectID &object_id) { - auto entry = object_id_refs_.find(object_id); - if (entry == object_id_refs_.end()) { - object_id_refs_[object_id] = std::make_pair(1, nullptr); - } else { - entry->second.first++; - } -} - -void ReferenceCounter::SetDependencies( - const ObjectID &object_id, std::shared_ptr> dependencies) { +void ReferenceCounter::AddOwnedObject( + const ObjectID &object_id, const TaskID &owner_id, const rpc::Address &owner_address, + std::shared_ptr> dependencies) { absl::MutexLock lock(&mutex_); - auto entry = object_id_refs_.find(object_id); - if (entry == object_id_refs_.end()) { - // If the entry doesn't exist, we initialize the direct reference count to zero - // because this corresponds to a submitted task whose return ObjectID will be created - // in the frontend language, incrementing the reference count. - object_id_refs_[object_id] = std::make_pair(0, dependencies); - } else { - RAY_CHECK(!entry->second.second); - entry->second.second = dependencies; - } - for (const ObjectID &dependency_id : *dependencies) { - AddReferenceInternal(dependency_id); + AddLocalReferenceInternal(dependency_id); } + + RAY_CHECK(object_id_refs_.count(object_id) == 0) + << "Cannot create an object that already exists. ObjectID: " << object_id; + // If the entry doesn't exist, we initialize the direct reference count to zero + // because this corresponds to a submitted task whose return ObjectID will be created + // in the frontend language, incrementing the reference count. + object_id_refs_.emplace(object_id, Reference(owner_id, owner_address, dependencies)); } -void ReferenceCounter::RemoveReference(const ObjectID &object_id, - std::vector *deleted) { +void ReferenceCounter::AddLocalReferenceInternal(const ObjectID &object_id) { + auto entry = object_id_refs_.find(object_id); + if (entry == object_id_refs_.end()) { + // TODO: Once ref counting is implemented, we should always know how the + // ObjectID was created, so there should always ben an entry. + entry = object_id_refs_.emplace(object_id, Reference()).first; + } + entry->second.local_ref_count++; +} + +void ReferenceCounter::AddLocalReference(const ObjectID &object_id) { + absl::MutexLock lock(&mutex_); + AddLocalReferenceInternal(object_id); +} + +void ReferenceCounter::RemoveLocalReference(const ObjectID &object_id, + std::vector *deleted) { absl::MutexLock lock(&mutex_); RemoveReferenceRecursive(object_id, deleted); } @@ -50,19 +56,36 @@ void ReferenceCounter::RemoveReferenceRecursive(const ObjectID &object_id, << object_id; return; } - if (--entry->second.first == 0) { + if (--entry->second.local_ref_count == 0) { // If the reference count reached 0, decrease the reference count for each dependency. - if (entry->second.second) { - for (const ObjectID &pending_task_object_id : *entry->second.second) { + if (entry->second.dependencies) { + for (const ObjectID &pending_task_object_id : *entry->second.dependencies) { RemoveReferenceRecursive(pending_task_object_id, deleted); } } - object_id_refs_.erase(object_id); + object_id_refs_.erase(entry); deleted->push_back(object_id); } } -bool ReferenceCounter::HasReference(const ObjectID &object_id) { +bool ReferenceCounter::GetOwner(const ObjectID &object_id, TaskID *owner_id, + rpc::Address *owner_address) const { + absl::MutexLock lock(&mutex_); + auto it = object_id_refs_.find(object_id); + if (it == object_id_refs_.end()) { + return false; + } + + if (it->second.owner.has_value()) { + *owner_id = it->second.owner.value().first; + *owner_address = it->second.owner.value().second; + return true; + } else { + return false; + } +} + +bool ReferenceCounter::HasReference(const ObjectID &object_id) const { absl::MutexLock lock(&mutex_); return object_id_refs_.find(object_id) != object_id_refs_.end(); } @@ -93,12 +116,12 @@ void ReferenceCounter::LogDebugString() const { for (const auto &entry : object_id_refs_) { RAY_LOG(DEBUG) << "\t" << entry.first.Hex(); - RAY_LOG(DEBUG) << "\t\treference count: " << entry.second.first; + RAY_LOG(DEBUG) << "\t\treference count: " << entry.second.local_ref_count; RAY_LOG(DEBUG) << "\t\tdependencies: "; - if (!entry.second.second) { + if (!entry.second.dependencies) { RAY_LOG(DEBUG) << "\t\t\tNULL"; } else { - for (const ObjectID &pending_task_object_id : *entry.second.second) { + for (const ObjectID &pending_task_object_id : *entry.second.dependencies) { RAY_LOG(DEBUG) << "\t\t\t" << pending_task_object_id.Hex(); } } diff --git a/src/ray/core_worker/reference_count.h b/src/ray/core_worker/reference_count.h index 96e5ff759..13f82560e 100644 --- a/src/ray/core_worker/reference_count.h +++ b/src/ray/core_worker/reference_count.h @@ -6,6 +6,7 @@ #include "absl/synchronization/mutex.h" #include "ray/common/id.h" +#include "ray/protobuf/common.pb.h" #include "ray/util/logging.h" namespace ray { @@ -19,30 +20,57 @@ class ReferenceCounter { ~ReferenceCounter() {} /// Increase the reference count for the ObjectID by one. If there is no - /// entry for the ObjectID, one will be created with no dependencies. - void AddReference(const ObjectID &object_id) LOCKS_EXCLUDED(mutex_); + /// entry for the ObjectID, one will be created. The object ID will not have + /// any owner information, since we don't know how it was created. + /// + /// \param[in] object_id The object to to increment the count for. + void AddLocalReference(const ObjectID &object_id) LOCKS_EXCLUDED(mutex_); /// Decrease the reference count for the ObjectID by one. If the reference count reaches /// zero, it will be erased from the map and the reference count for all of its /// dependencies will be decreased be one. /// /// \param[in] object_id The object to to decrement the count for. - /// \param[in] deleted List to store objects that hit zero ref count. - void RemoveReference(const ObjectID &object_id, std::vector *deleted) + /// \param[out] deleted List to store objects that hit zero ref count. + void RemoveLocalReference(const ObjectID &object_id, std::vector *deleted) LOCKS_EXCLUDED(mutex_); - /// Set the dependencies for the ObjectID. Dependencies for each ObjectID must be - /// set at most once. The direct reference count for the ObjectID is set to zero and the - /// reference count for each dependency is incremented. - void SetDependencies(const ObjectID &object_id, - std::shared_ptr> dependencies) + /// Add an object that we own. The object may depend on other objects. + /// Dependencies for each ObjectID must be set at most once. The direct + /// reference count for the ObjectID is set to zero and the reference count + /// for each dependency is incremented. + /// + /// TODO(swang): We could avoid copying the owner_id and owner_address since + /// we are the owner, but it is easier to store a copy for now, since the + /// owner ID will change for workers executing normal tasks and it is + /// possible to have leftover references after a task has finished. + /// + /// \param[in] object_id The ID of the object that we own. + /// \param[in] owner_id The ID of the object's owner. + /// \param[in] owner_address The address of the object's owner. + /// \param[in] dependencies The objects that the object depends on. + void AddOwnedObject(const ObjectID &object_id, const TaskID &owner_id, + const rpc::Address &owner_address, + std::shared_ptr> dependencies) LOCKS_EXCLUDED(mutex_); + /// Add an object that we are borrowing. + /// + /// \param[in] object_id The ID of the object that we are borrowing. + /// \param[in] owner_id The ID of the owner of the object. This is either the + /// task ID (for non-actors) or the actor ID of the owner. + /// \param[in] owner_address The owner's address. + void AddBorrowedObject(const ObjectID &object_id, const TaskID &owner_id, + const rpc::Address &owner_address) LOCKS_EXCLUDED(mutex_); + + bool GetOwner(const ObjectID &object_id, TaskID *owner_id, + rpc::Address *owner_address) const LOCKS_EXCLUDED(mutex_); + /// Returns the total number of ObjectIDs currently in scope. size_t NumObjectIDsInScope() const LOCKS_EXCLUDED(mutex_); /// Returns whether this object has an active reference. - bool HasReference(const ObjectID &object_id) LOCKS_EXCLUDED(mutex_); + bool HasReference(const ObjectID &object_id) const LOCKS_EXCLUDED(mutex_); /// Returns a set of all ObjectIDs currently in scope (i.e., nonzero reference count). std::unordered_set GetAllInScopeObjectIDs() const LOCKS_EXCLUDED(mutex_); @@ -51,9 +79,41 @@ class ReferenceCounter { void LogDebugString() const LOCKS_EXCLUDED(mutex_); private: + /// Metadata for an ObjectID reference in the language frontend. + struct Reference { + /// Constructor for a reference whose origin is unknown. + Reference() : owned_by_us(false) {} + /// Constructor for a reference that we created. + Reference(const TaskID &owner_id, const rpc::Address &owner_address, + std::shared_ptr> deps) + : dependencies(std::move(deps)), + owned_by_us(true), + owner({owner_id, owner_address}) {} + /// Constructor for a reference that was given to us. + Reference(const TaskID &owner_id, const rpc::Address &owner_address) + : owned_by_us(false), owner({owner_id, owner_address}) {} + /// The local ref count for the ObjectID in the language frontend. + size_t local_ref_count = 0; + /// The objects that this object depends on. Tracked only by the owner of + /// the object. Dependencies are stored as shared_ptrs because the same set + /// of dependencies can be shared among multiple entries. For example, when + /// a task has multiple return values, the entry for each return ObjectID + /// depends on all task dependencies. + std::shared_ptr> dependencies; + /// Whether we own the object. If we own the object, then we are + /// responsible for tracking the state of the task that creates the object + /// (see task_manager.h). + bool owned_by_us; + /// The object's owner, if we know it. This has no value if the object is + /// if we do not know the object's owner (because distributed ref counting + /// is not yet implemented). + const absl::optional> owner; + }; + /// Helper function with the same semantics as AddReference to allow adding a reference /// while already holding mutex_. - void AddReferenceInternal(const ObjectID &object_id) EXCLUSIVE_LOCKS_REQUIRED(mutex_); + void AddLocalReferenceInternal(const ObjectID &object_id) + EXCLUSIVE_LOCKS_REQUIRED(mutex_); /// Recursive helper function for decreasing reference counts. Will recursively call /// itself on any dependencies whose reference count reaches zero as a result of @@ -67,12 +127,8 @@ class ReferenceCounter { /// Protects access to the reference counting state. mutable absl::Mutex mutex_; - /// Holds all direct reference counts and dependency information for tracked ObjectIDs. - /// Dependencies are stored as shared_ptrs because the same set of dependencies can be - /// shared among multiple entries. For example, when a task has multiple return values, - /// the entry for each return ObjectID depends on all task dependencies. - absl::flat_hash_map>>> - object_id_refs_ GUARDED_BY(mutex_); + /// Holds all reference counts and dependency information for tracked ObjectIDs. + absl::flat_hash_map object_id_refs_ GUARDED_BY(mutex_); }; } // namespace ray diff --git a/src/ray/core_worker/reference_count_test.cc b/src/ray/core_worker/reference_count_test.cc index 004d74efb..288cfdc62 100644 --- a/src/ray/core_worker/reference_count_test.cc +++ b/src/ray/core_worker/reference_count_test.cc @@ -20,19 +20,19 @@ class ReferenceCountTest : public ::testing::Test { TEST_F(ReferenceCountTest, TestBasic) { std::vector out; ObjectID id = ObjectID::FromRandom(); - rc->AddReference(id); + rc->AddLocalReference(id); ASSERT_EQ(rc->NumObjectIDsInScope(), 1); - rc->AddReference(id); + rc->AddLocalReference(id); ASSERT_EQ(rc->NumObjectIDsInScope(), 1); - rc->AddReference(id); + rc->AddLocalReference(id); ASSERT_EQ(rc->NumObjectIDsInScope(), 1); - rc->RemoveReference(id, &out); + rc->RemoveLocalReference(id, &out); ASSERT_EQ(rc->NumObjectIDsInScope(), 1); ASSERT_EQ(out.size(), 0); - rc->RemoveReference(id, &out); + rc->RemoveLocalReference(id, &out); ASSERT_EQ(rc->NumObjectIDsInScope(), 1); ASSERT_EQ(out.size(), 0); - rc->RemoveReference(id, &out); + rc->RemoveLocalReference(id, &out); ASSERT_EQ(rc->NumObjectIDsInScope(), 0); ASSERT_EQ(out.size(), 1); } @@ -49,21 +49,21 @@ TEST_F(ReferenceCountTest, TestDependencies) { std::shared_ptr> deps = std::make_shared>(); deps->push_back(id2); deps->push_back(id3); - rc->SetDependencies(id1, deps); + rc->AddOwnedObject(id1, TaskID::Nil(), rpc::Address(), deps); - rc->AddReference(id1); - rc->AddReference(id1); - rc->AddReference(id3); + rc->AddLocalReference(id1); + rc->AddLocalReference(id1); + rc->AddLocalReference(id3); ASSERT_EQ(rc->NumObjectIDsInScope(), 3); - rc->RemoveReference(id1, &out); + rc->RemoveLocalReference(id1, &out); ASSERT_EQ(rc->NumObjectIDsInScope(), 3); ASSERT_EQ(out.size(), 0); - rc->RemoveReference(id1, &out); + rc->RemoveLocalReference(id1, &out); ASSERT_EQ(rc->NumObjectIDsInScope(), 1); ASSERT_EQ(out.size(), 2); - rc->RemoveReference(id3, &out); + rc->RemoveLocalReference(id3, &out); ASSERT_EQ(rc->NumObjectIDsInScope(), 0); ASSERT_EQ(out.size(), 3); } @@ -82,22 +82,22 @@ TEST_F(ReferenceCountTest, TestSharedDependencies) { std::shared_ptr> deps = std::make_shared>(); deps->push_back(id3); deps->push_back(id4); - rc->SetDependencies(id1, deps); - rc->SetDependencies(id2, deps); + rc->AddOwnedObject(id1, TaskID::Nil(), rpc::Address(), deps); + rc->AddOwnedObject(id2, TaskID::Nil(), rpc::Address(), deps); - rc->AddReference(id1); - rc->AddReference(id2); - rc->AddReference(id4); + rc->AddLocalReference(id1); + rc->AddLocalReference(id2); + rc->AddLocalReference(id4); ASSERT_EQ(rc->NumObjectIDsInScope(), 4); - rc->RemoveReference(id1, &out); + rc->RemoveLocalReference(id1, &out); ASSERT_EQ(rc->NumObjectIDsInScope(), 3); ASSERT_EQ(out.size(), 1); - rc->RemoveReference(id2, &out); + rc->RemoveLocalReference(id2, &out); ASSERT_EQ(rc->NumObjectIDsInScope(), 1); ASSERT_EQ(out.size(), 3); - rc->RemoveReference(id4, &out); + rc->RemoveLocalReference(id4, &out); ASSERT_EQ(rc->NumObjectIDsInScope(), 0); ASSERT_EQ(out.size(), 4); } @@ -113,34 +113,65 @@ TEST_F(ReferenceCountTest, TestRecursiveDependencies) { ObjectID id3 = ObjectID::FromRandom(); ObjectID id4 = ObjectID::FromRandom(); - std::shared_ptr> deps1 = - std::make_shared>(); - deps1->push_back(id2); - rc->SetDependencies(id1, deps1); - std::shared_ptr> deps2 = std::make_shared>(); deps2->push_back(id3); deps2->push_back(id4); - rc->SetDependencies(id2, deps2); + rc->AddOwnedObject(id2, TaskID::Nil(), rpc::Address(), deps2); - rc->AddReference(id1); - rc->AddReference(id2); - rc->AddReference(id4); + std::shared_ptr> deps1 = + std::make_shared>(); + deps1->push_back(id2); + rc->AddOwnedObject(id1, TaskID::Nil(), rpc::Address(), deps1); + + rc->AddLocalReference(id1); + rc->AddLocalReference(id2); + rc->AddLocalReference(id4); ASSERT_EQ(rc->NumObjectIDsInScope(), 4); - rc->RemoveReference(id2, &out); + rc->RemoveLocalReference(id2, &out); ASSERT_EQ(rc->NumObjectIDsInScope(), 4); ASSERT_EQ(out.size(), 0); - rc->RemoveReference(id1, &out); + rc->RemoveLocalReference(id1, &out); ASSERT_EQ(rc->NumObjectIDsInScope(), 1); ASSERT_EQ(out.size(), 3); - rc->RemoveReference(id4, &out); + rc->RemoveLocalReference(id4, &out); ASSERT_EQ(rc->NumObjectIDsInScope(), 0); ASSERT_EQ(out.size(), 4); } +// Tests that we can get the owner address correctly for objects that we own, +// objects that we borrowed via a serialized object ID, and objects whose +// origin we do not know. +TEST_F(ReferenceCountTest, TestOwnerAddress) { + auto object_id = ObjectID::FromRandom(); + TaskID task_id = TaskID::ForFakeTask(); + rpc::Address address; + address.set_ip_address("1234"); + auto deps = std::make_shared>(); + rc->AddOwnedObject(object_id, task_id, address, deps); + + TaskID added_id; + rpc::Address added_address; + ASSERT_TRUE(rc->GetOwner(object_id, &added_id, &added_address)); + ASSERT_EQ(task_id, added_id); + ASSERT_EQ(address.ip_address(), added_address.ip_address()); + + auto object_id2 = ObjectID::FromRandom(); + task_id = TaskID::ForFakeTask(); + address.set_ip_address("5678"); + rc->AddOwnedObject(object_id2, task_id, address, deps); + ASSERT_TRUE(rc->GetOwner(object_id2, &added_id, &added_address)); + ASSERT_EQ(task_id, added_id); + ASSERT_EQ(address.ip_address(), added_address.ip_address()); + + auto object_id3 = ObjectID::FromRandom(); + ASSERT_FALSE(rc->GetOwner(object_id3, &added_id, &added_address)); + rc->AddLocalReference(object_id3); + ASSERT_FALSE(rc->GetOwner(object_id3, &added_id, &added_address)); +} + // Tests that the ref counts are properly integrated into the local // object memory store. TEST(MemoryStoreIntegrationTest, TestSimple) { @@ -157,7 +188,7 @@ TEST(MemoryStoreIntegrationTest, TestSimple) { ASSERT_EQ(store.Size(), 0); // Tests ref counting overrides remove after get option. - rc->AddReference(id1); + rc->AddLocalReference(id1); RAY_CHECK_OK(store.Put(buffer, id1)); ASSERT_EQ(store.Size(), 1); std::vector> results; diff --git a/src/ray/core_worker/task_manager.cc b/src/ray/core_worker/task_manager.cc index b7df83802..7e3663f49 100644 --- a/src/ray/core_worker/task_manager.cc +++ b/src/ray/core_worker/task_manager.cc @@ -25,12 +25,8 @@ void TaskManager::CompletePendingTask(const TaskID &task_id, if (return_object.in_plasma()) { // Mark it as in plasma with a dummy object. - std::string meta = - std::to_string(static_cast(rpc::ErrorType::OBJECT_IN_PLASMA)); - auto metadata = - const_cast(reinterpret_cast(meta.data())); - auto meta_buffer = std::make_shared(metadata, meta.size()); - RAY_CHECK_OK(in_memory_store_->Put(RayObject(nullptr, meta_buffer), object_id)); + RAY_CHECK_OK( + in_memory_store_->Put(RayObject(rpc::ErrorType::OBJECT_IN_PLASMA), object_id)); } else { std::shared_ptr data_buffer; if (return_object.data().size() > 0) { @@ -70,10 +66,7 @@ void TaskManager::FailPendingTask(const TaskID &task_id, rpc::ErrorType error_ty const auto object_id = ObjectID::ForTaskReturn( task_id, /*index=*/i + 1, /*transport_type=*/static_cast(TaskTransportType::DIRECT)); - std::string meta = std::to_string(static_cast(error_type)); - auto metadata = const_cast(reinterpret_cast(meta.data())); - auto meta_buffer = std::make_shared(metadata, meta.size()); - RAY_CHECK_OK(in_memory_store_->Put(RayObject(nullptr, meta_buffer), object_id)); + RAY_CHECK_OK(in_memory_store_->Put(RayObject(error_type), object_id)); } } diff --git a/src/ray/protobuf/core_worker.proto b/src/ray/protobuf/core_worker.proto index 323624feb..ada27c5f8 100644 --- a/src/ray/protobuf/core_worker.proto +++ b/src/ray/protobuf/core_worker.proto @@ -87,6 +87,20 @@ message DirectActorCallArgWaitCompleteRequest { message DirectActorCallArgWaitCompleteReply { } +message GetObjectStatusRequest { + bytes owner_id = 1; + bytes object_id = 2; +} + +message GetObjectStatusReply { + enum ObjectStatus { + PENDING = 0; + CREATED = 1; + WRONG_OWNER = 2; + } + ObjectStatus status = 1; +} + service CoreWorkerService { // Push a task to a worker from the raylet. rpc AssignTask(AssignTaskRequest) returns (AssignTaskReply); @@ -95,4 +109,6 @@ service CoreWorkerService { // Reply from raylet that wait for direct actor call args has completed. rpc DirectActorCallArgWaitComplete(DirectActorCallArgWaitCompleteRequest) returns (DirectActorCallArgWaitCompleteReply); + // Ask the object's owner about the object's current status. + rpc GetObjectStatus(GetObjectStatusRequest) returns (GetObjectStatusReply); } diff --git a/src/ray/rpc/worker/core_worker_client.h b/src/ray/rpc/worker/core_worker_client.h index 40ad8f808..cf1a0e9d0 100644 --- a/src/ray/rpc/worker/core_worker_client.h +++ b/src/ray/rpc/worker/core_worker_client.h @@ -81,6 +81,13 @@ class CoreWorkerClientInterface { return Status::NotImplemented(""); } + /// Ask the owner of an object about the object's current status. + virtual ray::Status GetObjectStatus( + const GetObjectStatusRequest &request, + const ClientCallback &callback) { + return Status::NotImplemented(""); + } + virtual ~CoreWorkerClientInterface(){}; }; @@ -148,6 +155,14 @@ class CoreWorkerClient : public std::enable_shared_from_this, return call->GetStatus(); } + virtual ray::Status GetObjectStatus( + const GetObjectStatusRequest &request, + const ClientCallback &callback) { + auto call = client_call_manager_.CreateCall( + *stub_, &CoreWorkerService::Stub::PrepareAsyncGetObjectStatus, request, callback); + return call->GetStatus(); + } /// Send as many pending tasks as possible. This method is thread-safe. /// /// The client will guarantee no more than kMaxBytesInFlight bytes of RPCs are being