Skip to content

Commit cca931b

Browse files
filipecosta90filipecosta90
committed
Fixed flagged as "getkeys-api" during the registration ( AI.DAGRUN, AI.DAGRUN_RO, AI.MODELRUN, AI.SCRIPTRUN ) (#438)
* [wip] wip on fixing the commands being flagged as getkeys-api during the registration. ( AI.MODELRUN, AI.SCRIPTRUN, AI.DAGRUN, AI.DAGRUN_RO ) * [add] pytorch oss cluster tests passing * [add] TensorFlow lite tests passing on oss cluster * [add] ONNX tests passing on oss cluster * [add] TensorFlow tests passing on oss cluster * [wip] wip on dag * [add] DAG tests passing on oss cluster * [add] enabling oss cluster tests on CI * [add] bumping rmbuilder version from 6.0.1 to 6.0.5 on circleci * [fix] fixed potential invalid mem accesses on RedisAI_DagRunSyntaxParser, RAI_parseDAGPersistArgs * [fix] fixed RAI_FreeDagOp wrongfully calling RedisModule_Free on dagOp->inkeys, dagOp->outkeys * [fix] fixed dictKey allocation on RAI_parseDAGLoadArgs * [add] alter Sanitizer tests to accommodate keys on the same slot. * [add] extracted the methods to respond to RedisModule_IsKeysPositionRequest() from main code of dagrun,modelrun, and scriptrun to specific methods * [fix] Fixes per PR review * [fix] fixes per PR review * [fix] fix per CI ascii conversion error on test_dagro_modelrun_financialNet_no_writes_multiple_modelruns
1 parent 5683649 commit cca931b

21 files changed

+750
-646
lines changed

‎.circleci/config.yml

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,7 @@ commands:
9595
jobs:
9696
build-debian:
9797
docker:
98-
- image: redisfab/rmbuilder:6.0.1-x64-buster
98+
- image: redisfab/rmbuilder:6.0.5-x64-buster
9999
steps:
100100
- build-steps:
101101
platform: debian
@@ -112,7 +112,7 @@ jobs:
112112

113113
coverage:
114114
docker:
115-
- image: redisfab/rmbuilder:6.0.1-x64-buster
115+
- image: redisfab/rmbuilder:6.0.5-x64-buster
116116
steps:
117117
- checkout
118118
- run:
@@ -136,7 +136,7 @@ jobs:
136136
- run:
137137
name: Test with coverage
138138
command: |
139-
make -C opt test SHOW=1 COV=1
139+
make -C opt test SHOW=1 COV=1 CLUSTER=1
140140
make -C opt cov-upload
141141
no_output_timeout: 20m
142142

@@ -209,7 +209,7 @@ jobs:
209209
package:
210210
type: string
211211
docker:
212-
- image: redisfab/rmbuilder:6.0.1-x64-buster
212+
- image: redisfab/rmbuilder:6.0.5-x64-buster
213213
steps:
214214
- attach_workspace:
215215
at: workspace

‎opt/Makefile

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,12 +189,14 @@ endif
189189
export GEN ?= 1
190190
export SLAVES ?= 1
191191
export AOF ?= 1
192+
export CLUSTER ?= 1
192193

193194
test:
194195
$(COVERAGE_RESET)
195196
$(SHOW)\
196197
DEVICE=$(DEVICE) \
197198
MODULE=$(INSTALLED_TARGET) \
199+
CLUSTER=$(CLUSTER) \
198200
GEN=$(GEN) AOF=$(AOF) SLAVES=$(SLAVES) \
199201
VALGRIND=$(VALGRIND) \
200202
$(ROOT)/test/tests.sh

‎opt/readies

‎src/dag.c

Lines changed: 27 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -306,7 +306,6 @@ int RAI_parseDAGLoadArgs(RedisModuleCtx *ctx, RedisModuleString **argv,
306306
ctx, "ERR invalid or negative value found in number of keys to LOAD");
307307
return -1;
308308
}
309-
310309
int number_loaded_keys = 0;
311310
int separator_flag = 0;
312311
size_t argpos = 2;
@@ -318,8 +317,7 @@ int RAI_parseDAGLoadArgs(RedisModuleCtx *ctx, RedisModuleString **argv,
318317
} else {
319318
RAI_Tensor *t;
320319
RedisModuleKey *key;
321-
const int status = RAI_GetTensorFromKeyspace(ctx, argv[argpos], &key, &t,
322-
REDISMODULE_READ);
320+
const int status = RAI_GetTensorFromKeyspace(ctx, argv[argpos], &key, &t, REDISMODULE_READ);
323321
if (status == REDISMODULE_ERR) {
324322
RedisModule_Log(
325323
ctx, "warning",
@@ -383,10 +381,35 @@ int RAI_parseDAGPersistArgs(RedisModuleCtx *ctx, RedisModuleString **argv,
383381
return argpos;
384382
}
385383

384+
int RedisAI_DagRun_IsKeysPositionRequest_ReportKeys(RedisModuleCtx *ctx,
385+
RedisModuleString **argv, int argc){
386+
for (size_t argpos = 1; argpos < argc; argpos++){
387+
const char *arg_string = RedisModule_StringPtrLen(argv[argpos], NULL);
388+
if ( (!strcasecmp(arg_string, "LOAD") || !strcasecmp(arg_string, "PERSIST") ) && (argpos+1 < argc) ) {
389+
long long n_keys;
390+
argpos++;
391+
const int retval = RedisModule_StringToLongLong(argv[argpos], &n_keys);
392+
if(retval != REDISMODULE_OK){
393+
return REDISMODULE_ERR;
394+
}
395+
argpos++;
396+
if (n_keys > 0){
397+
size_t last_persist_argpos = n_keys+argpos;
398+
for (; argpos < last_persist_argpos && argpos < argc; argpos++){
399+
RedisModule_KeyAtPos(ctx, argpos);
400+
}
401+
}
402+
}
403+
}
404+
return REDISMODULE_OK;
405+
}
406+
386407
int RedisAI_DagRunSyntaxParser(RedisModuleCtx *ctx, RedisModuleString **argv,
387408
int argc, int dagMode) {
409+
if (RedisModule_IsKeysPositionRequest(ctx)) {
410+
return RedisAI_DagRun_IsKeysPositionRequest_ReportKeys(ctx, argv, argc);
411+
}
388412
if (argc < 4) return RedisModule_WrongArity(ctx);
389-
390413
RedisAI_RunInfo *rinfo = NULL;
391414
if (RAI_InitRunInfo(&rinfo) == REDISMODULE_ERR) {
392415
return RedisModule_ReplyWithError(

‎src/dag.h

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,21 @@ int RAI_parseDAGPersistArgs(RedisModuleCtx *ctx, RedisModuleString **argv,
7373
int argc, AI_dict **localContextDict,
7474
const char *chaining_operator);
7575

76+
/**
77+
* When a module command is called in order to obtain the position of
78+
* keys, since it was flagged as "getkeys-api" during the registration,
79+
* the command implementation checks for this special call using the
80+
* RedisModule_IsKeysPositionRequest() API and uses this function in
81+
* order to report keys.
82+
* No real execution is done on this special call.
83+
* @param ctx Context in which Redis modules operate
84+
* @param argv Redis command arguments, as an array of strings
85+
* @param argc Redis command number of arguments
86+
* @return
87+
*/
88+
int RedisAI_DagRun_IsKeysPositionRequest_ReportKeys(RedisModuleCtx *ctx,
89+
RedisModuleString **argv, int argc);
90+
7691
/**
7792
* DAGRUN and DAGRUN_RO parser, which reads the the sequence of
7893
* arguments and decides whether the sequence conforms to the syntax

‎src/model.c

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -548,6 +548,18 @@ int RAI_ModelSerialize(RAI_Model *model, char **buffer, size_t *len, RAI_Error *
548548
return ret;
549549
}
550550

551+
int RedisAI_ModelRun_IsKeysPositionRequest_ReportKeys(RedisModuleCtx *ctx,
552+
RedisModuleString **argv, int argc){
553+
RedisModule_KeyAtPos(ctx, 1);
554+
for (size_t argpos = 3; argpos < argc; argpos++){
555+
const char *str = RedisModule_StringPtrLen(argv[argpos], NULL);
556+
if (!strcasecmp(str, "OUTPUTS")) {
557+
continue;
558+
}
559+
RedisModule_KeyAtPos(ctx,argpos);
560+
}
561+
return REDISMODULE_OK;
562+
}
551563

552564
int RedisAI_Parse_ModelRun_RedisCommand(RedisModuleCtx *ctx,
553565
RedisModuleString **argv, int argc,
@@ -651,7 +663,6 @@ int RedisAI_Parse_ModelRun_RedisCommand(RedisModuleCtx *ctx,
651663
}
652664
}
653665
}
654-
655666
if ((*mto)->inputs && array_len((*mto)->inputs) != ninputs) {
656667
if (ctx == NULL) {
657668
RAI_SetError(

‎src/model.h

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -201,6 +201,21 @@ int RAI_ModelSerialize(RAI_Model* model, char** buffer, size_t* len,
201201
int RAI_GetModelFromKeyspace(RedisModuleCtx* ctx, RedisModuleString* keyName,
202202
RedisModuleKey** key, RAI_Model** model, int mode);
203203

204+
/**
205+
* When a module command is called in order to obtain the position of
206+
* keys, since it was flagged as "getkeys-api" during the registration,
207+
* the command implementation checks for this special call using the
208+
* RedisModule_IsKeysPositionRequest() API and uses this function in
209+
* order to report keys.
210+
* No real execution is done on this special call.
211+
* @param ctx Context in which Redis modules operate
212+
* @param argv Redis command arguments, as an array of strings
213+
* @param argc Redis command number of arguments
214+
* @return
215+
*/
216+
int RedisAI_ModelRun_IsKeysPositionRequest_ReportKeys(RedisModuleCtx *ctx,
217+
RedisModuleString **argv, int argc);
218+
204219
/**
205220
* Helper method to parse AI.MODELRUN arguments
206221
*

‎src/redisai.c

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -535,20 +535,20 @@ int RedisAI_ModelScan_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv
535535
*/
536536
int RedisAI_ModelRun_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv,
537537
int argc) {
538+
if (RedisModule_IsKeysPositionRequest(ctx)) {
539+
return RedisAI_ModelRun_IsKeysPositionRequest_ReportKeys(ctx, argv, argc);
540+
}
538541
if (argc < 3) return RedisModule_WrongArity(ctx);
539-
540542
RedisAI_RunInfo *rinfo = NULL;
541543
if (RAI_InitRunInfo(&rinfo) == REDISMODULE_ERR) {
542544
return RedisModule_ReplyWithError(ctx, "ERR Unable to allocate the memory and initialise the RedisAI_RunInfo structure");
543545
}
544-
545546
RAI_Model *mto;
546547
RedisModuleKey *modelKey;
547548
const int status = RAI_GetModelFromKeyspace(ctx, argv[1], &modelKey, &mto, REDISMODULE_READ);
548549
if(status==REDISMODULE_ERR){
549550
return REDISMODULE_ERR;
550551
}
551-
552552
RedisModule_RetainString(NULL, argv[1]);
553553
rinfo->runkey = argv[1];
554554
rinfo->mctx = RAI_ModelRunCtxCreate(mto);
@@ -582,16 +582,17 @@ int RedisAI_ModelRun_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv,
582582
* AI.SCRIPTRUN <key> <function> INPUTS <input> [input ...] OUTPUTS <output> [output ...]
583583
*/
584584
int RedisAI_ScriptRun_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
585+
if (RedisModule_IsKeysPositionRequest(ctx)) {
586+
return RedisAI_ScriptRun_IsKeysPositionRequest_ReportKeys(ctx, argv, argc);
587+
}
585588
if (argc < 6) return RedisModule_WrongArity(ctx);
586-
587589
RedisAI_RunInfo *rinfo = NULL;
588590
if (RAI_InitRunInfo(&rinfo) == REDISMODULE_ERR) {
589591
return RedisModule_ReplyWithError(
590592
ctx,
591593
"ERR Unable to allocate the memory and initialise the RedisAI_RunInfo"
592594
"structure");
593595
}
594-
595596
RedisModule_RetainString(NULL, argv[1]);
596597
rinfo->runkey = argv[1];
597598

@@ -1139,11 +1140,11 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc)
11391140
== REDISMODULE_ERR)
11401141
return REDISMODULE_ERR;
11411142

1142-
if (RedisModule_CreateCommand(ctx, "ai.dagrun", RedisAI_DagRun_RedisCommand, "write deny-oom", 3, 3, 1)
1143+
if (RedisModule_CreateCommand(ctx, "ai.dagrun", RedisAI_DagRun_RedisCommand, "write deny-oom getkeys-api", 3, 3, 1)
11431144
== REDISMODULE_ERR)
11441145
return REDISMODULE_ERR;
11451146

1146-
if (RedisModule_CreateCommand(ctx, "ai.dagrun_ro", RedisAI_DagRunRO_RedisCommand, "readonly", 3, 3, 1)
1147+
if (RedisModule_CreateCommand(ctx, "ai.dagrun_ro", RedisAI_DagRunRO_RedisCommand, "readonly getkeys-api", 3, 3, 1)
11471148
== REDISMODULE_ERR)
11481149
return REDISMODULE_ERR;
11491150

‎src/run_info.c

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,9 @@ int RAI_InitRunInfo(RedisAI_RunInfo **result) {
136136
void RAI_FreeDagOp(RedisModuleCtx *ctx, RAI_DagOp *dagOp) {
137137
if (dagOp) {
138138
RAI_FreeError(dagOp->err);
139+
if(dagOp->runkey){
140+
RedisModule_FreeString(ctx,dagOp->runkey);
141+
}
139142
if (dagOp->argv) {
140143
for (size_t i = 0; i < array_len(dagOp->argv); i++) {
141144
RedisModule_FreeString(ctx, dagOp->argv[i]);
@@ -205,9 +208,9 @@ void RAI_FreeRunInfo(RedisModuleCtx *ctx, struct RedisAI_RunInfo *rinfo) {
205208
}
206209
AI_dictReleaseIterator(iter);
207210

208-
RedisModule_Free(rinfo->dagTensorsContext);
209-
RedisModule_Free(rinfo->dagTensorsLoadedContext);
210-
RedisModule_Free(rinfo->dagTensorsPersistentContext);
211+
AI_dictRelease(rinfo->dagTensorsContext);
212+
AI_dictRelease(rinfo->dagTensorsLoadedContext);
213+
AI_dictRelease(rinfo->dagTensorsPersistentContext);
211214
}
212215

213216
if (rinfo->dagOps) {

‎src/script.c

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -267,6 +267,19 @@ int RAI_GetScriptFromKeyspace(RedisModuleCtx* ctx, RedisModuleString* keyName,
267267
return REDISMODULE_OK;
268268
}
269269

270+
int RedisAI_ScriptRun_IsKeysPositionRequest_ReportKeys(RedisModuleCtx *ctx,
271+
RedisModuleString **argv, int argc){
272+
RedisModule_KeyAtPos(ctx, 1);
273+
for (size_t argpos = 4; argpos < argc; argpos++){
274+
const char *str = RedisModule_StringPtrLen(argv[argpos], NULL);
275+
if (!strcasecmp(str, "OUTPUTS")) {
276+
continue;
277+
}
278+
RedisModule_KeyAtPos(ctx,argpos);
279+
}
280+
return REDISMODULE_OK;
281+
}
282+
270283
/**
271284
* AI.SCRIPTRUN <key> <function> INPUTS <input> [input ...] OUTPUTS <output> [output ...]
272285
*/

0 commit comments

Comments
 (0)