mirror of
https://github.com/wassname/ray.git
synced 2026-07-24 13:20:22 +08:00
Integrate credis with Ray & route task table entries into credis. (#1841)
This commit is contained in:
committed by
Robert Nishihara
parent
3509a33cf3
commit
fa97acbc89
@@ -0,0 +1,61 @@
|
||||
#ifndef RAY_CHAIN_MODULE_H_
|
||||
#define RAY_CHAIN_MODULE_H_
|
||||
|
||||
#include <functional>
|
||||
|
||||
#include "redismodule.h"
|
||||
|
||||
// NOTE(zongheng): this duplicated declaration serves as forward-declaration
|
||||
// only. The implementation is supposed to be linked in from credis. In
|
||||
// principle, we can expose a header from credis and simple include that header.
|
||||
// This is left as future work.
|
||||
//
|
||||
// Concrete definitions from credis (from an example commit):
|
||||
// https://github.com/ray-project/credis/blob/7eae7f2e58d16dfa1a95b5dfab02549f54b94e5d/src/member.cc#L41
|
||||
// https://github.com/ray-project/credis/blob/7eae7f2e58d16dfa1a95b5dfab02549f54b94e5d/src/master.cc#L36
|
||||
|
||||
// Typical usage to make an existing redismodule command chain-compatible:
|
||||
//
|
||||
// extern RedisChainModule module;
|
||||
// int MyCmd_RedisModuleCmd(...) {
|
||||
// return module.Mutate(..., NodeFunc, TailFunc);
|
||||
// }
|
||||
//
|
||||
// See, for instance, ChainTableAdd_RedisCommand in ray_redis_module.cc.
|
||||
class RedisChainModule {
|
||||
public:
|
||||
// A function that runs on every node in the chain. Type:
|
||||
// (context, argv, argc, (can be nullptr) mutated_key_str) -> int
|
||||
//
|
||||
// (Advanced) The optional fourth arg can be used in the following way:
|
||||
//
|
||||
// RedisModuleString* redis_key_str = nullptr;
|
||||
// node_func(ctx, argv, argc, &redis_key_str);
|
||||
// // "redis_key_str" now points to the RedisModuleString whose contents
|
||||
// // is mutated by "node_func".
|
||||
//
|
||||
// If the fourth arg is passed, NodeFunc *must* fill in the key being mutated.
|
||||
// It is okay for this NodeFunc to call "RM_FreeString(mutated_key_str)" after
|
||||
// assigning the fourth arg, since that call presumably only decrements a ref
|
||||
// count.
|
||||
using NodeFunc = std::function<
|
||||
int(RedisModuleCtx *, RedisModuleString **, int, RedisModuleString **)>;
|
||||
|
||||
// A function that (1) runs only after all NodeFunc's have run, and (2) runs
|
||||
// once on the tail. A typical usage is to publish a write.
|
||||
using TailFunc =
|
||||
std::function<int(RedisModuleCtx *, RedisModuleString **, int)>;
|
||||
|
||||
// TODO(zongheng): document the RM_Reply semantics.
|
||||
|
||||
// Runs "node_func" on every node in the chain; after the tail node has run it
|
||||
// too, finalizes the mutation by running "tail_func".
|
||||
// TODO(zongheng): currently only supports 1-node chain.
|
||||
int ChainReplicate(RedisModuleCtx *ctx,
|
||||
RedisModuleString **argv,
|
||||
int argc,
|
||||
NodeFunc node_func,
|
||||
TailFunc tail_func);
|
||||
};
|
||||
|
||||
#endif // RAY_CHAIN_MODULE_H_
|
||||
@@ -1,14 +1,26 @@
|
||||
#include "redismodule.h"
|
||||
#include <string.h>
|
||||
|
||||
#include "redis_string.h"
|
||||
|
||||
#include "common_protocol.h"
|
||||
#include "format/common_generated.h"
|
||||
#include "ray/gcs/format/gcs_generated.h"
|
||||
#include "ray/id.h"
|
||||
#include "redis_string.h"
|
||||
#include "redismodule.h"
|
||||
#include "task.h"
|
||||
|
||||
#if RAY_USE_NEW_GCS
|
||||
// Under this flag, ray-project/credis will be loaded. Specifically, via
|
||||
// "path/redis-server --loadmodule <credis module> --loadmodule <current
|
||||
// libray_redis_module>" (dlopen() under the hood) will a definition of "module"
|
||||
// be supplied.
|
||||
//
|
||||
// All commands in this file that depend on "module" must be wrapped by "#if
|
||||
// RAY_USE_NEW_GCS", until we switch to this launch configuration as the
|
||||
// default.
|
||||
#include "chain_module.h"
|
||||
extern RedisChainModule module;
|
||||
#endif
|
||||
|
||||
// Various tables are maintained in redis:
|
||||
//
|
||||
// == OBJECT TABLE ==
|
||||
@@ -88,9 +100,14 @@ RedisModuleString *FormatPubsubChannel(
|
||||
RedisModuleKey *OpenPrefixedKey(RedisModuleCtx *ctx,
|
||||
const char *prefix,
|
||||
RedisModuleString *keyname,
|
||||
int mode) {
|
||||
int mode,
|
||||
RedisModuleString **mutated_key_str) {
|
||||
RedisModuleString *prefixed_keyname =
|
||||
RedisString_Format(ctx, "%s%S", prefix, keyname);
|
||||
// Pass out the key being mutated, should the caller request so.
|
||||
if (mutated_key_str != nullptr) {
|
||||
*mutated_key_str = prefixed_keyname;
|
||||
}
|
||||
RedisModuleKey *key =
|
||||
(RedisModuleKey *) RedisModule_OpenKey(ctx, prefixed_keyname, mode);
|
||||
return key;
|
||||
@@ -99,7 +116,8 @@ RedisModuleKey *OpenPrefixedKey(RedisModuleCtx *ctx,
|
||||
RedisModuleKey *OpenPrefixedKey(RedisModuleCtx *ctx,
|
||||
RedisModuleString *prefix_enum,
|
||||
RedisModuleString *keyname,
|
||||
int mode) {
|
||||
int mode,
|
||||
RedisModuleString **mutated_key_str) {
|
||||
long long prefix_long;
|
||||
RAY_CHECK(RedisModule_StringToLongLong(prefix_enum, &prefix_long) ==
|
||||
REDISMODULE_OK)
|
||||
@@ -109,7 +127,24 @@ RedisModuleKey *OpenPrefixedKey(RedisModuleCtx *ctx,
|
||||
<< "This table has no prefix registered";
|
||||
RAY_CHECK(prefix >= TablePrefix_MIN && prefix <= TablePrefix_MAX)
|
||||
<< "Prefix must be a valid TablePrefix";
|
||||
return OpenPrefixedKey(ctx, table_prefixes[prefix], keyname, mode);
|
||||
return OpenPrefixedKey(ctx, table_prefixes[prefix], keyname, mode,
|
||||
mutated_key_str);
|
||||
}
|
||||
|
||||
RedisModuleKey *OpenPrefixedKey(RedisModuleCtx *ctx,
|
||||
const char *prefix,
|
||||
RedisModuleString *keyname,
|
||||
int mode) {
|
||||
return OpenPrefixedKey(ctx, prefix, keyname, mode,
|
||||
/*mutated_key_str=*/nullptr);
|
||||
}
|
||||
|
||||
RedisModuleKey *OpenPrefixedKey(RedisModuleCtx *ctx,
|
||||
RedisModuleString *prefix_enum,
|
||||
RedisModuleString *keyname,
|
||||
int mode) {
|
||||
return OpenPrefixedKey(ctx, prefix_enum, keyname, mode,
|
||||
/*mutated_key_str=*/nullptr);
|
||||
}
|
||||
|
||||
/// Open the key used to store the channels that should be published to when an
|
||||
@@ -444,11 +479,12 @@ bool PublishObjectNotification(RedisModuleCtx *ctx,
|
||||
|
||||
// NOTE(pcmoritz): This is a temporary redis command that will be removed once
|
||||
// the GCS uses https://github.com/pcmoritz/credis.
|
||||
int TaskTableAdd(RedisModuleCtx *ctx,
|
||||
RedisModuleString *id,
|
||||
RedisModuleString *data) {
|
||||
int PublishTaskTableAdd(RedisModuleCtx *ctx,
|
||||
RedisModuleString *id,
|
||||
RedisModuleString *data) {
|
||||
const char *buf = RedisModule_StringPtrLen(data, NULL);
|
||||
auto message = flatbuffers::GetRoot<TaskTableData>(buf);
|
||||
RAY_CHECK(message != nullptr);
|
||||
|
||||
if (message->scheduling_state() == SchedulingState_WAITING ||
|
||||
message->scheduling_state() == SchedulingState_SCHEDULED) {
|
||||
@@ -483,7 +519,6 @@ int TaskTableAdd(RedisModuleCtx *ctx,
|
||||
long long num_clients = RedisModule_CallReplyInteger(reply);
|
||||
RAY_CHECK(num_clients <= 1) << "Published to " << num_clients
|
||||
<< " clients.";
|
||||
|
||||
}
|
||||
return RedisModule_ReplyWithSimpleString(ctx, "OK");
|
||||
}
|
||||
@@ -542,47 +577,47 @@ int PublishTableAdd(RedisModuleCtx *ctx,
|
||||
return RedisModule_ReplyWithSimpleString(ctx, "OK");
|
||||
}
|
||||
|
||||
/// Add an entry at a key. This overwrites any existing data at the key.
|
||||
/// Publishes a notification about the update to all subscribers, if a pubsub
|
||||
/// channel is provided.
|
||||
///
|
||||
/// This is called from a client with the command:
|
||||
//
|
||||
/// RAY.TABLE_ADD <table_prefix> <pubsub_channel> <id> <data>
|
||||
///
|
||||
/// \param table_prefix The prefix string for keys in this table.
|
||||
/// \param pubsub_channel The pubsub channel name that notifications for
|
||||
/// this key should be published to. When publishing to a specific
|
||||
/// client, the channel name should be <pubsub_channel>:<client_id>.
|
||||
/// \param id The ID of the key to set.
|
||||
/// \param data The data to insert at the key.
|
||||
/// \return The current value at the key, or OK if there is no value.
|
||||
int TableAdd_RedisCommand(RedisModuleCtx *ctx,
|
||||
RedisModuleString **argv,
|
||||
int argc) {
|
||||
RedisModule_AutoMemory(ctx);
|
||||
// RAY.TABLE_ADD:
|
||||
// TableAdd_RedisCommand: the actual command handler.
|
||||
// (helper) TableAdd_DoWrite: performs the write to redis state.
|
||||
// (helper) TableAdd_DoPublish: performs a publish after the write.
|
||||
// ChainTableAdd_RedisCommand: the same command, chain-enabled.
|
||||
|
||||
int TableAdd_DoWrite(RedisModuleCtx *ctx,
|
||||
RedisModuleString **argv,
|
||||
int argc,
|
||||
RedisModuleString **mutated_key_str) {
|
||||
if (argc != 5) {
|
||||
return RedisModule_WrongArity(ctx);
|
||||
}
|
||||
|
||||
RedisModuleString *prefix_str = argv[1];
|
||||
RedisModuleString *id = argv[3];
|
||||
RedisModuleString *data = argv[4];
|
||||
|
||||
RedisModuleKey *key =
|
||||
OpenPrefixedKey(ctx, prefix_str, id, REDISMODULE_READ | REDISMODULE_WRITE,
|
||||
mutated_key_str);
|
||||
RedisModule_StringSet(key, data);
|
||||
return REDISMODULE_OK;
|
||||
}
|
||||
|
||||
int TableAdd_DoPublish(RedisModuleCtx *ctx,
|
||||
RedisModuleString **argv,
|
||||
int argc) {
|
||||
if (argc != 5) {
|
||||
return RedisModule_WrongArity(ctx);
|
||||
}
|
||||
RedisModuleString *pubsub_channel_str = argv[2];
|
||||
RedisModuleString *id = argv[3];
|
||||
RedisModuleString *data = argv[4];
|
||||
|
||||
// Set the keys in the table.
|
||||
RedisModuleKey *key = OpenPrefixedKey(ctx, prefix_str, id,
|
||||
REDISMODULE_READ | REDISMODULE_WRITE);
|
||||
RedisModule_StringSet(key, data);
|
||||
|
||||
// Publish a message on the requested pubsub channel if necessary.
|
||||
TablePubsub pubsub_channel = ParseTablePubsub(pubsub_channel_str);
|
||||
|
||||
if (pubsub_channel == TablePubsub_TASK) {
|
||||
// Publish the task to its subscribers.
|
||||
// TODO(swang): This is only necessary for legacy Ray and should be removed
|
||||
// once we switch to using the new GCS API for the task table.
|
||||
return TaskTableAdd(ctx, id, data);
|
||||
return PublishTaskTableAdd(ctx, id, data);
|
||||
} else if (pubsub_channel != TablePubsub_NO_PUBLISH) {
|
||||
// All other pubsub channels write the data back directly onto the channel.
|
||||
return PublishTableAdd(ctx, pubsub_channel_str, id, data);
|
||||
@@ -591,6 +626,38 @@ int TableAdd_RedisCommand(RedisModuleCtx *ctx,
|
||||
}
|
||||
}
|
||||
|
||||
/// Add an entry at a key. This overwrites any existing data at the key.
|
||||
/// Publishes a notification about the update to all subscribers, if a pubsub
|
||||
/// channel is provided.
|
||||
///
|
||||
/// This is called from a client with the command:
|
||||
///
|
||||
/// RAY.TABLE_ADD <table_prefix> <pubsub_channel> <id> <data>
|
||||
///
|
||||
/// \param table_prefix The prefix string for keys in this table.
|
||||
/// \param pubsub_channel The pubsub channel name that notifications for
|
||||
/// this key should be published to. When publishing to a specific
|
||||
/// client, the channel name should be <pubsub_channel>:<client_id>.
|
||||
/// \param id The ID of the key to set.
|
||||
/// \param data The data to insert at the key.
|
||||
/// \return The current value at the key, or OK if there is no value.
|
||||
int TableAdd_RedisCommand(RedisModuleCtx *ctx,
|
||||
RedisModuleString **argv,
|
||||
int argc) {
|
||||
RedisModule_AutoMemory(ctx);
|
||||
TableAdd_DoWrite(ctx, argv, argc, /*mutated_key_str=*/nullptr);
|
||||
return TableAdd_DoPublish(ctx, argv, argc);
|
||||
}
|
||||
|
||||
#if RAY_USE_NEW_GCS
|
||||
int ChainTableAdd_RedisCommand(RedisModuleCtx *ctx,
|
||||
RedisModuleString **argv,
|
||||
int argc) {
|
||||
return module.ChainReplicate(ctx, argv, argc, /*node_func=*/TableAdd_DoWrite,
|
||||
/*tail_func=*/TableAdd_DoPublish);
|
||||
}
|
||||
#endif
|
||||
|
||||
/// Append an entry to the log stored at a key. Publishes a notification about
|
||||
/// the update to all subscribers, if a pubsub channel is provided.
|
||||
///
|
||||
@@ -746,7 +813,6 @@ int TableLookup_RedisCommand(RedisModuleCtx *ctx,
|
||||
ctx, reinterpret_cast<const char *>(fbb.GetBufferPointer()),
|
||||
fbb.GetSize());
|
||||
}
|
||||
|
||||
return REDISMODULE_OK;
|
||||
}
|
||||
|
||||
@@ -1663,7 +1729,7 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx,
|
||||
}
|
||||
|
||||
if (RedisModule_CreateCommand(ctx, "ray.table_add", TableAdd_RedisCommand,
|
||||
"write", 0, 0, 0) == REDISMODULE_ERR) {
|
||||
"write pubsub", 0, 0, 0) == REDISMODULE_ERR) {
|
||||
return REDISMODULE_ERR;
|
||||
}
|
||||
|
||||
@@ -1763,6 +1829,15 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx,
|
||||
return REDISMODULE_ERR;
|
||||
}
|
||||
|
||||
#if RAY_USE_NEW_GCS
|
||||
// Chain-enabled commands that depend on ray-project/credis.
|
||||
if (RedisModule_CreateCommand(ctx, "ray.chain.table_add",
|
||||
ChainTableAdd_RedisCommand, "write pubsub", 0,
|
||||
0, 0) == REDISMODULE_ERR) {
|
||||
return REDISMODULE_ERR;
|
||||
}
|
||||
#endif
|
||||
|
||||
return REDISMODULE_OK;
|
||||
}
|
||||
|
||||
|
||||
@@ -1,6 +1,11 @@
|
||||
#ifndef RAY_REDIS_STRING_H_
|
||||
#define RAY_REDIS_STRING_H_
|
||||
|
||||
#include <stdarg.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <stdarg.h>
|
||||
|
||||
#include "redismodule.h"
|
||||
|
||||
/* Format a RedisModuleString.
|
||||
*
|
||||
@@ -65,3 +70,5 @@ RedisModuleString *RedisString_Format(RedisModuleCtx *ctx,
|
||||
va_end(ap);
|
||||
return result;
|
||||
}
|
||||
|
||||
#endif // RAY_REDIS_STRING_H_
|
||||
|
||||
@@ -3,12 +3,29 @@
|
||||
# This needs to be run in the build tree, which is normally ray/build
|
||||
|
||||
# Cause the script to exit if a single command fails.
|
||||
set -e
|
||||
set -ex
|
||||
|
||||
LaunchRedis() {
|
||||
port=$1
|
||||
if [[ "${RAY_USE_NEW_GCS}" = "on" ]]; then
|
||||
./src/credis/redis/src/redis-server \
|
||||
--loglevel warning \
|
||||
--loadmodule ./src/credis/build/src/libmember.so \
|
||||
--loadmodule ./src/common/redis_module/libray_redis_module.so \
|
||||
--port $port &
|
||||
else
|
||||
./src/common/thirdparty/redis/src/redis-server \
|
||||
--loglevel warning \
|
||||
--loadmodule ./src/common/redis_module/libray_redis_module.so \
|
||||
--port $port &
|
||||
fi
|
||||
sleep 1s
|
||||
}
|
||||
|
||||
|
||||
# Start the Redis shards.
|
||||
./src/common/thirdparty/redis/src/redis-server --loglevel warning --loadmodule ./src/common/redis_module/libray_redis_module.so --port 6379 &
|
||||
./src/common/thirdparty/redis/src/redis-server --loglevel warning --loadmodule ./src/common/redis_module/libray_redis_module.so --port 6380 &
|
||||
sleep 1s
|
||||
LaunchRedis 6379
|
||||
LaunchRedis 6380
|
||||
# Register the shard location with the primary shard.
|
||||
./src/common/thirdparty/redis/src/redis-cli set NumRedisShards 1
|
||||
./src/common/thirdparty/redis/src/redis-cli rpush RedisShards 127.0.0.1:6380
|
||||
|
||||
@@ -7,22 +7,21 @@ set -x
|
||||
# Cause the script to exit if a single command fails.
|
||||
set -e
|
||||
|
||||
# Start the Redis shards.
|
||||
./src/common/thirdparty/redis/src/redis-server --loglevel warning --loadmodule ./src/common/redis_module/libray_redis_module.so --port 6379 &
|
||||
./src/common/thirdparty/redis/src/redis-server --loglevel warning --loadmodule ./src/common/redis_module/libray_redis_module.so --port 6380 &
|
||||
sleep 1s
|
||||
# Register the shard location with the primary shard.
|
||||
./src/common/thirdparty/redis/src/redis-cli set NumRedisShards 1
|
||||
./src/common/thirdparty/redis/src/redis-cli rpush RedisShards 127.0.0.1:6380
|
||||
|
||||
if [ -z "$RAY_USE_NEW_GCS" ]; then
|
||||
# Start the Redis shards.
|
||||
./src/common/thirdparty/redis/src/redis-server --loglevel warning --loadmodule ./src/common/redis_module/libray_redis_module.so --port 6379 &
|
||||
./src/common/thirdparty/redis/src/redis-server --loglevel warning --loadmodule ./src/common/redis_module/libray_redis_module.so --port 6380 &
|
||||
sleep 1s
|
||||
# Register the shard location with the primary shard.
|
||||
./src/common/thirdparty/redis/src/redis-cli set NumRedisShards 1
|
||||
./src/common/thirdparty/redis/src/redis-cli rpush RedisShards 127.0.0.1:6380
|
||||
|
||||
valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/db_tests
|
||||
valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/io_tests
|
||||
valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/task_tests
|
||||
valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/redis_tests
|
||||
valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/task_table_tests
|
||||
valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/object_table_tests
|
||||
./src/common/thirdparty/redis/src/redis-cli shutdown
|
||||
./src/common/thirdparty/redis/src/redis-cli -p 6380 shutdown
|
||||
fi
|
||||
|
||||
./src/common/thirdparty/redis/src/redis-cli shutdown
|
||||
./src/common/thirdparty/redis/src/redis-cli -p 6380 shutdown
|
||||
|
||||
Reference in New Issue
Block a user