From 6114b6d20ee2e1b5c37862964b30e7949bba2a3c Mon Sep 17 00:00:00 2001 From: Stephanie Wang Date: Sun, 11 Mar 2018 19:17:18 -0700 Subject: [PATCH] Implement the client table for the new GCS (#1674) * Add subscription callback to CallbackData * Implement ClientTable * Hook up ClientTable to AsyncGCSClient * Add client_info to GCSClient Connect interface * client table callbacks * Unit test for client table * Doc * Fix idempotency check * Fix mac build * Fix memory issues in gcs client test * Fix disconnection bug * lint --- src/common/redis_module/ray_redis_module.cc | 44 ++++- src/global_scheduler/global_scheduler.cc | 8 +- src/local_scheduler/local_scheduler.cc | 8 +- src/plasma/plasma_manager.cc | 8 +- src/ray/CMakeLists.txt | 1 + src/ray/gcs/client.cc | 30 +-- src/ray/gcs/client.h | 37 ++-- src/ray/gcs/client_test.cc | 154 ++++++++++++---- src/ray/gcs/format/gcs.fbs | 18 +- src/ray/gcs/redis_context.h | 1 + src/ray/gcs/tables.cc | 194 ++++++++++++-------- src/ray/gcs/tables.h | 164 ++++++++++------- src/ray/gcs/task_table.cc | 68 +++++++ 13 files changed, 502 insertions(+), 233 deletions(-) create mode 100644 src/ray/gcs/task_table.cc diff --git a/src/common/redis_module/ray_redis_module.cc b/src/common/redis_module/ray_redis_module.cc index 9e05ba32d..7b5469665 100644 --- a/src/common/redis_module/ray_redis_module.cc +++ b/src/common/redis_module/ray_redis_module.cc @@ -420,9 +420,7 @@ int TableAdd_RedisCommand(RedisModuleCtx *ctx, // Publish a message on the requested pubsub channel if necessary. if (pubsub_channel == TablePubsub_TASK) { - size_t len = 0; - const char *buf = RedisModule_StringPtrLen(data, &len); - + const char *buf = RedisModule_StringPtrLen(data, NULL); auto message = flatbuffers::GetRoot(buf); if (message->scheduling_state() == SchedulingState_WAITING || @@ -462,6 +460,46 @@ int TableAdd_RedisCommand(RedisModuleCtx *ctx, RedisModule_FreeString(ctx, publish_message); RedisModule_FreeString(ctx, publish_topic); } + } else if (pubsub_channel == TablePubsub_CLIENT) { + const char *buf = RedisModule_StringPtrLen(data, NULL); + auto client_data = flatbuffers::GetRoot(buf); + + RedisModuleKey *clients_key = (RedisModuleKey *) RedisModule_OpenKey( + ctx, pubsub_channel_str, REDISMODULE_READ | REDISMODULE_WRITE); + // If this is a client addition, send all previous notifications, in order. + // NOTE(swang): This will go to all clients, so some clients will get + // duplicate notifications. + if (client_data->is_insertion() && + RedisModule_KeyType(clients_key) != REDISMODULE_KEYTYPE_EMPTY) { + // NOTE(swang): Sets are not implemented yet, so we use ZSETs instead. + CHECK_ERROR(RedisModule_ZsetFirstInScoreRange( + clients_key, REDISMODULE_NEGATIVE_INFINITE, + REDISMODULE_POSITIVE_INFINITE, 1, 1), + "Unable to initialize zset iterator"); + do { + RedisModuleString *message = + RedisModule_ZsetRangeCurrentElement(clients_key, NULL); + RedisModuleCallReply *reply = + RedisModule_Call(ctx, "PUBLISH", "ss", pubsub_channel_str, message); + if (reply == NULL) { + RedisModule_CloseKey(clients_key); + RedisModule_ReplyWithError(ctx, "error during PUBLISH"); + } + } while (RedisModule_ZsetRangeNext(clients_key)); + } + + // Append this notification to the past notifications so that it will get + // sent to new clients in the future. + size_t index = RedisModule_ValueLength(key); + RedisModule_ZsetAdd(clients_key, index, data, NULL); + // Publish the notification about this client. + RedisModuleCallReply *reply = + RedisModule_Call(ctx, "PUBLISH", "ss", pubsub_channel_str, data); + if (reply == NULL) { + RedisModule_ReplyWithError(ctx, "error during PUBLISH"); + } + + RedisModule_CloseKey(clients_key); } else if (pubsub_channel != TablePubsub_NO_PUBLISH) { // All other pubsub channels write the data back directly onto the channel. RedisModuleCallReply *reply = diff --git a/src/global_scheduler/global_scheduler.cc b/src/global_scheduler/global_scheduler.cc index b87c58cf9..f8c49be9c 100644 --- a/src/global_scheduler/global_scheduler.cc +++ b/src/global_scheduler/global_scheduler.cc @@ -140,8 +140,14 @@ GlobalSchedulerState *GlobalSchedulerState_init(event_loop *loop, "global_scheduler", node_ip_address, std::vector()); db_attach(state->db, loop, false); + + ClientTableDataT client_info; + client_info.client_id = get_db_client_id(state->db).binary(); + client_info.node_manager_address = std::string(node_ip_address); + client_info.local_scheduler_port = 0; + client_info.object_manager_port = 0; RAY_CHECK_OK(state->gcs_client.Connect(std::string(redis_primary_addr), - redis_primary_port)); + redis_primary_port, client_info)); RAY_CHECK_OK(state->gcs_client.context()->AttachToEventLoop(loop)); state->policy_state = GlobalSchedulerPolicyState_init(); return state; diff --git a/src/local_scheduler/local_scheduler.cc b/src/local_scheduler/local_scheduler.cc index 2ae75f0fe..4e2184f4a 100644 --- a/src/local_scheduler/local_scheduler.cc +++ b/src/local_scheduler/local_scheduler.cc @@ -374,8 +374,14 @@ LocalSchedulerState *LocalSchedulerState_init( state->db = db_connect(std::string(redis_primary_addr), redis_primary_port, "local_scheduler", node_ip_address, db_connect_args); db_attach(state->db, loop, false); + + ClientTableDataT client_info; + client_info.client_id = get_db_client_id(state->db).binary(); + client_info.node_manager_address = std::string(node_ip_address); + client_info.local_scheduler_port = 0; + client_info.object_manager_port = 0; RAY_CHECK_OK(state->gcs_client.Connect(std::string(redis_primary_addr), - redis_primary_port)); + redis_primary_port, client_info)); RAY_CHECK_OK(state->gcs_client.context()->AttachToEventLoop(loop)); } else { state->db = NULL; diff --git a/src/plasma/plasma_manager.cc b/src/plasma/plasma_manager.cc index fea58fc36..39c977a8e 100644 --- a/src/plasma/plasma_manager.cc +++ b/src/plasma/plasma_manager.cc @@ -486,8 +486,14 @@ PlasmaManagerState *PlasmaManagerState_init(const char *store_socket_name, state->db = db_connect(std::string(redis_primary_addr), redis_primary_port, "plasma_manager", manager_addr, db_connect_args); db_attach(state->db, state->loop, false); + + ClientTableDataT client_info; + client_info.client_id = get_db_client_id(state->db).binary(); + client_info.node_manager_address = std::string(manager_addr); + client_info.local_scheduler_port = 0; + client_info.object_manager_port = manager_port; RAY_CHECK_OK(state->gcs_client.Connect(std::string(redis_primary_addr), - redis_primary_port)); + redis_primary_port, client_info)); RAY_CHECK_OK(state->gcs_client.context()->AttachToEventLoop(state->loop)); } else { state->db = NULL; diff --git a/src/ray/CMakeLists.txt b/src/ray/CMakeLists.txt index 546a252b8..7a4166ac6 100644 --- a/src/ray/CMakeLists.txt +++ b/src/ray/CMakeLists.txt @@ -33,6 +33,7 @@ set(RAY_SRCS status.cc gcs/client.cc gcs/tables.cc + gcs/task_table.cc gcs/redis_context.cc gcs/asio.cc common/client_connection.cc diff --git a/src/ray/gcs/client.cc b/src/ray/gcs/client.cc index 074d1ea94..01bfc4566 100644 --- a/src/ray/gcs/client.cc +++ b/src/ray/gcs/client.cc @@ -10,11 +10,17 @@ AsyncGcsClient::AsyncGcsClient() {} AsyncGcsClient::~AsyncGcsClient() {} -Status AsyncGcsClient::Connect(const std::string &address, int port) { +Status AsyncGcsClient::Connect(const std::string &address, int port, + const ClientTableDataT &client_info) { context_.reset(new RedisContext()); RAY_RETURN_NOT_OK(context_->Connect(address, port)); object_table_.reset(new ObjectTable(context_, this)); task_table_.reset(new TaskTable(context_, this)); + client_table_.reset(new ClientTable(context_, this, client_info)); + // TODO(swang): Call the client table's Connect() method here. To do this, + // we need to make sure that we are attached to an event loop first. This + // currently isn't possible because the aeEventLoop, which we use for + // testing, requires us to connect to Redis first. return Status::OK(); } @@ -25,25 +31,21 @@ Status Attach(plasma::EventLoop &event_loop) { } Status AsyncGcsClient::Attach(boost::asio::io_service &io_service) { - asio_client_.reset(new RedisAsioClient(io_service, context_->async_context())); + asio_async_client_.reset(new RedisAsioClient(io_service, context_->async_context())); + asio_subscribe_client_.reset( + new RedisAsioClient(io_service, context_->subscribe_context())); return Status::OK(); } -ObjectTable &AsyncGcsClient::object_table() { - return *object_table_; -} +ObjectTable &AsyncGcsClient::object_table() { return *object_table_; } -TaskTable &AsyncGcsClient::task_table() { - return *task_table_; -} +TaskTable &AsyncGcsClient::task_table() { return *task_table_; } -FunctionTable &AsyncGcsClient::function_table() { - return *function_table_; -} +ClientTable &AsyncGcsClient::client_table() { return *client_table_; } -ClassTable &AsyncGcsClient::class_table() { - return *class_table_; -} +FunctionTable &AsyncGcsClient::function_table() { return *function_table_; } + +ClassTable &AsyncGcsClient::class_table() { return *class_table_; } } // namespace gcs diff --git a/src/ray/gcs/client.h b/src/ray/gcs/client.h index 552f96391..4b9b27719 100644 --- a/src/ray/gcs/client.h +++ b/src/ray/gcs/client.h @@ -22,7 +22,14 @@ class RAY_EXPORT AsyncGcsClient { AsyncGcsClient(); ~AsyncGcsClient(); - Status Connect(const std::string &address, int port); + /// Connect to the GCS. + /// + /// \param address The GCS IP address. + /// \param port The GCS port. + /// \param client_info Information about the local client to connect. + /// \return Status. + Status Connect(const std::string &address, int port, + const ClientTableDataT &client_info); /// Attach this client to a plasma event loop. Note that only /// one event loop should be attached at a time. Status Attach(plasma::EventLoop &event_loop); @@ -38,6 +45,7 @@ class RAY_EXPORT AsyncGcsClient { inline ConfigTable &config_table(); ObjectTable &object_table(); TaskTable &task_table(); + ClientTable &client_table(); inline ErrorTable &error_table(); // We also need something to export generic code to run on workers from the @@ -45,8 +53,7 @@ class RAY_EXPORT AsyncGcsClient { using GetExportCallback = std::function; Status AddExport(const std::string &driver_id, std::string &export_data); - Status GetExport(const std::string &driver_id, - int64_t export_index, + Status GetExport(const std::string &driver_id, int64_t export_index, const GetExportCallback &done_callback); std::shared_ptr context() { return context_; } @@ -56,29 +63,23 @@ class RAY_EXPORT AsyncGcsClient { std::unique_ptr class_table_; std::unique_ptr object_table_; std::unique_ptr task_table_; + std::unique_ptr client_table_; std::shared_ptr context_; - std::unique_ptr asio_client_; + std::unique_ptr asio_async_client_; + std::unique_ptr asio_subscribe_client_; }; class SyncGcsClient { - Status LogEvent(const std::string &key, - const std::string &value, - double timestamp); + Status LogEvent(const std::string &key, const std::string &value, double timestamp); Status NotifyError(const std::map &error_info); - Status RegisterFunction(const JobID &job_id, - const FunctionID &function_id, - const std::string &language, - const std::string &name, + Status RegisterFunction(const JobID &job_id, const FunctionID &function_id, + const std::string &language, const std::string &name, const std::string &data); - Status RetrieveFunction(const JobID &job_id, - const FunctionID &function_id, - std::string *name, - std::string *data); + Status RetrieveFunction(const JobID &job_id, const FunctionID &function_id, + std::string *name, std::string *data); Status AddExport(const std::string &driver_id, std::string &export_data); - Status GetExport(const std::string &driver_id, - int64_t export_index, - std::string *data); + Status GetExport(const std::string &driver_id, int64_t export_index, std::string *data); }; } // namespace gcs diff --git a/src/ray/gcs/client_test.cc b/src/ray/gcs/client_test.cc index f89fcf5d4..bff022076 100644 --- a/src/ray/gcs/client_test.cc +++ b/src/ray/gcs/client_test.cc @@ -2,9 +2,9 @@ // TODO(pcm): get rid of this and replace with the type safe plasma event loop extern "C" { +#include "hiredis/adapters/ae.h" #include "hiredis/async.h" #include "hiredis/hiredis.h" -#include "hiredis/adapters/ae.h" } #include "ray/gcs/client.h" @@ -12,22 +12,39 @@ extern "C" { namespace ray { +/* Flush redis. */ +static inline void flushall_redis(void) { + redisContext *context = redisConnect("127.0.0.1", 6379); + freeReplyObject(redisCommand(context, "FLUSHALL")); + redisFree(context); +} + class TestGcs : public ::testing::Test { public: TestGcs() { - RAY_CHECK_OK(client_.Connect("127.0.0.1", 6379)); - job_id_ = UniqueID::from_random(); + client_ = std::make_shared(); + ClientTableDataT client_info; + client_info.client_id = ClientID::from_random().binary(); + client_info.node_manager_address = "127.0.0.1"; + client_info.local_scheduler_port = 0; + client_info.object_manager_port = 0; + RAY_CHECK_OK(client_->Connect("127.0.0.1", 6379, client_info)); + + job_id_ = JobID::from_random(); } - virtual ~TestGcs(){}; + virtual ~TestGcs() { + // Clear all keys in the GCS. + flushall_redis(); + }; virtual void Start() = 0; virtual void Stop() = 0; protected: - gcs::AsyncGcsClient client_; - UniqueID job_id_; + std::shared_ptr client_; + JobID job_id_; }; TestGcs *test; @@ -36,9 +53,13 @@ class TestGcsWithAe : public TestGcs { public: TestGcsWithAe() { loop_ = aeCreateEventLoop(1024); - RAY_CHECK_OK(client_.context()->AttachToEventLoop(loop_)); + RAY_CHECK_OK(client_->context()->AttachToEventLoop(loop_)); + } + ~TestGcsWithAe() override { + // Destroy the client first since it has a reference to the event loop. + client_.reset(); + aeDeleteEventLoop(loop_); } - ~TestGcsWithAe() override { aeDeleteEventLoop(loop_); } void Start() override { aeMain(loop_); } void Stop() override { aeStop(loop_); } @@ -48,35 +69,42 @@ class TestGcsWithAe : public TestGcs { class TestGcsWithAsio : public TestGcs { public: - TestGcsWithAsio() { RAY_CHECK_OK(client_.Attach(io_service_)); } + TestGcsWithAsio() : TestGcs(), io_service_(), work_(io_service_) { + RAY_CHECK_OK(client_->Attach(io_service_)); + } + ~TestGcsWithAsio() { + // Destroy the client first since it has a reference to the event loop. + client_.reset(); + } void Start() override { io_service_.run(); } void Stop() override { io_service_.stop(); } private: boost::asio::io_service io_service_; + // Give the event loop some work so that it's forced to run until Stop() is + // called. + boost::asio::io_service::work work_; }; -void ObjectAdded(gcs::AsyncGcsClient *client, - const UniqueID &id, +void ObjectAdded(gcs::AsyncGcsClient *client, const UniqueID &id, std::shared_ptr data) { ASSERT_EQ(data->managers, std::vector({"A", "B"})); } -void Lookup(gcs::AsyncGcsClient *client, - const UniqueID &id, +void Lookup(gcs::AsyncGcsClient *client, const UniqueID &id, std::shared_ptr data) { // Check that the object entry was added. ASSERT_EQ(data->managers, std::vector({"A", "B"})); test->Stop(); } -void TestObjectTable(const UniqueID &job_id, gcs::AsyncGcsClient &client) { +void TestObjectTable(const JobID &job_id, std::shared_ptr client) { auto data = std::make_shared(); data->managers.push_back("A"); data->managers.push_back("B"); ObjectID object_id = ObjectID::from_random(); - RAY_CHECK_OK(client.object_table().Add(job_id, object_id, data, &ObjectAdded)); - RAY_CHECK_OK(client.object_table().Lookup(job_id, object_id, &Lookup)); + RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, &ObjectAdded)); + RAY_CHECK_OK(client->object_table().Lookup(job_id, object_id, &Lookup)); // Run the event loop. The loop will only stop if the Lookup callback is // called (or an assertion failure). test->Start(); @@ -92,41 +120,36 @@ TEST_F(TestGcsWithAsio, TestObjectTable) { TestObjectTable(job_id_, client_); } -void TaskAdded(gcs::AsyncGcsClient *client, - const TaskID &id, +void TaskAdded(gcs::AsyncGcsClient *client, const TaskID &id, std::shared_ptr data) { ASSERT_EQ(data->scheduling_state, SchedulingState_SCHEDULED); } -void TaskLookup(gcs::AsyncGcsClient *client, - const TaskID &id, +void TaskLookup(gcs::AsyncGcsClient *client, const TaskID &id, std::shared_ptr data) { ASSERT_EQ(data->scheduling_state, SchedulingState_SCHEDULED); } -void TaskLookupAfterUpdate(gcs::AsyncGcsClient *client, - const TaskID &id, +void TaskLookupAfterUpdate(gcs::AsyncGcsClient *client, const TaskID &id, std::shared_ptr data) { ASSERT_EQ(data->scheduling_state, SchedulingState_LOST); test->Stop(); } -void TaskUpdateCallback(gcs::AsyncGcsClient *client, - const TaskID &task_id, - const TaskTableDataT &task, - bool updated) { - RAY_CHECK_OK(client->task_table().Lookup(DriverID::nil(), task_id, - &TaskLookupAfterUpdate)); +void TaskUpdateCallback(gcs::AsyncGcsClient *client, const TaskID &task_id, + const TaskTableDataT &task, bool updated) { + RAY_CHECK_OK( + client->task_table().Lookup(DriverID::nil(), task_id, &TaskLookupAfterUpdate)); } -void TestTaskTable(const UniqueID &job_id, gcs::AsyncGcsClient &client) { +void TestTaskTable(const JobID &job_id, std::shared_ptr client) { auto data = std::make_shared(); data->scheduling_state = SchedulingState_SCHEDULED; ClientID local_scheduler_id = ClientID::from_binary("abcdefghijklmnopqrst"); data->scheduler_id = local_scheduler_id.binary(); TaskID task_id = TaskID::from_random(); - RAY_CHECK_OK(client.task_table().Add(job_id, task_id, data, &TaskAdded)); - RAY_CHECK_OK(client.task_table().Lookup(job_id, task_id, &TaskLookup)); + RAY_CHECK_OK(client->task_table().Add(job_id, task_id, data, &TaskAdded)); + RAY_CHECK_OK(client->task_table().Lookup(job_id, task_id, &TaskLookup)); auto update = std::make_shared(); update->test_scheduler_id = local_scheduler_id.binary(); update->test_state_bitmask = SchedulingState_SCHEDULED; @@ -134,7 +157,7 @@ void TestTaskTable(const UniqueID &job_id, gcs::AsyncGcsClient &client) { // After test-and-setting, the callback will lookup the current state of the // task. RAY_CHECK_OK( - client.task_table().TestAndUpdate(job_id, task_id, update, &TaskUpdateCallback)); + client->task_table().TestAndUpdate(job_id, task_id, update, &TaskUpdateCallback)); // Run the event loop. The loop will only stop if the lookup after the // test-and-set succeeds (or an assertion failure). test->Start(); @@ -155,11 +178,11 @@ void ObjectTableSubscribed(gcs::AsyncGcsClient *client, const UniqueID &id, test->Stop(); } -void TestSubscribeAll(const UniqueID &job_id, gcs::AsyncGcsClient &client) { +void TestSubscribeAll(const JobID &job_id, std::shared_ptr client) { // Subscribe to all object table notifications. The registered callback for // notifications will check whether the object below is added. - RAY_CHECK_OK(client.object_table().Subscribe(job_id, ClientID::nil(), &Lookup, - &ObjectTableSubscribed)); + RAY_CHECK_OK(client->object_table().Subscribe(job_id, ClientID::nil(), &Lookup, + &ObjectTableSubscribed)); // Run the event loop. The loop will only stop if the subscription succeeds. test->Start(); @@ -168,7 +191,7 @@ void TestSubscribeAll(const UniqueID &job_id, gcs::AsyncGcsClient &client) { data->managers.push_back("A"); data->managers.push_back("B"); ObjectID object_id = ObjectID::from_random(); - RAY_CHECK_OK(client.object_table().Add(job_id, object_id, data, &ObjectAdded)); + RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, &ObjectAdded)); // Run the event loop. The loop will only stop if the registered subscription // callback is called (or an assertion failure). test->Start(); @@ -184,4 +207,63 @@ TEST_F(TestGcsWithAsio, TestSubscribeAll) { TestSubscribeAll(job_id_, client_); } +void ClientTableNotification(gcs::AsyncGcsClient *client, const UniqueID &id, + std::shared_ptr data, bool is_insertion) { + ClientID added_id = client->client_table().GetLocalClientId(); + ASSERT_EQ(ClientID::from_binary(data->client_id), added_id); + ASSERT_EQ(data->is_insertion, is_insertion); + + auto cached_client = client->client_table().GetClient(added_id); + ASSERT_EQ(ClientID::from_binary(cached_client.client_id), added_id); + ASSERT_EQ(cached_client.is_insertion, is_insertion); +} + +void TestClientTableConnect(const JobID &job_id, + std::shared_ptr client) { + // Register callbacks for when a client gets added and removed. The latter + // event will stop the event loop. + client->client_table().RegisterClientAddedCallback( + [](gcs::AsyncGcsClient *client, const UniqueID &id, + std::shared_ptr data) { + ClientTableNotification(client, id, data, true); + test->Stop(); + }); + // Connect and disconnect to client table. We should receive notifications + // for the addition and removal of our own entry. + RAY_CHECK_OK(client->client_table().Connect()); + test->Start(); +} + +TEST_F(TestGcsWithAsio, TestClientTableConnect) { + test = this; + TestClientTableConnect(job_id_, client_); +} + +void TestClientTableDisconnect(const JobID &job_id, + std::shared_ptr client) { + // Register callbacks for when a client gets added and removed. The latter + // event will stop the event loop. + client->client_table().RegisterClientAddedCallback( + [](gcs::AsyncGcsClient *client, const UniqueID &id, + std::shared_ptr data) { + ClientTableNotification(client, id, data, true); + }); + client->client_table().RegisterClientRemovedCallback( + [](gcs::AsyncGcsClient *client, const UniqueID &id, + std::shared_ptr data) { + ClientTableNotification(client, id, data, false); + test->Stop(); + }); + // Connect and disconnect to client table. We should receive notifications + // for the addition and removal of our own entry. + RAY_CHECK_OK(client->client_table().Connect()); + RAY_CHECK_OK(client->client_table().Disconnect()); + test->Start(); +} + +TEST_F(TestGcsWithAsio, TestClientTableDisconnect) { + test = this; + TestClientTableDisconnect(job_id_, client_); +} + } // namespace diff --git a/src/ray/gcs/format/gcs.fbs b/src/ray/gcs/format/gcs.fbs index 3e122112b..386873e33 100644 --- a/src/ray/gcs/format/gcs.fbs +++ b/src/ray/gcs/format/gcs.fbs @@ -74,13 +74,6 @@ table CustomSerializerData { table ConfigTableData { } -table Resource { - // The type of the resource. - resource_name: string; - // The total capacity of this resource type. - resource_capacity: double; -} - table ClientTableData { // The client ID of the client that the message is about. client_id: string; @@ -88,17 +81,22 @@ table ClientTableData { node_manager_address: string; // The port at which the client's node manager is listening for TCP // connections from other node managers. - node_manager_port: int; + local_scheduler_port: int; // The port at which the client's object manager is listening for TCP // connections from other object managers. object_manager_port: int; - // The total resources of this client. - resources_total: [Resource]; // True if the message is about the addition of a client and false if it is // about the deletion of a client. is_insertion: bool; } +table Resource { + // The type of the resource. + resource_name: string; + // The total capacity of this resource type. + resource_capacity: double; +} + table NodeManagerHeartbeat { // The available resources on this node manager. This information may be // stale. diff --git a/src/ray/gcs/redis_context.h b/src/ray/gcs/redis_context.h index 8f10ab171..c118a1f94 100644 --- a/src/ray/gcs/redis_context.h +++ b/src/ray/gcs/redis_context.h @@ -56,6 +56,7 @@ class RedisContext { Status SubscribeAsync(const ClientID &client_id, const TablePubsub pubsub_channel, int64_t callback_index); redisAsyncContext *async_context() { return async_context_; } + redisAsyncContext *subscribe_context() { return subscribe_context_; }; private: redisContext *context_; diff --git a/src/ray/gcs/tables.cc b/src/ray/gcs/tables.cc index 66de6f942..7426e80d7 100644 --- a/src/ray/gcs/tables.cc +++ b/src/ray/gcs/tables.cc @@ -2,101 +2,139 @@ #include "ray/gcs/client.h" -#include "common_protocol.h" - -namespace { - -std::shared_ptr MakeTaskTableData(const TaskExecutionSpec &execution_spec, - const ClientID &local_scheduler_id, - SchedulingState scheduling_state) { - auto data = std::make_shared(); - data->scheduling_state = scheduling_state; - data->task_info = - std::string(execution_spec.Spec(), execution_spec.SpecSize()); - data->scheduler_id = local_scheduler_id.binary(); - - flatbuffers::FlatBufferBuilder fbb; - auto execution_dependencies = CreateTaskExecutionDependencies( - fbb, to_flatbuf(fbb, execution_spec.ExecutionDependencies())); - fbb.Finish(execution_dependencies); - - data->execution_dependencies = - std::string((const char *) fbb.GetBufferPointer(), fbb.GetSize()); - data->spillback_count = execution_spec.SpillbackCount(); - - return data; -} - -} // namespace - namespace ray { namespace gcs { -// TODO(pcm): This is a helper method that should go away once we get rid of -// the Task* datastructure and replace it with TaskTableDataT. -Status TaskTableAdd(AsyncGcsClient *gcs_client, Task *task) { - TaskExecutionSpec &execution_spec = *Task_task_execution_spec(task); - TaskSpec *spec = execution_spec.Spec(); - auto data = MakeTaskTableData(execution_spec, Task_local_scheduler(task), - static_cast(Task_state(task))); - return gcs_client->task_table().Add( - ray::JobID::nil(), TaskSpec_task_id(spec), data, - [](gcs::AsyncGcsClient *client, const TaskID &id, - std::shared_ptr data) {}); +void ClientTable::RegisterClientAddedCallback(const Callback &callback) { + client_added_callback_ = callback; + // Call the callback for any added clients that are cached. + for (const auto &entry : client_cache_) { + if (!entry.first.is_nil() && entry.second.is_insertion) { + auto data = std::make_shared(entry.second); + client_added_callback_(client_, entry.first, data); + } + } } -// TODO(pcm): This is a helper method that should go away once we get rid of -// the Task* datastructure and replace it with TaskTableDataT. -Status TaskTableTestAndUpdate(AsyncGcsClient *gcs_client, const TaskID &task_id, - const ClientID &local_scheduler_id, int test_state_bitmask, - SchedulingState update_state, - const TaskTable::TestAndUpdateCallback &callback) { - auto data = std::make_shared(); - data->test_scheduler_id = local_scheduler_id.binary(); - data->test_state_bitmask = test_state_bitmask; - data->update_state = update_state; - return gcs_client->task_table().TestAndUpdate(ray::JobID::nil(), task_id, - data, callback); +void ClientTable::RegisterClientRemovedCallback(const Callback &callback) { + client_removed_callback_ = callback; + // Call the callback for any removed clients that are cached. + for (const auto &entry : client_cache_) { + if (!entry.first.is_nil() && !entry.second.is_insertion) { + auto data = std::make_shared(entry.second); + client_removed_callback_(client_, entry.first, data); + } + } } -void ClientConnected(gcs::AsyncGcsClient *client, const ClientID &client_id, - std::shared_ptr data) {} +void ClientTable::HandleNotification(AsyncGcsClient *client, const ClientID &channel_id, + std::shared_ptr data) { + ClientID client_id = ClientID::from_binary(data->client_id); + // It's possible to get duplicate notifications from the client table, so + // check whether this notification is new. + auto entry = client_cache_.find(client_id); + bool is_new; + if (entry == client_cache_.end()) { + // If the entry is not in the cache, then the notification is new. + is_new = true; + } else { + // If the entry is in the cache, then the notification is new if the client + // was alive and is now dead. + bool was_inserted = entry->second.is_insertion; + bool is_deleted = !data->is_insertion; + is_new = (was_inserted && is_deleted); + // Once a client with a given ID has been removed, it should never be added + // again. If the entry was in the cache and the client was deleted, check + // that this new notification is not an insertion. + if (!entry->second.is_insertion) { + RAY_CHECK(!data->is_insertion) + << "Notification for addition of a client that was already removed:" + << client_id.hex(); + } + } -ClientTable::ClientTable(const std::shared_ptr &context, - AsyncGcsClient *client) - : Table(context, client), client_id_(UniqueID::from_random()) {} + // Add the notification to our cache. Notifications are idempotent. + client_cache_[client_id] = *data; -Status ClientTable::Connect(ClientID *client_id) { - auto data = std::make_shared(); - data->client_id = client_id_.binary(); - // TODO(swang): Get the address and port from somewhere. - data->node_manager_address = ""; - data->node_manager_port = 0; + // If the notification is new, call any registered callbacks. + if (is_new) { + if (data->is_insertion) { + if (client_added_callback_ != nullptr) { + client_added_callback_(client, client_id, data); + } + } else { + if (client_removed_callback_ != nullptr) { + client_removed_callback_(client, client_id, data); + } + } + } +} + +void ClientTable::HandleConnected(AsyncGcsClient *client, const ClientID &client_id, + std::shared_ptr data) { + RAY_CHECK(client_id == client_id_) << client_id.hex() << " " << client_id_.hex(); +} + +const ClientID &ClientTable::GetLocalClientId() { return client_id_; } + +const ClientTableDataT &ClientTable::GetLocalClient() { return local_client_; } + +Status ClientTable::Connect() { + RAY_CHECK(!disconnected_) << "Tried to reconnect a disconnected client."; + + auto data = std::make_shared(local_client_); data->is_insertion = true; - - // TODO(swang): - // - Add ourselves to the client table (easier if this is synchronous). - // - Subscribe to the client table. - // - Once subscription is complete, read all client table entries once. - - *client_id = client_id_; - return Status::OK(); + // Callback for a notification from the client table. + auto notification_callback = [this](AsyncGcsClient *client, const ClientID &channel_id, + std::shared_ptr data) { + return HandleNotification(client, channel_id, data); + }; + // Callback to handle our own successful connection once we've added + // ourselves. + auto add_callback = [this](AsyncGcsClient *client, const ClientID &id, + std::shared_ptr data) { + HandleConnected(client, id, data); + }; + // Callback to add ourselves once we've successfully subscribed. + auto subscription_callback = [this, data, add_callback]( + AsyncGcsClient *c, const ClientID &id, std::shared_ptr d) { + // Mark ourselves as deleted if we called Disconnect() since the last + // Connect() call. + if (disconnected_) { + data->is_insertion = false; + } + return Add(JobID::nil(), client_id_, data, add_callback); + }; + return Subscribe(JobID::nil(), ClientID::nil(), notification_callback, + subscription_callback); } Status ClientTable::Disconnect() { - auto data = std::make_shared(); - data->client_id = client_id_.binary(); - // TODO(swang): Get the address and port from somewhere. - data->node_manager_address = ""; - data->node_manager_port = 0; - data->is_insertion = false; - - // TODO(swang): - // - Add ourselves to the client table (easier if this is synchronous). + auto data = std::make_shared(local_client_); + data->is_insertion = true; + auto add_callback = [this](AsyncGcsClient *client, const ClientID &id, + std::shared_ptr data) { + HandleConnected(client, id, data); + }; + RAY_RETURN_NOT_OK(Add(JobID::nil(), client_id_, data, add_callback)); + // We successfully added the deletion entry. Mark ourselves as disconnected. + disconnected_ = true; return Status::OK(); } +const ClientTableDataT &ClientTable::GetClient(const ClientID &client_id) { + RAY_CHECK(!client_id.is_nil()); + auto entry = client_cache_.find(client_id); + if (entry != client_cache_.end()) { + return entry->second; + } else { + // If the requested client was not found, return a reference to the nil + // client entry. + return client_cache_[ClientID::nil()]; + } +} + } // namespace gcs } // namespace ray diff --git a/src/ray/gcs/tables.h b/src/ray/gcs/tables.h index 47a66c828..74d8712da 100644 --- a/src/ray/gcs/tables.h +++ b/src/ray/gcs/tables.h @@ -30,13 +30,16 @@ template class Table { public: using DataT = typename Data::NativeTableType; - using Callback = std::function< - void(AsyncGcsClient *client, const ID &id, std::shared_ptr data)>; + using Callback = std::function data)>; struct CallbackData { ID id; std::shared_ptr data; Callback callback; + // An optional callback to call for subscription operations, where the + // first message is a notification of subscription success. + Callback subscription_callback; Table *table; AsyncGcsClient *client; }; @@ -49,16 +52,19 @@ class Table { /// \param job_id The ID of the job (= driver). /// \param id The ID of the data that is added to the GCS. /// \param data Data that is added to the GCS. - /// \param done Callback that is called once the data has been written to the GCS. + /// \param done Callback that is called once the data has been written to the + /// GCS. /// \return Status - Status Add(const JobID &job_id, - const ID &id, - std::shared_ptr data, + Status Add(const JobID &job_id, const ID &id, std::shared_ptr data, const Callback &done) { auto d = std::shared_ptr( - new CallbackData({id, data, done, this, client_})); - int64_t callback_index = RedisCallbackManager::instance().add([d]( - const std::string &data) { (d->callback)(d->client, d->id, d->data); }); + new CallbackData({id, data, done, nullptr, this, client_})); + int64_t callback_index = + RedisCallbackManager::instance().add([d](const std::string &data) { + if (d->callback != nullptr) { + (d->callback)(d->client, d->id, d->data); + } + }); flatbuffers::FlatBufferBuilder fbb; fbb.ForceDefaults(true); fbb.Finish(Data::Pack(fbb, data.get())); @@ -75,13 +81,15 @@ class Table { /// \return Status Status Lookup(const JobID &job_id, const ID &id, const Callback &lookup) { auto d = std::shared_ptr( - new CallbackData({id, nullptr, lookup, this})); + new CallbackData({id, nullptr, lookup, nullptr, this, client_})); int64_t callback_index = RedisCallbackManager::instance().add([d](const std::string &data) { auto result = std::make_shared(); auto root = flatbuffers::GetRoot(data.data()); root->UnPackTo(result.get()); - (d->callback)(d->client, d->id, result); + if (d->callback != nullptr) { + (d->callback)(d->client, d->id, result); + } }); std::vector nil; RAY_RETURN_NOT_OK(context_->RunAsync("RAY.TABLE_LOOKUP", id, nil.data(), nil.size(), @@ -102,14 +110,17 @@ class Table { Status Subscribe(const JobID &job_id, const ClientID &client_id, const Callback &subscribe, const Callback &done) { auto d = std::shared_ptr( - new CallbackData({client_id, nullptr, subscribe, this})); + new CallbackData({client_id, nullptr, subscribe, done, this, client_})); int64_t callback_index = - RedisCallbackManager::instance().add([done, d](const std::string &data) { + RedisCallbackManager::instance().add([d](const std::string &data) { if (data.empty()) { // No data is provided. This is the callback for the initial // subscription request. - done(d->client, d->id, nullptr); + if (d->subscription_callback != nullptr) { + (d->subscription_callback)(d->client, d->id, nullptr); + } } else { + // Data is provided. This is the callback for a message. auto result = std::make_shared(); auto root = flatbuffers::GetRoot(data.data()); root->UnPackTo(result.get()); @@ -124,8 +135,7 @@ class Table { Status Remove(const JobID &job_id, const ID &id, const Callback &done); protected: - std::unordered_map, UniqueIDHasher> - callback_data_; + std::unordered_map, UniqueIDHasher> callback_data_; std::shared_ptr context_; AsyncGcsClient *client_; TablePubsub pubsub_channel_; @@ -149,10 +159,8 @@ class ObjectTable : public Table { /// \param done_callback Callback to be called when subscription is installed. /// This is only used for the tests. /// \return Status - Status SubscribeToNotifications(const JobID &job_id, - bool subscribe_all, - const Callback &object_available, - const Callback &done); + Status SubscribeToNotifications(const JobID &job_id, bool subscribe_all, + const Callback &object_available, const Callback &done); /// Request notifications about the availability of some objects from the /// object @@ -180,10 +188,9 @@ class TaskTable : public Table { pubsub_channel_ = TablePubsub_TASK; }; - using TestAndUpdateCallback = std::function; + using TestAndUpdateCallback = + std::function; using SubscribeToTaskCallback = std::function task)>; /// Update a task's scheduling information in the task table, if the current @@ -201,8 +208,7 @@ class TaskTable : public Table { /// with, if the current state matches test_state_bitmask. /// \param callback Function to be called when database returns result. /// \return Status - Status TestAndUpdate(const JobID &job_id, - const TaskID &id, + Status TestAndUpdate(const JobID &job_id, const TaskID &id, std::shared_ptr data, const TestAndUpdateCallback &callback) { int64_t callback_index = RedisCallbackManager::instance().add( @@ -213,7 +219,6 @@ class TaskTable : public Table { callback(client_, id, *result, root->updated()); }); flatbuffers::FlatBufferBuilder fbb; - TaskTableTestAndUpdateBuilder builder(fbb); fbb.Finish(TaskTableTestAndUpdate::Pack(fbb, data.get())); RAY_RETURN_NOT_OK(context_->RunAsync("RAY.TABLE_TEST_AND_UPDATE", id, fbb.GetBufferPointer(), fbb.GetSize(), @@ -231,7 +236,8 @@ class TaskTable : public Table { /// events we want to listen to. If you want to subscribe to updates /// from /// all local schedulers, pass in NIL_ID. - /// \param subscribe_callback Callback that will be called when the task table is + /// \param subscribe_callback Callback that will be called when the task table + /// is /// updated. /// \param state_filter Events we want to listen to. Can have values from the /// enum "scheduling_state" in task.h. @@ -257,50 +263,28 @@ Status TaskTableTestAndUpdate(AsyncGcsClient *gcs_client, const TaskID &task_id, SchedulingState update_state, const TaskTable::TestAndUpdateCallback &callback); -/// \class ClientInformation -/// -/// Represents information in the client table about a particular client. Each -/// client has an associated node manager. -class ClientInformation { - public: - /// Create a client information object. - /// - /// \param client_table_entry A serialized client table entry flatbuffer. - ClientInformation(const ClientTableData &client_table_entry); - - /// Get the client ID. - /// - /// \return The ID of this client. - const ClientID &GetClientId() const; - - /// Get the IP address of the client's node manager. - /// - /// \return The IP address of the client's node manager. - const std::string GetIpAddress() const; - - /// Get the port at which the client's node manager is listening for - /// TCP connections. - /// - /// \return The client's TCP port. - int GetPort() const; - - /// Get whether the client is alive. - /// - /// \return Whether the client is alive. - bool IsAlive() const; -}; - class ClientTable : private Table { public: - ClientTable(const std::shared_ptr &context, AsyncGcsClient *client); + ClientTable(const std::shared_ptr &context, AsyncGcsClient *client, + const ClientTableDataT &local_client) + : Table(context, client), + disconnected_(false), + client_id_(ClientID::from_binary(local_client.client_id)), + local_client_(local_client) { + pubsub_channel_ = TablePubsub_CLIENT; - /// Connect as a client to the GCS. This registers us in the client table and - /// begins subscription to client table notifications. + // Add a nil client to the cache so that we can serve requests for clients + // that we have not heard about. + ClientTableDataT nil_client; + nil_client.client_id = ClientID::nil().binary(); + client_cache_[ClientID::nil()] = nil_client; + }; + + /// Connect as a client to the GCS. This registers us in the client table + /// and begins subscription to client table notifications. /// - /// \param[out] client_id The assigned client ID will be written to this pointer. /// \return Status - // TODO(swang): Call this from AsyncGcsClient::Connect? - ray::Status Connect(ClientID *client_id); + ray::Status Connect(); /// Disconnect the client from the GCS. The client ID assigned during /// registration should never be reused after disconnecting. @@ -308,16 +292,54 @@ class ClientTable : private Table { /// \return Status ray::Status Disconnect(); - /// Get a client's information from the cache. + /// Register a callback to call when a new client is added. + /// + /// \param callback The callback to register. + void RegisterClientAddedCallback(const Callback &callback); + + /// Register a callback to call when a client is removed. + /// + /// \param callback The callback to register. + void RegisterClientRemovedCallback(const Callback &callback); + + /// Get a client's information from the cache. The cache only contains + /// information for clients that we've heard a notification for. /// /// \param client The client to get information about. - const ClientInformation &GetClientInformation(const ClientID &client); + /// \return A reference to the requested client. If the client is not in the + /// cache, then an entry with a nil ClientID will be returned. + const ClientTableDataT &GetClient(const ClientID &client); + + /// Get the local client's ID. + /// + /// \return The local client's ID. + const ClientID &GetLocalClientId(); + + /// Get the local client's information. + /// + /// \return The local client's information. + const ClientTableDataT &GetLocalClient(); private: + /// Handle a client table notification. + void HandleNotification(AsyncGcsClient *client, const ClientID &channel_id, + std::shared_ptr); + /// Handle this client's successful connection to the GCS. + void HandleConnected(AsyncGcsClient *client, const ClientID &client_id, + std::shared_ptr); + + /// Whether this client has called Disconnect(). + bool disconnected_; /// This client's ID. - ClientID client_id_; + const ClientID client_id_; + /// Information about this client. + ClientTableDataT local_client_; + /// The callback to call when a new client is added. + Callback client_added_callback_; + /// The callback to call when a client is removed. + Callback client_removed_callback_; /// A cache for information about all clients. - std::unordered_map client_cache_; + std::unordered_map client_cache_; }; } // namespace gcs diff --git a/src/ray/gcs/task_table.cc b/src/ray/gcs/task_table.cc new file mode 100644 index 000000000..a60ab148e --- /dev/null +++ b/src/ray/gcs/task_table.cc @@ -0,0 +1,68 @@ +#include "ray/gcs/tables.h" + +#include "ray/gcs/client.h" + +#include "common_protocol.h" +#include "task.h" + +// TODO(swang): This file extends tables.cc so that we can separate out the +// part that depends on the Task* datasturcture from the build. This should be +// merged with tables.cc once we get rid of the Task* datastructure. + +namespace { + +std::shared_ptr MakeTaskTableData(const TaskExecutionSpec &execution_spec, + const ClientID &local_scheduler_id, + SchedulingState scheduling_state) { + auto data = std::make_shared(); + data->scheduling_state = scheduling_state; + data->task_info = std::string(execution_spec.Spec(), execution_spec.SpecSize()); + data->scheduler_id = local_scheduler_id.binary(); + + flatbuffers::FlatBufferBuilder fbb; + auto execution_dependencies = CreateTaskExecutionDependencies( + fbb, to_flatbuf(fbb, execution_spec.ExecutionDependencies())); + fbb.Finish(execution_dependencies); + + data->execution_dependencies = + std::string((const char *)fbb.GetBufferPointer(), fbb.GetSize()); + data->spillback_count = execution_spec.SpillbackCount(); + + return data; +} + +} // namespace + +namespace ray { + +namespace gcs { + +// TODO(pcm): This is a helper method that should go away once we get rid of +// the Task* datastructure and replace it with TaskTableDataT. +Status TaskTableAdd(AsyncGcsClient *gcs_client, Task *task) { + TaskExecutionSpec &execution_spec = *Task_task_execution_spec(task); + TaskSpec *spec = execution_spec.Spec(); + auto data = MakeTaskTableData(execution_spec, Task_local_scheduler(task), + static_cast(Task_state(task))); + return gcs_client->task_table().Add(ray::JobID::nil(), TaskSpec_task_id(spec), data, + [](gcs::AsyncGcsClient *client, const TaskID &id, + std::shared_ptr data) {}); +} + +// TODO(pcm): This is a helper method that should go away once we get rid of +// the Task* datastructure and replace it with TaskTableDataT. +Status TaskTableTestAndUpdate(AsyncGcsClient *gcs_client, const TaskID &task_id, + const ClientID &local_scheduler_id, int test_state_bitmask, + SchedulingState update_state, + const TaskTable::TestAndUpdateCallback &callback) { + auto data = std::make_shared(); + data->test_scheduler_id = local_scheduler_id.binary(); + data->test_state_bitmask = test_state_bitmask; + data->update_state = update_state; + return gcs_client->task_table().TestAndUpdate(ray::JobID::nil(), task_id, data, + callback); +} + +} // namespace gcs + +} // namespace ray