Skip to content

Commit b9d4608

Browse files
alonre24alonre24
committed
Add SCRIPTRUN, PERSIST, TENSORSET ops to LLAPI.
test LLAPI (with/without errors) via test module (passed valgrind). Persist not checked (require enablement).
1 parent 7c696d3 commit b9d4608

File tree

10 files changed

+190
-29
lines changed

10 files changed

+190
-29
lines changed

‎src/DAG/dag.c‎

Lines changed: 0 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -757,14 +757,3 @@ void DAG_ReplyAndUnblock(RedisAI_OnFinishCtx *ctx, void *private_data) {
757757
if (rinfo->client)
758758
RedisModule_UnblockClient(rinfo->client, rinfo);
759759
}
760-
761-
void DAG_SetTensorsInLocalContext(RedisAI_RunInfo *rinfo) {
762-
for (size_t i = 0; i < rinfo->dagOpCount; i++) {
763-
RAI_DagOp *op = rinfo->dagOps[i];
764-
if (op->commandType == REDISAI_DAG_CMD_TENSORSET) {
765-
// Insert the tensor with its mangled (unique) name.
766-
void *t = (void *)RAI_TensorGetShallowCopy(op->outTensor);
767-
AI_dictReplace(rinfo->dagTensorsContext, (void *)op->outkeys[0], t);
768-
}
769-
}
770-
}

‎src/DAG/dag.h‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,4 @@ void RunInfo_FreeData(RedisModuleCtx *ctx, void *rinfo);
147147
*/
148148
void RedisAI_Disconnected(RedisModuleCtx *ctx, RedisModuleBlockedClient *bc);
149149

150-
void DAG_SetTensorsInLocalContext(RedisAI_RunInfo *rinfo);
151-
152150
#endif /* SRC_DAG_H_ */

‎src/DAG/dag_builder.c‎

Lines changed: 51 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,8 +54,7 @@ RAI_DAGRunCtx *RAI_DAGRunCtxCreate(void) {
5454
return (RAI_DAGRunCtx *)rinfo;
5555
}
5656

57-
RAI_DAGRunOp *RAI_DAGCreateModelRunOp(RAI_DAGRunCtx *run_info, RAI_Model *model) {
58-
RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)run_info;
57+
RAI_DAGRunOp *RAI_DAGCreateModelRunOp(RAI_Model *model) {
5958
RAI_ModelRunCtx *mctx = RAI_ModelRunCtxCreate(model);
6059
RAI_DagOp *op;
6160
RAI_InitDagOp(&op);
@@ -67,6 +66,18 @@ RAI_DAGRunOp *RAI_DAGCreateModelRunOp(RAI_DAGRunCtx *run_info, RAI_Model *model)
6766
return (RAI_DAGRunOp *)op;
6867
}
6968

69+
RAI_DAGRunOp *RAI_DAGCreateScriptRunOp(RAI_Script *script, const char *func_name) {
70+
RAI_ScriptRunCtx *sctx = RAI_ScriptRunCtxCreate(script, func_name);
71+
RAI_DagOp *op;
72+
RAI_InitDagOp(&op);
73+
74+
op->commandType = REDISAI_DAG_CMD_SCRIPTRUN;
75+
op->sctx = sctx;
76+
op->devicestr = script->devicestr;
77+
op->runkey = RAI_HoldString(NULL, (RedisModuleString *)script->infokey);
78+
return (RAI_DAGRunOp *)op;
79+
}
80+
7081
int RAI_DAGRunOpAddInput(RAI_DAGRunOp *DAGOp, const char *input) {
7182
RAI_DagOp *op = (RAI_DagOp *)DAGOp;
7283
RedisModuleString *inkey = RedisModule_CreateString(NULL, input, strlen(input));
@@ -115,6 +126,30 @@ int RAI_DAGLoadTensorRS(RAI_DAGRunCtx *run_info, RedisModuleString *t_name, RAI_
115126
return _RAI_DagLoadTensor(run_info, key_name, err);
116127
}
117128

129+
// todo: Persist tensors should not be part of dag reply, but before...
130+
int RAI_DAGAddPersistTensorRS(RAI_DAGRunCtx *run_info, RedisModuleString *t_name, RAI_Error *err) {
131+
132+
RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)run_info;
133+
if (AI_dictAdd(rinfo->dagTensorsPersistedContext, (void *)t_name, (void *)1) != DICT_OK) {
134+
RAI_SetError(err, RAI_EDAGBUILDER, "Tensor key to persist has already given");
135+
return REDISMODULE_ERR;
136+
}
137+
return REDISMODULE_OK;
138+
}
139+
140+
int RAI_DAGAddPersistTensor(RAI_DAGRunCtx *run_info, const char *t_name, RAI_Error *err) {
141+
142+
RedisModuleString *key_name = RedisModule_CreateString(NULL, t_name, strlen(t_name));
143+
RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)run_info;
144+
if (AI_dictAdd(rinfo->dagTensorsPersistedContext, (void *)key_name, (void *)1) != DICT_OK) {
145+
RAI_SetError(err, RAI_EDAGBUILDER, "Tensor key to persist has already given");
146+
RedisModule_FreeString(NULL, key_name);
147+
return REDISMODULE_ERR;
148+
}
149+
RedisModule_FreeString(NULL, key_name);
150+
return REDISMODULE_OK;
151+
}
152+
118153
int RAI_DAGAddTensorGet(RAI_DAGRunCtx *run_info, const char *t_name, RAI_Error *err) {
119154

120155
RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)run_info;
@@ -128,6 +163,20 @@ int RAI_DAGAddTensorGet(RAI_DAGRunCtx *run_info, const char *t_name, RAI_Error *
128163
return REDISMODULE_OK;
129164
}
130165

166+
int RAI_DAGAddTensorSet(RAI_DAGRunCtx *run_info, const char *t_name, RAI_Tensor *tensor) {
167+
168+
RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)run_info;
169+
RAI_DagOp *op;
170+
RAI_InitDagOp(&op);
171+
rinfo->dagOps = array_append(rinfo->dagOps, op);
172+
op->commandType = REDISAI_DAG_CMD_TENSORSET;
173+
op->devicestr = "CPU";
174+
RedisModuleString *name = RedisModule_CreateString(NULL, t_name, strlen(t_name));
175+
op->outkeys = array_append(op->outkeys, name);
176+
op->outTensor = RAI_TensorGetShallowCopy(tensor);
177+
return REDISMODULE_OK;
178+
}
179+
131180
void RAI_DAGRunOpFree(RAI_DAGRunOp *dagOp) {
132181
RAI_DagOp *op = (RAI_DagOp *)dagOp;
133182
RAI_FreeDagOp(op);

‎src/DAG/dag_builder.h‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,9 @@
44

55
RAI_DAGRunCtx *RAI_DAGRunCtxCreate(void);
66

7-
RAI_DAGRunOp *RAI_DAGCreateModelRunOp(RAI_DAGRunCtx *run_info, RAI_Model *model);
7+
RAI_DAGRunOp *RAI_DAGCreateModelRunOp(RAI_Model *model);
8+
9+
RAI_DAGRunOp *RAI_DAGCreateScriptRunOp(RAI_Script *script, const char *func_name);
810

911
int RAI_DAGRunOpAddInput(RAI_DAGRunOp *DAGOp, const char *input);
1012

@@ -16,6 +18,12 @@ int RAI_DAGLoadTensor(RAI_DAGRunCtx *run_info, const char *t_name, RAI_Error *er
1618

1719
int RAI_DAGLoadTensorRS(RAI_DAGRunCtx *run_info, RedisModuleString *t_name, RAI_Error *err);
1820

21+
int RAI_DAGAddPersistTensor(RAI_DAGRunCtx *run_info, const char *t_name, RAI_Error *err);
22+
23+
int RAI_DAGAddPersistTensorRS(RAI_DAGRunCtx *run_info, RedisModuleString *t_name, RAI_Error *err);
24+
25+
int RAI_DAGAddTensorSet(RAI_DAGRunCtx *run_info, const char *t_name, RAI_Tensor *tensor);
26+
1927
int RAI_DAGAddTensorGet(RAI_DAGRunCtx *run_info, const char *t_name, RAI_Error *err);
2028

2129
void RAI_DAGFree(RAI_DAGRunCtx *run_info);

‎src/DAG/dag_execute.c‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,17 @@
33
#include "background_workers.h"
44
#include "util/string_utils.h"
55

6+
void _DAG_SetTensorsInLocalContext(RedisAI_RunInfo *rinfo) {
7+
for (size_t i = 0; i < rinfo->dagOpCount; i++) {
8+
RAI_DagOp *op = rinfo->dagOps[i];
9+
if (op->commandType == REDISAI_DAG_CMD_TENSORSET) {
10+
// Insert the tensor with its mangled (unique) name.
11+
void *t = (void *)RAI_TensorGetShallowCopy(op->outTensor);
12+
AI_dictReplace(rinfo->dagTensorsContext, (void *)op->outkeys[0], t);
13+
}
14+
}
15+
}
16+
617
int MangleTensorsNames(RedisAI_RunInfo *rinfo) {
718

819
int res = REDISMODULE_ERR;
@@ -118,6 +129,9 @@ int MangleTensorsNames(RedisAI_RunInfo *rinfo) {
118129
rinfo->dagOps[i]->devicestr = "CPU";
119130
}
120131
}
132+
// Tensors from TENSORSET ops are ready to be put in DAG local context under their mangled
133+
// names.
134+
_DAG_SetTensorsInLocalContext(rinfo);
121135
res = REDISMODULE_OK;
122136

123137
cleanup : {

‎src/DAG/dag_parser.c‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -279,8 +279,6 @@ int ParseDAGRunCommand(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModuleS
279279
RedisModule_ReplyWithError(ctx, rinfo->err->detail_oneline);
280280
goto cleanup;
281281
}
282-
DAG_SetTensorsInLocalContext(rinfo);
283-
284282
return REDISMODULE_OK;
285283

286284
cleanup:

‎src/redisai.c‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -994,11 +994,15 @@ static int RedisAI_RegisterApi(RedisModuleCtx *ctx) {
994994

995995
REGISTER_API(DAGRunCtxCreate, ctx);
996996
REGISTER_API(DAGCreateModelRunOp, ctx);
997+
REGISTER_API(DAGCreateScriptRunOp, ctx);
997998
REGISTER_API(DAGRunOpAddInput, ctx);
998999
REGISTER_API(DAGRunOpAddOutput, ctx);
9991000
REGISTER_API(DAGAddRunOp, ctx);
10001001
REGISTER_API(DAGLoadTensor, ctx);
10011002
REGISTER_API(DAGLoadTensorRS, ctx);
1003+
REGISTER_API(DAGAddPersistTensor, ctx);
1004+
REGISTER_API(DAGAddPersistTensorRS, ctx);
1005+
REGISTER_API(DAGAddTensorSet, ctx);
10021006
REGISTER_API(DAGAddTensorGet, ctx);
10031007
REGISTER_API(DAGRun, ctx);
10041008
REGISTER_API(DAGNumOutputs, ctx);

‎src/redisai.h‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -127,8 +127,9 @@ int MODULE_API_FUNC(RedisAI_ScriptRunAsync)(RAI_ScriptRunCtx *sctx, RAI_OnFinish
127127
RAI_ScriptRunCtx *MODULE_API_FUNC(RedisAI_GetAsScriptRunCtx)(RAI_OnFinishCtx *ctx, RAI_Error *err);
128128

129129
RAI_DAGRunCtx *MODULE_API_FUNC(RedisAI_DAGRunCtxCreate)(void);
130-
RAI_DAGRunOp *MODULE_API_FUNC(RedisAI_DAGCreateModelRunOp)(RAI_DAGRunCtx *run_info,
131-
RAI_Model *model);
130+
RAI_DAGRunOp *MODULE_API_FUNC(RedisAI_DAGCreateModelRunOp)(RAI_Model *model);
131+
RAI_DAGRunOp *MODULE_API_FUNC(RedisAI_DAGCreateScriptRunOp)(RAI_Script *script,
132+
const char *func_name);
132133
int MODULE_API_FUNC(RedisAI_DAGRunOpAddInput)(RAI_DAGRunOp *DAGOp, const char *input);
133134
int MODULE_API_FUNC(RedisAI_DAGRunOpAddOutput)(RAI_DAGRunOp *DAGOp, const char *output);
134135
int MODULE_API_FUNC(RedisAI_DAGAddRunOp)(RAI_DAGRunCtx *run_info, RAI_DAGRunOp *DAGop,
@@ -137,6 +138,14 @@ int MODULE_API_FUNC(RedisAI_DAGLoadTensor)(RAI_DAGRunCtx *run_info, const char *
137138
RAI_Error *err);
138139
int MODULE_API_FUNC(RedisAI_DAGLoadTensorRS)(RAI_DAGRunCtx *run_info, RedisModuleString *t_name,
139140
RAI_Error *err);
141+
int MODULE_API_FUNC(RedisAI_DAGAddPersistTensor)(RAI_DAGRunCtx *run_info, const char *t_name,
142+
RAI_Error *err);
143+
144+
int MODULE_API_FUNC(RedisAI_DAGAddPersistTensorRS)(RAI_DAGRunCtx *run_info,
145+
RedisModuleString *t_name, RAI_Error *err);
146+
147+
int MODULE_API_FUNC(RedisAI_DAGAddTensorSet)(RAI_DAGRunCtx *run_info, const char *t_name,
148+
RAI_Tensor *tensor);
140149
int MODULE_API_FUNC(RedisAI_DAGAddTensorGet)(RAI_DAGRunCtx *run_info, const char *t_name,
141150
RAI_Error *err);
142151
int MODULE_API_FUNC(RedisAI_DAGRun)(RAI_DAGRunCtx *run_info, RAI_OnFinishCB DAGAsyncFinish,
@@ -233,11 +242,15 @@ static int RedisAI_Initialize(RedisModuleCtx *ctx) {
233242

234243
REDISAI_MODULE_INIT_FUNCTION(ctx, DAGRunCtxCreate);
235244
REDISAI_MODULE_INIT_FUNCTION(ctx, DAGCreateModelRunOp);
245+
REDISAI_MODULE_INIT_FUNCTION(ctx, DAGCreateScriptRunOp);
236246
REDISAI_MODULE_INIT_FUNCTION(ctx, DAGRunOpAddInput);
237247
REDISAI_MODULE_INIT_FUNCTION(ctx, DAGRunOpAddOutput);
238248
REDISAI_MODULE_INIT_FUNCTION(ctx, DAGAddRunOp);
239249
REDISAI_MODULE_INIT_FUNCTION(ctx, DAGLoadTensor);
240250
REDISAI_MODULE_INIT_FUNCTION(ctx, DAGLoadTensorRS);
251+
REDISAI_MODULE_INIT_FUNCTION(ctx, DAGAddPersistTensorRS);
252+
REDISAI_MODULE_INIT_FUNCTION(ctx, DAGAddPersistTensor);
253+
REDISAI_MODULE_INIT_FUNCTION(ctx, DAGAddTensorSet);
241254
REDISAI_MODULE_INIT_FUNCTION(ctx, DAGAddTensorGet);
242255
REDISAI_MODULE_INIT_FUNCTION(ctx, DAGRun);
243256
REDISAI_MODULE_INIT_FUNCTION(ctx, DAGNumOutputs);

‎tests/flow/tests_llapi.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -91,9 +91,15 @@ def test_dag_build_and_run(env):
9191

9292
with open(model_filename, 'rb') as f:
9393
model_pb = f.read()
94-
9594
ret = con.execute_command('AI.MODELSET', 'm{1}', 'TF', DEVICE,
9695
'INPUTS', 'a', 'b', 'OUTPUTS', 'mul', 'BLOB', model_pb)
9796
env.assertEqual(ret, b'OK')
97+
98+
script_filename = os.path.join(test_data_path, 'script.txt')
99+
with open(script_filename, 'rb') as f:
100+
script = f.read()
101+
ret = con.execute_command('AI.SCRIPTSET', 'myscript{1}', DEVICE, 'TAG', 'version1', 'SOURCE', script)
102+
env.assertEqual(ret, b'OK')
103+
98104
ret = con.execute_command("RAI_llapi.DAGrun")
99-
env.assertEqual(ret, b'DAG run success')
105+
env.assertEqual(ret, b'DAG run success')

‎tests/module/LLAPI_DAG.c‎

Lines changed: 89 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,14 +11,13 @@
1111
pthread_mutex_t global_lock = PTHREAD_MUTEX_INITIALIZER;
1212
pthread_cond_t global_cond = PTHREAD_COND_INITIALIZER;
1313

14-
static RAI_Model *_getModelFromKeySpace(RedisModuleCtx *ctx, const char *keyNameStr) {
14+
static void *_getFromKeySpace(RedisModuleCtx *ctx, const char *keyNameStr) {
1515

1616
RedisModuleString *keyRedisStr = RedisModule_CreateString(ctx, keyNameStr, strlen(keyNameStr));
1717
RedisModuleKey *key = RedisModule_OpenKey(ctx, keyRedisStr, REDISMODULE_READ);
18-
RAI_Model *model = RedisModule_ModuleTypeGetValue(key);
1918
RedisModule_FreeString(ctx, keyRedisStr);
2019
RedisModule_CloseKey(key);
21-
return model;
20+
return RedisModule_ModuleTypeGetValue(key);
2221
}
2322

2423
static void _DAGFinishFuncError(RAI_OnFinishCtx *onFinishCtx, void *private_data) {
@@ -58,13 +57,29 @@ static int _testLoadError(RAI_DAGRunCtx *run_info) {
5857
return LLAPIMODULE_ERR;
5958
}
6059

60+
static int _testPersistError(RAI_DAGRunCtx *run_info) {
61+
62+
RAI_Error *err;
63+
RedisAI_InitError(&err);
64+
int status = RedisAI_DAGAddPersistTensor(run_info, "t1", err);
65+
RedisModule_Assert(status == REDISMODULE_OK);
66+
status = RedisAI_DAGAddPersistTensor(run_info, "t1", err);
67+
if (strcmp(RedisAI_GetError(err), "Tensor key to persist has already given") == 0) {
68+
RedisModule_Assert(status == REDISMODULE_ERR);
69+
RedisAI_FreeError(err);
70+
return LLAPIMODULE_OK;
71+
}
72+
RedisAI_FreeError(err);
73+
return LLAPIMODULE_ERR;
74+
}
75+
6176
static int _testModelRunOpError(RedisModuleCtx *ctx, RAI_DAGRunCtx *run_info) {
6277

6378
RAI_Error *err;
6479
RedisAI_InitError(&err);
6580
// The model m{1} should exist in key space.
66-
RAI_Model *model = _getModelFromKeySpace(ctx, "m{1}");
67-
RAI_DAGRunOp *op = RedisAI_DAGCreateModelRunOp(run_info, model);
81+
RAI_Model *model = _getFromKeySpace(ctx, "m{1}");
82+
RAI_DAGRunOp *op = RedisAI_DAGCreateModelRunOp(model);
6883
RedisAI_DAGRunOpAddInput(op, "first_input");
6984

7085
// This model expect for 2 inputs not 1.
@@ -157,8 +172,8 @@ static int _testSimpleDAGRun(RedisModuleCtx *ctx, RAI_DAGRunCtx *run_info) {
157172
}
158173

159174
// The model m{1} should exist in key space.
160-
RAI_Model *model = _getModelFromKeySpace(ctx, "m{1}");
161-
RAI_DAGRunOp *op = RedisAI_DAGCreateModelRunOp(run_info, model);
175+
RAI_Model *model = _getFromKeySpace(ctx, "m{1}");
176+
RAI_DAGRunOp *op = RedisAI_DAGCreateModelRunOp(model);
162177
RedisAI_DAGRunOpAddInput(op, "a{1}");
163178
RedisAI_DAGRunOpAddInput(op, "b{1}");
164179
RedisAI_DAGRunOpAddOutput(op, "output");
@@ -175,6 +190,7 @@ static int _testSimpleDAGRun(RedisModuleCtx *ctx, RAI_DAGRunCtx *run_info) {
175190
}
176191
// Wait until the onFinish callback returns.
177192
pthread_cond_wait(&global_cond, &global_lock);
193+
pthread_mutex_unlock(&global_lock);
178194

179195
// Verify that we received the expected tensor at the end of the run.
180196
RedisModule_Assert(array_len(outputs) == 1);
@@ -199,6 +215,62 @@ static int _testSimpleDAGRun(RedisModuleCtx *ctx, RAI_DAGRunCtx *run_info) {
199215
return res;
200216
}
201217

218+
static int _testSimpleDAGRun2(RedisModuleCtx *ctx, RAI_DAGRunCtx *run_info) {
219+
220+
RAI_Error *err;
221+
RedisAI_InitError(&err);
222+
RAI_Tensor **outputs = array_new(RAI_Tensor *, 1);
223+
int res = LLAPIMODULE_ERR;
224+
225+
RAI_Tensor *tensor = _getFromKeySpace(ctx, "a{1}");
226+
RedisAI_DAGAddTensorSet(run_info, "input1", tensor);
227+
tensor = _getFromKeySpace(ctx, "b{1}");
228+
RedisAI_DAGAddTensorSet(run_info, "input2", tensor);
229+
230+
// The script myscript{1} should exist in key space.
231+
RAI_Script *script = _getFromKeySpace(ctx, "myscript{1}");
232+
RAI_DAGRunOp *op = RedisAI_DAGCreateScriptRunOp(script, "bar");
233+
RedisAI_DAGRunOpAddInput(op, "input1");
234+
RedisAI_DAGRunOpAddInput(op, "input2");
235+
RedisAI_DAGRunOpAddOutput(op, "output");
236+
int status = RedisAI_DAGAddRunOp(run_info, op, err);
237+
if (status != REDISMODULE_OK) {
238+
goto cleanup;
239+
}
240+
241+
RedisAI_DAGAddTensorGet(run_info, "output", err);
242+
pthread_mutex_lock(&global_lock);
243+
if (RedisAI_DAGRun(run_info, _DAGFinishFunc, &outputs, err) != REDISMODULE_OK) {
244+
pthread_mutex_unlock(&global_lock);
245+
goto cleanup;
246+
}
247+
// Wait until the onFinish callback returns.
248+
pthread_cond_wait(&global_cond, &global_lock);
249+
pthread_mutex_unlock(&global_lock);
250+
251+
// Verify that we received the expected tensor at the end of the run.
252+
RedisModule_Assert(array_len(outputs) == 1);
253+
RAI_Tensor *out_tensor = outputs[0];
254+
double expceted[4] = {4, 6, 4, 6};
255+
double val[4];
256+
for (long long i = 0; i < 4; i++) {
257+
if(RedisAI_TensorGetValueAsDouble(out_tensor, i, &val[i]) != 0) {
258+
goto cleanup;
259+
}
260+
if (expceted[i] != val[i]) {
261+
goto cleanup;
262+
}
263+
}
264+
RedisAI_TensorFree(out_tensor);
265+
res = LLAPIMODULE_OK;
266+
267+
cleanup:
268+
RedisAI_FreeError(err);
269+
array_free(outputs);
270+
RedisAI_DAGFree(run_info);
271+
return res;
272+
}
273+
202274
int RAI_llapi_DAGRun(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
203275
REDISMODULE_NOT_USED(argv);
204276

@@ -213,6 +285,11 @@ int RAI_llapi_DAGRun(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
213285
RedisAI_DAGFree(run_info);
214286
return RedisModule_ReplyWithSimpleString(ctx, "LOAD error test failed");
215287
}
288+
// Test the case of a failure due to persist two tensors with the same name.
289+
if(_testPersistError(run_info) != LLAPIMODULE_OK) {
290+
RedisAI_DAGFree(run_info);
291+
return RedisModule_ReplyWithSimpleString(ctx, "PERSIST error test failed");
292+
}
216293
// Test the case of a failure due to addition of a non compatible MODELRUN op.
217294
if(_testModelRunOpError(ctx, run_info) != LLAPIMODULE_OK) {
218295
RedisAI_DAGFree(run_info);
@@ -234,5 +311,10 @@ int RAI_llapi_DAGRun(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
234311
if(_testSimpleDAGRun(ctx, run_info) != LLAPIMODULE_OK) {
235312
return RedisModule_ReplyWithSimpleString(ctx, "Simple DAG run test failed");
236313
}
314+
run_info = RedisAI_DAGRunCtxCreate();
315+
// Test the case of building and running a DAG with TENSORSET, SCRIPTRUN and PERSIST ops.
316+
if(_testSimpleDAGRun2(ctx, run_info) != LLAPIMODULE_OK) {
317+
return RedisModule_ReplyWithSimpleString(ctx, "Simple DAG run2 test failed");
318+
}
237319
return RedisModule_ReplyWithSimpleString(ctx, "DAG run success");
238320
}

0 commit comments

Comments
 (0)