wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-12 10:41:38 +08:00
Code Issues Packages Projects Releases Wiki Activity

Implement the Subscribe call for the new GCS API (#1652)

Browse Source 
* Implement the Subscribe call for the new GCS API

* Document tests

* Upper case function name

* Fix build errors

* lint
This commit is contained in:
Stephanie Wang
2018-03-06 09:56:12 -08:00
committed by Philipp Moritz
parent 936bebef99
commit 0a6edb55a8
7 changed files with 266 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/common/redis_module/ray_redis_module.cc
+69 -45
View File
@@ -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<TablePubsub>(pubsub_channel_long);
RAY_CHECK(pubsub_channel >= TablePubsub_MIN &&
pubsub_channel <= TablePubsub_MAX)
<< "Pubsub channel must be a valid TablePubsub";
auto message = flatbuffers::GetRoot<TaskTableData>(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:<local scheduler ID>:<state>". 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<TaskTableData>(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:<local scheduler ID>:<state>".
* 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:<local scheduler ID>:<state>". 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:<local scheduler ID>:<state>". The message is a serialized
* SubscribeToTasksReply flatbuffer object. */
RedisModuleString *publish_topic = RedisString_Format(
ctx, "%s%S:%S", TASK_PREFIX, local_scheduler_id, state);
src/ray/gcs/client_test.cc
+43
View File
@@ -65,6 +65,7 @@ void ObjectAdded(gcs::AsyncGcsClient *client,
void Lookup(gcs::AsyncGcsClient *client,
const UniqueID &id,
std::shared_ptr<ObjectTableDataT> data) {
// Check that the object entry was added.
ASSERT_EQ(data->managers, std::vector<std::string>({"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<ObjectTableDataT> 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<ObjectTableDataT>();
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
src/ray/gcs/format/gcs.fbs
+9
View File
@@ -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;
src/ray/gcs/redis_context.cc
+91 -5
View File
@@ -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<int64_t>(privdata);
redisReply *reply = reinterpret_cast<redisReply *>(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<redisCallbackFn *>(&GlobalRedisCallback),
reinterpret_cast<void *>(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<redisCallbackFn *>(&GlobalRedisCallback),
reinterpret_cast<void *>(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<redisCallbackFn *>(&SubscribeRedisCallback),
reinterpret_cast<void *>(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<redisCallbackFn *>(&SubscribeRedisCallback),
reinterpret_cast<void *>(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
src/ray/gcs/redis_context.h
+10
View File
@@ -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
src/ray/gcs/tables.h
+42 -8
View File
@@ -42,7 +42,9 @@ class Table {
};
Table(const std::shared_ptr<RedisContext> &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<uint8_t> 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<CallbackData>(
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<DataT>();
auto root = flatbuffers::GetRoot<Data>(data.data());
root->UnPackTo(result.get());
(d->callback)(d->client, d->id, result);
}
});
std::vector<uint8_t> 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<RedisContext> context_;
AsyncGcsClient *client_;
TablePubsub pubsub_channel_;
};
class ObjectTable : public Table<ObjectID, ObjectTableData> {
public:
ObjectTable(const std::shared_ptr<RedisContext> &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<FunctionID, FunctionTableData>;
using ClassTable = Table<ClassID, ClassTableData>;
// TODO(swang): Set the pubsub channel for the actor table.
using ActorTable = Table<ActorID, ActorTableData>;
class TaskTable : public Table<TaskID, TaskTableData> {
public:
TaskTable(const std::shared_ptr<RedisContext> &context,
AsyncGcsClient *client)
: Table(context, client){};
: Table(context, client) {
pubsub_channel_ = TablePubsub_TASK;
};
using TestAndUpdateCallback = std::function<void(AsyncGcsClient *client,
const TaskID &id,
@@ -192,7 +226,7 @@ class TaskTable : public Table<TaskID, TaskTableData> {
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();
}
src/ray/id.h
+2
View File
@@ -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