Add test with gears, PR fixes.

alonre24 · alonre24 · commit 34c4e3eab210 · 2020-12-22T19:28:23.000+02:00
diff --git a/src/command_parser.c b/src/command_parser.c
@@ -102,8 +102,8 @@ static int _ModelRunCommand_ParseArgs(RedisModuleCtx *ctx, RedisModuleString **a
  * @return REDISMODULE_OK in case of success, REDISMODULE_ERR otherwise
  */
 
-static int ModelRunCtx_SetParams(RedisModuleCtx *ctx, RedisModuleString **inkeys,
-                                 RedisModuleString **outkeys, RAI_ModelRunCtx *mctx) {
+static int _ModelRunCtx_SetParams(RedisModuleCtx *ctx, RedisModuleString **inkeys,
+                                  RedisModuleString **outkeys, RAI_ModelRunCtx *mctx) {
 
     RAI_Model *model = mctx->model;
     RAI_Tensor *t;
@@ -152,7 +152,7 @@ int ParseModelRunCommand(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModul
     if (rinfo->single_op_dag) {
         rinfo->timeout = timeout;
         // Set params in ModelRunCtx, bring inputs from key space.
-        if (ModelRunCtx_SetParams(ctx, inkeys, outkeys, mctx) == REDISMODULE_ERR)
+        if (_ModelRunCtx_SetParams(ctx, inkeys, outkeys, mctx) == REDISMODULE_ERR)
             goto cleanup;
     }
     if (RAI_InitDagOp(&currentOp) == REDISMODULE_ERR) {
@@ -177,7 +177,7 @@ int ParseModelRunCommand(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModul
     if (runkey)
         RedisModule_FreeString(NULL, runkey);
     if (mctx)
-        RAI_ModelRunCtxFree(mctx, true);
+        RAI_ModelRunCtxFree(mctx);
     RAI_FreeRunInfo(rinfo);
     return REDISMODULE_ERR;
 }
diff --git a/src/modelRun_ctx.c b/src/modelRun_ctx.c
@@ -44,16 +44,14 @@ RAI_Tensor *RAI_ModelRunCtxOutputTensor(RAI_ModelRunCtx *mctx, size_t index) {
     return mctx->outputs[index].tensor;
 }
 
-void RAI_ModelRunCtxFree(RAI_ModelRunCtx *mctx, int freeTensors) {
-    if (freeTensors) {
-        for (size_t i = 0; i < array_len(mctx->inputs); ++i) {
-            RAI_TensorFree(mctx->inputs[i].tensor);
-        }
+void RAI_ModelRunCtxFree(RAI_ModelRunCtx *mctx) {
+    for (size_t i = 0; i < array_len(mctx->inputs); ++i) {
+        RAI_TensorFree(mctx->inputs[i].tensor);
+    }
 
-        for (size_t i = 0; i < array_len(mctx->outputs); ++i) {
-            if (mctx->outputs[i].tensor) {
-                RAI_TensorFree(mctx->outputs[i].tensor);
-            }
+    for (size_t i = 0; i < array_len(mctx->outputs); ++i) {
+        if (mctx->outputs[i].tensor) {
+            RAI_TensorFree(mctx->outputs[i].tensor);
         }
     }
 
diff --git a/src/modelRun_ctx.h b/src/modelRun_ctx.h
@@ -18,7 +18,7 @@ RAI_ModelRunCtx *RAI_ModelRunCtxCreate(RAI_Model *model);
  * @param mctx
  * @param freeTensors free input and output tensors or leave them allocated
  */
-void RAI_ModelRunCtxFree(RAI_ModelRunCtx *mctx, int freeTensors);
+void RAI_ModelRunCtxFree(RAI_ModelRunCtx *mctxs);
 
 /**
  * Allocates a RAI_ModelCtxParam data structure, and enforces a shallow copy of
diff --git a/src/redisai.c b/src/redisai.c
@@ -567,8 +567,7 @@ int RedisAI_ModelRun_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv,
     if (RedisModule_IsKeysPositionRequest(ctx)) {
         return RedisAI_ModelRun_IsKeysPositionRequest_ReportKeys(ctx, argv, argc);
     }
-    bool ro_dag = false;
-    return RedisAI_ExecuteCommand(ctx, argv, argc, CMD_MODELRUN, ro_dag);
+    return RedisAI_ExecuteCommand(ctx, argv, argc, CMD_MODELRUN, false);
 }
 
 /**
@@ -584,8 +583,7 @@ int RedisAI_ScriptRun_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv
         return RedisModule_WrongArity(ctx);
 
     // Convert The script run command into a DAG command that contains a single op.
-    bool ro_dag = false;
-    return RedisAI_ExecuteCommand(ctx, argv, argc, CMD_SCRIPTRUN, ro_dag);
+    return RedisAI_ExecuteCommand(ctx, argv, argc, CMD_SCRIPTRUN, false);
 }
 
 /**
@@ -892,8 +890,7 @@ int RedisAI_DagRun_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, i
     if (RedisModule_IsKeysPositionRequest(ctx)) {
         return RedisAI_DagRun_IsKeysPositionRequest_ReportKeys(ctx, argv, argc);
     }
-    bool ro_only = false;
-    return RedisAI_ExecuteCommand(ctx, argv, argc, CMD_DAGRUN, ro_only);
+    return RedisAI_ExecuteCommand(ctx, argv, argc, CMD_DAGRUN, false);
 }
 
 /**
diff --git a/src/redisai.h b/src/redisai.h
@@ -48,6 +48,7 @@ typedef void (*RAI_OnFinishCB)(RAI_OnFinishCtx *ctx, void *private_data);
 #define RedisAI_ErrorCode_ETENSORSET          14
 #define RedisAI_ErrorCode_ETENSORGET          15
 #define RedisAI_ErrorCode_EDAGRUN             16
+#define RedisAI_ErrorCode_EFINISHCTX          17
 
 enum RedisAI_DataFmt { REDISAI_DATA_BLOB = 0, REDISAI_DATA_VALUES, REDISAI_DATA_NONE };
 
@@ -93,7 +94,7 @@ int MODULE_API_FUNC(RedisAI_ModelRunCtxAddInput)(RAI_ModelRunCtx *mctx, const ch
 int MODULE_API_FUNC(RedisAI_ModelRunCtxAddOutput)(RAI_ModelRunCtx *mctx, const char *outputName);
 size_t MODULE_API_FUNC(RedisAI_ModelRunCtxNumOutputs)(RAI_ModelRunCtx *mctx);
 RAI_Tensor *MODULE_API_FUNC(RedisAI_ModelRunCtxOutputTensor)(RAI_ModelRunCtx *mctx, size_t index);
-void MODULE_API_FUNC(RedisAI_ModelRunCtxFree)(RAI_ModelRunCtx *mctx, bool freeTensors);
+void MODULE_API_FUNC(RedisAI_ModelRunCtxFree)(RAI_ModelRunCtx *mctx);
 int MODULE_API_FUNC(RedisAI_ModelRun)(RAI_ModelRunCtx **mctx, long long n, RAI_Error *err);
 RAI_Model *MODULE_API_FUNC(RedisAI_ModelGetShallowCopy)(RAI_Model *model);
 int MODULE_API_FUNC(RedisAI_ModelSerialize)(RAI_Model *model, char **buffer, size_t *len,
diff --git a/src/run_info.c b/src/run_info.c
@@ -179,7 +179,7 @@ void RAI_FreeDagOp(RAI_DagOp *dagOp) {
         array_free(dagOp->outTensors);
 
         if (dagOp->mctx) {
-            RAI_ModelRunCtxFree(dagOp->mctx, true);
+            RAI_ModelRunCtxFree(dagOp->mctx);
         }
         if (dagOp->sctx) {
             RAI_ScriptRunCtxFree(dagOp->sctx, true);
@@ -347,6 +347,10 @@ int RAI_RunInfoBatchable(struct RAI_DagOp *op1, struct RAI_DagOp *op2) {
 RAI_ModelRunCtx *RAI_GetAsModelRunCtx(RedisAI_RunInfo *rinfo, RAI_Error *err) {
 
     RAI_DagOp *op = rinfo->dagOps[0];
+    if (!rinfo->single_op_dag || !op->mctx) {
+        RAI_SetError(err, RedisAI_ErrorCode_EFINISHCTX, "Finish ctx is not a model run ctx");
+        return NULL;
+    }
     RAI_SetError(err, RAI_GetErrorCode(op->err), RAI_GetError(op->err));
     RAI_ModelRunCtx *mctx = op->mctx;
     rinfo->dagOps[0]->mctx = NULL;
diff --git a/tests/flow/tests_llapi.py b/tests/flow/tests_llapi.py
@@ -2,6 +2,7 @@
 
 from includes import *
 import os
+from functools import wraps
 
 '''
 python -m RLTest --test tests_llapi.py --module path/to/redisai.so
@@ -10,6 +11,25 @@
 goal_dir = os.path.join(os.getcwd(), "../module/LLAPI.so")
 TEST_MODULE_PATH = os.path.abspath(goal_dir)
 
+
+def skip_if_gears_not_loaded(f):
+    @wraps(f)
+    def wrapper(env, *args, **kwargs):
+        con = env.getConnection()
+        modules = con.execute_command("MODULE", "LIST")
+        if "rg" in [module[1] for module in modules]:
+            return f(env, *args, **kwargs)
+        try:
+            redisgears_path = os.path.join(os.path.dirname(__file__), '../../../RedisGears/redisgears.so')
+            ret = con.execute_command('MODULE', 'LOAD', redisgears_path)
+            env.assertEqual(ret, b'OK')
+            return f(env, *args, **kwargs)
+        except Exception as e:
+            env.debugPrint("skipping since RedisGears not loaded", force=True)
+            return
+    return wrapper
+
+
 def test_basic_check(env):
 
     con = env.getConnection()
@@ -38,3 +58,51 @@ 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')
+
+
+@skip_if_gears_not_loaded
+def test_model_run_async_via_gears(env):
+    script = '''
+import redisAI
+
+async def RedisAIModelRun(record):
+    keys = ['a{1}', 'b{1}']
+    tensors = redisAI.mgetTensorsFromKeyspace(keys)
+    modelRunner = redisAI.createModelRunner('m{1}')
+    redisAI.modelRunnerAddInput(modelRunner, 'a', tensors[0])
+    redisAI.modelRunnerAddInput(modelRunner, 'b', tensors[1])
+    redisAI.modelRunnerAddOutput(modelRunner, 'mul')
+    res = await redisAI.modelRunnerRunAsync(modelRunner)
+    if len(res[1]) > 0:
+        raise Exception(res[1][0])
+    redisAI.setTensorInKey('c{1}', res[0][0])
+    return "OK"
+
+GB("CommandReader").map(RedisAIModelRun).register(trigger="ModelRunAsyncTest")
+    '''
+    con = env.getConnection()
+    ret = con.execute_command('rg.pyexecute', script)
+    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')
+
+    with open(model_filename, 'rb') as f:
+        model_pb = f.read()
+
+    ret = con.execute_command('AI.MODELSET', 'm{1}', 'TF', DEVICE,
+                              'INPUTS', 'a', 'b', 'OUTPUTS', 'mul', 'BLOB', model_pb)
+    env.assertEqual(ret, b'OK')
+
+    ret = con.execute_command('AI.MODELGET', 'm{1}', 'META')
+    env.assertEqual(len(ret), 14)
+
+    con.execute_command('AI.TENSORSET', 'a{1}', 'FLOAT',
+                        2, 2, 'VALUES', 2, 3, 2, 3)
+    con.execute_command('AI.TENSORSET', 'b{1}', 'FLOAT',
+                        2, 2, 'VALUES', 2, 3, 2, 3)
+
+    ret = con.execute_command('rg.trigger', 'ModelRunAsyncTest')
+    env.assertEqual(ret[0], b'OK')
+    values = con.execute_command('AI.TENSORGET', 'c{1}', 'VALUES')
+    env.assertEqual(values, [b'4', b'9', b'4', b'9'])
diff --git a/tests/module/LLAPI.c b/tests/module/LLAPI.c
@@ -70,15 +70,15 @@ static int _ExecuteModelRunAsync(RedisModuleCtx *ctx, RAI_ModelRunCtx* mctx) {
 	pthread_mutex_lock(&global_lock);
 	if (RedisAI_ModelRunAsync(mctx, ModelFinishFunc, &status) != REDISMODULE_OK) {
 		pthread_mutex_unlock(&global_lock);
-		RedisAI_ModelRunCtxFree(mctx, true);
+		RedisAI_ModelRunCtxFree(mctx);
 		RedisModule_ReplyWithError(ctx, "Async run could not start");
 		return LLAPI_RUN_NONE;
 	}
 
 	// Wait until the onFinish callback returns.
 	pthread_cond_wait(&global_cond, &global_lock);
 	pthread_mutex_unlock(&global_lock);
-	RedisAI_ModelRunCtxFree(mctx, true);
+	RedisAI_ModelRunCtxFree(mctx);
 	return status;
 }
 

Original file line number	Diff line number	Diff line change
`@@ -567,8 +567,7 @@ int RedisAI_ModelRun_RedisCommand(RedisModuleCtx ctx, RedisModuleString *argv,`
`567`	`567`	`if (RedisModule_IsKeysPositionRequest(ctx)) {`
`568`	`568`	`return RedisAI_ModelRun_IsKeysPositionRequest_ReportKeys(ctx, argv, argc);`
`569`	`569`	`}`
`570`		`- bool ro_dag = false;`
`571`		`- return RedisAI_ExecuteCommand(ctx, argv, argc, CMD_MODELRUN, ro_dag);`
	`570`	`+ return RedisAI_ExecuteCommand(ctx, argv, argc, CMD_MODELRUN, false);`
`572`	`571`	`}`
`573`	`572`
`574`	`573`	`/**`
`@@ -584,8 +583,7 @@ int RedisAI_ScriptRun_RedisCommand(RedisModuleCtx ctx, RedisModuleString *argv`
`584`	`583`	`return RedisModule_WrongArity(ctx);`
`585`	`584`
`586`	`585`	`// Convert The script run command into a DAG command that contains a single op.`
`587`		`- bool ro_dag = false;`
`588`		`- return RedisAI_ExecuteCommand(ctx, argv, argc, CMD_SCRIPTRUN, ro_dag);`
	`586`	`+ return RedisAI_ExecuteCommand(ctx, argv, argc, CMD_SCRIPTRUN, false);`
`589`	`587`	`}`
`590`	`588`
`591`	`589`	`/**`
`@@ -892,8 +890,7 @@ int RedisAI_DagRun_RedisCommand(RedisModuleCtx ctx, RedisModuleString *argv, i`
`892`	`890`	`if (RedisModule_IsKeysPositionRequest(ctx)) {`
`893`	`891`	`return RedisAI_DagRun_IsKeysPositionRequest_ReportKeys(ctx, argv, argc);`
`894`	`892`	`}`
`895`		`- bool ro_only = false;`
`896`		`- return RedisAI_ExecuteCommand(ctx, argv, argc, CMD_DAGRUN, ro_only);`
	`893`	`+ return RedisAI_ExecuteCommand(ctx, argv, argc, CMD_DAGRUN, false);`
`897`	`894`	`}`
`898`	`895`
`899`	`896`	`/**`