Merge branch 'master' into Create_test_module_for_LLAPI

Author: alonre24

src/dag.c

@@ -1313,6 +1313,11 @@ void DAG_ReplyAndUnblock(RedisAI_OnFinishCtx *ctx, void *private_data) {
 int RedisAI_ProcessDagRunCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc,
                                  int dagMode) {
 
+    int flags = RedisModule_GetContextFlags(ctx);
+    bool blocking_not_allowed = (flags & (REDISMODULE_CTX_FLAGS_MULTI | REDISMODULE_CTX_FLAGS_LUA));
+    if (blocking_not_allowed)
+        return RedisModule_ReplyWithError(
+            ctx, "ERR Cannot run RedisAI command within a transaction or a LUA script");
     RedisAI_RunInfo *rinfo = NULL;
     if (RAI_InitRunInfo(&rinfo) == REDISMODULE_ERR) {
         RedisModule_ReplyWithError(

src/redisai.c

@@ -1,3 +1,8 @@
+
+#ifdef __linux__
+#define _GNU_SOURCE
+#endif
+
 #define REDISMODULE_MAIN
 #include "redismodule.h"
 #include "tensor.h"
@@ -32,6 +37,12 @@
 #define REDISAI_GIT_SHA "unknown"
 #endif
 
+#ifdef __linux__
+#ifndef RUSAGE_THREAD
+#define RUSAGE_THREAD 1
+#endif
+#endif
+
 int redisMajorVersion;
 int redisMinorVersion;
 int redisPatchVersion;
@@ -986,10 +997,31 @@ void RAI_moduleInfoFunc(RedisModuleInfoCtx *ctx, int for_crash_report) {
     RedisModule_InfoAddFieldLongLong(ctx, "inter_op_parallelism", getBackendsInterOpParallelism());
     RedisModule_InfoAddFieldLongLong(ctx, "intra_op_parallelism", getBackendsIntraOpParallelism());
     struct rusage self_ru, c_ru;
+
     // Return resource usage statistics for the calling process,
     // which is the sum of resources used by all threads in the
     // process
     getrusage(RUSAGE_SELF, &self_ru);
+
+    // Return resource usage statistics for the calling thread
+    // which in this case is Redis/RedisAI main thread
+    // RUSAGE_THREAD is Linux-specific.
+    sds main_thread_used_cpu_sys = sdsempty();
+    sds main_thread_used_cpu_user = sdsempty();
+#if (defined(__linux__) && defined(RUSAGE_THREAD))
+    struct rusage main_thread_ru;
+    getrusage(RUSAGE_THREAD, &main_thread_ru);
+    main_thread_used_cpu_sys =
+        sdscatprintf(main_thread_used_cpu_sys, "%ld.%06ld", (long)main_thread_ru.ru_stime.tv_sec,
+                     (long)self_ru.ru_stime.tv_usec);
+    main_thread_used_cpu_user =
+        sdscatprintf(main_thread_used_cpu_user, "%ld.%06ld", (long)main_thread_ru.ru_utime.tv_sec,
+                     (long)self_ru.ru_utime.tv_usec);
+#else
+    sdscatprintf(main_thread_used_cpu_sys, "N/A");
+    sdscatprintf(main_thread_used_cpu_user, "N/A");
+#endif
+
     // Return resource usage statistics for all of its
     // terminated child processes
     getrusage(RUSAGE_CHILDREN, &c_ru);
@@ -1006,6 +1038,8 @@ void RAI_moduleInfoFunc(RedisModuleInfoCtx *ctx, int for_crash_report) {
     RedisModule_InfoAddFieldCString(ctx, "self_used_cpu_user", self_used_cpu_user);
     RedisModule_InfoAddFieldCString(ctx, "children_used_cpu_sys", children_used_cpu_sys);
     RedisModule_InfoAddFieldCString(ctx, "children_used_cpu_user", children_used_cpu_user);
+    RedisModule_InfoAddFieldCString(ctx, "main_thread_used_cpu_sys", main_thread_used_cpu_sys);
+    RedisModule_InfoAddFieldCString(ctx, "main_thread_used_cpu_user", main_thread_used_cpu_user);
 
     AI_dictIterator *iter = AI_dictGetSafeIterator(run_queues);
     AI_dictEntry *entry = AI_dictNext(iter);
@@ -1018,7 +1052,7 @@ void RAI_moduleInfoFunc(RedisModuleInfoCtx *ctx, int for_crash_report) {
                 struct timespec ts;
                 clockid_t cid;
                 sds queue_used_cpu_total = sdscatprintf(
-                    sdsempty(), "queue_%s_bthread_#%d_used_cpu_total", queue_name, i + 1);
+                    sdsempty(), "queue_%s_bthread_n%d_used_cpu_total", queue_name, i + 1);
                 sds bthread_used_cpu_total = sdsempty();
 #if (!defined(_POSIX_C_SOURCE) && !defined(_XOPEN_SOURCE)) || defined(_DARWIN_C_SOURCE) ||         \
     defined(__cplusplus)
@@ -1034,7 +1068,7 @@ void RAI_moduleInfoFunc(RedisModuleInfoCtx *ctx, int for_crash_report) {
                     } else {
                         bthread_used_cpu_total =
                             sdscatprintf(bthread_used_cpu_total, "%ld.%06ld", (long)ts.tv_sec,
-                                         (long)(ts.tv_nsec / 1000000));
+                                         (long)(ts.tv_nsec / 1000));
                     }
                 }
                 RedisModule_InfoAddFieldCString(ctx, queue_used_cpu_total, bthread_used_cpu_total);

tests/flow/tests_common.py

@@ -301,4 +301,24 @@ def test_info_modules(env):
     env.assertEqual( 'ai_self_used_cpu_user' in ret, True )
     env.assertEqual( 'ai_children_used_cpu_sys' in ret, True )
     env.assertEqual( 'ai_children_used_cpu_user' in ret, True )
-    env.assertEqual( 'ai_queue_CPU_bthread_#1_used_cpu_total' in ret, True )
+    env.assertEqual( 'ai_queue_CPU_bthread_n1_used_cpu_total' in ret, True )
+
+def test_lua_multi(env):
+    con = env.getConnection()
+    ret = con.execute_command('MULTI')
+    env.assertEqual(ret, b'OK')
+    ret = con.execute_command('AI.MODELRUN', "no_model{1}", "INPUTS", "no_input{1}", "OUTPUTS", "no_output{1}")
+    env.assertEqual(ret, b'QUEUED')
+    try:
+        ret = con.execute_command('EXEC')
+    except Exception as e:
+        exception = e
+        env.assertEqual(type(exception), redis.exceptions.ResponseError)
+        env.assertEqual("ERR Cannot run RedisAI command within a transaction or a LUA script", exception.__str__())
+    try:
+        ret = con.execute_command('EVAL', "return redis.pcall('AI.MODELRUN', 'no_model{1}', 'INPUTS', 'NO_INPUT{1}',"
+                                          " 'OUTPUTS', 'NO_OUTPUT{1}')", 0)
+    except Exception as e:
+        exception = e
+        env.assertEqual(type(exception), redis.exceptions.ResponseError)
+        env.assertEqual("Cannot run RedisAI command within a transaction or a LUA script", exception.__str__())