- Setting a global error in DAG runInfo when one of its ops returns with an error. this is the one that is returned in LLAPI.

- Add macro to redisai.h to resolve double includes issues.
- Update arr.h, support creating dynamic arrays on stack (in addition).
- PR fixes, test module refactor

Author: alonre24

src/DAG/dag.c

@@ -571,6 +571,9 @@ void RedisAI_DagRunSessionStep(RedisAI_RunInfo *rinfo, const char *devicestr) {
 
     if (currentOp->result != REDISMODULE_OK) {
         __atomic_store_n(rinfo->dagError, 1, __ATOMIC_RELAXED);
+        RAI_ContextWriteLock(rinfo);
+        RAI_SetError(rinfo->err, RAI_GetErrorCode(currentOp->err), RAI_GetError(currentOp->err));
+        RAI_ContextUnlock(rinfo);
     }
 }
 
@@ -599,6 +602,8 @@ void RedisAI_BatchedDagRunSessionStep(RedisAI_RunInfo **batched_rinfo, const cha
 
         if (currentOp->result != REDISMODULE_OK) {
             __atomic_store_n(rinfo->dagError, 1, __ATOMIC_RELAXED);
+            RAI_SetError(rinfo->err, RAI_GetErrorCode(currentOp->err),
+                         RAI_GetError(currentOp->err));
         }
     }
     return;

src/DAG/dag_builder.c

@@ -4,6 +4,27 @@
 #include "string_utils.h"
 #include "modelRun_ctx.h"
 
+// Store the given arguments from the string in argv array and their amount in argc.
+int _StringToRMArray(const char *dag, RedisModuleString ***argv, int *argc, RAI_Error *err) {
+
+    char dag_string[strlen(dag) + 1];
+    strcpy(dag_string, dag);
+
+    char *token = strtok(dag_string, " ");
+    if (strcmp(token, "|>") != 0) {
+        RAI_SetError(err, RAI_EDAGBUILDER, "DAG op should start with: '|>' ");
+        return REDISMODULE_ERR;
+    }
+
+    while (token != NULL) {
+        RedisModuleString *RS_token = RedisModule_CreateString(NULL, token, strlen(token));
+        *argv = array_append(*argv, RS_token);
+        (*argc)++;
+        token = strtok(NULL, " ");
+    }
+    return REDISMODULE_OK;
+}
+
 int RAI_DAGLoadTensor(RAI_DAGRunCtx *run_info, const char *t_name, RAI_Tensor *tensor) {
 
     RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)run_info;
@@ -116,51 +137,43 @@ int RAI_DAGAddOpsFromString(RAI_DAGRunCtx *run_info, const char *dag, RAI_Error
 
     int res = REDISMODULE_ERR;
     RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)run_info;
+    array_new_on_stack(RAI_DagOp *, 10, new_ops);
+    array_new_on_stack(RedisModuleString *, 100, argv);
     int argc = 0;
-    char dag_string[strlen(dag) + 1];
-    strcpy(dag_string, dag);
-
-    char *token = strtok(dag_string, " ");
-    if (strcmp(token, "|>") != 0) {
-        RAI_SetError(err, RAI_EDAGBUILDER, "DAG op should start with: '|>' ");
-        return res;
-    }
-    RedisModuleString **argv = array_new(RedisModuleString *, 2);
-    while (token != NULL) {
-        RedisModuleString *RS_token = RedisModule_CreateString(NULL, token, strlen(token));
-        argv = array_append(argv, RS_token);
-        argc++;
-        token = strtok(NULL, " ");
+    if (_StringToRMArray(dag, &argv, &argc, err) != REDISMODULE_OK) {
+        goto cleanup;
     }
 
-    size_t num_ops_before = array_len(rinfo->dagOps);
-    size_t new_ops = 0;
     RAI_DagOp *op;
     for (size_t i = 0; i < argc; i++) {
         const char *arg_string = RedisModule_StringPtrLen(argv[i], NULL);
         if (strcmp(arg_string, "|>") == 0 && i < argc - 1) {
             RAI_InitDagOp(&op);
-            rinfo->dagOps = array_append(rinfo->dagOps, op);
-            new_ops++;
+            new_ops = array_append(new_ops, op);
             op->argv = &argv[i + 1];
         } else {
             op->argc++;
         }
     }
 
-    if (ParseDAGOps(rinfo, num_ops_before, new_ops) != REDISMODULE_OK) {
-        // Remove all ops that where added before the error and go back to the initial state.
+    if (ParseDAGOps(rinfo, new_ops) != REDISMODULE_OK) {
+        // Remove all ops that where created.
         RAI_SetError(err, RAI_GetErrorCode(rinfo->err), RAI_GetError(rinfo->err));
-        for (size_t i = num_ops_before; i < array_len(rinfo->dagOps); i++) {
-            RAI_FreeDagOp(rinfo->dagOps[i]);
+        for (size_t i = 0; i < array_len(new_ops); i++) {
+            RAI_FreeDagOp(new_ops[i]);
         }
-        rinfo->dagOps = array_trimm_len(rinfo->dagOps, num_ops_before);
         goto cleanup;
     }
+
+    // Copy the new op pointers to the DAG run info.
+    for (size_t i = 0; i < array_len(new_ops); i++) {
+        rinfo->dagOps = array_append(rinfo->dagOps, new_ops[i]);
+    }
     rinfo->dagOpCount = array_len(rinfo->dagOps);
     res = REDISMODULE_OK;
 
 cleanup:
+    array_free(new_ops);
     for (size_t i = 0; i < argc; i++) {
         RedisModule_FreeString(NULL, argv[i]);
     }

src/DAG/dag_execute.c

@@ -285,16 +285,11 @@ RAI_Tensor *RAI_DAGOutputTensor(RAI_OnFinishCtx *finish_ctx, size_t index) {
     return NULL;
 }
 
-int RAI_DAGRunError(RAI_OnFinishCtx *finish_ctx) {
+bool RAI_DAGRunError(RAI_OnFinishCtx *finish_ctx) {
     return *((RedisAI_RunInfo *)finish_ctx)->dagError;
 }
 
-RAI_Error *RAI_DAGCopyOpStatus(RAI_OnFinishCtx *finish_ctx, size_t index) {
+RAI_Error *RAI_DAGGetError(RAI_OnFinishCtx *finish_ctx) {
     RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)finish_ctx;
-    RedisModule_Assert(index < rinfo->dagOpCount);
-    RAI_Error *err;
-    RAI_InitError(&err);
-    RAI_SetError(err, RAI_GetErrorCode(rinfo->dagOps[index]->err),
-                 RAI_GetError(rinfo->dagOps[index]->err));
-    return err;
+    return rinfo->err;
 }

src/DAG/dag_execute.h

@@ -45,15 +45,14 @@ size_t RAI_DAGNumOutputs(RAI_OnFinishCtx *finish_ctx);
 RAI_Tensor *RAI_DAGOutputTensor(RAI_OnFinishCtx *finish_ctx, size_t index);
 
 /**
- * @brief Returns 1 if (at least) one of the DAG ops encountered an error.
+ * @brief Returns true if (at least) one of the DAG ops encountered an error.
  */
-int RAI_DAGRunError(RAI_OnFinishCtx *finish_ctx);
+bool RAI_DAGRunError(RAI_OnFinishCtx *finish_ctx);
 
 /**
- * @brief This can be called in the finish CB, returns the status of a certain in a DAG.
+ * @brief This can be called in the finish CB, to get DAG error details.
  * @param finish_ctx This represents the DAG runInfo at the end of the run.
- * @param index Index of a specific op in the DAG.
- * @retval returns an object that represents the i'th op status, from which a user can
+ * @retval returns an object that represents the DAG status, from which a user can
  * obtain the error code (error code is "OK" if no error has occurred) and error details.
  */
-RAI_Error *RAI_DAGCopyOpStatus(RAI_OnFinishCtx *finish_ctx, size_t index);
+RAI_Error *RAI_DAGGetError(RAI_OnFinishCtx *finish_ctx);

src/DAG/dag_parser.c

@@ -139,10 +139,10 @@ static int _parseTimeout(RedisModuleString **argv, int argc, long long *timeout,
     return REDISMODULE_OK;
 }
 
-static RAI_DagOp *_AddEmptyOp(RedisAI_RunInfo *rinfo) {
+static RAI_DagOp *_AddEmptyOp(RAI_DagOp ***ops) {
     RAI_DagOp *currentDagOp;
     RAI_InitDagOp(&currentDagOp);
-    rinfo->dagOps = array_append(rinfo->dagOps, currentDagOp);
+    *ops = array_append(*ops, currentDagOp);
     return currentDagOp;
 }
 
@@ -160,10 +160,10 @@ int _CollectOpArgs(RedisModuleString **argv, int argc, int arg_pos, RAI_DagOp *o
     return op->argc;
 }
 
-int ParseDAGOps(RedisAI_RunInfo *rinfo, size_t first_op_ind, size_t num_ops) {
+int ParseDAGOps(RedisAI_RunInfo *rinfo, RAI_DagOp **ops) {
 
-    for (long long i = 0; i < num_ops; i++) {
-        RAI_DagOp *currentOp = rinfo->dagOps[i + first_op_ind];
+    for (long long i = 0; i < array_len(ops); i++) {
+        RAI_DagOp *currentOp = ops[i];
         // The first op arg is the command name.
         const char *arg_string = RedisModule_StringPtrLen(currentOp->argv[0], NULL);
 
@@ -211,6 +211,7 @@ int ParseDAGOps(RedisAI_RunInfo *rinfo, size_t first_op_ind, size_t num_ops) {
 int ParseDAGRunCommand(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModuleString **argv,
                        int argc, bool dag_ro) {
 
+    int res = REDISMODULE_ERR;
     if (argc < 4) {
         if (dag_ro) {
             RAI_SetError(rinfo->err, RAI_EDAGBUILDER,
@@ -219,14 +220,15 @@ int ParseDAGRunCommand(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModuleS
             RAI_SetError(rinfo->err, RAI_EDAGBUILDER,
                          "ERR wrong number of arguments for 'AI.DAGRUN' command");
         }
-        goto cleanup;
+        return res;
     }
 
     int chainingOpCount = 0;
     int arg_pos = 1;
     bool load_complete = false;
     bool persist_complete = false;
     bool timeout_complete = false;
+    array_new_on_stack(RAI_DagOp *, 10, dag_ops);
 
     // The first arg is "AI.DAGRUN", so we go over from the next arg.
     while (arg_pos < argc) {
@@ -273,7 +275,7 @@ int ParseDAGRunCommand(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModuleS
         }
 
         if (!strcasecmp(arg_string, "|>") && arg_pos < argc - 1) {
-            RAI_DagOp *currentOp = _AddEmptyOp(rinfo);
+            RAI_DagOp *currentOp = _AddEmptyOp(&dag_ops);
             chainingOpCount++;
             int args_num = _CollectOpArgs(argv, argc, ++arg_pos, currentOp);
             arg_pos += args_num;
@@ -283,19 +285,30 @@ int ParseDAGRunCommand(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModuleS
         RAI_SetError(rinfo->err, RAI_EDAGBUILDER, "ERR Invalid DAGRUN command");
         goto cleanup;
     }
-    rinfo->dagOpCount = array_len(rinfo->dagOps);
-    if (rinfo->dagOpCount < 1) {
+
+    if (array_len(dag_ops) < 1) {
         RAI_SetError(rinfo->err, RAI_EDAGBUILDER, "ERR DAG is empty");
         goto cleanup;
     }
-    if (ParseDAGOps(rinfo, 0, rinfo->dagOpCount) != REDISMODULE_OK) {
+
+    if (ParseDAGOps(rinfo, dag_ops) != REDISMODULE_OK) {
+        for (size_t i = 0; i < array_len(dag_ops); i++) {
+            RAI_FreeDagOp(dag_ops[i]);
+        }
         goto cleanup;
     }
+    // After validating all the ops, insert them to the DAG.
+    for (size_t i = 0; i < array_len(dag_ops); i++) {
+        rinfo->dagOps = array_append(rinfo->dagOps, dag_ops[i]);
+    }
+    rinfo->dagOpCount = array_len(rinfo->dagOps);
+
     if (MangleTensorsNames(rinfo) != REDISMODULE_OK) {
         goto cleanup;
     }
-    return REDISMODULE_OK;
+    res = REDISMODULE_OK;
 
 cleanup:
-    return REDISMODULE_ERR;
+    array_free(dag_ops);
+    return res;
 }

src/DAG/dag_parser.h

@@ -15,14 +15,12 @@ int ParseDAGRunCommand(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModuleS
                        int argc, bool dag_ro);
 
 /**
- * @brief Parse the arguments of ops in the DAGRUN command and build (or extend) the DagOp object
- * accordingly.
- * @param rinfo The DAG run info with its op, where every op has an argv field that points to an
- * array of RedisModule strings the represents the op, and an argc field which is the number of
+ * @brief Parse the arguments of the given ops in the DAGRUN command and build every op accordingly.
+ * @param rinfo The DAG run info that will be populated with the ops if they are valid.
+ * with its op,
+ * @param ops A local array of ops, where every op has an argv field that points to an
+ * array of RedisModule strings arguments, and an argc field which is the number of
  * args.
- * @param first_op_ind The index of the first op in the for which we parse its argument and build
- * it.
- * @param num_ops The number of ops in the DAG the need to be parsed.
  * @return Returns REDISMODULE_OK if the command is valid, REDISMODULE_ERR otherwise.
  */
-int ParseDAGOps(RedisAI_RunInfo *rinfo, size_t first_op_ind, size_t num_ops);
+int ParseDAGOps(RedisAI_RunInfo *rinfo, RAI_DagOp **ops);

src/command_parser.c

@@ -128,7 +128,6 @@ int ParseModelRunCommand(RedisAI_RunInfo *rinfo, RAI_DagOp *currentOp, RedisModu
     int res = REDISMODULE_ERR;
     // Build a ModelRunCtx from command.
     RedisModuleCtx *ctx = RedisModule_GetThreadSafeContext(NULL);
-    // int lock_status = RedisModule_ThreadSafeContextTryLock(ctx);
     RAI_Model *model;
     long long timeout = 0;
     if (_ModelRunCommand_ParseArgs(argv, argc, ctx, &model, rinfo->err, &currentOp->inkeys,
@@ -157,9 +156,6 @@ int ParseModelRunCommand(RedisAI_RunInfo *rinfo, RAI_DagOp *currentOp, RedisModu
     res = REDISMODULE_OK;
 
 cleanup:
-    // if (lock_status == REDISMODULE_OK) {
-    // RedisModule_ThreadSafeContextUnlock(ctx);
-    //}
     RedisModule_FreeThreadSafeContext(ctx);
     return res;
 }
@@ -392,8 +388,7 @@ int RedisAI_ExecuteCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int ar
     rinfo->OnFinish = DAG_ReplyAndUnblock;
     rinfo->client = RedisModule_BlockClient(ctx, RedisAI_DagRun_Reply, NULL, RunInfo_FreeData, 0);
     if (DAG_InsertDAGToQueue(rinfo) != REDISMODULE_OK) {
-        RedisModule_ReplyWithError(ctx, rinfo->err->detail_oneline);
-        RAI_FreeRunInfo(rinfo);
+        RedisModule_UnblockClient(rinfo->client, rinfo);
         return REDISMODULE_ERR;
     }
     return REDISMODULE_OK;

src/redisai.c

@@ -953,6 +953,7 @@ static int RedisAI_RegisterApi(RedisModuleCtx *ctx) {
     REGISTER_API(GetError, ctx);
     REGISTER_API(GetErrorOneLine, ctx);
     REGISTER_API(GetErrorCode, ctx);
+    REGISTER_API(SetError, ctx);
 
     REGISTER_API(TensorCreate, ctx);
     REGISTER_API(TensorCreateByConcatenatingTensors, ctx);
@@ -975,6 +976,7 @@ static int RedisAI_RegisterApi(RedisModuleCtx *ctx) {
     REGISTER_API(ModelCreate, ctx);
     REGISTER_API(ModelFree, ctx);
     REGISTER_API(ModelRunCtxCreate, ctx);
+    REGISTER_API(GetModelFromKeyspace, ctx);
     REGISTER_API(ModelRunCtxAddInput, ctx);
     REGISTER_API(ModelRunCtxAddOutput, ctx);
     REGISTER_API(ModelRunCtxNumOutputs, ctx);
@@ -1017,7 +1019,7 @@ static int RedisAI_RegisterApi(RedisModuleCtx *ctx) {
     REGISTER_API(DAGNumOutputs, ctx);
     REGISTER_API(DAGOutputTensor, ctx);
     REGISTER_API(DAGRunError, ctx);
-    REGISTER_API(DAGCopyOpStatus, ctx);
+    REGISTER_API(DAGGetError, ctx);
     REGISTER_API(DAGRunOpFree, ctx);
     REGISTER_API(DAGFree, ctx);