Execute Redis commands in Torch Script

src/backends/torch.c

@@ -11,6 +11,20 @@ int RAI_InitBackendTorch(int (*get_api_fn)(const char *, void *)) {
     get_api_fn("RedisModule_Free", ((void **)&RedisModule_Free));
     get_api_fn("RedisModule_Realloc", ((void **)&RedisModule_Realloc));
     get_api_fn("RedisModule_Strdup", ((void **)&RedisModule_Strdup));
+    get_api_fn("RedisModule_CreateString", ((void **)&RedisModule_CreateString));
+    get_api_fn("RedisModule_FreeString", ((void **)&RedisModule_FreeString));
+    get_api_fn("RedisModule_Call", ((void **)&RedisModule_Call));
+    get_api_fn("RedisModule_CallReplyType", ((void **)&RedisModule_CallReplyType));
+    get_api_fn("RedisModule_CallReplyStringPtr", ((void **)&RedisModule_CallReplyStringPtr));
+    get_api_fn("RedisModule_CallReplyInteger", ((void **)&RedisModule_CallReplyInteger));
+    get_api_fn("RedisModule_CallReplyLength", ((void **)&RedisModule_CallReplyLength));
+    get_api_fn("RedisModule_CallReplyArrayElement", ((void **)&RedisModule_CallReplyArrayElement));
+    get_api_fn("RedisModule_FreeCallReply", ((void **)&RedisModule_FreeCallReply));
+    get_api_fn("RedisModule_GetThreadSafeContext", ((void **)&RedisModule_GetThreadSafeContext));
+    get_api_fn("RedisModule_ThreadSafeContextLock", ((void **)&RedisModule_ThreadSafeContextLock));
+    get_api_fn("RedisModule_ThreadSafeContextUnlock",
+               ((void **)&RedisModule_ThreadSafeContextUnlock));
+    get_api_fn("RedisModule_FreeThreadSafeContext", ((void **)&RedisModule_FreeThreadSafeContext));
 
     return REDISMODULE_OK;
 }

src/command_parser.c

@@ -166,7 +166,7 @@ static int _ScriptRunCommand_ParseArgs(RedisModuleCtx *ctx, RedisModuleString **
                                        RedisModuleString **runkey, char const **func_name,
                                        long long *timeout, int *variadic) {
 
-    if (argc < 5) {
+    if (argc < 3) {
         RAI_SetError(error, RAI_ESCRIPTRUN,
                      "ERR wrong number of arguments for 'AI.SCRIPTRUN' command");
         return REDISMODULE_ERR;
@@ -183,41 +183,51 @@ static int _ScriptRunCommand_ParseArgs(RedisModuleCtx *ctx, RedisModuleString **
     *runkey = argv[argpos];
 
     const char *arg_string = RedisModule_StringPtrLen(argv[++argpos], NULL);
-    if (!strcasecmp(arg_string, "TIMEOUT") || !strcasecmp(arg_string, "INPUTS")) {
+    if (!strcasecmp(arg_string, "TIMEOUT") || !strcasecmp(arg_string, "INPUTS") ||
+        !strcasecmp(arg_string, "OUTPUTS")) {
         RAI_SetError(error, RAI_ESCRIPTRUN, "ERR function name not specified");
         return REDISMODULE_ERR;
     }
     *func_name = arg_string;
-    arg_string = RedisModule_StringPtrLen(argv[++argpos], NULL);
 
-    // Parse timeout arg if given and store it in timeout
-    if (!strcasecmp(arg_string, "TIMEOUT")) {
-        if (_parseTimeout(argv[++argpos], error, timeout) == REDISMODULE_ERR)
-            return REDISMODULE_ERR;
-        arg_string = RedisModule_StringPtrLen(argv[++argpos], NULL);
-    }
-    if (strcasecmp(arg_string, "INPUTS") != 0) {
-        RAI_SetError(error, RAI_ESCRIPTRUN, "ERR INPUTS not specified");
-        return REDISMODULE_ERR;
-    }
-
-    bool is_input = true, is_output = false;
+    bool is_input = false;
+    bool is_output = false;
+    bool timeout_set = false;
     size_t ninputs = 0, noutputs = 0;
     int varidic_start_pos = -1;
 
     while (++argpos < argc) {
         arg_string = RedisModule_StringPtrLen(argv[argpos], NULL);
+
+        // Parse timeout arg if given and store it in timeout
+        if (!strcasecmp(arg_string, "TIMEOUT") && !timeout_set) {
+            if (_parseTimeout(argv[++argpos], error, timeout) == REDISMODULE_ERR)
+                return REDISMODULE_ERR;
+            timeout_set = true;
+            continue;
+        }
+
+        if (!strcasecmp(arg_string, "INPUTS") && !is_input) {
+            is_input = true;
+            is_output = false;
+            continue;
+        }
         if (!strcasecmp(arg_string, "OUTPUTS") && !is_output) {
             is_input = false;
             is_output = true;
-        } else if (!strcasecmp(arg_string, "$")) {
+            continue;
+        }
+        if (!strcasecmp(arg_string, "$")) {
             if (varidic_start_pos > -1) {
                 RAI_SetError(error, RAI_ESCRIPTRUN,
                              "ERR Already encountered a variable size list of tensors");
                 return REDISMODULE_ERR;
             }
             varidic_start_pos = ninputs;
-        } else {
+            continue;
+        }
+        // Parse argument name
+        {
             RAI_HoldString(NULL, argv[argpos]);
             if (is_input) {
                 ninputs++;

src/libtorch_c/CMakeLists.txt

@@ -1,3 +1,3 @@
-add_library(torch_c STATIC torch_c.cpp)
+add_library(torch_c STATIC torch_c.cpp torch_extensions/torch_redis.cpp)
 target_link_libraries(torch_c "${TORCH_LIBRARIES}")
 set_property(TARGET torch_c PROPERTY CXX_STANDARD 14)

src/libtorch_c/torch_c.cpp

@@ -8,6 +8,7 @@
 #include <iostream>
 #include <sstream>
 
+#include "torch_extensions/torch_redis.h"
 namespace {
 
 static DLDataType getDLDataType(const at::Tensor &t) {
@@ -246,6 +247,7 @@ void torchRunModule(ModuleContext *ctx, const char *fnName, int variadic, long n
     torch::DeviceType output_device_type = torch::kCPU;
     torch::Device output_device(output_device_type, -1);
 
+    if(nOutputs == 0) return;
     int count = 0;
     for (size_t i = 0; i < stack.size(); i++) {
         if (count > nOutputs - 1) {
@@ -304,28 +306,37 @@ extern "C" DLManagedTensor *torchNewTensor(DLDataType dtype, long ndims, int64_t
     return dl_tensor;
 }
 
-extern "C" void *torchCompileScript(const char *script, DLDeviceType device, int64_t device_id,
-                                    char **error, void *(*alloc)(size_t)) {
-    ModuleContext *ctx = new ModuleContext();
-    ctx->device = device;
-    ctx->device_id = device_id;
-    try {
-        auto cu = torch::jit::compile(script);
-        auto aten_device_type = getATenDeviceType(device);
-        if (aten_device_type == at::DeviceType::CUDA && !torch::cuda::is_available()) {
-            throw std::logic_error("GPU requested but Torch couldn't find CUDA");
-        }
-        ctx->cu = cu;
-        ctx->module = nullptr;
-    } catch (std::exception &e) {
-        size_t len = strlen(e.what()) + 1;
-        *error = (char *)alloc(len * sizeof(char));
-        strcpy(*error, e.what());
-        (*error)[len - 1] = '\0';
-        delete ctx;
-        return NULL;
+extern "C" void* torchCompileScript(const char* script, DLDeviceType device, int64_t device_id,
+                                    char **error, void* (*alloc)(size_t))
+{
+  ModuleContext* ctx = new ModuleContext();
+  ctx->device = device;
+  ctx->device_id = device_id;
+  try {
+    auto cu = std::make_shared<torch::jit::script::CompilationUnit>();
+    cu->define(
+        c10::nullopt,
+        script,
+        torch::jit::script::redisResolver(),
+        nullptr);
+    auto aten_device_type = getATenDeviceType(device);
+
+    if (aten_device_type == at::DeviceType::CUDA && !torch::cuda::is_available()) {
+      throw std::logic_error("GPU requested but Torch couldn't find CUDA");
     }
-    return ctx;
+    ctx->cu = cu;
+    ctx->module = nullptr;
+
+  }
+  catch(std::exception& e) {
+    size_t len = strlen(e.what()) +1;
+    *error = (char*)alloc(len * sizeof(char));
+    strcpy(*error, e.what());
+    (*error)[len-1] = '\0';
+    delete ctx;
+    return NULL;
+  }
+  return ctx;
 }
 
 extern "C" void *torchLoadModel(const char *graph, size_t graphlen, DLDeviceType device,

src/libtorch_c/torch_extensions/torch_redis.cpp

@@ -0,0 +1,69 @@
+#include <string>
+#include "torch_redis.h"
+#include "../../redismodule.h"
+
+torch::IValue IValueFromRedisReply(RedisModuleCallReply *reply){
+
+    int reply_type = RedisModule_CallReplyType(reply);
+    switch(reply_type) {
+        case REDISMODULE_REPLY_NULL: {
+            return torch::IValue();
+        }
+        case REDISMODULE_REPLY_STRING: {
+            size_t len;
+            const char *replyStr = RedisModule_CallReplyStringPtr(reply, &len);
+            std::string str = replyStr;
+            return torch::IValue(str.substr(0,len));
+        }
+        case REDISMODULE_REPLY_INTEGER: {
+            int intValue = RedisModule_CallReplyInteger(reply);
+            return torch::IValue(intValue);
+        }
+        case REDISMODULE_REPLY_ARRAY: {
+            c10::impl::GenericList vec = c10::impl::GenericList(c10::AnyType::create());
+            size_t len = RedisModule_CallReplyLength(reply);
+            for (auto i = 0; i < len; ++i) {
+                RedisModuleCallReply *subReply = RedisModule_CallReplyArrayElement(reply, i);
+                torch::IValue value = IValueFromRedisReply(subReply);
+                vec.push_back(value);
+            }
+            return torch::IValue(vec);
+        }
+        case REDISMODULE_REPLY_ERROR: {
+            size_t len;
+            const char *replyStr = RedisModule_CallReplyStringPtr(reply, &len);
+            throw std::runtime_error(replyStr);
+            break;
+        }
+        default:{
+            throw(std::runtime_error("Unsupported redis type"));
+        }
+    }
+}
+
+torch::IValue redisExecute(std::string fn_name, std::vector<std::string> args ) {
+    RedisModuleCtx* ctx = RedisModule_GetThreadSafeContext(NULL);
+    RedisModule_ThreadSafeContextLock(ctx);
+    size_t len = args.size();
+    RedisModuleString* arguments[len];
+    len = 0;
+    for (std::vector<std::string>::iterator it = args.begin(); it != args.end(); it++) {
+        const std::string arg = *it;
+        const char* str = arg.c_str();
+        arguments[len++] = RedisModule_CreateString(ctx, str, strlen(str));
+    }
+
+    RedisModuleCallReply *reply = RedisModule_Call(ctx, fn_name.c_str(), "!v", arguments, len);
+    RedisModule_ThreadSafeContextUnlock(ctx);
+    torch::IValue value = IValueFromRedisReply(reply);
+    RedisModule_FreeThreadSafeContext(ctx);
+    RedisModule_FreeCallReply(reply);
+    for(int i= 0; i < len; i++){
+        RedisModule_FreeString(NULL, arguments[i]);
+    }
+    return value;
+}
+
+torch::List<torch::IValue> asList(torch::IValue v) {
+    return v.toList();
+}

src/libtorch_c/torch_extensions/torch_redis.h

@@ -0,0 +1,33 @@
+#include "torch/jit.h"
+#include "torch/script.h"
+#include "torch/csrc/jit/frontend/resolver.h"
+
+namespace torch {
+namespace jit {
+namespace script {
+struct RedisResolver : public Resolver {
+
+    std::shared_ptr<SugaredValue> resolveValue(const std::string &name, Function &m,
+                                               const SourceRange &loc) override {
+        if (strcasecmp(name.c_str(), "torch") == 0) {
+            return std::make_shared<BuiltinModule>("aten");
+        } else if (strcasecmp(name.c_str(), "redis") == 0) {
+            return std::make_shared<BuiltinModule>("redis");
+        }
+        return nullptr;
+    }
+
+    TypePtr resolveType(const std::string &name, const SourceRange &loc) override {
+        return nullptr;
+    }
+};
+inline std::shared_ptr<RedisResolver> redisResolver() { return std::make_shared<RedisResolver>(); }
+} // namespace script
+} // namespace jit
+} // namespace torch
+
+torch::IValue redisExecute(std::string fn_name, std::vector<std::string> args);
+torch::List<torch::IValue> asList(torch::IValue);
+
+static auto registry =
+    torch::RegisterOperators("redis::execute", &redisExecute).op("redis::asList", &asList);

src/redisai.c

@@ -585,9 +585,6 @@ int RedisAI_ScriptRun_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv
         return RedisAI_ScriptRun_IsKeysPositionRequest_ReportKeys(ctx, argv, argc);
     }
 
-    if (argc < 6)
-        return RedisModule_WrongArity(ctx);
-
     // Convert The script run command into a DAG command that contains a single op.
     return RedisAI_ExecuteCommand(ctx, argv, argc, CMD_SCRIPTRUN, false);
 }

tests/flow/test_data/redis_scripts.py

@@ -0,0 +1,66 @@
+
+def redis_string_int_to_tensor(redis_value: Any):
+    return torch.tensor(int(str(redis_value)))
+
+
+def redis_string_float_to_tensor(redis_value: Any):
+    return torch.tensor(float(str((redis_value))))
+
+
+def redis_int_to_tensor(redis_value: int):
+    return torch.tensor(redis_value)
+
+
+def redis_int_list_to_tensor(redis_value: Any):
+    values = redis.asList(redis_value)
+    l  = [torch.tensor(int(str(v))).reshape(1,1) for v in values]
+    return torch.cat(l, dim=0)
+
+
+def redis_hash_to_tensor(redis_value: Any):
+    values = redis.asList(redis_value)
+    l  = [torch.tensor(int(str(v))).reshape(1,1) for v in values]
+    return torch.cat(l, dim=0)
+
+def test_redis_error():
+    redis.execute("SET", "x")
+
+def test_int_set_get():
+    redis.execute("SET", "x", "1")
+    res = redis.execute("GET", "x",)
+    redis.execute("DEL", "x")
+    return redis_string_int_to_tensor(res)
+
+def test_int_set_incr():
+    redis.execute("SET", "x", "1")
+    res = redis.execute("INCR", "x")
+    redis.execute("DEL", "x")
+    return redis_string_int_to_tensor(res)
+
+def test_float_set_get():
+    redis.execute("SET", "x", "1.1")
+    res = redis.execute("GET", "x",)
+    redis.execute("DEL", "x")
+    return redis_string_float_to_tensor(res)
+
+def test_int_list():
+    redis.execute("RPUSH", "x", "1")
+    redis.execute("RPUSH", "x", "2")
+    res = redis.execute("LRANGE", "x", "0", "2")
+    redis.execute("DEL", "x")
+    return redis_int_list_to_tensor(res)
+
+
+def test_hash():
+    redis.execute("HSET", "x", "field1", "1", "field2", "2")
+    res = redis.execute("HVALS", "x")
+    redis.execute("DEL", "x")
+    return redis_hash_to_tensor(res)
+
+
+def test_set_key():
+    redis.execute("SET", ["x{1}", "1"])
+
+
+def test_del_key():
+    redis.execute("DEL", ["x"])