Update TensorSet and TensorGet in DAG upon execution (#811)

* Set status of dag tensorSet op at parse time (when it is executed). Also, set tensorGet op status at reply time.

* Refactor validation of DAGCommand enum upon dag op execution

Author: alonre24

src/execution/DAG/dag.c

@@ -480,31 +480,17 @@ void Dag_SetTensorInGlobalCtx(RedisAI_RunInfo *rinfo, size_t index, RAI_Tensor *
 void RedisAI_DagRunSessionStep(RedisAI_RunInfo *rinfo, const char *devicestr) {
     RAI_DagOp *currentOp = RedisAI_DagCurrentOp(rinfo);
 
-    switch (currentOp->commandType) {
-    case REDISAI_DAG_CMD_TENSORSET: {
-        // TENSORSET op is done in parsing stage (consider removing it from dag ops).
-        currentOp->result = REDISMODULE_OK;
-        break;
-    }
-    case REDISAI_DAG_CMD_TENSORGET: {
-        // TENSORSET op is done when we finish (consider removing it from dag ops).
-        currentOp->result = REDISMODULE_OK;
-        break;
-    }
-    case REDISAI_DAG_CMD_MODELRUN: {
+    // Verify that the op type belongs to the DAGCommand enum.
+    VALIDATE_DAG_COMMAND(currentOp->commandType)
+
+    if (currentOp->commandType == REDISAI_DAG_CMD_MODELRUN) {
         RedisAI_DagRunSession_ModelRun_Step(rinfo, currentOp);
-        break;
-    }
-    case REDISAI_DAG_CMD_SCRIPTRUN: {
+    } else if (currentOp->commandType == REDISAI_DAG_CMD_SCRIPTRUN) {
         RedisAI_DagRunSession_ScriptRun_Step(rinfo, currentOp);
-        break;
-    }
-    default: {
-        /* unsupported DAG's command */
-        RAI_SetError(currentOp->err, RAI_EDAGRUN, "ERR unsupported command within DAG");
-        currentOp->result = REDISMODULE_ERR;
-        break;
-    }
+    } else {
+        // do nothing for tensorset (executed on parsing) and for tensorget (done
+        // on dag reply).
+        return;
     }
 
     if (currentOp->result != REDISMODULE_OK) {
@@ -568,26 +554,21 @@ int RedisAI_DagRun_Reply(RedisModuleCtx *ctx, RedisModuleString **argv, int argc
 
     for (size_t i = 0; i < n_dagOps; i++) {
         RAI_DagOp *currentOp = rinfo->dagOps[i];
+
         switch (currentOp->commandType) {
         case REDISAI_DAG_CMD_TENSORSET: {
             rinfo->dagReplyLength++;
-            if (currentOp->result == REDISMODULE_ERR) {
-                RedisModule_ReplyWithError(ctx, currentOp->err->detail_oneline);
-                dag_error = 1;
-            } else if (currentOp->result == -1) {
-                RedisModule_ReplyWithSimpleString(ctx, "NA");
-            } else {
-                RedisModule_ReplyWithSimpleString(ctx, "OK");
-            }
+            RedisModule_Assert(currentOp->result == REDISMODULE_OK);
+            RedisModule_ReplyWithSimpleString(ctx, "OK");
             break;
         }
 
         case REDISAI_DAG_CMD_TENSORGET: {
             rinfo->dagReplyLength++;
-            if (currentOp->result == -1) {
+            RAI_Tensor *t = Dag_GetTensorFromGlobalCtx(rinfo, currentOp->inkeys_indices[0]);
+            if (t == NULL) {
                 RedisModule_ReplyWithSimpleString(ctx, "NA");
             } else {
-                RAI_Tensor *t = Dag_GetTensorFromGlobalCtx(rinfo, currentOp->inkeys_indices[0]);
                 ReplyWithTensor(ctx, currentOp->fmt, t);
             }
             break;
@@ -636,8 +617,7 @@ int RedisAI_DagRun_Reply(RedisModuleCtx *ctx, RedisModuleString **argv, int argc
             break;
         }
         default:
-            /* no-op */
-            break;
+            RedisModule_Assert(false && "Dag reply - invalid op");
         }
     }
     if (dag_error) {

src/execution/DAG/dag_op.h

@@ -13,6 +13,9 @@ typedef enum DAGCommand {
     REDISAI_DAG_CMD_SCRIPTRUN
 } DAGCommand;
 
+#define VALIDATE_DAG_COMMAND(cmd)                                                                  \
+    RedisModule_Assert(cmd >= REDISAI_DAG_CMD_TENSORSET && cmd <= REDISAI_DAG_CMD_SCRIPTRUN);
+
 typedef struct RAI_DagOp {
     DAGCommand commandType;
     RedisModuleString *runkey;

src/execution/parsing/dag_parser.c

@@ -190,6 +190,7 @@ int ParseDAGExecuteOps(RedisAI_RunInfo *rinfo, RAI_DagOp **ops, bool ro) {
                                        rinfo->err) == -1) {
                 return REDISMODULE_ERR;
             }
+            currentOp->result = REDISMODULE_OK;
             continue;
         }
         if (!strcasecmp(arg_string, "AI.MODELEXECUTE")) {

src/execution/parsing/deprecated.c

@@ -599,8 +599,10 @@ int ParseDAGRunOps(RedisAI_RunInfo *rinfo, RAI_DagOp **ops) {
             RAI_HoldString(currentOp->argv[1]);
             currentOp->outkeys = array_append(currentOp->outkeys, currentOp->argv[1]);
             if (RAI_parseTensorSetArgs(currentOp->argv, currentOp->argc, &currentOp->outTensor, 0,
-                                       rinfo->err) == -1)
+                                       rinfo->err) == -1) {
                 goto cleanup;
+            }
+            currentOp->result = REDISMODULE_OK;
             continue;
         }
         if (!strcasecmp(arg_string, "AI.MODELRUN")) {

tests/flow/tests_sanitizer.py

@@ -31,15 +31,11 @@ def test_sanitizer_dagrun_mobilenet_v1(env):
         class_key = 'output{s}'
 
         ret = con.execute_command(
-            'AI.DAGRUN', '|>',
+            'AI.DAGEXECUTE', 'ROUTING', '{s}', '|>',
             'AI.TENSORSET', image_key, 'FLOAT', 1, 224, 224, 3, 'BLOB', img.tobytes(),
             '|>',
-            'AI.MODELRUN', model_name,
-                         'INPUTS', image_key,
-                         'OUTPUTS', class_key,
-                          '|>',
-            'AI.TENSORGET',  class_key, 'blob'
-        )
+            'AI.MODELEXECUTE', model_name, 'INPUTS', 1, image_key, 'OUTPUTS', 1, class_key,
+            '|>', 'AI.TENSORGET',  class_key, 'blob')
         env.assertEqual([b'OK', b'OK'], ret[:2])
         env.assertEqual(1001.0, len(ret[2])/4)