From 0a6edb55a86cc97bf2714bb81bcb87e383aed27e Mon Sep 17 00:00:00 2001 From: Stephanie Wang Date: Tue, 6 Mar 2018 09:56:12 -0800 Subject: [PATCH] Implement the Subscribe call for the new GCS API (#1652) * Implement the Subscribe call for the new GCS API * Document tests * Upper case function name * Fix build errors * lint --- src/common/redis_module/ray_redis_module.cc | 114 ++++++++++++-------- src/ray/gcs/client_test.cc | 43 ++++++++ src/ray/gcs/format/gcs.fbs | 9 ++ src/ray/gcs/redis_context.cc | 96 ++++++++++++++++- src/ray/gcs/redis_context.h | 10 ++ src/ray/gcs/tables.h | 50 +++++++-- src/ray/id.h | 2 + 7 files changed, 266 insertions(+), 58 deletions(-) diff --git a/src/common/redis_module/ray_redis_module.cc b/src/common/redis_module/ray_redis_module.cc index e65f6256c..9e05ba32d 100644 --- a/src/common/redis_module/ray_redis_module.cc +++ b/src/common/redis_module/ray_redis_module.cc @@ -394,58 +394,81 @@ bool PublishObjectNotification(RedisModuleCtx *ctx, int TableAdd_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { - if (argc != 3) { + if (argc != 4) { return RedisModule_WrongArity(ctx); } - RedisModuleString *id = argv[1]; - RedisModuleString *data = argv[2]; + RedisModuleString *pubsub_channel_str = argv[1]; + RedisModuleString *id = argv[2]; + RedisModuleString *data = argv[3]; - // Set the keys in the table + // Set the keys in the table. RedisModuleKey *key = OpenPrefixedKey(ctx, "T:", id, REDISMODULE_READ | REDISMODULE_WRITE); RedisModule_StringSet(key, data); RedisModule_CloseKey(key); - size_t len = 0; - const char *buf = RedisModule_StringPtrLen(data, &len); + // Get the requested pubsub channel. + long long pubsub_channel_long; + RAY_CHECK(RedisModule_StringToLongLong( + pubsub_channel_str, &pubsub_channel_long) == REDISMODULE_OK) + << "Pubsub channel must be a valid TablePubsub"; + auto pubsub_channel = static_cast(pubsub_channel_long); + RAY_CHECK(pubsub_channel >= TablePubsub_MIN && + pubsub_channel <= TablePubsub_MAX) + << "Pubsub channel must be a valid TablePubsub"; - auto message = flatbuffers::GetRoot(buf); + // 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); - if (message->scheduling_state() == SchedulingState_WAITING || - message->scheduling_state() == SchedulingState_SCHEDULED) { - /* Build the PUBLISH topic and message for task table subscribers. The topic - * is a string in the format "TASK_PREFIX::". The - * message is a serialized SubscribeToTasksReply flatbuffer object. */ - std::string state = std::to_string(message->scheduling_state()); - RedisModuleString *publish_topic = RedisString_Format( - ctx, "%s%b:%s", TASK_PREFIX, message->scheduler_id()->str().data(), - sizeof(DBClientID), state.c_str()); + auto message = flatbuffers::GetRoot(buf); - /* Construct the flatbuffers object for the payload. */ - flatbuffers::FlatBufferBuilder fbb; - /* Create the flatbuffers message. */ - auto msg = CreateTaskReply( - fbb, RedisStringToFlatbuf(fbb, id), message->scheduling_state(), - fbb.CreateString(message->scheduler_id()), - fbb.CreateString(message->execution_dependencies()), - fbb.CreateString(message->task_info()), message->spillback_count(), - true /* not used */); - fbb.Finish(msg); + if (message->scheduling_state() == SchedulingState_WAITING || + message->scheduling_state() == SchedulingState_SCHEDULED) { + /* Build the PUBLISH topic and message for task table subscribers. The + * topic + * is a string in the format "TASK_PREFIX::". + * The + * message is a serialized SubscribeToTasksReply flatbuffer object. */ + std::string state = std::to_string(message->scheduling_state()); + RedisModuleString *publish_topic = RedisString_Format( + ctx, "%s%b:%s", TASK_PREFIX, message->scheduler_id()->str().data(), + sizeof(DBClientID), state.c_str()); - RedisModuleString *publish_message = RedisModule_CreateString( - ctx, (const char *) fbb.GetBufferPointer(), fbb.GetSize()); + /* Construct the flatbuffers object for the payload. */ + flatbuffers::FlatBufferBuilder fbb; + /* Create the flatbuffers message. */ + auto msg = CreateTaskReply( + fbb, RedisStringToFlatbuf(fbb, id), message->scheduling_state(), + fbb.CreateString(message->scheduler_id()), + fbb.CreateString(message->execution_dependencies()), + fbb.CreateString(message->task_info()), message->spillback_count(), + true /* not used */); + fbb.Finish(msg); + RedisModuleString *publish_message = RedisModule_CreateString( + ctx, (const char *) fbb.GetBufferPointer(), fbb.GetSize()); + + RedisModuleCallReply *reply = RedisModule_Call( + ctx, "PUBLISH", "ss", publish_topic, publish_message); + + /* See how many clients received this publish. */ + long long num_clients = RedisModule_CallReplyInteger(reply); + RAY_CHECK(num_clients <= 1) << "Published to " << num_clients + << " clients."; + + RedisModule_FreeString(ctx, publish_message); + RedisModule_FreeString(ctx, publish_topic); + } + } else if (pubsub_channel != TablePubsub_NO_PUBLISH) { + // All other pubsub channels write the data back directly onto the channel. RedisModuleCallReply *reply = - RedisModule_Call(ctx, "PUBLISH", "ss", publish_topic, publish_message); - - /* See how many clients received this publish. */ - long long num_clients = RedisModule_CallReplyInteger(reply); - RAY_CHECK(num_clients <= 1) << "Published to " << num_clients - << " clients."; - - RedisModule_FreeString(ctx, publish_message); - RedisModule_FreeString(ctx, publish_topic); + RedisModule_Call(ctx, "PUBLISH", "ss", pubsub_channel_str, data); + if (reply == NULL) { + RedisModule_ReplyWithError(ctx, "error during PUBLISH"); + } } return RedisModule_ReplyWithSimpleString(ctx, "OK"); @@ -456,11 +479,11 @@ int TableAdd_RedisCommand(RedisModuleCtx *ctx, int TableLookup_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { - if (argc != 2) { + if (argc != 3) { return RedisModule_WrongArity(ctx); } - RedisModuleString *id = argv[1]; + RedisModuleString *id = argv[2]; RedisModuleKey *key = OpenPrefixedKey(ctx, "T:", id, REDISMODULE_READ); size_t len = 0; @@ -490,11 +513,11 @@ bool is_nil(const std::string &data) { int TableTestAndUpdate_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { - if (argc != 3) { + if (argc != 4) { return RedisModule_WrongArity(ctx); } - RedisModuleString *id = argv[1]; - RedisModuleString *update_data = argv[2]; + RedisModuleString *id = argv[2]; + RedisModuleString *update_data = argv[3]; RedisModuleKey *key = OpenPrefixedKey(ctx, "T:", id, REDISMODULE_READ | REDISMODULE_WRITE); @@ -1060,9 +1083,10 @@ int TaskTableWrite(RedisModuleCtx *ctx, if (state_value == TASK_STATUS_WAITING || state_value == TASK_STATUS_SCHEDULED) { - /* Build the PUBLISH topic and message for task table subscribers. The topic - * is a string in the format "TASK_PREFIX::". The - * message is a serialized SubscribeToTasksReply flatbuffer object. */ + /* Build the PUBLISH topic and message for task table subscribers. The + * topic is a string in the format + * "TASK_PREFIX::". The message is a serialized + * SubscribeToTasksReply flatbuffer object. */ RedisModuleString *publish_topic = RedisString_Format( ctx, "%s%S:%S", TASK_PREFIX, local_scheduler_id, state); diff --git a/src/ray/gcs/client_test.cc b/src/ray/gcs/client_test.cc index b717f7f77..47d9ffbe7 100644 --- a/src/ray/gcs/client_test.cc +++ b/src/ray/gcs/client_test.cc @@ -65,6 +65,7 @@ void ObjectAdded(gcs::AsyncGcsClient *client, 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(); } @@ -77,6 +78,8 @@ void TestObjectTable(const UniqueID &job_id, gcs::AsyncGcsClient &client) { 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(); } @@ -130,8 +133,12 @@ void TestTaskTable(const UniqueID &job_id, gcs::AsyncGcsClient &client) { update->test_scheduler_id = local_scheduler_id.binary(); update->test_state_bitmask = SchedulingState_SCHEDULED; update->update_state = SchedulingState_LOST; + // 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)); + // Run the event loop. The loop will only stop if the lookup after the + // test-and-set succeeds (or an assertion failure). test->Start(); } @@ -145,4 +152,40 @@ TEST_F(TestGcsWithAsio, TestTaskTable) { TestTaskTable(job_id_, client_); } +void ObjectTableSubscribed(gcs::AsyncGcsClient *client, + const UniqueID &id, + std::shared_ptr data) { + test->Stop(); +} + +void TestSubscribeAll(const UniqueID &job_id, gcs::AsyncGcsClient &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)); + // Run the event loop. The loop will only stop if the subscription succeeds. + test->Start(); + + // We have subscribed. Add an object table entry. + 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)); + // Run the event loop. The loop will only stop if the registered subscription + // callback is called (or an assertion failure). + test->Start(); +} + +TEST_F(TestGcsWithAe, TestSubscribeAll) { + test = this; + TestSubscribeAll(job_id_, client_); +} + +TEST_F(TestGcsWithAsio, TestSubscribeAll) { + test = this; + TestSubscribeAll(job_id_, client_); +} + } // namespace diff --git a/src/ray/gcs/format/gcs.fbs b/src/ray/gcs/format/gcs.fbs index 3ee10aa4c..0784e96e0 100644 --- a/src/ray/gcs/format/gcs.fbs +++ b/src/ray/gcs/format/gcs.fbs @@ -4,6 +4,15 @@ enum Language:int { JAVA = 2 } +// The channel that Add operations to the Table should be published on, if any. +enum TablePubsub:int { + NO_PUBLISH = 0, + TASK, + CLIENT, + OBJECT, + ACTOR +} + table FunctionTableData { language: Language; name: string; diff --git a/src/ray/gcs/redis_context.cc b/src/ray/gcs/redis_context.cc index bbefa78ce..27a27177a 100644 --- a/src/ray/gcs/redis_context.cc +++ b/src/ray/gcs/redis_context.cc @@ -28,6 +28,9 @@ void GlobalRedisCallback(void *c, void *r, void *privdata) { if (reply->type == REDIS_REPLY_NIL) { } else if (reply->type == REDIS_REPLY_STRING) { data = std::string(reply->str, reply->len); + } else if (reply->type == REDIS_REPLY_ARRAY) { + reply = reply->element[reply->elements - 1]; + data = std::string(reply->str, reply->len); } else if (reply->type == REDIS_REPLY_STATUS) { } else if (reply->type == REDIS_REPLY_ERROR) { RAY_LOG(ERROR) << "Redis error " << reply->str; @@ -36,6 +39,42 @@ void GlobalRedisCallback(void *c, void *r, void *privdata) { << " and with string " << reply->str; } RedisCallbackManager::instance().get(callback_index)(data); + // Delete the callback. + RedisCallbackManager::instance().remove(callback_index); +} + +void SubscribeRedisCallback(void *c, void *r, void *privdata) { + if (r == NULL) { + return; + } + int64_t callback_index = reinterpret_cast(privdata); + redisReply *reply = reinterpret_cast(r); + std::string data = ""; + if (reply->type == REDIS_REPLY_ARRAY) { + // Parse the message. + redisReply *message_type = reply->element[0]; + if (strcmp(message_type->str, "subscribe") == 0) { + // If the message is for the initial subscription call, do not fill in + // data. + } else if (strcmp(message_type->str, "message") == 0) { + // If the message is from a PUBLISH, make sure the data is nonempty. + redisReply *message = reply->element[reply->elements - 1]; + data = std::string(message->str, message->len); + RAY_CHECK(!data.empty()) << "Empty message received on subscribe channel"; + } else { + RAY_LOG(FATAL) << "Fatal redis error during subscribe" + << message_type->str; + } + + // NOTE(swang): We do not delete the callback after calling it since there + // may be more subscription messages. + RedisCallbackManager::instance().get(callback_index)(data); + } else if (reply->type == REDIS_REPLY_ERROR) { + RAY_LOG(ERROR) << "Redis error " << reply->str; + } else { + RAY_LOG(FATAL) << "Fatal redis error of type " << reply->type + << " and with string " << reply->str; + } } int64_t RedisCallbackManager::add(const RedisCallback &function) { @@ -49,6 +88,10 @@ RedisCallbackManager::RedisCallback &RedisCallbackManager::get( return *callbacks_[callback_index]; } +void RedisCallbackManager::remove(int64_t callback_index) { + callbacks_.erase(callback_index); +} + #define REDIS_CHECK_ERROR(CONTEXT, REPLY) \ if (REPLY == nullptr || REPLY->type == REDIS_REPLY_ERROR) { \ return Status::RedisError(CONTEXT->errstr); \ @@ -61,6 +104,9 @@ RedisContext::~RedisContext() { if (async_context_) { redisAsyncFree(async_context_); } + if (subscribe_context_) { + redisAsyncFree(subscribe_context_); + } } Status RedisContext::Connect(const std::string &address, int port) { @@ -95,11 +141,18 @@ Status RedisContext::Connect(const std::string &address, int port) { RAY_LOG(FATAL) << "Could not establish connection to redis " << address << ":" << port; } + // Connect to subscribe context + subscribe_context_ = redisAsyncConnect(address.c_str(), port); + if (subscribe_context_ == nullptr || subscribe_context_->err) { + RAY_LOG(FATAL) << "Could not establish subscribe connection to redis " + << address << ":" << port; + } return Status::OK(); } Status RedisContext::AttachToEventLoop(aeEventLoop *loop) { - if (redisAeAttach(loop, async_context_) != REDIS_OK) { + if (redisAeAttach(loop, async_context_) != REDIS_OK || + redisAeAttach(loop, subscribe_context_) != REDIS_OK) { return Status::RedisError("could not attach redis event loop"); } else { return Status::OK(); @@ -110,24 +163,25 @@ Status RedisContext::RunAsync(const std::string &command, const UniqueID &id, uint8_t *data, int64_t length, + const TablePubsub pubsub_channel, int64_t callback_index) { if (length > 0) { - std::string redis_command = command + " %b %b"; + std::string redis_command = command + " %d %b %b"; int status = redisAsyncCommand( async_context_, reinterpret_cast(&GlobalRedisCallback), reinterpret_cast(callback_index), redis_command.c_str(), - id.data(), id.size(), data, length); + pubsub_channel, id.data(), id.size(), data, length); if (status == REDIS_ERR) { return Status::RedisError(std::string(async_context_->errstr)); } } else { - std::string redis_command = command + " %b"; + std::string redis_command = command + " %d %b"; int status = redisAsyncCommand( async_context_, reinterpret_cast(&GlobalRedisCallback), reinterpret_cast(callback_index), redis_command.c_str(), - id.data(), id.size()); + pubsub_channel, id.data(), id.size()); if (status == REDIS_ERR) { return Status::RedisError(std::string(async_context_->errstr)); } @@ -135,6 +189,38 @@ Status RedisContext::RunAsync(const std::string &command, return Status::OK(); } +Status RedisContext::SubscribeAsync(const ClientID &client_id, + const TablePubsub pubsub_channel, + int64_t callback_index) { + RAY_CHECK(pubsub_channel != TablePubsub_NO_PUBLISH) + << "Client requested subscribe on a table that does not support pubsub"; + + int status = 0; + if (client_id.is_nil()) { + // Subscribe to all messages. + std::string redis_command = "SUBSCRIBE %d"; + status = redisAsyncCommand( + subscribe_context_, + reinterpret_cast(&SubscribeRedisCallback), + reinterpret_cast(callback_index), redis_command.c_str(), + pubsub_channel); + } else { + // Subscribe only to messages sent to this client. + // TODO(swang): Nobody sends on this channel yet. + std::string redis_command = "SUBSCRIBE %d:%b"; + status = redisAsyncCommand( + subscribe_context_, + reinterpret_cast(&SubscribeRedisCallback), + reinterpret_cast(callback_index), redis_command.c_str(), + pubsub_channel, client_id.data(), client_id.size()); + } + + if (status == REDIS_ERR) { + return Status::RedisError(std::string(subscribe_context_->errstr)); + } + return Status::OK(); +} + } // namespace gcs } // namespace ray diff --git a/src/ray/gcs/redis_context.h b/src/ray/gcs/redis_context.h index a731335b1..d38313665 100644 --- a/src/ray/gcs/redis_context.h +++ b/src/ray/gcs/redis_context.h @@ -9,6 +9,8 @@ #include "ray/status.h" #include "ray/util/logging.h" +#include "ray/gcs/format/gcs_generated.h" + struct redisContext; struct redisAsyncContext; struct aeEventLoop; @@ -30,6 +32,9 @@ class RedisCallbackManager { RedisCallback &get(int64_t callback_index); + /// Remove a callback. + void remove(int64_t callback_index); + private: RedisCallbackManager() : num_callbacks(0){}; @@ -49,12 +54,17 @@ class RedisContext { const UniqueID &id, uint8_t *data, int64_t length, + const TablePubsub pubsub_channel, int64_t callback_index); + Status SubscribeAsync(const ClientID &client_id, + const TablePubsub pubsub_channel, + int64_t callback_index); redisAsyncContext *async_context() { return async_context_; } private: redisContext *context_; redisAsyncContext *async_context_; + redisAsyncContext *subscribe_context_; }; } // namespace gcs diff --git a/src/ray/gcs/tables.h b/src/ray/gcs/tables.h index 7e67dd0b3..6d86d73f0 100644 --- a/src/ray/gcs/tables.h +++ b/src/ray/gcs/tables.h @@ -42,7 +42,9 @@ class Table { }; Table(const std::shared_ptr &context, AsyncGcsClient *client) - : context_(context), client_(client){}; + : context_(context), + client_(client), + pubsub_channel_(TablePubsub_NO_PUBLISH){}; /// Add an entry to the table. /// @@ -65,7 +67,7 @@ class Table { fbb.Finish(Data::Pack(fbb, data.get())); RAY_RETURN_NOT_OK(context_->RunAsync("RAY.TABLE_ADD", id, fbb.GetBufferPointer(), fbb.GetSize(), - callback_index)); + pubsub_channel_, callback_index)); return Status::OK(); } @@ -87,16 +89,42 @@ class Table { }); std::vector nil; RAY_RETURN_NOT_OK(context_->RunAsync("RAY.TABLE_LOOKUP", id, nil.data(), - nil.size(), callback_index)); + nil.size(), pubsub_channel_, + callback_index)); return Status::OK(); } /// Subscribe to updates of this table + /// + /// @param job_id The ID of the job (= driver). + /// @param client_id The type of update to listen to. If this is nil, then a + /// message for each Add to the table will be received. Else, only + /// messages for the given client will be received. + /// @param subscribe Callback that is called on each received message. + /// @param done Callback that is called when subscription is complete and we + /// are ready to receive messages.. + /// @return Status Status Subscribe(const JobID &job_id, - const ID &id, + const ClientID &client_id, const Callback &subscribe, const Callback &done) { - return Status::NotImplemented("Table::Subscribe is not implemented"); + auto d = std::shared_ptr( + new CallbackData({client_id, nullptr, subscribe, this})); + int64_t callback_index = RedisCallbackManager::instance().add( + [done, 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); + } else { + auto result = std::make_shared(); + auto root = flatbuffers::GetRoot(data.data()); + root->UnPackTo(result.get()); + (d->callback)(d->client, d->id, result); + } + }); + std::vector nil; + return context_->SubscribeAsync(client_id, pubsub_channel_, callback_index); } /// Remove and entry from the table @@ -107,13 +135,16 @@ class Table { callback_data_; std::shared_ptr context_; AsyncGcsClient *client_; + TablePubsub pubsub_channel_; }; class ObjectTable : public Table { public: ObjectTable(const std::shared_ptr &context, AsyncGcsClient *client) - : Table(context, client){}; + : Table(context, client) { + pubsub_channel_ = TablePubsub_OBJECT; + }; /// Set up a client-specific channel for receiving notifications about /// available @@ -147,13 +178,16 @@ using FunctionTable = Table; using ClassTable = Table; +// TODO(swang): Set the pubsub channel for the actor table. using ActorTable = Table; class TaskTable : public Table { public: TaskTable(const std::shared_ptr &context, AsyncGcsClient *client) - : Table(context, client){}; + : Table(context, client) { + pubsub_channel_ = TablePubsub_TASK; + }; using TestAndUpdateCallback = std::function { fbb.Finish(TaskTableTestAndUpdate::Pack(fbb, data.get())); RAY_RETURN_NOT_OK(context_->RunAsync("RAY.TABLE_TEST_AND_UPDATE", id, fbb.GetBufferPointer(), fbb.GetSize(), - callback_index)); + pubsub_channel_, callback_index)); return Status::OK(); } diff --git a/src/ray/id.h b/src/ray/id.h index 97761e04f..603e1fc8f 100644 --- a/src/ray/id.h +++ b/src/ray/id.h @@ -55,7 +55,9 @@ typedef UniqueID ActorID; typedef UniqueID ActorHandleID; typedef UniqueID WorkerID; typedef UniqueID DriverID; +// TODO(swang): Replace this with ClientID. typedef UniqueID DBClientID; +typedef UniqueID ClientID; typedef UniqueID ConfigID; } // namespace ray