From 5d47d02f8197f7bff743d72980da2ec354fc6f85 Mon Sep 17 00:00:00 2001 From: Tao Wang Date: Wed, 25 Nov 2020 11:36:44 +0800 Subject: [PATCH] [GCS]add callback for RegisterSelf api, make it done first (#12252) --- src/ray/gcs/accessor.h | 6 ++- .../gcs/gcs_client/service_based_accessor.cc | 31 +++++++------- .../gcs/gcs_client/service_based_accessor.h | 3 +- .../test/service_based_gcs_client_test.cc | 2 +- .../gcs_server/test/gcs_server_test_util.h | 3 +- src/ray/gcs/redis_accessor.cc | 9 +++- src/ray/gcs/redis_accessor.h | 3 +- .../test/object_manager_stress_test.cc | 2 +- .../test/object_manager_test.cc | 2 +- src/ray/raylet/raylet.cc | 42 ++++++++++--------- 10 files changed, 58 insertions(+), 45 deletions(-) diff --git a/src/ray/gcs/accessor.h b/src/ray/gcs/accessor.h index 15511012f..484318df4 100644 --- a/src/ray/gcs/accessor.h +++ b/src/ray/gcs/accessor.h @@ -435,11 +435,13 @@ class NodeInfoAccessor { public: virtual ~NodeInfoAccessor() = default; - /// Register local node to GCS synchronously. + /// Register local node to GCS asynchronously. /// /// \param node_info The information of node to register to GCS. + /// \param callback Callback that will be called when registration is complete. /// \return Status - virtual Status RegisterSelf(const rpc::GcsNodeInfo &local_node_info) = 0; + virtual Status RegisterSelf(const rpc::GcsNodeInfo &local_node_info, + const StatusCallback &callback) = 0; /// Cancel registration of local node to GCS synchronously. /// diff --git a/src/ray/gcs/gcs_client/service_based_accessor.cc b/src/ray/gcs/gcs_client/service_based_accessor.cc index 30fc67a7d..3f1c75058 100644 --- a/src/ray/gcs/gcs_client/service_based_accessor.cc +++ b/src/ray/gcs/gcs_client/service_based_accessor.cc @@ -418,7 +418,8 @@ ServiceBasedNodeInfoAccessor::ServiceBasedNodeInfoAccessor( ServiceBasedGcsClient *client_impl) : client_impl_(client_impl) {} -Status ServiceBasedNodeInfoAccessor::RegisterSelf(const GcsNodeInfo &local_node_info) { +Status ServiceBasedNodeInfoAccessor::RegisterSelf(const GcsNodeInfo &local_node_info, + const StatusCallback &callback) { auto node_id = NodeID::FromBinary(local_node_info.node_id()); RAY_LOG(DEBUG) << "Registering node info, node id = " << node_id << ", address is = " << local_node_info.node_manager_address(); @@ -427,22 +428,20 @@ Status ServiceBasedNodeInfoAccessor::RegisterSelf(const GcsNodeInfo &local_node_ rpc::RegisterNodeRequest request; request.mutable_node_info()->CopyFrom(local_node_info); - auto operation = [this, request, local_node_info, - node_id](const SequencerDoneCallback &done_callback) { - client_impl_->GetGcsRpcClient().RegisterNode( - request, [this, node_id, local_node_info, done_callback]( - const Status &status, const rpc::RegisterNodeReply &reply) { - if (status.ok()) { - local_node_info_.CopyFrom(local_node_info); - local_node_id_ = NodeID::FromBinary(local_node_info.node_id()); - } - RAY_LOG(DEBUG) << "Finished registering node info, status = " << status - << ", node id = " << node_id; - done_callback(); - }); - }; + client_impl_->GetGcsRpcClient().RegisterNode( + request, [this, node_id, local_node_info, callback]( + const Status &status, const rpc::RegisterNodeReply &reply) { + if (status.ok()) { + local_node_info_.CopyFrom(local_node_info); + local_node_id_ = NodeID::FromBinary(local_node_info.node_id()); + } + if (callback) { + callback(status); + } + RAY_LOG(DEBUG) << "Finished registering node info, status = " << status + << ", node id = " << node_id; + }); - sequencer_.Post(node_id, operation); return Status::OK(); } diff --git a/src/ray/gcs/gcs_client/service_based_accessor.h b/src/ray/gcs/gcs_client/service_based_accessor.h index 180f8f46b..e9840c536 100644 --- a/src/ray/gcs/gcs_client/service_based_accessor.h +++ b/src/ray/gcs/gcs_client/service_based_accessor.h @@ -147,7 +147,8 @@ class ServiceBasedNodeInfoAccessor : public NodeInfoAccessor { virtual ~ServiceBasedNodeInfoAccessor() = default; - Status RegisterSelf(const GcsNodeInfo &local_node_info) override; + Status RegisterSelf(const GcsNodeInfo &local_node_info, + const StatusCallback &callback) override; Status UnregisterSelf() override; diff --git a/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc b/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc index ab8c561ec..0700f60b5 100644 --- a/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc +++ b/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc @@ -288,7 +288,7 @@ class ServiceBasedGcsClientTest : public ::testing::Test { } bool RegisterSelf(const rpc::GcsNodeInfo &local_node_info) { - Status status = gcs_client_->Nodes().RegisterSelf(local_node_info); + Status status = gcs_client_->Nodes().RegisterSelf(local_node_info, nullptr); return status.ok(); } diff --git a/src/ray/gcs/gcs_server/test/gcs_server_test_util.h b/src/ray/gcs/gcs_server/test/gcs_server_test_util.h index 24b71c4ee..4e482b55a 100644 --- a/src/ray/gcs/gcs_server/test/gcs_server_test_util.h +++ b/src/ray/gcs/gcs_server/test/gcs_server_test_util.h @@ -305,7 +305,8 @@ struct GcsServerMocker { class MockedNodeInfoAccessor : public gcs::NodeInfoAccessor { public: - Status RegisterSelf(const rpc::GcsNodeInfo &local_node_info) override { + Status RegisterSelf(const rpc::GcsNodeInfo &local_node_info, + const gcs::StatusCallback &callback) override { return Status::NotImplemented(""); } diff --git a/src/ray/gcs/redis_accessor.cc b/src/ray/gcs/redis_accessor.cc index 06f001342..68c69e91c 100644 --- a/src/ray/gcs/redis_accessor.cc +++ b/src/ray/gcs/redis_accessor.cc @@ -497,9 +497,14 @@ RedisNodeInfoAccessor::RedisNodeInfoAccessor(RedisGcsClient *client_impl) resource_sub_executor_(client_impl_->resource_table()), heartbeat_batch_sub_executor_(client_impl->heartbeat_batch_table()) {} -Status RedisNodeInfoAccessor::RegisterSelf(const GcsNodeInfo &local_node_info) { +Status RedisNodeInfoAccessor::RegisterSelf(const GcsNodeInfo &local_node_info, + const StatusCallback &callback) { ClientTable &client_table = client_impl_->client_table(); - return client_table.Connect(local_node_info); + Status status = client_table.Connect(local_node_info); + if (callback != nullptr) { + callback(Status::OK()); + } + return status; } Status RedisNodeInfoAccessor::UnregisterSelf() { diff --git a/src/ray/gcs/redis_accessor.h b/src/ray/gcs/redis_accessor.h index 028ce015d..4e851e8a2 100644 --- a/src/ray/gcs/redis_accessor.h +++ b/src/ray/gcs/redis_accessor.h @@ -321,7 +321,8 @@ class RedisNodeInfoAccessor : public NodeInfoAccessor { virtual ~RedisNodeInfoAccessor() {} - Status RegisterSelf(const GcsNodeInfo &local_node_info) override; + Status RegisterSelf(const GcsNodeInfo &local_node_info, + const StatusCallback &callback) override; Status UnregisterSelf() override; diff --git a/src/ray/object_manager/test/object_manager_stress_test.cc b/src/ray/object_manager/test/object_manager_stress_test.cc index 14e02112b..c0e163f7c 100644 --- a/src/ray/object_manager/test/object_manager_stress_test.cc +++ b/src/ray/object_manager/test/object_manager_stress_test.cc @@ -70,7 +70,7 @@ class MockServer { node_info.set_node_manager_port(object_manager_port); node_info.set_object_manager_port(object_manager_port); - ray::Status status = gcs_client_->Nodes().RegisterSelf(node_info); + ray::Status status = gcs_client_->Nodes().RegisterSelf(node_info, nullptr); return status; } diff --git a/src/ray/object_manager/test/object_manager_test.cc b/src/ray/object_manager/test/object_manager_test.cc index 6779f7936..da8a80775 100644 --- a/src/ray/object_manager/test/object_manager_test.cc +++ b/src/ray/object_manager/test/object_manager_test.cc @@ -66,7 +66,7 @@ class MockServer { node_info.set_node_manager_port(object_manager_port); node_info.set_object_manager_port(object_manager_port); - ray::Status status = gcs_client_->Nodes().RegisterSelf(node_info); + ray::Status status = gcs_client_->Nodes().RegisterSelf(node_info, nullptr); return status; } diff --git a/src/ray/raylet/raylet.cc b/src/ray/raylet/raylet.cc index ab81ed48b..c7623f420 100644 --- a/src/ray/raylet/raylet.cc +++ b/src/ray/raylet/raylet.cc @@ -113,28 +113,32 @@ void Raylet::Stop() { } ray::Status Raylet::RegisterGcs() { - RAY_RETURN_NOT_OK(gcs_client_->Nodes().RegisterSelf(self_node_info_)); + auto register_callback = [this](const Status &status) { + RAY_CHECK_OK(status); + RAY_LOG(DEBUG) << "Node manager " << self_node_id_ << " started on " + << self_node_info_.node_manager_address() << ":" + << self_node_info_.node_manager_port() << " object manager at " + << self_node_info_.node_manager_address() << ":" + << self_node_info_.object_manager_port() << ", hostname " + << self_node_info_.node_manager_hostname(); - RAY_LOG(DEBUG) << "Node manager " << self_node_id_ << " started on " - << self_node_info_.node_manager_address() << ":" - << self_node_info_.node_manager_port() << " object manager at " - << self_node_info_.node_manager_address() << ":" - << self_node_info_.object_manager_port() << ", hostname " - << self_node_info_.node_manager_hostname(); + // Add resource information. + const NodeManagerConfig &node_manager_config = node_manager_.GetInitialConfig(); + std::unordered_map> resources; + for (const auto &resource_pair : + node_manager_config.resource_config.GetResourceMap()) { + auto resource = std::make_shared(); + resource->set_resource_capacity(resource_pair.second); + resources.emplace(resource_pair.first, resource); + } + RAY_CHECK_OK( + gcs_client_->Nodes().AsyncUpdateResources(self_node_id_, resources, nullptr)); + + RAY_CHECK_OK(node_manager_.RegisterGcs()); + }; - // Add resource information. - const NodeManagerConfig &node_manager_config = node_manager_.GetInitialConfig(); - std::unordered_map> resources; - for (const auto &resource_pair : node_manager_config.resource_config.GetResourceMap()) { - auto resource = std::make_shared(); - resource->set_resource_capacity(resource_pair.second); - resources.emplace(resource_pair.first, resource); - } RAY_RETURN_NOT_OK( - gcs_client_->Nodes().AsyncUpdateResources(self_node_id_, resources, nullptr)); - - RAY_RETURN_NOT_OK(node_manager_.RegisterGcs()); - + gcs_client_->Nodes().RegisterSelf(self_node_info_, register_callback)); return Status::OK(); }