Small refactor of RDB encode/decode (so we won't copy blob on rdb load) + small fixes

Author: alonre24

src/redis_ai_objects/tensor.c

@@ -136,8 +136,10 @@ static int _RAI_TensorParseStringsBlob(const char *tensor_blob, size_t blob_len,
         }
     }
     if (tensor_blob[blob_len - 1] != '\0' || elements_counter != tensor_len) {
-        RAI_SetError(err, RAI_ETENSORSET,
-                     "ERR Number of string elements in data blob does not match tensor length");
+        if (err) {
+            RAI_SetError(err, RAI_ETENSORSET,
+                         "ERR Number of string elements in data blob does not match tensor length");
+        }
         return REDISMODULE_ERR;
     }
     return REDISMODULE_OK;
@@ -647,11 +649,9 @@ int RAI_TensorGetValueAsCString(RAI_Tensor *t, long long i, const char **val) {
 int RAI_TensorSetData(RAI_Tensor *t, const char *data, size_t len) {
     DLDataType data_type = RAI_TensorDataType(t);
     if (data_type.code == kDLString) {
-        RAI_Error err = {0};
         if (_RAI_TensorParseStringsBlob(data, len, RAI_TensorLength(t),
                                         RAI_TensorStringElementsOffsets(t),
-                                        &err) != REDISMODULE_OK) {
-            RAI_ClearError(&err);
+                                        NULL) != REDISMODULE_OK) {
             return 0;
         }
         RedisModule_Free(RAI_TensorData(t));

src/serialization/RDB/decoder/current/v4/decode_v4.c

@@ -15,27 +15,36 @@ void *RAI_RDBLoadTensor_v4(RedisModuleIO *io) {
     DLDataType data_type = (DLDataType){.code = code, .bits = bits, .lanes = 1};
 
     int ndims = (int)RedisModule_LoadSigned(io);
-    int64_t shape[ndims];
+    size_t shape[ndims];
     for (size_t i = 0; i < ndims; ++i) {
         shape[i] = RedisModule_LoadSigned(io);
     }
 
-    RAI_Error err = {0};
+    RAI_Tensor *tensor = RAI_TensorNew(data_type, shape, ndims);
+
     size_t blob_len;
     char *data = RedisModule_LoadStringBuffer(io, &blob_len);
     if (RedisModule_IsIOError(io))
         goto error;
-    RAI_Tensor *t =
-        RAI_TensorCreateFromBlob(data_type, (const size_t *)shape, ndims, data, blob_len, &err);
-    RedisModule_Free(data);
-    if (t == NULL)
-        goto error;
 
-    return t;
+    tensor->blobSize = blob_len;
+    tensor->tensor.dl_tensor.data = data;
+
+    if (data_type.code == kDLString) {
+        for (size_t i = 0; i < RAI_TensorLength(tensor); i++) {
+            tensor->tensor.dl_tensor.elements_length[i] = RedisModule_LoadUnsigned(io);
+        }
+    }
+    if (RedisModule_IsIOError(io))
+        goto error;
+    return tensor;
 
 error:
-    RAI_ClearError(&err);
     RedisModule_LogIOError(io, "error", "Experienced a short read while reading a tensor from RDB");
+    RAI_TensorFree(tensor);
+    if (data) {
+        RedisModule_Free(data);
+    }
     return NULL;
 }
 

src/serialization/RDB/decoder/previous/v0/decode_v0.c

@@ -17,8 +17,8 @@ void *RAI_RDBLoadTensor_v0(RedisModuleIO *io) {
     dtype.code = RedisModule_LoadUnsigned(io);
     dtype.lanes = RedisModule_LoadUnsigned(io);
 
-    size_t ndims = RedisModule_LoadUnsigned(io);
-    int64_t shape[ndims];
+    int ndims = RedisModule_LoadUnsigned(io);
+    size_t shape[ndims];
     for (size_t i = 0; i < ndims; ++i) {
         shape[i] = RedisModule_LoadUnsigned(io);
     }
@@ -29,22 +29,18 @@ void *RAI_RDBLoadTensor_v0(RedisModuleIO *io) {
     }
     size_t byte_offset = RedisModule_LoadUnsigned(io);
 
-    RAI_Error err = {0};
     size_t blob_len;
     char *data = RedisModule_LoadStringBuffer(io, &blob_len);
     if (RedisModule_IsIOError(io))
         goto error;
 
-    RAI_Tensor *t =
-        RAI_TensorCreateFromBlob(dtype, (const size_t *)shape, (int)ndims, data, blob_len, &err);
-    RedisModule_Free(data);
-    if (t == NULL)
-        goto error;
+    RAI_Tensor *tensor = RAI_TensorNew(dtype, shape, ndims);
+    tensor->blobSize = blob_len;
+    tensor->tensor.dl_tensor.data = data;
 
-    return t;
+    return tensor;
 
 error:
-    RAI_ClearError(&err);
     RedisModule_LogIOError(io, "error", "Experienced a short read while reading a tensor from RDB");
     return NULL;
 }

src/serialization/RDB/encoder/v4/encode_v4.c

@@ -14,6 +14,12 @@ void RAI_RDBSaveTensor_v4(RedisModuleIO *io, void *value) {
 
     size_t size = RAI_TensorByteSize(tensor);
     RedisModule_SaveStringBuffer(io, tensor->tensor.dl_tensor.data, size);
+
+    if (tensor->tensor.dl_tensor.dtype.code == kDLString) {
+        for (size_t i = 0; i < RAI_TensorLength(tensor); i++) {
+            RedisModule_SaveUnsigned(io, tensor->tensor.dl_tensor.elements_length[i]);
+        }
+    }
 }
 
 void RAI_RDBSaveModel_v4(RedisModuleIO *io, void *value) {

tests/flow/test_serializations.py

@@ -322,7 +322,7 @@ def test_v4_onnx_model(self):
     def test_v4_tensor(self):
         key_name = "tensor{1}"
         con = get_connection(self.env, key_name)
-        tensor_rdb = b"\a\x81\x00\x8f\xd3\x10\xd4\x8eD\x04\x02\x00\x02 \x02\x02\x02\x02\x02\x01\x05\b\x01\x00\x00\x00\x02\x00\x00\x00\x00\t\x00Viy\xab4\xbe\xdd\x82"
+        tensor_rdb = b'\x07\x81\x00\x8f\xd3\x10\xd4\x8eD\x04\x02\x00\x02 \x02\x02\x02\x02\x02\x01\x05\x08\x01\x00\x00\x00\x02\x00\x00\x00\x00\t\x00Viy\xab4\xbe\xdd\x82'
         self.env.assertEqual(con.execute_command('FLUSHALL'), True)
         con.restore(key_name, 0, tensor_rdb, True)
         _, tensor_type, _, tensor_shape = con.execute_command('AI.TENSORGET', key_name, 'META')
@@ -331,7 +331,7 @@ def test_v4_tensor(self):
         self.env.assertEqual(values, [1, 2])
 
         # test RDB load of string tensor
-        str_tensor_rdb = b"\a\x81\x00\x8f\xd3\x10\xd4\x8eD\x04\x02\a\x02\b\x02\x01\x02\x02\x05\x12str_val1\x00str_val2\x00\x00\t\x00\xf9;\xe0\xd12.\x06z"
+        str_tensor_rdb = b'\x07\x81\x00\x8f\xd3\x10\xd4\x8eD\x04\x02\x07\x02\x08\x02\x01\x02\x02\x05\x12str_val1\x00str_val2\x00\x02\x00\x02\t\x00\t\x00\x8b\x05Z\x0f:\x877O'
         con.restore('string_tensor{1}', 0, str_tensor_rdb, True)
         _, tensor_type, _, tensor_shape = con.execute_command('AI.TENSORGET', 'string_tensor{1}', 'META')
         self.env.assertEqual([tensor_type, tensor_shape], [b"STRING", [2]])

tests/flow/tests_onnx.py

@@ -90,6 +90,7 @@ def test_onnx_string_tensors(env):
     env.assertEqual(tensor_values, [b'input11', b'input12', b'input21', b'input22'])
 
     if env.useSlaves:
+        ensureSlaveSynced(con, env)
         slave_con = env.getSlaveConnection()
         slave_tensor_values = slave_con.execute_command('AI.TENSORGET', 'out_tensor{1}', 'VALUES')
         env.assertEqual(tensor_values, slave_tensor_values)

tests/flow/tests_tensorflow.py

@@ -726,6 +726,7 @@ def test_tf_string_tensors(env):
                                     b'here we want to test a string longer than 24 chars, to force heap alloc in tf'])
 
     if env.useSlaves:
+        ensureSlaveSynced(con, env)
         slave_con = env.getSlaveConnection()
         slave_tensor_values = slave_con.execute_command('AI.TENSORGET', 'out_tensor{1}', 'VALUES')
         env.assertEqual(tensor_values, slave_tensor_values)