@@ -125,52 +125,41 @@ static int _ModelRunCtx_SetParams(RedisModuleCtx *ctx, RedisModuleString **inkey
125125int ParseModelRunCommand (RedisAI_RunInfo * rinfo , RedisModuleCtx * ctx , RedisModuleString * * argv ,
126126 int argc ) {
127127
128+ RAI_DagOp * currentOp ;
129+ RAI_InitDagOp (& currentOp );
130+ rinfo -> dagOps = array_append (rinfo -> dagOps , currentOp );
131+
128132 // Build a ModelRunCtx from command.
129- RAI_Error error = {0 };
130133 RAI_Model * model ;
131- RedisModuleString * * inkeys = array_new (RedisModuleString * , 1 );
132- RedisModuleString * * outkeys = array_new (RedisModuleString * , 1 );
133- RedisModuleString * runkey = NULL ;
134- RAI_ModelRunCtx * mctx = NULL ;
135- RAI_DagOp * currentOp ;
136134
137135 long long timeout = 0 ;
138- if (_ModelRunCommand_ParseArgs (ctx , argv , argc , & model , & error , & inkeys , & outkeys , & runkey ,
136+ if (_ModelRunCommand_ParseArgs (ctx , argv , argc , & model , currentOp -> err , & currentOp -> inkeys ,
137+ & currentOp -> outkeys , & currentOp -> runkey ,
139138 & timeout ) == REDISMODULE_ERR ) {
140- RedisModule_ReplyWithError (ctx , RAI_GetErrorOneLine (& error ));
139+ RedisModule_ReplyWithError (ctx , RAI_GetErrorOneLine (currentOp -> err ));
140+ goto cleanup ;
141+ }
142+
143+ if (timeout > 0 && !rinfo -> single_op_dag ) {
144+ RedisModule_ReplyWithError (ctx , "ERR TIMEOUT not allowed within a DAG command" );
141145 goto cleanup ;
142146 }
143- mctx = RAI_ModelRunCtxCreate (model );
144147
148+ RAI_ModelRunCtx * mctx = RAI_ModelRunCtxCreate (model );
145149 if (rinfo -> single_op_dag ) {
146150 rinfo -> timeout = timeout ;
147151 // Set params in ModelRunCtx, bring inputs from key space.
148- if (_ModelRunCtx_SetParams (ctx , inkeys , outkeys , mctx ) == REDISMODULE_ERR )
152+ if (_ModelRunCtx_SetParams (ctx , currentOp -> inkeys , currentOp -> outkeys , mctx ) ==
153+ REDISMODULE_ERR )
149154 goto cleanup ;
150155 }
151- if (RAI_InitDagOp (& currentOp ) == REDISMODULE_ERR ) {
152- RedisModule_ReplyWithError (
153- ctx , "ERR Unable to allocate the memory and initialise the RAI_dagOp structure" );
154- goto cleanup ;
155- }
156+
156157 currentOp -> commandType = REDISAI_DAG_CMD_MODELRUN ;
157- Dag_PopulateOp ( currentOp , mctx , inkeys , outkeys , runkey ) ;
158- rinfo -> dagOps = array_append ( rinfo -> dagOps , currentOp ) ;
158+ currentOp -> mctx = mctx ;
159+ currentOp -> devicestr = mctx -> model -> devicestr ;
159160 return REDISMODULE_OK ;
160161
161162cleanup :
162- for (size_t i = 0 ; i < array_len (inkeys ); i ++ ) {
163- RedisModule_FreeString (NULL , inkeys [i ]);
164- }
165- array_free (inkeys );
166- for (size_t i = 0 ; i < array_len (outkeys ); i ++ ) {
167- RedisModule_FreeString (NULL , outkeys [i ]);
168- }
169- array_free (outkeys );
170- if (runkey )
171- RedisModule_FreeString (NULL , runkey );
172- if (mctx )
173- RAI_ModelRunCtxFree (mctx );
174163 RAI_FreeRunInfo (rinfo );
175164 return REDISMODULE_ERR ;
176165}
@@ -283,55 +272,44 @@ static int _ScriptRunCtx_SetParams(RedisModuleCtx *ctx, RedisModuleString **inke
283272int ParseScriptRunCommand (RedisAI_RunInfo * rinfo , RedisModuleCtx * ctx , RedisModuleString * * argv ,
284273 int argc ) {
285274
275+ RAI_DagOp * currentOp ;
276+ RAI_InitDagOp (& currentOp );
277+ rinfo -> dagOps = array_append (rinfo -> dagOps , currentOp );
278+
286279 // Build a ScriptRunCtx from command.
287- RAI_Error error = {0 };
288280 RAI_Script * script ;
289- RedisModuleString * * inkeys = array_new (RedisModuleString * , 1 );
290- RedisModuleString * * outkeys = array_new (RedisModuleString * , 1 );
291- RedisModuleString * runkey = NULL ;
292281 const char * func_name = NULL ;
293- RAI_ScriptRunCtx * sctx = NULL ;
294- RAI_DagOp * currentOp ;
295282
296283 long long timeout = 0 ;
297284 int variadic = -1 ;
298- if (_ScriptRunCommand_ParseArgs (ctx , argv , argc , & script , & error , & inkeys , & outkeys , & runkey ,
299- & func_name , & timeout , & variadic ) == REDISMODULE_ERR ) {
300- RedisModule_ReplyWithError (ctx , RAI_GetErrorOneLine (& error ));
285+ if (_ScriptRunCommand_ParseArgs (ctx , argv , argc , & script , currentOp -> err , & currentOp -> inkeys ,
286+ & currentOp -> outkeys , & currentOp -> runkey , & func_name , & timeout ,
287+ & variadic ) == REDISMODULE_ERR ) {
288+ RedisModule_ReplyWithError (ctx , RAI_GetErrorOneLine (currentOp -> err ));
301289 goto cleanup ;
302290 }
303- sctx = RAI_ScriptRunCtxCreate (script , func_name );
291+ if (timeout > 0 && !rinfo -> single_op_dag ) {
292+ RedisModule_ReplyWithError (ctx , "ERR TIMEOUT not allowed within a DAG command" );
293+ goto cleanup ;
294+ }
295+
296+ RAI_ScriptRunCtx * sctx = RAI_ScriptRunCtxCreate (script , func_name );
304297 sctx -> variadic = variadic ;
305298
306299 if (rinfo -> single_op_dag ) {
307300 rinfo -> timeout = timeout ;
308301 // Set params in ScriptRunCtx, bring inputs from key space.
309- if (_ScriptRunCtx_SetParams (ctx , inkeys , outkeys , sctx ) == REDISMODULE_ERR )
302+ if (_ScriptRunCtx_SetParams (ctx , currentOp -> inkeys , currentOp -> outkeys , sctx ) ==
303+ REDISMODULE_ERR )
310304 goto cleanup ;
311305 }
312- if (RAI_InitDagOp (& currentOp ) == REDISMODULE_ERR ) {
313- RedisModule_ReplyWithError (
314- ctx , "ERR Unable to allocate the memory and initialise the RAI_dagOp structure" );
315- goto cleanup ;
316- }
306+ currentOp -> sctx = sctx ;
317307 currentOp -> commandType = REDISAI_DAG_CMD_SCRIPTRUN ;
318- Dag_PopulateOp ( currentOp , sctx , inkeys , outkeys , runkey ) ;
319- rinfo -> dagOps = array_append ( rinfo -> dagOps , currentOp );
308+ currentOp -> devicestr = sctx -> script -> devicestr ;
309+
320310 return REDISMODULE_OK ;
321311
322312cleanup :
323- for (size_t i = 0 ; i < array_len (inkeys ); i ++ ) {
324- RedisModule_FreeString (NULL , inkeys [i ]);
325- }
326- array_free (inkeys );
327- for (size_t i = 0 ; i < array_len (outkeys ); i ++ ) {
328- RedisModule_FreeString (NULL , outkeys [i ]);
329- }
330- array_free (outkeys );
331- if (runkey )
332- RedisModule_FreeString (NULL , runkey );
333- if (sctx )
334- RAI_ScriptRunCtxFree (sctx );
335313 RAI_FreeRunInfo (rinfo );
336314 return REDISMODULE_ERR ;
337315}
0 commit comments