wip

Author: DvirDukhan

src/libtorch_c/torch_c.cpp

@@ -7,6 +7,40 @@
 #include <iostream>
 #include <sstream>
 
+
+#include "torch/csrc/jit/frontend/resolver.h"
+#include "torch/script.h"
+#include "torch/jit.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 {
 
 static DLDataType getDLDataType(const at::Tensor &t) {

src/libtorch_c/torch_extensions/torch_redis_value.h

@@ -0,0 +1,44 @@
+#pragma once
+#include "torch/script.h"
+#include "torch/custom_class.h"
+#include "../../redismodule.h"
+
+
+
+struct RedisValue: torch::CustomClassHolder {
+private:
+    union {
+        int intValue;
+        std::string stringValue;
+        std::vector<RedisValue*> arrayValue;
+    };
+public:
+    RedisValue(RedisModuleCallReply *reply) {
+        if(RedisModule_CallReplyType(reply) == REDISMODULE_REPLY_ARRAY) {
+        size_t len = RedisModule_CallReplyLength(reply);
+        for(auto i = 0 ; i < len ; ++i){
+            RedisModuleCallReply *subReply = RedisModule_CallReplyArrayElement(reply, i);
+            RedisValue value(subReply);
+            arrayValue.push_back(value);
+        }
+    }
+
+    if(RedisModule_CallReplyType(reply) == REDISMODULE_REPLY_STRING ||
+            RedisModule_CallReplyType(reply) == REDISMODULE_REPLY_ERROR){
+        size_t len;
+        const char* replyStr = RedisModule_CallReplyStringPtr(reply, &len);
+        PyObject* ret = PyUnicode_FromStringAndSize(replyStr, len);
+        if(!ret){
+            PyErr_Clear();
+            ret = PyByteArray_FromStringAndSize(replyStr, len);
+        }
+        return ret;
+    }
+
+    if(RedisModule_CallReplyType(reply) == REDISMODULE_REPLY_INTEGER){
+        long long val = RedisModule_CallReplyInteger(reply);
+        return PyLong_FromLongLong(val);
+    }
+    }
+
+};

test/test_data/redis_scripts.txt

@@ -0,0 +1,71 @@
+def redis_string_to_int(redis_value: RedisValue):
+    return int(redis_value.stringValue())
+
+def redis_string_to_float(redis_value: RedisValue):
+    return float(redis_value.stringValue())
+
+def redis_string_int_to_tensor(redis_value: RedisValue):
+    return tensor(redis_string_to_int(redis_value))
+
+def redis_string_float_to_tensor(redis_value: RedisValue):
+    return tensor(redis_string_to_float(redis_value))
+
+def redis_int_to_tensor(redis_value: RedisValue):
+    return tensor(redis_value.intValue())
+
+def redis_int_list_to_tensor(redis_value: RedisValue):
+    len = len(redis_value.getList())
+    l = []
+    for v in redis_value.getList():
+        l.append(redis_string_to_int(v))
+    return torch.cat(l, dim=0)
+
+def redis_float_list_to_tensor(redis_value: RedisValue):
+    len = len(redis_value.getList())
+    l = []
+    for v in redis_value.getList():
+        l.append(redis_string_to_float(v))
+    return torch.cat(l, dim=0)
+
+def redis_hash_to_tensor(redis_value: RedisValue):
+    len = len(redis_value.getList())
+    l = []
+    for v in redis_value.getList():
+        l.append(redis_string_to_float(v.getList()[1]))
+    return torch.cat(l, dim=0)
+
+def test_redis_error():
+    res = redis.executeCommand("SET", "x")
+    return tensor(res.getValueType())
+
+def test_int_set_get():
+    redis.executeCommand("SET", "x", "1")
+    res = redis.executeCommand("GET", "x",)
+    redis.executeCommand("DEL", "x")
+    return redis_string_int_to_tensor(res)
+
+def test_float_set_get():
+    redis.executeCommand("SET", "x", "1.1")
+    res = redis.executeCommand("GET", "x",)
+    redis.executeCommand("DEL", "x")
+    return redis_string_int_to_tensor(res)
+
+def test_int_list():
+    redis.executeCommand("LPUSH", "x", "1")
+    redis.executeCommand("LPUSH", "x", "2")
+    res = redis.executeCommand("LRANGE", "x")
+    redis.executeCommand("DEL", "x")
+    return redis_int_list_to_tensor(res)
+
+def test_float_list():
+    redis.executeCommand("LPUSH", "x", "1.1")
+    redis.executeCommand("LPUSH", "x", "2.2")
+    res = redis.executeCommand("LRANGE", "x")
+    redis.executeCommand("DEL", "x")
+    return redis_float_list_to_tensor(res)
+
+def test_hash():
+    redis.executeCommand("HSET", "x", "1", "2.2)
+    res = redis.executeCommand("HGETALL", "x")
+    redis.executeCommand("DEL", "x")
+    return redis_float_list_to_tensor(res)

test/test_torchscript_extensions.py

@@ -0,0 +1,35 @@
+import redis
+
+from includes import *
+
+'''
+python -m RLTest --test tests_torchscript_extensions.py --module path/to/redisai.so
+'''
+
+
+class test_torch_script_extesions:
+
+    def __init__(self):
+        self.env = Env()
+        if not TEST_PT:
+            self.env.debugPrint("skipping {} since TEST_PT=0".format(sys._getframe().f_code.co_name), force=True)
+            self.env.skip() 
+
+        self.con = self.env.getConnect()
+        script_filename = os.path.join(test_data_path, 'redis_scripts.txt')
+        with open(script_filename, 'rb') as f:
+            script = f.read()
+
+        ret = self.con.execute_command('AI.SCRIPTSET', 'redis_scripts', DEVICE, 'SOURCE', script)
+        self.env.assertEqual(ret, b'OK')
+        self.env.ensureSlaveSynced(self.con, self.env)
+    
+    def test_int_get_set(self):
+        self.con.execute_command('AI.SCRIPTRUN', 'redis_scripts', 'test_int_set_get', 'INPUTS', 'OUTPUTS', 'y')
+        y = self.con.execute_command('AI.TENSORGET', 'y', 'meta' ,'VALUES')
+        self.env.assertEqual(y, ["dtype", "INT64", "shape", [0, 1], "VALUES", "1"] )
+
+    def test_float_get_set(self):
+        self.con.execute_command('AI.SCRIPTRUN', 'redis_scripts', 'test_float_set_get', 'INPUTS', 'OUTPUTS', 'y')
+        y = self.con.execute_command('AI.TENSORGET', 'y', 'meta' ,'VALUES')
+        self.env.assertEqual(y, ["dtype", "FLOAT", "shape", [0, 1], "VALUES", "1.1"])