Validate that we don't load more than one tensor under the same name to DAG in LLAPI.

Author: alonre24

src/DAG/dag_builder.c

@@ -30,6 +30,12 @@ int RAI_DAGLoadTensor(RAI_DAGRunCtx *run_info, const char *t_name, RAI_Tensor *t
     RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)run_info;
     RedisModuleString *key_name = RedisModule_CreateString(NULL, t_name, strlen(t_name));
 
+    // Cannot load more than one tensor under the same name
+    if (AI_dictFind(rinfo->tensorsNamesToIndices, key_name) != NULL) {
+        RedisModule_FreeString(NULL, key_name);
+        return REDISMODULE_ERR;
+    }
+
     // Add the tensor to the DAG shared tensors and map its name to the relevant index.
     size_t index = array_len(rinfo->dagSharedTensors);
     AI_dictAdd(rinfo->tensorsNamesToIndices, (void *)key_name, (void *)index);

src/DAG/dag_parser.c

@@ -246,8 +246,8 @@ int ParseDAGRunCommand(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModuleS
         const char *arg_string = RedisModule_StringPtrLen(argv[arg_pos], NULL);
 
         if (!strcasecmp(arg_string, "LOAD") && !load_complete && chainingOpCount == 0) {
-            /* Load the required tensors from key space and store them in both
-               dagTensorsLoadedContext and dagTensorsContext dicts. */
+            /* Load the required tensors from key space to the dag shared tensors
+             * array, and save a mapping of their names to the corresponding indices. */
             const int parse_result =
                 _ParseDAGLoadArgs(ctx, &argv[arg_pos], argc - arg_pos, rinfo->tensorsNamesToIndices,
                                   &rinfo->dagSharedTensors, "|>", rinfo->err);
@@ -263,8 +263,9 @@ int ParseDAGRunCommand(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModuleS
                              "ERR PERSIST cannot be specified in a read-only DAG");
                 goto cleanup;
             }
-            /* Store the keys to persist in dagTensorsPersistedContext dict.
-               These keys will be populated later on with actual tensors. */
+            /* Store the keys to persist in persistTensors dict, these keys will
+             * be mapped later to the indices in the dagSharedTensors array in which the
+             * tensors to persist will be found by the end of the DAG run. */
             const int parse_result = _ParseDAGPersistArgs(&argv[arg_pos], argc - arg_pos,
                                                           rinfo->persistTensors, "|>", rinfo->err);
             if (parse_result <= 0)

tests/module/DAG_utils.c

@@ -72,6 +72,26 @@ static void _DAGFinishFunc(RAI_OnFinishCtx *onFinishCtx, void *private_data) {
     pthread_cond_signal(&global_cond);
 }
 
+int testLoadTensor(RedisModuleCtx *ctx) {
+    RAI_DAGRunCtx *run_info = RedisAI_DAGRunCtxCreate();
+    int res = LLAPIMODULE_ERR;
+    RAI_Tensor *t = (RAI_Tensor *)_getFromKeySpace(ctx, "a{1}");
+    if (RedisAI_DAGLoadTensor(run_info, "input", t) != REDISMODULE_OK) {
+        goto cleanup;
+    }
+    t = (RAI_Tensor *)_getFromKeySpace(ctx, "b{1}");
+
+    // cannot load more than one tensor under the same name.
+    if (RedisAI_DAGLoadTensor(run_info, "input", t) != REDISMODULE_ERR) {
+        goto cleanup;
+    }
+    res = LLAPIMODULE_OK;
+
+    cleanup:
+    RedisAI_DAGFree(run_info);
+    return res;
+}
+
 int testModelRunOpError(RedisModuleCtx *ctx) {
 
     RAI_DAGRunCtx *run_info = RedisAI_DAGRunCtxCreate();

tests/module/DAG_utils.h

@@ -11,6 +11,8 @@ typedef struct RAI_RunResults {
     RAI_Error *error;
 } RAI_RunResults;
 
+int testLoadTensor(RedisModuleCtx *ctx);
+
 int testModelRunOpError(RedisModuleCtx *ctx);
 
 int testEmptyDAGError(RedisModuleCtx *ctx);

tests/module/LLAPI.c

@@ -243,6 +243,10 @@ int RAI_llapi_DAGRun(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
         return REDISMODULE_OK;
     }
 
+    // Test the case a successful and failure tensor load input to DAG.
+    if(testLoadTensor(ctx) != LLAPIMODULE_OK) {
+        return RedisModule_ReplyWithSimpleString(ctx, "LOAD tensor test failed");
+    }
     // Test the case of a failure due to addition of a non compatible MODELRUN op.
     if(testModelRunOpError(ctx) != LLAPIMODULE_OK) {
         return RedisModule_ReplyWithSimpleString(ctx, "MODELRUN op error test failed");