Support Async script run via LLAPI

alonre24 · alonre24 · commit 34d153bace9e · 2020-12-28T16:45:30.000+02:00
diff --git a/src/DAG/dag.c b/src/DAG/dag.c
@@ -257,7 +257,7 @@ void RedisAI_DagRunSession_ScriptRun_Step(RedisAI_RunInfo *rinfo, RAI_DagOp *cur
         RAI_ContextUnlock(rinfo);
 
         for (uint i = 0; i < n_inkeys; i++) {
-            RAI_ScriptRunCtxAddInput(currentOp->sctx, inputTensors[i]);
+            RAI_ScriptRunCtxAddInput(currentOp->sctx, inputTensors[i], currentOp->err);
         }
         for (uint i = 0; i < n_outkeys; i++) {
             RAI_ScriptRunCtxAddOutput(currentOp->sctx);
@@ -593,16 +593,16 @@ static void _PersistTensors(RedisModuleCtx *ctx, RedisAI_RunInfo *rinfo) {
                                        "ERR specified persistent key that was not used in DAG");
             rinfo->dagReplyLength++;
             RedisModule_Log(ctx, "warning",
-                            "on DAGRUN's PERSIST pecified persistent key (%s) that "
+                            "on DAGRUN's PERSIST specified persistent key (%s) that "
                             "was not used on DAG. Logging all local context keys",
-                            persist_key_name);
+                            RedisModule_StringPtrLen(persist_key_name, NULL));
             AI_dictIterator *local_iter = AI_dictGetSafeIterator(rinfo->dagTensorsContext);
             AI_dictEntry *local_entry = AI_dictNext(local_iter);
 
             while (local_entry) {
                 RedisModuleString *localcontext_key_name = AI_dictGetKey(local_entry);
                 RedisModule_Log(ctx, "warning", "DAG's local context key (%s)",
-                                localcontext_key_name);
+                                RedisModule_StringPtrLen(localcontext_key_name, NULL));
                 local_entry = AI_dictNext(local_iter);
             }
             AI_dictReleaseIterator(local_iter);
diff --git a/src/command_parser.c b/src/command_parser.c
@@ -263,6 +263,7 @@ static int _ScriptRunCtx_SetParams(RedisModuleCtx *ctx, RedisModuleString **inke
 
     RAI_Tensor *t;
     RedisModuleKey *key;
+    RAI_Error *err;
     size_t ninputs = array_len(inkeys), noutputs = array_len(outkeys);
     for (size_t i = 0; i < ninputs; i++) {
         const int status = RAI_GetTensorFromKeyspace(ctx, inkeys[i], &key, &t, REDISMODULE_READ);
@@ -271,7 +272,7 @@ static int _ScriptRunCtx_SetParams(RedisModuleCtx *ctx, RedisModuleString **inke
                             RedisModule_StringPtrLen(inkeys[i], NULL));
             return REDISMODULE_ERR;
         }
-        RAI_ScriptRunCtxAddInput(sctx, t);
+        RAI_ScriptRunCtxAddInput(sctx, t, err);
     }
     for (size_t i = 0; i < noutputs; i++) {
         RAI_ScriptRunCtxAddOutput(sctx);
diff --git a/src/redisai.c b/src/redisai.c
@@ -969,7 +969,7 @@ static int RedisAI_RegisterApi(RedisModuleCtx *ctx) {
     REGISTER_API(ModelGetShallowCopy, ctx);
     REGISTER_API(ModelRedisType, ctx);
     REGISTER_API(ModelRunAsync, ctx);
-    REGISTER_API(GetAsModelRunCtx, ctx)
+    REGISTER_API(GetAsModelRunCtx, ctx);
 
     REGISTER_API(ScriptCreate, ctx);
     REGISTER_API(ScriptFree, ctx);
@@ -983,6 +983,8 @@ static int RedisAI_RegisterApi(RedisModuleCtx *ctx) {
     REGISTER_API(ScriptRun, ctx);
     REGISTER_API(ScriptGetShallowCopy, ctx);
     REGISTER_API(ScriptRedisType, ctx);
+    REGISTER_API(ScriptRunAsync, ctx);
+    REGISTER_API(GetAsScriptRunCtx, ctx);
 
     return REDISMODULE_OK;
 }
diff --git a/src/redisai.h b/src/redisai.h
@@ -121,6 +121,9 @@ void MODULE_API_FUNC(RedisAI_ScriptRunCtxFree)(RAI_ScriptRunCtx *sctx);
 int MODULE_API_FUNC(RedisAI_ScriptRun)(RAI_ScriptRunCtx *sctx, RAI_Error *err);
 RAI_Script *MODULE_API_FUNC(RedisAI_ScriptGetShallowCopy)(RAI_Script *script);
 RedisModuleType *MODULE_API_FUNC(RedisAI_ScriptRedisType)(void);
+int MODULE_API_FUNC(RedisAI_ScriptRunAsync)(RAI_ScriptRunCtx *sctx, RAI_OnFinishCB DAGAsyncFinish,
+                                            void *private_data);
+RAI_ScriptRunCtx *MODULE_API_FUNC(RedisAI_GetAsScriptRunCtx)(RAI_OnFinishCtx *ctx, RAI_Error *err);
 
 int MODULE_API_FUNC(RedisAI_GetLLAPIVersion)();
 
@@ -204,6 +207,8 @@ static int RedisAI_Initialize(RedisModuleCtx *ctx) {
     REDISAI_MODULE_INIT_FUNCTION(ctx, ScriptRun);
     REDISAI_MODULE_INIT_FUNCTION(ctx, ScriptGetShallowCopy);
     REDISAI_MODULE_INIT_FUNCTION(ctx, ScriptRedisType);
+    REDISAI_MODULE_INIT_FUNCTION(ctx, ScriptRunAsync);
+    REDISAI_MODULE_INIT_FUNCTION(ctx, GetAsScriptRunCtx);
 
     if (RedisAI_GetLLAPIVersion() < REDISAI_LLAPI_VERSION) {
         return REDISMODULE_ERR;
diff --git a/src/run_info.c b/src/run_info.c
@@ -329,6 +329,7 @@ int RAI_RunInfoBatchable(struct RAI_DagOp *op1, struct RAI_DagOp *op2) {
 
     return 1;
 }
+
 RAI_ModelRunCtx *RAI_GetAsModelRunCtx(RedisAI_RunInfo *rinfo, RAI_Error *err) {
 
     RAI_DagOp *op = rinfo->dagOps[0];
@@ -342,3 +343,17 @@ RAI_ModelRunCtx *RAI_GetAsModelRunCtx(RedisAI_RunInfo *rinfo, RAI_Error *err) {
     RAI_FreeRunInfo(rinfo);
     return mctx;
 }
+
+RAI_ScriptRunCtx *RAI_GetAsScriptRunCtx(RedisAI_RunInfo *rinfo, RAI_Error *err) {
+
+    RAI_DagOp *op = rinfo->dagOps[0];
+    if (!rinfo->single_op_dag || !op->sctx) {
+        RAI_SetError(err, RedisAI_ErrorCode_EFINISHCTX, "Finish ctx is not a script run ctx");
+        return NULL;
+    }
+    RAI_SetError(err, RAI_GetErrorCode(op->err), RAI_GetError(op->err));
+    RAI_ScriptRunCtx *sctx = op->sctx;
+    rinfo->dagOps[0]->sctx = NULL;
+    RAI_FreeRunInfo(rinfo);
+    return sctx;
+}
diff --git a/src/run_info.h b/src/run_info.h
@@ -188,6 +188,14 @@ int RAI_RunInfoBatchable(struct RAI_DagOp *op1, struct RAI_DagOp *op2);
  */
 RAI_ModelRunCtx *RAI_GetAsModelRunCtx(RedisAI_RunInfo *rinfo, RAI_Error *err);
 
+/**
+ * Retreive the ScriptRunCtx of a DAG runInfo that contains a single op of type
+ * SCRIPTRUN.
+ * @param DAG runInfo.
+ * @return Pointer to the ScriptRunCtx in DAG's single op.
+ */
+RAI_ScriptRunCtx *RAI_GetAsScriptRunCtx(RedisAI_RunInfo *rinfo, RAI_Error *err);
+
 #ifdef __cplusplus
 } // extern "C"
 #endif
diff --git a/src/script.c b/src/script.c
@@ -7,7 +7,8 @@
  */
 
 #include "script.h"
-
+#include "run_info.h"
+#include "DAG/dag.h"
 #include "backends.h"
 #include "rmutil/alloc.h"
 #include "script_struct.h"
@@ -164,7 +165,7 @@ static int _Script_RunCtxAddParam(RAI_ScriptCtxParam **paramArr, RAI_Tensor *ten
     return 1;
 }
 
-int RAI_ScriptRunCtxAddInput(RAI_ScriptRunCtx *sctx, RAI_Tensor *inputTensor) {
+int RAI_ScriptRunCtxAddInput(RAI_ScriptRunCtx *sctx, RAI_Tensor *inputTensor, RAI_Error *error) {
     // Even if variadic is set, we still allow to add inputs in the LLAPI
     return _Script_RunCtxAddParam(&sctx->inputs, inputTensor);
 }
@@ -352,3 +353,26 @@ int RedisAI_Parse_ScriptRun_RedisCommand(RedisModuleCtx *ctx, RedisModuleString
 }
 
 RedisModuleType *RAI_ScriptRedisType(void) { return RedisAI_ScriptType; }
+
+int RAI_ScriptRunAsync(RAI_ScriptRunCtx *sctx, RAI_OnFinishCB ScriptAsyncFinish,
+                       void *private_data) {
+
+    RedisAI_RunInfo *rinfo = NULL;
+    if (RAI_InitRunInfo(&rinfo) == REDISMODULE_ERR) {
+        return REDISMODULE_ERR;
+    }
+    rinfo->single_op_dag = 1;
+    rinfo->OnFinish = (RedisAI_OnFinishCB)ScriptAsyncFinish;
+    rinfo->private_data = private_data;
+
+    RAI_DagOp *op;
+    if (RAI_InitDagOp(&op) == REDISMODULE_ERR) {
+        return REDISMODULE_ERR;
+    }
+    op->commandType = REDISAI_DAG_CMD_SCRIPTRUN;
+    Dag_PopulateOp(op, sctx, NULL, NULL, NULL);
+
+    rinfo->dagOps = array_append(rinfo->dagOps, op);
+    rinfo->dagOpCount = 1;
+    return DAG_InsertDAGToQueue(rinfo);
+}
diff --git a/src/script.h b/src/script.h
@@ -14,6 +14,7 @@
 #include "redismodule.h"
 #include "script_struct.h"
 #include "tensor.h"
+#include "run_info.h"
 
 extern RedisModuleType *RedisAI_ScriptType;
 
@@ -68,7 +69,7 @@ RAI_ScriptRunCtx *RAI_ScriptRunCtxCreate(RAI_Script *script, const char *fnname)
  * @param inputTensor input tensor structure
  * @return returns 1 on success, 0 in case of error.
  */
-int RAI_ScriptRunCtxAddInput(RAI_ScriptRunCtx *sctx, RAI_Tensor *inputTensor);
+int RAI_ScriptRunCtxAddInput(RAI_ScriptRunCtx *sctx, RAI_Tensor *inputTensor, RAI_Error *error);
 
 /**
  * For each Allocates a RAI_ScriptCtxParam data structure, and enforces a
@@ -212,4 +213,16 @@ void RedisAI_ReplyOrSetError(RedisModuleCtx *ctx, RAI_Error *error, RAI_ErrorCod
  */
 RedisModuleType *RAI_ScriptRedisType(void);
 
+/**
+ * Insert the ScriptRunCtx to the run queues so it will run asynchronously.
+ *
+ * @param sctx SodelRunCtx to execute
+ * @param ScriptAsyncFinish A callback that will be called when the execution is finished.
+ * @param private_data This is going to be sent to to the ScriptAsyncFinish.
+ * @return REDISMODULE_OK if the sctx was insert to the queues successfully, REDISMODULE_ERR
+ * otherwise.
+ */
+int RAI_ScriptRunAsync(RAI_ScriptRunCtx *sctx, RAI_OnFinishCB ScriptAsyncFinish,
+                       void *private_data);
+
 #endif /* SRC_SCRIPT_H_ */
diff --git a/tests/flow/tests_llapi.py b/tests/flow/tests_llapi.py
@@ -2,30 +2,45 @@
 
 from includes import *
 import os
+from functools import wraps
 
 '''
 python -m RLTest --test tests_llapi.py --module path/to/redisai.so
 '''
 
-goal_dir = os.path.join(os.getcwd(), "../module/LLAPI.so")
-TEST_MODULE_PATH = os.path.abspath(goal_dir)
 
+def ensure_test_module_loaded(f):
+    @wraps(f)
+    def wrapper(env, *args, **kwargs):
+        goal_dir = os.path.join(os.getcwd(), "../module/LLAPI.so")
+        TEST_MODULE_PATH = os.path.abspath(goal_dir)
+        con = env.getConnection()
+        modules = con.execute_command("MODULE", "LIST")
+        if b'RAI_llapi' in [module[1] for module in modules]:
+            return f(env, *args, **kwargs)
+        try:
+            ret = con.execute_command('MODULE', 'LOAD', TEST_MODULE_PATH)
+            env.assertEqual(ret, b'OK')
+            return f(env, *args, **kwargs)
+        except Exception as e:
+            env.assertFalse(True)
+            env.debugPrint(str(e), force=True)
+            return
+    return wrapper
 
+
+@ensure_test_module_loaded
 def test_basic_check(env):
 
     con = env.getConnection()
-    ret = con.execute_command("MODULE", "LOAD", TEST_MODULE_PATH)
-    env.assertEqual(ret, b'OK')
     ret = con.execute_command("RAI_llapi.basic_check")
     env.assertEqual(ret, b'OK')
 
 
+@ensure_test_module_loaded
 def test_model_run_async(env):
 
     con = env.getConnection()
-    ret = con.execute_command("MODULE", "LOAD", TEST_MODULE_PATH)
-    env.assertEqual(ret, b'OK')
-
     test_data_path = os.path.join(os.path.dirname(__file__), 'test_data')
     model_filename = os.path.join(test_data_path, 'graph.pb')
 
@@ -39,3 +54,25 @@ def test_model_run_async(env):
     con.execute_command('AI.TENSORSET', 'b{1}', 'FLOAT', 2, 2, 'VALUES', 2, 3, 2, 3)
     ret = con.execute_command("RAI_llapi.modelRun")
     env.assertEqual(ret, b'Async run success')
+
+
+@ensure_test_module_loaded
+def test_script_run_async(env):
+
+    con = env.getConnection()
+    test_data_path = os.path.join(os.path.dirname(__file__), 'test_data')
+    script_filename = os.path.join(test_data_path, 'script.txt')
+
+    with open(script_filename, 'rb') as f:
+        script = f.read()
+
+    ret = con.execute_command('AI.SCRIPTSET', 'myscript{1}', DEVICE, 'TAG', 'version1', 'SOURCE', script)
+    env.assertEqual(ret, b'OK')
+
+    ret = con.execute_command('AI.TENSORSET', 'a{1}', 'FLOAT', 2, 2, 'VALUES', 2, 3, 2, 3)
+    env.assertEqual(ret, b'OK')
+    ret = con.execute_command('AI.TENSORSET', 'b{1}', 'FLOAT', 2, 2, 'VALUES', 2, 3, 2, 3)
+    env.assertEqual(ret, b'OK')
+
+    ret = con.execute_command("RAI_llapi.scriptRun")
+    env.assertEqual(ret, b'Async run success')
diff --git a/tests/module/LLAPI.c b/tests/module/LLAPI.c
