Merge pull request #334 from RedisAI/dagro

Add DAGRUNRO command

Author: filipecosta90

src/redisai.c

@@ -986,6 +986,115 @@ int RedisAI_DagRun_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv,
   return REDISMODULE_OK;
 }
 
+/**
+ * AI.DAGRUNRO [LOAD <nkeys> key1 key2... ] |> [COMMAND1] |> [COMMAND2] |> [COMMANDN]
+ *
+ * Read-only (no PERSIST) DAG execution.
+ * The request is queued and evaded asynchronously from a separate thread. The
+ * client blocks until the computation finishes.
+ */
+int RedisAI_DagRunRO_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv,
+                                int argc) {
+  if (argc < 4) return RedisModule_WrongArity(ctx);
+
+  RedisAI_RunInfo *rinfo = NULL;
+  if (RAI_InitRunInfo(&rinfo) == REDISMODULE_ERR) {
+    return RedisModule_ReplyWithError(ctx, "ERR Unable to allocate the memory and initialise the RedisAI_RunInfo structure");
+  }
+  rinfo->use_local_context = 1;
+  RAI_DagOp* currentDagOp = NULL;
+  RAI_InitDagOp(&currentDagOp);
+  array_append(rinfo->dagOps,currentDagOp);
+
+  int loadFlag=0;
+  int chainingOpCount=0;
+  const char* deviceStr = NULL;
+
+  for (size_t argpos = 1; argpos <= argc - 1; argpos++) {
+    const char *arg_string = RedisModule_StringPtrLen(argv[argpos], NULL);
+    if (!strcasecmp(arg_string, "LOAD")) {
+      loadFlag=1;
+      const int parse_result = RAI_parseDAGLoadArgs(
+          ctx, &argv[argpos], argc - argpos,&(rinfo->dagTensorsLoadedContext), &(rinfo->dagTensorsContext), "|>");
+      if (parse_result > 0) {
+        argpos += parse_result - 1;
+      } else {
+        RAI_FreeRunInfo(ctx,rinfo);
+        return REDISMODULE_ERR;
+      }
+    } else if (!strcasecmp(arg_string, "PERSIST")) {
+      return RedisModule_ReplyWithError(ctx,"ERR PERSIST cannot be specified in a read-only DAG");;
+    } else if (!strcasecmp(arg_string, "|>")) {
+      // on the first pipe operator, if LOAD or PERSIST were used, we've already
+      // allocated memory
+      if (!(loadFlag == 1 && chainingOpCount == 0)) {
+        rinfo->dagNumberCommands++;
+        RAI_DagOp *currentDagOp = NULL;
+        RAI_InitDagOp(&currentDagOp);
+        array_append(rinfo->dagOps, currentDagOp);
+      }
+      chainingOpCount++;
+    } else {
+      if (!strcasecmp(arg_string, "AI.TENSORGET")) {
+        rinfo->dagOps[rinfo->dagNumberCommands]->commandType = REDISAI_DAG_CMD_TENSORGET;
+      }
+      if (!strcasecmp(arg_string, "AI.TENSORSET")) {
+        rinfo->dagOps[rinfo->dagNumberCommands]->commandType = REDISAI_DAG_CMD_TENSORSET;
+      }
+      if (!strcasecmp(arg_string, "AI.MODELRUN")) {
+        if (argc - 2 < argpos) {
+          return RedisModule_WrongArity(ctx);
+        }
+        rinfo->dagOps[rinfo->dagNumberCommands]->commandType = REDISAI_DAG_CMD_MODELRUN;
+        RAI_Model *mto;
+        RedisModuleKey *modelKey;
+        const int status = RAI_GetModelFromKeyspace(ctx, argv[argpos+1], &modelKey,
+                                                    &mto, REDISMODULE_READ);
+        if (status == REDISMODULE_ERR) {
+          RAI_FreeRunInfo(ctx,rinfo);
+          return REDISMODULE_ERR;
+        } 
+        if (deviceStr==NULL){
+          deviceStr=mto->devicestr;
+        }else{
+          // If the device strings are not equivalent, reply with error ( for now )
+          if(strcasecmp(mto->devicestr, deviceStr)!=0){            
+            RAI_FreeRunInfo(ctx,rinfo);
+            return RedisModule_ReplyWithError(ctx,"ERR multi-device DAGs not supported yet");;
+          }
+        }
+        rinfo->dagOps[rinfo->dagNumberCommands]->runkey = argv[argpos];
+        rinfo->dagOps[rinfo->dagNumberCommands]->mctx =
+            RAI_ModelRunCtxCreate(mto);
+      }
+      RedisModule_RetainString(NULL, argv[argpos]);
+      array_append(rinfo->dagOps[rinfo->dagNumberCommands]->argv, argv[argpos]);
+      rinfo->dagOps[rinfo->dagNumberCommands]->argc++;
+    }
+  }
+
+  RunQueueInfo *run_queue_info = NULL;
+  // If there was no MODELRUN on the DAG, we default all ops to CPU
+  if(deviceStr==NULL){
+    deviceStr="CPU";
+  }
+  // If the queue does not exist, initialize it
+  if (ensureRunQueue(deviceStr,&run_queue_info) == REDISMODULE_ERR) {
+    RAI_FreeRunInfo(ctx,rinfo);
+    return RedisModule_ReplyWithError(
+        ctx, "ERR Queue not initialized for device");
+  }
+
+  rinfo->client = RedisModule_BlockClient(ctx, RedisAI_DagRun_Reply, NULL,
+                                          NULL, 0);
+
+  pthread_mutex_lock(&run_queue_info->run_queue_mutex);
+  queuePush(run_queue_info->run_queue, rinfo);
+  pthread_cond_signal(&run_queue_info->queue_condition_var);
+  pthread_mutex_unlock(&run_queue_info->run_queue_mutex);
+
+  return REDISMODULE_OK;
+}
 #define EXECUTION_PLAN_FREE_MSG 100
 
 #define REGISTER_API(name, ctx) \
@@ -1135,6 +1244,10 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
       == REDISMODULE_ERR)
     return REDISMODULE_ERR;
 
+  if (RedisModule_CreateCommand(ctx, "ai.dagrunro", RedisAI_DagRunRO_RedisCommand, "readonly", 3, 3, 1)
+      == REDISMODULE_ERR)
+    return REDISMODULE_ERR;
+
   // Default configs
   RAI_BackendsPath = NULL;
   perqueueThreadPoolSize = REDISAI_DEFAULT_THREADS_PER_QUEUE;

test/tests_dag.py

@@ -98,6 +98,42 @@ def test_dag_common_errors(env):
         env.assertEqual("invalid or negative value found in number of keys to LOAD",exception.__str__())
 
 
+def test_dagro_common_errors(env):
+    con = env.getConnection()
+
+    # ERR unsupported command within DAG
+    try:
+        command = "AI.DAGRUNRO |> "\
+                "AI.DONTEXIST tensor1 FLOAT 1 2 VALUES 5 10"
+
+        ret = con.execute_command(command)
+    except Exception as e:
+        exception = e
+        env.assertEqual(type(exception), redis.exceptions.ResponseError)
+        env.assertEqual("ERR unsupported command within DAG",exception.__str__())
+
+    # ERR wrong number of arguments for 'AI.DAGRUN' command
+    try:
+        command = "AI.DAGRUNRO "
+
+        ret = con.execute_command(command)
+    except Exception as e:
+        exception = e
+        env.assertEqual(type(exception), redis.exceptions.ResponseError)
+        env.assertEqual("wrong number of arguments for 'AI.DAGRUNRO' command",exception.__str__())
+
+    # ERR invalid or negative value found in number of keys to LOAD
+    try:
+        command = "AI.DAGRUNRO LOAD notnumber |> "\
+                "AI.TENSORSET tensor1 FLOAT 1 2 VALUES 5 10"
+
+        ret = con.execute_command(command)
+    except Exception as e:
+        exception = e
+        env.assertEqual(type(exception), redis.exceptions.ResponseError)
+        env.assertEqual("invalid or negative value found in number of keys to LOAD",exception.__str__())
+
+
 def test_dag_modelrun_financialNet_errors(env):
     con = env.getConnection()
 
@@ -112,7 +148,6 @@ def test_dag_modelrun_financialNet_errors(env):
                                   'FLOAT', 1, 256,
                                   'BLOB', creditcard_referencedata[0].tobytes())
     env.assertEqual(ret, b'OK')
-    
 
     # ERR wrong number of inputs
     try:
@@ -131,7 +166,6 @@ def test_dag_modelrun_financialNet_errors(env):
         env.assertEqual(type(exception), redis.exceptions.ResponseError)
         env.assertEqual("ERR unsupported command within DAG",exception.__str__())
 
-        
 
 def test_dag_local_tensorset(env):
     con = env.getConnection()
@@ -147,6 +181,22 @@ def test_dag_local_tensorset(env):
     ret = con.execute_command("EXISTS volatile_tensor")
     env.assertEqual(ret, 0 )
 
+
+def test_dagro_local_tensorset(env):
+    con = env.getConnection()
+
+    command = "AI.DAGRUNRO "\
+        "AI.TENSORSET volatile_tensor1 FLOAT 1 2 VALUES 5 10 |> "\
+        "AI.TENSORSET volatile_tensor2 FLOAT 1 2 VALUES 5 10 "
+
+    ret = con.execute_command(command)
+    env.assertEqual(ret, [b'OK',b'OK'])
+
+    # assert that transaction tensor does not exist
+    ret = con.execute_command("EXISTS volatile_tensor")
+    env.assertEqual(ret, 0 )
+
+
 def test_dag_local_tensorset_persist(env):
     con = env.getConnection()
 
@@ -165,6 +215,21 @@ def test_dag_local_tensorset_persist(env):
     env.assertEqual(ret, [b'dtype', b'FLOAT', b'shape', [1, 2], b'values', [b'5', b'10']])
 
 
+def test_dagro_local_tensorset_persist(env):
+    con = env.getConnection()
+
+    command = "AI.DAGRUNRO "\
+        "PERSIST 1 tensor1 |> "\
+        "AI.TENSORSET tensor1 FLOAT 1 2 VALUES 5 10"
+
+    try:
+        con.execute_command(command)
+    except Exception as e:
+        exception = e
+        env.assertEqual(type(exception), redis.exceptions.ResponseError)
+        env.assertEqual("PERSIST cannot be specified in a read-only DAG", exception.__str__())
+
+
 def test_dag_multilocal_tensorset_persist(env):
     con = env.getConnection()
 
@@ -197,6 +262,7 @@ def test_dag_multilocal_tensorset_persist(env):
     ret = con.execute_command("AI.TENSORGET tensor3 META VALUES")
     env.assertEqual(ret, [b'dtype', b'FLOAT', b'shape', [1, 2], b'values', [b'5', b'10']])
 
+
 def test_dag_local_tensorset_tensorget_persist(env):
     con = env.getConnection()
 
@@ -282,6 +348,20 @@ def test_dag_keyspace_tensorget(env):
     env.assertEqual(ret, [[b'5', b'10']])
 
 
+def test_dagro_keyspace_tensorget(env):
+    con = env.getConnection()
+
+    ret = con.execute_command(
+        "AI.TENSORSET persisted_tensor FLOAT 1 2 VALUES 5 10")
+    env.assertEqual(ret, b'OK')
+
+    command = "AI.DAGRUNRO LOAD 1 persisted_tensor |> "\
+        "AI.TENSORGET persisted_tensor VALUES"
+
+    ret = con.execute_command(command)
+    env.assertEqual(ret, [[b'5', b'10']])
+
+
 def test_dag_keyspace_and_localcontext_tensorget(env):
     con = env.getConnection()
 
@@ -337,6 +417,7 @@ def test_dag_modelrun_financialNet_separate_tensorget(env):
         env.assertEqual(ret, 0 )
         tensor_number = tensor_number + 1
 
+
 def test_dag_modelrun_financialNet(env):
     con = env.getConnection()
 
@@ -373,6 +454,7 @@ def test_dag_modelrun_financialNet(env):
         env.assertEqual(ret, 0 )
         tensor_number = tensor_number + 1
 
+
 def test_dag_modelrun_financialNet_no_writes(env):
     con = env.getConnection()
 
@@ -422,7 +504,7 @@ def test_dag_modelrun_financialNet_no_writes(env):
         tensor_number = tensor_number + 1
 
 
-def test_dag_modelrun_financialNet_no_writes_multiple_modelruns(env):
+def test_dagro_modelrun_financialNet_no_writes_multiple_modelruns(env):
     con = env.getConnection()
 
     model_pb, creditcard_transactions, creditcard_referencedata = load_creditcardfraud_data(
@@ -442,7 +524,7 @@ def test_dag_modelrun_financialNet_no_writes_multiple_modelruns(env):
     tensor_number = 1
     for transaction_tensor in creditcard_transactions:
         ret = con.execute_command(
-            'AI.DAGRUN', 'LOAD', '1', 'referenceTensor:{}'.format(tensor_number), '|>',
+            'AI.DAGRUNRO', 'LOAD', '1', 'referenceTensor:{}'.format(tensor_number), '|>',
             'AI.TENSORSET', 'transactionTensor:{}'.format(tensor_number), 'FLOAT', 1, 30,'BLOB', transaction_tensor.tobytes(), '|>',
             'AI.MODELRUN', 'financialNet', 
                            'INPUTS', 'transactionTensor:{}'.format(tensor_number), 'referenceTensor:{}'.format(tensor_number),