[fix] fixed reference count on ai.dagrun and ai.dagrunro for tensor structure. Added AI_dictType AI_dictTypeTensorVals with proper valDestructor

Author: filipecosta90

src/dag.c

@@ -91,6 +91,16 @@ void *RedisAI_DagRunSession(RedisAI_RunInfo *rinfo) {
               currentOp->result = REDISMODULE_ERR;
             }
           }
+          // since we've increased the reference count prior modelrun we need to decrease it
+          const size_t ninputs = RAI_ModelRunCtxNumInputs(currentOp->mctx);
+          for (size_t inputNumber = 0; inputNumber < ninputs; inputNumber++) {
+            RAI_Tensor *tensor =
+                RAI_ModelRunCtxInputTensor(currentOp->mctx, inputNumber);
+            if (tensor) {
+              RAI_TensorFree(tensor);
+            }
+          }
+
         } else {
           currentOp->result = REDISMODULE_ERR;
         }
@@ -243,7 +253,6 @@ int RedisAI_DagRun_Reply(RedisModuleCtx *ctx, RedisModuleString **argv,
       }
       RedisModule_CloseKey(key);
       RedisAI_ReplicateTensorSet(ctx, tensor_keyname, tensor);
-      // TODO: free Tensor
     } else {
       RedisModule_ReplyWithError(
           ctx, "ERR specified persistent key that was not used on DAG");

src/run_info.c

@@ -16,6 +16,40 @@
 #include "util/arr_rm_alloc.h"
 #include "util/dict.h"
 
+
+static uint64_t RAI_TensorDictKeyHashFunction(const void *key){
+    return AI_dictGenHashFunction(key, strlen((char*)key));
+}
+
+static int RAI_TensorDictKeyStrcmp(void *privdata, const void *key1, const void *key2){
+    const char* strKey1 = key1;
+    const char* strKey2 = key2;
+    return strcmp(strKey1, strKey2) == 0;
+}
+
+static void RAI_TensorDictKeyFree(void *privdata, void *key){
+    RedisModule_Free(key);
+}
+
+static void* RAI_TensorDictKeyDup(void *privdata, const void *key){
+    return RedisModule_Strdup((char*)key);
+}
+
+static void RAI_TensorDictValFree(void *privdata, const void *obj){
+    return RAI_TensorFree((RAI_Tensor*)obj);
+}
+
+
+AI_dictType AI_dictTypeTensorVals = {
+        .hashFunction = RAI_TensorDictKeyHashFunction,
+        .keyDup = RAI_TensorDictKeyDup,
+        .valDup = NULL,
+        .keyCompare = RAI_TensorDictKeyStrcmp,
+        .keyDestructor = RAI_TensorDictKeyFree,
+        .valDestructor = RAI_TensorDictValFree,
+};
+
+
 /**
  * Allocate the memory and initialise the RAI_DagOp.
  * @param result Output parameter to capture allocated RAI_DagOp.
@@ -76,7 +110,7 @@ int RAI_InitRunInfo(RedisAI_RunInfo **result) {
     return REDISMODULE_ERR;
   }
   rinfo->use_local_context = 0;
-  rinfo->dagTensorsContext = AI_dictCreate(&AI_dictTypeHeapStrings, NULL);
+  rinfo->dagTensorsContext = AI_dictCreate(&AI_dictTypeTensorVals, NULL);
   if (!(rinfo->dagTensorsContext)) {
     return REDISMODULE_ERR;
   }
@@ -148,17 +182,16 @@ void RAI_FreeRunInfo(RedisModuleCtx *ctx, struct RedisAI_RunInfo *rinfo) {
       tensor = AI_dictGetVal(entry);
       char *key = (char *)AI_dictGetKey(entry);
 
-      if (tensor&&key!=NULL) {
+      if (tensor && key != NULL) {
         // if the key is persistent then we should not delete it
         AI_dictEntry *persistent_entry =
             AI_dictFind(rinfo->dagTensorsPersistentContext, key);
-        // if the key was loaded from the keyspace then we should not delete
-        // it
+        // if the key was loaded from the keyspace then we should not delete it
         AI_dictEntry *loaded_entry =
             AI_dictFind(rinfo->dagTensorsLoadedContext, key);
 
         if (persistent_entry == NULL && loaded_entry == NULL) {
-          RAI_TensorFree(tensor);
+          AI_dictDelete(rinfo->dagTensorsContext, key);
         }
 
         if (persistent_entry) {

test/tests_sanitizer.py

@@ -27,17 +27,17 @@ def test_sanitizer_dagrun_mobilenet_v1(env):
     env.assertEqual(ret, b'OK')
 
     for opnumber in range(1, MAX_ITERATIONS):
-        image_key = 'image'
-        temp_key1 = 'temp_key1'
-        temp_key2 = 'temp_key2'
+        image_key = 'image{}'.format(opnumber)
         class_key = 'output'
 
         ret = con.execute_command(
             'AI.DAGRUN', '|>',
-            'AI.TENSORSET', image_key, 'FLOAT', 1, 224, 224, 3, 'BLOB', img.tobytes(), '|>',
+            'AI.TENSORSET', image_key, 'FLOAT', 1, 224, 224, 3, 'BLOB', img.tobytes(), 
+            '|>',
             'AI.MODELRUN', model_name,
                          'INPUTS', image_key,
-                         'OUTPUTS', class_key,  '|>',
+                         'OUTPUTS', class_key, 
+                          '|>',
             'AI.TENSORGET',  class_key, 'blob'
         )
         env.assertEqual([b'OK', b'OK'], ret[:2])

test/tests_tensorflow.py

@@ -625,8 +625,8 @@ def test_tensorflow_modelrun_with_batch_and_minbatch(env):
 
     con.execute_command('AI.MODELSET', model_name, 'TF', DEVICE,
                         'BATCHSIZE', batch_size, 'MINBATCHSIZE', minbatch_size,
-                        'INPUTS', inputvar,
-                        'OUTPUTS', outputvar,
+                        'INPUTS', input_var,
+                        'OUTPUTS', output_var,
                         'BLOB', model_pb)
     con.execute_command('AI.TENSORSET', 'input',
                         'FLOAT', 1, img.shape[1], img.shape[0], img.shape[2],
@@ -649,8 +649,8 @@ def run(name=model_name, output_name='output'):
 
     con.execute_command('AI.MODELSET', another_model_name, 'TF', DEVICE,
                         'BATCHSIZE', batch_size, 'MINBATCHSIZE', minbatch_size,
-                        'INPUTS', inputvar,
-                        'OUTPUTS', outputvar,
+                        'INPUTS', input_var,
+                        'OUTPUTS', output_var,
                         'BLOB', model_pb)
 
     p1b = mp.Process(target=run, args=(another_model_name, 'final1'))