diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 2686bfdad..b40ca36b0 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -29,6 +29,8 @@ ADD_LIBRARY(redisai_obj OBJECT config.c DAG/dag.c DAG/dag_parser.c + DAG/dag_builder.c + DAG/dag_execute.c modelRun_ctx.c backends.c backends/util.c diff --git a/src/DAG/dag.c b/src/DAG/dag.c index 6dbdab546..d8ef2e0d7 100644 --- a/src/DAG/dag.c +++ b/src/DAG/dag.c @@ -57,7 +57,7 @@ static void Dag_LoadInputsToModelRunCtx(RedisAI_RunInfo *rinfo, RAI_DagOp *curre for (uint i = 0; i < n_inkeys; i++) { RAI_Tensor *inputTensor; const int get_result = RAI_getTensorFromLocalContext( - NULL, rinfo->dagTensorsContext, currentOp->inkeys[i], &inputTensor, currentOp->err); + rinfo->dagTensorsContext, currentOp->inkeys[i], &inputTensor, currentOp->err); if (get_result == REDISMODULE_ERR) { // We check for this outside the function // this check cannot be covered by tests @@ -101,6 +101,113 @@ static void Dag_StoreOutputsFromModelRunCtx(RedisAI_RunInfo *rinfo, RAI_DagOp *c RAI_ContextUnlock(rinfo); } +static int _StoreTensorInKeySpace(RedisModuleCtx *ctx, RAI_Tensor *tensor, + RedisModuleString *persist_key_name, bool mangled_name) { + + int ret = REDISMODULE_ERR; + RedisModuleKey *key; + size_t persist_key_len; + const char *persist_key_str = RedisModule_StringPtrLen(persist_key_name, &persist_key_len); + + RedisModuleString *demangled_key_name; + if (mangled_name) { + demangled_key_name = RedisModule_CreateString(NULL, persist_key_str, persist_key_len - 4); + } else { + demangled_key_name = RedisModule_CreateString(NULL, persist_key_str, persist_key_len); + } + + const int status = + RAI_OpenKey_Tensor(ctx, demangled_key_name, &key, REDISMODULE_READ | REDISMODULE_WRITE); + if (status == REDISMODULE_ERR) { + RedisModule_ReplyWithError(ctx, "ERR could not save tensor"); + goto clean_up; + } + if (RedisModule_ModuleTypeSetValue(key, RedisAI_TensorType, RAI_TensorGetShallowCopy(tensor)) != + REDISMODULE_OK) { + RedisModule_ReplyWithError(ctx, "ERR could not save tensor"); + RedisModule_CloseKey(key); + goto clean_up; + } + // Only if we got until here, tensor is saved in keyspace. + RedisAI_ReplicateTensorSet(ctx, demangled_key_name, tensor); + RedisModule_CloseKey(key); + ret = REDISMODULE_OK; + +clean_up: + RedisModule_FreeString(NULL, demangled_key_name); + return ret; +} + +static void _DAG_PersistTensors(RedisModuleCtx *ctx, RedisAI_RunInfo *rinfo) { + + AI_dictIterator *persist_iter = AI_dictGetSafeIterator(rinfo->dagTensorsPersistedContext); + AI_dictEntry *persist_entry = AI_dictNext(persist_iter); + + while (persist_entry) { + RedisModuleString *persist_key_name = AI_dictGetKey(persist_entry); + AI_dictEntry *tensor_entry = AI_dictFind(rinfo->dagTensorsContext, persist_key_name); + RedisModule_Assert(tensor_entry); + RAI_Tensor *tensor = AI_dictGetVal(tensor_entry); + if (tensor == NULL) { + persist_entry = AI_dictNext(persist_iter); + continue; + } + if (_StoreTensorInKeySpace(ctx, tensor, persist_key_name, true) == REDISMODULE_ERR) { + *rinfo->dagError = 1; + RedisModule_Log(ctx, "warning", + "Could not persist tensor under the key (%s) after executing DAGRUN " + "command, persist stopped", + RedisModule_StringPtrLen(persist_key_name, NULL)); + AI_dictReleaseIterator(persist_iter); + return; + } + persist_entry = AI_dictNext(persist_iter); + } + AI_dictReleaseIterator(persist_iter); +} + +static void _ModelSingleOp_PersistTensors(RedisModuleCtx *ctx, RAI_DagOp *op) { + + const size_t noutputs = RAI_ModelRunCtxNumOutputs(op->mctx); + for (size_t outputNumber = 0; outputNumber < noutputs; outputNumber++) { + RedisModuleString *persist_key_name = op->outkeys[outputNumber]; + RAI_Tensor *tensor = RAI_ModelRunCtxOutputTensor(op->mctx, outputNumber); + tensor = tensor ? RAI_TensorGetShallowCopy(tensor) : NULL; + if (!tensor) + continue; + + if (_StoreTensorInKeySpace(ctx, tensor, persist_key_name, false) == REDISMODULE_ERR) { + RedisModule_Log(ctx, "warning", + "Could not persist tensor under the key (%s) after executing DAGRUN " + "command, persist stopped", + RedisModule_StringPtrLen(persist_key_name, NULL)); + op->result = REDISMODULE_ERR; + return; + } + } +} + +static void _ScriptSingleOp_PersistTensors(RedisModuleCtx *ctx, RAI_DagOp *op) { + + const size_t noutputs = RAI_ScriptRunCtxNumOutputs(op->sctx); + for (size_t outputNumber = 0; outputNumber < noutputs; outputNumber++) { + RedisModuleString *persist_key_name = op->outkeys[outputNumber]; + RAI_Tensor *tensor = RAI_ScriptRunCtxOutputTensor(op->sctx, outputNumber); + tensor = tensor ? RAI_TensorGetShallowCopy(tensor) : NULL; + if (!tensor) + continue; + + if (_StoreTensorInKeySpace(ctx, tensor, persist_key_name, false) == REDISMODULE_ERR) { + RedisModule_Log(ctx, "warning", + "Could not persist tensor under the key (%s) after executing DAGRUN " + "command, persist stopped", + RedisModule_StringPtrLen(persist_key_name, NULL)); + op->result = REDISMODULE_ERR; + return; + } + } +} + /** * Execution of a MODELRUN DAG step. * If an error occurs, it is recorded in the DagOp struct. @@ -201,7 +308,7 @@ void RedisAI_DagRunSession_ScriptRun_Step(RedisAI_RunInfo *rinfo, RAI_DagOp *cur for (uint i = 0; i < n_inkeys; i++) { RAI_Tensor *inputTensor; const int get_result = RAI_getTensorFromLocalContext( - NULL, rinfo->dagTensorsContext, currentOp->inkeys[i], &inputTensor, currentOp->err); + rinfo->dagTensorsContext, currentOp->inkeys[i], &inputTensor, currentOp->err); if (get_result == REDISMODULE_ERR) { // We check for this outside the function // this check cannot be covered by tests @@ -258,8 +365,7 @@ size_t RAI_DagOpBatchSize(RAI_DagOp *op, RedisAI_RunInfo *rinfo) { if (rinfo->single_op_dag) { input = op->mctx->inputs[i].tensor; } else { - RAI_getTensorFromLocalContext(NULL, rinfo->dagTensorsContext, op->inkeys[i], &input, - op->err); + RAI_getTensorFromLocalContext(rinfo->dagTensorsContext, op->inkeys[i], &input, op->err); } // We are expecting input != NULL, because we only reach this function if all inputs // are available in context for the current dagOp. We could be more defensive @@ -307,14 +413,14 @@ int RAI_DagOpBatchable(RAI_DagOp *op1, RedisAI_RunInfo *rinfo1, RAI_DagOp *op2, if (rinfo1->single_op_dag == 1) { input1 = op1->mctx->inputs[i].tensor; } else { - RAI_getTensorFromLocalContext(NULL, rinfo1->dagTensorsContext, op1->inkeys[i], &input1, + RAI_getTensorFromLocalContext(rinfo1->dagTensorsContext, op1->inkeys[i], &input1, op1->err); } RAI_Tensor *input2; if (rinfo2->single_op_dag == 1) { input2 = op2->mctx->inputs[i].tensor; } else { - RAI_getTensorFromLocalContext(NULL, rinfo2->dagTensorsContext, op2->inkeys[i], &input2, + RAI_getTensorFromLocalContext(rinfo2->dagTensorsContext, op2->inkeys[i], &input2, op2->err); } if (input1 == NULL || input2 == NULL) { @@ -464,7 +570,11 @@ void RedisAI_DagRunSessionStep(RedisAI_RunInfo *rinfo, const char *devicestr) { } if (currentOp->result != REDISMODULE_OK) { - __atomic_store_n(rinfo->dagError, 1, __ATOMIC_RELAXED); + // If this is the first op with error, save the error in the DAG runInfo. + if (__sync_val_compare_and_swap(rinfo->dagError, 0, 1) == 0) { + RAI_SetError(rinfo->err, RAI_GetErrorCode(currentOp->err), + RAI_GetError(currentOp->err)); + } } } @@ -472,16 +582,12 @@ void RedisAI_BatchedDagRunSessionStep(RedisAI_RunInfo **batched_rinfo, const cha // Assumption: ops are guaranteed to be all MODELRUN int n_ops = array_len(batched_rinfo); - assert(n_ops > 1); - RAI_DagOp *currentOps[n_ops]; for (int i = 0; i < n_ops; i++) { RedisAI_RunInfo *rinfo = batched_rinfo[i]; - RAI_DagOp *currentOp = RedisAI_DagCurrentOp(rinfo); - currentOps[i] = currentOp; } @@ -492,109 +598,12 @@ void RedisAI_BatchedDagRunSessionStep(RedisAI_RunInfo **batched_rinfo, const cha RAI_DagOp *currentOp = currentOps[i]; if (currentOp->result != REDISMODULE_OK) { - __atomic_store_n(rinfo->dagError, 1, __ATOMIC_RELAXED); - } - } - return; -} - -static int _StoreTensorInKeySpace(RedisModuleCtx *ctx, RAI_Tensor *tensor, - RedisModuleString *persist_key_name, bool mangled_name) { - - int ret = REDISMODULE_ERR; - RedisModuleKey *key; - size_t persist_key_len; - const char *persist_key_str = RedisModule_StringPtrLen(persist_key_name, &persist_key_len); - - RedisModuleString *demangled_key_name; - if (mangled_name) { - demangled_key_name = RedisModule_CreateString(NULL, persist_key_str, persist_key_len - 4); - } else { - demangled_key_name = RedisModule_CreateString(NULL, persist_key_str, persist_key_len); - } - - const int status = - RAI_OpenKey_Tensor(ctx, demangled_key_name, &key, REDISMODULE_READ | REDISMODULE_WRITE); - if (status == REDISMODULE_ERR) { - RedisModule_ReplyWithError(ctx, "ERR could not save tensor"); - goto clean_up; - } else { - if (RedisModule_ModuleTypeSetValue(key, RedisAI_TensorType, - RAI_TensorGetShallowCopy(tensor)) != REDISMODULE_OK) { - RedisModule_ReplyWithError(ctx, "ERR could not save tensor"); - goto clean_up; - } - } - ret = REDISMODULE_OK; - -clean_up: - RedisModule_CloseKey(key); - RedisAI_ReplicateTensorSet(ctx, demangled_key_name, tensor); - RedisModule_FreeString(NULL, demangled_key_name); - return ret; -} - -static void _PersistTensors(RedisModuleCtx *ctx, RedisAI_RunInfo *rinfo) { - - AI_dictIterator *persist_iter = AI_dictGetSafeIterator(rinfo->dagTensorsPersistedContext); - AI_dictEntry *persist_entry = AI_dictNext(persist_iter); - - while (persist_entry) { - RedisModuleString *persist_key_name = AI_dictGetKey(persist_entry); - AI_dictEntry *tensor_entry = AI_dictFind(rinfo->dagTensorsContext, persist_key_name); - if (tensor_entry) { - RAI_Tensor *tensor = AI_dictGetVal(tensor_entry); - if (tensor == NULL) { - persist_entry = AI_dictNext(persist_iter); - continue; - } - if (_StoreTensorInKeySpace(ctx, tensor, persist_key_name, true) == REDISMODULE_ERR) - rinfo->dagReplyLength++; - - } else { - RedisModule_ReplyWithError(ctx, - "ERR specified persistent key that was not used in DAG"); - rinfo->dagReplyLength++; - RedisModule_Log(ctx, "warning", - "on DAGRUN's PERSIST specified persistent key (%s) that " - "was not used on DAG. Logging all local context keys", - 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)", - RedisModule_StringPtrLen(localcontext_key_name, NULL)); - local_entry = AI_dictNext(local_iter); - } - AI_dictReleaseIterator(local_iter); - - for (size_t opN = 0; opN < array_len(rinfo->dagOps); opN++) { - RedisModule_Log(ctx, "warning", "DAG's op n# %zu - cmdType %d ( argc %d )", opN, - rinfo->dagOps[opN]->commandType, rinfo->dagOps[opN]->argc); + // If this is the first op with error, save the error in the DAG runInfo. + if (__sync_val_compare_and_swap(rinfo->dagError, 0, 1) == 0) { + RAI_SetError(rinfo->err, RAI_GetErrorCode(currentOp->err), + RAI_GetError(currentOp->err)); } } - persist_entry = AI_dictNext(persist_iter); - } - AI_dictReleaseIterator(persist_iter); -} - -static void _ModelSingleOp_PersistTensors(RedisModuleCtx *ctx, RAI_DagOp *op) { - const size_t noutputs = RAI_ModelRunCtxNumOutputs(op->mctx); - for (size_t outputNumber = 0; outputNumber < noutputs; outputNumber++) { - RAI_Tensor *tensor = RAI_ModelRunCtxOutputTensor(op->mctx, outputNumber); - if (tensor) - _StoreTensorInKeySpace(ctx, tensor, op->outkeys[outputNumber], false); - } -} - -static void _ScriptSingleOp_PersistTensors(RedisModuleCtx *ctx, RAI_DagOp *op) { - const size_t noutputs = RAI_ScriptRunCtxNumOutputs(op->sctx); - for (size_t outputNumber = 0; outputNumber < noutputs; outputNumber++) { - RAI_Tensor *tensor = RAI_ScriptRunCtxOutputTensor(op->sctx, outputNumber); - if (tensor) - _StoreTensorInKeySpace(ctx, tensor, op->outkeys[outputNumber], false); } } @@ -603,20 +612,18 @@ int RedisAI_DagRun_Reply(RedisModuleCtx *ctx, RedisModuleString **argv, int argc REDISMODULE_NOT_USED(argc); RedisAI_RunInfo *rinfo = RedisModule_GetBlockedClientPrivateData(ctx); + if (*rinfo->timedOut) { + RedisModule_ReplyWithSimpleString(ctx, "TIMEDOUT"); + return REDISMODULE_OK; + } + if (RAI_GetErrorCode(rinfo->err) == RAI_EDAGRUN) { RedisModule_ReplyWithError(ctx, RAI_GetErrorOneLine(rinfo->err)); - return REDISMODULE_ERR; + return REDISMODULE_OK; } int dag_error = 0; - char *detail_oneline; - size_t n_dagOps = array_len(rinfo->dagOps); - if (*rinfo->timedOut) { - RedisModule_ReplyWithSimpleString(ctx, "TIMEDOUT"); - return REDISMODULE_OK; - } - if (!rinfo->single_op_dag) { RedisModule_ReplyWithArray(ctx, REDISMODULE_POSTPONED_ARRAY_LEN); } @@ -640,8 +647,8 @@ int RedisAI_DagRun_Reply(RedisModuleCtx *ctx, RedisModuleString **argv, int argc case REDISAI_DAG_CMD_TENSORGET: { rinfo->dagReplyLength++; RAI_Tensor *t; - int res = RAI_getTensorFromLocalContext(NULL, rinfo->dagTensorsContext, - currentOp->inkeys[0], &t, currentOp->err); + int res = RAI_getTensorFromLocalContext(rinfo->dagTensorsContext, currentOp->inkeys[0], + &t, currentOp->err); if (res != REDISMODULE_OK) { RedisModule_ReplyWithError(ctx, currentOp->err->detail_oneline); dag_error = 1; @@ -701,16 +708,10 @@ int RedisAI_DagRun_Reply(RedisModuleCtx *ctx, RedisModuleString **argv, int argc } if (dag_error) { - if (rinfo->single_op_dag == 0) { - RedisModule_ReplySetArrayLength(ctx, rinfo->dagReplyLength); - } - return REDISMODULE_ERR; + goto cleanup; } - if (!rinfo->single_op_dag) { - // Save the required tensors in redis key space. - _PersistTensors(ctx, rinfo); - RedisModule_ReplySetArrayLength(ctx, rinfo->dagReplyLength); + _DAG_PersistTensors(ctx, rinfo); } else { if (rinfo->dagOps[0]->commandType == REDISAI_DAG_CMD_MODELRUN) { _ModelSingleOp_PersistTensors(ctx, rinfo->dagOps[0]); @@ -720,6 +721,10 @@ int RedisAI_DagRun_Reply(RedisModuleCtx *ctx, RedisModuleString **argv, int argc } } +cleanup: + if (!rinfo->single_op_dag) { + RedisModule_ReplySetArrayLength(ctx, rinfo->dagReplyLength); + } return REDISMODULE_OK; } @@ -749,87 +754,10 @@ int RedisAI_DagRun_IsKeysPositionRequest_ReportKeys(RedisModuleCtx *ctx, RedisMo void RunInfo_FreeData(RedisModuleCtx *ctx, void *rinfo) { RAI_FreeRunInfo(rinfo); } -// Add Shallow copies of the DAG run info to the devices' queues. -// Return REDISMODULE_OK in case of success, REDISMODULE_ERR if (at least) one insert op had -// failed. -int DAG_InsertDAGToQueue(RedisAI_RunInfo *rinfo) { - const char **devices = array_new(const char *, 10); - - for (long long i = 0; i < array_len(rinfo->dagOps); i++) { - const char *devicestr = rinfo->dagOps[i]->devicestr; - bool found = false; - for (long long j = 0; j < array_len(devices); j++) { - if (strcasecmp(devicestr, devices[j]) == 0) { - found = true; - break; - } - } - if (!found) { - devices = array_append(devices, devicestr); - } - } - - size_t ndevices = array_len(devices); - if (ndevices == 1) - rinfo->single_device_dag = 1; - RedisAI_RunInfo **rinfo_copies = array_new(RedisAI_RunInfo *, ndevices); - - for (long long i = 0; i < ndevices; i++) { - RedisAI_RunInfo *rinfo_copy; - RAI_ShallowCopyDagRunInfo(&rinfo_copy, rinfo); - rinfo_copies = array_append(rinfo_copies, rinfo_copy); - } - - for (long long i = 0; i < ndevices; i++) { - RedisAI_RunInfo *rinfo_copy = rinfo_copies[i]; - for (long long j = 0; j < rinfo_copy->dagOpCount; j++) { - if (strcasecmp(rinfo_copy->dagOps[j]->devicestr, devices[i]) == 0) { - rinfo_copy->dagDeviceOps = - array_append(rinfo_copy->dagDeviceOps, rinfo_copy->dagOps[j]); - } - } - rinfo_copy->dagDeviceOpCount = array_len(rinfo_copy->dagDeviceOps); - } - - RunQueueInfo **run_queues_info = array_new(RunQueueInfo *, ndevices); - for (long long i = 0; i < ndevices; i++) { - const char *devicestr = devices[i]; - RunQueueInfo *run_queue_info = NULL; - if (ensureRunQueue(devicestr, &run_queue_info) == REDISMODULE_ERR) { - // A device run queue was not created properly, so we free everything, - // set an error and finish. - array_free(devices); - for (int j = 0; j < ndevices; j++) { - RAI_DagRunInfoFreeShallowCopy(rinfo_copies[j]); - } - array_free(rinfo_copies); - array_free(run_queues_info); - RAI_SetError(rinfo->err, RAI_EDAGRUN, "ERR Queue not initialized for device"); - rinfo->OnFinish((RedisAI_OnFinishCtx *)rinfo, rinfo->private_data); - return REDISMODULE_ERR; - } - run_queues_info = array_append(run_queues_info, run_queue_info); - } - for (long long i = 0; i < ndevices; i++) { - RedisAI_RunInfo *rinfo_copy = rinfo_copies[i]; - RunQueueInfo *run_queue_info = run_queues_info[i]; - gettimeofday(&rinfo_copy->queuingTime, NULL); - - pthread_mutex_lock(&run_queue_info->run_queue_mutex); - queuePush(run_queue_info->run_queue, rinfo_copy); - pthread_cond_signal(&run_queue_info->queue_condition_var); - pthread_mutex_unlock(&run_queue_info->run_queue_mutex); - } - - array_free(devices); - array_free(rinfo_copies); - array_free(run_queues_info); - return REDISMODULE_OK; -} - void DAG_ReplyAndUnblock(RedisAI_OnFinishCtx *ctx, void *private_data) { RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)ctx; - if (rinfo->client) + if (rinfo->client) { RedisModule_UnblockClient(rinfo->client, rinfo); + } } diff --git a/src/DAG/dag_builder.c b/src/DAG/dag_builder.c new file mode 100644 index 000000000..83e0a0a07 --- /dev/null +++ b/src/DAG/dag_builder.c @@ -0,0 +1,188 @@ +#include "dag_builder.h" +#include "run_info.h" +#include "dag_parser.h" +#include "string_utils.h" +#include "modelRun_ctx.h" + +// Store the given arguments from the string in argv array and their amount in argc. +int _StringToRMArray(const char *dag, RedisModuleString ***argv, int *argc, RAI_Error *err) { + + char dag_string[strlen(dag) + 1]; + strcpy(dag_string, dag); + + char *token = strtok(dag_string, " "); + if (strcmp(token, "|>") != 0) { + RAI_SetError(err, RAI_EDAGBUILDER, "DAG op should start with: '|>' "); + return REDISMODULE_ERR; + } + + while (token != NULL) { + RedisModuleString *RS_token = RedisModule_CreateString(NULL, token, strlen(token)); + *argv = array_append(*argv, RS_token); + (*argc)++; + token = strtok(NULL, " "); + } + return REDISMODULE_OK; +} + +int RAI_DAGLoadTensor(RAI_DAGRunCtx *run_info, const char *t_name, RAI_Tensor *tensor) { + + RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)run_info; + RedisModuleString *key_name = RedisModule_CreateString(NULL, t_name, strlen(t_name)); + // Add the tensor under its "mangled" key name to the DAG local context dict. + char buf[16]; + sprintf(buf, "%04d", 1); + RedisModule_StringAppendBuffer(NULL, key_name, buf, strlen(buf)); + AI_dictAdd(rinfo->dagTensorsContext, (void *)key_name, + (void *)RAI_TensorGetShallowCopy(tensor)); + RedisModule_FreeString(NULL, key_name); + + return REDISMODULE_OK; +} + +RAI_DAGRunCtx *RAI_DAGRunCtxCreate(void) { + RedisAI_RunInfo *rinfo; + RAI_InitRunInfo(&rinfo); + return (RAI_DAGRunCtx *)rinfo; +} + +RAI_DAGRunOp *RAI_DAGCreateModelRunOp(RAI_Model *model) { + RAI_ModelRunCtx *mctx = RAI_ModelRunCtxCreate(model); + RAI_DagOp *op; + RAI_InitDagOp(&op); + + op->commandType = REDISAI_DAG_CMD_MODELRUN; + op->mctx = mctx; + op->devicestr = model->devicestr; + op->runkey = RAI_HoldString(NULL, (RedisModuleString *)model->infokey); + return (RAI_DAGRunOp *)op; +} + +RAI_DAGRunOp *RAI_DAGCreateScriptRunOp(RAI_Script *script, const char *func_name) { + RAI_ScriptRunCtx *sctx = RAI_ScriptRunCtxCreate(script, func_name); + RAI_DagOp *op; + RAI_InitDagOp(&op); + + op->commandType = REDISAI_DAG_CMD_SCRIPTRUN; + op->sctx = sctx; + op->devicestr = script->devicestr; + op->runkey = RAI_HoldString(NULL, (RedisModuleString *)script->infokey); + return (RAI_DAGRunOp *)op; +} + +int RAI_DAGRunOpAddInput(RAI_DAGRunOp *DAGOp, const char *input) { + RAI_DagOp *op = (RAI_DagOp *)DAGOp; + RedisModuleString *inkey = RedisModule_CreateString(NULL, input, strlen(input)); + op->inkeys = array_append(op->inkeys, inkey); + return REDISMODULE_OK; +} + +int RAI_DAGRunOpAddOutput(RAI_DAGRunOp *DAGOp, const char *output) { + RAI_DagOp *op = (RAI_DagOp *)DAGOp; + RedisModuleString *outkey = RedisModule_CreateString(NULL, output, strlen(output)); + op->outkeys = array_append(op->outkeys, outkey); + return REDISMODULE_OK; +} + +int RAI_DAGAddRunOp(RAI_DAGRunCtx *run_info, RAI_DAGRunOp *DAGop, RAI_Error *err) { + + RAI_DagOp *op = (RAI_DagOp *)DAGop; + RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)run_info; + if (op->mctx) { + RAI_Model *model = op->mctx->model; + if (ModelGetNumInputs(model) != array_len(op->inkeys)) { + RAI_SetError(err, RAI_EDAGBUILDER, + "Number of keys given as INPUTS does not match model definition"); + return REDISMODULE_ERR; + } + if (ModelGetNumOutputs(model) != array_len(op->outkeys)) { + RAI_SetError(err, RAI_EDAGBUILDER, + "Number of keys given as OUTPUTS does not match model definition"); + return REDISMODULE_ERR; + } + } + rinfo->dagOps = array_append(rinfo->dagOps, op); + + return REDISMODULE_OK; +} + +int RAI_DAGAddTensorGet(RAI_DAGRunCtx *run_info, const char *t_name, RAI_Error *err) { + + RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)run_info; + RAI_DagOp *op; + RAI_InitDagOp(&op); + rinfo->dagOps = array_append(rinfo->dagOps, op); + op->commandType = REDISAI_DAG_CMD_TENSORGET; + op->devicestr = "CPU"; + RedisModuleString *name = RedisModule_CreateString(NULL, t_name, strlen(t_name)); + op->inkeys = array_append(op->inkeys, name); + return REDISMODULE_OK; +} + +int RAI_DAGAddTensorSet(RAI_DAGRunCtx *run_info, const char *t_name, RAI_Tensor *tensor) { + + RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)run_info; + RAI_DagOp *op; + RAI_InitDagOp(&op); + rinfo->dagOps = array_append(rinfo->dagOps, op); + op->commandType = REDISAI_DAG_CMD_TENSORSET; + op->devicestr = "CPU"; + RedisModuleString *name = RedisModule_CreateString(NULL, t_name, strlen(t_name)); + op->outkeys = array_append(op->outkeys, name); + op->outTensor = RAI_TensorGetShallowCopy(tensor); + return REDISMODULE_OK; +} + +int RAI_DAGAddOpsFromString(RAI_DAGRunCtx *run_info, const char *dag, RAI_Error *err) { + + int res = REDISMODULE_ERR; + RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)run_info; + array_new_on_stack(RAI_DagOp *, 10, new_ops); + array_new_on_stack(RedisModuleString *, 100, argv); + int argc = 0; + if (_StringToRMArray(dag, &argv, &argc, err) != REDISMODULE_OK) { + goto cleanup; + } + + RAI_DagOp *op; + for (size_t i = 0; i < argc; i++) { + const char *arg_string = RedisModule_StringPtrLen(argv[i], NULL); + if (strcmp(arg_string, "|>") == 0 && i < argc - 1) { + RAI_InitDagOp(&op); + new_ops = array_append(new_ops, op); + op->argv = &argv[i + 1]; + } else { + op->argc++; + } + } + + if (ParseDAGOps(rinfo, new_ops) != REDISMODULE_OK) { + RAI_SetError(err, RAI_GetErrorCode(rinfo->err), RAI_GetError(rinfo->err)); + goto cleanup; + } + rinfo->dagOpCount = array_len(rinfo->dagOps); + res = REDISMODULE_OK; + +cleanup: + array_free(new_ops); + for (size_t i = 0; i < argc; i++) { + RedisModule_FreeString(NULL, argv[i]); + } + array_free(argv); + return res; +} + +size_t RAI_DAGNumOps(RAI_DAGRunCtx *run_info) { + RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)run_info; + return array_len(rinfo->dagOps); +} + +void RAI_DAGRunOpFree(RAI_DAGRunOp *dagOp) { + RAI_DagOp *op = (RAI_DagOp *)dagOp; + RAI_FreeDagOp(op); +} + +void RAI_DAGFree(RAI_DAGRunCtx *run_info) { + RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)run_info; + RAI_FreeRunInfo(rinfo); +} diff --git a/src/DAG/dag_builder.h b/src/DAG/dag_builder.h new file mode 100644 index 000000000..f03b67fb1 --- /dev/null +++ b/src/DAG/dag_builder.h @@ -0,0 +1,92 @@ +#pragma once + +#include "redisai.h" + +/** + * @brief Create a new empty DAG runInfo object. + */ +RAI_DAGRunCtx *RAI_DAGRunCtxCreate(void); + +/** + * @brief Create a new MODELRUN op for a DAG. + * @param model The model to run. + */ +RAI_DAGRunOp *RAI_DAGCreateModelRunOp(RAI_Model *model); + +/** + * @brief Create a new SCRIPTRUN op for a DAG. + * @param script The script to run. + * @param func_name The specific function to run in the given script. + */ +RAI_DAGRunOp *RAI_DAGCreateScriptRunOp(RAI_Script *script, const char *func_name); + +/** + * @brief Add an input key to a DAG run op (before inserting it to the DAG). + * @param DAGop The DAG run op (MODELRUN / SCRIPTRUN). + * @param input The tensor input name (this name should appear in a previous op of the DAG). + */ +int RAI_DAGRunOpAddInput(RAI_DAGRunOp *DAGOp, const char *input); + +/** + * @brief Add an output key to a DAG run op (before inserting it to the DAG). + * @param DAGop The DAG run op (MODELRUN / SCRIPTRUN). + * @param output The tensor output name (this name may appear in one of the following ops of the + * DAG). + */ +int RAI_DAGRunOpAddOutput(RAI_DAGRunOp *DAGOp, const char *output); + +/** + * @brief Add a run op (MODELRUN/SCRIPTRUN) to a DAG. + * @param runInfo The DAG to insert the op to. + * @param DAGop The DAG run op (MODELRUN / SCRIPTRUN). + * @param err Error is returned in case of a MODELRUN op if the number of inputs and outputs + * given to the op does not match to the number of inputs and outputs in the model definition. + */ +int RAI_DAGAddRunOp(RAI_DAGRunCtx *run_info, RAI_DAGRunOp *DAGop, RAI_Error *err); + +/** + * @brief Load a given tensor to the DAG local context. + * @param runInfo The DAG to load the tensor into. + * @param tname The tensor key. + * @param tensor The tensor to load to the DAG (we load a shallow copy). + */ +int RAI_DAGLoadTensor(RAI_DAGRunCtx *run_info, const char *t_name, RAI_Tensor *tensor); + +/** + * @brief Append a TENSORSET op to a DAG (can use to load an intermediate tensors) + * @param runInfo The DAG to append this op into. + * @param tensor The tensor to set. + */ +int RAI_DAGAddTensorSet(RAI_DAGRunCtx *run_info, const char *t_name, RAI_Tensor *tensor); + +/** + * @brief Append a TENSORGET op to a DAG (can use to output intermediate and final tensors) + * @param runInfo The DAG to append this op into. + * @param tensor The tensor to set. + */ +int RAI_DAGAddTensorGet(RAI_DAGRunCtx *run_info, const char *t_name, RAI_Error *err); + +/** + * @brief Add ops to a DAG from string (according to the command syntax). In case of a valid + * string, the ops are added to the DAG run info, and otherwise all the ops are discarded. + * @param runInfo The DAG to insert the ops into. + * @param dag The string representing the DAG ops to add. + * @param err Error is returned in case of a MODELRUN op if the number of inputs and outputs + * given to the op does not match to the number of inputs and outputs in the model definition. + */ +int RAI_DAGAddOpsFromString(RAI_DAGRunCtx *run_info, const char *dag, RAI_Error *err); + +/** + * @brief Returns the number of ops in a DAG. + */ +size_t RAI_DAGNumOps(RAI_DAGRunCtx *run_info); + +/** + * @brief Free DAG's runInfo and all its internal ops. + */ +void RAI_DAGFree(RAI_DAGRunCtx *run_info); + +/** + * @brief Free a specific DAG run op (MODELRUN/SCRIPTRUN). + */ +void RAI_DAGRunOpFree(RAI_DAGRunOp *dagOp); diff --git a/src/DAG/dag_execute.c b/src/DAG/dag_execute.c new file mode 100644 index 000000000..71072ce25 --- /dev/null +++ b/src/DAG/dag_execute.c @@ -0,0 +1,295 @@ +#include "dag_execute.h" +#include "run_info.h" +#include "background_workers.h" +#include "util/string_utils.h" + +void _DAG_SetTensorsInLocalContext(RedisAI_RunInfo *rinfo) { + for (size_t i = 0; i < rinfo->dagOpCount; i++) { + RAI_DagOp *op = rinfo->dagOps[i]; + if (op->commandType == REDISAI_DAG_CMD_TENSORSET) { + // Insert the tensor with its mangled (unique) name. + void *t = (void *)RAI_TensorGetShallowCopy(op->outTensor); + AI_dictReplace(rinfo->dagTensorsContext, (void *)op->outkeys[0], t); + } + } +} + +int MangleTensorsNames(RedisAI_RunInfo *rinfo) { + + int res = REDISMODULE_ERR; + AI_dict *mangled_tensors = AI_dictCreate(&AI_dictTypeHeapRStrings, NULL); + + { + AI_dictIterator *iter = AI_dictGetSafeIterator(rinfo->dagTensorsContext); + AI_dictEntry *entry = AI_dictNext(iter); + while (entry) { + RedisModuleString *key = (RedisModuleString *)AI_dictGetKey(entry); + size_t key_len; + const char *key_str = RedisModule_StringPtrLen(key, &key_len); + RedisModuleString *demangled_key = RedisModule_CreateString(NULL, key_str, key_len - 4); + int *instance = RedisModule_Alloc(sizeof(int)); + *instance = 1; + AI_dictAdd(mangled_tensors, (void *)demangled_key, (void *)instance); + RedisModule_FreeString(NULL, demangled_key); + entry = AI_dictNext(iter); + } + AI_dictReleaseIterator(iter); + } + + for (long long i = 0; i < array_len(rinfo->dagOps); i++) { + RAI_DagOp *currentOp = rinfo->dagOps[i]; + + RedisModuleString **mangled_inkeys = + array_new(RedisModuleString *, array_len(currentOp->inkeys)); + for (long long j = 0; j < array_len(currentOp->inkeys); j++) { + RedisModuleString *key = currentOp->inkeys[j]; + AI_dictEntry *entry = AI_dictFind(mangled_tensors, key); + if (!entry) { + array_free(mangled_inkeys); + RAI_SetError(rinfo->err, RAI_EDAGRUN, "ERR INPUT key cannot be found in DAG"); + goto cleanup; + } + int *instance = AI_dictGetVal(entry); + char buf[16]; + sprintf(buf, "%04d", *instance); + RedisModuleString *mangled_key = RedisModule_CreateStringFromString(NULL, key); + RedisModule_StringAppendBuffer(NULL, mangled_key, buf, strlen(buf)); + mangled_inkeys = array_append(mangled_inkeys, mangled_key); + } + + RedisModuleString **mangled_outkeys = + array_new(RedisModuleString *, array_len(currentOp->outkeys)); + for (long long j = 0; j < array_len(currentOp->outkeys); j++) { + RedisModuleString *key = currentOp->outkeys[j]; + AI_dictEntry *entry = AI_dictFind(mangled_tensors, key); + int *instance = NULL; + if (entry) { + instance = AI_dictGetVal(entry); + *instance += 1; + } else { + instance = RedisModule_Alloc(sizeof(int)); + *instance = 1; + AI_dictAdd(mangled_tensors, (void *)key, (void *)instance); + } + char buf[16]; + sprintf(buf, "%04d", *instance); + RedisModuleString *mangled_key = RedisModule_CreateStringFromString(NULL, key); + RedisModule_StringAppendBuffer(NULL, mangled_key, buf, strlen(buf)); + mangled_outkeys = array_append(mangled_outkeys, mangled_key); + } + + if (currentOp->inkeys) { + for (size_t j = 0; j < array_len(currentOp->inkeys); j++) { + RedisModule_FreeString(NULL, currentOp->inkeys[j]); + } + array_free(currentOp->inkeys); + } + + if (currentOp->outkeys) { + for (size_t j = 0; j < array_len(currentOp->outkeys); j++) { + RedisModule_FreeString(NULL, currentOp->outkeys[j]); + } + array_free(currentOp->outkeys); + } + + currentOp->inkeys = mangled_inkeys; + currentOp->outkeys = mangled_outkeys; + } + + AI_dict *mangled_persisted = AI_dictCreate(&AI_dictTypeHeapRStrings, NULL); + { + AI_dictIterator *iter = AI_dictGetSafeIterator(rinfo->dagTensorsPersistedContext); + AI_dictEntry *entry = AI_dictNext(iter); + while (entry) { + RedisModuleString *key = (RedisModuleString *)AI_dictGetKey(entry); + AI_dictEntry *mangled_entry = AI_dictFind(mangled_tensors, key); + if (!mangled_entry) { + AI_dictRelease(mangled_persisted); + AI_dictReleaseIterator(iter); + RAI_SetError(rinfo->err, RAI_EDAGRUN, "ERR PERSIST key cannot be found in DAG"); + goto cleanup; + } + if (AI_dictFind(mangled_persisted, key) != NULL) { + AI_dictRelease(mangled_persisted); + AI_dictReleaseIterator(iter); + RAI_SetError(rinfo->err, RAI_EDAGRUN, "ERR PERSIST keys must be unique"); + goto cleanup; + } + int *instance = AI_dictGetVal(mangled_entry); + char buf[16]; + sprintf(buf, "%04d", *instance); + RedisModuleString *mangled_key = RedisModule_CreateStringFromString(NULL, key); + RedisModule_StringAppendBuffer(NULL, mangled_key, buf, strlen(buf)); + AI_dictAdd(mangled_persisted, (void *)mangled_key, (void *)1); + RedisModule_FreeString(NULL, mangled_key); + entry = AI_dictNext(iter); + } + AI_dictReleaseIterator(iter); + } + + AI_dictRelease(rinfo->dagTensorsPersistedContext); + rinfo->dagTensorsPersistedContext = mangled_persisted; + + for (long long i = 0; i < array_len(rinfo->dagOps); i++) { + if (rinfo->dagOps[i]->devicestr == NULL) { + rinfo->dagOps[i]->devicestr = "CPU"; + } + } + // Tensors from TENSORSET ops are ready to be put in DAG local context under their mangled + // names. + _DAG_SetTensorsInLocalContext(rinfo); + res = REDISMODULE_OK; + +cleanup : { + AI_dictIterator *iter = AI_dictGetSafeIterator(mangled_tensors); + AI_dictEntry *entry = AI_dictNext(iter); + while (entry) { + int *val = (int *)AI_dictGetVal(entry); + RedisModule_Free(val); + entry = AI_dictNext(iter); + } + AI_dictReleaseIterator(iter); +} + AI_dictRelease(mangled_tensors); + return res; +} + +// Add Shallow copies of the DAG run info to the devices' queues. +// Return REDISMODULE_OK in case of success, REDISMODULE_ERR if (at least) one insert op had +// failed. +int DAG_InsertDAGToQueue(RedisAI_RunInfo *rinfo) { + const char **devices = array_new(const char *, 10); + + for (long long i = 0; i < array_len(rinfo->dagOps); i++) { + const char *devicestr = rinfo->dagOps[i]->devicestr; + bool found = false; + for (long long j = 0; j < array_len(devices); j++) { + if (strcasecmp(devicestr, devices[j]) == 0) { + found = true; + break; + } + } + if (!found) { + devices = array_append(devices, devicestr); + } + } + + size_t ndevices = array_len(devices); + if (ndevices == 1) + rinfo->single_device_dag = 1; + RedisAI_RunInfo **rinfo_copies = array_new(RedisAI_RunInfo *, ndevices); + + for (long long i = 0; i < ndevices; i++) { + RedisAI_RunInfo *rinfo_copy; + RAI_ShallowCopyDagRunInfo(&rinfo_copy, rinfo); + rinfo_copies = array_append(rinfo_copies, rinfo_copy); + } + + for (long long i = 0; i < ndevices; i++) { + RedisAI_RunInfo *rinfo_copy = rinfo_copies[i]; + for (long long j = 0; j < rinfo_copy->dagOpCount; j++) { + if (strcasecmp(rinfo_copy->dagOps[j]->devicestr, devices[i]) == 0) { + rinfo_copy->dagDeviceOps = + array_append(rinfo_copy->dagDeviceOps, rinfo_copy->dagOps[j]); + } + } + rinfo_copy->dagDeviceOpCount = array_len(rinfo_copy->dagDeviceOps); + } + + RunQueueInfo **run_queues_info = array_new(RunQueueInfo *, ndevices); + for (long long i = 0; i < ndevices; i++) { + const char *devicestr = devices[i]; + RunQueueInfo *run_queue_info = NULL; + if (ensureRunQueue(devicestr, &run_queue_info) == REDISMODULE_ERR) { + // A device run queue was not created properly, so we free everything, + // set an error and finish. + array_free(devices); + for (int j = 0; j < ndevices; j++) { + RAI_DagRunInfoFreeShallowCopy(rinfo_copies[j]); + } + array_free(rinfo_copies); + array_free(run_queues_info); + RAI_SetError(rinfo->err, RAI_EDAGRUN, "ERR Queue not initialized for device"); + return REDISMODULE_ERR; + } + run_queues_info = array_append(run_queues_info, run_queue_info); + } + for (long long i = 0; i < ndevices; i++) { + RedisAI_RunInfo *rinfo_copy = rinfo_copies[i]; + RunQueueInfo *run_queue_info = run_queues_info[i]; + gettimeofday(&rinfo_copy->queuingTime, NULL); + + pthread_mutex_lock(&run_queue_info->run_queue_mutex); + queuePush(run_queue_info->run_queue, rinfo_copy); + pthread_cond_signal(&run_queue_info->queue_condition_var); + pthread_mutex_unlock(&run_queue_info->run_queue_mutex); + } + + array_free(devices); + array_free(rinfo_copies); + array_free(run_queues_info); + return REDISMODULE_OK; +} + +int RAI_DAGRun(RAI_DAGRunCtx *run_info, RAI_OnFinishCB DAGAsyncFinish, void *private_data, + RAI_Error *err) { + + RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)run_info; + rinfo->dagOpCount = array_len(rinfo->dagOps); + if (rinfo->dagOpCount < 1) { + RAI_SetError(err, RAI_EDAGRUN, "ERR DAG is empty"); + return REDISMODULE_ERR; + } + // Make the inkeys and outkeys of the DAG ops unique, to ensure that the operations + // will be execute in the right order. + if (MangleTensorsNames(rinfo) != REDISMODULE_OK) { + RAI_SetError(err, rinfo->err->code, rinfo->err->detail); + return REDISMODULE_ERR; + } + rinfo->OnFinish = (RedisAI_OnFinishCB)DAGAsyncFinish; + rinfo->private_data = private_data; + if (DAG_InsertDAGToQueue(rinfo) != REDISMODULE_OK) { + RAI_SetError(err, rinfo->err->code, rinfo->err->detail); + return REDISMODULE_ERR; + } + return REDISMODULE_OK; +} + +size_t RAI_DAGNumOutputs(RAI_OnFinishCtx *finish_ctx) { + size_t n_outputs = 0; + RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)finish_ctx; + for (size_t i = 0; i < rinfo->dagOpCount; i++) { + if (rinfo->dagOps[i]->commandType == REDISAI_DAG_CMD_TENSORGET) { + n_outputs++; + } + } + return n_outputs; +} + +const RAI_Tensor *RAI_DAGOutputTensor(RAI_OnFinishCtx *finish_ctx, size_t index) { + size_t tensor_get_op_ind = -1; + RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)finish_ctx; + for (size_t i = 0; i < rinfo->dagOpCount; i++) { + RAI_DagOp *op = rinfo->dagOps[i]; + if (op->commandType == REDISAI_DAG_CMD_TENSORGET) { + tensor_get_op_ind++; + if (tensor_get_op_ind == index) { + RAI_Tensor *t; + int res = RAI_getTensorFromLocalContext(rinfo->dagTensorsContext, op->inkeys[0], &t, + op->err); + RedisModule_Assert(res == REDISMODULE_OK); + return t; + } + } + } + return NULL; +} + +bool RAI_DAGRunError(RAI_OnFinishCtx *finish_ctx) { + return *((RedisAI_RunInfo *)finish_ctx)->dagError; +} + +const RAI_Error *RAI_DAGGetError(RAI_OnFinishCtx *finish_ctx) { + RedisAI_RunInfo *rinfo = (RedisAI_RunInfo *)finish_ctx; + return rinfo->err; +} diff --git a/src/DAG/dag_execute.h b/src/DAG/dag_execute.h new file mode 100644 index 000000000..c4b5b1b99 --- /dev/null +++ b/src/DAG/dag_execute.h @@ -0,0 +1,58 @@ +#pragma once + +#include "redismodule.h" +#include "run_info.h" + +/** + * @brief We are given a DAG runInfo of a sequence of operations, each with its own + input and output keys. The names of the keys will be used to look whether the + inputs to a DAG operation have all been realized by previous operations (or if + they are available as part of LOADed keys from keyspace). + This strategy is fine if keys are not aliased, that is, if a command's output + overwrites the key of a previous command. This would trick DAG operations into + thinking that their input is ready when it's not. + To overcome this, we make key names unique, so that names are not aliased. We + mangle the names by appending a numerical suffix ":0001". After computing, we + demangle the keys in order to persist them.*/ +int MangleTensorsNames(RedisAI_RunInfo *rinfo); + +/** + * @brief Run asynchronously a DAG. This will validate that the sequence of DAG ops + * is valid and generate a unique key to the tensor that flow in the DAG (mangleTensorsNames) + * Then, DAG is sent to the devices' run queues and will be execute by a workung thread. + * @param DAGAsyncFinish This is a callback that will be called after the whole DAG finish its run. + * @param private_data This is an input to the DAGAsyncFinish callback. Can be used to save the + * results and errors + * @param err Error is returned in case that the validation failed, and the DAG wasn't inserted to + * the queues. + */ +int RAI_DAGRun(RAI_DAGRunCtx *run_info, RAI_OnFinishCB DAGAsyncFinish, void *private_data, + RAI_Error *err); + +/** + * @brief This can be called in the finish CB, returns the number of outputs (TENSORGET ops). + * @param finish_ctx This represents the DAG runInfo at the end of the run. + */ +size_t RAI_DAGNumOutputs(RAI_OnFinishCtx *finish_ctx); + +/** + * @brief This can be called in the finish CB, returns a specific output tensor (result of a + * TENSORGET op). + * @param finish_ctx This represents the DAG runInfo at the end of the run. + * @param index The index of the TENSORGET op in the DAG. + * @retval returns the tensor that the i'th TENSORGET op outputs. + */ +const RAI_Tensor *RAI_DAGOutputTensor(RAI_OnFinishCtx *finish_ctx, size_t index); + +/** + * @brief Returns true if (at least) one of the DAG ops encountered an error. + */ +bool RAI_DAGRunError(RAI_OnFinishCtx *finish_ctx); + +/** + * @brief This can be called in the finish CB, to get DAG error details. + * @param finish_ctx This represents the DAG runInfo at the end of the run. + * @retval returns an object that represents the DAG status, from which a user can + * obtain the error code (error code is "OK" if no error has occurred) and error details. + */ +const RAI_Error *RAI_DAGGetError(RAI_OnFinishCtx *finish_ctx); diff --git a/src/DAG/dag_parser.c b/src/DAG/dag_parser.c index 968038da9..ad7e53890 100644 --- a/src/DAG/dag_parser.c +++ b/src/DAG/dag_parser.c @@ -4,159 +4,10 @@ #include "dag_parser.h" #include "command_parser.h" #include "modelRun_ctx.h" +#include "dag_execute.h" #include -#include "util/string_utils.h" - -void _SetTensorsInDagLocalContext(RedisAI_RunInfo *rinfo) { - for (size_t i = 0; i < rinfo->dagOpCount; i++) { - RAI_DagOp *op = rinfo->dagOps[i]; - if (op->commandType == REDISAI_DAG_CMD_TENSORSET) { - // Insert the tensor with its mangled (unique) name. - void *t = (void *)RAI_TensorGetShallowCopy(op->outTensor); - AI_dictReplace(rinfo->dagTensorsContext, (void *)op->outkeys[0], t); - } - } -} - -/* At this point, we have built a sequence of DAG operations, each with its own - input and output keys. The names of the keys will be used to look whether the - inputs to a DAG operation have all been realized by previous operations (or if - they are available as part of LOADed keys from keyspace). - This strategy is fine if keys are not aliased, that is, if a command's output - overwrites the key of a previous command. This would trick DAG operations into - thinking that their input is ready when it's not. - To overcome this, we make key names unique, so that names are not aliased. We - mangle the names by appending a numerical suffix ":0001". After computing, we - demangle the keys in order to persist them.*/ -int _MangleTensorsNames(RedisModuleCtx *ctx, RedisAI_RunInfo *rinfo) { - - int res = REDISMODULE_ERR; - AI_dict *occurrences_counter = AI_dictCreate(&AI_dictTypeHeapRStrings, NULL); - - { - // We first save the tensors' names that were indicated in the LOAD phase. - // These tensors where loaded and kept in dagTensorsContext with their "mangled" name. - AI_dictIterator *iter = AI_dictGetSafeIterator(rinfo->dagTensorsContext); - AI_dictEntry *entry = AI_dictNext(iter); - while (entry) { - RedisModuleString *key = (RedisModuleString *)AI_dictGetKey(entry); - size_t key_len; - const char *key_str = RedisModule_StringPtrLen(key, &key_len); - RedisModuleString *demangled_key = RedisModule_CreateString(NULL, key_str, key_len - 4); - int *instance = RedisModule_Alloc(sizeof(int)); - *instance = 1; - AI_dictAdd(occurrences_counter, (void *)demangled_key, (void *)instance); - RedisModule_FreeString(NULL, demangled_key); - entry = AI_dictNext(iter); - } - AI_dictReleaseIterator(iter); - } - - for (long long i = 0; i < array_len(rinfo->dagOps); i++) { - RAI_DagOp *currentOp = rinfo->dagOps[i]; - - RedisModuleString **mangled_inkeys = - array_new(RedisModuleString *, array_len(currentOp->inkeys)); - for (long long j = 0; j < array_len(currentOp->inkeys); j++) { - RedisModuleString *key = currentOp->inkeys[j]; - AI_dictEntry *entry = AI_dictFind(occurrences_counter, key); - if (!entry) { - array_free(mangled_inkeys); - RedisModule_ReplyWithError(ctx, "ERR INPUT key cannot be found in DAG"); - goto cleanup; - } - int *instance = AI_dictGetVal(entry); - char buf[16]; - sprintf(buf, "%04d", *instance); - RedisModuleString *mangled_key = RedisModule_CreateStringFromString(NULL, key); - RedisModule_StringAppendBuffer(NULL, mangled_key, buf, strlen(buf)); - mangled_inkeys = array_append(mangled_inkeys, mangled_key); - } - - RedisModuleString **mangled_outkeys = - array_new(RedisModuleString *, array_len(currentOp->outkeys)); - for (long long j = 0; j < array_len(currentOp->outkeys); j++) { - RedisModuleString *key = currentOp->outkeys[j]; - AI_dictEntry *entry = AI_dictFind(occurrences_counter, key); - int *instance = NULL; - if (entry) { - instance = AI_dictGetVal(entry); - *instance += 1; - } else { - instance = RedisModule_Alloc(sizeof(int)); - *instance = 1; - AI_dictAdd(occurrences_counter, (void *)key, (void *)instance); - } - char buf[16]; - sprintf(buf, "%04d", *instance); - RedisModuleString *mangled_key = RedisModule_CreateStringFromString(NULL, key); - RedisModule_StringAppendBuffer(NULL, mangled_key, buf, strlen(buf)); - mangled_outkeys = array_append(mangled_outkeys, mangled_key); - } - - for (size_t j = 0; j < array_len(currentOp->inkeys); j++) { - RedisModule_FreeString(NULL, currentOp->inkeys[j]); - } - array_free(currentOp->inkeys); - - for (size_t j = 0; j < array_len(currentOp->outkeys); j++) { - RedisModule_FreeString(NULL, currentOp->outkeys[j]); - } - array_free(currentOp->outkeys); - - currentOp->inkeys = mangled_inkeys; - currentOp->outkeys = mangled_outkeys; - } - - // If we need to persist a certain tensor under a specified key, we need to take it - // from the last op in which this key appears (that is, the tensor associated with - // the "maximal" mangled name generated from that key). - AI_dict *mangled_persisted = AI_dictCreate(&AI_dictTypeHeapRStrings, NULL); - { - AI_dictIterator *iter = AI_dictGetSafeIterator(rinfo->dagTensorsPersistedContext); - AI_dictEntry *entry = AI_dictNext(iter); - while (entry) { - RedisModuleString *key = (RedisModuleString *)AI_dictGetKey(entry); - AI_dictEntry *mangled_entry = AI_dictFind(occurrences_counter, key); - if (!mangled_entry) { - AI_dictRelease(mangled_persisted); - AI_dictReleaseIterator(iter); - RedisModule_ReplyWithError(ctx, "ERR PERSIST key cannot be found in DAG"); - goto cleanup; - } - int *instance = AI_dictGetVal(mangled_entry); - char buf[16]; - sprintf(buf, "%04d", *instance); - RedisModule_StringAppendBuffer(NULL, key, buf, strlen(buf)); - AI_dictAdd(mangled_persisted, (void *)key, (void *)1); - entry = AI_dictNext(iter); - } - AI_dictReleaseIterator(iter); - } - - AI_dictRelease(rinfo->dagTensorsPersistedContext); - rinfo->dagTensorsPersistedContext = mangled_persisted; - - for (long long i = 0; i < array_len(rinfo->dagOps); i++) { - if (rinfo->dagOps[i]->devicestr == NULL) { - rinfo->dagOps[i]->devicestr = "CPU"; - } - } - res = REDISMODULE_OK; - -cleanup : { - AI_dictIterator *iter = AI_dictGetSafeIterator(occurrences_counter); - AI_dictEntry *entry = AI_dictNext(iter); - while (entry) { - int *val = (int *)AI_dictGetVal(entry); - RedisModule_Free(val); - entry = AI_dictNext(iter); - } - AI_dictReleaseIterator(iter); -} - AI_dictRelease(occurrences_counter); - return res; -} +#include "dag.h" +#include "string_utils.h" /** * DAGRUN Building Block to parse [LOAD key1 key2... ] @@ -173,17 +24,19 @@ cleanup : { * @return processed number of arguments on success, or -1 if the parsing failed */ static int _ParseDAGLoadArgs(RedisModuleCtx *ctx, RedisModuleString **argv, int argc, - AI_dict **localContextDict, const char *chaining_operator) { + AI_dict **localContextDict, const char *chaining_operator, + RAI_Error *err) { if (argc < 3) { - RedisModule_WrongArity(ctx); + RAI_SetError(err, RAI_EDAGBUILDER, + "ERR wrong number of arguments for LOAD in 'AI.DAGRUN' command"); return -1; } long long n_keys; const int retval = RedisModule_StringToLongLong(argv[1], &n_keys); if (retval != REDISMODULE_OK || n_keys <= 0) { - RedisModule_ReplyWithError(ctx, - "ERR invalid or negative value found in number of keys to LOAD"); + RAI_SetError(err, RAI_EDAGBUILDER, + "ERR invalid or negative value found in number of keys to LOAD"); return -1; } @@ -198,7 +51,8 @@ static int _ParseDAGLoadArgs(RedisModuleCtx *ctx, RedisModuleString **argv, int break; RAI_Tensor *t; RedisModuleKey *key; - const int status = RAI_GetTensorFromKeyspace(ctx, key_name, &key, &t, REDISMODULE_READ); + const int status = + RAI_GetTensorFromKeyspace(ctx, key_name, &key, &t, REDISMODULE_READ, err); if (status == REDISMODULE_ERR) { RedisModule_Log(ctx, "warning", "on DAGRUN's LOAD could not load tensor %s from keyspace", arg_string); @@ -214,7 +68,9 @@ static int _ParseDAGLoadArgs(RedisModuleCtx *ctx, RedisModuleString **argv, int } if (number_loaded_keys != n_keys) { - RedisModule_WrongArity(ctx); + RAI_SetError(err, RAI_EDAGBUILDER, + "ERR number of keys to LOAD in AI.DAGRUN command does not match the number of " + "given arguments"); return -1; } return number_loaded_keys + 2; @@ -233,17 +89,19 @@ static int _ParseDAGLoadArgs(RedisModuleCtx *ctx, RedisModuleString **argv, int * @return processed number of arguments on success, or -1 if the parsing failed */ static int _ParseDAGPersistArgs(RedisModuleCtx *ctx, RedisModuleString **argv, int argc, - AI_dict **persistContextDict, const char *chaining_operator) { + AI_dict **persistContextDict, const char *chaining_operator, + RAI_Error *err) { if (argc < 3) { - RedisModule_WrongArity(ctx); + RAI_SetError(err, RAI_EDAGBUILDER, + "ERR wrong number of arguments for PERSIST in 'AI.DAGRUN' command"); return -1; } long long n_keys; const int retval = RedisModule_StringToLongLong(argv[1], &n_keys); if (retval != REDISMODULE_OK || n_keys <= 0) { - RedisModule_ReplyWithError( - ctx, "ERR invalid or negative value found in number of keys to PERSIST"); + RAI_SetError(err, RAI_EDAGBUILDER, + "ERR invalid or negative value found in number of keys to PERSIST"); return -1; } @@ -259,31 +117,32 @@ static int _ParseDAGPersistArgs(RedisModuleCtx *ctx, RedisModuleString **argv, i } } if (number_keys_to_persist != n_keys) { - RedisModule_WrongArity(ctx); + RAI_SetError(err, RAI_EDAGBUILDER, + "ERR number of keys to PERSIST in AI.DAGRUN command does not match the number " + "of given arguments"); return -1; } return number_keys_to_persist + 2; } -static int _parseTimeout(RedisModuleCtx *ctx, RedisModuleString **argv, int argc, - long long *timeout) { +static int _parseTimeout(RedisModuleString **argv, int argc, long long *timeout, RAI_Error *err) { - if (argc == 0) { - RedisModule_ReplyWithError(ctx, "ERR No value provided for TIMEOUT"); + if (argc < 2) { + RAI_SetError(err, RAI_EDAGBUILDER, "ERR No value provided for TIMEOUT"); return REDISMODULE_ERR; } const int retval = RedisModule_StringToLongLong(argv[1], timeout); if (retval != REDISMODULE_OK || timeout <= 0) { - RedisModule_ReplyWithError(ctx, "ERR Invalid value for TIMEOUT"); + RAI_SetError(err, RAI_EDAGBUILDER, "ERR Invalid value for TIMEOUT"); return REDISMODULE_ERR; } return REDISMODULE_OK; } -static RAI_DagOp *_AddEmptyOp(RedisAI_RunInfo *rinfo) { +static RAI_DagOp *_AddEmptyOp(RAI_DagOp ***ops) { RAI_DagOp *currentDagOp; RAI_InitDagOp(¤tDagOp); - rinfo->dagOps = array_append(rinfo->dagOps, currentDagOp); + *ops = array_append(*ops, currentDagOp); return currentDagOp; } @@ -301,10 +160,10 @@ int _CollectOpArgs(RedisModuleString **argv, int argc, int arg_pos, RAI_DagOp *o return op->argc; } -int _ParseDAGOps(RedisModuleCtx *ctx, RedisAI_RunInfo *rinfo) { +int ParseDAGOps(RedisAI_RunInfo *rinfo, RAI_DagOp **ops) { - for (long long i = 0; i < array_len(rinfo->dagOps); i++) { - RAI_DagOp *currentOp = rinfo->dagOps[i]; + for (long long i = 0; i < array_len(ops); i++) { + RAI_DagOp *currentOp = ops[i]; // The first op arg is the command name. const char *arg_string = RedisModule_StringPtrLen(currentOp->argv[0], NULL); @@ -313,9 +172,9 @@ int _ParseDAGOps(RedisModuleCtx *ctx, RedisAI_RunInfo *rinfo) { currentOp->devicestr = "CPU"; RAI_HoldString(NULL, currentOp->argv[1]); currentOp->inkeys = array_append(currentOp->inkeys, currentOp->argv[1]); - currentOp->fmt = ParseTensorGetArgs(ctx, currentOp->argv, currentOp->argc); + currentOp->fmt = ParseTensorGetArgs(rinfo->err, currentOp->argv, currentOp->argc); if (currentOp->fmt == TENSOR_NONE) - return REDISMODULE_ERR; + goto cleanup; continue; } if (!strcasecmp(arg_string, "AI.TENSORSET")) { @@ -323,39 +182,57 @@ int _ParseDAGOps(RedisModuleCtx *ctx, RedisAI_RunInfo *rinfo) { currentOp->devicestr = "CPU"; RAI_HoldString(NULL, currentOp->argv[1]); currentOp->outkeys = array_append(currentOp->outkeys, currentOp->argv[1]); - if (RAI_parseTensorSetArgs(ctx, currentOp->argv, currentOp->argc, ¤tOp->outTensor, - 0, currentOp->err) == -1) - return REDISMODULE_ERR; + if (RAI_parseTensorSetArgs(currentOp->argv, currentOp->argc, ¤tOp->outTensor, 0, + rinfo->err) == -1) + goto cleanup; continue; } if (!strcasecmp(arg_string, "AI.MODELRUN")) { - if (ParseModelRunCommand(rinfo, currentOp, ctx, currentOp->argv, currentOp->argc) != + if (ParseModelRunCommand(rinfo, currentOp, currentOp->argv, currentOp->argc) != REDISMODULE_OK) { - return REDISMODULE_ERR; + goto cleanup; } continue; } if (!strcasecmp(arg_string, "AI.SCRIPTRUN")) { - if (ParseScriptRunCommand(rinfo, currentOp, ctx, currentOp->argv, currentOp->argc) != + if (ParseScriptRunCommand(rinfo, currentOp, currentOp->argv, currentOp->argc) != REDISMODULE_OK) { - return REDISMODULE_ERR; + goto cleanup; } continue; } // If none of the cases match, we have an invalid op. - RedisModule_ReplyWithError(ctx, "unsupported command within DAG"); - return REDISMODULE_ERR; + RAI_SetError(rinfo->err, RAI_EDAGBUILDER, "unsupported command within DAG"); + goto cleanup; + } + + // After validating all the ops, insert them to the DAG. + for (size_t i = 0; i < array_len(ops); i++) { + rinfo->dagOps = array_append(rinfo->dagOps, ops[i]); } + rinfo->dagOpCount = array_len(rinfo->dagOps); return REDISMODULE_OK; + +cleanup: + for (size_t i = 0; i < array_len(ops); i++) { + RAI_FreeDagOp(ops[i]); + } + return REDISMODULE_ERR; } -// Parse the DAG run command and return REDISMODULE_OK only if it is a valid command to execute. int ParseDAGRunCommand(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModuleString **argv, int argc, bool dag_ro) { + int res = REDISMODULE_ERR; if (argc < 4) { - RedisModule_WrongArity(ctx); - goto cleanup; + if (dag_ro) { + RAI_SetError(rinfo->err, RAI_EDAGBUILDER, + "ERR wrong number of arguments for 'AI.DAGRUN_RO' command"); + } else { + RAI_SetError(rinfo->err, RAI_EDAGBUILDER, + "ERR wrong number of arguments for 'AI.DAGRUN' command"); + } + return res; } int chainingOpCount = 0; @@ -363,6 +240,7 @@ int ParseDAGRunCommand(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModuleS bool load_complete = false; bool persist_complete = false; bool timeout_complete = false; + array_new_on_stack(RAI_DagOp *, 10, dag_ops); // The first arg is "AI.DAGRUN", so we go over from the next arg. while (arg_pos < argc) { @@ -371,8 +249,8 @@ int ParseDAGRunCommand(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModuleS if (!strcasecmp(arg_string, "LOAD") && !load_complete && chainingOpCount == 0) { /* Load the required tensors from key space and store them in both dagTensorsLoadedContext and dagTensorsContext dicts. */ - const int parse_result = _ParseDAGLoadArgs(ctx, &argv[arg_pos], argc - arg_pos, - &(rinfo->dagTensorsContext), "|>"); + const int parse_result = _ParseDAGLoadArgs( + ctx, &argv[arg_pos], argc - arg_pos, &(rinfo->dagTensorsContext), "|>", rinfo->err); if (parse_result <= 0) goto cleanup; arg_pos += parse_result; @@ -381,14 +259,15 @@ int ParseDAGRunCommand(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModuleS } if (!strcasecmp(arg_string, "PERSIST") && !persist_complete && chainingOpCount == 0) { if (dag_ro) { - RedisModule_ReplyWithError(ctx, - "ERR PERSIST cannot be specified in a read-only DAG"); + RAI_SetError(rinfo->err, RAI_EDAGBUILDER, + "ERR PERSIST cannot be specified in a read-only DAG"); goto cleanup; } /* Store the keys to persist in dagTensorsPersistedContext dict. These keys will be populated later on with actual tensors. */ - const int parse_result = _ParseDAGPersistArgs( - ctx, &argv[arg_pos], argc - arg_pos, &(rinfo->dagTensorsPersistedContext), "|>"); + const int parse_result = + _ParseDAGPersistArgs(ctx, &argv[arg_pos], argc - arg_pos, + &(rinfo->dagTensorsPersistedContext), "|>", rinfo->err); if (parse_result <= 0) goto cleanup; @@ -398,7 +277,8 @@ int ParseDAGRunCommand(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModuleS } if (!strcasecmp(arg_string, "TIMEOUT") && !timeout_complete && chainingOpCount == 0) { long long timeout; - if (_parseTimeout(ctx, &argv[arg_pos], argc - arg_pos, &timeout) == REDISMODULE_ERR) + if (_parseTimeout(&argv[arg_pos], argc - arg_pos, &timeout, rinfo->err) == + REDISMODULE_ERR) goto cleanup; rinfo->timeout = timeout; arg_pos += 2; @@ -407,27 +287,32 @@ int ParseDAGRunCommand(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModuleS } if (!strcasecmp(arg_string, "|>") && arg_pos < argc - 1) { - RAI_DagOp *currentOp = _AddEmptyOp(rinfo); + RAI_DagOp *currentOp = _AddEmptyOp(&dag_ops); chainingOpCount++; int args_num = _CollectOpArgs(argv, argc, ++arg_pos, currentOp); arg_pos += args_num; continue; } // If none of the cases match, we have an invalid op. - RedisModule_ReplyWithError(ctx, "ERR Invalid DAGRUN command"); + RAI_SetError(rinfo->err, RAI_EDAGBUILDER, "ERR Invalid DAGRUN command"); goto cleanup; } - rinfo->dagOpCount = array_len(rinfo->dagOps); - if (rinfo->dagOpCount < 1) + + if (array_len(dag_ops) < 1) { + RAI_SetError(rinfo->err, RAI_EDAGBUILDER, "ERR DAG is empty"); goto cleanup; - if (_ParseDAGOps(ctx, rinfo) != REDISMODULE_OK) + } + + if (ParseDAGOps(rinfo, dag_ops) != REDISMODULE_OK) { goto cleanup; - if (_MangleTensorsNames(ctx, rinfo) != REDISMODULE_OK) + } + + if (MangleTensorsNames(rinfo) != REDISMODULE_OK) { goto cleanup; - _SetTensorsInDagLocalContext(rinfo); - return REDISMODULE_OK; + } + res = REDISMODULE_OK; cleanup: - RAI_FreeRunInfo(rinfo); - return REDISMODULE_ERR; + array_free(dag_ops); + return res; } diff --git a/src/DAG/dag_parser.h b/src/DAG/dag_parser.h index 8160f890b..bc3d7a30c 100644 --- a/src/DAG/dag_parser.h +++ b/src/DAG/dag_parser.h @@ -13,3 +13,14 @@ */ int ParseDAGRunCommand(RedisAI_RunInfo *rinfo, RedisModuleCtx *ctx, RedisModuleString **argv, int argc, bool dag_ro); + +/** + * @brief Parse the arguments of the given ops in the DAGRUN command and build every op accordingly. + * @param rinfo The DAG run info that will be populated with the ops if they are valid. + * with its op, + * @param ops A local array of ops, where every op has an argv field that points to an + * array of RedisModule strings arguments, and an argc field which is the number of + * args. + * @return Returns REDISMODULE_OK if the command is valid, REDISMODULE_ERR otherwise. + */ +int ParseDAGOps(RedisAI_RunInfo *rinfo, RAI_DagOp **ops); diff --git a/src/command_parser.c b/src/command_parser.c index 3310b256a..33570f212 100644 --- a/src/command_parser.c +++ b/src/command_parser.c @@ -17,7 +17,7 @@ static int _parseTimeout(RedisModuleString *timeout_arg, RAI_Error *error, long return REDISMODULE_OK; } -static int _ModelRunCommand_ParseArgs(RedisModuleCtx *ctx, RedisModuleString **argv, int argc, +static int _ModelRunCommand_ParseArgs(RedisModuleCtx *ctx, int argc, RedisModuleString **argv, RAI_Model **model, RAI_Error *error, RedisModuleString ***inkeys, RedisModuleString ***outkeys, RedisModuleString **runkey, long long *timeout) { @@ -93,7 +93,8 @@ static int _ModelRunCommand_ParseArgs(RedisModuleCtx *ctx, RedisModuleString **a */ static int _ModelRunCtx_SetParams(RedisModuleCtx *ctx, RedisModuleString **inkeys, - RedisModuleString **outkeys, RAI_ModelRunCtx *mctx) { + RedisModuleString **outkeys, RAI_ModelRunCtx *mctx, + RAI_Error *err) { RAI_Model *model = mctx->model; RAI_Tensor *t; @@ -101,9 +102,10 @@ static int _ModelRunCtx_SetParams(RedisModuleCtx *ctx, RedisModuleString **inkey char *opname = NULL; 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); + const int status = + RAI_GetTensorFromKeyspace(ctx, inkeys[i], &key, &t, REDISMODULE_READ, err); if (status == REDISMODULE_ERR) { - RedisModule_Log(ctx, "warning", "could not load tensor %s from keyspace", + RedisModule_Log(ctx, "warning", "could not load input tensor %s from keyspace", RedisModule_StringPtrLen(inkeys[i], NULL)); return REDISMODULE_ERR; } @@ -111,6 +113,7 @@ static int _ModelRunCtx_SetParams(RedisModuleCtx *ctx, RedisModuleString **inkey opname = model->inputs[i]; RAI_ModelRunCtxAddInput(mctx, opname, t); } + for (size_t i = 0; i < noutputs; i++) { if (model->outputs) opname = model->outputs[i]; @@ -119,22 +122,22 @@ static int _ModelRunCtx_SetParams(RedisModuleCtx *ctx, RedisModuleString **inkey return REDISMODULE_OK; } -int ParseModelRunCommand(RedisAI_RunInfo *rinfo, RAI_DagOp *currentOp, RedisModuleCtx *ctx, - RedisModuleString **argv, int argc) { +int ParseModelRunCommand(RedisAI_RunInfo *rinfo, RAI_DagOp *currentOp, RedisModuleString **argv, + int argc) { + int res = REDISMODULE_ERR; // Build a ModelRunCtx from command. + RedisModuleCtx *ctx = RedisModule_GetThreadSafeContext(NULL); RAI_Model *model; - long long timeout = 0; - if (_ModelRunCommand_ParseArgs(ctx, argv, argc, &model, currentOp->err, ¤tOp->inkeys, + if (_ModelRunCommand_ParseArgs(ctx, argc, argv, &model, rinfo->err, ¤tOp->inkeys, ¤tOp->outkeys, ¤tOp->runkey, &timeout) == REDISMODULE_ERR) { - RedisModule_ReplyWithError(ctx, RAI_GetErrorOneLine(currentOp->err)); goto cleanup; } if (timeout > 0 && !rinfo->single_op_dag) { - RedisModule_ReplyWithError(ctx, "ERR TIMEOUT not allowed within a DAG command"); + RAI_SetError(rinfo->err, RAI_EDAGBUILDER, "ERR TIMEOUT not allowed within a DAG command"); goto cleanup; } @@ -146,16 +149,15 @@ int ParseModelRunCommand(RedisAI_RunInfo *rinfo, RAI_DagOp *currentOp, RedisModu if (rinfo->single_op_dag) { rinfo->timeout = timeout; // Set params in ModelRunCtx, bring inputs from key space. - if (_ModelRunCtx_SetParams(ctx, currentOp->inkeys, currentOp->outkeys, mctx) == + if (_ModelRunCtx_SetParams(ctx, currentOp->inkeys, currentOp->outkeys, mctx, rinfo->err) == REDISMODULE_ERR) goto cleanup; } - return REDISMODULE_OK; + res = REDISMODULE_OK; cleanup: - if (rinfo->single_op_dag) - RAI_FreeRunInfo(rinfo); - return REDISMODULE_ERR; + RedisModule_FreeThreadSafeContext(ctx); + return res; } static int _ScriptRunCommand_ParseArgs(RedisModuleCtx *ctx, RedisModuleString **argv, int argc, @@ -172,9 +174,8 @@ static int _ScriptRunCommand_ParseArgs(RedisModuleCtx *ctx, RedisModuleString ** size_t argpos = 1; RedisModuleKey *scriptKey; const int status = - RAI_GetScriptFromKeyspace(ctx, argv[argpos], &scriptKey, script, REDISMODULE_READ); + RAI_GetScriptFromKeyspace(ctx, argv[argpos], &scriptKey, script, REDISMODULE_READ, error); if (status == REDISMODULE_ERR) { - RAI_SetError(error, RAI_ESCRIPTRUN, "ERR Script not found"); return REDISMODULE_ERR; } RAI_HoldString(NULL, argv[argpos]); @@ -276,16 +277,17 @@ static int _ScriptRunCommand_ParseArgs(RedisModuleCtx *ctx, RedisModuleString ** */ static int _ScriptRunCtx_SetParams(RedisModuleCtx *ctx, RedisModuleString **inkeys, - RedisModuleString **outkeys, RAI_ScriptRunCtx *sctx) { + RedisModuleString **outkeys, RAI_ScriptRunCtx *sctx, + RAI_Error *err) { 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); + const int status = + RAI_GetTensorFromKeyspace(ctx, inkeys[i], &key, &t, REDISMODULE_READ, err); if (status == REDISMODULE_ERR) { - RedisModule_Log(ctx, "warning", "could not load tensor %s from keyspace", + RedisModule_Log(ctx, "warning", "could not load input tensor %s from keyspace", RedisModule_StringPtrLen(inkeys[i], NULL)); return REDISMODULE_ERR; } @@ -297,23 +299,25 @@ static int _ScriptRunCtx_SetParams(RedisModuleCtx *ctx, RedisModuleString **inke return REDISMODULE_OK; } -int ParseScriptRunCommand(RedisAI_RunInfo *rinfo, RAI_DagOp *currentOp, RedisModuleCtx *ctx, - RedisModuleString **argv, int argc) { +int ParseScriptRunCommand(RedisAI_RunInfo *rinfo, RAI_DagOp *currentOp, RedisModuleString **argv, + int argc) { + int res = REDISMODULE_ERR; // Build a ScriptRunCtx from command. + RedisModuleCtx *ctx = RedisModule_GetThreadSafeContext(NULL); + // int lock_status = RedisModule_ThreadSafeContextTryLock(ctx); RAI_Script *script; const char *func_name = NULL; long long timeout = 0; int variadic = -1; - if (_ScriptRunCommand_ParseArgs(ctx, argv, argc, &script, currentOp->err, ¤tOp->inkeys, + if (_ScriptRunCommand_ParseArgs(ctx, argv, argc, &script, rinfo->err, ¤tOp->inkeys, ¤tOp->outkeys, ¤tOp->runkey, &func_name, &timeout, &variadic) == REDISMODULE_ERR) { - RedisModule_ReplyWithError(ctx, RAI_GetErrorOneLine(currentOp->err)); goto cleanup; } if (timeout > 0 && !rinfo->single_op_dag) { - RedisModule_ReplyWithError(ctx, "ERR TIMEOUT not allowed within a DAG command"); + RAI_SetError(rinfo->err, RAI_EDAGBUILDER, "ERR TIMEOUT not allowed within a DAG command"); goto cleanup; } @@ -326,16 +330,18 @@ int ParseScriptRunCommand(RedisAI_RunInfo *rinfo, RAI_DagOp *currentOp, RedisMod if (rinfo->single_op_dag) { rinfo->timeout = timeout; // Set params in ScriptRunCtx, bring inputs from key space. - if (_ScriptRunCtx_SetParams(ctx, currentOp->inkeys, currentOp->outkeys, sctx) == + if (_ScriptRunCtx_SetParams(ctx, currentOp->inkeys, currentOp->outkeys, sctx, rinfo->err) == REDISMODULE_ERR) goto cleanup; } - return REDISMODULE_OK; + res = REDISMODULE_OK; cleanup: - if (rinfo->single_op_dag) - RAI_FreeRunInfo(rinfo); - return REDISMODULE_ERR; + // if (lock_status == REDISMODULE_OK) { + // RedisModule_ThreadSafeContextUnlock(ctx); + //} + RedisModule_FreeThreadSafeContext(ctx); + return res; } int RedisAI_ExecuteCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc, @@ -357,14 +363,14 @@ int RedisAI_ExecuteCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int ar RAI_DagOp *modelRunOp; RAI_InitDagOp(&modelRunOp); rinfo->dagOps = array_append(rinfo->dagOps, modelRunOp); - status = ParseModelRunCommand(rinfo, modelRunOp, ctx, argv, argc); + status = ParseModelRunCommand(rinfo, modelRunOp, argv, argc); break; case CMD_SCRIPTRUN: rinfo->single_op_dag = 1; RAI_DagOp *scriptRunOp; RAI_InitDagOp(&scriptRunOp); rinfo->dagOps = array_append(rinfo->dagOps, scriptRunOp); - status = ParseScriptRunCommand(rinfo, scriptRunOp, ctx, argv, argc); + status = ParseScriptRunCommand(rinfo, scriptRunOp, argv, argc); break; case CMD_DAGRUN: status = ParseDAGRunCommand(rinfo, ctx, argv, argc, ro_dag); @@ -373,11 +379,17 @@ int RedisAI_ExecuteCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int ar break; } if (status == REDISMODULE_ERR) { + RedisModule_ReplyWithError(ctx, RAI_GetErrorOneLine(rinfo->err)); + RAI_FreeRunInfo(rinfo); return REDISMODULE_OK; } rinfo->dagOpCount = array_len(rinfo->dagOps); rinfo->OnFinish = DAG_ReplyAndUnblock; rinfo->client = RedisModule_BlockClient(ctx, RedisAI_DagRun_Reply, NULL, RunInfo_FreeData, 0); - return DAG_InsertDAGToQueue(rinfo); + if (DAG_InsertDAGToQueue(rinfo) != REDISMODULE_OK) { + RedisModule_UnblockClient(rinfo->client, rinfo); + return REDISMODULE_ERR; + } + return REDISMODULE_OK; } diff --git a/src/command_parser.h b/src/command_parser.h index a71e02f19..74fd6c792 100644 --- a/src/command_parser.h +++ b/src/command_parser.h @@ -12,8 +12,8 @@ typedef enum RunCommand { CMD_MODELRUN = 0, CMD_SCRIPTRUN, CMD_DAGRUN } RunComma * is saved as well (if given, otherwise it is zero). * @return Returns REDISMODULE_OK if the command is valid, REDISMODULE_ERR otherwise. */ -int ParseModelRunCommand(RedisAI_RunInfo *rinfo, RAI_DagOp *currentOp, RedisModuleCtx *ctx, - RedisModuleString **argv, int argc); +int ParseModelRunCommand(RedisAI_RunInfo *rinfo, RAI_DagOp *currentOp, RedisModuleString **argv, + int argc); /** * @brief Parse and validate SCRIPTRUN command: create a scriptRunCtx based on the script obtained @@ -22,8 +22,8 @@ int ParseModelRunCommand(RedisAI_RunInfo *rinfo, RAI_DagOp *currentOp, RedisModu * runkey, and the given timeout is saved as well (if given, otherwise it is zero). * @return Returns REDISMODULE_OK if the command is valid, REDISMODULE_ERR otherwise. */ -int ParseScriptRunCommand(RedisAI_RunInfo *rinfo, RAI_DagOp *currentOp, RedisModuleCtx *ctx, - RedisModuleString **argv, int argc); +int ParseScriptRunCommand(RedisAI_RunInfo *rinfo, RAI_DagOp *currentOp, RedisModuleString **argv, + int argc); /** * @brief Parse and execute RedisAI run command. After parsing and validation, the resulted diff --git a/src/err.c b/src/err.c index 9a4f9db49..56f541c11 100644 --- a/src/err.c +++ b/src/err.c @@ -29,6 +29,13 @@ const char *RAI_GetErrorOneLine(RAI_Error *err) { return err->detail_oneline; } RAI_ErrorCode RAI_GetErrorCode(RAI_Error *err) { return err->code; } +void RAI_CloneError(RAI_Error *dest, const RAI_Error *src) { + dest->code = src->code; + RedisModule_Assert(!dest->detail); + dest->detail = RedisModule_Strdup(src->detail); + dest->detail_oneline = RAI_Chomp(dest->detail); +} + void RAI_SetError(RAI_Error *err, RAI_ErrorCode code, const char *detail) { if (!err) { return; diff --git a/src/err.h b/src/err.h index c3f3a7f9c..e06d17b27 100644 --- a/src/err.h +++ b/src/err.h @@ -25,7 +25,9 @@ typedef enum { RAI_ESCRIPTFREE, RAI_ETENSORSET, RAI_ETENSORGET, + RAI_EDAGBUILDER, RAI_EDAGRUN, + RAI_EFINISHCTX } RAI_ErrorCode; typedef struct RAI_Error { @@ -79,6 +81,14 @@ const char *RAI_GetErrorOneLine(RAI_Error *err); */ RAI_ErrorCode RAI_GetErrorCode(RAI_Error *err); +/** + * Make dest a clone of src + * + * @param dest An allocated error + * @param src The error to copy + */ +void RAI_CloneError(RAI_Error *dest, const RAI_Error *src); + /** * Resets an previously used/allocated RAI_Error * diff --git a/src/model.c b/src/model.c index 3a598ee8f..665058710 100644 --- a/src/model.c +++ b/src/model.c @@ -26,16 +26,16 @@ * Return REDISMODULE_OK if the model value stored at key was correctly * returned and available at *model variable. */ int RAI_GetModelFromKeyspace(RedisModuleCtx *ctx, RedisModuleString *keyName, RedisModuleKey **key, - RAI_Model **model, int mode, RAI_Error *error) { + RAI_Model **model, int mode, RAI_Error *err) { *key = RedisModule_OpenKey(ctx, keyName, mode); if (RedisModule_KeyType(*key) == REDISMODULE_KEYTYPE_EMPTY) { RedisModule_CloseKey(*key); - RAI_SetError(error, REDISMODULE_KEYTYPE_EMPTY, "ERR model key is empty"); + RAI_SetError(err, RAI_EMODELRUN, "ERR model key is empty"); return REDISMODULE_ERR; } if (RedisModule_ModuleTypeGetType(*key) != RedisAI_ModelType) { RedisModule_CloseKey(*key); - RAI_SetError(error, REDISMODULE_ERR, REDISMODULE_ERRORMSG_WRONGTYPE); + RAI_SetError(err, RAI_EMODELRUN, REDISMODULE_ERRORMSG_WRONGTYPE); return REDISMODULE_ERR; } *model = RedisModule_ModuleTypeGetValue(*key); @@ -253,6 +253,10 @@ int RedisAI_ModelRun_IsKeysPositionRequest_ReportKeys(RedisModuleCtx *ctx, Redis RedisModuleType *RAI_ModelRedisType(void) { return RedisAI_ModelType; } +size_t ModelGetNumInputs(RAI_Model *model) { return model->ninputs; } + +size_t ModelGetNumOutputs(RAI_Model *model) { return model->noutputs; } + int RAI_ModelRunAsync(RAI_ModelRunCtx *mctx, RAI_OnFinishCB ModelAsyncFinish, void *private_data) { RedisAI_RunInfo *rinfo = NULL; @@ -270,5 +274,9 @@ int RAI_ModelRunAsync(RAI_ModelRunCtx *mctx, RAI_OnFinishCB ModelAsyncFinish, vo rinfo->dagOps = array_append(rinfo->dagOps, op); rinfo->dagOpCount = 1; - return DAG_InsertDAGToQueue(rinfo); + if (DAG_InsertDAGToQueue(rinfo) != REDISMODULE_OK) { + RAI_FreeRunInfo(rinfo); + return REDISMODULE_ERR; + } + return REDISMODULE_OK; } diff --git a/src/model.h b/src/model.h index 049c19f04..b96082bda 100644 --- a/src/model.h +++ b/src/model.h @@ -122,7 +122,7 @@ int RAI_ModelSerialize(RAI_Model *model, char **buffer, size_t *len, RAI_Error * * an error getting the Model */ int RAI_GetModelFromKeyspace(RedisModuleCtx *ctx, RedisModuleString *keyName, RedisModuleKey **key, - RAI_Model **model, int mode, RAI_Error *error); + RAI_Model **model, int mode, RAI_Error *err); /** * When a module command is called in order to obtain the position of @@ -145,6 +145,15 @@ int RedisAI_ModelRun_IsKeysPositionRequest_ReportKeys(RedisModuleCtx *ctx, Redis */ RedisModuleType *RAI_ModelRedisType(void); +/** + * @brief Returns the number of inputs in the model definition. + */ +size_t ModelGetNumInputs(RAI_Model *model); + +/** + * @brief Returns the number of outputs in the model definition. + */ +size_t ModelGetNumOutputs(RAI_Model *model); /** * Insert the ModelRunCtx to the run queues so it will run asynchronously. * diff --git a/src/modelRun_ctx.c b/src/modelRun_ctx.c index 3b94b259c..5024bf130 100644 --- a/src/modelRun_ctx.c +++ b/src/modelRun_ctx.c @@ -68,57 +68,3 @@ void RAI_ModelRunCtxFree(RAI_ModelRunCtx *mctx) { } RedisModule_Free(mctx); } - -int RedisAI_Parse_ModelRun_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc, - RAI_ModelRunCtx **mctx, RedisModuleString ***inkeys, - RedisModuleString ***outkeys, RAI_Model **mto, - RAI_Error *error) { - if (argc < 3) { - RAI_SetError(error, RAI_EMODELRUN, - "ERR wrong number of arguments for 'AI.MODELRUN' command"); - return -1; - } - - const char *inputstr = RedisModule_StringPtrLen(argv[2], NULL); - if (strcasecmp(inputstr, "INPUTS") != 0) { - RAI_SetError(error, RAI_EMODELRUN, "ERR INPUTS not specified"); - return -1; - } - - int is_input = 0; - size_t ninputs = 0; - size_t noutputs = 0; - int outputs_flag_count = 0; - size_t argpos = 3; - - for (; argpos <= argc - 1; argpos++) { - const char *arg_string = RedisModule_StringPtrLen(argv[argpos], NULL); - if (!strcasecmp(arg_string, "OUTPUTS") && outputs_flag_count == 0) { - is_input = 1; - outputs_flag_count = 1; - } else { - RAI_HoldString(NULL, argv[argpos]); - if (is_input == 0) { - *inkeys = array_append(*inkeys, argv[argpos]); - ninputs++; - } else { - *outkeys = array_append(*outkeys, argv[argpos]); - noutputs++; - } - } - } - if ((*mto)->inputs && array_len((*mto)->inputs) != ninputs) { - RAI_SetError(error, RAI_EMODELRUN, - "Number of names given as INPUTS during MODELSET and keys given as " - "INPUTS here do not match"); - return -1; - } - - if ((*mto)->outputs && array_len((*mto)->outputs) != noutputs) { - RAI_SetError(error, RAI_EMODELRUN, - "Number of names given as OUTPUTS during MODELSET and keys given as " - "OUTPUTS here do not match"); - return -1; - } - return argpos; -} diff --git a/src/modelRun_ctx.h b/src/modelRun_ctx.h index 3fc62c75b..243a0cd3a 100644 --- a/src/modelRun_ctx.h +++ b/src/modelRun_ctx.h @@ -76,22 +76,3 @@ RAI_Tensor *RAI_ModelRunCtxInputTensor(RAI_ModelRunCtx *mctx, size_t index); * @return RAI_Tensor */ RAI_Tensor *RAI_ModelRunCtxOutputTensor(RAI_ModelRunCtx *mctx, size_t index); - -/** - * Helper method to parse AI.MODELRUN arguments - * - * @param ctx Context in which Redis modules operate - * @param argv Redis command arguments, as an array of strings - * @param argc Redis command number of arguments - * @param mctx Destination Model context to store the parsed data - * @param outkeys array to store the parsed output keys - * @param mto model to run the session from - * @param error error data structure to store error message in the case of - * parsing failures - * @return processed number of arguments on success, or -1 if the parsing failed - */ -// todo: remove this after DAG LLAPI is done. -int RedisAI_Parse_ModelRun_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc, - RAI_ModelRunCtx **mctx, RedisModuleString ***inkeys, - RedisModuleString ***outkeys, RAI_Model **mto, - RAI_Error *error); diff --git a/src/redisai.c b/src/redisai.c index 730b7e4cf..7b1922d68 100644 --- a/src/redisai.c +++ b/src/redisai.c @@ -11,6 +11,8 @@ #include "backends/util.h" #include "background_workers.h" #include "DAG/dag.h" +#include "DAG/dag_builder.h" +#include "DAG/dag_execute.h" #include "model.h" #include "modelRun_ctx.h" #include "script.h" @@ -105,12 +107,14 @@ int RedisAI_TensorSet_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv } RAI_Tensor *t = NULL; - RAI_Error err; - const int parse_result = RAI_parseTensorSetArgs(ctx, argv, argc, &t, 1, &err); + RAI_Error err = {0}; + const int parse_result = RAI_parseTensorSetArgs(argv, argc, &t, 1, &err); // if the number of parsed args is negative something went wrong if (parse_result < 0) { RedisModule_CloseKey(key); + RedisModule_ReplyWithError(ctx, RAI_GetErrorOneLine(&err)); + RAI_ClearError(&err); return REDISMODULE_ERR; } @@ -134,13 +138,17 @@ int RedisAI_TensorGet_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv RAI_Tensor *t; RedisModuleKey *key; - const int status = RAI_GetTensorFromKeyspace(ctx, argv[1], &key, &t, REDISMODULE_READ); + RAI_Error err = {0}; + const int status = RAI_GetTensorFromKeyspace(ctx, argv[1], &key, &t, REDISMODULE_READ, &err); if (status == REDISMODULE_ERR) { + RedisModule_ReplyWithError(ctx, RAI_GetErrorOneLine(&err)); + RAI_ClearError(&err); return REDISMODULE_ERR; } - - uint fmt = ParseTensorGetArgs(ctx, argv, argc); + uint fmt = ParseTensorGetArgs(&err, argv, argc); if (fmt == TENSOR_NONE) { + RedisModule_ReplyWithError(ctx, RAI_GetErrorOneLine(&err)); + RAI_ClearError(&err); // This means that args are invalid. return REDISMODULE_ERR; } @@ -410,12 +418,12 @@ int RedisAI_ModelGet_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, if (argc < 2 || argc > 4) return RedisModule_WrongArity(ctx); + RAI_Error err = {0}; RAI_Model *mto; RedisModuleKey *key; - RAI_Error err = {0}; const int status = RAI_GetModelFromKeyspace(ctx, argv[1], &key, &mto, REDISMODULE_READ, &err); if (status == REDISMODULE_ERR) { - RedisModule_ReplyWithError(ctx, err.detail); + RedisModule_ReplyWithError(ctx, RAI_GetErrorOneLine(&err)); RAI_ClearError(&err); return REDISMODULE_ERR; } @@ -512,13 +520,13 @@ int RedisAI_ModelDel_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, if (argc != 2) return RedisModule_WrongArity(ctx); - RAI_Error err = {0}; RAI_Model *mto; RedisModuleKey *key; + RAI_Error err = {0}; const int status = RAI_GetModelFromKeyspace(ctx, argv[1], &key, &mto, REDISMODULE_READ | REDISMODULE_WRITE, &err); if (status == REDISMODULE_ERR) { - RedisModule_ReplyWithError(ctx, err.detail); + RedisModule_ReplyWithError(ctx, RAI_GetErrorOneLine(&err)); RAI_ClearError(&err); return REDISMODULE_ERR; } @@ -598,8 +606,11 @@ int RedisAI_ScriptGet_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv RAI_Script *sto; RedisModuleKey *key; - const int status = RAI_GetScriptFromKeyspace(ctx, argv[1], &key, &sto, REDISMODULE_READ); + RAI_Error err = {0}; + const int status = RAI_GetScriptFromKeyspace(ctx, argv[1], &key, &sto, REDISMODULE_READ, &err); if (status == REDISMODULE_ERR) { + RedisModule_ReplyWithError(ctx, RAI_GetErrorOneLine(&err)); + RAI_ClearError(&err); return REDISMODULE_ERR; } @@ -646,8 +657,11 @@ int RedisAI_ScriptDel_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv RAI_Script *sto; RedisModuleKey *key; - const int status = RAI_GetScriptFromKeyspace(ctx, argv[1], &key, &sto, REDISMODULE_WRITE); + RAI_Error err = {0}; + const int status = RAI_GetScriptFromKeyspace(ctx, argv[1], &key, &sto, REDISMODULE_WRITE, &err); if (status == REDISMODULE_ERR) { + RedisModule_ReplyWithError(ctx, RAI_GetErrorOneLine(&err)); + RAI_ClearError(&err); return REDISMODULE_ERR; } key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_WRITE); @@ -939,6 +953,7 @@ static int RedisAI_RegisterApi(RedisModuleCtx *ctx) { REGISTER_API(GetError, ctx); REGISTER_API(GetErrorOneLine, ctx); REGISTER_API(GetErrorCode, ctx); + REGISTER_API(CloneError, ctx); REGISTER_API(TensorCreate, ctx); REGISTER_API(TensorCreateByConcatenatingTensors, ctx); @@ -961,6 +976,7 @@ static int RedisAI_RegisterApi(RedisModuleCtx *ctx) { REGISTER_API(ModelCreate, ctx); REGISTER_API(ModelFree, ctx); REGISTER_API(ModelRunCtxCreate, ctx); + REGISTER_API(GetModelFromKeyspace, ctx); REGISTER_API(ModelRunCtxAddInput, ctx); REGISTER_API(ModelRunCtxAddOutput, ctx); REGISTER_API(ModelRunCtxNumOutputs, ctx); @@ -974,6 +990,7 @@ static int RedisAI_RegisterApi(RedisModuleCtx *ctx) { REGISTER_API(GetAsModelRunCtx, ctx); REGISTER_API(ScriptCreate, ctx); + REGISTER_API(GetScriptFromKeyspace, ctx); REGISTER_API(ScriptFree, ctx); REGISTER_API(ScriptRunCtxCreate, ctx); REGISTER_API(ScriptRunCtxAddInput, ctx); @@ -988,6 +1005,25 @@ static int RedisAI_RegisterApi(RedisModuleCtx *ctx) { REGISTER_API(ScriptRunAsync, ctx); REGISTER_API(GetAsScriptRunCtx, ctx); + REGISTER_API(DAGRunCtxCreate, ctx); + REGISTER_API(DAGCreateModelRunOp, ctx); + REGISTER_API(DAGCreateScriptRunOp, ctx); + REGISTER_API(DAGRunOpAddInput, ctx); + REGISTER_API(DAGRunOpAddOutput, ctx); + REGISTER_API(DAGAddRunOp, ctx); + REGISTER_API(DAGLoadTensor, ctx); + REGISTER_API(DAGAddTensorSet, ctx); + REGISTER_API(DAGAddTensorGet, ctx); + REGISTER_API(DAGAddOpsFromString, ctx); + REGISTER_API(DAGNumOps, ctx); + REGISTER_API(DAGRun, ctx); + REGISTER_API(DAGNumOutputs, ctx); + REGISTER_API(DAGOutputTensor, ctx); + REGISTER_API(DAGRunError, ctx); + REGISTER_API(DAGGetError, ctx); + REGISTER_API(DAGRunOpFree, ctx); + REGISTER_API(DAGFree, ctx); + return REDISMODULE_OK; } diff --git a/src/redisai.h b/src/redisai.h index 5ccd429d1..0cd2cffcd 100644 --- a/src/redisai.h +++ b/src/redisai.h @@ -7,6 +7,14 @@ #define REDISAI_LLAPI_VERSION 1 #define MODULE_API_FUNC(x) (*x) +#ifdef REDISAI_EXTERN +#define REDISAI_API extern +#endif + +#ifndef REDISAI_API +#define REDISAI_API +#endif + #ifndef REDISAI_H_INCLUDE typedef struct RAI_Tensor RAI_Tensor; typedef struct RAI_Model RAI_Model; @@ -15,6 +23,7 @@ typedef struct RAI_Script RAI_Script; typedef struct RAI_ModelRunCtx RAI_ModelRunCtx; typedef struct RAI_ScriptRunCtx RAI_ScriptRunCtx; typedef struct RAI_DAGRunCtx RAI_DAGRunCtx; +typedef struct RAI_DAGRunOp RAI_DAGRunOp; typedef struct RAI_Error RAI_Error; typedef struct RAI_ModelOpts RAI_ModelOpts; typedef struct RAI_OnFinishCtx RAI_OnFinishCtx; @@ -30,100 +39,158 @@ typedef void (*RAI_OnFinishCB)(RAI_OnFinishCtx *ctx, void *private_data); #define REDISAI_DEVICE_CPU 0 #define REDISAI_DEVICE_GPU 1 -#define RedisAI_ErrorCode int -#define RedisAI_ErrorCode_OK 0 -#define RedisAI_ErrorCode_EMODELIMPORT 1 -#define RedisAI_ErrorCode_EMODELCONFIGURE 2 -#define RedisAI_ErrorCode_EMODELCREATE 3 -#define RedisAI_ErrorCode_EMODELRUN 4 -#define RedisAI_ErrorCode_EMODELSERIALIZE 5 -#define RedisAI_ErrorCode_EMODELFREE 6 -#define RedisAI_ErrorCode_ESCRIPTIMPORT 7 -#define RedisAI_ErrorCode_ESCRIPTCONFIGURE 8 -#define RedisAI_ErrorCode_ESCRIPTCREATE 9 -#define RedisAI_ErrorCode_ESCRIPTRUN 10 -#define RedisAI_ErrorCode_EUNSUPPORTEDBACKEND 11 -#define RedisAI_ErrorCode_EBACKENDNOTLOADED 12 -#define RedisAI_ErrorCode_ESCRIPTFREE 13 -#define RedisAI_ErrorCode_ETENSORSET 14 -#define RedisAI_ErrorCode_ETENSORGET 15 -#define RedisAI_ErrorCode_EDAGRUN 16 -#define RedisAI_ErrorCode_EFINISHCTX 17 - -int MODULE_API_FUNC(RedisAI_InitError)(RAI_Error **err); -void MODULE_API_FUNC(RedisAI_ClearError)(RAI_Error *err); -void MODULE_API_FUNC(RedisAI_FreeError)(RAI_Error *err); -const char *MODULE_API_FUNC(RedisAI_GetError)(RAI_Error *err); -const char *MODULE_API_FUNC(RedisAI_GetErrorOneLine)(RAI_Error *err); -RedisAI_ErrorCode MODULE_API_FUNC(RedisAI_GetErrorCode)(RAI_Error *err); - -RAI_Tensor *MODULE_API_FUNC(RedisAI_TensorCreate)(const char *dataTypeStr, long long *dims, - int ndims); -RAI_Tensor *MODULE_API_FUNC(RedisAI_TensorCreateByConcatenatingTensors)(RAI_Tensor **ts, - long long n); -RAI_Tensor *MODULE_API_FUNC(RedisAI_TensorCreateBySlicingTensor)(RAI_Tensor *t, long long offset, - long long len); -size_t MODULE_API_FUNC(RedisAI_TensorLength)(RAI_Tensor *t); -size_t MODULE_API_FUNC(RedisAI_TensorDataSize)(RAI_Tensor *t); -size_t MODULE_API_FUNC(RedisAI_TensorDataType)(RAI_Tensor *t); -void MODULE_API_FUNC(RedisAI_TensorFree)(RAI_Tensor *t); -int MODULE_API_FUNC(RedisAI_TensorSetData)(RAI_Tensor *tensor, const char *data, size_t len); -int MODULE_API_FUNC(RedisAI_TensorSetValueFromLongLong)(RAI_Tensor *tensor, long long i, - long long val); -int MODULE_API_FUNC(RedisAI_TensorSetValueFromDouble)(RAI_Tensor *tensor, long long i, double val); -int MODULE_API_FUNC(RedisAI_TensorGetValueAsDouble)(RAI_Tensor *t, long long i, double *val); -int MODULE_API_FUNC(RedisAI_TensorGetValueAsLongLong)(RAI_Tensor *t, long long i, long long *val); -RAI_Tensor *MODULE_API_FUNC(RedisAI_TensorGetShallowCopy)(RAI_Tensor *t); -int MODULE_API_FUNC(RedisAI_TensorNumDims)(RAI_Tensor *t); -long long MODULE_API_FUNC(RedisAI_TensorDim)(RAI_Tensor *t, int dim); -size_t MODULE_API_FUNC(RedisAI_TensorByteSize)(RAI_Tensor *t); -char *MODULE_API_FUNC(RedisAI_TensorData)(RAI_Tensor *t); -RedisModuleType *MODULE_API_FUNC(RedisAI_TensorRedisType)(void); - -RAI_Model *MODULE_API_FUNC(RedisAI_ModelCreate)(int backend, char *devicestr, char *tag, - RAI_ModelOpts opts, size_t ninputs, - const char **inputs, size_t noutputs, - const char **outputs, const char *modeldef, - size_t modellen, RAI_Error *err); -void MODULE_API_FUNC(RedisAI_ModelFree)(RAI_Model *model, RAI_Error *err); -RAI_ModelRunCtx *MODULE_API_FUNC(RedisAI_ModelRunCtxCreate)(RAI_Model *model); -int MODULE_API_FUNC(RedisAI_ModelRunCtxAddInput)(RAI_ModelRunCtx *mctx, const char *inputName, - RAI_Tensor *inputTensor); -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); -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, - RAI_Error *err); -RedisModuleType *MODULE_API_FUNC(RedisAI_ModelRedisType)(void); -int MODULE_API_FUNC(RedisAI_ModelRunAsync)(RAI_ModelRunCtx *mctxs, RAI_OnFinishCB DAGAsyncFinish, - void *private_data); -RAI_ModelRunCtx *MODULE_API_FUNC(RedisAI_GetAsModelRunCtx)(RAI_OnFinishCtx *ctx, RAI_Error *err); - -RAI_Script *MODULE_API_FUNC(RedisAI_ScriptCreate)(char *devicestr, char *tag, const char *scriptdef, - RAI_Error *err); -void MODULE_API_FUNC(RedisAI_ScriptFree)(RAI_Script *script, RAI_Error *err); -RAI_ScriptRunCtx *MODULE_API_FUNC(RedisAI_ScriptRunCtxCreate)(RAI_Script *script, - const char *fnname); -int MODULE_API_FUNC(RedisAI_ScriptRunCtxAddInput)(RAI_ScriptRunCtx *sctx, RAI_Tensor *inputTensor, +typedef enum RedisAI_ErrorCode { + RedisAI_ErrorCode_OK = 0, + RedisAI_ErrorCode_EMODELIMPORT, + RedisAI_ErrorCode_EMODELCONFIGURE, + RedisAI_ErrorCode_EMODELCREATE, + RedisAI_ErrorCode_EMODELRUN, + RedisAI_ErrorCode_EMODELSERIALIZE, + RedisAI_ErrorCode_EMODELFREE, + RedisAI_ErrorCode_ESCRIPTIMPORT, + RedisAI_ErrorCode_ESCRIPTCONFIGURE, + RedisAI_ErrorCode_ESCRIPTCREATE, + RedisAI_ErrorCode_ESCRIPTRUN, + RedisAI_ErrorCode_EUNSUPPORTEDBACKEND, + RedisAI_ErrorCode_EBACKENDNOTLOADED, + RedisAI_ErrorCode_ESCRIPTFREE, + RedisAI_ErrorCode_ETENSORSET, + RedisAI_ErrorCode_ETENSORGET, + RedisAI_ErrorCode_EDAGBUILDER, + RedisAI_ErrorCode_EDAGRUN, + RedisAI_ErrorCode_EFINISHCTX +} RedisAI_ErrorCode; + +REDISAI_API int MODULE_API_FUNC(RedisAI_InitError)(RAI_Error **err); +REDISAI_API void MODULE_API_FUNC(RedisAI_ClearError)(RAI_Error *err); +REDISAI_API void MODULE_API_FUNC(RedisAI_FreeError)(RAI_Error *err); +REDISAI_API const char *MODULE_API_FUNC(RedisAI_GetError)(RAI_Error *err); +REDISAI_API const char *MODULE_API_FUNC(RedisAI_GetErrorOneLine)(RAI_Error *err); +REDISAI_API RedisAI_ErrorCode MODULE_API_FUNC(RedisAI_GetErrorCode)(RAI_Error *err); +REDISAI_API void MODULE_API_FUNC(RedisAI_CloneError)(RAI_Error *dest, const RAI_Error *src); + +REDISAI_API RAI_Tensor *MODULE_API_FUNC(RedisAI_TensorCreate)(const char *dataTypeStr, + long long *dims, int ndims); +REDISAI_API RAI_Tensor *MODULE_API_FUNC(RedisAI_TensorCreateByConcatenatingTensors)(RAI_Tensor **ts, + long long n); +REDISAI_API RAI_Tensor *MODULE_API_FUNC(RedisAI_TensorCreateBySlicingTensor)(RAI_Tensor *t, + long long offset, + long long len); +REDISAI_API size_t MODULE_API_FUNC(RedisAI_TensorLength)(RAI_Tensor *t); +REDISAI_API size_t MODULE_API_FUNC(RedisAI_TensorDataSize)(RAI_Tensor *t); +REDISAI_API size_t MODULE_API_FUNC(RedisAI_TensorDataType)(RAI_Tensor *t); +REDISAI_API void MODULE_API_FUNC(RedisAI_TensorFree)(RAI_Tensor *t); +REDISAI_API int MODULE_API_FUNC(RedisAI_TensorSetData)(RAI_Tensor *tensor, const char *data, + size_t len); +REDISAI_API int MODULE_API_FUNC(RedisAI_TensorSetValueFromLongLong)(RAI_Tensor *tensor, long long i, + long long val); +REDISAI_API int MODULE_API_FUNC(RedisAI_TensorSetValueFromDouble)(RAI_Tensor *tensor, long long i, + double val); +REDISAI_API int MODULE_API_FUNC(RedisAI_TensorGetValueAsDouble)(RAI_Tensor *t, long long i, + double *val); +REDISAI_API int MODULE_API_FUNC(RedisAI_TensorGetValueAsLongLong)(RAI_Tensor *t, long long i, + long long *val); +REDISAI_API RAI_Tensor *MODULE_API_FUNC(RedisAI_TensorGetShallowCopy)(RAI_Tensor *t); +REDISAI_API int MODULE_API_FUNC(RedisAI_TensorNumDims)(RAI_Tensor *t); +REDISAI_API long long MODULE_API_FUNC(RedisAI_TensorDim)(RAI_Tensor *t, int dim); +REDISAI_API size_t MODULE_API_FUNC(RedisAI_TensorByteSize)(RAI_Tensor *t); +REDISAI_API char *MODULE_API_FUNC(RedisAI_TensorData)(RAI_Tensor *t); +REDISAI_API RedisModuleType *MODULE_API_FUNC(RedisAI_TensorRedisType)(void); + +REDISAI_API RAI_Model *MODULE_API_FUNC(RedisAI_ModelCreate)(int backend, char *devicestr, char *tag, + RAI_ModelOpts opts, size_t ninputs, + const char **inputs, size_t noutputs, + const char **outputs, + const char *modeldef, size_t modellen, + RAI_Error *err); +REDISAI_API void MODULE_API_FUNC(RedisAI_ModelFree)(RAI_Model *model, RAI_Error *err); +REDISAI_API RAI_ModelRunCtx *MODULE_API_FUNC(RedisAI_ModelRunCtxCreate)(RAI_Model *model); +REDISAI_API int MODULE_API_FUNC(RedisAI_GetModelFromKeyspace)(RedisModuleCtx *ctx, + RedisModuleString *keyName, + RedisModuleKey **key, + RAI_Model **model, int mode, + RAI_Error *err); +REDISAI_API int MODULE_API_FUNC(RedisAI_ModelRunCtxAddInput)(RAI_ModelRunCtx *mctx, + const char *inputName, + RAI_Tensor *inputTensor); +REDISAI_API int MODULE_API_FUNC(RedisAI_ModelRunCtxAddOutput)(RAI_ModelRunCtx *mctx, + const char *outputName); +REDISAI_API size_t MODULE_API_FUNC(RedisAI_ModelRunCtxNumOutputs)(RAI_ModelRunCtx *mctx); +REDISAI_API RAI_Tensor *MODULE_API_FUNC(RedisAI_ModelRunCtxOutputTensor)(RAI_ModelRunCtx *mctx, + size_t index); +REDISAI_API void MODULE_API_FUNC(RedisAI_ModelRunCtxFree)(RAI_ModelRunCtx *mctx); +REDISAI_API int MODULE_API_FUNC(RedisAI_ModelRun)(RAI_ModelRunCtx **mctx, long long n, RAI_Error *err); -int MODULE_API_FUNC(RedisAI_ScriptRunCtxAddInputList)(RAI_ScriptRunCtx *sctx, - RAI_Tensor **inputTensors, size_t len, - RAI_Error *err); -int MODULE_API_FUNC(RedisAI_ScriptRunCtxAddOutput)(RAI_ScriptRunCtx *sctx); -size_t MODULE_API_FUNC(RedisAI_ScriptRunCtxNumOutputs)(RAI_ScriptRunCtx *sctx); -RAI_Tensor *MODULE_API_FUNC(RedisAI_ScriptRunCtxOutputTensor)(RAI_ScriptRunCtx *sctx, size_t index); -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)(); +REDISAI_API RAI_Model *MODULE_API_FUNC(RedisAI_ModelGetShallowCopy)(RAI_Model *model); +REDISAI_API int MODULE_API_FUNC(RedisAI_ModelSerialize)(RAI_Model *model, char **buffer, + size_t *len, RAI_Error *err); +REDISAI_API RedisModuleType *MODULE_API_FUNC(RedisAI_ModelRedisType)(void); +REDISAI_API int MODULE_API_FUNC(RedisAI_ModelRunAsync)(RAI_ModelRunCtx *mctxs, + RAI_OnFinishCB DAGAsyncFinish, + void *private_data); +REDISAI_API RAI_ModelRunCtx *MODULE_API_FUNC(RedisAI_GetAsModelRunCtx)(RAI_OnFinishCtx *ctx, + RAI_Error *err); + +REDISAI_API RAI_Script *MODULE_API_FUNC(RedisAI_ScriptCreate)(char *devicestr, char *tag, + const char *scriptdef, + RAI_Error *err); +REDISAI_API int MODULE_API_FUNC(RedisAI_GetScriptFromKeyspace)(RedisModuleCtx *ctx, + RedisModuleString *keyName, + RedisModuleKey **key, + RAI_Script **script, int mode, + RAI_Error *err); +REDISAI_API void MODULE_API_FUNC(RedisAI_ScriptFree)(RAI_Script *script, RAI_Error *err); +REDISAI_API RAI_ScriptRunCtx *MODULE_API_FUNC(RedisAI_ScriptRunCtxCreate)(RAI_Script *script, + const char *fnname); +REDISAI_API int MODULE_API_FUNC(RedisAI_ScriptRunCtxAddInput)(RAI_ScriptRunCtx *sctx, + RAI_Tensor *inputTensor, + RAI_Error *err); +REDISAI_API int MODULE_API_FUNC(RedisAI_ScriptRunCtxAddInputList)(RAI_ScriptRunCtx *sctx, + RAI_Tensor **inputTensors, + size_t len, RAI_Error *err); +REDISAI_API int MODULE_API_FUNC(RedisAI_ScriptRunCtxAddOutput)(RAI_ScriptRunCtx *sctx); +REDISAI_API size_t MODULE_API_FUNC(RedisAI_ScriptRunCtxNumOutputs)(RAI_ScriptRunCtx *sctx); +REDISAI_API RAI_Tensor *MODULE_API_FUNC(RedisAI_ScriptRunCtxOutputTensor)(RAI_ScriptRunCtx *sctx, + size_t index); +REDISAI_API void MODULE_API_FUNC(RedisAI_ScriptRunCtxFree)(RAI_ScriptRunCtx *sctx); +REDISAI_API int MODULE_API_FUNC(RedisAI_ScriptRun)(RAI_ScriptRunCtx *sctx, RAI_Error *err); +REDISAI_API RAI_Script *MODULE_API_FUNC(RedisAI_ScriptGetShallowCopy)(RAI_Script *script); +REDISAI_API RedisModuleType *MODULE_API_FUNC(RedisAI_ScriptRedisType)(void); +REDISAI_API int MODULE_API_FUNC(RedisAI_ScriptRunAsync)(RAI_ScriptRunCtx *sctx, + RAI_OnFinishCB DAGAsyncFinish, + void *private_data); +REDISAI_API RAI_ScriptRunCtx *MODULE_API_FUNC(RedisAI_GetAsScriptRunCtx)(RAI_OnFinishCtx *ctx, + RAI_Error *err); + +REDISAI_API RAI_DAGRunCtx *MODULE_API_FUNC(RedisAI_DAGRunCtxCreate)(void); +REDISAI_API RAI_DAGRunOp *MODULE_API_FUNC(RedisAI_DAGCreateModelRunOp)(RAI_Model *model); +REDISAI_API RAI_DAGRunOp *MODULE_API_FUNC(RedisAI_DAGCreateScriptRunOp)(RAI_Script *script, + const char *func_name); +REDISAI_API int MODULE_API_FUNC(RedisAI_DAGRunOpAddInput)(RAI_DAGRunOp *DAGOp, const char *input); +REDISAI_API int MODULE_API_FUNC(RedisAI_DAGRunOpAddOutput)(RAI_DAGRunOp *DAGOp, const char *output); +REDISAI_API int MODULE_API_FUNC(RedisAI_DAGAddRunOp)(RAI_DAGRunCtx *run_info, RAI_DAGRunOp *DAGop, + RAI_Error *err); +REDISAI_API int MODULE_API_FUNC(RedisAI_DAGLoadTensor)(RAI_DAGRunCtx *run_info, const char *t_name, + RAI_Tensor *tensor); +REDISAI_API int MODULE_API_FUNC(RedisAI_DAGAddTensorSet)(RAI_DAGRunCtx *run_info, + const char *t_name, RAI_Tensor *tensor); +REDISAI_API int MODULE_API_FUNC(RedisAI_DAGAddTensorGet)(RAI_DAGRunCtx *run_info, + const char *t_name, RAI_Error *err); +REDISAI_API int MODULE_API_FUNC(RedisAI_DAGAddOpsFromString)(RAI_DAGRunCtx *run_info, + const char *dag, RAI_Error *err); +REDISAI_API size_t MODULE_API_FUNC(RedisAI_DAGNumOps)(RAI_DAGRunCtx *run_info); +REDISAI_API int MODULE_API_FUNC(RedisAI_DAGRun)(RAI_DAGRunCtx *run_info, + RAI_OnFinishCB DAGAsyncFinish, void *private_data, + RAI_Error *err); +REDISAI_API size_t MODULE_API_FUNC(RedisAI_DAGNumOutputs)(RAI_OnFinishCtx *finish_ctx); +REDISAI_API const RAI_Tensor *MODULE_API_FUNC(RedisAI_DAGOutputTensor)(RAI_OnFinishCtx *finish_ctx, + size_t index); +REDISAI_API int MODULE_API_FUNC(RedisAI_DAGRunError)(RAI_OnFinishCtx *finish_ctx); +REDISAI_API const RAI_Error *MODULE_API_FUNC(RedisAI_DAGGetError)(RAI_OnFinishCtx *finish_ctx); +REDISAI_API void MODULE_API_FUNC(RedisAI_DAGRunOpFree)(RAI_DAGRunOp *dagOp); +REDISAI_API void MODULE_API_FUNC(RedisAI_DAGFree)(RAI_DAGRunCtx *run_info); + +REDISAI_API int MODULE_API_FUNC(RedisAI_GetLLAPIVersion)(); #ifndef __cplusplus #define REDISAI_MODULE_INIT_FUNCTION(ctx, name) \ @@ -157,6 +224,7 @@ static int RedisAI_Initialize(RedisModuleCtx *ctx) { REDISAI_MODULE_INIT_FUNCTION(ctx, GetError); REDISAI_MODULE_INIT_FUNCTION(ctx, GetErrorOneLine); REDISAI_MODULE_INIT_FUNCTION(ctx, GetErrorCode); + REDISAI_MODULE_INIT_FUNCTION(ctx, CloneError); REDISAI_MODULE_INIT_FUNCTION(ctx, GetLLAPIVersion); @@ -181,6 +249,7 @@ static int RedisAI_Initialize(RedisModuleCtx *ctx) { REDISAI_MODULE_INIT_FUNCTION(ctx, ModelCreate); REDISAI_MODULE_INIT_FUNCTION(ctx, ModelFree); REDISAI_MODULE_INIT_FUNCTION(ctx, ModelRunCtxCreate); + REDISAI_MODULE_INIT_FUNCTION(ctx, GetModelFromKeyspace); REDISAI_MODULE_INIT_FUNCTION(ctx, ModelRunCtxAddInput); REDISAI_MODULE_INIT_FUNCTION(ctx, ModelRunCtxAddOutput); REDISAI_MODULE_INIT_FUNCTION(ctx, ModelRunCtxNumOutputs); @@ -194,6 +263,7 @@ static int RedisAI_Initialize(RedisModuleCtx *ctx) { REDISAI_MODULE_INIT_FUNCTION(ctx, GetAsModelRunCtx); REDISAI_MODULE_INIT_FUNCTION(ctx, ScriptCreate); + REDISAI_MODULE_INIT_FUNCTION(ctx, GetScriptFromKeyspace); REDISAI_MODULE_INIT_FUNCTION(ctx, ScriptFree); REDISAI_MODULE_INIT_FUNCTION(ctx, ScriptRunCtxCreate); REDISAI_MODULE_INIT_FUNCTION(ctx, ScriptRunCtxAddInput); @@ -208,6 +278,25 @@ static int RedisAI_Initialize(RedisModuleCtx *ctx) { REDISAI_MODULE_INIT_FUNCTION(ctx, ScriptRunAsync); REDISAI_MODULE_INIT_FUNCTION(ctx, GetAsScriptRunCtx); + REDISAI_MODULE_INIT_FUNCTION(ctx, DAGRunCtxCreate); + REDISAI_MODULE_INIT_FUNCTION(ctx, DAGCreateModelRunOp); + REDISAI_MODULE_INIT_FUNCTION(ctx, DAGCreateScriptRunOp); + REDISAI_MODULE_INIT_FUNCTION(ctx, DAGRunOpAddInput); + REDISAI_MODULE_INIT_FUNCTION(ctx, DAGRunOpAddOutput); + REDISAI_MODULE_INIT_FUNCTION(ctx, DAGAddRunOp); + REDISAI_MODULE_INIT_FUNCTION(ctx, DAGLoadTensor); + REDISAI_MODULE_INIT_FUNCTION(ctx, DAGAddTensorSet); + REDISAI_MODULE_INIT_FUNCTION(ctx, DAGAddTensorGet); + REDISAI_MODULE_INIT_FUNCTION(ctx, DAGAddOpsFromString); + REDISAI_MODULE_INIT_FUNCTION(ctx, DAGNumOps); + REDISAI_MODULE_INIT_FUNCTION(ctx, DAGRun); + REDISAI_MODULE_INIT_FUNCTION(ctx, DAGNumOutputs); + REDISAI_MODULE_INIT_FUNCTION(ctx, DAGOutputTensor); + REDISAI_MODULE_INIT_FUNCTION(ctx, DAGRunError); + REDISAI_MODULE_INIT_FUNCTION(ctx, DAGGetError); + REDISAI_MODULE_INIT_FUNCTION(ctx, DAGRunOpFree); + REDISAI_MODULE_INIT_FUNCTION(ctx, DAGFree); + if (RedisAI_GetLLAPIVersion() < REDISAI_LLAPI_VERSION) { return REDISMODULE_ERR; } diff --git a/src/run_info.c b/src/run_info.c index 405c813c1..ac6f4ddf3 100644 --- a/src/run_info.c +++ b/src/run_info.c @@ -179,7 +179,8 @@ void RAI_ContextReadLock(RedisAI_RunInfo *rinfo) { if (rinfo->single_op_dag || rinfo->single_device_dag) { return; } - pthread_rwlock_rdlock(rinfo->dagLock); + // This is a temporary solution + pthread_rwlock_wrlock(rinfo->dagLock); } void RAI_ContextWriteLock(RedisAI_RunInfo *rinfo) { diff --git a/src/script.c b/src/script.c index 5e79ad9f3..5315ccfe0 100644 --- a/src/script.c +++ b/src/script.c @@ -155,16 +155,16 @@ RAI_Script *RAI_ScriptGetShallowCopy(RAI_Script *script) { * Return REDISMODULE_OK if the model value stored at key was correctly * returned and available at *model variable. */ int RAI_GetScriptFromKeyspace(RedisModuleCtx *ctx, RedisModuleString *keyName, RedisModuleKey **key, - RAI_Script **script, int mode) { + RAI_Script **script, int mode, RAI_Error *err) { *key = RedisModule_OpenKey(ctx, keyName, mode); if (RedisModule_KeyType(*key) == REDISMODULE_KEYTYPE_EMPTY) { RedisModule_CloseKey(*key); - RedisModule_ReplyWithError(ctx, "ERR script key is empty"); + RAI_SetError(err, RAI_ESCRIPTRUN, "ERR script key is empty"); return REDISMODULE_ERR; } if (RedisModule_ModuleTypeGetType(*key) != RedisAI_ScriptType) { RedisModule_CloseKey(*key); - RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE); + RAI_SetError(err, RAI_ESCRIPTRUN, REDISMODULE_ERRORMSG_WRONGTYPE); return REDISMODULE_ERR; } *script = RedisModule_ModuleTypeGetValue(*key); @@ -197,69 +197,6 @@ int RedisAI_ScriptRun_IsKeysPositionRequest_ReportKeys(RedisModuleCtx *ctx, return REDISMODULE_OK; } -/** - * AI.SCRIPTRUN INPUTS [input ...] OUTPUTS - * [output ...] - */ -int RedisAI_Parse_ScriptRun_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc, - RedisModuleString ***inkeys, RedisModuleString ***outkeys, - int *variadic, RAI_Error *error) { - if (argc < 6) { - RAI_SetError(error, RAI_ESCRIPTRUN, - "ERR wrong number of arguments for 'AI.SCRIPTRUN' command"); - return -1; - } - - const char *inputstr = RedisModule_StringPtrLen(argv[2], NULL); - if (strcasecmp(inputstr, "INPUTS") == 0) { - RAI_SetError(error, RAI_ESCRIPTRUN, "ERR function name not specified"); - return -1; - } - - inputstr = RedisModule_StringPtrLen(argv[3], NULL); - if (strcasecmp(inputstr, "INPUTS")) { - RAI_SetError(error, RAI_ESCRIPTRUN, "ERR INPUTS not specified"); - return -1; - } - - // parsing aux vars - int is_output = 0; - int outputs_flag_count = 0; - int variadic_ = *variadic; - size_t argpos = 4; - for (; argpos <= argc - 1; argpos++) { - const char *arg_string = RedisModule_StringPtrLen(argv[argpos], NULL); - if (!arg_string) { - RAI_SetError(error, RAI_ESCRIPTRUN, "ERR NULL argument on SCRIPTRUN"); - return -1; - } - if (!strcasecmp(arg_string, "OUTPUTS") && outputs_flag_count == 0) { - is_output = 1; - outputs_flag_count = 1; - } else { - if (!strcasecmp(arg_string, "$")) { - if (variadic_ > -1) { - RAI_SetError(error, RAI_ESCRIPTRUN, - "ERR Already encountered a variable size list of tensors"); - return -1; - } - variadic_ = argpos - 4; - continue; - } - RedisModule_RetainString(ctx, argv[argpos]); - if (is_output == 0) { - *inkeys = array_append(*inkeys, argv[argpos]); - } else { - *outkeys = array_append(*outkeys, argv[argpos]); - } - } - } - - *variadic = variadic_; - - return argpos; -} - RedisModuleType *RAI_ScriptRedisType(void) { return RedisAI_ScriptType; } int RAI_ScriptRunAsync(RAI_ScriptRunCtx *sctx, RAI_OnFinishCB ScriptAsyncFinish, @@ -281,5 +218,9 @@ int RAI_ScriptRunAsync(RAI_ScriptRunCtx *sctx, RAI_OnFinishCB ScriptAsyncFinish, rinfo->dagOps = array_append(rinfo->dagOps, op); rinfo->dagOpCount = 1; - return DAG_InsertDAGToQueue(rinfo); + if (DAG_InsertDAGToQueue(rinfo) != REDISMODULE_OK) { + RAI_FreeRunInfo(rinfo); + return REDISMODULE_ERR; + } + return REDISMODULE_OK; } diff --git a/src/script.h b/src/script.h index 33da8ef87..d1b236d16 100644 --- a/src/script.h +++ b/src/script.h @@ -162,7 +162,7 @@ RAI_Script *RAI_ScriptGetShallowCopy(RAI_Script *script); * an error getting the Script */ int RAI_GetScriptFromKeyspace(RedisModuleCtx *ctx, RedisModuleString *keyName, RedisModuleKey **key, - RAI_Script **script, int mode); + RAI_Script **script, int mode, RAI_Error *err); /** * When a module command is called in order to obtain the position of @@ -179,23 +179,6 @@ int RAI_GetScriptFromKeyspace(RedisModuleCtx *ctx, RedisModuleString *keyName, R int RedisAI_ScriptRun_IsKeysPositionRequest_ReportKeys(RedisModuleCtx *ctx, RedisModuleString **argv, int argc); -/** - * Helper method to parse AI.SCRIPTRUN arguments - * - * @param ctx Context in which Redis modules operate - * @param argv Redis command arguments, as an array of strings - * @param argc Redis command number of arguments - * @param outkeys array to store the parsed input keys - * @param outkeys array to store the parsed output keys - * @param variadic int to store the variadic input location - * @param error error data structure to store error message in the case of - * parsing failures - * @return processed number of arguments on success, or -1 if the parsing failed - */ -int RedisAI_Parse_ScriptRun_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc, - RedisModuleString ***inkeys, RedisModuleString ***outkeys, - int *variadic, RAI_Error *error); - #if 0 /** * Helper method to reply if the ctx is not NULL or fallback and set the error in the RAI_Error structure diff --git a/src/tensor.c b/src/tensor.c index 2f6c2e293..fd4682dbe 100644 --- a/src/tensor.c +++ b/src/tensor.c @@ -486,12 +486,12 @@ int RAI_TensorGetValueAsDouble(RAI_Tensor *t, long long i, double *val) { *val = ((double *)data)[i]; break; default: - return 1; + return REDISMODULE_ERR; } } else { - return 1; + return REDISMODULE_ERR; } - return 0; + return REDISMODULE_OK; } int RAI_TensorGetValueAsLongLong(RAI_Tensor *t, long long i, long long *val) { @@ -515,7 +515,7 @@ int RAI_TensorGetValueAsLongLong(RAI_Tensor *t, long long i, long long *val) { *val = ((int64_t *)data)[i]; break; default: - return 0; + return REDISMODULE_ERR; } } else if (dtype.code == kDLUInt) { switch (dtype.bits) { @@ -532,12 +532,12 @@ int RAI_TensorGetValueAsLongLong(RAI_Tensor *t, long long i, long long *val) { *val = ((uint64_t *)data)[i]; break; default: - return 0; + return REDISMODULE_ERR; } } else { - return 0; + return REDISMODULE_ERR; } - return 1; + return REDISMODULE_OK; } RAI_Tensor *RAI_TensorGetShallowCopy(RAI_Tensor *t) { @@ -573,20 +573,17 @@ int RAI_OpenKey_Tensor(RedisModuleCtx *ctx, RedisModuleString *keyName, RedisMod return REDISMODULE_OK; } -/* Return REDISMODULE_ERR if there was an error getting the Tensor. - * Return REDISMODULE_OK if the tensor value stored at key was correctly - * returned and available at *tensor variable. */ int RAI_GetTensorFromKeyspace(RedisModuleCtx *ctx, RedisModuleString *keyName, RedisModuleKey **key, - RAI_Tensor **tensor, int mode) { + RAI_Tensor **tensor, int mode, RAI_Error *err) { *key = RedisModule_OpenKey(ctx, keyName, mode); if (RedisModule_KeyType(*key) == REDISMODULE_KEYTYPE_EMPTY) { RedisModule_CloseKey(*key); - RedisModule_ReplyWithError(ctx, "ERR tensor key is empty"); + RAI_SetError(err, RAI_ETENSORGET, "ERR tensor key is empty"); return REDISMODULE_ERR; } if (RedisModule_ModuleTypeGetType(*key) != RedisAI_TensorType) { RedisModule_CloseKey(*key); - RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE); + RAI_SetError(err, RAI_ETENSORGET, REDISMODULE_ERRORMSG_WRONGTYPE); return REDISMODULE_ERR; } *tensor = RedisModule_ModuleTypeGetValue(*key); @@ -597,20 +594,15 @@ int RAI_GetTensorFromKeyspace(RedisModuleCtx *ctx, RedisModuleString *keyName, R /* Return REDISMODULE_ERR if there was an error getting the Tensor. * Return REDISMODULE_OK if the tensor value is present at the localContextDict. */ -int RAI_getTensorFromLocalContext(RedisModuleCtx *ctx, AI_dict *localContextDict, - RedisModuleString *localContextKey, RAI_Tensor **tensor, - RAI_Error *error) { +int RAI_getTensorFromLocalContext(AI_dict *localContextDict, RedisModuleString *localContextKey, + RAI_Tensor **tensor, RAI_Error *error) { int result = REDISMODULE_ERR; AI_dictEntry *tensor_entry = AI_dictFind(localContextDict, localContextKey); if (tensor_entry) { *tensor = AI_dictGetVal(tensor_entry); result = REDISMODULE_OK; } else { - if (ctx == NULL) { - RAI_SetError(error, RAI_ETENSORGET, "ERR tensor key is empty"); - } else { - RedisModule_ReplyWithError(ctx, "ERR tensor key is empty"); - } + RAI_SetError(error, RAI_ETENSORGET, "ERR tensor key is empty"); } return result; } @@ -639,21 +631,17 @@ void RedisAI_ReplicateTensorSet(RedisModuleCtx *ctx, RedisModuleString *key, RAI } } -int RAI_parseTensorSetArgs(RedisModuleCtx *ctx, RedisModuleString **argv, int argc, RAI_Tensor **t, - int enforceArity, RAI_Error *error) { +int RAI_parseTensorSetArgs(RedisModuleString **argv, int argc, RAI_Tensor **t, int enforceArity, + RAI_Error *error) { if (argc < 4) { - RedisModule_WrongArity(ctx); + RAI_SetError(error, RAI_ETENSORSET, "wrong number of arguments for 'AI.TENSORSET' command"); return -1; } // get the tensor datatype const char *typestr = RedisModule_StringPtrLen(argv[2], NULL); size_t datasize = RAI_TensorDataSizeFromString(typestr); if (!datasize) { - if (ctx == NULL) { - RAI_SetError(error, RAI_ETENSORSET, "ERR invalid data type"); - } else { - RedisModule_ReplyWithError(ctx, "ERR invalid data type"); - } + RAI_SetError(error, RAI_ETENSORSET, "ERR invalid data type"); return -1; } const char *fmtstr; @@ -677,12 +665,8 @@ int RAI_parseTensorSetArgs(RedisModuleCtx *ctx, RedisModuleString **argv, int ar // check right away if the arity is correct if (remaining_args != 1 && enforceArity == 1) { array_free(dims); - if (ctx == NULL) { - RAI_SetError(error, RAI_ETENSORSET, - "ERR wrong number of arguments for 'AI.TENSORSET' command"); - } else { - RedisModule_WrongArity(ctx); - } + RAI_SetError(error, RAI_ETENSORSET, + "ERR wrong number of arguments for 'AI.TENSORSET' command"); return -1; } argpos++; @@ -694,12 +678,8 @@ int RAI_parseTensorSetArgs(RedisModuleCtx *ctx, RedisModuleString **argv, int ar // check right away if the arity is correct if (remaining_args != len && enforceArity == 1) { array_free(dims); - if (ctx == NULL) { - RAI_SetError(error, RAI_ETENSORSET, - "ERR wrong number of arguments for 'AI.TENSORSET' command"); - } else { - RedisModule_WrongArity(ctx); - } + RAI_SetError(error, RAI_ETENSORSET, + "ERR wrong number of arguments for 'AI.TENSORSET' command"); return -1; } argpos++; @@ -709,13 +689,8 @@ int RAI_parseTensorSetArgs(RedisModuleCtx *ctx, RedisModuleString **argv, int ar const int retval = RedisModule_StringToLongLong(argv[argpos], &dimension); if (retval != REDISMODULE_OK || dimension <= 0) { array_free(dims); - if (ctx == NULL) { - RAI_SetError(error, RAI_ETENSORSET, - "ERR invalid or negative value found in tensor shape"); - } else { - RedisModule_ReplyWithError( - ctx, "ERR invalid or negative value found in tensor shape"); - } + RAI_SetError(error, RAI_ETENSORSET, + "ERR invalid or negative value found in tensor shape"); return -1; } ndims++; @@ -738,11 +713,7 @@ int RAI_parseTensorSetArgs(RedisModuleCtx *ctx, RedisModuleString **argv, int ar if (!t) { array_free(dims); - if (ctx == NULL) { - RAI_SetError(error, RAI_ETENSORSET, "ERR could not create tensor"); - } else { - RedisModule_ReplyWithError(ctx, "ERR could not create tensor"); - } + RAI_SetError(error, RAI_ETENSORSET, "ERR could not create tensor"); return -1; } long i = 0; @@ -754,24 +725,15 @@ int RAI_parseTensorSetArgs(RedisModuleCtx *ctx, RedisModuleString **argv, int ar if (retval != REDISMODULE_OK) { RAI_TensorFree(*t); array_free(dims); - if (ctx == NULL) { - RAI_SetError(error, RAI_ETENSORSET, "ERR invalid value"); - } else { - RedisModule_ReplyWithError(ctx, "ERR invalid value"); - } + RAI_SetError(error, RAI_ETENSORSET, "ERR invalid value"); return -1; } const int retset = RAI_TensorSetValueFromDouble(*t, i, val); if (retset == -1) { RAI_TensorFree(*t); array_free(dims); - if (ctx == NULL) { - RAI_SetError(error, RAI_ETENSORSET, - "ERR cannot specify values for this datatype"); - } else { - RedisModule_ReplyWithError(ctx, - "ERR cannot specify values for this datatype"); - } + RAI_SetError(error, RAI_ETENSORSET, + "ERR cannot specify values for this datatype"); return -1; } } else { @@ -780,24 +742,15 @@ int RAI_parseTensorSetArgs(RedisModuleCtx *ctx, RedisModuleString **argv, int ar if (retval != REDISMODULE_OK) { RAI_TensorFree(*t); array_free(dims); - if (ctx == NULL) { - RAI_SetError(error, RAI_ETENSORSET, "ERR invalid value"); - } else { - RedisModule_ReplyWithError(ctx, "ERR invalid value"); - } + RAI_SetError(error, RAI_ETENSORSET, "ERR invalid value"); return -1; } const int retset = RAI_TensorSetValueFromLongLong(*t, i, val); if (retset == -1) { RAI_TensorFree(*t); array_free(dims); - if (ctx == NULL) { - RAI_SetError(error, RAI_ETENSORSET, - "ERR cannot specify values for this datatype"); - } else { - RedisModule_ReplyWithError(ctx, - "ERR cannot specify values for this datatype"); - } + RAI_SetError(error, RAI_ETENSORSET, + "ERR cannot specify values for this datatype"); return -1; } } @@ -824,7 +777,7 @@ int RAI_TensorReplyWithValues(RedisModuleCtx *ctx, RAI_Tensor *t) { double val; for (i = 0; i < len; i++) { int ret = RAI_TensorGetValueAsDouble(t, i, &val); - if (ret == 1) { + if (ret == REDISMODULE_ERR) { RedisModule_ReplyWithError(ctx, "ERR cannot get values for this datatype"); return -1; } @@ -834,7 +787,7 @@ int RAI_TensorReplyWithValues(RedisModuleCtx *ctx, RAI_Tensor *t) { long long val; for (i = 0; i < len; i++) { int ret = RAI_TensorGetValueAsLongLong(t, i, &val); - if (!ret) { + if (ret == REDISMODULE_ERR) { RedisModule_ReplyWithError(ctx, "ERR cannot get values for this datatype"); return -1; } @@ -845,10 +798,10 @@ int RAI_TensorReplyWithValues(RedisModuleCtx *ctx, RAI_Tensor *t) { return 0; } -uint ParseTensorGetArgs(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { +uint ParseTensorGetArgs(RAI_Error *err, RedisModuleString **argv, int argc) { uint fmt = TENSOR_NONE; if (argc < 2 || argc > 4) { - RedisModule_WrongArity(ctx); + RAI_SetError(err, RAI_EDAGBUILDER, "wrong number of arguments for 'AI.TENSORGET' command"); return fmt; } if (argc == 2) { @@ -863,13 +816,13 @@ uint ParseTensorGetArgs(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) } else if (!strcasecmp(fmtstr, "META")) { fmt |= TENSOR_META; } else { - RedisModule_ReplyWithError(ctx, "ERR unsupported data format"); + RAI_SetError(err, RAI_EDAGBUILDER, "ERR unsupported data format"); return TENSOR_NONE; } } if (fmt == TENSOR_ILLEGAL_VALUES_BLOB) { - RedisModule_ReplyWithError(ctx, "ERR both BLOB and VALUES specified"); + RAI_SetError(err, RAI_EDAGBUILDER, "ERR both BLOB and VALUES specified"); return TENSOR_NONE; } return fmt; diff --git a/src/tensor.h b/src/tensor.h index 6ee6e4cd0..e5587dc40 100644 --- a/src/tensor.h +++ b/src/tensor.h @@ -315,8 +315,8 @@ int RAI_OpenKey_Tensor(RedisModuleCtx *ctx, RedisModuleString *keyName, RedisMod int mode); /** - * Helper method to get Tensor from keyspace. In the case of failure the key is - * closed and the error is replied ( no cleaning actions required ) + * Helper method to get Tensor from keyspace. In case of a failure an + * error is documented. * * @param ctx Context in which Redis modules operate * @param keyName key name @@ -329,12 +329,11 @@ int RAI_OpenKey_Tensor(RedisModuleCtx *ctx, RedisModuleString *keyName, RedisMod * an error getting the Tensor */ int RAI_GetTensorFromKeyspace(RedisModuleCtx *ctx, RedisModuleString *keyName, RedisModuleKey **key, - RAI_Tensor **tensor, int mode); + RAI_Tensor **tensor, int mode, RAI_Error *err); /** * Helper method to get Tensor from local context ( no keyspace access ) * - * @param ctx Context in which Redis modules operate * @param localContextDict local non-blocking hash table containing DAG's * tensors * @param localContextKey key name @@ -343,9 +342,8 @@ int RAI_GetTensorFromKeyspace(RedisModuleCtx *ctx, RedisModuleString *keyName, R * failure * @return REDISMODULE_OK on success, or REDISMODULE_ERR if failed */ -int RAI_getTensorFromLocalContext(RedisModuleCtx *ctx, AI_dict *localContextDict, - RedisModuleString *localContextKey, RAI_Tensor **tensor, - RAI_Error *error); +int RAI_getTensorFromLocalContext(AI_dict *localContextDict, RedisModuleString *localContextKey, + RAI_Tensor **tensor, RAI_Error *error); /** * Helper method to replicate a tensor via an AI.TENSORSET command to the @@ -363,7 +361,6 @@ void RedisAI_ReplicateTensorSet(RedisModuleCtx *ctx, RedisModuleString *key, RAI /** * Helper method to parse AI.TENSORGET arguments * - * @param ctx Context in which Redis modules operate * @param argv Redis command arguments, as an array of strings * @param argc Redis command number of arguments * @param t Destination tensor to store the parsed data @@ -372,19 +369,20 @@ void RedisAI_ReplicateTensorSet(RedisModuleCtx *ctx, RedisModuleString *key, RAI * parsing failures * @return processed number of arguments on success, or -1 if the parsing failed */ -int RAI_parseTensorSetArgs(RedisModuleCtx *ctx, RedisModuleString **argv, int argc, RAI_Tensor **t, - int enforceArity, RAI_Error *error); +int RAI_parseTensorSetArgs(RedisModuleString **argv, int argc, RAI_Tensor **t, int enforceArity, + RAI_Error *error); /** * Helper method to parse AI.TENSORGET arguments * - * @param ctx Context in which Redis modules operate + * @param error error data structure to store error message in the case of + * parsing failures * @param argv Redis command arguments, as an array of strings * @param argc Redis command number of arguments * @return The format in which tensor is returned. */ -uint ParseTensorGetArgs(RedisModuleCtx *ctx, RedisModuleString **argv, int argc); +uint ParseTensorGetArgs(RAI_Error *error, RedisModuleString **argv, int argc); /** * Helper method to return a tensor to the client in a response to AI.TENSORGET diff --git a/src/util/arr.h b/src/util/arr.h index 46e7927b5..1881d7b38 100644 --- a/src/util/arr.h +++ b/src/util/arr.h @@ -38,6 +38,7 @@ #endif typedef struct { + bool on_stack; uint32_t len; // TODO: optimize memory by making cap a 16-bit delta from len, and elem_sz 16 bit as well. This // makes the whole header fit in 64 bit @@ -57,7 +58,8 @@ typedef void *array_t; /* Initialize a new array with a given element size and capacity. Should not be used directly - use * array_new instead */ static array_t array_new_sz(uint32_t elem_sz, uint32_t cap, uint32_t len) { - array_hdr_t *hdr = (array_hdr_t *)array_alloc_fn(sizeof(array_hdr_t) + cap * elem_sz); + array_hdr_t *hdr = array_alloc_fn(sizeof(array_hdr_t) + cap * elem_sz); + hdr->on_stack = false; hdr->cap = cap; hdr->elem_sz = elem_sz; hdr->len = len; @@ -73,16 +75,48 @@ static array_t array_new_sz(uint32_t elem_sz, uint32_t cap, uint32_t len) { * */ #define array_new(T, cap) (array_new_sz(sizeof(T), cap, 0)) +#define array_new_on_stack(T, capacity, varName) \ + struct { \ + bool on_stack; \ + uint32_t len; \ + uint32_t cap; \ + uint32_t elem_sz; \ + char buf[(capacity) * sizeof(T)]; \ + } varName##_struct; \ + varName##_struct.on_stack = true; \ + varName##_struct.len = 0; \ + varName##_struct.cap = (capacity); \ + varName##_struct.elem_sz = sizeof(T); \ + T *varName = (T *)varName##_struct.buf; + /* Initialize an array for a given type T with a given length. The capacity allocated is identical * to the length * */ #define array_newlen(T, len) (array_new_sz(sizeof(T), len, len)) +static inline array_t array_persist(array_t arr) { + array_hdr_t *hdr = array_hdr(arr); + if (!hdr->on_stack) + return arr; + array_t tmp = array_new_sz(hdr->elem_sz, hdr->len, hdr->len); + memcpy(tmp, arr, hdr->elem_sz * hdr->len); + return tmp; +} + static inline array_t array_ensure_cap(array_t arr, uint32_t cap) { array_hdr_t *hdr = array_hdr(arr); + if (hdr->on_stack) { + // On stack and still place, return the array. + if (cap <= hdr->cap) + return (array_t)hdr->buf; + // Move to heap. + array_t tmp = array_new_sz(hdr->elem_sz, MAX(hdr->cap * 2, cap), hdr->len); + memcpy(tmp, arr, hdr->elem_sz * hdr->len); + return tmp; + } if (cap > hdr->cap) { hdr->cap = MAX(hdr->cap * 2, cap); - hdr = (array_hdr_t *)array_realloc_fn(hdr, array_sizeof(hdr)); + hdr = array_realloc_fn(hdr, array_sizeof(hdr)); } return (array_t)hdr->buf; } @@ -108,14 +142,15 @@ static inline uint32_t array_len(array_t arr) { return arr ? array_hdr(arr)->len static inline void *array_trimm(array_t arr, uint32_t len, uint32_t cap) { array_hdr_t *arr_hdr = array_hdr(arr); - assert(len >= 0 && "trimming len is negative"); - assert((cap == -1 || cap > 0 || len == cap) && "trimming capacity is illegal"); - assert((cap == -1 || cap >= len) && "trimming len is greater then capacity"); - assert((len <= arr_hdr->len) && "trimming len is greater then current len"); + RedisModule_Assert(len >= 0 && "trimming len is negative"); + RedisModule_Assert((cap == -1 || cap > 0 || len == cap) && "trimming capacity is illegal"); + RedisModule_Assert((cap == -1 || cap >= len) && "trimming len is greater then capacity"); + RedisModule_Assert((len <= arr_hdr->len) && "trimming len is greater then current len"); arr_hdr->len = len; if (cap != -1) { arr_hdr->cap = cap; - arr_hdr = (array_hdr_t *)array_realloc_fn(arr_hdr, array_sizeof(arr_hdr)); + if (!arr_hdr->on_stack) + arr_hdr = array_realloc_fn(arr_hdr, array_sizeof(arr_hdr)); } return arr_hdr->buf; } @@ -126,7 +161,11 @@ static inline void *array_trimm(array_t arr, uint32_t len, uint32_t cap) { #define array_clear(arr) array_hdr(arr)->len = 0 /* Free the array, without dealing with individual elements */ -static void array_free(array_t arr) { array_free_fn(array_hdr(arr)); } +static void array_free(array_t arr) { + array_hdr_t *arr_hdr = array_hdr(arr); + if (!arr_hdr->on_stack) + array_free_fn(array_hdr(arr)); +} /* Repeate the code in "blk" for each element in the array, and give it the name of "as". * e.g: @@ -157,7 +196,7 @@ static void array_free(array_t arr) { array_free_fn(array_hdr(arr)); } /* Pop the top element from the array, reduce the size and return it */ #define array_pop(arr) \ ({ \ - assert(array_hdr(arr)->len > 0); \ + RedisModule_Assert(array_hdr(arr)->len > 0); \ arr[--(array_hdr(arr)->len)]; \ }) diff --git a/tests/flow/tests_dag.py b/tests/flow/tests_dag.py index 62a7ea078..83b517985 100644 --- a/tests/flow/tests_dag.py +++ b/tests/flow/tests_dag.py @@ -24,9 +24,43 @@ def test_dag_load(env): ret = con.execute_command(command) env.assertEqual(ret, [b'OK']) + def test_dag_load_errors(env): con = env.getConnection() + # ERR wrong number of arguments for LOAD + try: + command = "AI.DAGRUN PERSIST 1 no_tensor{1} LOAD 1" + + 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 LOAD in 'AI.DAGRUN' command",exception.__str__()) + + # ERR invalid or negative number of arguments for LOAD + try: + command = "AI.DAGRUN LOAD notnumber{{1}} |> AI.TENSORGET 'no_tensor{1}'" + + 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__()) + + con.execute_command('AI.TENSORSET', 'a{{1}}', 'FLOAT', 2, 2, 'VALUES', 2, 3, 2, 3) + + # ERR number of keys to LOAD does not match the number of given arguments. + try: + command = "AI.DAGRUN LOAD 2 a{{1}}" + + ret = con.execute_command(command) + except Exception as e: + exception = e + env.assertEqual(type(exception), redis.exceptions.ResponseError) + env.assertEqual("number of keys to LOAD in AI.DAGRUN command does not match the number of " + "given arguments", exception.__str__()) + # ERR tensor key is empty try: command = "AI.DAGRUN "\ @@ -55,51 +89,103 @@ def test_dag_load_errors(env): env.assertEqual("WRONGTYPE Operation against a key holding the wrong kind of value",exception.__str__()) -def test_dag_common_errors(env): +def test_dag_persist_errors(env): con = env.getConnection() + con.execute_command('AI.TENSORSET', 'a{{1}}', 'FLOAT', 2, 2, 'VALUES', 2, 3, 2, 3) - # ERR unsupported command within DAG + # ERR wrong number of arguments for PERSIST try: - command = "AI.DAGRUN |> "\ - "AI.DONTEXIST tensor1{{1}} FLOAT 1 2 VALUES 5 10" + command = "AI.DAGRUN LOAD 1 a{{1}} PERSIST 1" ret = con.execute_command(command) except Exception as e: exception = e env.assertEqual(type(exception), redis.exceptions.ResponseError) - env.assertEqual("unsupported command within DAG",exception.__str__()) + env.assertEqual("wrong number of arguments for PERSIST in 'AI.DAGRUN' command",exception.__str__()) - # ERR wrong number of arguments for 'AI.DAGRUN' command + # ERR invalid or negative value found in number of keys to PERSIST try: - command = "AI.DAGRUN " + command = "AI.DAGRUN PERSIST notnumber{{1}} |> " \ + "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("wrong number of arguments for 'AI.DAGRUN' command",exception.__str__()) + env.assertEqual("invalid or negative value found in number of keys to PERSIST",exception.__str__()) - # ERR invalid or negative value found in number of keys to PERSIST + # ERR number of keys to PERSIST does not match the number of given arguments. + try: + command = "AI.DAGRUN PERSIST 2 tensor1{1}" + + ret = con.execute_command(command) + except Exception as e: + exception = e + env.assertEqual(type(exception), redis.exceptions.ResponseError) + env.assertEqual("number of keys to PERSIST in AI.DAGRUN command does not match the number of " + "given arguments", exception.__str__()) + + +def test_dag_timeout_errors(env): + con = env.getConnection() + + # ERR no value provided for timeout try: - command = "AI.DAGRUN PERSIST notnumber{{1}} |> "\ + command = "AI.DAGRUN PERSIST 1 no_tensor{1} TIMEOUT" + + ret = con.execute_command(command) + except Exception as e: + exception = e + env.assertEqual(type(exception), redis.exceptions.ResponseError) + env.assertEqual("No value provided for TIMEOUT",exception.__str__()) + + # ERR invalid timeout value + try: + command = "AI.DAGRUN TIMEOUT notnumber{{1}} |> " \ "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 PERSIST",exception.__str__()) + env.assertEqual("Invalid value for TIMEOUT",exception.__str__()) - # ERR invalid or negative value found in number of keys to LOAD + +def test_dag_common_errors(env): + con = env.getConnection() + + # ERR unsupported command within DAG try: - command = "AI.DAGRUN LOAD notnumber{{1}} |> "\ - "AI.TENSORSET tensor1 FLOAT 1 2 VALUES 5 10" + command = "AI.DAGRUN |> "\ + "AI.DONTEXIST tensor1{{1}} 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__()) + env.assertEqual("unsupported command within DAG",exception.__str__()) + + # ERR wrong number of arguments for 'AI.DAGRUN' command + try: + command = "AI.DAGRUN " + + 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.DAGRUN' command",exception.__str__()) + + # ERR DAG with no ops + command = "AI.TENSORSET volatile_tensor1 FLOAT 1 2 VALUES 5 10" + ret = con.execute_command(command) + env.assertEqual(ret, b'OK') + try: + command = "AI.DAGRUN LOAD 1 volatile_tensor1" + ret = con.execute_command(command) + except Exception as e: + exception = e + env.assertEqual(type(exception), redis.exceptions.ResponseError) + env.assertEqual("DAG is empty",exception.__str__()) def test_dagro_common_errors(env): diff --git a/tests/flow/tests_llapi.py b/tests/flow/tests_llapi.py index 60319bff5..5500447c5 100644 --- a/tests/flow/tests_llapi.py +++ b/tests/flow/tests_llapi.py @@ -21,11 +21,11 @@ def wrapper(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 f(env, *args, **kwargs) return wrapper @@ -76,3 +76,74 @@ def test_script_run_async(env): ret = con.execute_command("RAI_llapi.scriptRun") env.assertEqual(ret, b'Async run success') + + +@ensure_test_module_loaded +def test_dag_build_and_run(env): + con = env.getConnection() + + 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) + 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') + + 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("RAI_llapi.DAGrun") + env.assertEqual(ret, b'DAG run success') + + +@ensure_test_module_loaded +def test_dagrun_multidevice_resnet(env): + con = env.getConnection() + + model_name_0 = 'imagenet_model1:{{1}}' + model_name_1 = 'imagenet_model2:{{1}}' + script_name_0 = 'imagenet_script1:{{1}}' + script_name_1 = 'imagenet_script2:{{1}}' + inputvar = 'images' + outputvar = 'output' + image_key = 'image:{{1}}' + temp_key1 = 'temp_key1:{{1}}' + temp_key2_0 = 'temp_key2_0' + temp_key2_1 = 'temp_key2_1' + class_key_0 = 'output0:{{1}}' + class_key_1 = 'output1:{{1}}' + + model_pb, script, labels, img = load_resnet_test_data() + + device_0 = 'CPU:1' + device_1 = DEVICE + + ret = con.execute_command('AI.MODELSET', model_name_0, 'TF', device_0, + 'INPUTS', inputvar, + 'OUTPUTS', outputvar, + 'BLOB', model_pb) + env.assertEqual(ret, b'OK') + + ret = con.execute_command('AI.MODELSET', model_name_1, 'TF', device_1, + 'INPUTS', inputvar, + 'OUTPUTS', outputvar, + 'BLOB', model_pb) + env.assertEqual(ret, b'OK') + ret = con.execute_command('AI.SCRIPTSET', script_name_0, device_0, 'SOURCE', script) + env.assertEqual(ret, b'OK') + ret = con.execute_command('AI.SCRIPTSET', script_name_1, device_1, 'SOURCE', script) + env.assertEqual(ret, b'OK') + ret = con.execute_command('AI.TENSORSET', image_key, 'UINT8', img.shape[1], img.shape[0], 3, 'BLOB', img.tobytes()) + env.assertEqual(ret, b'OK') + + ret = con.execute_command("RAI_llapi.DAG_resnet") + env.assertEqual(ret, b'DAG resnet success') diff --git a/tests/module/DAG_utils.c b/tests/module/DAG_utils.c new file mode 100644 index 000000000..0b8486b29 --- /dev/null +++ b/tests/module/DAG_utils.c @@ -0,0 +1,419 @@ +#define REDISAI_EXTERN 1 +#include "DAG_utils.h" +#include +#include +#include +#include +#include "util/arr.h" + +pthread_mutex_t global_lock = PTHREAD_MUTEX_INITIALIZER; +pthread_cond_t global_cond = PTHREAD_COND_INITIALIZER; + +static void _InitRunResults(RAI_RunResults *results) { + RedisAI_InitError(&results->error); + results->outputs = array_new(RAI_Tensor *, 1); +} + +static void _FreeRunResults(RAI_RunResults *results) { + RedisAI_FreeError(results->error); + for (size_t i = 0; i < array_len(results->outputs); i++) { + RedisAI_TensorFree(results->outputs[i]); + } + array_free(results->outputs); +} + +static size_t _ResultsNumOutputs(RAI_RunResults results) { + return array_len(results.outputs); +} +static void *_getFromKeySpace(RedisModuleCtx *ctx, const char *keyNameStr) { + + RedisModuleString *keyRedisStr = RedisModule_CreateString(ctx, keyNameStr, strlen(keyNameStr)); + RedisModuleKey *key = RedisModule_OpenKey(ctx, keyRedisStr, REDISMODULE_READ); + RedisModule_FreeString(ctx, keyRedisStr); + void *value = RedisModule_ModuleTypeGetValue(key); + RedisModule_CloseKey(key); + return value; +} + +static bool _assertError(RAI_Error *err, int status, const char *error_msg) { + if (status != REDISMODULE_ERR) { + return false; + } + // return true only if err contains the specific error_msg + return strcmp(RedisAI_GetError(err), error_msg) == 0; +} + +static void _DAGFinishFuncError(RAI_OnFinishCtx *onFinishCtx, void *private_data) { + REDISMODULE_NOT_USED(onFinishCtx); + REDISMODULE_NOT_USED(private_data); + //Do nothing, this callback should not be used... + RedisModule_Assert(false); +} + +static void _DAGFinishFunc(RAI_OnFinishCtx *onFinishCtx, void *private_data) { + + RAI_RunResults *results = (RAI_RunResults *)private_data; + if (RedisAI_DAGRunError(onFinishCtx)) { + const RAI_Error *error = RedisAI_DAGGetError(onFinishCtx); + RedisAI_CloneError(results->error, error); + pthread_cond_signal(&global_cond); + return; + } + size_t n_outputs = RedisAI_DAGNumOutputs(onFinishCtx); + for (size_t i = 0; i < n_outputs; i++) { + RAI_Tensor *t = RedisAI_DAGOutputTensor(onFinishCtx, i); + RedisModule_Assert(t != NULL); + results->outputs = array_append(results->outputs, RedisAI_TensorGetShallowCopy(t)); + } + + // Verify that we return NULL as output for an index out of range. + RAI_Tensor *t = RedisAI_DAGOutputTensor(onFinishCtx, n_outputs); + RedisModule_Assert(t == NULL); + pthread_cond_signal(&global_cond); +} + +int testModelRunOpError(RedisModuleCtx *ctx) { + + RAI_DAGRunCtx *run_info = RedisAI_DAGRunCtxCreate(); + RAI_Error *err; + RedisAI_InitError(&err); + int res = LLAPIMODULE_ERR; + // The model m{1} should exist in key space. + RAI_Model *model = (RAI_Model *)_getFromKeySpace(ctx, "m{1}"); + RAI_DAGRunOp *op = RedisAI_DAGCreateModelRunOp(model); + RedisAI_DAGRunOpAddInput(op, "first_input"); + + // This model expect for 2 inputs not 1. + int status = RedisAI_DAGAddRunOp(run_info, op, err); + if (!_assertError(err, status, "Number of keys given as INPUTS does not match model definition")) { + goto cleanup; + } + RedisAI_ClearError(err); + RedisAI_DAGRunOpAddInput(op, "second_input"); + status = RedisAI_DAGAddRunOp(run_info, op, err); + + // We still get an error since the model expects for an output as well. + if (!_assertError(err, status, "Number of keys given as OUTPUTS does not match model definition")) { + goto cleanup; + } + res = LLAPIMODULE_OK; + + cleanup: + RedisAI_FreeError(err); + RedisAI_DAGRunOpFree(op); + RedisAI_DAGFree(run_info); + return res; +} + +int testEmptyDAGError(RedisModuleCtx *ctx) { + + RAI_DAGRunCtx *run_info = RedisAI_DAGRunCtxCreate(); + RAI_Error* err; + RedisAI_InitError(&err); + int res = LLAPIMODULE_ERR; + + RAI_Tensor *t = (RAI_Tensor *)_getFromKeySpace(ctx, "a{1}"); + RedisAI_DAGLoadTensor(run_info, "input", t); + + int status = RedisAI_DAGRun(run_info, _DAGFinishFuncError, NULL, err); + if(!_assertError(err, status, "ERR DAG is empty")) { + goto cleanup; + } + res = LLAPIMODULE_OK; + + cleanup: + RedisAI_FreeError(err); + RedisAI_DAGFree(run_info); + return res; +} + +int testKeysMismatchError(RedisModuleCtx *ctx) { + + RAI_DAGRunCtx *run_info = RedisAI_DAGRunCtxCreate(); + RAI_Error* err; + RedisAI_InitError(&err); + int res = LLAPIMODULE_ERR; + + RAI_Tensor *t = (RAI_Tensor *)_getFromKeySpace(ctx, "a{1}"); + RedisAI_DAGLoadTensor(run_info, "input", t); + + RedisAI_DAGAddTensorGet(run_info, "non existing tensor", err); + int status = RedisAI_DAGRun(run_info, _DAGFinishFuncError, NULL, err); + if(!_assertError(err, status, "ERR INPUT key cannot be found in DAG")) { + goto cleanup; + } + res = LLAPIMODULE_OK; + + cleanup: + RedisAI_FreeError(err); + RedisAI_DAGFree(run_info); + return res; +} + +int testBuildDAGFromString(RedisModuleCtx *ctx) { + + RAI_DAGRunCtx *run_info = RedisAI_DAGRunCtxCreate(); + RAI_RunResults results; + _InitRunResults(&results); + int res = LLAPIMODULE_ERR; + + RAI_Tensor *t = (RAI_Tensor *)_getFromKeySpace(ctx, "a{1}"); + RedisAI_DAGLoadTensor(run_info, "input1", t); + + const char *dag_string = "bad string"; + int status = RedisAI_DAGAddOpsFromString(run_info, dag_string, results.error); + if (!_assertError(results.error, status, "DAG op should start with: '|>' ")) { + goto cleanup; + } + RedisAI_ClearError(results.error); + + t = (RAI_Tensor *)_getFromKeySpace(ctx, "b{1}"); + RedisAI_DAGAddTensorSet(run_info, "input2", t); + + dag_string = "|> AI.MODELRUN m{1} INPUTS input1 input2 OUTPUTS output |> bad_op no_tensor"; + status = RedisAI_DAGAddOpsFromString(run_info, dag_string, results.error); + if (!_assertError(results.error, status, "unsupported command within DAG")) { + goto cleanup; + } + RedisAI_ClearError(results.error); + RedisModule_Assert(RedisAI_DAGNumOps(run_info) == 1); + + dag_string = "|> AI.MODELRUN m{1} INPUTS input1 input2 OUTPUTS output |> AI.TENSORGET output"; + if (RedisAI_DAGAddOpsFromString(run_info, dag_string, results.error) != REDISMODULE_OK) { + goto cleanup; + } + RedisModule_Assert(RedisAI_DAGNumOps(run_info) == 3); + RedisAI_DAGAddTensorGet(run_info, "input1", results.error); + RedisModule_Assert(RedisAI_DAGNumOps(run_info) == 4); + + pthread_mutex_lock(&global_lock); + if (RedisAI_DAGRun(run_info, _DAGFinishFunc, &results, results.error) != REDISMODULE_OK) { + pthread_mutex_unlock(&global_lock); + goto cleanup; + } + // Wait until the onFinish callback returns. + pthread_cond_wait(&global_cond, &global_lock); + pthread_mutex_unlock(&global_lock); + + // Verify that we received the expected tensor at the end of the run. + RedisModule_Assert(_ResultsNumOutputs(results) == 2); + res = LLAPIMODULE_OK; + + cleanup: + _FreeRunResults(&results); + RedisAI_DAGFree(run_info); + return res; +} + +int testSimpleDAGRun(RedisModuleCtx *ctx) { + + RAI_DAGRunCtx *run_info = RedisAI_DAGRunCtxCreate(); + RAI_RunResults results; + _InitRunResults(&results); + int res = LLAPIMODULE_ERR; + + RAI_Tensor *t = (RAI_Tensor *)_getFromKeySpace(ctx, "a{1}"); + RedisAI_DAGLoadTensor(run_info, "input1", t); + t = (RAI_Tensor *)_getFromKeySpace(ctx, "b{1}"); + RedisAI_DAGLoadTensor(run_info, "input2", t); + + // The model m{1} should exist in key space. + RAI_Model *model = (RAI_Model *)_getFromKeySpace(ctx, "m{1}"); + RAI_DAGRunOp *op = RedisAI_DAGCreateModelRunOp(model); + RedisAI_DAGRunOpAddInput(op, "input1"); + RedisAI_DAGRunOpAddInput(op, "input2"); + RedisAI_DAGRunOpAddOutput(op, "output"); + if (RedisAI_DAGAddRunOp(run_info, op, results.error) != REDISMODULE_OK) { + goto cleanup; + } + + RedisAI_DAGAddTensorGet(run_info, "output", results.error); + pthread_mutex_lock(&global_lock); + if (RedisAI_DAGRun(run_info, _DAGFinishFunc, &results, results.error) != REDISMODULE_OK) { + pthread_mutex_unlock(&global_lock); + goto cleanup; + } + // Wait until the onFinish callback returns. + pthread_cond_wait(&global_cond, &global_lock); + pthread_mutex_unlock(&global_lock); + + // Verify that we received the expected tensor at the end of the run. + RedisModule_Assert(_ResultsNumOutputs(results) == 1); + RAI_Tensor *out_tensor = results.outputs[0]; + double expceted[4] = {4, 9, 4, 9}; + double val; + for (long long i = 0; i < 4; i++) { + if(RedisAI_TensorGetValueAsDouble(out_tensor, i, &val) != 0) { + goto cleanup; + } + if (expceted[i] != val) { + goto cleanup; + } + } + res = LLAPIMODULE_OK; + + cleanup: + _FreeRunResults(&results); + RedisAI_DAGFree(run_info); + return res; +} + +int testSimpleDAGRun2(RedisModuleCtx *ctx) { + + RAI_DAGRunCtx *run_info = RedisAI_DAGRunCtxCreate(); + RAI_RunResults results; + _InitRunResults(&results); + int res = LLAPIMODULE_ERR; + + RAI_Tensor *tensor = (RAI_Tensor*)_getFromKeySpace(ctx, "a{1}"); + RedisAI_DAGAddTensorSet(run_info, "input1", tensor); + tensor = (RAI_Tensor*)_getFromKeySpace(ctx, "b{1}"); + RedisAI_DAGAddTensorSet(run_info, "input2", tensor); + + // The script myscript{1} should exist in key space. + RAI_Script *script = (RAI_Script*)_getFromKeySpace(ctx, "myscript{1}"); + RAI_DAGRunOp *op = RedisAI_DAGCreateScriptRunOp(script, "bar"); + RedisAI_DAGRunOpAddInput(op, "input1"); + RedisAI_DAGRunOpAddInput(op, "input2"); + RedisAI_DAGRunOpAddOutput(op, "output"); + if (RedisAI_DAGAddRunOp(run_info, op, results.error) != REDISMODULE_OK) { + goto cleanup; + } + + RedisAI_DAGAddTensorGet(run_info, "output", results.error); + pthread_mutex_lock(&global_lock); + if (RedisAI_DAGRun(run_info, _DAGFinishFunc, &results, results.error) != REDISMODULE_OK) { + pthread_mutex_unlock(&global_lock); + goto cleanup; + } + // Wait until the onFinish callback returns. + pthread_cond_wait(&global_cond, &global_lock); + pthread_mutex_unlock(&global_lock); + + // Verify that we received the expected tensor at the end of the run. + RedisModule_Assert(_ResultsNumOutputs(results) == 1); + RAI_Tensor *out_tensor = results.outputs[0]; + double expceted[4] = {4, 6, 4, 6}; + double val; + for (long long i = 0; i < 4; i++) { + if(RedisAI_TensorGetValueAsDouble(out_tensor, i, &val) != 0) { + goto cleanup; + } + if (expceted[i] != val) { + goto cleanup; + } + } + res = LLAPIMODULE_OK; + + cleanup: + _FreeRunResults(&results); + RedisAI_DAGFree(run_info); + return res; +} + +int testSimpleDAGRun2Error(RedisModuleCtx *ctx) { + + RAI_DAGRunCtx *run_info = RedisAI_DAGRunCtxCreate(); + RAI_RunResults results; + _InitRunResults(&results); + int res = LLAPIMODULE_ERR; + + RAI_Tensor *tensor = (RAI_Tensor*) _getFromKeySpace(ctx, "a{1}"); + RedisAI_DAGAddTensorSet(run_info, "input1", tensor); + + // The script myscript{1} should exist in key space. + RAI_Script *script = (RAI_Script*)_getFromKeySpace(ctx, "myscript{1}"); + RAI_DAGRunOp *op = RedisAI_DAGCreateScriptRunOp(script, "no_function"); + RedisAI_DAGRunOpAddInput(op, "input1"); + RedisAI_DAGRunOpAddOutput(op, "output"); + if (RedisAI_DAGAddRunOp(run_info, op, results.error) != REDISMODULE_OK) { + goto cleanup; + } + + RedisAI_DAGAddTensorGet(run_info, "output", results.error); + pthread_mutex_lock(&global_lock); + if (RedisAI_DAGRun(run_info, _DAGFinishFunc, &results, results.error) != REDISMODULE_OK) { + pthread_mutex_unlock(&global_lock); + goto cleanup; + } + // Wait until the onFinish callback returns. + pthread_cond_wait(&global_cond, &global_lock); + pthread_mutex_unlock(&global_lock); + + // Verify that we received an error in SCRIPTRUN op. + RedisModule_Assert(_ResultsNumOutputs(results) == 0); + if (RedisAI_GetErrorCode(results.error) != RedisAI_ErrorCode_ESCRIPTRUN) { + goto cleanup; + } + res = LLAPIMODULE_OK; + + cleanup: + _FreeRunResults(&results); + RedisAI_DAGFree(run_info); + return res; +} + +int testDAGResnet(RedisModuleCtx *ctx) { + RAI_RunResults results; + _InitRunResults(&results); + int res = LLAPIMODULE_ERR; + + // Build the DAG with LOAD->SCRIPTRUN->MODELRUN->MODELRUN-SCRIPTRUN->SCRIPTRUN->TENSORGET + RAI_DAGRunCtx *run_info = RedisAI_DAGRunCtxCreate(); + RAI_Tensor *t = (RAI_Tensor *)_getFromKeySpace(ctx, "image:{{1}}"); + RedisAI_DAGLoadTensor(run_info, "image:{{1}}", t); + + RAI_Script *script = (RAI_Script *)_getFromKeySpace(ctx, "imagenet_script1:{{1}}"); + RAI_DAGRunOp *script_op = RedisAI_DAGCreateScriptRunOp(script, "pre_process_3ch"); + RedisAI_DAGRunOpAddInput(script_op, "image:{{1}}"); + RedisAI_DAGRunOpAddOutput(script_op, "tmp_key:{{1}}"); + RedisAI_DAGAddRunOp(run_info, script_op, results.error); + + RAI_Model *model = (RAI_Model *)_getFromKeySpace(ctx, "imagenet_model1:{{1}}"); + RAI_DAGRunOp *model_op = RedisAI_DAGCreateModelRunOp(model); + RedisAI_DAGRunOpAddInput(model_op, "tmp_key:{{1}}"); + RedisAI_DAGRunOpAddOutput(model_op, "tmp_key2_0"); + if (RedisAI_DAGAddRunOp(run_info, model_op, results.error) != REDISMODULE_OK) goto cleanup; + + model = (RAI_Model *)_getFromKeySpace(ctx, "imagenet_model2:{{1}}"); + model_op = RedisAI_DAGCreateModelRunOp(model); + RedisAI_DAGRunOpAddInput(model_op, "tmp_key:{{1}}"); + RedisAI_DAGRunOpAddOutput(model_op, "tmp_key2_1"); + if (RedisAI_DAGAddRunOp(run_info, model_op, results.error) != REDISMODULE_OK) goto cleanup; + + script_op = RedisAI_DAGCreateScriptRunOp(script, "ensemble"); + RedisAI_DAGRunOpAddInput(script_op, "tmp_key2_0"); + RedisAI_DAGRunOpAddInput(script_op, "tmp_key2_1"); + RedisAI_DAGRunOpAddOutput(script_op, "tmp_key_1:{{1}}"); + RedisAI_DAGAddRunOp(run_info, script_op, results.error); + + script_op = RedisAI_DAGCreateScriptRunOp(script, "post_process"); + RedisAI_DAGRunOpAddInput(script_op, "tmp_key_1:{{1}}"); + RedisAI_DAGRunOpAddOutput(script_op, "output:{{1}}"); + RedisAI_DAGAddRunOp(run_info, script_op, results.error); + + RedisAI_DAGAddTensorGet(run_info, "output:{{1}}", results.error); + + pthread_mutex_lock(&global_lock); + if (RedisAI_DAGRun(run_info, _DAGFinishFunc, &results, results.error) != REDISMODULE_OK) { + pthread_mutex_unlock(&global_lock); + goto cleanup; + } + // Wait until the onFinish callback returns. + pthread_cond_wait(&global_cond, &global_lock); + pthread_mutex_unlock(&global_lock); + + // Verify that we received the expected output. + RedisModule_Assert(_ResultsNumOutputs(results) == 1); + RAI_Tensor *out_tensor = results.outputs[0]; + long long val; + if(RedisAI_TensorGetValueAsLongLong(out_tensor, 0, &val) != 0) goto cleanup; + if (0 <= val && val <= 1000) { + res = LLAPIMODULE_OK; + } + + cleanup: + _FreeRunResults(&results); + RedisAI_DAGFree(run_info); + return res; +} diff --git a/tests/module/DAG_utils.h b/tests/module/DAG_utils.h new file mode 100644 index 000000000..fbaa40c3b --- /dev/null +++ b/tests/module/DAG_utils.h @@ -0,0 +1,28 @@ +#pragma once +#include "redisai.h" +#include +#include + +#define LLAPIMODULE_OK 0 +#define LLAPIMODULE_ERR 1 + +typedef struct RAI_RunResults { + RAI_Tensor **outputs; + RAI_Error *error; +} RAI_RunResults; + +int testModelRunOpError(RedisModuleCtx *ctx); + +int testEmptyDAGError(RedisModuleCtx *ctx); + +int testKeysMismatchError(RedisModuleCtx *ctx); + +int testBuildDAGFromString(RedisModuleCtx *ctx); + +int testSimpleDAGRun(RedisModuleCtx *ctx); + +int testSimpleDAGRun2(RedisModuleCtx *ctx); + +int testSimpleDAGRun2Error(RedisModuleCtx *ctx); + +int testDAGResnet(RedisModuleCtx *ctx); diff --git a/tests/module/LLAPI.c b/tests/module/LLAPI.c index 0bf12a070..16e4563ac 100644 --- a/tests/module/LLAPI.c +++ b/tests/module/LLAPI.c @@ -1,14 +1,10 @@ #define REDISMODULE_MAIN -#include "../../src/redisai.h" +#include "DAG_utils.h" #include #include -#include -#include -pthread_mutex_t global_lock = PTHREAD_MUTEX_INITIALIZER; -pthread_cond_t global_cond = PTHREAD_COND_INITIALIZER; typedef enum LLAPI_status {LLAPI_RUN_NONE = 0, LLAPI_RUN_SUCCESS, @@ -16,6 +12,9 @@ typedef enum LLAPI_status {LLAPI_RUN_NONE = 0, LLAPI_NUM_OUTPUTS_ERROR } LLAPI_status; +extern pthread_mutex_t global_lock; +extern pthread_cond_t global_cond; + int RAI_llapi_basic_check(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { REDISMODULE_NOT_USED(argv); @@ -236,6 +235,60 @@ int RAI_llapi_scriptRun(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) return RedisModule_ReplyWithSimpleString(ctx, "Async run success"); } +int RAI_llapi_DAGRun(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + REDISMODULE_NOT_USED(argv); + + if(argc > 1) { + RedisModule_WrongArity(ctx); + return REDISMODULE_OK; + } + + // Test the case of a failure due to addition of a non compatible MODELRUN op. + if(testModelRunOpError(ctx) != LLAPIMODULE_OK) { + return RedisModule_ReplyWithSimpleString(ctx, "MODELRUN op error test failed"); + } + // Test the case of a failure due an empty DAG. + if(testEmptyDAGError(ctx) != LLAPIMODULE_OK) { + return RedisModule_ReplyWithSimpleString(ctx, "DAG keys mismatch error test failed"); + } + // Test the case of a failure due to an op within a DAG whose inkey does not exist in the DAG. + if(testKeysMismatchError(ctx) != LLAPIMODULE_OK) { + return RedisModule_ReplyWithSimpleString(ctx, "DAG keys mismatch error test failed"); + } + // Test the case of building and running a DAG with LOAD, TENSORGET and MODELRUN ops. + if(testSimpleDAGRun(ctx) != LLAPIMODULE_OK) { + return RedisModule_ReplyWithSimpleString(ctx, "Simple DAG run test failed"); + } + // Test the case of building and running a DAG with TENSORSET, SCRIPTRUN and TENSORGET ops. + if(testSimpleDAGRun2(ctx) != LLAPIMODULE_OK) { + return RedisModule_ReplyWithSimpleString(ctx, "Simple DAG run2 test failed"); + } + // Test the case of building the same DAG as in previous test, but when this time it should return with an error. + if(testSimpleDAGRun2Error(ctx) != LLAPIMODULE_OK) { + return RedisModule_ReplyWithSimpleString(ctx, "Simple DAG run2 error test failed"); + } + // Test the case of building DAG ops from string. + if(testBuildDAGFromString(ctx) != LLAPIMODULE_OK) { + return RedisModule_ReplyWithSimpleString(ctx, "Build DAG from string test failed"); + } + return RedisModule_ReplyWithSimpleString(ctx, "DAG run success"); +} + +int RAI_llapi_DAG_resnet(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + REDISMODULE_NOT_USED(argv); + + if(argc > 1) { + RedisModule_WrongArity(ctx); + return REDISMODULE_OK; + } + + if (testDAGResnet(ctx) != LLAPIMODULE_OK) { + return RedisModule_ReplyWithSimpleString(ctx, "DAG resnet failed"); + } + return RedisModule_ReplyWithSimpleString(ctx, "DAG resnet success"); +} + + int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { REDISMODULE_NOT_USED(argv); REDISMODULE_NOT_USED(argc); @@ -259,5 +312,17 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) if(RedisModule_CreateCommand(ctx, "RAI_llapi.scriptRun", RAI_llapi_scriptRun, "", 0, 0, 0) == REDISMODULE_ERR) return REDISMODULE_ERR; + + + if(RedisModule_CreateCommand(ctx, "RAI_llapi.DAGRun", RAI_llapi_DAGRun, "", + 0, 0, 0) == REDISMODULE_ERR) { + return REDISMODULE_ERR; + } + + if(RedisModule_CreateCommand(ctx, "RAI_llapi.DAG_resnet", RAI_llapi_DAG_resnet, "", + 0, 0, 0) == REDISMODULE_ERR) { + return REDISMODULE_ERR; + } + return REDISMODULE_OK; } diff --git a/tests/module/Makefile b/tests/module/Makefile index 22f57ee62..2ad5b1984 100644 --- a/tests/module/Makefile +++ b/tests/module/Makefile @@ -29,8 +29,8 @@ all: $(TEST_MODULES) %.o: %.c $(CC) $(DEBUGFLAGS) -I../../src -DREDIS_MODULE_TARGET -DREDISMODULE_EXPERIMENTAL_API $(SHOBJ_CFLAGS) -fPIC -c $< -o $@ -%.so: %.o - $(CC) -o $@ LLAPI.o $(SHOBJ_LDFLAGS) -lc -lm +%.so: %.o DAG_utils.o + $(CC) -o $@ $^ $(SHOBJ_LDFLAGS) -lc -lm chmod +x LLAPI.so .PHONY: clean diff --git a/tests/unit/unit_tests_err.cpp b/tests/unit/unit_tests_err.cpp index 5b73f514f..36129338e 100644 --- a/tests/unit/unit_tests_err.cpp +++ b/tests/unit/unit_tests_err.cpp @@ -6,7 +6,6 @@ extern "C" { #include "src/err.h" } - class ErrorStructTest : public ::testing::Test { protected: static void SetUpTestCase() {