From 970cd78701786b5d3314370293c8517de8c57e4d Mon Sep 17 00:00:00 2001 From: micafan <550435771@qq.com> Date: Fri, 3 Jan 2020 14:07:37 +0800 Subject: [PATCH] [GCS] refactor the GCS Client Dynamic Resource Interface (#6266) --- BUILD.bazel | 16 + src/ray/gcs/accessor.h | 39 +++ src/ray/gcs/entry_change_notification.h | 17 +- src/ray/gcs/redis_accessor.cc | 52 ++++ src/ray/gcs/redis_accessor.h | 18 ++ src/ray/gcs/redis_gcs_client.h | 5 +- src/ray/gcs/subscription_executor.cc | 2 + src/ray/gcs/tables.cc | 10 +- src/ray/gcs/tables.h | 4 +- src/ray/gcs/test/redis_gcs_client_test.cc | 290 +++++++++--------- .../gcs/test/redis_node_info_accessor_test.cc | 165 ++++++++++ src/ray/raylet/node_manager.cc | 62 ++-- src/ray/raylet/raylet.cc | 4 +- 13 files changed, 497 insertions(+), 187 deletions(-) create mode 100644 src/ray/gcs/test/redis_node_info_accessor_test.cc diff --git a/BUILD.bazel b/BUILD.bazel index 021bfaccd..a1f6cfd54 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -891,6 +891,22 @@ cc_test( ], ) +cc_test( + name = "redis_node_info_accessor_test", + srcs = ["src/ray/gcs/test/redis_node_info_accessor_test.cc"], + args = ["$(location redis-server) $(location redis-cli) $(location libray_redis_module.so)"], + copts = COPTS, + data = [ + "//:libray_redis_module.so", + "//:redis-cli", + "//:redis-server", + ], + deps = [ + ":gcs", + "@com_google_googletest//:gtest_main", + ], +) + cc_test( name = "asio_test", srcs = ["src/ray/gcs/test/asio_test.cc"], diff --git a/src/ray/gcs/accessor.h b/src/ray/gcs/accessor.h index c3003d117..2f0b0251a 100644 --- a/src/ray/gcs/accessor.h +++ b/src/ray/gcs/accessor.h @@ -352,6 +352,45 @@ class NodeInfoAccessor { /// \return Whether the node is removed. virtual bool IsRemoved(const ClientID &node_id) const = 0; + // TODO(micafan) Define ResourceMap in GCS proto. + typedef std::unordered_map> + ResourceMap; + + /// Get node's resources from GCS asynchronously. + /// + /// \param node_id The ID of node to lookup dynamic resources. + /// \param callback Callback that will be called after lookup finishes. + /// \return Status + virtual Status AsyncGetResources(const ClientID &node_id, + const OptionalItemCallback &callback) = 0; + + /// Update resources of node in GCS asynchronously. + /// + /// \param node_id The ID of node to update dynamic resources. + /// \param resources The dynamic resources of node to be updated. + /// \param callback Callback that will be called after update finishes. + virtual Status AsyncUpdateResources(const ClientID &node_id, + const ResourceMap &resources, + const StatusCallback &callback) = 0; + + /// Delete resources of a node from GCS asynchronously. + /// + /// \param node_id The ID of node to delete resources from GCS. + /// \param resource_names The names of resource to be deleted. + /// \param callback Callback that will be called after delete finishes. + virtual Status AsyncDeleteResources(const ClientID &node_id, + const std::vector &resource_names, + const StatusCallback &callback) = 0; + + /// Subscribe to node resource changes. + /// + /// \param subscribe Callback that will be called when any resource is updated. + /// \param done Callback that will be called when subscription is complete. + /// \return Status + virtual Status AsyncSubscribeToResources( + const SubscribeCallback &subscribe, + const StatusCallback &done) = 0; + /// Report heartbeat of a node to GCS asynchronously. /// /// \param data_ptr The heartbeat that will be reported to GCS. diff --git a/src/ray/gcs/entry_change_notification.h b/src/ray/gcs/entry_change_notification.h index 8e1bb4336..f14ea91b3 100644 --- a/src/ray/gcs/entry_change_notification.h +++ b/src/ray/gcs/entry_change_notification.h @@ -14,7 +14,7 @@ namespace gcs { template class EntryChangeNotification { public: - EntryChangeNotification(rpc::GcsChangeMode change_mode, std::vector data) + EntryChangeNotification(rpc::GcsChangeMode change_mode, Data data) : change_mode_(change_mode), data_(std::move(data)) {} EntryChangeNotification(EntryChangeNotification &&other) { @@ -41,14 +41,23 @@ class EntryChangeNotification { /// Get data of this notification. /// /// \return Data - const std::vector &GetData() const { return data_; } + const Data &GetData() const { return data_; } private: rpc::GcsChangeMode change_mode_; - std::vector data_; + Data data_; }; -typedef EntryChangeNotification ObjectChangeNotification; +template +using ArrayNotification = EntryChangeNotification>; + +typedef ArrayNotification ObjectChangeNotification; + +template +using MapNotification = + EntryChangeNotification>>; + +typedef MapNotification ResourceChangeNotification; } // namespace gcs diff --git a/src/ray/gcs/redis_accessor.cc b/src/ray/gcs/redis_accessor.cc index 1902ee1fa..a84da3923 100644 --- a/src/ray/gcs/redis_accessor.cc +++ b/src/ray/gcs/redis_accessor.cc @@ -367,6 +367,7 @@ Status RedisObjectInfoAccessor::AsyncUnsubscribeToLocations(const ObjectID &obje RedisNodeInfoAccessor::RedisNodeInfoAccessor(RedisGcsClient *client_impl) : client_impl_(client_impl), + resource_sub_executor_(client_impl_->resource_table()), heartbeat_sub_executor_(client_impl->heartbeat_table()), heartbeat_batch_sub_executor_(client_impl->heartbeat_batch_table()) {} @@ -502,6 +503,57 @@ Status RedisNodeInfoAccessor::AsyncSubscribeBatchHeartbeat( done); } +Status RedisNodeInfoAccessor::AsyncGetResources( + const ClientID &node_id, const OptionalItemCallback &callback) { + RAY_CHECK(callback != nullptr); + auto on_done = [callback](RedisGcsClient *client, const ClientID &id, + const ResourceMap &data) { + boost::optional result; + if (!data.empty()) { + result = data; + } + callback(Status::OK(), result); + }; + + DynamicResourceTable &resource_table = client_impl_->resource_table(); + return resource_table.Lookup(JobID::Nil(), node_id, on_done); +} + +Status RedisNodeInfoAccessor::AsyncUpdateResources(const ClientID &node_id, + const ResourceMap &resources, + const StatusCallback &callback) { + Hash::HashCallback on_done = nullptr; + if (callback != nullptr) { + on_done = [callback](RedisGcsClient *client, const ClientID &node_id, + const ResourceMap &resources) { callback(Status::OK()); }; + } + + DynamicResourceTable &resource_table = client_impl_->resource_table(); + return resource_table.Update(JobID::Nil(), node_id, resources, on_done); +} + +Status RedisNodeInfoAccessor::AsyncDeleteResources( + const ClientID &node_id, const std::vector &resource_names, + const StatusCallback &callback) { + Hash::HashRemoveCallback on_done = nullptr; + if (callback != nullptr) { + on_done = [callback](RedisGcsClient *client, const ClientID &node_id, + const std::vector &resource_names) { + callback(Status::OK()); + }; + } + + DynamicResourceTable &resource_table = client_impl_->resource_table(); + return resource_table.RemoveEntries(JobID::Nil(), node_id, resource_names, on_done); +} + +Status RedisNodeInfoAccessor::AsyncSubscribeToResources( + const SubscribeCallback &subscribe, + const StatusCallback &done) { + RAY_CHECK(subscribe != nullptr); + return resource_sub_executor_.AsyncSubscribeAll(ClientID::Nil(), subscribe, done); +} + } // namespace gcs } // namespace ray diff --git a/src/ray/gcs/redis_accessor.h b/src/ray/gcs/redis_accessor.h index fb520ef7d..d82259e96 100644 --- a/src/ray/gcs/redis_accessor.h +++ b/src/ray/gcs/redis_accessor.h @@ -229,6 +229,20 @@ class RedisNodeInfoAccessor : public NodeInfoAccessor { bool IsRemoved(const ClientID &node_id) const override; + Status AsyncGetResources(const ClientID &node_id, + const OptionalItemCallback &callback) override; + + Status AsyncUpdateResources(const ClientID &node_id, const ResourceMap &resources, + const StatusCallback &callback) override; + + Status AsyncDeleteResources(const ClientID &node_id, + const std::vector &resource_names, + const StatusCallback &callback) override; + + Status AsyncSubscribeToResources( + const SubscribeCallback &subscribe, + const StatusCallback &done) override; + Status AsyncReportHeartbeat(const std::shared_ptr &data_ptr, const StatusCallback &callback) override; @@ -247,6 +261,10 @@ class RedisNodeInfoAccessor : public NodeInfoAccessor { private: RedisGcsClient *client_impl_{nullptr}; + typedef SubscriptionExecutor + DynamicResourceSubscriptionExecutor; + DynamicResourceSubscriptionExecutor resource_sub_executor_; + typedef SubscriptionExecutor HeartbeatSubscriptionExecutor; HeartbeatSubscriptionExecutor heartbeat_sub_executor_; diff --git a/src/ray/gcs/redis_gcs_client.h b/src/ray/gcs/redis_gcs_client.h index 4c0b3c218..8e76f7804 100644 --- a/src/ray/gcs/redis_gcs_client.h +++ b/src/ray/gcs/redis_gcs_client.h @@ -30,6 +30,7 @@ class RAY_EXPORT RedisGcsClient : public GcsClient { friend class TaskTableTestHelper; friend class ClientTableTestHelper; friend class SetTestHelper; + friend class HashTableTestHelper; friend class ActorCheckpointIdTable; public: @@ -65,7 +66,6 @@ class RAY_EXPORT RedisGcsClient : public GcsClient { TaskLeaseTable &task_lease_table(); ErrorTable &error_table(); ProfileTable &profile_table(); - DynamicResourceTable &resource_table(); /// Used only for direct calls. Tasks submitted through the raylet transport /// should use Actors(), which has a requirement on the order in which /// entries can be appended to the log. @@ -100,10 +100,11 @@ class RAY_EXPORT RedisGcsClient : public GcsClient { JobTable &job_table(); /// This method will be deprecated, use method Objects() instead ObjectTable &object_table(); - /// The following three methods will be deprecated, use method Nodes() instead. + /// The following four methods will be deprecated, use method Nodes() instead. ClientTable &client_table(); HeartbeatTable &heartbeat_table(); HeartbeatBatchTable &heartbeat_batch_table(); + DynamicResourceTable &resource_table(); /// This method will be deprecated, use method Tasks() instead. raylet::TaskTable &raylet_task_table(); diff --git a/src/ray/gcs/subscription_executor.cc b/src/ray/gcs/subscription_executor.cc index 32c5214c0..775acf854 100644 --- a/src/ray/gcs/subscription_executor.cc +++ b/src/ray/gcs/subscription_executor.cc @@ -189,6 +189,8 @@ template class SubscriptionExecutor; template class SubscriptionExecutor; template class SubscriptionExecutor; template class SubscriptionExecutor; +template class SubscriptionExecutor; template class SubscriptionExecutor; template class SubscriptionExecutor; diff --git a/src/ray/gcs/tables.cc b/src/ray/gcs/tables.cc index 2008f4be3..6abf60deb 100644 --- a/src/ray/gcs/tables.cc +++ b/src/ray/gcs/tables.cc @@ -368,8 +368,8 @@ Status Set::Subscribe(const JobID &job_id, const ClientID &client_id, auto on_subscribe = [subscribe](RedisGcsClient *client, const ID &id, const GcsChangeMode change_mode, const std::vector &data) { - EntryChangeNotification change_notification(change_mode, data); - std::vector> notification_vec; + ArrayNotification change_notification(change_mode, data); + std::vector> notification_vec; notification_vec.emplace_back(std::move(change_notification)); subscribe(client, id, notification_vec); }; @@ -498,7 +498,11 @@ Status Hash::Subscribe(const JobID &job_id, const ClientID &client_id, data_map.emplace(key, std::move(value)); } } - subscribe(client_, id, gcs_entry.change_mode(), data_map); + MapNotification notification(gcs_entry.change_mode(), + data_map); + std::vector> notification_vec; + notification_vec.emplace_back(std::move(notification)); + subscribe(client_, id, notification_vec); } } }; diff --git a/src/ray/gcs/tables.h b/src/ray/gcs/tables.h index e618093a0..99388b316 100644 --- a/src/ray/gcs/tables.h +++ b/src/ray/gcs/tables.h @@ -473,7 +473,7 @@ class Set : private Log, using NotificationCallback = std::function &data)>; + const std::vector> &data)>; /// Subscribe to any add or remove operations to this table. /// /// \param job_id The ID of the job. @@ -540,7 +540,7 @@ class HashInterface { /// \return Void using HashNotificationCallback = std::function; + const std::vector> &data)>; /// Add entries of a hash table. /// diff --git a/src/ray/gcs/test/redis_gcs_client_test.cc b/src/ray/gcs/test/redis_gcs_client_test.cc index 1f5d16033..8f45e55c3 100644 --- a/src/ray/gcs/test/redis_gcs_client_test.cc +++ b/src/ray/gcs/test/redis_gcs_client_test.cc @@ -60,8 +60,7 @@ ClientID local_client_id = ClientID::FromRandom(); class TestGcsWithAsio : public TestGcs { public: TestGcsWithAsio(CommandType command_type) - : TestGcs(command_type), io_service_(), work_(io_service_) { - } + : TestGcs(command_type), io_service_(), work_(io_service_) {} TestGcsWithAsio() : TestGcsWithAsio(CommandType::kRegular) {} @@ -1302,153 +1301,162 @@ TEST_F(TestGcsWithAsio, TestClientTableMarkDisconnected) { ClientTableTestHelper::TestClientTableMarkDisconnected(job_id_, client_); } -void TestHashTable(const JobID &job_id, std::shared_ptr client) { - const int expected_count = 14; - ClientID client_id = ClientID::FromRandom(); - // Prepare the first resource map: data_map1. - DynamicResourceTable::DataMap data_map1; - auto cpu_data = std::make_shared(); - cpu_data->set_resource_capacity(100); - data_map1.emplace("CPU", cpu_data); - auto gpu_data = std::make_shared(); - gpu_data->set_resource_capacity(2); - data_map1.emplace("GPU", gpu_data); - // Prepare the second resource map: data_map2 which decreases CPU, - // increases GPU and add a new CUSTOM compared to data_map1. - DynamicResourceTable::DataMap data_map2; - auto data_cpu = std::make_shared(); - data_cpu->set_resource_capacity(50); - data_map2.emplace("CPU", data_cpu); - auto data_gpu = std::make_shared(); - data_gpu->set_resource_capacity(10); - data_map2.emplace("GPU", data_gpu); - auto data_custom = std::make_shared(); - data_custom->set_resource_capacity(2); - data_map2.emplace("CUSTOM", data_custom); - data_map2["CPU"]->set_resource_capacity(50); - // This is a common comparison function for the test. - auto compare_test = [](const DynamicResourceTable::DataMap &data1, - const DynamicResourceTable::DataMap &data2) { - ASSERT_EQ(data1.size(), data2.size()); - for (const auto &data : data1) { - auto iter = data2.find(data.first); - ASSERT_TRUE(iter != data2.end()); - ASSERT_EQ(iter->second->resource_capacity(), data.second->resource_capacity()); - } - }; - auto subscribe_callback = [](RedisGcsClient *client) { - ASSERT_TRUE(true); - test->IncrementNumCallbacks(); - }; - auto notification_callback = [data_map1, data_map2, compare_test]( - RedisGcsClient *client, const ClientID &id, - const GcsChangeMode change_mode, - const DynamicResourceTable::DataMap &data) { - if (change_mode == GcsChangeMode::REMOVE) { - ASSERT_EQ(data.size(), 2); - ASSERT_TRUE(data.find("GPU") != data.end()); - ASSERT_TRUE(data.find("CUSTOM") != data.end() || data.find("CPU") != data.end()); - // The key "None-Existent" will not appear in the notification. - } else { - if (data.size() == 2) { - compare_test(data_map1, data); - } else if (data.size() == 3) { - compare_test(data_map2, data); - } else { - ASSERT_TRUE(false); +class HashTableTestHelper { + public: + static void TestHashTable(const JobID &job_id, + std::shared_ptr client) { + const int expected_count = 14; + ClientID client_id = ClientID::FromRandom(); + // Prepare the first resource map: data_map1. + DynamicResourceTable::DataMap data_map1; + auto cpu_data = std::make_shared(); + cpu_data->set_resource_capacity(100); + data_map1.emplace("CPU", cpu_data); + auto gpu_data = std::make_shared(); + gpu_data->set_resource_capacity(2); + data_map1.emplace("GPU", gpu_data); + // Prepare the second resource map: data_map2 which decreases CPU, + // increases GPU and add a new CUSTOM compared to data_map1. + DynamicResourceTable::DataMap data_map2; + auto data_cpu = std::make_shared(); + data_cpu->set_resource_capacity(50); + data_map2.emplace("CPU", data_cpu); + auto data_gpu = std::make_shared(); + data_gpu->set_resource_capacity(10); + data_map2.emplace("GPU", data_gpu); + auto data_custom = std::make_shared(); + data_custom->set_resource_capacity(2); + data_map2.emplace("CUSTOM", data_custom); + data_map2["CPU"]->set_resource_capacity(50); + // This is a common comparison function for the test. + auto compare_test = [](const DynamicResourceTable::DataMap &data1, + const DynamicResourceTable::DataMap &data2) { + ASSERT_EQ(data1.size(), data2.size()); + for (const auto &data : data1) { + auto iter = data2.find(data.first); + ASSERT_TRUE(iter != data2.end()); + ASSERT_EQ(iter->second->resource_capacity(), data.second->resource_capacity()); } - } - test->IncrementNumCallbacks(); - // It is not sure which of the notification or lookup callback will come first. - if (test->NumCallbacks() == expected_count) { - test->Stop(); - } - }; - // Step 0: Subscribe the change of the hash table. - RAY_CHECK_OK(client->resource_table().Subscribe( - job_id, ClientID::Nil(), notification_callback, subscribe_callback)); - RAY_CHECK_OK(client->resource_table().RequestNotifications(job_id, client_id, - local_client_id, nullptr)); + }; + auto subscribe_callback = [](RedisGcsClient *client) { + ASSERT_TRUE(true); + test->IncrementNumCallbacks(); + }; + auto notification_callback = + [data_map1, data_map2, compare_test]( + RedisGcsClient *client, const ClientID &id, + const std::vector &result) { + RAY_CHECK(result.size() == 1); + const ResourceChangeNotification ¬ification = result.back(); + if (notification.IsRemoved()) { + ASSERT_EQ(notification.GetData().size(), 2); + ASSERT_TRUE(notification.GetData().find("GPU") != + notification.GetData().end()); + ASSERT_TRUE( + notification.GetData().find("CUSTOM") != notification.GetData().end() || + notification.GetData().find("CPU") != notification.GetData().end()); + // The key "None-Existent" will not appear in the notification. + } else { + if (notification.GetData().size() == 2) { + compare_test(data_map1, notification.GetData()); + } else if (notification.GetData().size() == 3) { + compare_test(data_map2, notification.GetData()); + } else { + ASSERT_TRUE(false); + } + } + test->IncrementNumCallbacks(); + // It is not sure which of the notification or lookup callback will come first. + if (test->NumCallbacks() == expected_count) { + test->Stop(); + } + }; + // Step 0: Subscribe the change of the hash table. + RAY_CHECK_OK(client->resource_table().Subscribe( + job_id, ClientID::Nil(), notification_callback, subscribe_callback)); + RAY_CHECK_OK(client->resource_table().RequestNotifications(job_id, client_id, + local_client_id, nullptr)); - // Step 1: Add elements to the hash table. - auto update_callback1 = [data_map1, compare_test]( - RedisGcsClient *client, const ClientID &id, - const DynamicResourceTable::DataMap &callback_data) { - compare_test(data_map1, callback_data); - test->IncrementNumCallbacks(); - }; - RAY_CHECK_OK( - client->resource_table().Update(job_id, client_id, data_map1, update_callback1)); - auto lookup_callback1 = [data_map1, compare_test]( - RedisGcsClient *client, const ClientID &id, - const DynamicResourceTable::DataMap &callback_data) { - compare_test(data_map1, callback_data); - test->IncrementNumCallbacks(); - }; - RAY_CHECK_OK(client->resource_table().Lookup(job_id, client_id, lookup_callback1)); + // Step 1: Add elements to the hash table. + auto update_callback1 = [data_map1, compare_test]( + RedisGcsClient *client, const ClientID &id, + const DynamicResourceTable::DataMap &callback_data) { + compare_test(data_map1, callback_data); + test->IncrementNumCallbacks(); + }; + RAY_CHECK_OK( + client->resource_table().Update(job_id, client_id, data_map1, update_callback1)); + auto lookup_callback1 = [data_map1, compare_test]( + RedisGcsClient *client, const ClientID &id, + const DynamicResourceTable::DataMap &callback_data) { + compare_test(data_map1, callback_data); + test->IncrementNumCallbacks(); + }; + RAY_CHECK_OK(client->resource_table().Lookup(job_id, client_id, lookup_callback1)); - // Step 2: Decrease one element, increase one and add a new one. - RAY_CHECK_OK(client->resource_table().Update(job_id, client_id, data_map2, nullptr)); - auto lookup_callback2 = [data_map2, compare_test]( - RedisGcsClient *client, const ClientID &id, - const DynamicResourceTable::DataMap &callback_data) { - compare_test(data_map2, callback_data); - test->IncrementNumCallbacks(); - }; - RAY_CHECK_OK(client->resource_table().Lookup(job_id, client_id, lookup_callback2)); - std::vector delete_keys({"GPU", "CUSTOM", "None-Existent"}); - auto remove_callback = [delete_keys](RedisGcsClient *client, const ClientID &id, - const std::vector &callback_data) { - for (size_t i = 0; i < callback_data.size(); ++i) { - // All deleting keys exist in this argument even if the key doesn't exist. - ASSERT_EQ(callback_data[i], delete_keys[i]); - } - test->IncrementNumCallbacks(); - }; - RAY_CHECK_OK(client->resource_table().RemoveEntries(job_id, client_id, delete_keys, - remove_callback)); - DynamicResourceTable::DataMap data_map3(data_map2); - data_map3.erase("GPU"); - data_map3.erase("CUSTOM"); - auto lookup_callback3 = [data_map3, compare_test]( - RedisGcsClient *client, const ClientID &id, - const DynamicResourceTable::DataMap &callback_data) { - compare_test(data_map3, callback_data); - test->IncrementNumCallbacks(); - }; - RAY_CHECK_OK(client->resource_table().Lookup(job_id, client_id, lookup_callback3)); + // Step 2: Decrease one element, increase one and add a new one. + RAY_CHECK_OK(client->resource_table().Update(job_id, client_id, data_map2, nullptr)); + auto lookup_callback2 = [data_map2, compare_test]( + RedisGcsClient *client, const ClientID &id, + const DynamicResourceTable::DataMap &callback_data) { + compare_test(data_map2, callback_data); + test->IncrementNumCallbacks(); + }; + RAY_CHECK_OK(client->resource_table().Lookup(job_id, client_id, lookup_callback2)); + std::vector delete_keys({"GPU", "CUSTOM", "None-Existent"}); + auto remove_callback = [delete_keys](RedisGcsClient *client, const ClientID &id, + const std::vector &callback_data) { + for (size_t i = 0; i < callback_data.size(); ++i) { + // All deleting keys exist in this argument even if the key doesn't exist. + ASSERT_EQ(callback_data[i], delete_keys[i]); + } + test->IncrementNumCallbacks(); + }; + RAY_CHECK_OK(client->resource_table().RemoveEntries(job_id, client_id, delete_keys, + remove_callback)); + DynamicResourceTable::DataMap data_map3(data_map2); + data_map3.erase("GPU"); + data_map3.erase("CUSTOM"); + auto lookup_callback3 = [data_map3, compare_test]( + RedisGcsClient *client, const ClientID &id, + const DynamicResourceTable::DataMap &callback_data) { + compare_test(data_map3, callback_data); + test->IncrementNumCallbacks(); + }; + RAY_CHECK_OK(client->resource_table().Lookup(job_id, client_id, lookup_callback3)); - // Step 3: Reset the the resources to data_map1. - RAY_CHECK_OK( - client->resource_table().Update(job_id, client_id, data_map1, update_callback1)); - auto lookup_callback4 = [data_map1, compare_test]( - RedisGcsClient *client, const ClientID &id, - const DynamicResourceTable::DataMap &callback_data) { - compare_test(data_map1, callback_data); - test->IncrementNumCallbacks(); - }; - RAY_CHECK_OK(client->resource_table().Lookup(job_id, client_id, lookup_callback4)); + // Step 3: Reset the the resources to data_map1. + RAY_CHECK_OK( + client->resource_table().Update(job_id, client_id, data_map1, update_callback1)); + auto lookup_callback4 = [data_map1, compare_test]( + RedisGcsClient *client, const ClientID &id, + const DynamicResourceTable::DataMap &callback_data) { + compare_test(data_map1, callback_data); + test->IncrementNumCallbacks(); + }; + RAY_CHECK_OK(client->resource_table().Lookup(job_id, client_id, lookup_callback4)); - // Step 4: Removing all elements will remove the home Hash table from GCS. - RAY_CHECK_OK(client->resource_table().RemoveEntries( - job_id, client_id, {"GPU", "CPU", "CUSTOM", "None-Existent"}, nullptr)); - auto lookup_callback5 = [](RedisGcsClient *client, const ClientID &id, - const DynamicResourceTable::DataMap &callback_data) { - ASSERT_EQ(callback_data.size(), 0); - test->IncrementNumCallbacks(); - // It is not sure which of notification or lookup callback will come first. - if (test->NumCallbacks() == expected_count) { - test->Stop(); - } - }; - RAY_CHECK_OK(client->resource_table().Lookup(job_id, client_id, lookup_callback5)); - test->Start(); - ASSERT_EQ(test->NumCallbacks(), expected_count); -} + // Step 4: Removing all elements will remove the home Hash table from GCS. + RAY_CHECK_OK(client->resource_table().RemoveEntries( + job_id, client_id, {"GPU", "CPU", "CUSTOM", "None-Existent"}, nullptr)); + auto lookup_callback5 = [](RedisGcsClient *client, const ClientID &id, + const DynamicResourceTable::DataMap &callback_data) { + ASSERT_EQ(callback_data.size(), 0); + test->IncrementNumCallbacks(); + // It is not sure which of notification or lookup callback will come first. + if (test->NumCallbacks() == expected_count) { + test->Stop(); + } + }; + RAY_CHECK_OK(client->resource_table().Lookup(job_id, client_id, lookup_callback5)); + test->Start(); + ASSERT_EQ(test->NumCallbacks(), expected_count); + } +}; TEST_F(TestGcsWithAsio, TestHashTable) { test = this; - TestHashTable(job_id_, client_); + HashTableTestHelper::TestHashTable(job_id_, client_); } #undef TEST_TASK_TABLE_MACRO diff --git a/src/ray/gcs/test/redis_node_info_accessor_test.cc b/src/ray/gcs/test/redis_node_info_accessor_test.cc new file mode 100644 index 000000000..166705630 --- /dev/null +++ b/src/ray/gcs/test/redis_node_info_accessor_test.cc @@ -0,0 +1,165 @@ +#include +#include "gtest/gtest.h" +#include "ray/gcs/redis_accessor.h" +#include "ray/gcs/redis_gcs_client.h" +#include "ray/gcs/test/accessor_test_base.h" + +namespace ray { + +namespace gcs { + +class NodeDynamicResourceTest : public AccessorTestBase { + protected: + typedef NodeInfoAccessor::ResourceMap ResourceMap; + virtual void GenTestData() { + for (size_t node_index = 0; node_index < node_number_; ++node_index) { + ClientID id = ClientID::FromRandom(); + ResourceMap resource_map; + for (size_t rs_index = 0; rs_index < resource_type_number_; ++rs_index) { + std::shared_ptr rs_data = + std::make_shared(); + rs_data->set_resource_capacity(rs_index); + std::string resource_name = std::to_string(rs_index); + resource_map[resource_name] = rs_data; + if (resource_to_delete_.empty()) { + resource_to_delete_.emplace_back(resource_name); + } + } + id_to_resource_map_[id] = std::move(resource_map); + } + } + + std::unordered_map id_to_resource_map_; + + size_t node_number_{100}; + size_t resource_type_number_{5}; + + std::vector resource_to_delete_; + + std::atomic sub_pending_count_{0}; + std::atomic do_sub_pending_count_{0}; +}; + +TEST_F(NodeDynamicResourceTest, UpdateAndGet) { + NodeInfoAccessor &node_accessor = gcs_client_->Nodes(); + for (const auto &node_rs : id_to_resource_map_) { + ++pending_count_; + const ClientID &id = node_rs.first; + // Update + Status status = node_accessor.AsyncUpdateResources( + node_rs.first, node_rs.second, [this, &node_accessor, id](Status status) { + RAY_CHECK_OK(status); + auto get_callback = [this, id](Status status, + const boost::optional &result) { + --pending_count_; + RAY_CHECK_OK(status); + const auto it = id_to_resource_map_.find(id); + ASSERT_TRUE(result); + ASSERT_EQ(it->second.size(), result->size()); + }; + // Get + status = node_accessor.AsyncGetResources(id, get_callback); + RAY_CHECK_OK(status); + }); + } + WaitPendingDone(wait_pending_timeout_); +} + +TEST_F(NodeDynamicResourceTest, Delete) { + NodeInfoAccessor &node_accessor = gcs_client_->Nodes(); + for (const auto &node_rs : id_to_resource_map_) { + ++pending_count_; + // Update + Status status = node_accessor.AsyncUpdateResources(node_rs.first, node_rs.second, + [this](Status status) { + RAY_CHECK_OK(status); + --pending_count_; + }); + } + WaitPendingDone(wait_pending_timeout_); + + for (const auto &node_rs : id_to_resource_map_) { + ++pending_count_; + const ClientID &id = node_rs.first; + // Delete + Status status = node_accessor.AsyncDeleteResources( + id, resource_to_delete_, [this, &node_accessor, id](Status status) { + RAY_CHECK_OK(status); + // Get + status = node_accessor.AsyncGetResources( + id, [this, id](Status status, const boost::optional &result) { + --pending_count_; + RAY_CHECK_OK(status); + const auto it = id_to_resource_map_.find(id); + ASSERT_TRUE(result); + ASSERT_EQ(it->second.size() - resource_to_delete_.size(), result->size()); + }); + }); + } + WaitPendingDone(wait_pending_timeout_); +} + +TEST_F(NodeDynamicResourceTest, Subscribe) { + NodeInfoAccessor &node_accessor = gcs_client_->Nodes(); + for (const auto &node_rs : id_to_resource_map_) { + ++pending_count_; + // Update + Status status = node_accessor.AsyncUpdateResources(node_rs.first, node_rs.second, + [this](Status status) { + RAY_CHECK_OK(status); + --pending_count_; + }); + } + WaitPendingDone(wait_pending_timeout_); + + auto subscribe = [this](const ClientID &id, + const ResourceChangeNotification ¬ification) { + RAY_LOG(INFO) << "receive client id=" << id; + auto it = id_to_resource_map_.find(id); + ASSERT_TRUE(it != id_to_resource_map_.end()); + if (notification.IsAdded()) { + ASSERT_EQ(notification.GetData().size(), it->second.size()); + } else { + ASSERT_EQ(notification.GetData().size(), resource_to_delete_.size()); + } + --sub_pending_count_; + }; + + auto done = [this](Status status) { + RAY_CHECK_OK(status); + --pending_count_; + }; + + // Subscribe + ++pending_count_; + Status status = node_accessor.AsyncSubscribeToResources(subscribe, done); + RAY_CHECK_OK(status); + + for (const auto &node_rs : id_to_resource_map_) { + // Delete + ++pending_count_; + ++sub_pending_count_; + Status status = node_accessor.AsyncDeleteResources(node_rs.first, resource_to_delete_, + [this](Status status) { + RAY_CHECK_OK(status); + --pending_count_; + }); + RAY_CHECK_OK(status); + } + + WaitPendingDone(wait_pending_timeout_); + WaitPendingDone(sub_pending_count_, wait_pending_timeout_); +} + +} // namespace gcs + +} // namespace ray + +int main(int argc, char **argv) { + ::testing::InitGoogleTest(&argc, argv); + RAY_CHECK(argc == 4); + ray::REDIS_SERVER_EXEC_PATH = argv[1]; + ray::REDIS_CLIENT_EXEC_PATH = argv[2]; + ray::REDIS_MODULE_LIBRARY_PATH = argv[3]; + return RUN_ALL_TESTS(); +} diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index 788bdf581..9ddf9c4c0 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -183,31 +183,27 @@ ray::Status NodeManager::RegisterGcs() { // Subscribe to resource changes. const auto &resources_changed = - [this]( - gcs::RedisGcsClient *client, const ClientID &id, - const gcs::GcsChangeMode change_mode, - const std::unordered_map> - &data) { - if (change_mode == gcs::GcsChangeMode::APPEND_OR_ADD) { + [this](const ClientID &id, + const gcs::ResourceChangeNotification &resource_notification) { + if (resource_notification.IsAdded()) { ResourceSet resource_set; - for (auto &entry : data) { + for (auto &entry : resource_notification.GetData()) { resource_set.AddOrUpdateResource(entry.first, entry.second->resource_capacity()); } ResourceCreateUpdated(id, resource_set); - } - if (change_mode == gcs::GcsChangeMode::REMOVE) { + } else { + RAY_CHECK(resource_notification.IsRemoved()); std::vector resource_names; - for (auto &entry : data) { + for (auto &entry : resource_notification.GetData()) { resource_names.push_back(entry.first); } ResourceDeleted(id, resource_names); } }; - RAY_RETURN_NOT_OK( - gcs_client_->resource_table().Subscribe(JobID::Nil(), ClientID::Nil(), - /*subscribe_callback=*/resources_changed, - /*done_callback=*/nullptr)); + RAY_RETURN_NOT_OK(gcs_client_->Nodes().AsyncSubscribeToResources( + /*subscribe_callback=*/resources_changed, + /*done_callback=*/nullptr)); // Subscribe to heartbeat batches from the monitor. const auto &heartbeat_batch_added = @@ -471,17 +467,18 @@ void NodeManager::NodeAdded(const GcsNodeInfo &node_info) { remote_node_manager_clients_.emplace(node_id, std::move(client)); // Fetch resource info for the remote client and update cluster resource map. - RAY_CHECK_OK(gcs_client_->resource_table().Lookup( - JobID::Nil(), node_id, - [this](gcs::RedisGcsClient *client, const ClientID &node_id, - const std::unordered_map> &pairs) { - ResourceSet resource_set; - for (auto &resource_entry : pairs) { - resource_set.AddOrUpdateResource(resource_entry.first, - resource_entry.second->resource_capacity()); + RAY_CHECK_OK(gcs_client_->Nodes().AsyncGetResources( + node_id, + [this, node_id](Status status, + const boost::optional &data) { + if (data) { + ResourceSet resource_set; + for (auto &resource_entry : *data) { + resource_set.AddOrUpdateResource(resource_entry.first, + resource_entry.second->resource_capacity()); + } + ResourceCreateUpdated(node_id, resource_set); } - ResourceCreateUpdated(node_id, resource_set); })); } @@ -1704,15 +1701,15 @@ void NodeManager::ProcessSetResourceRequest( double const &capacity = message->capacity(); bool is_deletion = capacity <= 0; - ClientID client_id = from_flatbuf(*message->client_id()); + ClientID node_id = from_flatbuf(*message->client_id()); - // If the python arg was null, set client_id to the local client - if (client_id.IsNil()) { - client_id = self_node_id_; + // If the python arg was null, set node_id to the local node id. + if (node_id.IsNil()) { + node_id = self_node_id_; } if (is_deletion && - cluster_resource_map_[client_id].GetTotalResources().GetResourceMap().count( + cluster_resource_map_[node_id].GetTotalResources().GetResourceMap().count( resource_name) == 0) { // Resource does not exist in the cluster resource map, thus nothing to delete. // Return.. @@ -1724,15 +1721,14 @@ void NodeManager::ProcessSetResourceRequest( // Submit to the resource table. This calls the ResourceCreateUpdated or ResourceDeleted // callback, which updates cluster_resource_map_. if (is_deletion) { - RAY_CHECK_OK(gcs_client_->resource_table().RemoveEntries(JobID::Nil(), client_id, - {resource_name}, nullptr)); + RAY_CHECK_OK( + gcs_client_->Nodes().AsyncDeleteResources(node_id, {resource_name}, nullptr)); } else { std::unordered_map> data_map; auto resource_table_data = std::make_shared(); resource_table_data->set_resource_capacity(capacity); data_map.emplace(resource_name, resource_table_data); - RAY_CHECK_OK( - gcs_client_->resource_table().Update(JobID::Nil(), client_id, data_map, nullptr)); + RAY_CHECK_OK(gcs_client_->Nodes().AsyncUpdateResources(node_id, data_map, nullptr)); } } diff --git a/src/ray/raylet/raylet.cc b/src/ray/raylet/raylet.cc index 17076670d..936b69fd0 100644 --- a/src/ray/raylet/raylet.cc +++ b/src/ray/raylet/raylet.cc @@ -103,8 +103,8 @@ ray::Status Raylet::RegisterGcs() { resource->set_resource_capacity(resource_pair.second); resources.emplace(resource_pair.first, resource); } - RAY_RETURN_NOT_OK(gcs_client_->resource_table().Update(JobID::Nil(), self_node_id_, - resources, nullptr)); + RAY_RETURN_NOT_OK( + gcs_client_->Nodes().AsyncUpdateResources(self_node_id_, resources, nullptr)); RAY_RETURN_NOT_OK(node_manager_.RegisterGcs());