From eb307f93f8097becbde6d0c3b92cf33393248137 Mon Sep 17 00:00:00 2001 From: Zhijun Fu <37800433+zhijunfu@users.noreply.github.com> Date: Tue, 30 Jul 2019 17:47:17 +0800 Subject: [PATCH] Support direct actor call (#5183) --- BUILD.bazel | 55 +- src/ray/common/buffer.h | 9 +- src/ray/common/grpc_util.h | 24 +- src/ray/common/task/task_execution_spec.cc | 8 +- src/ray/common/task/task_spec.cc | 59 +- src/ray/common/task/task_spec.h | 10 +- src/ray/common/task/task_util.h | 43 +- src/ray/core_worker/common.h | 2 +- src/ray/core_worker/core_worker.cc | 27 +- src/ray/core_worker/core_worker.h | 27 +- src/ray/core_worker/core_worker_test.cc | 505 ----------- src/ray/core_worker/object_interface.h | 6 + .../store_provider/local_plasma_provider.cc | 7 +- .../store_provider/store_provider.h | 3 + src/ray/core_worker/task_execution.cc | 9 +- src/ray/core_worker/task_interface.cc | 49 +- src/ray/core_worker/task_interface.h | 32 +- src/ray/core_worker/test/core_worker_test.cc | 827 ++++++++++++++++++ src/ray/core_worker/{ => test}/mock_worker.cc | 11 +- .../transport/direct_actor_transport.cc | 222 +++++ .../transport/direct_actor_transport.h | 146 ++++ src/ray/core_worker/transport/transport.h | 4 + src/ray/gcs/gcs_client_interface.h | 2 +- src/ray/protobuf/core_worker.proto | 3 + src/ray/protobuf/direct_actor.proto | 30 + src/ray/protobuf/gcs.proto | 4 + src/ray/raylet/node_manager.cc | 23 +- src/ray/raylet/node_manager.h | 6 +- src/ray/rpc/worker/direct_actor_client.h | 58 ++ src/ray/rpc/worker/direct_actor_server.h | 68 ++ src/ray/util/test_util.h | 33 + 31 files changed, 1684 insertions(+), 628 deletions(-) delete mode 100644 src/ray/core_worker/core_worker_test.cc create mode 100644 src/ray/core_worker/test/core_worker_test.cc rename src/ray/core_worker/{ => test}/mock_worker.cc (87%) create mode 100644 src/ray/core_worker/transport/direct_actor_transport.cc create mode 100644 src/ray/core_worker/transport/direct_actor_transport.h create mode 100644 src/ray/protobuf/direct_actor.proto create mode 100644 src/ray/rpc/worker/direct_actor_client.h create mode 100644 src/ray/rpc/worker/direct_actor_server.h create mode 100644 src/ray/util/test_util.h diff --git a/BUILD.bazel b/BUILD.bazel index 64b380751..fbbae2ad9 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -101,6 +101,17 @@ cc_proto_library( deps = ["core_worker_proto"], ) +proto_library( + name = "direct_actor_proto", + srcs = ["src/ray/protobuf/direct_actor.proto"], + deps = [":common_proto"], +) + +cc_proto_library( + name = "direct_actor_cc_proto", + deps = ["direct_actor_proto"], +) + # === End of protobuf definitions === # === Begin of rpc definitions === @@ -207,7 +218,15 @@ cc_grpc_library( deps = [":worker_cc_proto"], ) -# Worker server and client. +# direct actor gRPC lib. +cc_grpc_library( + name = "direct_actor_cc_grpc", + srcs = [":direct_actor_proto"], + grpc_only = True, + deps = [":direct_actor_cc_proto"], +) + +# worker server and client. cc_library( name = "worker_rpc", hdrs = glob([ @@ -215,6 +234,7 @@ cc_library( ]), copts = COPTS, deps = [ + "direct_actor_cc_grpc", ":grpc_common_lib", ":ray_common", ":worker_cc_grpc", @@ -388,21 +408,36 @@ cc_library( # should only depend on `raylet_client`, instead of the whole `raylet_lib`. ":raylet_lib", ":worker_rpc", + ":gcs", + ], +) + +cc_library( + name = "mock_worker_lib", + srcs = ["src/ray/core_worker/test/mock_worker.cc"], + hdrs = glob([ + "src/ray/core_worker/test/*.h", + ]), + copts = COPTS, + deps = [ + ":core_worker_lib", ], ) cc_binary( name = "mock_worker", - srcs = ["src/ray/core_worker/mock_worker.cc"], copts = COPTS, deps = [ - ":core_worker_lib", + ":mock_worker_lib", ], ) -cc_binary( - name = "core_worker_test", - srcs = ["src/ray/core_worker/core_worker_test.cc"], +cc_library( + name = "core_worker_test_lib", + srcs = ["src/ray/core_worker/test/core_worker_test.cc"], + hdrs = glob([ + "src/ray/core_worker/test/*.h", + ]), copts = COPTS, deps = [ ":core_worker_lib", @@ -411,6 +446,14 @@ cc_binary( ], ) +cc_binary( + name = "core_worker_test", + copts = COPTS, + deps = [ + ":core_worker_test_lib", + ], +) + cc_test( name = "lineage_cache_test", srcs = ["src/ray/raylet/lineage_cache_test.cc"], diff --git a/src/ray/common/buffer.h b/src/ray/common/buffer.h index d14cb3fbb..994254ce4 100644 --- a/src/ray/common/buffer.h +++ b/src/ray/common/buffer.h @@ -23,7 +23,11 @@ class Buffer { virtual ~Buffer(){}; bool operator==(const Buffer &rhs) const { - return this->Data() == rhs.Data() && this->Size() == rhs.Size(); + if (this->Size() != rhs.Size()) { + return false; + } + + return this->Size() == 0 || memcmp(Data(), rhs.Data(), Size()) == 0; } }; @@ -54,7 +58,8 @@ class LocalMemoryBuffer : public Buffer { std::vector buffer_; }; -/// Represents a byte buffer for plasma object. +/// Represents a byte buffer for plasma object. This can be used to hold the +/// reference to a plasma object (via the underlying plasma::PlasmaBuffer). class PlasmaBuffer : public Buffer { public: PlasmaBuffer(std::shared_ptr buffer) : buffer_(buffer) {} diff --git a/src/ray/common/grpc_util.h b/src/ray/common/grpc_util.h index ec8c9ab3c..edcb58b5e 100644 --- a/src/ray/common/grpc_util.h +++ b/src/ray/common/grpc_util.h @@ -13,30 +13,40 @@ template class MessageWrapper { public: /// Construct an empty message wrapper. This should not be used directly. - MessageWrapper() {} + MessageWrapper() : message_(std::make_shared()) {} /// Construct from a protobuf message object. /// The input message will be **copied** into this object. /// /// \param message The protobuf message. - explicit MessageWrapper(const Message message) : message_(std::move(message)) {} + explicit MessageWrapper(const Message message) + : message_(std::make_shared(std::move(message))) {} + + /// Construct from a protobuf message shared_ptr. + /// + /// \param message The protobuf message. + explicit MessageWrapper(std::shared_ptr message) : message_(message) {} /// Construct from protobuf-serialized binary. /// /// \param serialized_binary Protobuf-serialized binary. - explicit MessageWrapper(const std::string &serialized_binary) { - message_.ParseFromString(serialized_binary); + explicit MessageWrapper(const std::string &serialized_binary) + : message_(std::make_shared()) { + message_->ParseFromString(serialized_binary); } + /// Get const reference of the protobuf message. + const Message &GetMessage() const { return *message_; } + /// Get reference of the protobuf message. - const Message &GetMessage() const { return message_; } + Message &GetMutableMessage() const { return *message_; } /// Serialize the message to a string. - const std::string Serialize() const { return message_.SerializeAsString(); } + const std::string Serialize() const { return message_->SerializeAsString(); } protected: /// The wrapped message. - Message message_; + std::shared_ptr message_; }; /// Helper function that converts a ray status to gRPC status. diff --git a/src/ray/common/task/task_execution_spec.cc b/src/ray/common/task/task_execution_spec.cc index 3f1af7551..3be3d9042 100644 --- a/src/ray/common/task/task_execution_spec.cc +++ b/src/ray/common/task/task_execution_spec.cc @@ -4,15 +4,17 @@ namespace ray { -size_t TaskExecutionSpecification::NumForwards() const { return message_.num_forwards(); } +size_t TaskExecutionSpecification::NumForwards() const { + return message_->num_forwards(); +} void TaskExecutionSpecification::IncrementNumForwards() { - message_.set_num_forwards(message_.num_forwards() + 1); + message_->set_num_forwards(message_->num_forwards() + 1); } std::string TaskExecutionSpecification::DebugString() const { std::ostringstream stream; - stream << "num_forwards=" << message_.num_forwards(); + stream << "num_forwards=" << message_->num_forwards(); return stream.str(); } diff --git a/src/ray/common/task/task_spec.cc b/src/ray/common/task/task_spec.cc index b0c01b418..b9183a9f4 100644 --- a/src/ray/common/task/task_spec.cc +++ b/src/ray/common/task/task_spec.cc @@ -6,9 +6,9 @@ namespace ray { void TaskSpecification::ComputeResources() { - auto required_resources = MapFromProtobuf(message_.required_resources()); + auto required_resources = MapFromProtobuf(message_->required_resources()); auto required_placement_resources = - MapFromProtobuf(message_.required_placement_resources()); + MapFromProtobuf(message_->required_placement_resources()); if (required_placement_resources.empty()) { required_placement_resources = required_resources; } @@ -18,24 +18,24 @@ void TaskSpecification::ComputeResources() { // Task specification getter methods. TaskID TaskSpecification::TaskId() const { - return TaskID::FromBinary(message_.task_id()); + return TaskID::FromBinary(message_->task_id()); } -JobID TaskSpecification::JobId() const { return JobID::FromBinary(message_.job_id()); } +JobID TaskSpecification::JobId() const { return JobID::FromBinary(message_->job_id()); } TaskID TaskSpecification::ParentTaskId() const { - return TaskID::FromBinary(message_.parent_task_id()); + return TaskID::FromBinary(message_->parent_task_id()); } -size_t TaskSpecification::ParentCounter() const { return message_.parent_counter(); } +size_t TaskSpecification::ParentCounter() const { return message_->parent_counter(); } std::vector TaskSpecification::FunctionDescriptor() const { - return VectorFromProtobuf(message_.function_descriptor()); + return VectorFromProtobuf(message_->function_descriptor()); } -size_t TaskSpecification::NumArgs() const { return message_.args_size(); } +size_t TaskSpecification::NumArgs() const { return message_->args_size(); } -size_t TaskSpecification::NumReturns() const { return message_.num_returns(); } +size_t TaskSpecification::NumReturns() const { return message_->num_returns(); } ObjectID TaskSpecification::ReturnId(size_t return_index) const { return ObjectID::ForTaskReturn(TaskId(), return_index + 1); @@ -46,19 +46,19 @@ bool TaskSpecification::ArgByRef(size_t arg_index) const { } size_t TaskSpecification::ArgIdCount(size_t arg_index) const { - return message_.args(arg_index).object_ids_size(); + return message_->args(arg_index).object_ids_size(); } ObjectID TaskSpecification::ArgId(size_t arg_index, size_t id_index) const { - return ObjectID::FromBinary(message_.args(arg_index).object_ids(id_index)); + return ObjectID::FromBinary(message_->args(arg_index).object_ids(id_index)); } const uint8_t *TaskSpecification::ArgVal(size_t arg_index) const { - return reinterpret_cast(message_.args(arg_index).data().data()); + return reinterpret_cast(message_->args(arg_index).data().data()); } size_t TaskSpecification::ArgValLength(size_t arg_index) const { - return message_.args(arg_index).data().size(); + return message_->args(arg_index).data().size(); } const ResourceSet TaskSpecification::GetRequiredResources() const { @@ -74,64 +74,65 @@ bool TaskSpecification::IsDriverTask() const { return FunctionDescriptor().empty(); } -Language TaskSpecification::GetLanguage() const { return message_.language(); } +Language TaskSpecification::GetLanguage() const { return message_->language(); } bool TaskSpecification::IsNormalTask() const { - return message_.type() == TaskType::NORMAL_TASK; + return message_->type() == TaskType::NORMAL_TASK; } bool TaskSpecification::IsActorCreationTask() const { - return message_.type() == TaskType::ACTOR_CREATION_TASK; + return message_->type() == TaskType::ACTOR_CREATION_TASK; } bool TaskSpecification::IsActorTask() const { - return message_.type() == TaskType::ACTOR_TASK; + return message_->type() == TaskType::ACTOR_TASK; } // === Below are getter methods specific to actor creation tasks. ActorID TaskSpecification::ActorCreationId() const { RAY_CHECK(IsActorCreationTask()); - return ActorID::FromBinary(message_.actor_creation_task_spec().actor_id()); + return ActorID::FromBinary(message_->actor_creation_task_spec().actor_id()); } uint64_t TaskSpecification::MaxActorReconstructions() const { RAY_CHECK(IsActorCreationTask()); - return message_.actor_creation_task_spec().max_actor_reconstructions(); + return message_->actor_creation_task_spec().max_actor_reconstructions(); } std::vector TaskSpecification::DynamicWorkerOptions() const { RAY_CHECK(IsActorCreationTask()); - return VectorFromProtobuf(message_.actor_creation_task_spec().dynamic_worker_options()); + return VectorFromProtobuf( + message_->actor_creation_task_spec().dynamic_worker_options()); } // === Below are getter methods specific to actor tasks. ActorID TaskSpecification::ActorId() const { RAY_CHECK(IsActorTask()); - return ActorID::FromBinary(message_.actor_task_spec().actor_id()); + return ActorID::FromBinary(message_->actor_task_spec().actor_id()); } ActorHandleID TaskSpecification::ActorHandleId() const { RAY_CHECK(IsActorTask()); - return ActorHandleID::FromBinary(message_.actor_task_spec().actor_handle_id()); + return ActorHandleID::FromBinary(message_->actor_task_spec().actor_handle_id()); } uint64_t TaskSpecification::ActorCounter() const { RAY_CHECK(IsActorTask()); - return message_.actor_task_spec().actor_counter(); + return message_->actor_task_spec().actor_counter(); } ObjectID TaskSpecification::ActorCreationDummyObjectId() const { RAY_CHECK(IsActorTask()); return ObjectID::FromBinary( - message_.actor_task_spec().actor_creation_dummy_object_id()); + message_->actor_task_spec().actor_creation_dummy_object_id()); } ObjectID TaskSpecification::PreviousActorTaskDummyObjectId() const { RAY_CHECK(IsActorTask()); return ObjectID::FromBinary( - message_.actor_task_spec().previous_actor_task_dummy_object_id()); + message_->actor_task_spec().previous_actor_task_dummy_object_id()); } ObjectID TaskSpecification::ActorDummyObject() const { @@ -142,17 +143,17 @@ ObjectID TaskSpecification::ActorDummyObject() const { std::vector TaskSpecification::NewActorHandles() const { RAY_CHECK(IsActorTask()); return IdVectorFromProtobuf( - message_.actor_task_spec().new_actor_handles()); + message_->actor_task_spec().new_actor_handles()); } std::string TaskSpecification::DebugString() const { std::ostringstream stream; - stream << "Type=" << TaskType_Name(message_.type()) - << ", Language=" << Language_Name(message_.language()) + stream << "Type=" << TaskType_Name(message_->type()) + << ", Language=" << Language_Name(message_->language()) << ", function_descriptor="; // Print function descriptor. - const auto list = VectorFromProtobuf(message_.function_descriptor()); + const auto list = VectorFromProtobuf(message_->function_descriptor()); // The 4th is the code hash which is binary bits. No need to output it. const size_t size = std::min(static_cast(3), list.size()); for (int i = 0; i < size; ++i) { diff --git a/src/ray/common/task/task_spec.h b/src/ray/common/task/task_spec.h index 353abb7b1..d1fae90ac 100644 --- a/src/ray/common/task/task_spec.h +++ b/src/ray/common/task/task_spec.h @@ -27,7 +27,15 @@ class TaskSpecification : public MessageWrapper { /// The input message will be **copied** into this object. /// /// \param message The protobuf message. - explicit TaskSpecification(rpc::TaskSpec message) : MessageWrapper(std::move(message)) { + explicit TaskSpecification(rpc::TaskSpec message) : MessageWrapper(message) { + ComputeResources(); + } + + /// Construct from a protobuf message shared_ptr. + /// + /// \param message The protobuf message. + explicit TaskSpecification(std::shared_ptr message) + : MessageWrapper(message) { ComputeResources(); } diff --git a/src/ray/common/task/task_util.h b/src/ray/common/task/task_util.h index d530b691e..c331d928c 100644 --- a/src/ray/common/task/task_util.h +++ b/src/ray/common/task/task_util.h @@ -9,11 +9,13 @@ namespace ray { /// Helper class for building a `TaskSpecification` object. class TaskSpecBuilder { public: + TaskSpecBuilder() : message_(std::make_shared()) {} + /// Build the `TaskSpecification` object. TaskSpecification Build() { return TaskSpecification(message_); } /// Get a reference to the internal protobuf message object. - const rpc::TaskSpec &GetMessage() const { return message_; } + const rpc::TaskSpec &GetMessage() const { return *message_; } /// Set the common attributes of the task spec. /// See `common.proto` for meaning of the arguments. @@ -25,19 +27,20 @@ class TaskSpecBuilder { uint64_t num_returns, const std::unordered_map &required_resources, const std::unordered_map &required_placement_resources) { - message_.set_type(TaskType::NORMAL_TASK); - message_.set_language(language); + message_->set_type(TaskType::NORMAL_TASK); + message_->set_language(language); for (const auto &fd : function_descriptor) { - message_.add_function_descriptor(fd); + message_->add_function_descriptor(fd); } - message_.set_job_id(job_id.Binary()); - message_.set_task_id(GenerateTaskId(job_id, parent_task_id, parent_counter).Binary()); - message_.set_parent_task_id(parent_task_id.Binary()); - message_.set_parent_counter(parent_counter); - message_.set_num_returns(num_returns); - message_.mutable_required_resources()->insert(required_resources.begin(), - required_resources.end()); - message_.mutable_required_placement_resources()->insert( + message_->set_job_id(job_id.Binary()); + message_->set_task_id( + GenerateTaskId(job_id, parent_task_id, parent_counter).Binary()); + message_->set_parent_task_id(parent_task_id.Binary()); + message_->set_parent_counter(parent_counter); + message_->set_num_returns(num_returns); + message_->mutable_required_resources()->insert(required_resources.begin(), + required_resources.end()); + message_->mutable_required_placement_resources()->insert( required_placement_resources.begin(), required_placement_resources.end()); return *this; } @@ -47,7 +50,7 @@ class TaskSpecBuilder { /// \param arg_id Id of the argument. /// \return Reference to the builder object itself. TaskSpecBuilder &AddByRefArg(const ObjectID &arg_id) { - message_.add_args()->add_object_ids(arg_id.Binary()); + message_->add_args()->add_object_ids(arg_id.Binary()); return *this; } @@ -56,7 +59,7 @@ class TaskSpecBuilder { /// \param data String object that contains the data. /// \return Reference to the builder object itself. TaskSpecBuilder &AddByValueArg(const std::string &data) { - message_.add_args()->set_data(data); + message_->add_args()->set_data(data); return *this; } @@ -66,7 +69,7 @@ class TaskSpecBuilder { /// \param size Size of the data. /// \return Reference to the builder object itself. TaskSpecBuilder &AddByValueArg(const void *data, size_t size) { - message_.add_args()->set_data(data, size); + message_->add_args()->set_data(data, size); return *this; } @@ -77,8 +80,8 @@ class TaskSpecBuilder { TaskSpecBuilder &SetActorCreationTaskSpec( const ActorID &actor_id, uint64_t max_reconstructions = 0, const std::vector &dynamic_worker_options = {}) { - message_.set_type(TaskType::ACTOR_CREATION_TASK); - auto actor_creation_spec = message_.mutable_actor_creation_task_spec(); + message_->set_type(TaskType::ACTOR_CREATION_TASK); + auto actor_creation_spec = message_->mutable_actor_creation_task_spec(); actor_creation_spec->set_actor_id(actor_id.Binary()); actor_creation_spec->set_max_actor_reconstructions(max_reconstructions); for (const auto &option : dynamic_worker_options) { @@ -96,8 +99,8 @@ class TaskSpecBuilder { const ObjectID &actor_creation_dummy_object_id, const ObjectID &previous_actor_task_dummy_object_id, uint64_t actor_counter, const std::vector &new_handle_ids = {}) { - message_.set_type(TaskType::ACTOR_TASK); - auto actor_spec = message_.mutable_actor_task_spec(); + message_->set_type(TaskType::ACTOR_TASK); + auto actor_spec = message_->mutable_actor_task_spec(); actor_spec->set_actor_id(actor_id.Binary()); actor_spec->set_actor_handle_id(actor_handle_id.Binary()); actor_spec->set_actor_creation_dummy_object_id( @@ -112,7 +115,7 @@ class TaskSpecBuilder { } private: - rpc::TaskSpec message_; + std::shared_ptr message_; }; } // namespace ray diff --git a/src/ray/core_worker/common.h b/src/ray/core_worker/common.h index 5f5e439a2..47b40a44b 100644 --- a/src/ray/core_worker/common.h +++ b/src/ray/core_worker/common.h @@ -76,7 +76,7 @@ struct TaskInfo { enum class StoreProviderType { LOCAL_PLASMA, PLASMA }; -enum class TaskTransportType { RAYLET }; +enum class TaskTransportType { RAYLET, DIRECT_ACTOR }; } // namespace ray diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc index 4a7dc4e59..71df4451b 100644 --- a/src/ray/core_worker/core_worker.cc +++ b/src/ray/core_worker/core_worker.cc @@ -6,20 +6,29 @@ namespace ray { CoreWorker::CoreWorker( const WorkerType worker_type, const Language language, const std::string &store_socket, const std::string &raylet_socket, - const JobID &job_id, + const JobID &job_id, const gcs::GcsClientOptions &gcs_options, const CoreWorkerTaskExecutionInterface::TaskExecutor &execution_callback) : worker_type_(worker_type), language_(language), raylet_socket_(raylet_socket), worker_context_(worker_type, job_id), - task_interface_(worker_context_, raylet_client_), - object_interface_(worker_context_, raylet_client_, store_socket) { + io_work_(io_service_) { + // Initialize gcs client + gcs_client_ = + std::unique_ptr(new gcs::RedisGcsClient(gcs_options)); + RAY_CHECK_OK(gcs_client_->Connect(io_service_)); + + object_interface_ = std::unique_ptr( + new CoreWorkerObjectInterface(worker_context_, raylet_client_, store_socket)); + task_interface_ = std::unique_ptr(new CoreWorkerTaskInterface( + worker_context_, raylet_client_, *object_interface_, io_service_, *gcs_client_)); + int rpc_server_port = 0; if (worker_type_ == WorkerType::WORKER) { RAY_CHECK(execution_callback != nullptr); task_execution_interface_ = std::unique_ptr( new CoreWorkerTaskExecutionInterface(worker_context_, raylet_client_, - object_interface_, execution_callback)); + *object_interface_, execution_callback)); rpc_server_port = task_execution_interface_->worker_server_.GetPort(); } // TODO(zhijunfu): currently RayletClient would crash in its constructor if it cannot @@ -30,6 +39,16 @@ CoreWorker::CoreWorker( raylet_socket_, WorkerID::FromBinary(worker_context_.GetWorkerID().Binary()), (worker_type_ == ray::WorkerType::WORKER), worker_context_.GetCurrentJobID(), language_, rpc_server_port)); + + io_thread_ = std::thread(&CoreWorker::StartIOService, this); } +CoreWorker::~CoreWorker() { + gcs_client_->Disconnect(); + io_service_.stop(); + io_thread_.join(); +} + +void CoreWorker::StartIOService() { io_service_.run(); } + } // namespace ray diff --git a/src/ray/core_worker/core_worker.h b/src/ray/core_worker/core_worker.h index 447a1f922..a5baa76d0 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/object_interface.h" #include "ray/core_worker/task_execution.h" #include "ray/core_worker/task_interface.h" +#include "ray/gcs/redis_gcs_client.h" #include "ray/rpc/raylet/raylet_client.h" namespace ray { @@ -26,9 +27,11 @@ class CoreWorker { /// NOTE(zhijunfu): the constructor would throw if a failure happens. CoreWorker(const WorkerType worker_type, const Language language, const std::string &store_socket, const std::string &raylet_socket, - const JobID &job_id, + const JobID &job_id, const gcs::GcsClientOptions &gcs_options, const CoreWorkerTaskExecutionInterface::TaskExecutor &execution_callback); + ~CoreWorker(); + /// Type of this worker. WorkerType GetWorkerType() const { return worker_type_; } @@ -37,11 +40,11 @@ class CoreWorker { /// Return the `CoreWorkerTaskInterface` that contains the methods related to task /// submisson. - CoreWorkerTaskInterface &Tasks() { return task_interface_; } + CoreWorkerTaskInterface &Tasks() { return *task_interface_; } /// Return the `CoreWorkerObjectInterface` that contains methods related to object /// store. - CoreWorkerObjectInterface &Objects() { return object_interface_; } + CoreWorkerObjectInterface &Objects() { return *object_interface_; } /// Return the `CoreWorkerTaskExecutionInterface` that contains methods related to /// task execution. @@ -51,6 +54,8 @@ class CoreWorker { } private: + void StartIOService(); + /// Type of this worker. const WorkerType worker_type_; @@ -63,14 +68,26 @@ class CoreWorker { /// Worker context. WorkerContext worker_context_; + /// event loop where the IO events are handled. e.g. async GCS operations. + boost::asio::io_service io_service_; + + /// keeps io_service_ alive. + boost::asio::io_service::work io_work_; + + /// The thread to handle IO events. + std::thread io_thread_; + /// Raylet client. std::unique_ptr raylet_client_; + /// GCS client. + std::unique_ptr gcs_client_; + /// The `CoreWorkerTaskInterface` instance. - CoreWorkerTaskInterface task_interface_; + std::unique_ptr task_interface_; /// The `CoreWorkerObjectInterface` instance. - CoreWorkerObjectInterface object_interface_; + std::unique_ptr object_interface_; /// The `CoreWorkerTaskExecutionInterface` instance. /// This is only available if it's not a driver. diff --git a/src/ray/core_worker/core_worker_test.cc b/src/ray/core_worker/core_worker_test.cc deleted file mode 100644 index c0ce4e9b9..000000000 --- a/src/ray/core_worker/core_worker_test.cc +++ /dev/null @@ -1,505 +0,0 @@ -#include -#include "gmock/gmock.h" -#include "gtest/gtest.h" - -#include "ray/common/buffer.h" -#include "ray/core_worker/context.h" -#include "ray/core_worker/core_worker.h" -#include "ray/rpc/raylet/raylet_client.h" - -#include -#include -#include - -#include "ray/thirdparty/hiredis/async.h" -#include "ray/thirdparty/hiredis/hiredis.h" - -namespace ray { - -std::string store_executable; -std::string raylet_executable; -std::string mock_worker_executable; - -ray::ObjectID RandomObjectID() { return ObjectID::FromRandom(); } - -static void flushall_redis(void) { - redisContext *context = redisConnect("127.0.0.1", 6379); - freeReplyObject(redisCommand(context, "FLUSHALL")); - freeReplyObject(redisCommand(context, "SET NumRedisShards 1")); - freeReplyObject(redisCommand(context, "LPUSH RedisShards 127.0.0.1:6380")); - redisFree(context); -} - -class CoreWorkerTest : public ::testing::Test { - public: - CoreWorkerTest(int num_nodes) { - // flush redis first. - flushall_redis(); - - RAY_CHECK(num_nodes >= 0); - if (num_nodes > 0) { - raylet_socket_names_.resize(num_nodes); - raylet_store_socket_names_.resize(num_nodes); - } - - // start plasma store. - for (auto &store_socket : raylet_store_socket_names_) { - store_socket = StartStore(); - } - - // start raylet on each node. Assign each node with different resources so that - // a task can be scheduled to the desired node. - for (int i = 0; i < num_nodes; i++) { - raylet_socket_names_[i] = - StartRaylet(raylet_store_socket_names_[i], "127.0.0.1", "127.0.0.1", - "\"CPU,4.0,resource" + std::to_string(i) + ",10\""); - } - } - - ~CoreWorkerTest() { - for (const auto &raylet_socket : raylet_socket_names_) { - StopRaylet(raylet_socket); - } - - for (const auto &store_socket : raylet_store_socket_names_) { - StopStore(store_socket); - } - } - - std::string StartStore() { - std::string store_socket_name = "/tmp/store" + RandomObjectID().Hex(); - std::string store_pid = store_socket_name + ".pid"; - std::string plasma_command = store_executable + " -m 10000000 -s " + - store_socket_name + - " 1> /dev/null 2> /dev/null & echo $! > " + store_pid; - RAY_LOG(DEBUG) << plasma_command; - RAY_CHECK(system(plasma_command.c_str()) == 0); - usleep(200 * 1000); - return store_socket_name; - } - - void StopStore(std::string store_socket_name) { - std::string store_pid = store_socket_name + ".pid"; - std::string kill_9 = "kill -9 `cat " + store_pid + "`"; - RAY_LOG(DEBUG) << kill_9; - ASSERT_TRUE(system(kill_9.c_str()) == 0); - ASSERT_TRUE(system(("rm -rf " + store_socket_name).c_str()) == 0); - ASSERT_TRUE(system(("rm -rf " + store_socket_name + ".pid").c_str()) == 0); - } - - std::string StartRaylet(std::string store_socket_name, std::string node_ip_address, - std::string redis_address, std::string resource) { - std::string raylet_socket_name = "/tmp/raylet" + RandomObjectID().Hex(); - std::string ray_start_cmd = raylet_executable; - ray_start_cmd.append(" --raylet_socket_name=" + raylet_socket_name) - .append(" --store_socket_name=" + store_socket_name) - .append(" --object_manager_port=0 --node_manager_port=0") - .append(" --node_ip_address=" + node_ip_address) - .append(" --redis_address=" + redis_address) - .append(" --redis_port=6379") - .append(" --num_initial_workers=1") - .append(" --maximum_startup_concurrency=10") - .append(" --static_resource_list=" + resource) - .append(" --python_worker_command=\"" + mock_worker_executable + " " + - store_socket_name + " " + raylet_socket_name + "\"") - .append(" & echo $! > " + raylet_socket_name + ".pid"); - - RAY_LOG(DEBUG) << "Ray Start command: " << ray_start_cmd; - RAY_CHECK(system(ray_start_cmd.c_str()) == 0); - usleep(200 * 1000); - return raylet_socket_name; - } - - void StopRaylet(std::string raylet_socket_name) { - std::string raylet_pid = raylet_socket_name + ".pid"; - std::string kill_9 = "kill -9 `cat " + raylet_pid + "`"; - RAY_LOG(DEBUG) << kill_9; - ASSERT_TRUE(system(kill_9.c_str()) == 0); - ASSERT_TRUE(system(("rm -rf " + raylet_socket_name).c_str()) == 0); - ASSERT_TRUE(system(("rm -rf " + raylet_socket_name + ".pid").c_str()) == 0); - } - - void SetUp() {} - - void TearDown() {} - - void TestNormalTask(const std::unordered_map &resources) { - CoreWorker driver(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[0], - raylet_socket_names_[0], JobID::FromInt(1), nullptr); - - // Test pass by value. - { - uint8_t array1[] = {1, 2, 3, 4, 5, 6, 7, 8}; - - auto buffer1 = std::make_shared(array1, sizeof(array1)); - - RayFunction func{Language::PYTHON, {}}; - std::vector args; - args.emplace_back(TaskArg::PassByValue(buffer1)); - - TaskOptions options; - - std::vector return_ids; - RAY_CHECK_OK(driver.Tasks().SubmitTask(func, args, options, &return_ids)); - - ASSERT_EQ(return_ids.size(), 1); - - std::vector> results; - RAY_CHECK_OK(driver.Objects().Get(return_ids, -1, &results)); - - ASSERT_EQ(results.size(), 1); - ASSERT_EQ(results[0]->GetData()->Size(), buffer1->Size()); - ASSERT_EQ(memcmp(results[0]->GetData()->Data(), buffer1->Data(), buffer1->Size()), - 0); - } - - // Test pass by reference. - { - uint8_t array1[] = {10, 11, 12, 13, 14, 15}; - auto buffer1 = std::make_shared(array1, sizeof(array1)); - - ObjectID object_id; - RAY_CHECK_OK(driver.Objects().Put(RayObject(buffer1, nullptr), &object_id)); - - std::vector args; - args.emplace_back(TaskArg::PassByReference(object_id)); - - RayFunction func{Language::PYTHON, {}}; - TaskOptions options; - - std::vector return_ids; - RAY_CHECK_OK(driver.Tasks().SubmitTask(func, args, options, &return_ids)); - - ASSERT_EQ(return_ids.size(), 1); - - std::vector> results; - RAY_CHECK_OK(driver.Objects().Get(return_ids, -1, &results)); - - ASSERT_EQ(results.size(), 1); - ASSERT_EQ(results[0]->GetData()->Size(), buffer1->Size()); - ASSERT_EQ(memcmp(results[0]->GetData()->Data(), buffer1->Data(), buffer1->Size()), - 0); - } - } - - void TestActorTask(const std::unordered_map &resources) { - CoreWorker driver(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[0], - raylet_socket_names_[0], JobID::FromInt(1), nullptr); - - std::unique_ptr actor_handle; - - // Test creating actor. - { - uint8_t array[] = {1, 2, 3}; - auto buffer = std::make_shared(array, sizeof(array)); - - RayFunction func{Language::PYTHON, {}}; - std::vector args; - args.emplace_back(TaskArg::PassByValue(buffer)); - - ActorCreationOptions actor_options{0, resources}; - - // Create an actor. - RAY_CHECK_OK(driver.Tasks().CreateActor(func, args, actor_options, &actor_handle)); - } - - // Test submitting a task for that actor. - { - uint8_t array1[] = {1, 2, 3, 4, 5, 6, 7, 8}; - uint8_t array2[] = {10, 11, 12, 13, 14, 15}; - - auto buffer1 = std::make_shared(array1, sizeof(array1)); - auto buffer2 = std::make_shared(array2, sizeof(array2)); - - ObjectID object_id; - RAY_CHECK_OK(driver.Objects().Put(RayObject(buffer1, nullptr), &object_id)); - - // Create arguments with PassByRef and PassByValue. - std::vector args; - args.emplace_back(TaskArg::PassByReference(object_id)); - args.emplace_back(TaskArg::PassByValue(buffer2)); - - TaskOptions options{1, resources}; - std::vector return_ids; - RayFunction func{Language::PYTHON, {}}; - RAY_CHECK_OK(driver.Tasks().SubmitActorTask(*actor_handle, func, args, options, - &return_ids)); - RAY_CHECK(return_ids.size() == 1); - - std::vector> results; - RAY_CHECK_OK(driver.Objects().Get(return_ids, -1, &results)); - - ASSERT_EQ(results.size(), 1); - ASSERT_EQ(results[0]->GetData()->Size(), buffer1->Size() + buffer2->Size()); - ASSERT_EQ(memcmp(results[0]->GetData()->Data(), buffer1->Data(), buffer1->Size()), - 0); - ASSERT_EQ(memcmp(results[0]->GetData()->Data() + buffer1->Size(), buffer2->Data(), - buffer2->Size()), - 0); - } - } - - protected: - std::vector raylet_socket_names_; - std::vector raylet_store_socket_names_; -}; - -class ZeroNodeTest : public CoreWorkerTest { - public: - ZeroNodeTest() : CoreWorkerTest(0) {} -}; - -class SingleNodeTest : public CoreWorkerTest { - public: - SingleNodeTest() : CoreWorkerTest(1) {} -}; - -class TwoNodeTest : public CoreWorkerTest { - public: - TwoNodeTest() : CoreWorkerTest(2) {} -}; - -TEST_F(ZeroNodeTest, TestTaskArg) { - // Test by-reference argument. - ObjectID id = ObjectID::FromRandom(); - TaskArg by_ref = TaskArg::PassByReference(id); - ASSERT_TRUE(by_ref.IsPassedByReference()); - ASSERT_EQ(by_ref.GetReference(), id); - // Test by-value argument. - std::shared_ptr buffer = - std::make_shared(static_cast(0), 0); - TaskArg by_value = TaskArg::PassByValue(buffer); - ASSERT_FALSE(by_value.IsPassedByReference()); - auto data = by_value.GetValue(); - ASSERT_TRUE(data != nullptr); - ASSERT_EQ(*data, *buffer); -} - -TEST_F(ZeroNodeTest, TestWorkerContext) { - auto job_id = JobID::JobID::FromInt(1); - - WorkerContext context(WorkerType::WORKER, job_id); - ASSERT_TRUE(context.GetCurrentTaskID().IsNil()); - ASSERT_EQ(context.GetNextTaskIndex(), 1); - ASSERT_EQ(context.GetNextTaskIndex(), 2); - ASSERT_EQ(context.GetNextPutIndex(), 1); - ASSERT_EQ(context.GetNextPutIndex(), 2); - - auto thread_func = [&context]() { - // Verify that task_index, put_index are thread-local. - ASSERT_TRUE(!context.GetCurrentTaskID().IsNil()); - ASSERT_EQ(context.GetNextTaskIndex(), 1); - ASSERT_EQ(context.GetNextPutIndex(), 1); - }; - - std::thread async_thread(thread_func); - async_thread.join(); - - // Verify that these fields are thread-local. - ASSERT_EQ(context.GetNextTaskIndex(), 3); - ASSERT_EQ(context.GetNextPutIndex(), 3); -} - -TEST_F(ZeroNodeTest, TestActorHandle) { - ActorHandle handle1(ActorID::FromRandom(), ActorHandleID::FromRandom(), Language::JAVA, - {"org.ray.exampleClass", "exampleMethod", "exampleSignature"}); - - auto forkedHandle1 = handle1.Fork(); - ASSERT_EQ(1, handle1.NumForks()); - ASSERT_EQ(handle1.ActorID(), forkedHandle1.ActorID()); - ASSERT_NE(handle1.ActorHandleID(), forkedHandle1.ActorHandleID()); - ASSERT_EQ(handle1.ActorLanguage(), forkedHandle1.ActorLanguage()); - ASSERT_EQ(handle1.ActorCreationTaskFunctionDescriptor(), - forkedHandle1.ActorCreationTaskFunctionDescriptor()); - ASSERT_EQ(handle1.ActorCursor(), forkedHandle1.ActorCursor()); - ASSERT_EQ(0, forkedHandle1.TaskCounter()); - ASSERT_EQ(0, forkedHandle1.NumForks()); - auto forkedHandle2 = handle1.Fork(); - ASSERT_EQ(2, handle1.NumForks()); - ASSERT_EQ(0, forkedHandle2.TaskCounter()); - ASSERT_EQ(0, forkedHandle2.NumForks()); - - std::string buffer; - handle1.Serialize(&buffer); - auto handle2 = ActorHandle::Deserialize(buffer); - ASSERT_EQ(handle1.ActorID(), handle2.ActorID()); - ASSERT_EQ(handle1.ActorHandleID(), handle2.ActorHandleID()); - ASSERT_EQ(handle1.ActorLanguage(), handle2.ActorLanguage()); - ASSERT_EQ(handle1.ActorCreationTaskFunctionDescriptor(), - handle2.ActorCreationTaskFunctionDescriptor()); - ASSERT_EQ(handle1.ActorCursor(), handle2.ActorCursor()); - ASSERT_EQ(handle1.TaskCounter(), handle2.TaskCounter()); - ASSERT_EQ(handle1.NumForks(), handle2.NumForks()); -} - -TEST_F(SingleNodeTest, TestObjectInterface) { - CoreWorker core_worker(WorkerType::DRIVER, Language::PYTHON, - raylet_store_socket_names_[0], raylet_socket_names_[0], - JobID::FromInt(1), nullptr); - - uint8_t array1[] = {1, 2, 3, 4, 5, 6, 7, 8}; - uint8_t array2[] = {10, 11, 12, 13, 14, 15}; - - std::vector buffers; - buffers.emplace_back(std::make_shared(array1, sizeof(array1)), - std::make_shared(array1, sizeof(array1) / 2)); - buffers.emplace_back(std::make_shared(array2, sizeof(array2)), - std::make_shared(array2, sizeof(array2) / 2)); - - std::vector ids(buffers.size()); - for (size_t i = 0; i < ids.size(); i++) { - RAY_CHECK_OK(core_worker.Objects().Put(buffers[i], &ids[i])); - } - - // Test Get(). - std::vector> results; - RAY_CHECK_OK(core_worker.Objects().Get(ids, -1, &results)); - - ASSERT_EQ(results.size(), 2); - for (size_t i = 0; i < ids.size(); i++) { - ASSERT_EQ(results[i]->GetData()->Size(), buffers[i].GetData()->Size()); - ASSERT_EQ(memcmp(results[i]->GetData()->Data(), buffers[i].GetData()->Data(), - buffers[i].GetData()->Size()), - 0); - ASSERT_EQ(results[i]->GetMetadata()->Size(), buffers[i].GetMetadata()->Size()); - ASSERT_EQ(memcmp(results[i]->GetMetadata()->Data(), buffers[i].GetMetadata()->Data(), - buffers[i].GetMetadata()->Size()), - 0); - } - - // Test Wait(). - ObjectID non_existent_id = ObjectID::FromRandom(); - std::vector all_ids(ids); - all_ids.push_back(non_existent_id); - - std::vector wait_results; - RAY_CHECK_OK(core_worker.Objects().Wait(all_ids, 2, -1, &wait_results)); - ASSERT_EQ(wait_results.size(), 3); - ASSERT_EQ(wait_results, std::vector({true, true, false})); - - RAY_CHECK_OK(core_worker.Objects().Wait(all_ids, 3, 100, &wait_results)); - ASSERT_EQ(wait_results.size(), 3); - ASSERT_EQ(wait_results, std::vector({true, true, false})); - - // Test Delete(). - // clear the reference held by PlasmaBuffer. - results.clear(); - RAY_CHECK_OK(core_worker.Objects().Delete(ids, true, false)); - - // Note that Delete() calls RayletClient::FreeObjects and would not - // wait for objects being deleted, so wait a while for plasma store - // to process the command. - usleep(200 * 1000); - RAY_CHECK_OK(core_worker.Objects().Get(ids, 0, &results)); - ASSERT_EQ(results.size(), 2); - ASSERT_TRUE(!results[0]); - ASSERT_TRUE(!results[1]); -} - -TEST_F(TwoNodeTest, TestObjectInterfaceCrossNodes) { - CoreWorker worker1(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[0], - raylet_socket_names_[0], JobID::FromInt(1), nullptr); - - CoreWorker worker2(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[1], - raylet_socket_names_[1], JobID::FromInt(1), nullptr); - - uint8_t array1[] = {1, 2, 3, 4, 5, 6, 7, 8}; - uint8_t array2[] = {10, 11, 12, 13, 14, 15}; - - std::vector buffers; - buffers.emplace_back(array1, sizeof(array1)); - buffers.emplace_back(array2, sizeof(array2)); - - std::vector ids(buffers.size()); - for (size_t i = 0; i < ids.size(); i++) { - RAY_CHECK_OK(worker1.Objects().Put( - RayObject(std::make_shared(buffers[i]), nullptr), &ids[i])); - } - - // Test Get() from remote node. - std::vector> results; - RAY_CHECK_OK(worker2.Objects().Get(ids, -1, &results)); - - ASSERT_EQ(results.size(), 2); - for (size_t i = 0; i < ids.size(); i++) { - ASSERT_EQ(results[i]->GetData()->Size(), buffers[i].Size()); - ASSERT_EQ(memcmp(results[i]->GetData()->Data(), buffers[i].Data(), buffers[i].Size()), - 0); - } - - // Test Wait() from remote node. - ObjectID non_existent_id = ObjectID::FromRandom(); - std::vector all_ids(ids); - all_ids.push_back(non_existent_id); - - std::vector wait_results; - RAY_CHECK_OK(worker2.Objects().Wait(all_ids, 2, -1, &wait_results)); - ASSERT_EQ(wait_results.size(), 3); - ASSERT_EQ(wait_results, std::vector({true, true, false})); - - RAY_CHECK_OK(worker2.Objects().Wait(all_ids, 3, 100, &wait_results)); - ASSERT_EQ(wait_results.size(), 3); - ASSERT_EQ(wait_results, std::vector({true, true, false})); - - // Test Delete() from all machines. - // clear the reference held by PlasmaBuffer. - results.clear(); - RAY_CHECK_OK(worker2.Objects().Delete(ids, false, false)); - - // Note that Delete() calls RayletClient::FreeObjects and would not - // wait for objects being deleted, so wait a while for plasma store - // to process the command. - usleep(1000 * 1000); - // Verify objects are deleted from both machines. - RAY_CHECK_OK(worker2.Objects().Get(ids, 0, &results)); - ASSERT_EQ(results.size(), 2); - ASSERT_TRUE(!results[0]); - ASSERT_TRUE(!results[1]); - - RAY_CHECK_OK(worker1.Objects().Get(ids, 0, &results)); - ASSERT_EQ(results.size(), 2); - ASSERT_TRUE(!results[0]); - ASSERT_TRUE(!results[1]); -} - -TEST_F(SingleNodeTest, TestNormalTaskLocal) { - std::unordered_map resources; - TestNormalTask(resources); -} - -TEST_F(TwoNodeTest, TestNormalTaskCrossNodes) { - std::unordered_map resources; - resources.emplace("resource1", 1); - TestNormalTask(resources); -} - -TEST_F(SingleNodeTest, TestActorTaskLocal) { - std::unordered_map resources; - TestActorTask(resources); -} - -TEST_F(TwoNodeTest, TestActorTaskCrossNodes) { - std::unordered_map resources; - resources.emplace("resource1", 1); - TestActorTask(resources); -} - -TEST_F(SingleNodeTest, TestCoreWorkerConstructorFailure) { - try { - CoreWorker core_worker(WorkerType::DRIVER, Language::PYTHON, "", - raylet_socket_names_[0], JobID::FromInt(1), nullptr); - } catch (const std::exception &e) { - std::cout << "Caught exception when constructing core worker: " << e.what(); - } -} - -} // namespace ray - -int main(int argc, char **argv) { - ::testing::InitGoogleTest(&argc, argv); - RAY_CHECK(argc == 4); - ray::store_executable = std::string(argv[1]); - ray::raylet_executable = std::string(argv[2]); - ray::mock_worker_executable = std::string(argv[3]); - return RUN_ALL_TESTS(); -} diff --git a/src/ray/core_worker/object_interface.h b/src/ray/core_worker/object_interface.h index 2f7ddaacd..9d5009f62 100644 --- a/src/ray/core_worker/object_interface.h +++ b/src/ray/core_worker/object_interface.h @@ -88,6 +88,12 @@ class CoreWorkerObjectInterface { store_providers_; friend class CoreWorkerTaskInterface; + + /// TODO(zhijunfu): This is necessary as direct call task submitter needs to create + /// a local plasma store provider, later we can refactor ObjectInterface to add a + /// `ObjectProviderLayer`, which will encapsulate the functionalities to get or create + /// a specific `StoreProvider`, and this can be removed then. + friend class CoreWorkerDirectActorTaskSubmitter; }; } // namespace ray diff --git a/src/ray/core_worker/store_provider/local_plasma_provider.cc b/src/ray/core_worker/store_provider/local_plasma_provider.cc index 758ebd72a..8ea565603 100644 --- a/src/ray/core_worker/store_provider/local_plasma_provider.cc +++ b/src/ray/core_worker/store_provider/local_plasma_provider.cc @@ -8,12 +8,7 @@ namespace ray { CoreWorkerLocalPlasmaStoreProvider::CoreWorkerLocalPlasmaStoreProvider( const std::string &store_socket) { - auto status = store_client_.Connect(store_socket); - if (!status.ok()) { - RAY_LOG(ERROR) << "Connecting plasma store failed when trying to construct" - << " core worker: " << status.message(); - throw std::runtime_error(status.message()); - } + RAY_ARROW_CHECK_OK(store_client_.Connect(store_socket)); } Status CoreWorkerLocalPlasmaStoreProvider::Put(const RayObject &object, diff --git a/src/ray/core_worker/store_provider/store_provider.h b/src/ray/core_worker/store_provider/store_provider.h index 2533fe8b9..7d2504a03 100644 --- a/src/ray/core_worker/store_provider/store_provider.h +++ b/src/ray/core_worker/store_provider/store_provider.h @@ -24,6 +24,9 @@ class RayObject { /// Return the metadata of the ray object. const std::shared_ptr &GetMetadata() const { return metadata_; }; + /// Whether this object has metadata. + bool HasMetadata() const { return metadata_ != nullptr && metadata_->Size() > 0; } + private: /// Data of the ray object. const std::shared_ptr data_; diff --git a/src/ray/core_worker/task_execution.cc b/src/ray/core_worker/task_execution.cc index c3ddee166..20fe1df1b 100644 --- a/src/ray/core_worker/task_execution.cc +++ b/src/ray/core_worker/task_execution.cc @@ -1,6 +1,7 @@ #include "ray/core_worker/task_execution.h" #include "ray/core_worker/context.h" #include "ray/core_worker/core_worker.h" +#include "ray/core_worker/transport/direct_actor_transport.h" #include "ray/core_worker/transport/raylet_transport.h" namespace ray { @@ -21,6 +22,11 @@ CoreWorkerTaskExecutionInterface::CoreWorkerTaskExecutionInterface( TaskTransportType::RAYLET, std::unique_ptr(new CoreWorkerRayletTaskReceiver( raylet_client, object_interface_, main_service_, worker_server_, func))); + task_receivers_.emplace( + TaskTransportType::DIRECT_ACTOR, + std::unique_ptr( + new CoreWorkerDirectActorTaskReceiver(object_interface_, main_service_, + worker_server_, func))); // Start RPC server after all the task receivers are properly initialized. worker_server_.Run(); @@ -50,8 +56,7 @@ Status CoreWorkerTaskExecutionInterface::ExecuteTask( 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, this logic only - // applies to task submitted via raylet. + // Decrease to account for the dummy object id. num_returns--; } diff --git a/src/ray/core_worker/task_interface.cc b/src/ray/core_worker/task_interface.cc index 47345d79b..43defedf2 100644 --- a/src/ray/core_worker/task_interface.cc +++ b/src/ray/core_worker/task_interface.cc @@ -2,13 +2,14 @@ #include "ray/core_worker/context.h" #include "ray/core_worker/core_worker.h" #include "ray/core_worker/task_interface.h" +#include "ray/core_worker/transport/direct_actor_transport.h" #include "ray/core_worker/transport/raylet_transport.h" namespace ray { ActorHandle::ActorHandle( const class ActorID &actor_id, const class ActorHandleID &actor_handle_id, - const Language actor_language, + const Language actor_language, bool is_direct_call, const std::vector &actor_creation_task_function_descriptor) { inner_.set_actor_id(actor_id.Data(), actor_id.Size()); inner_.set_actor_handle_id(actor_handle_id.Data(), actor_handle_id.Size()); @@ -17,6 +18,7 @@ ActorHandle::ActorHandle( actor_creation_task_function_descriptor.begin(), actor_creation_task_function_descriptor.end()}; inner_.set_actor_cursor(actor_id.Data(), actor_id.Size()); + inner_.set_is_direct_call(is_direct_call); } ActorHandle::ActorHandle(const ActorHandle &other) @@ -44,6 +46,8 @@ int64_t ActorHandle::TaskCounter() const { return inner_.task_counter(); }; int64_t ActorHandle::NumForks() const { return inner_.num_forks(); }; +bool ActorHandle::IsDirectCallActor() const { return inner_.is_direct_call(); } + ActorHandle ActorHandle::Fork() { ActorHandle new_handle; std::unique_lock guard(mutex_); @@ -91,19 +95,25 @@ std::vector ActorHandle::NewActorHandles() const { void ActorHandle::ClearNewActorHandles() { new_actor_handles_.clear(); } CoreWorkerTaskInterface::CoreWorkerTaskInterface( - WorkerContext &worker_context, std::unique_ptr &raylet_client) + WorkerContext &worker_context, std::unique_ptr &raylet_client, + CoreWorkerObjectInterface &object_interface, boost::asio::io_service &io_service, + gcs::RedisGcsClient &gcs_client) : worker_context_(worker_context) { task_submitters_.emplace(TaskTransportType::RAYLET, std::unique_ptr( new CoreWorkerRayletTaskSubmitter(raylet_client))); + task_submitters_.emplace(TaskTransportType::DIRECT_ACTOR, + std::unique_ptr( + new CoreWorkerDirectActorTaskSubmitter( + io_service, gcs_client, object_interface))); } -TaskSpecBuilder CoreWorkerTaskInterface::BuildCommonTaskSpec( - const RayFunction &function, const std::vector &args, uint64_t num_returns, +void CoreWorkerTaskInterface::BuildCommonTaskSpec( + TaskSpecBuilder &builder, const RayFunction &function, + const std::vector &args, uint64_t num_returns, const std::unordered_map &required_resources, const std::unordered_map &required_placement_resources, std::vector *return_ids) { - TaskSpecBuilder builder; auto next_task_index = worker_context_.GetNextTaskIndex(); // Build common task spec. builder.SetCommonTaskSpec( @@ -125,15 +135,15 @@ TaskSpecBuilder CoreWorkerTaskInterface::BuildCommonTaskSpec( for (int i = 0; i < num_returns; i++) { (*return_ids)[i] = ObjectID::ForTaskReturn(task_id, i + 1); } - return builder; } Status CoreWorkerTaskInterface::SubmitTask(const RayFunction &function, const std::vector &args, const TaskOptions &task_options, std::vector *return_ids) { - auto builder = BuildCommonTaskSpec(function, args, task_options.num_returns, - task_options.resources, {}, return_ids); + TaskSpecBuilder builder; + BuildCommonTaskSpec(builder, function, args, task_options.num_returns, + task_options.resources, {}, return_ids); return task_submitters_[TaskTransportType::RAYLET]->SubmitTask(builder.Build()); } @@ -142,15 +152,17 @@ Status CoreWorkerTaskInterface::CreateActor( const ActorCreationOptions &actor_creation_options, std::unique_ptr *actor_handle) { std::vector return_ids; - auto builder = BuildCommonTaskSpec(function, args, 1, actor_creation_options.resources, - actor_creation_options.resources, &return_ids); + TaskSpecBuilder builder; + BuildCommonTaskSpec(builder, function, args, 1, actor_creation_options.resources, + actor_creation_options.resources, &return_ids); const ActorID actor_id = ActorID::FromBinary(return_ids[0].Binary()); builder.SetActorCreationTaskSpec(actor_id, actor_creation_options.max_reconstructions, {}); *actor_handle = std::unique_ptr(new ActorHandle( - actor_id, ActorHandleID::Nil(), function.language, function.function_descriptor)); + actor_id, ActorHandleID::Nil(), function.language, + actor_creation_options.is_direct_call, function.function_descriptor)); (*actor_handle)->IncreaseTaskCounter(); (*actor_handle)->SetActorCursor(return_ids[0]); @@ -162,12 +174,13 @@ Status CoreWorkerTaskInterface::SubmitActorTask(ActorHandle &actor_handle, const std::vector &args, const TaskOptions &task_options, std::vector *return_ids) { - // Add one for actor cursor object id. + // Add one for actor cursor object id for tasks. const auto num_returns = task_options.num_returns + 1; // Build common task spec. - auto builder = BuildCommonTaskSpec(function, args, num_returns, task_options.resources, - {}, return_ids); + TaskSpecBuilder builder; + BuildCommonTaskSpec(builder, function, args, num_returns, task_options.resources, {}, + return_ids); std::unique_lock guard(actor_handle.mutex_); // Build actor task spec. @@ -183,13 +196,17 @@ Status CoreWorkerTaskInterface::SubmitActorTask(ActorHandle &actor_handle, auto actor_cursor = (*return_ids).back(); actor_handle.SetActorCursor(actor_cursor); actor_handle.ClearNewActorHandles(); + guard.unlock(); // Submit task. - auto status = task_submitters_[TaskTransportType::RAYLET]->SubmitTask(builder.Build()); - + const bool is_direct_call = actor_handle.IsDirectCallActor(); + const auto transport_type = + is_direct_call ? TaskTransportType::DIRECT_ACTOR : TaskTransportType::RAYLET; + auto status = task_submitters_[transport_type]->SubmitTask(builder.Build()); // Remove cursor from return ids. (*return_ids).pop_back(); + return status; } diff --git a/src/ray/core_worker/task_interface.h b/src/ray/core_worker/task_interface.h index e261802fe..42e31007f 100644 --- a/src/ray/core_worker/task_interface.h +++ b/src/ray/core_worker/task_interface.h @@ -10,7 +10,9 @@ #include "ray/common/task/task_util.h" #include "ray/core_worker/common.h" #include "ray/core_worker/context.h" +#include "ray/core_worker/object_interface.h" #include "ray/core_worker/transport/transport.h" +#include "ray/gcs/redis_gcs_client.h" #include "ray/protobuf/core_worker.pb.h" namespace ray { @@ -34,13 +36,18 @@ struct TaskOptions { /// Options of an actor creation task. struct ActorCreationOptions { ActorCreationOptions() {} - ActorCreationOptions(uint64_t max_reconstructions, + ActorCreationOptions(uint64_t max_reconstructions, bool is_direct_call, const std::unordered_map &resources) - : max_reconstructions(max_reconstructions), resources(resources) {} + : max_reconstructions(max_reconstructions), + is_direct_call(is_direct_call), + resources(resources) {} /// Maximum number of times that the actor should be reconstructed when it dies /// unexpectedly. It must be non-negative. If it's 0, the actor won't be reconstructed. const uint64_t max_reconstructions = 0; + /// Whether to use direct actor call. If this is set to true, callers will submit + /// tasks directly to the created actor without going through raylet. + const bool is_direct_call = false; /// Resources required by the whole lifetime of this actor. const std::unordered_map resources; }; @@ -49,7 +56,7 @@ struct ActorCreationOptions { class ActorHandle { public: ActorHandle(const ActorID &actor_id, const ActorHandleID &actor_handle_id, - const Language actor_language, + const Language actor_language, bool is_direct_call, const std::vector &actor_creation_task_function_descriptor); ActorHandle(const ActorHandle &other); @@ -77,6 +84,10 @@ class ActorHandle { /// It's used to make sure ids of actor handles are unique. int64_t NumForks() const; + /// Whether direct call is used. If this is true, then the tasks + /// are submitted directly to the actor without going through raylet. + bool IsDirectCallActor() const; + ActorHandle Fork(); void Serialize(std::string *output); @@ -116,7 +127,10 @@ class ActorHandle { class CoreWorkerTaskInterface { public: CoreWorkerTaskInterface(WorkerContext &worker_context, - std::unique_ptr &raylet_client); + std::unique_ptr &raylet_client, + CoreWorkerObjectInterface &object_interface, + boost::asio::io_service &io_service, + gcs::RedisGcsClient &gcs_client); /// Submit a normal task. /// @@ -155,6 +169,7 @@ class CoreWorkerTaskInterface { private: /// Build common attributes of the task spec, and compute return ids. /// + /// \param[in] builder Builder to build a `TaskSpec`. /// \param[in] function The remote function to execute. /// \param[in] args Arguments of this task. /// \param[in] num_returns Number of returns. @@ -162,9 +177,10 @@ class CoreWorkerTaskInterface { /// \param[in] required_placement_resources Resources required by placing this task on a /// node. /// \param[out] return_ids Return IDs. - /// \return A `TaskSpecBuilder`. - TaskSpecBuilder BuildCommonTaskSpec( - const RayFunction &function, const std::vector &args, uint64_t num_returns, + /// \return Void. + void BuildCommonTaskSpec( + TaskSpecBuilder &builder, const RayFunction &function, + const std::vector &args, uint64_t num_returns, const std::unordered_map &required_resources, const std::unordered_map &required_placement_resources, std::vector *return_ids); @@ -175,6 +191,8 @@ class CoreWorkerTaskInterface { /// All the task submitters supported. EnumUnorderedMap> task_submitters_; + + friend class CoreWorkerTest; }; } // namespace ray diff --git a/src/ray/core_worker/test/core_worker_test.cc b/src/ray/core_worker/test/core_worker_test.cc new file mode 100644 index 000000000..560728e30 --- /dev/null +++ b/src/ray/core_worker/test/core_worker_test.cc @@ -0,0 +1,827 @@ +#include +#include "gmock/gmock.h" +#include "gtest/gtest.h" + +#include "ray/common/buffer.h" +#include "ray/core_worker/context.h" +#include "ray/core_worker/core_worker.h" +#include "ray/core_worker/transport/direct_actor_transport.h" +#include "ray/rpc/raylet/raylet_client.h" +#include "src/ray/util/test_util.h" + +#include "src/ray/protobuf/direct_actor.grpc.pb.h" +#include "src/ray/protobuf/direct_actor.pb.h" + +#include +#include +#include + +#include "ray/thirdparty/hiredis/async.h" +#include "ray/thirdparty/hiredis/hiredis.h" + +namespace ray { + +std::string store_executable; +std::string raylet_executable; +std::string mock_worker_executable; + +ray::ObjectID RandomObjectID() { return ObjectID::FromRandom(); } + +static void flushall_redis(void) { + redisContext *context = redisConnect("127.0.0.1", 6379); + freeReplyObject(redisCommand(context, "FLUSHALL")); + freeReplyObject(redisCommand(context, "SET NumRedisShards 1")); + freeReplyObject(redisCommand(context, "LPUSH RedisShards 127.0.0.1:6380")); + redisFree(context); +} + +std::shared_ptr GenerateRandomBuffer() { + auto seed = std::chrono::high_resolution_clock::now().time_since_epoch().count(); + std::mt19937 gen(seed); + std::uniform_int_distribution<> dis(1, 10); + std::uniform_int_distribution<> value_dis(1, 255); + + std::vector arg1(dis(gen), value_dis(gen)); + return std::make_shared(arg1.data(), arg1.size(), true); +} + +std::unique_ptr CreateActorHelper( + CoreWorker &worker, const std::unordered_map &resources, + bool is_direct_call, uint64_t max_reconstructions) { + std::unique_ptr actor_handle; + + // Test creating actor. + uint8_t array[] = {1, 2, 3}; + auto buffer = std::make_shared(array, sizeof(array)); + + RayFunction func{ray::Language::PYTHON, {"actor creation task"}}; + std::vector args; + args.emplace_back(TaskArg::PassByValue(buffer)); + + ActorCreationOptions actor_options{max_reconstructions, is_direct_call, resources}; + + // Create an actor. + RAY_CHECK_OK(worker.Tasks().CreateActor(func, args, actor_options, &actor_handle)); + return actor_handle; +} + +class CoreWorkerTest : public ::testing::Test { + public: + CoreWorkerTest(int num_nodes) : gcs_options_("127.0.0.1", 6379, "") { + // flush redis first. + flushall_redis(); + + RAY_CHECK(num_nodes >= 0); + if (num_nodes > 0) { + raylet_socket_names_.resize(num_nodes); + raylet_store_socket_names_.resize(num_nodes); + } + + // start plasma store. + for (auto &store_socket : raylet_store_socket_names_) { + store_socket = StartStore(); + } + + // start raylet on each node. Assign each node with different resources so that + // a task can be scheduled to the desired node. + for (int i = 0; i < num_nodes; i++) { + raylet_socket_names_[i] = + StartRaylet(raylet_store_socket_names_[i], "127.0.0.1", "127.0.0.1", + "\"CPU,4.0,resource" + std::to_string(i) + ",10\""); + } + } + + ~CoreWorkerTest() { + for (const auto &raylet_socket : raylet_socket_names_) { + StopRaylet(raylet_socket); + } + + for (const auto &store_socket : raylet_store_socket_names_) { + StopStore(store_socket); + } + } + + std::string StartStore() { + std::string store_socket_name = "/tmp/store" + RandomObjectID().Hex(); + std::string store_pid = store_socket_name + ".pid"; + std::string plasma_command = store_executable + " -m 10000000 -s " + + store_socket_name + + " 1> /dev/null 2> /dev/null & echo $! > " + store_pid; + RAY_LOG(DEBUG) << plasma_command; + RAY_CHECK(system(plasma_command.c_str()) == 0); + usleep(200 * 1000); + return store_socket_name; + } + + void StopStore(std::string store_socket_name) { + std::string store_pid = store_socket_name + ".pid"; + std::string kill_9 = "kill -9 `cat " + store_pid + "`"; + RAY_LOG(DEBUG) << kill_9; + ASSERT_EQ(system(kill_9.c_str()), 0); + ASSERT_EQ(system(("rm -rf " + store_socket_name).c_str()), 0); + ASSERT_EQ(system(("rm -rf " + store_socket_name + ".pid").c_str()), 0); + } + + std::string StartRaylet(std::string store_socket_name, std::string node_ip_address, + std::string redis_address, std::string resource) { + std::string raylet_socket_name = "/tmp/raylet" + RandomObjectID().Hex(); + std::string ray_start_cmd = raylet_executable; + ray_start_cmd.append(" --raylet_socket_name=" + raylet_socket_name) + .append(" --store_socket_name=" + store_socket_name) + .append(" --object_manager_port=0 --node_manager_port=0") + .append(" --node_ip_address=" + node_ip_address) + .append(" --redis_address=" + redis_address) + .append(" --redis_port=6379") + .append(" --num_initial_workers=1") + .append(" --maximum_startup_concurrency=10") + .append(" --static_resource_list=" + resource) + .append(" --python_worker_command=\"" + mock_worker_executable + " " + + store_socket_name + " " + raylet_socket_name + "\"") + .append(" & echo $! > " + raylet_socket_name + ".pid"); + + RAY_LOG(DEBUG) << "Ray Start command: " << ray_start_cmd; + RAY_CHECK(system(ray_start_cmd.c_str()) == 0); + usleep(200 * 1000); + return raylet_socket_name; + } + + void StopRaylet(std::string raylet_socket_name) { + std::string raylet_pid = raylet_socket_name + ".pid"; + std::string kill_9 = "kill -9 `cat " + raylet_pid + "`"; + RAY_LOG(DEBUG) << kill_9; + ASSERT_TRUE(system(kill_9.c_str()) == 0); + ASSERT_TRUE(system(("rm -rf " + raylet_socket_name).c_str()) == 0); + ASSERT_TRUE(system(("rm -rf " + raylet_socket_name + ".pid").c_str()) == 0); + } + + void SetUp() {} + + void TearDown() {} + + // Test normal tasks. + void TestNormalTask(const std::unordered_map &resources); + + // Test actor tasks. + void TestActorTask(const std::unordered_map &resources, + bool is_direct_call); + + // Test actor failure case, verify that the tasks would either succeed or + // fail with exceptions, in that case the return objects fetched from `Get` + // contain errors. + void TestActorFailure(const std::unordered_map &resources, + bool is_direct_call); + + // Test actor failover case. Verify that actor can be reconstructed successfully, + // and as long as we wait for actor reconstruction before submitting new tasks, + // it is guaranteed that all tasks are successfully completed. + void TestActorReconstruction(const std::unordered_map &resources, + bool is_direct_call); + + protected: + bool WaitForDirectCallActorState(CoreWorker &worker, const ActorID &actor_id, + bool wait_alive, int timeout_ms); + + std::vector raylet_socket_names_; + std::vector raylet_store_socket_names_; + gcs::GcsClientOptions gcs_options_; +}; + +bool CoreWorkerTest::WaitForDirectCallActorState(CoreWorker &worker, + const ActorID &actor_id, bool wait_alive, + int timeout_ms) { + auto condition_func = [&worker, actor_id, wait_alive]() -> bool { + auto &task_submitters = worker.Tasks().task_submitters_; + RAY_CHECK(task_submitters.count(TaskTransportType::DIRECT_ACTOR) > 0); + auto submitter = + worker.Tasks().task_submitters_[TaskTransportType::DIRECT_ACTOR].get(); + auto direct_actor_submitter = + dynamic_cast(submitter); + RAY_CHECK(direct_actor_submitter != nullptr); + bool actor_alive = direct_actor_submitter->IsActorAlive(actor_id); + return wait_alive ? actor_alive : !actor_alive; + }; + + return WaitForCondition(condition_func, timeout_ms); +} + +void CoreWorkerTest::TestNormalTask( + const std::unordered_map &resources) { + CoreWorker driver(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[0], + raylet_socket_names_[0], JobID::FromInt(1), gcs_options_, nullptr); + + // Test for tasks with by-value and by-ref args. + { + const int num_tasks = 100; + for (int i = 0; i < num_tasks; i++) { + auto buffer1 = GenerateRandomBuffer(); + auto buffer2 = GenerateRandomBuffer(); + + ObjectID object_id; + RAY_CHECK_OK(driver.Objects().Put(RayObject(buffer2, nullptr), &object_id)); + + std::vector args; + args.emplace_back(TaskArg::PassByValue(buffer1)); + args.emplace_back(TaskArg::PassByReference(object_id)); + + RayFunction func{ray::Language::PYTHON, {}}; + TaskOptions options; + + std::vector return_ids; + RAY_CHECK_OK(driver.Tasks().SubmitTask(func, args, options, &return_ids)); + + ASSERT_EQ(return_ids.size(), 1); + + std::vector> results; + RAY_CHECK_OK(driver.Objects().Get(return_ids, -1, &results)); + + ASSERT_EQ(results.size(), 1); + ASSERT_EQ(results[0]->GetData()->Size(), buffer1->Size() + buffer2->Size()); + ASSERT_EQ(memcmp(results[0]->GetData()->Data(), buffer1->Data(), buffer1->Size()), + 0); + ASSERT_EQ(memcmp(results[0]->GetData()->Data() + buffer1->Size(), buffer2->Data(), + buffer2->Size()), + 0); + } + } +} + +void CoreWorkerTest::TestActorTask( + const std::unordered_map &resources, bool is_direct_call) { + CoreWorker driver(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[0], + raylet_socket_names_[0], JobID::FromInt(1), gcs_options_, nullptr); + + auto actor_handle = CreateActorHelper(driver, resources, is_direct_call, 1000); + + // Test submitting some tasks with by-value args for that actor. + { + const int num_tasks = 100; + for (int i = 0; i < num_tasks; i++) { + auto buffer1 = GenerateRandomBuffer(); + auto buffer2 = GenerateRandomBuffer(); + + // Create arguments with PassByRef and PassByValue. + std::vector args; + args.emplace_back(TaskArg::PassByValue(buffer1)); + args.emplace_back(TaskArg::PassByValue(buffer2)); + + TaskOptions options{1, resources}; + std::vector return_ids; + RayFunction func{ray::Language::PYTHON, {}}; + + RAY_CHECK_OK(driver.Tasks().SubmitActorTask(*actor_handle, func, args, options, + &return_ids)); + ASSERT_EQ(return_ids.size(), 1); + + std::vector> results; + RAY_CHECK_OK(driver.Objects().Get(return_ids, -1, &results)); + + ASSERT_EQ(results.size(), 1); + ASSERT_EQ(results[0]->GetData()->Size(), buffer1->Size() + buffer2->Size()); + ASSERT_EQ(memcmp(results[0]->GetData()->Data(), buffer1->Data(), buffer1->Size()), + 0); + ASSERT_EQ(memcmp(results[0]->GetData()->Data() + buffer1->Size(), buffer2->Data(), + buffer2->Size()), + 0); + } + } + + // Test submitting a task with both by-value and by-ref args for that actor. + { + uint8_t array1[] = {1, 2, 3, 4, 5, 6, 7, 8}; + uint8_t array2[] = {10, 11, 12, 13, 14, 15}; + + auto buffer1 = std::make_shared(array1, sizeof(array1)); + auto buffer2 = std::make_shared(array2, sizeof(array2)); + + ObjectID object_id; + RAY_CHECK_OK(driver.Objects().Put(RayObject(buffer1, nullptr), &object_id)); + + // Create arguments with PassByRef and PassByValue. + std::vector args; + args.emplace_back(TaskArg::PassByReference(object_id)); + args.emplace_back(TaskArg::PassByValue(buffer2)); + + TaskOptions options{1, resources}; + std::vector return_ids; + RayFunction func{ray::Language::PYTHON, {}}; + auto status = + driver.Tasks().SubmitActorTask(*actor_handle, func, args, options, &return_ids); + if (is_direct_call) { + // For direct actor call, submitting a task with by-reference arguments + // would fail. + ASSERT_TRUE(!status.ok()); + return; + } + + ASSERT_EQ(return_ids.size(), 1); + + std::vector> results; + RAY_CHECK_OK(driver.Objects().Get(return_ids, -1, &results)); + + ASSERT_EQ(results.size(), 1); + ASSERT_EQ(results[0]->GetData()->Size(), buffer1->Size() + buffer2->Size()); + ASSERT_EQ(memcmp(results[0]->GetData()->Data(), buffer1->Data(), buffer1->Size()), 0); + ASSERT_EQ(memcmp(results[0]->GetData()->Data() + buffer1->Size(), buffer2->Data(), + buffer2->Size()), + 0); + } +} + +void CoreWorkerTest::TestActorReconstruction( + const std::unordered_map &resources, bool is_direct_call) { + CoreWorker driver(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[0], + raylet_socket_names_[0], JobID::FromInt(1), gcs_options_, nullptr); + + // creating actor. + auto actor_handle = CreateActorHelper(driver, resources, is_direct_call, 1000); + + // Wait for actor alive event. + ASSERT_TRUE(WaitForDirectCallActorState(driver, actor_handle->ActorID(), true, + 30 * 1000 /* 30s */)); + RAY_LOG(INFO) << "actor has been created"; + + // Test submitting some tasks with by-value args for that actor. + { + const int num_tasks = 100; + const int task_index_to_kill_worker = (num_tasks + 1) / 2; + std::vector>> all_results; + for (int i = 0; i < num_tasks; i++) { + if (i == task_index_to_kill_worker) { + RAY_LOG(INFO) << "killing worker"; + system("pkill mock_worker"); + + // Wait for actor restruction event, and then for alive event. + ASSERT_TRUE(WaitForDirectCallActorState(driver, actor_handle->ActorID(), false, + 30 * 1000 /* 30s */)); + ASSERT_TRUE(WaitForDirectCallActorState(driver, actor_handle->ActorID(), true, + 30 * 1000 /* 30s */)); + + RAY_LOG(INFO) << "actor has been reconstructed"; + } + + // wait for actor being reconstructed. + auto buffer1 = GenerateRandomBuffer(); + + // Create arguments with PassByValue. + std::vector args; + args.emplace_back(TaskArg::PassByValue(buffer1)); + + TaskOptions options{1, resources}; + std::vector return_ids; + RayFunction func{ray::Language::PYTHON, {}}; + + auto status = + driver.Tasks().SubmitActorTask(*actor_handle, func, args, options, &return_ids); + RAY_CHECK_OK(status); + ASSERT_EQ(return_ids.size(), 1); + // Verify if it's expected data. + std::vector> results; + + RAY_CHECK_OK(driver.Objects().Get(return_ids, -1, &results)); + ASSERT_EQ(results[0]->GetData()->Size(), buffer1->Size()); + ASSERT_EQ(*results[0]->GetData(), *buffer1); + } + } +} + +void CoreWorkerTest::TestActorFailure( + const std::unordered_map &resources, bool is_direct_call) { + CoreWorker driver(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[0], + raylet_socket_names_[0], JobID::FromInt(1), gcs_options_, nullptr); + + // creating actor. + auto actor_handle = + CreateActorHelper(driver, resources, is_direct_call, 0 /* not reconstructable */); + + // Test submitting some tasks with by-value args for that actor. + { + const int num_tasks = 3000; + const int task_index_to_kill_worker = (num_tasks + 1) / 2; + std::vector>> all_results; + for (int i = 0; i < num_tasks; i++) { + if (i == task_index_to_kill_worker) { + RAY_LOG(INFO) << "killing worker"; + system("pkill mock_worker"); + } + + // wait for actor being reconstructed. + auto buffer1 = GenerateRandomBuffer(); + + // Create arguments with PassByRef and PassByValue. + std::vector args; + args.emplace_back(TaskArg::PassByValue(buffer1)); + + TaskOptions options{1, resources}; + std::vector return_ids; + RayFunction func{ray::Language::PYTHON, {}}; + + auto status = + driver.Tasks().SubmitActorTask(*actor_handle, func, args, options, &return_ids); + if (i < task_index_to_kill_worker) { + RAY_CHECK_OK(status); + } + + ASSERT_EQ(return_ids.size(), 1); + all_results.emplace_back(std::make_pair(return_ids[0], buffer1)); + } + + for (int i = 0; i < num_tasks; i++) { + const auto &entry = all_results[i]; + std::vector return_ids; + return_ids.push_back(entry.first); + std::vector> results; + RAY_CHECK_OK(driver.Objects().Get(return_ids, -1, &results)); + ASSERT_EQ(results.size(), 1); + + if (results[0]->HasMetadata()) { + // Verify if this is the desired error. + std::string meta = std::to_string(static_cast(rpc::ErrorType::ACTOR_DIED)); + ASSERT_TRUE(memcmp(results[0]->GetMetadata()->Data(), meta.data(), meta.size()) == + 0); + } else { + // Verify if it's expected data. + ASSERT_EQ(*results[0]->GetData(), *entry.second); + } + } + } +} + +class ZeroNodeTest : public CoreWorkerTest { + public: + ZeroNodeTest() : CoreWorkerTest(0) {} +}; + +class SingleNodeTest : public CoreWorkerTest { + public: + SingleNodeTest() : CoreWorkerTest(1) {} +}; + +class TwoNodeTest : public CoreWorkerTest { + public: + TwoNodeTest() : CoreWorkerTest(2) {} +}; + +TEST_F(ZeroNodeTest, TestTaskArg) { + // Test by-reference argument. + ObjectID id = ObjectID::FromRandom(); + TaskArg by_ref = TaskArg::PassByReference(id); + ASSERT_TRUE(by_ref.IsPassedByReference()); + ASSERT_EQ(by_ref.GetReference(), id); + // Test by-value argument. + std::shared_ptr buffer = + std::make_shared(static_cast(0), 0); + TaskArg by_value = TaskArg::PassByValue(buffer); + ASSERT_FALSE(by_value.IsPassedByReference()); + auto data = by_value.GetValue(); + ASSERT_TRUE(data != nullptr); + ASSERT_EQ(*data, *buffer); +} + +// Performance batchmark for `PushTaskRequest` creation. +TEST_F(ZeroNodeTest, TestTaskSpecPerf) { + // Create a dummy actor handle, and then create a number of `TaskSpec` + // to benchmark performance. + uint8_t array[] = {1, 2, 3}; + auto buffer = std::make_shared(array, sizeof(array)); + RayFunction function{ray::Language::PYTHON, {}}; + std::vector args; + args.emplace_back(TaskArg::PassByValue(buffer)); + + std::unordered_map resources; + ActorCreationOptions actor_options{0, /* is_direct_call */ true, resources}; + + ActorHandle actor_handle(ActorID::FromRandom(), ActorHandleID::Nil(), function.language, + true, function.function_descriptor); + + // Manually create `num_tasks` task specs, and for each of them create a + // `PushTaskRequest`, this is to batch performance of TaskSpec + // creation/copy/destruction. + int64_t start_ms = current_time_ms(); + const auto num_tasks = 10000 * 10; + RAY_LOG(INFO) << "start creating " << num_tasks << " PushTaskRequests"; + for (int i = 0; i < num_tasks; i++) { + TaskOptions options{1, resources}; + std::vector return_ids; + auto num_returns = options.num_returns; + + TaskSpecBuilder builder; + builder.SetCommonTaskSpec(function.language, function.function_descriptor, + JobID::FromInt(1), TaskID::FromRandom(), 0, num_returns, + resources, resources); + // Set task arguments. + for (const auto &arg : args) { + if (arg.IsPassedByReference()) { + builder.AddByRefArg(arg.GetReference()); + } else { + builder.AddByValueArg(arg.GetValue()->Data(), arg.GetValue()->Size()); + } + } + + const auto actor_creation_dummy_object_id = + ObjectID::FromBinary(actor_handle.ActorID().Binary()); + builder.SetActorTaskSpec( + actor_handle.ActorID(), actor_handle.ActorHandleID(), + actor_creation_dummy_object_id, + /*previous_actor_task_dummy_object_id=*/actor_handle.ActorCursor(), 0, {}); + + const auto &task_spec = builder.Build(); + + ASSERT_TRUE(task_spec.IsActorTask()); + auto request = std::unique_ptr(new rpc::PushTaskRequest); + request->mutable_task_spec()->Swap(&task_spec.GetMutableMessage()); + } + RAY_LOG(INFO) << "Finish creating " << num_tasks << " PushTaskRequests" + << ", which takes " << current_time_ms() - start_ms << " ms"; +} + +TEST_F(SingleNodeTest, TestDirectActorTaskSubmissionPerf) { + CoreWorker driver(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[0], + raylet_socket_names_[0], JobID::FromInt(1), gcs_options_, nullptr); + + std::unique_ptr actor_handle; + + // Test creating actor. + uint8_t array[] = {1, 2, 3}; + auto buffer = std::make_shared(array, sizeof(array)); + + RayFunction func{ray::Language::PYTHON, {}}; + std::vector args; + args.emplace_back(TaskArg::PassByValue(buffer)); + + std::unordered_map resources; + ActorCreationOptions actor_options{0, /* is_direct_call */ true, resources}; + + // Create an actor. + RAY_CHECK_OK(driver.Tasks().CreateActor(func, args, actor_options, &actor_handle)); + + // wait for actor creation finish. + ASSERT_TRUE(WaitForDirectCallActorState(driver, actor_handle->ActorID(), true, + 30 * 1000 /* 30s */)); + + // Test submitting some tasks with by-value args for that actor. + int64_t start_ms = current_time_ms(); + const int num_tasks = 10000; + RAY_LOG(INFO) << "start submitting " << num_tasks << " tasks"; + for (int i = 0; i < num_tasks; i++) { + // Create arguments with PassByValue. + std::vector args; + args.emplace_back(TaskArg::PassByValue(buffer)); + + TaskOptions options{1, resources}; + std::vector return_ids; + RayFunction func{ray::Language::PYTHON, {}}; + + RAY_CHECK_OK( + driver.Tasks().SubmitActorTask(*actor_handle, func, args, options, &return_ids)); + ASSERT_EQ(return_ids.size(), 1); + } + RAY_LOG(INFO) << "finish submitting " << num_tasks << " tasks" + << ", which takes " << current_time_ms() - start_ms << " ms"; +} + +TEST_F(ZeroNodeTest, TestWorkerContext) { + auto job_id = JobID::JobID::FromInt(1); + + WorkerContext context(WorkerType::WORKER, job_id); + ASSERT_TRUE(context.GetCurrentTaskID().IsNil()); + ASSERT_EQ(context.GetNextTaskIndex(), 1); + ASSERT_EQ(context.GetNextTaskIndex(), 2); + ASSERT_EQ(context.GetNextPutIndex(), 1); + ASSERT_EQ(context.GetNextPutIndex(), 2); + + auto thread_func = [&context]() { + // Verify that task_index, put_index are thread-local. + ASSERT_TRUE(!context.GetCurrentTaskID().IsNil()); + ASSERT_EQ(context.GetNextTaskIndex(), 1); + ASSERT_EQ(context.GetNextPutIndex(), 1); + }; + + std::thread async_thread(thread_func); + async_thread.join(); + + // Verify that these fields are thread-local. + ASSERT_EQ(context.GetNextTaskIndex(), 3); + ASSERT_EQ(context.GetNextPutIndex(), 3); +} + +TEST_F(ZeroNodeTest, TestActorHandle) { + ActorHandle handle1(ActorID::FromRandom(), ActorHandleID::FromRandom(), Language::JAVA, + false, + {"org.ray.exampleClass", "exampleMethod", "exampleSignature"}); + + auto forkedHandle1 = handle1.Fork(); + ASSERT_EQ(1, handle1.NumForks()); + ASSERT_EQ(handle1.ActorID(), forkedHandle1.ActorID()); + ASSERT_NE(handle1.ActorHandleID(), forkedHandle1.ActorHandleID()); + ASSERT_EQ(handle1.ActorLanguage(), forkedHandle1.ActorLanguage()); + ASSERT_EQ(handle1.ActorCreationTaskFunctionDescriptor(), + forkedHandle1.ActorCreationTaskFunctionDescriptor()); + ASSERT_EQ(handle1.ActorCursor(), forkedHandle1.ActorCursor()); + ASSERT_EQ(0, forkedHandle1.TaskCounter()); + ASSERT_EQ(0, forkedHandle1.NumForks()); + auto forkedHandle2 = handle1.Fork(); + ASSERT_EQ(2, handle1.NumForks()); + ASSERT_EQ(0, forkedHandle2.TaskCounter()); + ASSERT_EQ(0, forkedHandle2.NumForks()); + + std::string buffer; + handle1.Serialize(&buffer); + auto handle2 = ActorHandle::Deserialize(buffer); + ASSERT_EQ(handle1.ActorID(), handle2.ActorID()); + ASSERT_EQ(handle1.ActorHandleID(), handle2.ActorHandleID()); + ASSERT_EQ(handle1.ActorLanguage(), handle2.ActorLanguage()); + ASSERT_EQ(handle1.ActorCreationTaskFunctionDescriptor(), + handle2.ActorCreationTaskFunctionDescriptor()); + ASSERT_EQ(handle1.ActorCursor(), handle2.ActorCursor()); + ASSERT_EQ(handle1.TaskCounter(), handle2.TaskCounter()); + ASSERT_EQ(handle1.NumForks(), handle2.NumForks()); +} + +TEST_F(SingleNodeTest, TestObjectInterface) { + CoreWorker core_worker(WorkerType::DRIVER, Language::PYTHON, + raylet_store_socket_names_[0], raylet_socket_names_[0], + JobID::FromInt(1), gcs_options_, nullptr); + + uint8_t array1[] = {1, 2, 3, 4, 5, 6, 7, 8}; + uint8_t array2[] = {10, 11, 12, 13, 14, 15}; + + std::vector buffers; + buffers.emplace_back(std::make_shared(array1, sizeof(array1)), + std::make_shared(array1, sizeof(array1) / 2)); + buffers.emplace_back(std::make_shared(array2, sizeof(array2)), + std::make_shared(array2, sizeof(array2) / 2)); + + std::vector ids(buffers.size()); + for (size_t i = 0; i < ids.size(); i++) { + RAY_CHECK_OK(core_worker.Objects().Put(buffers[i], &ids[i])); + } + + // Test Get(). + std::vector> results; + RAY_CHECK_OK(core_worker.Objects().Get(ids, -1, &results)); + + ASSERT_EQ(results.size(), 2); + for (size_t i = 0; i < ids.size(); i++) { + ASSERT_EQ(*results[i]->GetData(), *buffers[i].GetData()); + ASSERT_EQ(*results[i]->GetMetadata(), *buffers[i].GetMetadata()); + } + + // Test Wait(). + ObjectID non_existent_id = ObjectID::FromRandom(); + std::vector all_ids(ids); + all_ids.push_back(non_existent_id); + + std::vector wait_results; + RAY_CHECK_OK(core_worker.Objects().Wait(all_ids, 2, -1, &wait_results)); + ASSERT_EQ(wait_results.size(), 3); + ASSERT_EQ(wait_results, std::vector({true, true, false})); + + RAY_CHECK_OK(core_worker.Objects().Wait(all_ids, 3, 100, &wait_results)); + ASSERT_EQ(wait_results.size(), 3); + ASSERT_EQ(wait_results, std::vector({true, true, false})); + + // Test Delete(). + // clear the reference held by PlasmaBuffer. + results.clear(); + RAY_CHECK_OK(core_worker.Objects().Delete(ids, true, false)); + + // Note that Delete() calls RayletClient::FreeObjects and would not + // wait for objects being deleted, so wait a while for plasma store + // to process the command. + usleep(200 * 1000); + RAY_CHECK_OK(core_worker.Objects().Get(ids, 0, &results)); + ASSERT_EQ(results.size(), 2); + ASSERT_TRUE(!results[0]); + ASSERT_TRUE(!results[1]); +} + +TEST_F(TwoNodeTest, TestObjectInterfaceCrossNodes) { + CoreWorker worker1(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[0], + raylet_socket_names_[0], JobID::FromInt(1), gcs_options_, nullptr); + + CoreWorker worker2(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[1], + raylet_socket_names_[1], JobID::FromInt(1), gcs_options_, nullptr); + + uint8_t array1[] = {1, 2, 3, 4, 5, 6, 7, 8}; + uint8_t array2[] = {10, 11, 12, 13, 14, 15}; + + std::vector buffers; + buffers.emplace_back(array1, sizeof(array1)); + buffers.emplace_back(array2, sizeof(array2)); + + std::vector ids(buffers.size()); + for (size_t i = 0; i < ids.size(); i++) { + RAY_CHECK_OK(worker1.Objects().Put( + RayObject(std::make_shared(buffers[i]), nullptr), &ids[i])); + } + + // Test Get() from remote node. + std::vector> results; + RAY_CHECK_OK(worker2.Objects().Get(ids, -1, &results)); + + ASSERT_EQ(results.size(), 2); + for (size_t i = 0; i < ids.size(); i++) { + ASSERT_EQ(results[i]->GetData()->Size(), buffers[i].Size()); + ASSERT_EQ(*(results[i]->GetData()), buffers[i]); + } + + // Test Wait() from remote node. + ObjectID non_existent_id = ObjectID::FromRandom(); + std::vector all_ids(ids); + all_ids.push_back(non_existent_id); + + std::vector wait_results; + RAY_CHECK_OK(worker2.Objects().Wait(all_ids, 2, -1, &wait_results)); + ASSERT_EQ(wait_results.size(), 3); + ASSERT_EQ(wait_results, std::vector({true, true, false})); + + RAY_CHECK_OK(worker2.Objects().Wait(all_ids, 3, 100, &wait_results)); + ASSERT_EQ(wait_results.size(), 3); + ASSERT_EQ(wait_results, std::vector({true, true, false})); + + // Test Delete() from all machines. + // clear the reference held by PlasmaBuffer. + results.clear(); + RAY_CHECK_OK(worker2.Objects().Delete(ids, false, false)); + + // Note that Delete() calls RayletClient::FreeObjects and would not + // wait for objects being deleted, so wait a while for plasma store + // to process the command. + usleep(1000 * 1000); + // Verify objects are deleted from both machines. + RAY_CHECK_OK(worker2.Objects().Get(ids, 0, &results)); + ASSERT_EQ(results.size(), 2); + ASSERT_TRUE(!results[0]); + ASSERT_TRUE(!results[1]); + + RAY_CHECK_OK(worker1.Objects().Get(ids, 0, &results)); + ASSERT_EQ(results.size(), 2); + ASSERT_TRUE(!results[0]); + ASSERT_TRUE(!results[1]); +} + +TEST_F(SingleNodeTest, TestNormalTaskLocal) { + std::unordered_map resources; + TestNormalTask(resources); +} + +TEST_F(TwoNodeTest, TestNormalTaskCrossNodes) { + std::unordered_map resources; + resources.emplace("resource1", 1); + TestNormalTask(resources); +} + +TEST_F(SingleNodeTest, TestActorTaskLocal) { + std::unordered_map resources; + TestActorTask(resources, false); +} + +TEST_F(TwoNodeTest, TestActorTaskCrossNodes) { + std::unordered_map resources; + resources.emplace("resource1", 1); + TestActorTask(resources, false); +} + +TEST_F(SingleNodeTest, TestDirectActorTaskLocal) { + std::unordered_map resources; + TestActorTask(resources, true); +} + +TEST_F(TwoNodeTest, TestDirectActorTaskCrossNodes) { + std::unordered_map resources; + resources.emplace("resource1", 1); + TestActorTask(resources, true); +} + +TEST_F(SingleNodeTest, TestDirectActorTaskLocalReconstruction) { + std::unordered_map resources; + TestActorReconstruction(resources, true); +} + +TEST_F(TwoNodeTest, TestDirectActorTaskCrossNodesReconstruction) { + std::unordered_map resources; + resources.emplace("resource1", 1); + TestActorReconstruction(resources, true); +} + +TEST_F(SingleNodeTest, TestDirectActorTaskLocalFailure) { + std::unordered_map resources; + TestActorFailure(resources, true); +} + +TEST_F(TwoNodeTest, TestDirectActorTaskCrossNodesFailure) { + std::unordered_map resources; + resources.emplace("resource1", 1); + TestActorFailure(resources, true); +} + +} // namespace ray + +int main(int argc, char **argv) { + ::testing::InitGoogleTest(&argc, argv); + RAY_CHECK(argc == 4); + ray::store_executable = std::string(argv[1]); + ray::raylet_executable = std::string(argv[2]); + ray::mock_worker_executable = std::string(argv[3]); + return RUN_ALL_TESTS(); +} diff --git a/src/ray/core_worker/mock_worker.cc b/src/ray/core_worker/test/mock_worker.cc similarity index 87% rename from src/ray/core_worker/mock_worker.cc rename to src/ray/core_worker/test/mock_worker.cc index 7117f4ec4..d4f7a7abd 100644 --- a/src/ray/core_worker/mock_worker.cc +++ b/src/ray/core_worker/test/mock_worker.cc @@ -1,7 +1,9 @@ +#define BOOST_BIND_NO_PLACEHOLDERS #include "ray/core_worker/context.h" #include "ray/core_worker/core_worker.h" #include "ray/core_worker/store_provider/store_provider.h" #include "ray/core_worker/task_execution.h" +#include "src/ray/util/test_util.h" using namespace std::placeholders; @@ -18,9 +20,10 @@ namespace ray { /// for more details on how this class is used. class MockWorker { public: - MockWorker(const std::string &store_socket, const std::string &raylet_socket) + 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::JobID::FromInt(1), + JobID::FromInt(1), gcs_options, std::bind(&MockWorker::ExecuteTask, this, _1, _2, _3, _4, _5)) {} void Run() { @@ -49,6 +52,7 @@ class MockWorker { for (int i = 0; i < num_returns; i++) { results->push_back(std::make_shared(memory_buffer, nullptr)); } + return Status::OK(); } @@ -62,7 +66,8 @@ int main(int argc, char **argv) { auto store_socket = std::string(argv[1]); auto raylet_socket = std::string(argv[2]); - ray::MockWorker worker(store_socket, raylet_socket); + ray::gcs::GcsClientOptions gcs_options("127.0.0.1", 6379, ""); + ray::MockWorker worker(store_socket, raylet_socket, gcs_options); worker.Run(); return 0; } diff --git a/src/ray/core_worker/transport/direct_actor_transport.cc b/src/ray/core_worker/transport/direct_actor_transport.cc new file mode 100644 index 000000000..fdcabd3a7 --- /dev/null +++ b/src/ray/core_worker/transport/direct_actor_transport.cc @@ -0,0 +1,222 @@ + +#include "ray/core_worker/transport/direct_actor_transport.h" +#include "ray/common/task/task.h" + +using ray::rpc::ActorTableData; + +namespace ray { + +bool HasByReferenceArgs(const TaskSpecification &spec) { + for (int i = 0; i < spec.NumArgs(); ++i) { + if (spec.ArgIdCount(i) > 0) { + return true; + } + } + return false; +} + +CoreWorkerDirectActorTaskSubmitter::CoreWorkerDirectActorTaskSubmitter( + boost::asio::io_service &io_service, gcs::RedisGcsClient &gcs_client, + CoreWorkerObjectInterface &object_interface) + : io_service_(io_service), + gcs_client_(gcs_client), + client_call_manager_(io_service), + store_provider_( + object_interface.CreateStoreProvider(StoreProviderType::LOCAL_PLASMA)) { + RAY_CHECK_OK(SubscribeActorUpdates()); +} + +Status CoreWorkerDirectActorTaskSubmitter::SubmitTask( + const TaskSpecification &task_spec) { + if (HasByReferenceArgs(task_spec)) { + return Status::Invalid("direct actor call only supports by-value arguments"); + } + + RAY_CHECK(task_spec.IsActorTask()); + const auto &actor_id = task_spec.ActorId(); + + const auto task_id = task_spec.TaskId(); + const auto num_returns = task_spec.NumReturns(); + + auto request = std::unique_ptr(new rpc::PushTaskRequest); + request->mutable_task_spec()->Swap(&task_spec.GetMutableMessage()); + + std::unique_lock guard(mutex_); + auto iter = actor_states_.find(actor_id); + if (iter == actor_states_.end() || + iter->second.state_ == ActorTableData::RECONSTRUCTING) { + // Actor is not yet created, or is being reconstructed, cache the request + // and submit after actor is alive. + // TODO(zhijunfu): it might be possible for a user to specify an invalid + // actor handle (e.g. from unpickling), in that case it might be desirable + // to have a timeout to mark it as invalid if it doesn't show up in the + // specified time. + pending_requests_[actor_id].emplace_back(std::move(request)); + return Status::OK(); + } else if (iter->second.state_ == ActorTableData::ALIVE) { + // Actor is alive, submit the request. + if (rpc_clients_.count(actor_id) == 0) { + // If rpc client is not available, then create it. + ConnectAndSendPendingTasks(actor_id, iter->second.location_.first, + iter->second.location_.second); + } + + // Submit request. + auto &client = rpc_clients_[actor_id]; + return PushTask(*client, *request, task_id, num_returns); + } else { + // Actor is dead, treat the task as failure. + RAY_CHECK(iter->second.state_ == ActorTableData::DEAD); + TreatTaskAsFailed(task_id, num_returns, rpc::ErrorType::ACTOR_DIED); + return Status::IOError("Actor is dead."); + } +} + +Status CoreWorkerDirectActorTaskSubmitter::SubscribeActorUpdates() { + // Register a callback to handle actor notifications. + auto actor_notification_callback = [this](const ActorID &actor_id, + const ActorTableData &actor_data) { + std::unique_lock guard(mutex_); + actor_states_.erase(actor_id); + actor_states_.emplace( + actor_id, + ActorStateData(actor_data.state(), actor_data.ip_address(), actor_data.port())); + + if (actor_data.state() == ActorTableData::ALIVE) { + // Check if this actor is the one that we're interested, if we already have + // a connection to the actor, or have pending requests for it, we should + // create a new connection. + if (pending_requests_.count(actor_id) > 0) { + ConnectAndSendPendingTasks(actor_id, actor_data.ip_address(), actor_data.port()); + } + } else { + // Remove rpc client if it's dead or being reconstructed. + rpc_clients_.erase(actor_id); + } + + RAY_LOG(INFO) << "received notification on actor, state=" + << static_cast(actor_data.state()) << ", actor_id: " << actor_id + << ", ip address: " << actor_data.ip_address() + << ", port: " << actor_data.port(); + }; + + return gcs_client_.Actors().AsyncSubscribe(actor_notification_callback, nullptr); +} + +void CoreWorkerDirectActorTaskSubmitter::ConnectAndSendPendingTasks( + const ActorID &actor_id, std::string ip_address, int port) { + std::unique_ptr grpc_client( + new rpc::DirectActorClient(ip_address, port, client_call_manager_)); + RAY_CHECK(rpc_clients_.emplace(actor_id, std::move(grpc_client)).second); + + // Submit all pending requests. + auto &client = rpc_clients_[actor_id]; + auto &requests = pending_requests_[actor_id]; + while (!requests.empty()) { + const auto &request = *requests.front(); + auto status = + PushTask(*client, request, TaskID::FromBinary(request.task_spec().task_id()), + request.task_spec().num_returns()); + requests.pop_front(); + } +} + +Status CoreWorkerDirectActorTaskSubmitter::PushTask(rpc::DirectActorClient &client, + const rpc::PushTaskRequest &request, + const TaskID &task_id, + int num_returns) { + auto status = client.PushTask( + request, + [this, task_id, num_returns](Status status, const rpc::PushTaskReply &reply) { + if (!status.ok()) { + TreatTaskAsFailed(task_id, num_returns, rpc::ErrorType::ACTOR_DIED); + return; + } + for (int i = 0; i < reply.return_objects_size(); i++) { + const auto &return_object = reply.return_objects(i); + ObjectID object_id = ObjectID::FromBinary(return_object.object_id()); + std::shared_ptr data_buffer; + if (return_object.data().size() > 0) { + data_buffer = std::make_shared( + const_cast( + reinterpret_cast(return_object.data().data())), + return_object.data().size()); + } + std::shared_ptr metadata_buffer; + if (return_object.metadata().size() > 0) { + metadata_buffer = std::make_shared( + const_cast( + reinterpret_cast(return_object.metadata().data())), + return_object.metadata().size()); + } + store_provider_->Put(RayObject(data_buffer, metadata_buffer), object_id); + } + }); + return status; +} + +void CoreWorkerDirectActorTaskSubmitter::TreatTaskAsFailed( + const TaskID &task_id, int num_returns, const rpc::ErrorType &error_type) { + for (int i = 0; i < num_returns; i++) { + const auto object_id = ObjectID::ForTaskReturn(task_id, i + 1); + 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()); + store_provider_->Put(RayObject(nullptr, meta_buffer), object_id); + } +} + +bool CoreWorkerDirectActorTaskSubmitter::IsActorAlive(const ActorID &actor_id) const { + std::unique_lock guard(mutex_); + auto iter = actor_states_.find(actor_id); + return (iter != actor_states_.end() && iter->second.state_ == ActorTableData::ALIVE); +} + +CoreWorkerDirectActorTaskReceiver::CoreWorkerDirectActorTaskReceiver( + CoreWorkerObjectInterface &object_interface, boost::asio::io_service &io_service, + rpc::GrpcServer &server, const TaskHandler &task_handler) + : object_interface_(object_interface), + task_service_(io_service, *this), + task_handler_(task_handler) { + server.RegisterService(task_service_); +} + +void CoreWorkerDirectActorTaskReceiver::HandlePushTask( + const rpc::PushTaskRequest &request, rpc::PushTaskReply *reply, + rpc::SendReplyCallback send_reply_callback) { + const TaskSpecification task_spec(request.task_spec()); + if (HasByReferenceArgs(task_spec)) { + send_reply_callback( + Status::Invalid("direct actor call only supports by value arguments"), nullptr, + nullptr); + return; + } + + auto num_returns = task_spec.NumReturns(); + RAY_CHECK(task_spec.IsActorCreationTask() || task_spec.IsActorTask()); + RAY_CHECK(num_returns > 0); + // Decrease to account for the dummy object id. + num_returns--; + + std::vector> results; + auto status = task_handler_(task_spec, &results); + RAY_CHECK(results.size() == num_returns) << results.size() << " " << num_returns; + + for (int i = 0; i < results.size(); i++) { + auto return_object = (*reply).add_return_objects(); + ObjectID id = ObjectID::ForTaskReturn(task_spec.TaskId(), i + 1); + return_object->set_object_id(id.Binary()); + const auto &result = results[i]; + if (result->GetData() != nullptr) { + return_object->set_data(result->GetData()->Data(), result->GetData()->Size()); + } + if (result->GetMetadata() != nullptr) { + return_object->set_metadata(result->GetMetadata()->Data(), + result->GetMetadata()->Size()); + } + } + + send_reply_callback(status, nullptr, nullptr); +} + +} // namespace ray diff --git a/src/ray/core_worker/transport/direct_actor_transport.h b/src/ray/core_worker/transport/direct_actor_transport.h new file mode 100644 index 000000000..90a256900 --- /dev/null +++ b/src/ray/core_worker/transport/direct_actor_transport.h @@ -0,0 +1,146 @@ +#ifndef RAY_CORE_WORKER_DIRECT_ACTOR_TRANSPORT_H +#define RAY_CORE_WORKER_DIRECT_ACTOR_TRANSPORT_H + +#include + +#include "ray/core_worker/object_interface.h" +#include "ray/core_worker/transport/transport.h" +#include "ray/gcs/redis_gcs_client.h" +#include "ray/rpc/worker/direct_actor_client.h" +#include "ray/rpc/worker/direct_actor_server.h" + +namespace ray { + +/// In direct actor call task submitter and receiver, a task is directly submitted +/// to the actor that will execute it. + +/// The state data for an actor. +struct ActorStateData { + ActorStateData(gcs::ActorTableData::ActorState state, const std::string &ip, int port) + : state_(state), location_(std::make_pair(ip, port)) {} + + /// Actor's state (e.g. alive, dead, reconstrucing). + gcs::ActorTableData::ActorState state_; + + /// IP address and port that the actor is listening on. + std::pair location_; +}; + +class CoreWorkerDirectActorTaskSubmitter : public CoreWorkerTaskSubmitter { + public: + CoreWorkerDirectActorTaskSubmitter(boost::asio::io_service &io_service, + gcs::RedisGcsClient &gcs_client, + CoreWorkerObjectInterface &object_interface); + + /// Submit a task to an actor for execution. + /// + /// \param[in] task The task spec to submit. + /// \return Status. + Status SubmitTask(const TaskSpecification &task_spec) override; + + private: + /// Subscribe to all actor updates. + Status SubscribeActorUpdates(); + + /// Helper function to push a task to an actor. + /// + /// \param[in] client The RPC client to send tasks to an actor. + /// \param[in] request The request to send. + /// \param[in] task_id The ID of a task. + /// \param[in] num_returns Number of return objects. + /// \return Status. + Status PushTask(rpc::DirectActorClient &client, const rpc::PushTaskRequest &request, + const TaskID &task_id, int num_returns); + + /// Treat a task as failed. + /// + /// \param[in] task_id The ID of a task. + /// \param[in] num_returns Number of return objects. + /// \param[in] error_type The type of the specific error. + /// \return Void. + void TreatTaskAsFailed(const TaskID &task_id, int num_returns, + const rpc::ErrorType &error_type); + + /// Create connection to actor and send all pending tasks. + /// Note that this function doesn't take lock, the caller is expected to hold + /// `mutex_` before calling this function. + /// + /// \param[in] actor_id Actor ID. + /// \param[in] ip_address The ip address of the node that the actor is running on. + /// \param[in] port The port that the actor is listening on. + /// \return Void. + void ConnectAndSendPendingTasks(const ActorID &actor_id, std::string ip_address, + int port); + + /// Whether the specified actor is alive. + /// + /// \param[in] actor_id The actor ID. + /// \return Whether this actor is alive. + bool IsActorAlive(const ActorID &actor_id) const; + + /// The IO event loop. + boost::asio::io_service &io_service_; + + /// Gcs client. + gcs::RedisGcsClient &gcs_client_; + + /// The `ClientCallManager` object that is shared by all `DirectActorClient`s. + rpc::ClientCallManager client_call_manager_; + + /// Mutex to proect the various maps below. + mutable std::mutex mutex_; + + /// Map from actor id to actor state. This currently includes all actors in the system. + /// + /// TODO(zhijunfu): this map currently keeps track of all the actors in the system, + /// like `actor_registry_` in raylet. Later after new GCS client interface supports + /// subscribing updates for a specific actor, this will be updated to only include + /// entries for actors that the transport submits tasks to. + std::unordered_map actor_states_; + + /// Map from actor id to rpc client. This only includes actors that we send tasks to. + /// + /// TODO(zhijunfu): this will be moved into `actor_states_` later when we can + /// subscribe updates for a specific actor. + std::unordered_map> rpc_clients_; + + /// Map from actor id to the actor's pending requests. + std::unordered_map>> + pending_requests_; + + /// The store provider. + std::unique_ptr store_provider_; + + friend class CoreWorkerTest; +}; + +class CoreWorkerDirectActorTaskReceiver : public CoreWorkerTaskReceiver, + public rpc::DirectActorHandler { + public: + CoreWorkerDirectActorTaskReceiver(CoreWorkerObjectInterface &object_interface, + boost::asio::io_service &io_service, + rpc::GrpcServer &server, + const TaskHandler &task_handler); + + /// Handle a `PushTask` request. + /// The implementation can handle this request asynchronously. When hanling is done, the + /// `done_callback` should be called. + /// + /// \param[in] request The request message. + /// \param[out] reply The reply message. + /// \param[in] done_callback The callback to be called when the request is done. + void HandlePushTask(const rpc::PushTaskRequest &request, rpc::PushTaskReply *reply, + rpc::SendReplyCallback send_reply_callback) override; + + private: + // Object interface. + CoreWorkerObjectInterface &object_interface_; + /// The rpc service for `DirectActorService`. + rpc::DirectActorGrpcService task_service_; + /// The callback function to process a task. + TaskHandler task_handler_; +}; + +} // namespace ray + +#endif // RAY_CORE_WORKER_DIRECT_ACTOR_TRANSPORT_H diff --git a/src/ray/core_worker/transport/transport.h b/src/ray/core_worker/transport/transport.h index 4fd3e7763..34353841f 100644 --- a/src/ray/core_worker/transport/transport.h +++ b/src/ray/core_worker/transport/transport.h @@ -28,6 +28,8 @@ class CoreWorkerTaskSubmitter { /// \param[in] task The task spec to submit. /// \return Status. virtual Status SubmitTask(const TaskSpecification &task_spec) = 0; + + virtual ~CoreWorkerTaskSubmitter() {} }; /// This class receives tasks for execution. @@ -36,6 +38,8 @@ class CoreWorkerTaskReceiver { using TaskHandler = std::function> *results)>; + + virtual ~CoreWorkerTaskReceiver() {} }; } // namespace ray diff --git a/src/ray/gcs/gcs_client_interface.h b/src/ray/gcs/gcs_client_interface.h index 69c6fa814..f62ac9f39 100644 --- a/src/ray/gcs/gcs_client_interface.h +++ b/src/ray/gcs/gcs_client_interface.h @@ -69,7 +69,7 @@ class GcsClientOptions { /// Before exit, `Disconnect()` must be called. class GcsClientInterface : public std::enable_shared_from_this { public: - virtual ~GcsClientInterface() { RAY_CHECK(!is_connected_); } + virtual ~GcsClientInterface() {} /// Connect to GCS Service. Non-thread safe. /// This function must be called before calling other functions. diff --git a/src/ray/protobuf/core_worker.proto b/src/ray/protobuf/core_worker.proto index 0b238f456..77c807fbb 100644 --- a/src/ray/protobuf/core_worker.proto +++ b/src/ray/protobuf/core_worker.proto @@ -27,4 +27,7 @@ message ActorHandle { // The number of times that this actor handle has been forked. // It's used to make sure ids of actor handles are unique. int64 num_forks = 7; + + // Whether direct actor call is used. + bool is_direct_call = 8; } diff --git a/src/ray/protobuf/direct_actor.proto b/src/ray/protobuf/direct_actor.proto new file mode 100644 index 000000000..373af60f1 --- /dev/null +++ b/src/ray/protobuf/direct_actor.proto @@ -0,0 +1,30 @@ +syntax = "proto3"; + +package ray.rpc; + +import "src/ray/protobuf/common.proto"; + +message ReturnObject { + // Object ID. + bytes object_id = 1; + // Data of the object. + bytes data = 2; + // Metaata of the object. + bytes metadata = 3; +} + +message PushTaskRequest { + // The task to be pushed. + TaskSpec task_spec = 1; +} + +message PushTaskReply { + // The returned objects. + repeated ReturnObject return_objects = 1; +} + +// Service for direct actor. +service DirectActorService { + // Push a task to a worker. + rpc PushTask(PushTaskRequest) returns (PushTaskReply); +} diff --git a/src/ray/protobuf/gcs.proto b/src/ray/protobuf/gcs.proto index ab59c1b0d..218e1acd5 100644 --- a/src/ray/protobuf/gcs.proto +++ b/src/ray/protobuf/gcs.proto @@ -105,6 +105,10 @@ message ActorTableData { uint64 max_reconstructions = 7; // Remaining number of reconstructions. uint64 remaining_reconstructions = 8; + // The IP address of the node that the actor is running on. + string ip_address = 9; + // The port that the actor is listening on. + int32 port = 10; } message ErrorTableData { diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index 93dd2ceec..20efb4985 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -1826,7 +1826,7 @@ void NodeManager::FinishAssignedTask(Worker &worker) { } std::shared_ptr NodeManager::CreateActorTableDataFromCreationTask( - const TaskSpecification &task_spec) { + const TaskSpecification &task_spec, int port) { RAY_CHECK(task_spec.IsActorCreationTask()); auto actor_id = task_spec.ActorCreationId(); auto actor_entry = actor_registry_.find(actor_id); @@ -1859,6 +1859,11 @@ std::shared_ptr NodeManager::CreateActorTableDataFromCreationTas actor_info_ptr->remaining_reconstructions() - 1); } + // Set the ip address & port, which could change after reconstruction. + actor_info_ptr->set_ip_address( + gcs_client_->client_table().GetLocalClient().node_manager_address()); + actor_info_ptr->set_port(port); + // Set the new fields for the actor's state to indicate that the actor is // now alive on this node manager. actor_info_ptr->set_node_manager_id( @@ -1889,10 +1894,11 @@ void NodeManager::FinishAssignedActorTask(Worker &worker, const Task &task) { // Lookup the parent actor id. auto parent_task_id = task_spec.ParentTaskId(); RAY_CHECK(actor_handle_id.IsNil()); + int port = worker.Port(); RAY_CHECK_OK(gcs_client_->raylet_task_table().Lookup( JobID::Nil(), parent_task_id, /*success_callback=*/ - [this, task_spec, resumed_from_checkpoint]( + [this, task_spec, resumed_from_checkpoint, port]( ray::gcs::RedisGcsClient *client, const TaskID &parent_task_id, const TaskTableData &parent_task_data) { // The task was in the GCS task table. Use the stored task spec to @@ -1905,11 +1911,11 @@ void NodeManager::FinishAssignedActorTask(Worker &worker, const Task &task) { parent_actor_id = parent_task.GetTaskSpecification().ActorId(); } FinishAssignedActorCreationTask(parent_actor_id, task_spec, - resumed_from_checkpoint); + resumed_from_checkpoint, port); }, /*failure_callback=*/ - [this, task_spec, resumed_from_checkpoint](ray::gcs::RedisGcsClient *client, - const TaskID &parent_task_id) { + [this, task_spec, resumed_from_checkpoint, port](ray::gcs::RedisGcsClient *client, + const TaskID &parent_task_id) { // The parent task was not in the GCS task table. It should most likely be in // the // lineage cache. @@ -1932,7 +1938,7 @@ void NodeManager::FinishAssignedActorTask(Worker &worker, const Task &task) { << "ray.init(redis_max_memory=)."; } FinishAssignedActorCreationTask(parent_actor_id, task_spec, - resumed_from_checkpoint); + resumed_from_checkpoint, port); })); } else { // The actor was not resumed from a checkpoint. We extend the actor's @@ -1967,10 +1973,11 @@ void NodeManager::ExtendActorFrontier(const ObjectID &dummy_object, void NodeManager::FinishAssignedActorCreationTask(const ActorID &parent_actor_id, const TaskSpecification &task_spec, - bool resumed_from_checkpoint) { + bool resumed_from_checkpoint, + int port) { // Notify the other node managers that the actor has been created. const ActorID actor_id = task_spec.ActorCreationId(); - auto new_actor_info = CreateActorTableDataFromCreationTask(task_spec); + auto new_actor_info = CreateActorTableDataFromCreationTask(task_spec, port); new_actor_info->set_parent_actor_id(parent_actor_id.Binary()); auto update_callback = [actor_id](Status status) { if (!status.ok()) { diff --git a/src/ray/raylet/node_manager.h b/src/ray/raylet/node_manager.h index acafc256b..224a701ba 100644 --- a/src/ray/raylet/node_manager.h +++ b/src/ray/raylet/node_manager.h @@ -287,8 +287,9 @@ class NodeManager : public rpc::NodeManagerServiceHandler, /// /// \param task_spec Task specification of the actor creation task that created the /// actor. + /// \param worker The port that the actor is listening on. std::shared_ptr CreateActorTableDataFromCreationTask( - const TaskSpecification &task_spec); + const TaskSpecification &task_spec, int port); /// Handle a worker finishing an assigned actor task or actor creation task. /// \param worker The worker that finished the task. /// \param task The actor task or actor creation task. @@ -302,10 +303,11 @@ class NodeManager : public rpc::NodeManagerServiceHandler, /// \param task_spec Task specification of the actor creation task that created the /// actor. /// \param resumed_from_checkpoint If the actor was resumed from a checkpoint. + /// \param port Rpc server port that the actor is listening on. /// \return Void. void FinishAssignedActorCreationTask(const ActorID &parent_actor_id, const TaskSpecification &task_spec, - bool resumed_from_checkpoint); + bool resumed_from_checkpoint, int port); /// Extend actor frontier after an actor task or actor creation task executes. /// /// \param dummy_object Dummy object corresponding to the task. diff --git a/src/ray/rpc/worker/direct_actor_client.h b/src/ray/rpc/worker/direct_actor_client.h new file mode 100644 index 000000000..51e74ffd6 --- /dev/null +++ b/src/ray/rpc/worker/direct_actor_client.h @@ -0,0 +1,58 @@ +#ifndef RAY_RPC_DIRECT_ACTOR_CLIENT_H +#define RAY_RPC_DIRECT_ACTOR_CLIENT_H + +#include + +#include + +#include "ray/common/status.h" +#include "ray/rpc/client_call.h" +#include "ray/util/logging.h" +#include "src/ray/protobuf/direct_actor.grpc.pb.h" +#include "src/ray/protobuf/direct_actor.pb.h" + +namespace ray { +namespace rpc { + +/// Client used for communicating with a direct actor server. +class DirectActorClient { + public: + /// Constructor. + /// + /// \param[in] address Address of the direct actor server. + /// \param[in] port Port of the direct actor server. + /// \param[in] client_call_manager The `ClientCallManager` used for managing requests. + DirectActorClient(const std::string &address, const int port, + ClientCallManager &client_call_manager) + : client_call_manager_(client_call_manager) { + std::shared_ptr channel = grpc::CreateChannel( + address + ":" + std::to_string(port), grpc::InsecureChannelCredentials()); + stub_ = DirectActorService::NewStub(channel); + }; + + /// Push a task. + /// + /// \param[in] request The request message. + /// \param[in] callback The callback function that handles reply. + /// \return if the rpc call succeeds + ray::Status PushTask(const PushTaskRequest &request, + const ClientCallback &callback) { + auto call = client_call_manager_ + .CreateCall( + *stub_, &DirectActorService::Stub::PrepareAsyncPushTask, request, + callback); + return call->GetStatus(); + } + + private: + /// The gRPC-generated stub. + std::unique_ptr stub_; + + /// The `ClientCallManager` used for managing requests. + ClientCallManager &client_call_manager_; +}; + +} // namespace rpc +} // namespace ray + +#endif // RAY_RPC_DIRECT_ACTOR_CLIENT_H diff --git a/src/ray/rpc/worker/direct_actor_server.h b/src/ray/rpc/worker/direct_actor_server.h new file mode 100644 index 000000000..a0670aa2a --- /dev/null +++ b/src/ray/rpc/worker/direct_actor_server.h @@ -0,0 +1,68 @@ +#ifndef RAY_RPC_DIRECT_ACTOR_SERVER_H +#define RAY_RPC_DIRECT_ACTOR_SERVER_H + +#include "ray/rpc/grpc_server.h" +#include "ray/rpc/server_call.h" + +#include "src/ray/protobuf/direct_actor.grpc.pb.h" +#include "src/ray/protobuf/direct_actor.pb.h" + +namespace ray { +namespace rpc { + +/// Interface of the `DirectActorService`, see `src/ray/protobuf/direct_actor.proto`. +class DirectActorHandler { + public: + /// Handle a `PushTask` request. + /// The implementation can handle this request asynchronously. When hanling is done, the + /// `done_callback` should be called. + /// + /// \param[in] request The request message. + /// \param[out] reply The reply message. + /// \param[in] done_callback The callback to be called when the request is done. + virtual void HandlePushTask(const PushTaskRequest &request, PushTaskReply *reply, + SendReplyCallback send_reply_callback) = 0; +}; + +/// The `GrpcServer` for `WorkerService`. +class DirectActorGrpcService : public GrpcService { + public: + /// Constructor. + /// + /// \param[in] main_service See super class. + /// \param[in] handler The service handler that actually handle the requests. + DirectActorGrpcService(boost::asio::io_service &main_service, + DirectActorHandler &service_handler) + : GrpcService(main_service), service_handler_(service_handler){}; + + protected: + grpc::Service &GetGrpcService() override { return service_; } + + void InitServerCallFactories( + const std::unique_ptr &cq, + std::vector, int>> + *server_call_factories_and_concurrencies) override { + // Initialize the Factory for `PushTask` requests. + std::unique_ptr push_task_call_Factory( + new ServerCallFactoryImpl( + service_, &DirectActorService::AsyncService::RequestPushTask, + service_handler_, &DirectActorHandler::HandlePushTask, cq, main_service_)); + + // Set `PushTask`'s accept concurrency to 100. + server_call_factories_and_concurrencies->emplace_back( + std::move(push_task_call_Factory), 100); + } + + private: + /// The grpc async service object. + DirectActorService::AsyncService service_; + + /// The service handler that actually handle the requests. + DirectActorHandler &service_handler_; +}; + +} // namespace rpc +} // namespace ray + +#endif diff --git a/src/ray/util/test_util.h b/src/ray/util/test_util.h new file mode 100644 index 000000000..6221b774b --- /dev/null +++ b/src/ray/util/test_util.h @@ -0,0 +1,33 @@ +#ifndef RAY_UTIL_TEST_UTIL_H +#define RAY_UTIL_TEST_UTIL_H + +#include + +namespace ray { + +/// Wait until the condition is met, or timeout is reached. +/// +/// \param[in] condition The condition to wait for. +/// \param[in] timeout_ms Timeout in milliseconds to wait for for. +/// \return Whether the condition is met. +bool WaitForCondition(std::function condition, int timeout_ms) { + int wait_time = 0; + while (true) { + if (condition()) { + return true; + } + + // sleep 100ms. + const int wait_interval_ms = 100; + usleep(wait_interval_ms * 1000); + wait_time += wait_interval_ms; + if (wait_time > timeout_ms) { + break; + } + } + return false; +} + +} // namespace ray + +#endif // RAY_UTIL_TEST_UTIL_H